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

Commit 8669fe0

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 9263839 commit 8669fe0

10 files changed

Lines changed: 239 additions & 33 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 side-channel or configuration information to the backend, such as a user ID for a
272+
* parameterized 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: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@
3232
import com.google.cloud.spanner.SessionClient.SessionOption;
3333
import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl;
3434
import com.google.cloud.spanner.spi.v1.SpannerRpc;
35-
import com.google.common.base.Strings;
3635
import com.google.common.base.Ticker;
3736
import com.google.common.collect.Lists;
3837
import com.google.common.util.concurrent.MoreExecutors;
@@ -182,6 +181,10 @@ ErrorHandler getErrorHandler() {
182181
return this.errorHandler;
183182
}
184183

184+
SpannerImpl getSpanner() {
185+
return spanner;
186+
}
187+
185188
void setCurrentSpan(ISpan span) {
186189
currentSpan = span;
187190
}
@@ -486,9 +489,20 @@ ApiFuture<Transaction> beginTransactionAsync(
486489
if (sessionReference.getIsMultiplexed() && mutation != null) {
487490
requestBuilder.setMutationKey(mutation);
488491
}
489-
if (sessionReference.getIsMultiplexed() && !Strings.isNullOrEmpty(transactionOptions.tag())) {
490-
requestBuilder.setRequestOptions(
491-
RequestOptions.newBuilder().setTransactionTag(transactionOptions.tag()).build());
492+
RequestOptions requestOptions = transactionOptions.toRequestOptionsProto(true);
493+
RequestOptions.ClientContext defaultClientContext = spanner.getOptions().getClientContext();
494+
if (defaultClientContext != null) {
495+
RequestOptions.ClientContext.Builder builder = defaultClientContext.toBuilder();
496+
if (requestOptions.hasClientContext()) {
497+
builder.mergeFrom(requestOptions.getClientContext());
498+
}
499+
requestOptions = requestOptions.toBuilder().setClientContext(builder.build()).build();
500+
}
501+
if (!sessionReference.getIsMultiplexed()) {
502+
requestOptions = requestOptions.toBuilder().clearTransactionTag().build();
503+
}
504+
if (!requestOptions.equals(RequestOptions.getDefaultInstance())) {
505+
requestBuilder.setRequestOptions(requestOptions);
492506
}
493507
final BeginTransactionRequest request = requestBuilder.build();
494508
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: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,10 @@ String getTransactionTag() {
138138
public void setup() {
139139
SessionImpl session = mock(SessionImpl.class);
140140
when(session.getName()).thenReturn("session-1");
141+
SpannerImpl spanner = mock(SpannerImpl.class);
142+
SpannerOptions spannerOptions = mock(SpannerOptions.class);
143+
when(spanner.getOptions()).thenReturn(spannerOptions);
144+
when(session.getSpanner()).thenReturn(spanner);
141145
TestReadContextBuilder builder = new TestReadContextBuilder();
142146
context =
143147
builder
@@ -322,6 +326,10 @@ public void executeSqlRequestBuilderWithRequestOptions() {
322326
public void executeSqlRequestBuilderWithRequestOptionsWithTxnTag() {
323327
SessionImpl session = mock(SessionImpl.class);
324328
when(session.getName()).thenReturn("session-1");
329+
SpannerImpl spanner = mock(SpannerImpl.class);
330+
SpannerOptions spannerOptions = mock(SpannerOptions.class);
331+
when(spanner.getOptions()).thenReturn(spannerOptions);
332+
when(session.getSpanner()).thenReturn(spanner);
325333
TestReadContextWithTagBuilder builder = new TestReadContextWithTagBuilder();
326334
TestReadContextWithTag contextWithTag =
327335
builder
@@ -345,6 +353,18 @@ public void executeSqlRequestBuilderWithRequestOptionsWithTxnTag() {
345353
assertThat(request.getRequestOptions().getTransactionTag()).isEqualTo("app=spanner,env=test");
346354
}
347355

356+
@Test
357+
public void testBuildRequestOptionsWithClientContext() {
358+
RequestOptions.ClientContext clientContext =
359+
RequestOptions.ClientContext.newBuilder()
360+
.putSecureContext(
361+
"key", com.google.protobuf.Value.newBuilder().setStringValue("value").build())
362+
.build();
363+
RequestOptions requestOptions =
364+
context.buildRequestOptions(Options.fromQueryOptions(Options.clientContext(clientContext)));
365+
assertEquals(clientContext, requestOptions.getClientContext());
366+
}
367+
348368
@Test
349369
public void testGetExecuteSqlRequestBuilderWithDirectedReadOptions() {
350370
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(

0 commit comments

Comments
 (0)