From bbf3c14b4d45534c6acced8381607de7995b9500 Mon Sep 17 00:00:00 2001
From: Mattie Fu <mattiefu@google.com>
Date: Tue, 30 Dec 2025 12:29:44 -0500
Subject: [PATCH 1/2] fix: create stub with BigtableClientContext so otels are
 closed

---
 .../clirr-ignored-differences.xml             |   7 ++
 .../bigtable/data/v2/BigtableDataClient.java  |   4 +-
 .../data/v2/BigtableDataClientFactory.java    |  12 +-
 .../data/v2/stub/BigtableClientContext.java   |   5 +
 .../data/v2/stub/EnhancedBigtableStub.java    | 111 +++++++++++-------
 .../metrics/BigtableTracerCallableTest.java   |   8 +-
 .../metrics/BuiltinMetricsTracerTest.java     |   4 +-
 .../v2/stub/metrics/MetricsTracerTest.java    |   4 +-
 8 files changed, 104 insertions(+), 51 deletions(-)

diff --git a/google-cloud-bigtable/clirr-ignored-differences.xml b/google-cloud-bigtable/clirr-ignored-differences.xml
index c3a0fa05e186..c9a6f3762cc9 100644
--- a/google-cloud-bigtable/clirr-ignored-differences.xml
+++ b/google-cloud-bigtable/clirr-ignored-differences.xml
@@ -34,6 +34,13 @@
         <className>com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub</className>
         <method>*</method>
     </difference>
+    <difference>
+        <!-- change method types is ok because EnhancedBigtableStub is InternalApi -->
+        <differenceType>7005</differenceType>
+        <className>com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub</className>
+        <method>*</method>
+        <to>*</to>
+    </difference>
     <difference>
         <!-- method name change is ok because EnhancedBigtableStub is InternalApi -->
         <differenceType>7002</differenceType>
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java
index 889d36e383e5..cef5e58f3a29 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java
@@ -25,7 +25,6 @@
 import com.google.api.gax.batching.Batcher;
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.rpc.ApiExceptions;
-import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStream;
 import com.google.api.gax.rpc.ServerStreamingCallable;
@@ -55,6 +54,7 @@
 import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement;
 import com.google.cloud.bigtable.data.v2.models.sql.ResultSet;
 import com.google.cloud.bigtable.data.v2.models.sql.SqlType;
+import com.google.cloud.bigtable.data.v2.stub.BigtableClientContext;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
 import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream;
 import com.google.common.util.concurrent.MoreExecutors;
@@ -186,7 +186,7 @@ public static BigtableDataClient create(BigtableDataSettings settings) throws IO
    * BigtableDataClientFactory#close()} is called.
    */
   static BigtableDataClient createWithClientContext(
-      BigtableDataSettings settings, ClientContext context) throws IOException {
+      BigtableDataSettings settings, BigtableClientContext context) throws IOException {
     EnhancedBigtableStub stub =
         EnhancedBigtableStub.createWithClientContext(settings.getStubSettings(), context);
     return new BigtableDataClient(stub);
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java
index cddea20c7d1e..5628d58417d1 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java
@@ -114,7 +114,8 @@ public BigtableDataClient createDefault() {
                       defaultSettings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
               .build();
 
-      return BigtableDataClient.createWithClientContext(defaultSettings, clientContext);
+      return BigtableDataClient.createWithClientContext(
+          defaultSettings, sharedClientContext.setClientContext(clientContext));
     } catch (IOException e) {
       // Should never happen because the connection has been established already
       throw new RuntimeException(
@@ -141,7 +142,8 @@ public BigtableDataClient createForAppProfile(@Nonnull String appProfileId) thro
                 EnhancedBigtableStub.createBigtableTracerFactory(
                     settings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
             .build();
-    return BigtableDataClient.createWithClientContext(settings, clientContext);
+    return BigtableDataClient.createWithClientContext(
+        settings, sharedClientContext.setClientContext(clientContext));
   }
 
   /**
@@ -169,7 +171,8 @@ public BigtableDataClient createForInstance(@Nonnull String projectId, @Nonnull
                     settings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
             .build();
 
-    return BigtableDataClient.createWithClientContext(settings, clientContext);
+    return BigtableDataClient.createWithClientContext(
+        settings, sharedClientContext.setClientContext(clientContext));
   }
 
   /**
@@ -196,6 +199,7 @@ public BigtableDataClient createForInstance(
                 EnhancedBigtableStub.createBigtableTracerFactory(
                     settings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
             .build();
-    return BigtableDataClient.createWithClientContext(settings, clientContext);
+    return BigtableDataClient.createWithClientContext(
+        settings, sharedClientContext.setClientContext(clientContext));
   }
 }
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java
index bac1ec4a06d5..f67ca1b8c02b 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java
@@ -197,6 +197,11 @@ public ClientContext getClientContext() {
     return this.clientContext;
   }
 
+  public BigtableClientContext setClientContext(ClientContext clientContext) {
+    return new BigtableClientContext(
+        clientContext, openTelemetry, internalOpenTelemetry, metricsProvider);
+  }
+
   public void close() throws Exception {
     for (BackgroundResource resource : clientContext.getBackgroundResources()) {
       resource.close();
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 5f6b69dea89e..1231c4e6dc05 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -27,7 +27,6 @@
 import com.google.api.gax.batching.Batcher;
 import com.google.api.gax.batching.BatcherImpl;
 import com.google.api.gax.batching.FlowController;
-import com.google.api.gax.core.BackgroundResource;
 import com.google.api.gax.grpc.GaxGrpcProperties;
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.grpc.GrpcCallSettings;
@@ -170,7 +169,7 @@ public class EnhancedBigtableStub implements AutoCloseable {
   private static final String CLIENT_NAME = "Bigtable";
   private static final long FLOW_CONTROL_ADJUSTING_INTERVAL_MS = TimeUnit.SECONDS.toMillis(20);
   private final EnhancedBigtableStubSettings settings;
-  private final ClientContext clientContext;
+  private final BigtableClientContext bigtableClientContext;
 
   private final boolean closeClientContext;
   private final RequestContext requestContext;
@@ -209,11 +208,13 @@ public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
         bigtableClientContext.getClientContext().toBuilder()
             .setTracerFactory(createBigtableTracerFactory(settings, openTelemetry))
             .build();
-    return new EnhancedBigtableStub(settings, contextWithTracer);
+    bigtableClientContext = bigtableClientContext.setClientContext(contextWithTracer);
+    return new EnhancedBigtableStub(settings, bigtableClientContext);
   }
 
   public static EnhancedBigtableStub createWithClientContext(
-      EnhancedBigtableStubSettings settings, ClientContext clientContext) throws IOException {
+      EnhancedBigtableStubSettings settings, BigtableClientContext clientContext)
+      throws IOException {
 
     return new EnhancedBigtableStub(settings, clientContext, false);
   }
@@ -289,16 +290,17 @@ static Attributes createBuiltinAttributes(EnhancedBigtableStubSettings settings)
         "bigtable-java/" + Version.VERSION);
   }
 
-  public EnhancedBigtableStub(EnhancedBigtableStubSettings settings, ClientContext clientContext) {
+  public EnhancedBigtableStub(
+      EnhancedBigtableStubSettings settings, BigtableClientContext clientContext) {
     this(settings, clientContext, true);
   }
 
   public EnhancedBigtableStub(
       EnhancedBigtableStubSettings settings,
-      ClientContext clientContext,
+      BigtableClientContext clientContext,
       boolean closeClientContext) {
     this.settings = settings;
-    this.clientContext = clientContext;
+    this.bigtableClientContext = clientContext;
     this.closeClientContext = closeClientContext;
     this.requestContext =
         RequestContext.create(
@@ -347,7 +349,7 @@ public EnhancedBigtableStub(
   public <RowT> ServerStreamingCallable<ReadRowsRequest, RowT> createReadRowsRawCallable(
       RowAdapter<RowT> rowAdapter) {
     return createReadRowsBaseCallable(settings.readRowsSettings(), rowAdapter)
-        .withDefaultCallContext(clientContext.getDefaultCallContext());
+        .withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext());
   }
 
   /**
@@ -375,10 +377,13 @@ public <RowT> ServerStreamingCallable<Query, RowT> createReadRowsCallable(
     SpanName span = getSpanName("ReadRows");
     ServerStreamingCallable<Query, RowT> traced =
         new TracedServerStreamingCallable<>(
-            readRowsUserCallable, clientContext.getTracerFactory(), span);
+            readRowsUserCallable,
+            bigtableClientContext.getClientContext().getTracerFactory(),
+            span);
 
     return traced.withDefaultCallContext(
-        clientContext
+        bigtableClientContext
+            .getClientContext()
             .getDefaultCallContext()
             .withRetrySettings(settings.readRowsSettings().getRetrySettings()));
   }
@@ -398,6 +403,7 @@ public <RowT> ServerStreamingCallable<Query, RowT> createReadRowsCallable(
    * </ul>
    */
   public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT> rowAdapter) {
+    ClientContext clientContext = bigtableClientContext.getClientContext();
     if (!settings.getEnableSkipTrailers()) {
       ServerStreamingCallable<ReadRowsRequest, RowT> readRowsCallable =
           createReadRowsBaseCallable(
@@ -506,7 +512,7 @@ private <ReqT, RowT> ServerStreamingCallable<ReadRowsRequest, RowT> createReadRo
             .build();
 
     ServerStreamingCallable<ReadRowsRequest, RowT> watched =
-        Callables.watched(merging, innerSettings, clientContext);
+        Callables.watched(merging, innerSettings, bigtableClientContext.getClientContext());
 
     ServerStreamingCallable<ReadRowsRequest, RowT> withBigtableTracer =
         new BigtableTracerStreamingCallable<>(watched);
@@ -582,7 +588,7 @@ public <ReqT, RowT> ServerStreamingCallable<Query, RowT> createSkipLargeRowsCall
             .build();
 
     ServerStreamingCallable<ReadRowsRequest, RowT> watched =
-        Callables.watched(merging, innerSettings, clientContext);
+        Callables.watched(merging, innerSettings, bigtableClientContext.getClientContext());
 
     ServerStreamingCallable<ReadRowsRequest, RowT> withBigtableTracer =
         new BigtableTracerStreamingCallable<>(watched);
@@ -604,10 +610,13 @@ public <ReqT, RowT> ServerStreamingCallable<Query, RowT> createSkipLargeRowsCall
     SpanName span = getSpanName("ReadRows");
     ServerStreamingCallable<Query, RowT> traced =
         new TracedServerStreamingCallable<>(
-            readRowsUserCallable, clientContext.getTracerFactory(), span);
+            readRowsUserCallable,
+            bigtableClientContext.getClientContext().getTracerFactory(),
+            span);
 
     return traced.withDefaultCallContext(
-        clientContext
+        bigtableClientContext
+            .getClientContext()
             .getDefaultCallContext()
             .withRetrySettings(readRowsSettings.getRetrySettings()));
   }
@@ -644,10 +653,12 @@ private <RowT> UnaryCallable<Query, List<RowT>> createBulkReadRowsCallable(
         new TracedBatcherUnaryCallable<>(readRowsUserCallable.all());
 
     UnaryCallable<Query, List<RowT>> traced =
-        new TracedUnaryCallable<>(tracedBatcher, clientContext.getTracerFactory(), span);
+        new TracedUnaryCallable<>(
+            tracedBatcher, bigtableClientContext.getClientContext().getTracerFactory(), span);
 
     return traced.withDefaultCallContext(
-        clientContext
+        bigtableClientContext
+            .getClientContext()
             .getDefaultCallContext()
             .withRetrySettings(settings.readRowsSettings().getRetrySettings()));
   }
@@ -717,7 +728,8 @@ public ApiFuture<List<KeyOffset>> futureCall(String s, ApiCallContext apiCallCon
         methodName,
         new SampleRowKeysCallableWithRequest(retryable, requestContext)
             .withDefaultCallContext(
-                clientContext
+                bigtableClientContext
+                    .getClientContext()
                     .getDefaultCallContext()
                     .withRetrySettings(settings.sampleRowKeysSettings().getRetrySettings())));
   }
@@ -760,6 +772,7 @@ private UnaryCallable<RowMutation, Void> createMutateRowCallable() {
    * <p>This function should not be exposed to external users, as it could cause a data loss.
    */
   private UnaryCallable<BulkMutation, MutateRowsAttemptResult> createMutateRowsBaseCallable() {
+    ClientContext clientContext = bigtableClientContext.getClientContext();
     ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> base =
         GrpcRawCallableFactory.createServerStreamingCallable(
             GrpcCallSettings.<MutateRowsRequest, MutateRowsResponse>newBuilder()
@@ -874,9 +887,10 @@ public Batcher<RowMutationEntry, Void> newMutateRowsBatcher(
         bulkMutateRowsCallable,
         BulkMutation.create(tableId),
         settings.bulkMutateRowsSettings().getBatchingSettings(),
-        clientContext.getExecutor(),
+        bigtableClientContext.getClientContext().getExecutor(),
         bulkMutationFlowController,
-        MoreObjects.firstNonNull(ctx, clientContext.getDefaultCallContext()));
+        MoreObjects.firstNonNull(
+            ctx, bigtableClientContext.getClientContext().getDefaultCallContext()));
   }
 
   /**
@@ -905,9 +919,10 @@ public Batcher<RowMutationEntry, Void> newMutateRowsBatcher(
         bulkMutateRowsCallable,
         BulkMutation.create(targetId),
         settings.bulkMutateRowsSettings().getBatchingSettings(),
-        clientContext.getExecutor(),
+        bigtableClientContext.getClientContext().getExecutor(),
         bulkMutationFlowController,
-        MoreObjects.firstNonNull(ctx, clientContext.getDefaultCallContext()));
+        MoreObjects.firstNonNull(
+            ctx, bigtableClientContext.getClientContext().getDefaultCallContext()));
   }
 
   /**
@@ -933,9 +948,10 @@ public Batcher<ByteString, Row> newBulkReadRowsBatcher(
         bulkReadRowsCallable,
         query,
         settings.bulkReadRowsSettings().getBatchingSettings(),
-        clientContext.getExecutor(),
+        bigtableClientContext.getClientContext().getExecutor(),
         null,
-        MoreObjects.firstNonNull(ctx, clientContext.getDefaultCallContext()));
+        MoreObjects.firstNonNull(
+            ctx, bigtableClientContext.getClientContext().getDefaultCallContext()));
   }
 
   /**
@@ -995,6 +1011,7 @@ private UnaryCallable<ReadModifyWriteRow, Row> createReadModifyWriteRowCallable(
    */
   private ServerStreamingCallable<String, ByteStringRange>
       createGenerateInitialChangeStreamPartitionsCallable() {
+    ClientContext clientContext = bigtableClientContext.getClientContext();
     ServerStreamingCallable<
             GenerateInitialChangeStreamPartitionsRequest,
             GenerateInitialChangeStreamPartitionsResponse>
@@ -1075,6 +1092,7 @@ private UnaryCallable<ReadModifyWriteRow, Row> createReadModifyWriteRowCallable(
       ServerStreamingCallable<ReadChangeStreamQuery, ChangeStreamRecordT>
           createReadChangeStreamCallable(
               ChangeStreamRecordAdapter<ChangeStreamRecordT> changeStreamRecordAdapter) {
+    ClientContext clientContext = bigtableClientContext.getClientContext();
     ServerStreamingCallable<ReadChangeStreamRequest, ReadChangeStreamResponse> base =
         GrpcRawCallableFactory.createServerStreamingCallable(
             GrpcCallSettings.<ReadChangeStreamRequest, ReadChangeStreamResponse>newBuilder()
@@ -1151,6 +1169,7 @@ private UnaryCallable<ReadModifyWriteRow, Row> createReadModifyWriteRowCallable(
    */
   @InternalApi("For internal use only")
   public ExecuteQueryCallable createExecuteQueryCallable() {
+    ClientContext clientContext = bigtableClientContext.getClientContext();
     ServerStreamingCallable<ExecuteQueryRequest, ExecuteQueryResponse> base =
         GrpcRawCallableFactory.createServerStreamingCallable(
             GrpcCallSettings.<ExecuteQueryRequest, ExecuteQueryResponse>newBuilder()
@@ -1245,9 +1264,13 @@ private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUserFacin
       String methodName, UnaryCallable<RequestT, ResponseT> inner) {
 
     UnaryCallable<RequestT, ResponseT> traced =
-        new TracedUnaryCallable<>(inner, clientContext.getTracerFactory(), getSpanName(methodName));
+        new TracedUnaryCallable<>(
+            inner,
+            bigtableClientContext.getClientContext().getTracerFactory(),
+            getSpanName(methodName));
 
-    return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
+    return traced.withDefaultCallContext(
+        bigtableClientContext.getClientContext().getDefaultCallContext());
   }
 
   private Map<String, String> composeRequestParams(
@@ -1314,11 +1337,14 @@ public ApiFuture<RespT> futureCall(ReqT reqT, ApiCallContext apiCallContext) {
     UnaryCallable<ReqT, RespT> traced =
         new TracedUnaryCallable<>(
             transformed,
-            clientContext.getTracerFactory(),
+            bigtableClientContext.getClientContext().getTracerFactory(),
             getSpanName(methodDescriptor.getBareMethodName()));
 
     return traced.withDefaultCallContext(
-        clientContext.getDefaultCallContext().withRetrySettings(callSettings.getRetrySettings()));
+        bigtableClientContext
+            .getClientContext()
+            .getDefaultCallContext()
+            .withRetrySettings(callSettings.getRetrySettings()));
   }
 
   private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnaryCallableNew(
@@ -1347,8 +1373,11 @@ private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnar
 
     return new BigtableUnaryOperationCallable<>(
         transformed,
-        clientContext.getDefaultCallContext().withRetrySettings(callSettings.getRetrySettings()),
-        clientContext.getTracerFactory(),
+        bigtableClientContext
+            .getClientContext()
+            .getDefaultCallContext()
+            .withRetrySettings(callSettings.getRetrySettings()),
+        bigtableClientContext.getClientContext().getTracerFactory(),
         getSpanName(methodDescriptor.getBareMethodName()),
         /* allowNoResponse= */ false);
   }
@@ -1371,9 +1400,11 @@ private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> withRetries(
     if (settings.getEnableRetryInfo()) {
       retrying =
           com.google.cloud.bigtable.gaxx.retrying.Callables.retrying(
-              innerCallable, unaryCallSettings, clientContext);
+              innerCallable, unaryCallSettings, bigtableClientContext.getClientContext());
     } else {
-      retrying = Callables.retrying(innerCallable, unaryCallSettings, clientContext);
+      retrying =
+          Callables.retrying(
+              innerCallable, unaryCallSettings, bigtableClientContext.getClientContext());
     }
     if (settings.getEnableRoutingCookie()) {
       return new CookiesUnaryCallable<>(retrying);
@@ -1389,9 +1420,11 @@ private <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT> withR
     if (settings.getEnableRetryInfo()) {
       retrying =
           com.google.cloud.bigtable.gaxx.retrying.Callables.retrying(
-              innerCallable, serverStreamingCallSettings, clientContext);
+              innerCallable, serverStreamingCallSettings, bigtableClientContext.getClientContext());
     } else {
-      retrying = Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext);
+      retrying =
+          Callables.retrying(
+              innerCallable, serverStreamingCallSettings, bigtableClientContext.getClientContext());
     }
     if (settings.getEnableRoutingCookie()) {
       return new CookiesServerStreamingCallable<>(retrying);
@@ -1408,7 +1441,7 @@ private <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT> large
     ServerStreamingCallable<RequestT, ResponseT> retrying;
     retrying =
         com.google.cloud.bigtable.gaxx.retrying.Callables.retryingForLargeRows(
-            innerCallable, serverStreamingCallSettings, clientContext);
+            innerCallable, serverStreamingCallSettings, bigtableClientContext.getClientContext());
     if (settings.getEnableRoutingCookie()) {
       return new CookiesServerStreamingCallable<>(retrying);
     }
@@ -1507,12 +1540,10 @@ private SpanName getSpanName(String methodName) {
   @Override
   public void close() {
     if (closeClientContext) {
-      for (BackgroundResource backgroundResource : clientContext.getBackgroundResources()) {
-        try {
-          backgroundResource.close();
-        } catch (Exception e) {
-          throw new IllegalStateException("Failed to close resource", e);
-        }
+      try {
+        bigtableClientContext.close();
+      } catch (Exception e) {
+        throw new IllegalStateException("failed to close client context", e);
       }
     }
   }
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java
index b0966a216649..1a09951b77aa 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java
@@ -139,7 +139,9 @@ public void sendHeaders(Metadata headers) {
                     null))
             .build();
     attempts = settings.getStubSettings().readRowsSettings().getRetrySettings().getMaxAttempts();
-    stub = new EnhancedBigtableStub(settings.getStubSettings(), clientContext);
+    stub =
+        new EnhancedBigtableStub(
+            settings.getStubSettings(), bigtableClientContext.setClientContext(clientContext));
 
     // Create another server without injecting the server-timing header and another stub that
     // connects to it.
@@ -164,7 +166,9 @@ public void sendHeaders(Metadata headers) {
                     null))
             .build();
     noHeaderStub =
-        new EnhancedBigtableStub(noHeaderSettings.getStubSettings(), noHeaderClientContext);
+        new EnhancedBigtableStub(
+            noHeaderSettings.getStubSettings(),
+            noHeaderBigtableClientContext.setClientContext(noHeaderClientContext));
   }
 
   @After
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java
index a59cf8475135..864a801e6ee1 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java
@@ -50,7 +50,6 @@
 import com.google.api.gax.batching.BatchingSettings;
 import com.google.api.gax.batching.FlowControlSettings;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
-import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.NotFoundException;
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.StreamController;
@@ -71,6 +70,7 @@
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
 import com.google.cloud.bigtable.data.v2.models.RowMutationEntry;
 import com.google.cloud.bigtable.data.v2.models.TableId;
+import com.google.cloud.bigtable.data.v2.stub.BigtableClientContext;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
 import com.google.common.base.Stopwatch;
@@ -276,7 +276,7 @@ public void sendHeaders(Metadata headers) {
         });
     stubSettingsBuilder.setTransportChannelProvider(channelProvider.build());
     EnhancedBigtableStubSettings stubSettings = stubSettingsBuilder.build();
-    stub = new EnhancedBigtableStub(stubSettings, ClientContext.create(stubSettings));
+    stub = new EnhancedBigtableStub(stubSettings, BigtableClientContext.create(stubSettings));
   }
 
   @After
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java
index a9f3aa038ba8..8446f80f8736 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java
@@ -133,7 +133,9 @@ public void setUp() throws Exception {
                     localStats.getStatsRecorder(),
                     null))
             .build();
-    stub = new EnhancedBigtableStub(settings.getStubSettings(), clientContext);
+    stub =
+        new EnhancedBigtableStub(
+            settings.getStubSettings(), bigtableClientContext.setClientContext(clientContext));
   }
 
   @After

From f9b81491434cad505e624413c421a3176ed937b4 Mon Sep 17 00:00:00 2001
From: Mattie Fu <mattiefu@google.com>
Date: Tue, 6 Jan 2026 11:41:38 -0500
Subject: [PATCH 2/2] address comment

---
 .../cloud/bigtable/data/v2/BigtableDataClientFactory.java | 8 ++++----
 .../bigtable/data/v2/stub/BigtableClientContext.java      | 2 +-
 .../cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java | 2 +-
 .../data/v2/stub/metrics/BigtableTracerCallableTest.java  | 4 ++--
 .../bigtable/data/v2/stub/metrics/MetricsTracerTest.java  | 2 +-
 5 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java
index 5628d58417d1..599dce9f31af 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java
@@ -115,7 +115,7 @@ public BigtableDataClient createDefault() {
               .build();
 
       return BigtableDataClient.createWithClientContext(
-          defaultSettings, sharedClientContext.setClientContext(clientContext));
+          defaultSettings, sharedClientContext.withClientContext(clientContext));
     } catch (IOException e) {
       // Should never happen because the connection has been established already
       throw new RuntimeException(
@@ -143,7 +143,7 @@ public BigtableDataClient createForAppProfile(@Nonnull String appProfileId) thro
                     settings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
             .build();
     return BigtableDataClient.createWithClientContext(
-        settings, sharedClientContext.setClientContext(clientContext));
+        settings, sharedClientContext.withClientContext(clientContext));
   }
 
   /**
@@ -172,7 +172,7 @@ public BigtableDataClient createForInstance(@Nonnull String projectId, @Nonnull
             .build();
 
     return BigtableDataClient.createWithClientContext(
-        settings, sharedClientContext.setClientContext(clientContext));
+        settings, sharedClientContext.withClientContext(clientContext));
   }
 
   /**
@@ -200,6 +200,6 @@ public BigtableDataClient createForInstance(
                     settings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
             .build();
     return BigtableDataClient.createWithClientContext(
-        settings, sharedClientContext.setClientContext(clientContext));
+        settings, sharedClientContext.withClientContext(clientContext));
   }
 }
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java
index f67ca1b8c02b..c26f16b305fc 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java
@@ -197,7 +197,7 @@ public ClientContext getClientContext() {
     return this.clientContext;
   }
 
-  public BigtableClientContext setClientContext(ClientContext clientContext) {
+  public BigtableClientContext withClientContext(ClientContext clientContext) {
     return new BigtableClientContext(
         clientContext, openTelemetry, internalOpenTelemetry, metricsProvider);
   }
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 1231c4e6dc05..79cfee703a98 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -208,7 +208,7 @@ public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
         bigtableClientContext.getClientContext().toBuilder()
             .setTracerFactory(createBigtableTracerFactory(settings, openTelemetry))
             .build();
-    bigtableClientContext = bigtableClientContext.setClientContext(contextWithTracer);
+    bigtableClientContext = bigtableClientContext.withClientContext(contextWithTracer);
     return new EnhancedBigtableStub(settings, bigtableClientContext);
   }
 
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java
index 1a09951b77aa..8c3746144fa1 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java
@@ -141,7 +141,7 @@ public void sendHeaders(Metadata headers) {
     attempts = settings.getStubSettings().readRowsSettings().getRetrySettings().getMaxAttempts();
     stub =
         new EnhancedBigtableStub(
-            settings.getStubSettings(), bigtableClientContext.setClientContext(clientContext));
+            settings.getStubSettings(), bigtableClientContext.withClientContext(clientContext));
 
     // Create another server without injecting the server-timing header and another stub that
     // connects to it.
@@ -168,7 +168,7 @@ public void sendHeaders(Metadata headers) {
     noHeaderStub =
         new EnhancedBigtableStub(
             noHeaderSettings.getStubSettings(),
-            noHeaderBigtableClientContext.setClientContext(noHeaderClientContext));
+            noHeaderBigtableClientContext.withClientContext(noHeaderClientContext));
   }
 
   @After
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java
index 8446f80f8736..5c4161d0e342 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java
@@ -135,7 +135,7 @@ public void setUp() throws Exception {
             .build();
     stub =
         new EnhancedBigtableStub(
-            settings.getStubSettings(), bigtableClientContext.setClientContext(clientContext));
+            settings.getStubSettings(), bigtableClientContext.withClientContext(clientContext));
   }
 
   @After
