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():