From 5d789b137e0a42c7eaa5b02a70ede13fcd1390df Mon Sep 17 00:00:00 2001 From: Mackenzie Starr Date: Fri, 3 Jan 2020 13:02:39 -0500 Subject: [PATCH] fix(bigtable): add ability to use single-row transactions (#10021) * fix: add ability to use single-row transactions fixes #10018 * fix: add unit tests for supporting single-row transactions --- google/cloud/bigtable/row.py | 6 +++++- tests/unit/test_row.py | 14 ++++++++++---- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/google/cloud/bigtable/row.py b/google/cloud/bigtable/row.py index 8d1f91f32..079ba6c8f 100644 --- a/google/cloud/bigtable/row.py +++ b/google/cloud/bigtable/row.py @@ -582,6 +582,7 @@ def commit(self): table_name=self._table.name, row_key=self._row_key, predicate_filter=self._filter.to_pb(), + app_profile_id=self._table._app_profile_id, true_mutations=true_mutations, false_mutations=false_mutations, ) @@ -908,7 +909,10 @@ def commit(self): data_client = self._table._instance._client.table_data_client row_response = data_client.read_modify_write_row( - table_name=self._table.name, row_key=self._row_key, rules=self._rule_pb_list + table_name=self._table.name, + row_key=self._row_key, + rules=self._rule_pb_list, + app_profile_id=self._table._app_profile_id, ) # Reset modifications after commit-ing request. diff --git a/tests/unit/test_row.py b/tests/unit/test_row.py index b4aaefb86..47424d910 100644 --- a/tests/unit/test_row.py +++ b/tests/unit/test_row.py @@ -409,6 +409,7 @@ def test_commit(self): project_id = "project-id" row_key = b"row_key" table_name = "projects/more-stuff" + app_profile_id = "app_profile_id" column_family_id1 = u"column_family_id1" column_family_id2 = u"column_family_id2" column_family_id3 = u"column_family_id3" @@ -420,7 +421,7 @@ def test_commit(self): client = self._make_client( project=project_id, credentials=credentials, admin=True ) - table = _Table(table_name, client=client) + table = _Table(table_name, client=client, app_profile_id=app_profile_id) row_filter = RowSampleFilter(0.33) row = self._make_one(row_key, table, filter_=row_filter) @@ -444,6 +445,8 @@ def test_commit(self): row.delete_cell(column_family_id2, column2, state=True) row.delete_cells(column_family_id3, row.ALL_COLUMNS, state=True) result = row.commit() + call_args = api.transport.check_and_mutate_row.call_args.args[0] + self.assertEqual(app_profile_id, call_args.app_profile_id) self.assertEqual(result, expected_result) self.assertEqual(row._true_pb_mutations, []) self.assertEqual(row._false_pb_mutations, []) @@ -564,6 +567,7 @@ def test_commit(self): project_id = "project-id" row_key = b"row_key" table_name = "projects/more-stuff" + app_profile_id = "app_profile_id" column_family_id = u"column_family_id" column = b"column" @@ -572,7 +576,7 @@ def test_commit(self): client = self._make_client( project=project_id, credentials=credentials, admin=True ) - table = _Table(table_name, client=client) + table = _Table(table_name, client=client, app_profile_id=app_profile_id) row = self._make_one(row_key, table) # Create request_pb @@ -593,7 +597,8 @@ def mock_parse_rmw_row_response(row_response): with _Monkey(MUT, _parse_rmw_row_response=mock_parse_rmw_row_response): row.append_cell_value(column_family_id, column, value) result = row.commit() - + call_args = api.transport.read_modify_write_row.call_args.args[0] + self.assertEqual(app_profile_id, call_args.app_profile_id) self.assertEqual(result, expected_result) self.assertEqual(row._rule_pb_list, []) @@ -819,9 +824,10 @@ def __init__(self, client=None): class _Table(object): - def __init__(self, name, client=None): + def __init__(self, name, client=None, app_profile_id=None): self.name = name self._instance = _Instance(client) + self._app_profile_id = app_profile_id self.client = client self.mutated_rows = []