From f7f6aec9aae74f963385da60ddd0cc0b0bd9628e Mon Sep 17 00:00:00 2001
From: Brent Shaffer <betterbrent@google.com>
Date: Thu, 20 Jun 2024 07:31:04 -0700
Subject: [PATCH] feat(Bigtable): add bigtable-attempt header (#7414)

---
 .../BigtableDataOperationException.php        |  2 +-
 Bigtable/src/ResumableStream.php              | 19 +++++++----
 Bigtable/tests/Unit/SmartRetriesTest.php      | 34 ++++++++++++++++++-
 3 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/Bigtable/src/Exception/BigtableDataOperationException.php b/Bigtable/src/Exception/BigtableDataOperationException.php
index bb717ac9c7ee..960f30caa8c3 100644
--- a/Bigtable/src/Exception/BigtableDataOperationException.php
+++ b/Bigtable/src/Exception/BigtableDataOperationException.php
@@ -41,7 +41,7 @@ public function __construct(
     ) {
         $this->metadata = $metadata;
 
-        parent::__construct($message, $code);
+        parent::__construct($message, $code ?? 0);
     }
 
     /**
diff --git a/Bigtable/src/ResumableStream.php b/Bigtable/src/ResumableStream.php
index 597a8fa391a0..fb344157a4b5 100644
--- a/Bigtable/src/ResumableStream.php
+++ b/Bigtable/src/ResumableStream.php
@@ -126,19 +126,25 @@ public function __construct(
      */
     public function readAll()
     {
-        $tries = 0;
-        $argumentFunction = $this->argumentFunction;
-        $retryFunction = $this->retryFunction;
+        $attempt = 0;
         do {
             $ex = null;
-            list($this->request, $this->callOptions) = $argumentFunction($this->request, $this->callOptions);
+            list($this->request, $this->callOptions) =
+                ($this->argumentFunction)($this->request, $this->callOptions);
 
             $completed = $this->pluck('requestCompleted', $this->callOptions, false);
 
             if ($completed !== true) {
+                // Send in "bigtable-attempt" header on retry request
+                $headers = $this->callOptions['headers'] ?? [];
+                if ($attempt > 0) {
+                    $headers['bigtable-attempt'] = [(string) $attempt];
+                }
+                $attempt++;
+
                 $stream = call_user_func_array(
                     [$this->gapicClient, $this->method],
-                    [$this->request, $this->callOptions]
+                    [$this->request, ['headers' => $headers] + $this->callOptions]
                 );
 
                 try {
@@ -148,8 +154,7 @@ public function readAll()
                 } catch (\Exception $ex) {
                 }
             }
-            $tries++;
-        } while ((!$this->retryFunction || $retryFunction($ex)) && $tries <= $this->retries);
+        } while ((!$this->retryFunction || ($this->retryFunction)($ex)) && $attempt <= $this->retries);
         if ($ex !== null) {
             throw $ex;
         }
diff --git a/Bigtable/tests/Unit/SmartRetriesTest.php b/Bigtable/tests/Unit/SmartRetriesTest.php
index f6c52136da17..c29accd4060a 100644
--- a/Bigtable/tests/Unit/SmartRetriesTest.php
+++ b/Bigtable/tests/Unit/SmartRetriesTest.php
@@ -136,6 +136,38 @@ public function testReadRowsShouldRetryForProvidedAttempts()
         $iterator->getIterator()->current();
     }
 
+    public function testReadRowsContainsAttemptHeader()
+    {
+        $attempt = 0;
+        $expectedArgs = $this->options;
+        $retryingApiException = $this->retryingApiException;
+        $this->serverStream->readAll()
+            ->shouldBeCalledTimes(2)
+            ->will(function () use (&$attempt, $retryingApiException) {
+                // throw a retriable exception on the first call
+                return 0 === $attempt++ ? throw $retryingApiException : [];
+            });
+        $this->bigtableClient->readRows(
+            Argument::type(ReadRowsRequest::class),
+            Argument::that(function ($callOptions) use (&$attempt) {
+                $attemptHeader = $callOptions['headers']['bigtable-attempt'][0] ?? null;
+                if ($attempt === 0) {
+                    $this->assertNull($attemptHeader);
+                } else {
+                    $this->assertSame((string) $attempt, $attemptHeader);
+                }
+
+                return true;
+            })
+        )->shouldBeCalledTimes(2)
+            ->willReturn(
+                $this->serverStream->reveal()
+            );
+
+        $iterator = $this->table->readRows();
+        $iterator->getIterator()->current();
+    }
+
     public function testReadRowsPartialSuccess()
     {
         $expectedArgs = $this->options;
@@ -186,7 +218,7 @@ public function testReadRowsWithRowsLimit()
                     $this->generateRowsResponse(3, 4)
                 )
             );
-        
+
         $allowedRowsLimit = ['5' => 1, '3' => 1];
         $this->bigtableClient->readRows(
             Argument::that(function ($request) use (&$allowedRowsLimit) {