Skip to content
This repository was archived by the owner on Apr 7, 2026. It is now read-only.

Commit 4652157

Browse files
aseeringaseering
committed
feat: Add ClientContext support to Connection API
This change adds support for setting and propagating ClientContext in the Spanner Connection API. ClientContext allows propagating client-scoped session state (e.g., secure parameters) to Spanner RPCs. - Added setClientContext/getClientContext to Connection interface and implementation. - Implemented state propagation from Connection to UnitOfWork and its implementations (ReadWriteTransaction, SingleUseTransaction). - Fixed accidental import removal in OptionsTest.java. - Fixed TransactionRunnerImplTest to correctly verify ClientContext propagation. - Added ClientContextMockServerTest for end-to-end verification.
1 parent 8669fe0 commit 4652157

9 files changed

Lines changed: 451 additions & 13 deletions

File tree

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ abstract class AbstractBaseUnitOfWork implements UnitOfWork {
8080
protected final List<TransactionRetryListener> transactionRetryListeners;
8181
protected final boolean excludeTxnFromChangeStreams;
8282
protected final RpcPriority rpcPriority;
83+
protected final com.google.spanner.v1.RequestOptions.ClientContext clientContext;
8384
protected final Span span;
8485

8586
/** Class for keeping track of the stacktrace of the caller of an async statement. */
@@ -117,6 +118,7 @@ abstract static class Builder<B extends Builder<?, T>, T extends AbstractBaseUni
117118

118119
private boolean excludeTxnFromChangeStreams;
119120
private RpcPriority rpcPriority;
121+
private com.google.spanner.v1.RequestOptions.ClientContext clientContext;
120122
private Span span;
121123

122124
Builder() {}
@@ -163,6 +165,11 @@ B setRpcPriority(@Nullable RpcPriority rpcPriority) {
163165
return self();
164166
}
165167

168+
B setClientContext(@Nullable com.google.spanner.v1.RequestOptions.ClientContext clientContext) {
169+
this.clientContext = clientContext;
170+
return self();
171+
}
172+
166173
B setSpan(@Nullable Span span) {
167174
this.span = span;
168175
return self();
@@ -179,6 +186,7 @@ B setSpan(@Nullable Span span) {
179186
this.transactionRetryListeners = builder.transactionRetryListeners;
180187
this.excludeTxnFromChangeStreams = builder.excludeTxnFromChangeStreams;
181188
this.rpcPriority = builder.rpcPriority;
189+
this.clientContext = builder.clientContext;
182190
this.span = Preconditions.checkNotNull(builder.span);
183191
}
184192

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,25 @@ default String getStatementTag() {
449449
throw new UnsupportedOperationException();
450450
}
451451

452+
/**
453+
* Sets the client context to use for the statements that are executed. The client context
454+
* persists until it is changed or cleared.
455+
*
456+
* @param clientContext The client context to use with the statements that will be executed on
457+
* this connection.
458+
*/
459+
default void setClientContext(com.google.spanner.v1.RequestOptions.ClientContext clientContext) {
460+
throw new UnsupportedOperationException();
461+
}
462+
463+
/**
464+
* @return The client context that will be used with the statements that are executed on this
465+
* connection.
466+
*/
467+
default com.google.spanner.v1.RequestOptions.ClientContext getClientContext() {
468+
throw new UnsupportedOperationException();
469+
}
470+
452471
/**
453472
* Sets whether the next transaction should be excluded from all change streams with the DDL
454473
* option `allow_txn_exclusion=true`

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@
9494
import com.google.common.util.concurrent.MoreExecutors;
9595
import com.google.spanner.v1.DirectedReadOptions;
9696
import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
97+
import com.google.spanner.v1.RequestOptions;
9798
import com.google.spanner.v1.ResultSetStats;
9899
import com.google.spanner.v1.TransactionOptions.IsolationLevel;
99100
import com.google.spanner.v1.TransactionOptions.ReadWrite.ReadLockMode;
@@ -299,6 +300,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
299300
private IsolationLevel transactionIsolationLevel;
300301
private String transactionTag;
301302
private String statementTag;
303+
private RequestOptions.ClientContext clientContext;
302304
private boolean excludeTxnFromChangeStreams;
303305
private byte[] protoDescriptors;
304306
private String protoDescriptorsFilePath;
@@ -536,6 +538,7 @@ private void reset(Context context, boolean inTransaction) {
536538
this.connectionState.resetValue(SAVEPOINT_SUPPORT, context, inTransaction);
537539
this.protoDescriptors = null;
538540
this.protoDescriptorsFilePath = null;
541+
this.clientContext = null;
539542

540543
if (!isTransactionStarted()) {
541544
setDefaultTransactionOptions(getDefaultIsolationLevel());
@@ -955,6 +958,18 @@ public String getTransactionTag() {
955958
return transactionTag;
956959
}
957960

961+
@Override
962+
public void setClientContext(RequestOptions.ClientContext clientContext) {
963+
ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
964+
this.clientContext = clientContext;
965+
}
966+
967+
@Override
968+
public RequestOptions.ClientContext getClientContext() {
969+
ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
970+
return clientContext;
971+
}
972+
958973
@Override
959974
public void setTransactionTag(String tag) {
960975
ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
@@ -2026,6 +2041,9 @@ private QueryOption[] mergeQueryRequestOptions(
20262041
options =
20272042
appendQueryOption(options, Options.priority(getConnectionPropertyValue(RPC_PRIORITY)));
20282043
}
2044+
if (clientContext != null) {
2045+
options = appendQueryOption(options, Options.clientContext(clientContext));
2046+
}
20292047
if (currentUnitOfWork != null
20302048
&& currentUnitOfWork.supportsDirectedReads(parsedStatement)
20312049
&& getConnectionPropertyValue(DIRECTED_READ) != null) {
@@ -2070,6 +2088,14 @@ private UpdateOption[] mergeUpdateRequestOptions(UpdateOption... options) {
20702088
options[options.length - 1] = Options.priority(getConnectionPropertyValue(RPC_PRIORITY));
20712089
}
20722090
}
2091+
if (clientContext != null) {
2092+
if (options == null || options.length == 0) {
2093+
options = new UpdateOption[] {Options.clientContext(clientContext)};
2094+
} else {
2095+
options = Arrays.copyOf(options, options.length + 1);
2096+
options[options.length - 1] = Options.clientContext(clientContext);
2097+
}
2098+
}
20732099
return options;
20742100
}
20752101

@@ -2299,6 +2325,7 @@ UnitOfWork createNewUnitOfWork(
22992325
createSpanForUnitOfWork(
23002326
statementType == StatementType.DDL ? DDL_STATEMENT : SINGLE_USE_TRANSACTION))
23012327
.setProtoDescriptors(getProtoDescriptors())
2328+
.setClientContext(clientContext)
23022329
.build();
23032330
if (!isInternalMetadataQuery && !forceSingleUse) {
23042331
// Reset the transaction options after starting a single-use transaction.
@@ -2317,6 +2344,7 @@ UnitOfWork createNewUnitOfWork(
23172344
.setTransactionTag(transactionTag)
23182345
.setRpcPriority(getConnectionPropertyValue(RPC_PRIORITY))
23192346
.setSpan(createSpanForUnitOfWork(READ_ONLY_TRANSACTION))
2347+
.setClientContext(clientContext)
23202348
.build();
23212349
case READ_WRITE_TRANSACTION:
23222350
return ReadWriteTransaction.newBuilder()
@@ -2340,6 +2368,7 @@ UnitOfWork createNewUnitOfWork(
23402368
.setExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams)
23412369
.setRpcPriority(getConnectionPropertyValue(RPC_PRIORITY))
23422370
.setSpan(createSpanForUnitOfWork(READ_WRITE_TRANSACTION))
2371+
.setClientContext(clientContext)
23432372
.build();
23442373
case DML_BATCH:
23452374
// A DML batch can run inside the current transaction. It should therefore only
@@ -2359,6 +2388,7 @@ UnitOfWork createNewUnitOfWork(
23592388
.setRpcPriority(getConnectionPropertyValue(RPC_PRIORITY))
23602389
// Use the transaction Span for the DML batch.
23612390
.setSpan(transactionStack.peek().getSpan())
2391+
.setClientContext(clientContext)
23622392
.build();
23632393
case DDL_BATCH:
23642394
return DdlBatch.newBuilder()
@@ -2369,6 +2399,7 @@ UnitOfWork createNewUnitOfWork(
23692399
.setSpan(createSpanForUnitOfWork(DDL_BATCH))
23702400
.setProtoDescriptors(getProtoDescriptors())
23712401
.setConnectionState(connectionState)
2402+
.setClientContext(clientContext)
23722403
.build();
23732404
default:
23742405
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,9 @@ private TransactionOption[] extractOptions(Builder builder) {
350350
if (this.readLockMode != ReadLockMode.READ_LOCK_MODE_UNSPECIFIED) {
351351
numOptions++;
352352
}
353+
if (this.clientContext != null) {
354+
numOptions++;
355+
}
353356
TransactionOption[] options = new TransactionOption[numOptions];
354357
int index = 0;
355358
if (builder.returnCommitStats) {
@@ -373,6 +376,9 @@ private TransactionOption[] extractOptions(Builder builder) {
373376
if (this.readLockMode != ReadLockMode.READ_LOCK_MODE_UNSPECIFIED) {
374377
options[index++] = Options.readLockMode(this.readLockMode);
375378
}
379+
if (this.clientContext != null) {
380+
options[index++] = Options.clientContext(this.clientContext);
381+
}
376382
return options;
377383
}
378384

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -520,6 +520,9 @@ private TransactionRunner createWriteTransaction() {
520520
!= ReadLockMode.READ_LOCK_MODE_UNSPECIFIED) {
521521
numOptions++;
522522
}
523+
if (this.clientContext != null) {
524+
numOptions++;
525+
}
523526
if (numOptions == 0) {
524527
return dbClient.readWriteTransaction();
525528
}
@@ -547,6 +550,9 @@ private TransactionRunner createWriteTransaction() {
547550
!= ReadLockMode.READ_LOCK_MODE_UNSPECIFIED) {
548551
options[index++] = Options.readLockMode(connectionState.getValue(READ_LOCK_MODE).getValue());
549552
}
553+
if (this.clientContext != null) {
554+
options[index++] = Options.clientContext(this.clientContext);
555+
}
550556
return dbClient.readWriteTransaction(options);
551557
}
552558

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,13 @@
3535
import com.google.spanner.v1.ReadRequest.LockHint;
3636
import com.google.spanner.v1.ReadRequest.OrderBy;
3737
import com.google.spanner.v1.RequestOptions;
38+
import com.google.spanner.v1.RequestOptions.Priority;
39+
import com.google.spanner.v1.TransactionOptions.IsolationLevel;
40+
import com.google.spanner.v1.TransactionOptions.ReadWrite;
41+
import com.google.spanner.v1.TransactionOptions.ReadWrite.ReadLockMode;
42+
import org.junit.Test;
43+
import org.junit.runner.RunWith;
44+
import org.junit.runners.JUnit4;
3845

3946
/** Unit tests for {@link Options}. */
4047
@RunWith(JUnit4.class)

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -108,28 +108,28 @@ public void testCommitWithClientContext() {
108108
.putSecureContext(
109109
"key", com.google.protobuf.Value.newBuilder().setStringValue("value").build())
110110
.build();
111-
Options options =
112-
Options.fromTransactionOptions(
111+
when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
112+
when(session.newTransaction(any(Options.class), any())).thenReturn(txn);
113+
Mockito.clearInvocations(session);
114+
transactionRunner =
115+
new TransactionRunnerImpl(
116+
session,
113117
Options.priority(Options.RpcPriority.HIGH),
114118
Options.tag("tag"),
115119
Options.clientContext(clientContext));
116-
transactionRunner = new TransactionRunnerImpl(session, options);
117-
when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
118-
when(session.newTransaction(any(Options.class), any())).thenReturn(txn);
120+
transactionRunner.setSpan(span);
119121

120122
transactionRunner.run(
121123
transaction -> {
122124
return null;
123125
});
124126

125-
ArgumentCaptor<CommitRequest> commitRequestCaptor =
126-
ArgumentCaptor.forClass(CommitRequest.class);
127-
verify(rpc).commitAsync(commitRequestCaptor.capture(), anyMap());
128-
CommitRequest request = commitRequestCaptor.getValue();
129-
RequestOptions requestOptions = request.getRequestOptions();
130-
assertEquals(RequestOptions.Priority.PRIORITY_HIGH, requestOptions.getPriority());
131-
assertEquals("tag", requestOptions.getTransactionTag());
132-
assertEquals(clientContext, requestOptions.getClientContext());
127+
ArgumentCaptor<Options> optionsCaptor = ArgumentCaptor.forClass(Options.class);
128+
verify(session).newTransaction(optionsCaptor.capture(), any());
129+
Options capturedOptions = optionsCaptor.getValue();
130+
assertEquals(RequestOptions.Priority.PRIORITY_HIGH, capturedOptions.priority());
131+
assertEquals("tag", capturedOptions.tag());
132+
assertEquals(clientContext, capturedOptions.clientContext());
133133
}
134134

135135
@Mock private SpannerRpc rpc;

0 commit comments

Comments
 (0)