fix: using context from call in ReadRowsRetryingCallable (#2560)

isha97 · web-flow · commit eeb19b72629f · 2024-07-19T11:19:58.000-07:00
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/readrows/ReadRowsRetryingCallable.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/stub/readrows/ReadRowsRetryingCallable.java
@@ -55,15 +55,16 @@ public void call(
       ReadRowsRequest request,
       final ResponseObserver<ReadRowsResponse> responseObserver,
       ApiCallContext context) {
+    ApiCallContext actualContext = this.context.merge(context);
     ReadRowsAttemptCallable attemptCallable =
         new ReadRowsAttemptCallable(
             innerCallable,
             resumptionStrategyPrototype.createNew(),
             request,
-            this.context,
+            actualContext,
             responseObserver);
 
-    RetryingFuture<Void> retryingFuture = executor.createFuture(attemptCallable, this.context);
+    RetryingFuture<Void> retryingFuture = executor.createFuture(attemptCallable, actualContext);
     attemptCallable.setExternalFuture(retryingFuture);
     attemptCallable.start();
 
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/readrows/ReadRowsRetryingCallable.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/readrows/ReadRowsRetryingCallable.java
@@ -55,15 +55,16 @@ public void call(
       ReadRowsRequest request,
       final ResponseObserver<ReadRowsResponse> responseObserver,
       ApiCallContext context) {
+    ApiCallContext actualContext = this.context.merge(context);
     ReadRowsAttemptCallable attemptCallable =
         new ReadRowsAttemptCallable(
             innerCallable,
             resumptionStrategyPrototype.createNew(),
             request,
-            this.context,
+            actualContext,
             responseObserver);
 
-    RetryingFuture<Void> retryingFuture = executor.createFuture(attemptCallable, this.context);
+    RetryingFuture<Void> retryingFuture = executor.createFuture(attemptCallable, actualContext);
     attemptCallable.setExternalFuture(retryingFuture);
     attemptCallable.start();
 
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/stub/readrows/ReadRowsRetryingCallable.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/stub/readrows/ReadRowsRetryingCallable.java
@@ -55,15 +55,16 @@ public void call(
       ReadRowsRequest request,
       final ResponseObserver<ReadRowsResponse> responseObserver,
       ApiCallContext context) {
+    ApiCallContext actualContext = this.context.merge(context);
     ReadRowsAttemptCallable attemptCallable =
         new ReadRowsAttemptCallable(
             innerCallable,
             resumptionStrategyPrototype.createNew(),
             request,
-            this.context,
+            actualContext,
             responseObserver);
 
-    RetryingFuture<Void> retryingFuture = executor.createFuture(attemptCallable, this.context);
+    RetryingFuture<Void> retryingFuture = executor.createFuture(attemptCallable, actualContext);
     attemptCallable.setExternalFuture(retryingFuture);
     attemptCallable.start();
 
