googleapis · Mariatta · Mar 30, 2023 · Mar 24, 2023
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():