fix: Pass the "retry" when calling read_rows. (#759)

googleapis · Mar 30, 2023 · 505273b · 505273b
1 parent 45d3e43
commit 505273b
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 5 deletions.
diff --git a/google/cloud/bigtable/row_data.py b/google/cloud/bigtable/row_data.py
@@ -157,7 +157,9 @@ def __init__(self, read_method, request, retry=DEFAULT_RETRY_READ_ROWS):
         # Otherwise there is a risk of entering an infinite loop that resets
         # the timeout counter just before it being triggered. The increment
         # by 1 second here is customary but should not be much less than that.
-        self.response_iterator = read_method(request, timeout=self.retry._deadline + 1)
+        self.response_iterator = read_method(
+            request, timeout=self.retry._deadline + 1, retry=self.retry
+        )
 
         self.rows = {}
 

diff --git a/tests/unit/test_row_data.py b/tests/unit/test_row_data.py
@@ -446,7 +446,9 @@ def test_partial_rows_data_constructor_with_retry():
         client._data_stub.ReadRows, request, retry
     )
     partial_rows_data.read_method.assert_called_once_with(
-        request, timeout=DEFAULT_RETRY_READ_ROWS.deadline + 1
+        request,
+        timeout=DEFAULT_RETRY_READ_ROWS.deadline + 1,
+        retry=DEFAULT_RETRY_READ_ROWS,
     )
     assert partial_rows_data.request is request
     assert partial_rows_data.rows == {}

diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py
@@ -643,6 +643,7 @@ def _table_read_row_helper(chunks, expected_result, app_profile_id=None):
     from google.cloud.bigtable import table as MUT
     from google.cloud.bigtable.row_set import RowSet
     from google.cloud.bigtable.row_filters import RowSampleFilter
+    from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS
 
     credentials = _make_credentials()
     client = _make_client(project="project-id", credentials=credentials, admin=True)
@@ -691,7 +692,9 @@ def mock_create_row_request(table_name, **kwargs):
     assert result == expected_result
     assert mock_created == expected_request
 
-    data_api.read_rows.assert_called_once_with(request_pb, timeout=61.0)
+    data_api.read_rows.assert_called_once_with(
+        request_pb, timeout=61.0, retry=DEFAULT_RETRY_READ_ROWS
+    )
 
 
 def test_table_read_row_miss_no__responses():
@@ -906,7 +909,7 @@ def mock_create_row_request(table_name, **kwargs):
     }
     assert mock_created == [(table.name, created_kwargs)]
 
-    data_api.read_rows.assert_called_once_with(request_pb, timeout=61.0)
+    data_api.read_rows.assert_called_once_with(request_pb, timeout=61.0, retry=retry)
 
 
 def test_table_read_retry_rows():
@@ -1082,6 +1085,7 @@ def test_table_yield_rows_with_row_set():
     from google.cloud.bigtable.row_set import RowSet
     from google.cloud.bigtable.row_set import RowRange
     from google.cloud.bigtable.table import _create_row_request
+    from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS
 
     credentials = _make_credentials()
     client = _make_client(project="project-id", credentials=credentials, admin=True)
@@ -1149,7 +1153,9 @@ def test_table_yield_rows_with_row_set():
         end_key=ROW_KEY_2,
     )
     expected_request.rows.row_keys.append(ROW_KEY_3)
-    data_api.read_rows.assert_called_once_with(expected_request, timeout=61.0)
+    data_api.read_rows.assert_called_once_with(
+        expected_request, timeout=61.0, retry=DEFAULT_RETRY_READ_ROWS
+    )
 
 
 def test_table_sample_row_keys():