fix(Spanner): retry logic of ExecuteStreamingSql (#6149)

googleapis · May 3, 2023 · fbd16a8 · fbd16a8
1 parent a3af7ed
commit fbd16a8
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 1 deletion.
diff --git a/Spanner/src/Result.php b/Spanner/src/Result.php
@@ -178,6 +178,7 @@ public function rows($format = self::RETURN_ASSOCIATIVE)
         $call = $this->call;
         $generator = null;
         $shouldRetry = false;
+        $isResultsYielded = false;
         $backoff = new ExponentialBackoff($this->retries, function ($ex) {
             if ($ex instanceof ServiceException) {
                 return $ex->getCode() === Grpc\STATUS_UNAVAILABLE;
@@ -209,13 +210,13 @@ public function rows($format = self::RETURN_ASSOCIATIVE)
                         list($yieldableRows, $chunkedResult) = $this->parseRowsFromBufferedResults($bufferedResults);
 
                         foreach ($yieldableRows as $row) {
+                            $isResultsYielded = true;
                             yield $this->mapper->decodeValues($this->columns, $row, $format);
                         }
                     }
 
                     // Now that we've yielded all available rows, flush the buffer.
                     $bufferedResults = [];
-                    $shouldRetry = $hasResumeToken;
 
                     // If the last item in the buffer had a chunked value let's
                     // hold on to it so we can stitch it together into a yieldable
@@ -225,6 +226,8 @@ public function rows($format = self::RETURN_ASSOCIATIVE)
                     }
                 }
 
+                // retry without resume token when results have not yielded
+                $shouldRetry = !$isResultsYielded || $hasResumeToken;
                 $generator->next();
                 $valid = $generator->valid();
             } catch (ServiceException $ex) {

diff --git a/Spanner/tests/Unit/ResultTest.php b/Spanner/tests/Unit/ResultTest.php
@@ -159,6 +159,77 @@ function () use ($chunks, &$timesCalled) {
         $this->assertEquals(2, $timesCalled);
     }
 
+    public function testRowsRetriesWithoutResumeTokenWhenNotYieldedRows()
+    {
+        $timesCalled = 0;
+        $chunks = [
+            [
+                'metadata' => $this->metadata,
+                'values' => ['a']
+            ],
+            [
+                'values' => ['b']
+            ],
+            [
+                'values' => ['c']
+            ]
+        ];
+
+        $result = $this->getResultClass(
+            null,
+            'r',
+            null,
+            function () use ($chunks, &$timesCalled) {
+                $timesCalled++;
+                foreach ($chunks as $key => $chunk) {
+                    if ($key === 1 && $timesCalled < 2) {
+                        throw new ServiceException('Unavailable', 14);
+                    }
+                    yield $chunk;
+                }
+            }
+        );
+
+        iterator_to_array($result->rows());
+        $this->assertEquals(2, $timesCalled);
+    }
+
+    public function testRowsRetriesWithResumeTokenWhenNotYieldedRows()
+    {
+        $timesCalled = 0;
+        $chunks = [
+            [
+                'metadata' => $this->metadata,
+                'values' => ['a'],
+                'resumeToken' => 'abc'
+            ],
+            [
+                'values' => ['b']
+            ],
+            [
+                'values' => ['c']
+            ]
+        ];
+
+        $result = $this->getResultClass(
+            null,
+            'r',
+            null,
+            function () use ($chunks, &$timesCalled) {
+                $timesCalled++;
+                foreach ($chunks as $key => $chunk) {
+                    if ($key === 1 && $timesCalled < 2) {
+                        throw new ServiceException('Unavailable', 14);
+                    }
+                    yield $chunk;
+                }
+            }
+        );
+
+        iterator_to_array($result->rows());
+        $this->assertEquals(2, $timesCalled);
+    }
+
     public function testThrowsExceptionWhenCannotRetry()
     {
         $this->expectException(ServiceException::class);