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

Commit 8fa9f39

Browse files
aseeringaseering
committed
feat: Add ClientContext to Options and propagate to RPCs
This change adds support for ClientContext in Options and ensures it is propagated to ExecuteSql, Read, Commit, and BeginTransaction requests. It aligns with go/spanner-client-scoped-session-state design. - Added RequestOptions.ClientContext to Options. - Refactored request option building to Options.toRequestOptionsProto. - Updated AbstractReadContext, TransactionRunnerImpl, and SessionImpl to use the shared logic. - Added tests.
1 parent 3ad105a commit 8fa9f39

9 files changed

Lines changed: 223 additions & 32 deletions

File tree

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java‎

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -684,22 +684,20 @@ QueryOptions buildQueryOptions(QueryOptions requestOptions) {
684684
}
685685

686686
RequestOptions buildRequestOptions(Options options) {
687-
// Shortcut for the most common return value.
688-
if (!(options.hasPriority() || options.hasTag() || getTransactionTag() != null)) {
689-
return RequestOptions.getDefaultInstance();
690-
}
691-
692-
RequestOptions.Builder builder = RequestOptions.newBuilder();
693-
if (options.hasPriority()) {
694-
builder.setPriority(options.priority());
695-
}
696-
if (options.hasTag()) {
697-
builder.setRequestTag(options.tag());
687+
RequestOptions requestOptions = options.toRequestOptionsProto(false);
688+
RequestOptions.ClientContext defaultClientContext =
689+
session.getSpanner().getOptions().getClientContext();
690+
if (defaultClientContext != null) {
691+
RequestOptions.ClientContext.Builder builder = defaultClientContext.toBuilder();
692+
if (requestOptions.hasClientContext()) {
693+
builder.mergeFrom(requestOptions.getClientContext());
694+
}
695+
requestOptions = requestOptions.toBuilder().setClientContext(builder.build()).build();
698696
}
699697
if (getTransactionTag() != null) {
700-
builder.setTransactionTag(getTransactionTag());
698+
return requestOptions.toBuilder().setTransactionTag(getTransactionTag()).build();
701699
}
702-
return builder.build();
700+
return requestOptions;
703701
}
704702

705703
ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.google.spanner.v1.DirectedReadOptions;
2121
import com.google.spanner.v1.ReadRequest.LockHint;
2222
import com.google.spanner.v1.ReadRequest.OrderBy;
23+
import com.google.spanner.v1.RequestOptions;
2324
import com.google.spanner.v1.RequestOptions.Priority;
2425
import com.google.spanner.v1.TransactionOptions.IsolationLevel;
2526
import com.google.spanner.v1.TransactionOptions.ReadWrite.ReadLockMode;
@@ -265,6 +266,37 @@ public static ReadQueryUpdateTransactionOption priority(RpcPriority priority) {
265266
return new PriorityOption(priority);
266267
}
267268

269+
/**
270+
* Specifying this will add the given client context to the request. The client context is used to
271+
* pass opaque side-channel information to the backend, such as a user ID for a parameterized
272+
* secure view.
273+
*/
274+
public static ReadQueryUpdateTransactionOption clientContext(
275+
RequestOptions.ClientContext clientContext) {
276+
return new ClientContextOption(clientContext);
277+
}
278+
279+
RequestOptions toRequestOptionsProto(boolean isTransactionOption) {
280+
if (!hasPriority() && !hasTag() && !hasClientContext()) {
281+
return RequestOptions.getDefaultInstance();
282+
}
283+
RequestOptions.Builder builder = RequestOptions.newBuilder();
284+
if (hasPriority()) {
285+
builder.setPriority(priority());
286+
}
287+
if (hasTag()) {
288+
if (isTransactionOption) {
289+
builder.setTransactionTag(tag());
290+
} else {
291+
builder.setRequestTag(tag());
292+
}
293+
}
294+
if (hasClientContext()) {
295+
builder.setClientContext(clientContext());
296+
}
297+
return builder.build();
298+
}
299+
268300
public static TransactionOption maxCommitDelay(Duration maxCommitDelay) {
269301
Preconditions.checkArgument(!maxCommitDelay.isNegative(), "maxCommitDelay should be positive");
270302
return new MaxCommitDelayOption(maxCommitDelay);
@@ -462,6 +494,20 @@ void appendToOptions(Options options) {
462494
}
463495
}
464496

497+
static final class ClientContextOption extends InternalOption
498+
implements ReadQueryUpdateTransactionOption {
499+
private final RequestOptions.ClientContext clientContext;
500+
501+
ClientContextOption(RequestOptions.ClientContext clientContext) {
502+
this.clientContext = clientContext;
503+
}
504+
505+
@Override
506+
void appendToOptions(Options options) {
507+
options.clientContext = clientContext;
508+
}
509+
}
510+
465511
static final class TagOption extends InternalOption implements ReadQueryUpdateTransactionOption {
466512
private final String tag;
467513

@@ -574,6 +620,7 @@ void appendToOptions(Options options) {
574620
private String filter;
575621
private RpcPriority priority;
576622
private String tag;
623+
private RequestOptions.ClientContext clientContext;
577624
private String etag;
578625
private Boolean validateOnly;
579626
private Boolean withExcludeTxnFromChangeStreams;
@@ -666,6 +713,14 @@ Priority priority() {
666713
return priority == null ? null : priority.proto;
667714
}
668715

716+
boolean hasClientContext() {
717+
return clientContext != null;
718+
}
719+
720+
RequestOptions.ClientContext clientContext() {
721+
return clientContext;
722+
}
723+
669724
boolean hasTag() {
670725
return tag != null;
671726
}
@@ -777,6 +832,9 @@ public String toString() {
777832
if (priority != null) {
778833
b.append("priority: ").append(priority).append(' ');
779834
}
835+
if (clientContext != null) {
836+
b.append("clientContext: ").append(clientContext).append(' ');
837+
}
780838
if (tag != null) {
781839
b.append("tag: ").append(tag).append(' ');
782840
}
@@ -850,6 +908,7 @@ public boolean equals(Object o) {
850908
&& Objects.equals(pageToken(), that.pageToken())
851909
&& Objects.equals(filter(), that.filter())
852910
&& Objects.equals(priority(), that.priority())
911+
&& Objects.equals(clientContext(), that.clientContext())
853912
&& Objects.equals(tag(), that.tag())
854913
&& Objects.equals(etag(), that.etag())
855914
&& Objects.equals(validateOnly(), that.validateOnly())
@@ -894,6 +953,9 @@ public int hashCode() {
894953
if (priority != null) {
895954
result = 31 * result + priority.hashCode();
896955
}
956+
if (clientContext != null) {
957+
result = 31 * result + clientContext.hashCode();
958+
}
897959
if (tag != null) {
898960
result = 31 * result + tag.hashCode();
899961
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,10 @@ ErrorHandler getErrorHandler() {
177177
return this.errorHandler;
178178
}
179179

180+
SpannerImpl getSpanner() {
181+
return spanner;
182+
}
183+
180184
void setCurrentSpan(ISpan span) {
181185
currentSpan = span;
182186
}
@@ -481,9 +485,21 @@ ApiFuture<Transaction> beginTransactionAsync(
481485
if (sessionReference.getIsMultiplexed() && mutation != null) {
482486
requestBuilder.setMutationKey(mutation);
483487
}
484-
if (sessionReference.getIsMultiplexed() && !Strings.isNullOrEmpty(transactionOptions.tag())) {
485-
requestBuilder.setRequestOptions(
486-
RequestOptions.newBuilder().setTransactionTag(transactionOptions.tag()).build());
488+
RequestOptions requestOptions = transactionOptions.toRequestOptionsProto(true);
489+
RequestOptions.ClientContext defaultClientContext =
490+
spanner.getOptions().getClientContext();
491+
if (defaultClientContext != null) {
492+
RequestOptions.ClientContext.Builder builder = defaultClientContext.toBuilder();
493+
if (requestOptions.hasClientContext()) {
494+
builder.mergeFrom(requestOptions.getClientContext());
495+
}
496+
requestOptions = requestOptions.toBuilder().setClientContext(builder.build()).build();
497+
}
498+
if (!sessionReference.getIsMultiplexed()) {
499+
requestOptions = requestOptions.toBuilder().clearTransactionTag().build();
500+
}
501+
if (!requestOptions.equals(RequestOptions.getDefaultInstance())) {
502+
requestBuilder.setRequestOptions(requestOptions);
487503
}
488504
final BeginTransactionRequest request = requestBuilder.build();
489505
final ApiFuture<Transaction> requestFuture;

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@
6767
import com.google.spanner.v1.DirectedReadOptions;
6868
import com.google.spanner.v1.ExecuteSqlRequest;
6969
import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
70+
import com.google.spanner.v1.RequestOptions;
7071
import com.google.spanner.v1.SpannerGrpc;
7172
import com.google.spanner.v1.TransactionOptions;
7273
import com.google.spanner.v1.TransactionOptions.IsolationLevel;
@@ -257,6 +258,7 @@ public static GcpChannelPoolOptions createDefaultDynamicChannelPoolOptions() {
257258
private final boolean enableEndToEndTracing;
258259
private final String monitoringHost;
259260
private final TransactionOptions defaultTransactionOptions;
261+
private final RequestOptions.ClientContext clientContext;
260262

261263
enum TracingFramework {
262264
OPEN_CENSUS,
@@ -922,13 +924,19 @@ protected SpannerOptions(Builder builder) {
922924
enableEndToEndTracing = builder.enableEndToEndTracing;
923925
monitoringHost = builder.monitoringHost;
924926
defaultTransactionOptions = builder.defaultTransactionOptions;
927+
clientContext = builder.clientContext;
925928
}
926929

927930
private String getResolvedUniverseDomain() {
928931
String universeDomain = getUniverseDomain();
929932
return Strings.isNullOrEmpty(universeDomain) ? GOOGLE_DEFAULT_UNIVERSE : universeDomain;
930933
}
931934

935+
/** Returns the default {@link RequestOptions.ClientContext} for this {@link SpannerOptions}. */
936+
public RequestOptions.ClientContext getClientContext() {
937+
return clientContext;
938+
}
939+
932940
/**
933941
* The environment to read configuration values from. The default implementation uses environment
934942
* variables.
@@ -1142,6 +1150,7 @@ public static class Builder
11421150
private String experimentalHost = null;
11431151
private boolean usePlainText = false;
11441152
private TransactionOptions defaultTransactionOptions = TransactionOptions.getDefaultInstance();
1153+
private RequestOptions.ClientContext clientContext;
11451154

11461155
private static String createCustomClientLibToken(String token) {
11471156
return token + " " + ServiceOptions.getGoogApiClientLibName();
@@ -1243,6 +1252,7 @@ protected Builder() {
12431252
this.enableEndToEndTracing = options.enableEndToEndTracing;
12441253
this.monitoringHost = options.monitoringHost;
12451254
this.defaultTransactionOptions = options.defaultTransactionOptions;
1255+
this.clientContext = options.clientContext;
12461256
}
12471257

12481258
@Override
@@ -1977,6 +1987,12 @@ public Builder setDefaultTransactionOptions(
19771987
return this;
19781988
}
19791989

1990+
/** Sets the default {@link RequestOptions.ClientContext} for all requests. */
1991+
public Builder setClientContext(RequestOptions.ClientContext clientContext) {
1992+
this.clientContext = clientContext;
1993+
return this;
1994+
}
1995+
19801996
@SuppressWarnings("rawtypes")
19811997
@Override
19821998
public SpannerOptions build() {

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -464,15 +464,9 @@ public void run() {
464464
waitForTransactionTimeoutMillis, TimeUnit.MILLISECONDS)
465465
: transactionId);
466466
}
467-
if (options.hasPriority() || getTransactionTag() != null) {
468-
RequestOptions.Builder requestOptionsBuilder = RequestOptions.newBuilder();
469-
if (options.hasPriority()) {
470-
requestOptionsBuilder.setPriority(options.priority());
471-
}
472-
if (getTransactionTag() != null) {
473-
requestOptionsBuilder.setTransactionTag(getTransactionTag());
474-
}
475-
requestBuilder.setRequestOptions(requestOptionsBuilder.build());
467+
RequestOptions requestOptions = options.toRequestOptionsProto(true);
468+
if (!requestOptions.equals(RequestOptions.getDefaultInstance())) {
469+
requestBuilder.setRequestOptions(requestOptions);
476470
}
477471
if (session.getIsMultiplexed() && getLatestPrecommitToken() != null) {
478472
// Set the precommit token in the CommitRequest for multiplexed sessions.

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,18 @@ public void executeSqlRequestBuilderWithRequestOptionsWithTxnTag() {
345345
assertThat(request.getRequestOptions().getTransactionTag()).isEqualTo("app=spanner,env=test");
346346
}
347347

348+
@Test
349+
public void testBuildRequestOptionsWithClientContext() {
350+
RequestOptions.ClientContext clientContext =
351+
RequestOptions.ClientContext.newBuilder()
352+
.putSecureContext(
353+
"key", com.google.protobuf.Value.newBuilder().setStringValue("value").build())
354+
.build();
355+
RequestOptions requestOptions =
356+
context.buildRequestOptions(Options.fromQueryOptions(Options.clientContext(clientContext)));
357+
assertEquals(clientContext, requestOptions.getClientContext());
358+
}
359+
348360
@Test
349361
public void testGetExecuteSqlRequestBuilderWithDirectedReadOptions() {
350362
ExecuteSqlRequest.Builder request =

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

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -34,17 +34,37 @@
3434
import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection;
3535
import com.google.spanner.v1.ReadRequest.LockHint;
3636
import com.google.spanner.v1.ReadRequest.OrderBy;
37-
import com.google.spanner.v1.RequestOptions.Priority;
38-
import com.google.spanner.v1.TransactionOptions.IsolationLevel;
39-
import com.google.spanner.v1.TransactionOptions.ReadWrite;
40-
import com.google.spanner.v1.TransactionOptions.ReadWrite.ReadLockMode;
41-
import org.junit.Test;
42-
import org.junit.runner.RunWith;
43-
import org.junit.runners.JUnit4;
37+
import com.google.spanner.v1.RequestOptions;
4438

4539
/** Unit tests for {@link Options}. */
4640
@RunWith(JUnit4.class)
4741
public class OptionsTest {
42+
@Test
43+
public void testToRequestOptionsProto() {
44+
RequestOptions.ClientContext clientContext =
45+
RequestOptions.ClientContext.newBuilder()
46+
.putSecureContext(
47+
"key", com.google.protobuf.Value.newBuilder().setStringValue("value").build())
48+
.build();
49+
Options options =
50+
Options.fromQueryOptions(
51+
Options.priority(RpcPriority.HIGH),
52+
Options.tag("tag"),
53+
Options.clientContext(clientContext));
54+
55+
RequestOptions protoForStatement = options.toRequestOptionsProto(false);
56+
assertEquals(RequestOptions.Priority.PRIORITY_HIGH, protoForStatement.getPriority());
57+
assertEquals("tag", protoForStatement.getRequestTag());
58+
assertEquals("", protoForStatement.getTransactionTag());
59+
assertEquals(clientContext, protoForStatement.getClientContext());
60+
61+
RequestOptions protoForTransaction = options.toRequestOptionsProto(true);
62+
assertEquals(RequestOptions.Priority.PRIORITY_HIGH, protoForTransaction.getPriority());
63+
assertEquals("", protoForTransaction.getRequestTag());
64+
assertEquals("tag", protoForTransaction.getTransactionTag());
65+
assertEquals(clientContext, protoForTransaction.getClientContext());
66+
}
67+
4868
private static final DirectedReadOptions DIRECTED_READ_OPTIONS =
4969
DirectedReadOptions.newBuilder()
5070
.setIncludeReplicas(

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,10 @@
4747
import com.google.spanner.v1.CommitResponse;
4848
import com.google.spanner.v1.Mutation.Write;
4949
import com.google.spanner.v1.PartialResultSet;
50+
import com.google.spanner.v1.RequestOptions;
51+
import com.google.spanner.v1.ResultSet;
5052
import com.google.spanner.v1.ResultSetMetadata;
53+
import com.google.spanner.v1.ResultSetStats;
5154
import com.google.spanner.v1.RollbackRequest;
5255
import com.google.spanner.v1.Session;
5356
import com.google.spanner.v1.Transaction;
@@ -77,6 +80,42 @@
7780
/** Unit tests for {@link com.google.cloud.spanner.SessionImpl}. */
7881
@RunWith(JUnit4.class)
7982
public class SessionImplTest {
83+
@Test
84+
public void testBeginTransactionWithClientContext() {
85+
RequestOptions.ClientContext clientContext =
86+
RequestOptions.ClientContext.newBuilder()
87+
.putSecureContext(
88+
"key", com.google.protobuf.Value.newBuilder().setStringValue("value").build())
89+
.build();
90+
Mockito.when(
91+
rpc.beginTransactionAsync(
92+
Mockito.any(BeginTransactionRequest.class), anyMap(), eq(true)))
93+
.thenReturn(
94+
ApiFutures.immediateFuture(
95+
Transaction.newBuilder().setId(ByteString.copyFromUtf8("tx")).build()));
96+
97+
((SessionImpl) session)
98+
.beginTransactionAsync(
99+
Options.fromTransactionOptions(
100+
Options.priority(Options.RpcPriority.HIGH),
101+
Options.tag("tag"),
102+
Options.clientContext(clientContext)),
103+
true,
104+
Collections.emptyMap(),
105+
null,
106+
null);
107+
108+
ArgumentCaptor<BeginTransactionRequest> requestCaptor =
109+
ArgumentCaptor.forClass(BeginTransactionRequest.class);
110+
Mockito.verify(rpc).beginTransactionAsync(requestCaptor.capture(), anyMap(), eq(true));
111+
BeginTransactionRequest request = requestCaptor.getValue();
112+
RequestOptions requestOptions = request.getRequestOptions();
113+
assertEquals(RequestOptions.Priority.PRIORITY_HIGH, requestOptions.getPriority());
114+
// TransactionTag should NOT be set because session is not multiplexed.
115+
assertEquals("", requestOptions.getTransactionTag());
116+
assertEquals(clientContext, requestOptions.getClientContext());
117+
}
118+
80119
@Mock private SpannerRpc rpc;
81120
@Mock private SpannerOptions spannerOptions;
82121
private com.google.cloud.spanner.Session session;

0 commit comments

Comments
 (0)