From eeb19b72629f28fd740c61f60cc8856577f4b8a9 Mon Sep 17 00:00:00 2001
From: Isha Tarte <tarteisha97@gmail.com>
Date: Fri, 19 Jul 2024 11:19:58 -0700
Subject: [PATCH] fix: using context from call in ReadRowsRetryingCallable
 (#2560)

---
 .../storage/v1/stub/readrows/ReadRowsRetryingCallable.java   | 5 +++--
 .../v1beta1/stub/readrows/ReadRowsRetryingCallable.java      | 5 +++--
 .../v1beta2/stub/readrows/ReadRowsRetryingCallable.java      | 5 +++--
 3 files changed, 9 insertions(+), 6 deletions(-)

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
index d498766b3a..7009f39112 100644
--- 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
index 2b4308ae70..179f14c23e 100644
--- 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
index 8d3384de81..7095778691 100644
--- 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();