Skip to content

Commit

Permalink
fix: Pass the "retry" when calling read_rows. (#759)
Browse files Browse the repository at this point in the history
  • Loading branch information
Mariatta authored Mar 30, 2023
1 parent 45d3e43 commit 505273b
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 5 deletions.
4 changes: 3 additions & 1 deletion google/cloud/bigtable/row_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {}

Expand Down
4 changes: 3 additions & 1 deletion tests/unit/test_row_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 == {}
Expand Down
12 changes: 9 additions & 3 deletions tests/unit/test_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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():
Expand Down

0 comments on commit 505273b

Please sign in to comment.