From 2f8ba7a4801b17b5afb6180a7ace1327a2d05a52 Mon Sep 17 00:00:00 2001 From: Mariatta Wijaya Date: Thu, 24 Feb 2022 13:44:16 -0800 Subject: [PATCH] fix: Pass app_profile_id when building updated request (#512) - Pass the app_profile_id from the message - Update unittest Fixes internal bug #214449800 --- google/cloud/bigtable/row_data.py | 4 ++++ tests/unit/test_row_data.py | 10 ++++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/google/cloud/bigtable/row_data.py b/google/cloud/bigtable/row_data.py index 6ab1188a8..1cdd99026 100644 --- a/google/cloud/bigtable/row_data.py +++ b/google/cloud/bigtable/row_data.py @@ -624,9 +624,13 @@ 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, } if self.message.rows_limit != 0: diff --git a/tests/unit/test_row_data.py b/tests/unit/test_row_data.py index 06fd2f016..e48893df5 100644 --- a/tests/unit/test_row_data.py +++ b/tests/unit/test_row_data.py @@ -893,7 +893,10 @@ def test_RRRM_build_updated_request(rrrm_data): row_filter = RowSampleFilter(0.33) last_scanned_key = b"row_key25" request = _ReadRowsRequestPB( - filter=row_filter.to_pb(), rows_limit=8, table_name=TABLE_NAME + filter=row_filter.to_pb(), + rows_limit=8, + table_name=TABLE_NAME, + app_profile_id="app-profile-id-1", ) request.rows.row_ranges.append(row_range1.get_range_kwargs()) @@ -902,7 +905,10 @@ def test_RRRM_build_updated_request(rrrm_data): result = request_manager.build_updated_request() expected_result = _ReadRowsRequestPB( - table_name=TABLE_NAME, filter=row_filter.to_pb(), rows_limit=6 + table_name=TABLE_NAME, + filter=row_filter.to_pb(), + rows_limit=6, + app_profile_id="app-profile-id-1", ) row_range1 = types.RowRange(