Skip to content

Commit 86658ba

Browse files
aseeringaseering
authored
feat: ClientContext and secure parameters support (#4316)
* 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. * 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. * Address code-review comments
1 parent e8bc9ec commit 86658ba

17 files changed

+701
-25
lines changed

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

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -693,17 +693,15 @@ QueryOptions buildQueryOptions(QueryOptions requestOptions) {
693693
}
694694

695695
RequestOptions buildRequestOptions(Options options) {
696-
// Shortcut for the most common return value.
697-
if (!(options.hasPriority() || options.hasTag() || getTransactionTag() != null)) {
698-
return RequestOptions.getDefaultInstance();
699-
}
700-
701-
RequestOptions.Builder builder = RequestOptions.newBuilder();
702-
if (options.hasPriority()) {
703-
builder.setPriority(options.priority());
704-
}
705-
if (options.hasTag()) {
706-
builder.setRequestTag(options.tag());
696+
RequestOptions.Builder builder = options.toRequestOptionsProto(false).toBuilder();
697+
RequestOptions.ClientContext defaultClientContext =
698+
session.getSpanner().getOptions().getClientContext();
699+
if (defaultClientContext != null) {
700+
RequestOptions.ClientContext.Builder clientContextBuilder = defaultClientContext.toBuilder();
701+
if (builder.hasClientContext()) {
702+
clientContextBuilder.mergeFrom(builder.getClientContext());
703+
}
704+
builder.setClientContext(clientContextBuilder.build());
707705
}
708706
if (getTransactionTag() != null) {
709707
builder.setTransactionTag(getTransactionTag());

‎java-spanner/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
}

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

Lines changed: 20 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,22 @@ 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.Builder optionsBuilder =
493+
transactionOptions.toRequestOptionsProto(true).toBuilder();
494+
RequestOptions.ClientContext defaultClientContext = spanner.getOptions().getClientContext();
495+
if (defaultClientContext != null) {
496+
RequestOptions.ClientContext.Builder builder = defaultClientContext.toBuilder();
497+
if (optionsBuilder.hasClientContext()) {
498+
builder.mergeFrom(optionsBuilder.getClientContext());
499+
}
500+
optionsBuilder.setClientContext(builder.build());
501+
}
502+
if (!sessionReference.getIsMultiplexed()) {
503+
optionsBuilder.clearTransactionTag();
504+
}
505+
RequestOptions requestOptions = optionsBuilder.build();
506+
if (!requestOptions.equals(RequestOptions.getDefaultInstance())) {
507+
requestBuilder.setRequestOptions(requestOptions);
492508
}
493509
final BeginTransactionRequest request = requestBuilder.build();
494510
final ApiFuture<Transaction> requestFuture;

‎java-spanner/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
@@ -68,6 +68,7 @@
6868
import com.google.spanner.v1.DirectedReadOptions;
6969
import com.google.spanner.v1.ExecuteSqlRequest;
7070
import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
71+
import com.google.spanner.v1.RequestOptions;
7172
import com.google.spanner.v1.SpannerGrpc;
7273
import com.google.spanner.v1.TransactionOptions;
7374
import com.google.spanner.v1.TransactionOptions.IsolationLevel;
@@ -260,6 +261,7 @@ public static GcpChannelPoolOptions createDefaultDynamicChannelPoolOptions() {
260261
private final boolean enableEndToEndTracing;
261262
private final String monitoringHost;
262263
private final TransactionOptions defaultTransactionOptions;
264+
private final RequestOptions.ClientContext clientContext;
263265

264266
enum TracingFramework {
265267
OPEN_CENSUS,
@@ -927,13 +929,19 @@ protected SpannerOptions(Builder builder) {
927929
enableEndToEndTracing = builder.enableEndToEndTracing;
928930
monitoringHost = builder.monitoringHost;
929931
defaultTransactionOptions = builder.defaultTransactionOptions;
932+
clientContext = builder.clientContext;
930933
}
931934

932935
private String getResolvedUniverseDomain() {
933936
String universeDomain = getUniverseDomain();
934937
return Strings.isNullOrEmpty(universeDomain) ? GOOGLE_DEFAULT_UNIVERSE : universeDomain;
935938
}
936939

940+
/** Returns the default {@link RequestOptions.ClientContext} for this {@link SpannerOptions}. */
941+
public RequestOptions.ClientContext getClientContext() {
942+
return clientContext;
943+
}
944+
937945
/**
938946
* The environment to read configuration values from. The default implementation uses environment
939947
* variables.
@@ -1161,6 +1169,7 @@ public static class Builder
11611169
private String experimentalHost = null;
11621170
private boolean usePlainText = false;
11631171
private TransactionOptions defaultTransactionOptions = TransactionOptions.getDefaultInstance();
1172+
private RequestOptions.ClientContext clientContext;
11641173

11651174
private static String createCustomClientLibToken(String token) {
11661175
return token + " " + ServiceOptions.getGoogApiClientLibName();
@@ -1264,6 +1273,7 @@ protected Builder() {
12641273
this.enableEndToEndTracing = options.enableEndToEndTracing;
12651274
this.monitoringHost = options.monitoringHost;
12661275
this.defaultTransactionOptions = options.defaultTransactionOptions;
1276+
this.clientContext = options.clientContext;
12671277
}
12681278

12691279
@Override
@@ -2016,6 +2026,12 @@ public Builder setDefaultTransactionOptions(
20162026
return this;
20172027
}
20182028

2029+
/** Sets the default {@link RequestOptions.ClientContext} for all requests. */
2030+
public Builder setDefaultClientContext(RequestOptions.ClientContext clientContext) {
2031+
this.clientContext = clientContext;
2032+
return this;
2033+
}
2034+
20192035
@SuppressWarnings("rawtypes")
20202036
@Override
20212037
public SpannerOptions build() {

‎java-spanner/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.

‎java-spanner/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

‎java-spanner/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`

0 commit comments

Comments
 (0)