diff --git a/google/cloud/bigtable/row_data.py b/google/cloud/bigtable/row_data.py index 1cdd99026..3abb6c68c 100644 --- a/google/cloud/bigtable/row_data.py +++ b/google/cloud/bigtable/row_data.py @@ -625,16 +625,12 @@ def __init__(self, message, last_scanned_key, rows_read_so_far): def build_updated_request(self): """Updates the given message request as per last scanned key""" - # TODO: Generalize this to ensure fields don't get rewritten when retrying the request - - r_kwargs = { - "table_name": self.message.table_name, - "filter": self.message.filter, - "app_profile_id": self.message.app_profile_id, - } + resume_request = data_messages_v2_pb2.ReadRowsRequest() + data_messages_v2_pb2.ReadRowsRequest.copy_from(resume_request, self.message) if self.message.rows_limit != 0: - r_kwargs["rows_limit"] = max( + # TODO: Throw an error if rows_limit - read_so_far is 0 or negative. + resume_request.rows_limit = max( 1, self.message.rows_limit - self.rows_read_so_far ) @@ -643,14 +639,14 @@ def build_updated_request(self): # to request only rows that have not been returned yet if "rows" not in self.message: row_range = data_v2_pb2.RowRange(start_key_open=self.last_scanned_key) - r_kwargs["rows"] = data_v2_pb2.RowSet(row_ranges=[row_range]) + resume_request.rows = data_v2_pb2.RowSet(row_ranges=[row_range]) else: row_keys = self._filter_rows_keys() row_ranges = self._filter_row_ranges() - r_kwargs["rows"] = data_v2_pb2.RowSet( + resume_request.rows = data_v2_pb2.RowSet( row_keys=row_keys, row_ranges=row_ranges ) - return data_messages_v2_pb2.ReadRowsRequest(**r_kwargs) + return resume_request def _filter_rows_keys(self): """ Helper for :meth:`build_updated_request`""" diff --git a/setup.py b/setup.py index 78bfefcec..15a66b6a3 100644 --- a/setup.py +++ b/setup.py @@ -38,7 +38,7 @@ # https://github.com/googleapis/google-cloud-python/issues/10566 "google-cloud-core >= 1.4.1, <3.0.0dev", "grpc-google-iam-v1 >= 0.12.3, < 0.13dev", - "proto-plus >= 1.15.0", + "proto-plus >= 1.18.0", ] extras = {"libcst": "libcst >= 0.2.5"} diff --git a/testing/constraints-3.6.txt b/testing/constraints-3.6.txt index 3d010787d..92345bed8 100644 --- a/testing/constraints-3.6.txt +++ b/testing/constraints-3.6.txt @@ -8,5 +8,5 @@ google-api-core==1.31.5 google-cloud-core==1.4.1 grpc-google-iam-v1==0.12.3 -proto-plus==1.15.0 +proto-plus==1.18.0 libcst==0.2.5 diff --git a/tests/unit/test_row_data.py b/tests/unit/test_row_data.py index e48893df5..60a138800 100644 --- a/tests/unit/test_row_data.py +++ b/tests/unit/test_row_data.py @@ -928,7 +928,7 @@ def test_RRRM_build_updated_request_full_table(): request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) result = request_manager.build_updated_request() - expected_result = _ReadRowsRequestPB(table_name=TABLE_NAME, filter={}) + expected_result = _ReadRowsRequestPB(table_name=TABLE_NAME) row_range1 = types.RowRange(start_key_open=last_scanned_key) expected_result.rows.row_ranges.append(row_range1) assert expected_result == result @@ -1021,7 +1021,7 @@ def test_RRRM_build_updated_request_rows_limit(): request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) result = request_manager.build_updated_request() - expected_result = _ReadRowsRequestPB(table_name=TABLE_NAME, filter={}, rows_limit=8) + expected_result = _ReadRowsRequestPB(table_name=TABLE_NAME, rows_limit=8) row_range1 = types.RowRange(start_key_open=last_scanned_key) expected_result.rows.row_ranges.append(row_range1) assert expected_result == result