feat(spanner): add client-level CallContextConfigurator to SpannerOptions

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

@@ -332,6 +332,7 @@ static GcpChannelPoolOptions mergeWithDefaultChannelPoolOptions(
   private final Map<DatabaseId, QueryOptions> mergedQueryOptions;
 
   private final CallCredentialsProvider callCredentialsProvider;
+  private final CallContextConfigurator callContextConfigurator;
   private final CloseableExecutorProvider asyncExecutorProvider;
   private final String compressorName;
   private final String emulatorHost;
@@ -386,9 +387,11 @@ public interface CallCredentialsProvider {
 
   /**
    * {@link CallContextConfigurator} can be used to modify the {@link ApiCallContext} for one or
-   * more specific RPCs. This can be used to set specific timeout value for RPCs or use specific
-   * {@link CallCredentials} for an RPC. The {@link CallContextConfigurator} must be set as a value
-   * on the {@link Context} using the {@link SpannerOptions#CALL_CONTEXT_CONFIGURATOR_KEY} key.
+   * more specific RPCs. This can be used to set specific timeout values for RPCs or use specific
+   * {@link CallCredentials} for an RPC. The {@link CallContextConfigurator} can be configured at
+   * the client level using {@link Builder#setCallContextConfigurator(CallContextConfigurator)}, or
+   * on a per-call basis as a value on the {@link Context} using the {@link
+   * SpannerOptions#CALL_CONTEXT_CONFIGURATOR_KEY} key.
    *
    * <p>This API is meant for advanced users. Most users should instead use the {@link
    * SpannerCallContextTimeoutConfigurator} for setting timeouts per RPC.
@@ -517,8 +520,10 @@ static <ReqT, RespT> SpannerMethod valueOf(ReqT request, MethodDescriptor<ReqT,
 
   /**
    * Helper class to configure timeouts for specific Spanner RPCs. The {@link
-   * SpannerCallContextTimeoutConfigurator} must be set as a value on the {@link Context} using the
-   * {@link SpannerOptions#CALL_CONTEXT_CONFIGURATOR_KEY} key.
+   * SpannerCallContextTimeoutConfigurator} can be set client-wide via {@link
+   * Builder#setCallContextConfigurator(CallContextConfigurator)} or on individual requests as a
+   * value on the {@link Context} using the {@link SpannerOptions#CALL_CONTEXT_CONFIGURATOR_KEY}
+   * key.
    *
    * <p>Example usage:
    *
@@ -574,43 +579,40 @@ public <ReqT, RespT> ApiCallContext configure(
       if (spannerMethod == null) {
         return null;
       }
-      switch (SpannerMethod.valueOf(request, method)) {
+      ApiCallContext callContext = context == null ? GrpcCallContext.createDefault() : context;
+      switch (spannerMethod) {
         case BATCH_UPDATE:
           return batchUpdateTimeout == null
               ? null
-              : GrpcCallContext.createDefault().withTimeoutDuration(batchUpdateTimeout);
+              : callContext.withTimeoutDuration(batchUpdateTimeout);
         case COMMIT:
-          return commitTimeout == null
-              ? null
-              : GrpcCallContext.createDefault().withTimeoutDuration(commitTimeout);
+          return commitTimeout == null ? null : callContext.withTimeoutDuration(commitTimeout);
         case EXECUTE_QUERY:
           return executeQueryTimeout == null
               ? null
-              : GrpcCallContext.createDefault()
+              : callContext
                   .withTimeoutDuration(executeQueryTimeout)
                   .withStreamWaitTimeoutDuration(executeQueryTimeout);
         case EXECUTE_UPDATE:
           return executeUpdateTimeout == null
               ? null
-              : GrpcCallContext.createDefault().withTimeoutDuration(executeUpdateTimeout);
+              : callContext.withTimeoutDuration(executeUpdateTimeout);
         case PARTITION_QUERY:
           return partitionQueryTimeout == null
               ? null
-              : GrpcCallContext.createDefault().withTimeoutDuration(partitionQueryTimeout);
+              : callContext.withTimeoutDuration(partitionQueryTimeout);
         case PARTITION_READ:
           return partitionReadTimeout == null
               ? null
-              : GrpcCallContext.createDefault().withTimeoutDuration(partitionReadTimeout);
+              : callContext.withTimeoutDuration(partitionReadTimeout);
         case READ:
           return readTimeout == null
               ? null
-              : GrpcCallContext.createDefault()
+              : callContext
                   .withTimeoutDuration(readTimeout)
                   .withStreamWaitTimeoutDuration(readTimeout);
         case ROLLBACK:
-          return rollbackTimeout == null
-              ? null
-              : GrpcCallContext.createDefault().withTimeoutDuration(rollbackTimeout);
+          return rollbackTimeout == null ? null : callContext.withTimeoutDuration(rollbackTimeout);
         default:
       }
       return null;
@@ -1030,6 +1032,7 @@ protected SpannerOptions(Builder builder) {
       this.mergedQueryOptions = ImmutableMap.copyOf(merged);
     }
     callCredentialsProvider = builder.callCredentialsProvider;
+    callContextConfigurator = builder.callContextConfigurator;
     asyncExecutorProvider = builder.asyncExecutorProvider;
     compressorName = builder.compressorName;
     emulatorHost = builder.emulatorHost;
@@ -1379,6 +1382,7 @@ private static Builder prepareBuilder(Builder builder) {
     private Duration grpcKeepAliveTime = Duration.ofSeconds(120);
     private Duration grpcKeepAliveTimeout = Duration.ofSeconds(20);
     private CallCredentialsProvider callCredentialsProvider;
+    private CallContextConfigurator callContextConfigurator;
     private CloseableExecutorProvider asyncExecutorProvider;
     private String compressorName;
     private String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST");
@@ -1488,6 +1492,7 @@ protected Builder() {
       this.enableGrpcGcpOtelMetrics = options.enableGrpcGcpOtelMetrics;
       this.defaultQueryOptions = options.defaultQueryOptions;
       this.callCredentialsProvider = options.callCredentialsProvider;
+      this.callContextConfigurator = options.callContextConfigurator;
       this.grpcKeepAliveTime = options.grpcKeepAliveTime;
       this.grpcKeepAliveTimeout = options.grpcKeepAliveTimeout;
       this.asyncExecutorProvider = options.asyncExecutorProvider;
@@ -1874,6 +1879,84 @@ public Builder setCallCredentialsProvider(CallCredentialsProvider callCredential
       return this;
     }
 
+    /**
+     * Configures a client-level {@link CallContextConfigurator} to apply custom gRPC options,
+     * timeouts, or credentials to RPCs executed by this Spanner client.
+     *
+     * <p>By default, Spanner clients allow customizing call options on individual requests using
+     * gRPC's thread-local {@link io.grpc.Context} with {@link #CALL_CONTEXT_CONFIGURATOR_KEY}.
+     * While useful for fine-grained per-RPC overrides, managing thread-local context can be
+     * cumbersome or error-prone in asynchronous, reactive, or multi-threaded pipelines where
+     * operations jump across threads. Setting a {@link CallContextConfigurator} here applies
+     * client-wide across all requests executed by this client instance without requiring
+     * thread-local context propagation.
+     *
+     * <p>This configurator applies to all RPCs executed by {@link DatabaseClient}, {@link Spanner},
+     * {@link DatabaseAdminClient}, and {@link InstanceAdminClient} instances obtained from this
+     * client library. Note that raw GAPIC generated clients (such as {@link
+     * Spanner#createDatabaseAdminClient()} and {@link Spanner#createInstanceAdminClient()}) bypass
+     * this configurator and should be configured via {@link #setDatabaseAdminStubSettings} and
+     * {@link #setInstanceAdminStubSettings}.
+     *
+     * <p>Implementations of {@link CallContextConfigurator} configured at the client level must be
+     * thread-safe as they are shared across all concurrent operations executed by this client.
+     *
+     * <p>If both a client-level configurator and a thread-local configurator (via {@link
+     * #CALL_CONTEXT_CONFIGURATOR_KEY}) are present when an RPC is executed:
+     *
+     * <ol>
+     *   <li>The client-level configurator is evaluated first to establish the baseline call
+     *       context.
+     *   <li>The thread-local configurator is evaluated next using that baseline context.
+     *   <li>Any options returned by the thread-local configurator are merged on top of the
+     *       client-level options, allowing per-call configurations to override or extend
+     *       client-level defaults.
+     * </ol>
+     *
+     * <p>Example: Configure a client-level stream wait timeout of 30 seconds for streaming SQL
+     * queries to detect stalled streams faster:
+     *
+     * <pre>{@code
+     * SpannerOptions options =
+     *     SpannerOptions.newBuilder()
+     *         .setProjectId("my-project")
+     *         .setCallContextConfigurator(
+     *             new CallContextConfigurator() {
+     *               @Override
+     *               public <ReqT, RespT> ApiCallContext configure(
+     *                   ApiCallContext context, ReqT request, MethodDescriptor<ReqT, RespT> method) {
+     *                 if (method == SpannerGrpc.getExecuteStreamingSqlMethod()) {
+     *                   return context.withStreamWaitTimeoutDuration(Duration.ofSeconds(30));
+     *                 }
+     *                 return null;
+     *               }
+     *             })
+     *         .build();
+     * }</pre>
+     *
+     * <p>You can also use {@link SpannerCallContextTimeoutConfigurator} if you only need to adjust
+     * standard timeouts across RPC types:
+     *
+     * <pre>{@code
+     * SpannerOptions options =
+     *     SpannerOptions.newBuilder()
+     *         .setProjectId("my-project")
+     *         .setCallContextConfigurator(
+     *             SpannerCallContextTimeoutConfigurator.create()
+     *                 .withExecuteQueryTimeoutDuration(Duration.ofSeconds(30)))
+     *         .build();
+     * }</pre>
+     *
+     * @param callContextConfigurator the configurator to apply to all RPCs, or {@code null} to
+     *     clear
+     * @return this {@link Builder} instance
+     */
+    public Builder setCallContextConfigurator(
+        @Nullable CallContextConfigurator callContextConfigurator) {
+      this.callContextConfigurator = callContextConfigurator;
+      return this;
+    }
+
     /**
      * Sets the compression to use for all gRPC calls. The compressor must be a valid name known in
      * the {@link CompressorRegistry}. This will enable compression both from the client to the
@@ -2681,6 +2764,15 @@ public CallCredentialsProvider getCallCredentialsProvider() {
     return callCredentialsProvider;
   }
 
+  /**
+   * Returns the client-level {@link CallContextConfigurator} configured for this {@link
+   * SpannerOptions}, or {@code null} if none is set.
+   */
+  @Nullable
+  public CallContextConfigurator getCallContextConfigurator() {
+    return callContextConfigurator;
+  }
+
   private boolean usesNoCredentials() {
     // When JMH is enabled, we need to enable built-in metrics
     if (System.getProperty("jmh.enabled") != null

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -268,6 +268,9 @@ public class GapicSpannerRpc implements SpannerRpc {
 
   private static final String API_FILE = "grpc-gcp-apiconfig.json";
 
+  private static final CallOptions.Key<Boolean> BASE_CONTEXT_MARKER_KEY =
+      CallOptions.Key.create("BASE_CONTEXT_MARKER_KEY");
+
   private final RequestIdCreator requestIdCreator = new RequestIdCreatorImpl();
   private boolean rpcIsClosed;
   private final SpannerStub spannerStub;
@@ -286,6 +289,7 @@ public class GapicSpannerRpc implements SpannerRpc {
   private final String projectName;
   private final SpannerMetadataProvider metadataProvider;
   private final CallCredentialsProvider callCredentialsProvider;
+  private final CallContextConfigurator callContextConfigurator;
   private final String compressorName;
   private final Duration waitTimeout =
       systemProperty(PROPERTY_TIMEOUT_SECONDS, DEFAULT_TIMEOUT_SECONDS);
@@ -363,6 +367,7 @@ public GapicSpannerRpc(final SpannerOptions options) {
             headerProviderWithUserAgent.getHeaders(),
             internalHeaderProviderBuilder.getResourceHeaderKey());
     this.callCredentialsProvider = options.getCallCredentialsProvider();
+    this.callContextConfigurator = options.getCallContextConfigurator();
     this.compressorName = options.getCompressorName();
     this.leaderAwareRoutingEnabled = options.isLeaderAwareRoutingEnabled();
     this.endToEndTracingEnabled = options.isEndToEndTracingEnabled();
@@ -2111,7 +2116,7 @@ public StreamingCall read(
             requestId,
             request.getSession(),
             request,
-            SpannerGrpc.getReadMethod(),
+            SpannerGrpc.getStreamingReadMethod(),
             routeToLeader);
     SpannerResponseObserver responseObserver = new SpannerResponseObserver(consumer);
     spannerStub.streamingReadCallable().call(request, responseObserver, context);
@@ -2448,6 +2453,24 @@ <ReqT, RespT> GrpcCallContext newCallContext(
       MethodDescriptor<ReqT, RespT> method,
       boolean routeToLeader) {
     GrpcCallContext context = this.baseGrpcCallContext;
+    if (callCredentialsProvider != null) {
+      CallCredentials callCredentials = callCredentialsProvider.getCallCredentials();
+      if (callCredentials != null) {
+        context =
+            context.withCallOptions(context.getCallOptions().withCallCredentials(callCredentials));
+      }
+    }
+
+    // 1. Sequentially evaluate client-level and thread-level configurators.
+    // The thread-level configurator receives the context after client-level modifications,
+    // allowing thread-scoped settings to override or extend client defaults.
+    context = applyConfigurators(context, request, method);
+
+    // 2. Attach Spanner-internal routing options and headers to the final context.
+    // Doing this AFTER configurator evaluation guarantees that internal options (request ID,
+    // channel affinity) cannot be wiped out by configurators returning
+    // GrpcCallContext.createDefault(),
+    // and internal headers (resource prefix, route-to-leader) cannot be duplicated.
     Long affinity = options == null ? null : Option.CHANNEL_HINT.getLong(options);
     ChannelAffinityRef channelAffinityRef =
         options == null ? null : Option.CHANNEL_ID_AFFINITY.getChannelAffinityRef(options);
@@ -2489,19 +2512,78 @@ <ReqT, RespT> GrpcCallContext newCallContext(
     if (routeToLeader && leaderAwareRoutingEnabled) {
       context = context.withExtraHeaders(metadataProvider.newRouteToLeaderHeader());
     }
-    if (callCredentialsProvider != null) {
-      CallCredentials callCredentials = callCredentialsProvider.getCallCredentials();
-      if (callCredentials != null) {
-        context =
-            context.withCallOptions(context.getCallOptions().withCallCredentials(callCredentials));
-      }
+    if (compressorName != null && context.getCallOptions().getCompressor() == null) {
+      context = context.withCallOptions(context.getCallOptions().withCompression(compressorName));
     }
-    CallContextConfigurator configurator = SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY.get();
-    ApiCallContext apiCallContextFromContext = null;
-    if (configurator != null) {
-      apiCallContextFromContext = configurator.configure(context, request, method);
+    return context;
+  }
+
+  private <ReqT, RespT> GrpcCallContext applyConfigurators(
+      GrpcCallContext context, ReqT request, MethodDescriptor<ReqT, RespT> method) {
+    if (method == null) {
+      return context;
     }
-    return (GrpcCallContext) context.merge(apiCallContextFromContext);
+    CallContextConfigurator threadConfigurator = SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY.get();
+    if (this.callContextConfigurator == null && threadConfigurator == null) {
+      return context;
+    }
+    GrpcCallContext callContext =
+        context.withCallOptions(
+            context.getCallOptions().withOption(BASE_CONTEXT_MARKER_KEY, Boolean.TRUE));
+    if (this.callContextConfigurator != null) {
+      callContext =
+          applySingleConfigurator(callContext, this.callContextConfigurator, request, method);
+    }
+    if (threadConfigurator != null) {
+      callContext = applySingleConfigurator(callContext, threadConfigurator, request, method);
+    }
+    return callContext;
+  }
+
+  private static <ReqT, RespT> GrpcCallContext applySingleConfigurator(
+      GrpcCallContext base,
+      CallContextConfigurator configurator,
+      ReqT request,
+      MethodDescriptor<ReqT, RespT> method) {
+    ApiCallContext configured;
+    try {
+      configured = configurator.configure(base, request, method);
+    } catch (Throwable t) {
+      throw SpannerExceptionFactory.asSpannerException(t);
+    }
+    if (configured == null || configured == base) {
+      return base;
+    }
+    if (!(configured instanceof GrpcCallContext)) {
+      throw new IllegalArgumentException(
+          "context must be an instance of GrpcCallContext, but found "
+              + configured.getClass().getName());
+    }
+    GrpcCallContext overlay = (GrpcCallContext) configured;
+
+    // Check whether overlay was derived from base (retaining BASE_CONTEXT_MARKER_KEY)
+    // or is a standalone delta context (such as one created via GrpcCallContext.createDefault()).
+    boolean isDerived =
+        Boolean.TRUE.equals(overlay.getCallOptions().getOption(BASE_CONTEXT_MARKER_KEY));
+    if (isDerived) {
+      return overlay;
+    }
+
+    // Overlay is a standalone delta context. Merge it onto base.
+    GrpcCallContext merged = (GrpcCallContext) base.merge(overlay);
+
+    // If the delta context did not set custom CallOptions, GAX's merge would replace base's
+    // CallOptions with CallOptions.DEFAULT. In that case, preserve base's CallOptions.
+    if (overlay.getCallOptions().equals(CallOptions.DEFAULT)
+        && !base.getCallOptions().equals(CallOptions.DEFAULT)) {
+      merged = merged.withCallOptions(base.getCallOptions());
+    } else if (!Boolean.TRUE.equals(merged.getCallOptions().getOption(BASE_CONTEXT_MARKER_KEY))) {
+      merged =
+          merged.withCallOptions(
+              merged.getCallOptions().withOption(BASE_CONTEXT_MARKER_KEY, Boolean.TRUE));
+    }
+
+    return merged;
   }
 
   @Override