diff --git a/Bigtable/src/ChunkFormatter.php b/Bigtable/src/ChunkFormatter.php index af7ffc441292..6e2ef5a5ee5e 100644 --- a/Bigtable/src/ChunkFormatter.php +++ b/Bigtable/src/ChunkFormatter.php @@ -142,7 +142,7 @@ public function __construct( $this->gapicClient = $gapicClient; $this->request = $request; $this->options = $options; - + if ($request->getRowsLimit()) { $this->originalRowsLimit = $request->getRowsLimit(); } @@ -228,6 +228,9 @@ private function updateReadRowsRequest($request, $options, $prevRowKey) { if ($this->originalRowsLimit) { $request->setRowsLimit($this->originalRowsLimit - $this->numberOfRowsRead); + if ($this->numberOfRowsRead === $this->originalRowsLimit) { + $options['requestCompleted'] = true; + } } if ($request->hasRows()) { diff --git a/Bigtable/tests/Unit/SmartRetriesTest.php b/Bigtable/tests/Unit/SmartRetriesTest.php index c29accd4060a..a8c6078eef41 100644 --- a/Bigtable/tests/Unit/SmartRetriesTest.php +++ b/Bigtable/tests/Unit/SmartRetriesTest.php @@ -247,6 +247,34 @@ public function testReadRowsWithRowsLimit() $this->assertEquals($expectedRows, $rows); } + public function testReadRowsWithRowsLimitAndExceptionThrownAfterSuccess() + { + $args = ['rowsLimit' => 1]; + $expectedArgs = $args + $this->options; + $this->serverStream->readAll() + ->shouldBeCalledTimes(1) + ->willReturn( + $this->arrayAsGeneratorWithException( + $this->generateRowsResponse(1, 1), + $this->retryingApiException + ) + ); + $this->bigtableClient->readRows( + Argument::that(function ($request) use (&$allowedRowsLimit) { + return $request->getRowsLimit() === 1; + }), + Argument::type('array') + ) + ->shouldBeCalledTimes(1) + ->willReturn( + $this->serverStream->reveal() + ); + + $rows = $this->table->readRows($args); + $expectedRows = $this->generateExpectedRows(1, 1); + $this->assertEquals($expectedRows, iterator_to_array($rows)); + } + public function testReadRowsWithRowKeys() { $args = ['rowKeys' => ['rk1', 'rk2', 'rk3', 'rk4']];