From 4e61e15aa88c48a8fb9e34a48e5beab1e09641af Mon Sep 17 00:00:00 2001 From: HemangChothani Date: Thu, 7 May 2020 13:55:52 +0530 Subject: [PATCH 1/3] feat(bigtable): add support read row by prefix --- docs/snippets_table.py | 30 +++++++++++++++++++++++++++++- google/cloud/bigtable/row_set.py | 18 ++++++++++++++++++ tests/system.py | 32 ++++++++++++++++++++++++++++++++ tests/unit/test_row_set.py | 5 +++++ 4 files changed, 84 insertions(+), 1 deletion(-) diff --git a/docs/snippets_table.py b/docs/snippets_table.py index 702cf31b1..b465f82cb 100644 --- a/docs/snippets_table.py +++ b/docs/snippets_table.py @@ -620,7 +620,7 @@ def test_bigtable_create_update_delete_column_family(): assert column_family_id not in column_families -def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): +def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by_prefix_from_keys(): row_keys = [ b"row_key_1", b"row_key_2", @@ -631,6 +631,7 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): b"row_key_7", b"row_key_8", b"row_key_9", + b"sample_row_key_1", ] rows = [] @@ -692,6 +693,33 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): found_row_keys = [row.row_key for row in read_rows] assert found_row_keys == expected_row_keys + # [START bigtable_row_range_by_prefix_from_keys] + from google.cloud.bigtable import Client + from google.cloud.bigtable.row_set import RowSet + + client = Client(admin=True) + instance = client.instance(INSTANCE_ID) + table = instance.table(TABLE_ID) + + row_set = RowSet() + row_set.add_row_range_by_prefix_from_keys("row") + # [END bigtable_row_range_by_prefix_from_keys] + + read_rows = table.read_rows(row_set=row_set) + expected_row_keys = [ + b"row_key_1", + b"row_key_2", + b"row_key_3", + b"row_key_4", + b"row_key_5", + b"row_key_6", + b"row_key_7", + b"row_key_8", + b"row_key_9", + ] + found_row_keys = [row.row_key for row in read_rows] + assert found_row_keys == expected_row_keys + table.truncate(timeout=200) diff --git a/google/cloud/bigtable/row_set.py b/google/cloud/bigtable/row_set.py index 0cb6443b0..d9b1fe235 100644 --- a/google/cloud/bigtable/row_set.py +++ b/google/cloud/bigtable/row_set.py @@ -109,6 +109,24 @@ def add_row_range_from_keys( row_range = RowRange(start_key, end_key, start_inclusive, end_inclusive) self.row_ranges.append(row_range) + def add_row_range_by_prefix_from_keys(self, row_key_prefix): + """Add row range to row_ranges list that start with the row_key_prefix from the row keys + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_row_range_by_prefix_from_keys] + :end-before: [END bigtable_row_range_by_prefix_from_keys] + + :type row_key_prefix: str + :param row_key_prefix: Add all rows that start with this row key prefix. + Prefix cannot be zero length.""" + + end_key = row_key_prefix[:-1] + chr(ord(row_key_prefix[-1]) + 1) + self.add_row_range_from_keys( + row_key_prefix.encode("utf-8"), end_key.encode("utf-8") + ) + def _update_message_request(self, message): """Add row keys and row range to given request message diff --git a/tests/system.py b/tests/system.py index e9e3ab791..736197254 100644 --- a/tests/system.py +++ b/tests/system.py @@ -1013,6 +1013,38 @@ def test_yield_rows_with_row_set(self): found_row_keys = [row.row_key for row in read_rows] self.assertEqual(found_row_keys, expected_row_keys) + def test_add_row_range_by_prefix_from_keys(self): + row_keys = [ + b"row_key_1", + b"row_key_2", + b"row_key_3", + b"row_key_4", + b"sample_row_key_1", + b"sample_row_key_2", + ] + + rows = [] + for row_key in row_keys: + row = self._table.row(row_key) + row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) + rows.append(row) + self.rows_to_delete.append(row) + self._table.mutate_rows(rows) + + row_set = RowSet() + row_set.add_row_range_by_prefix_from_keys("row") + + read_rows = self._table.yield_rows(row_set=row_set) + + expected_row_keys = [ + b"row_key_1", + b"row_key_2", + b"row_key_3", + b"row_key_4", + ] + found_row_keys = [row.row_key for row in read_rows] + self.assertEqual(found_row_keys, expected_row_keys) + def test_read_large_cell_limit(self): row = self._table.row(ROW_KEY) self.rows_to_delete.append(row) diff --git a/tests/unit/test_row_set.py b/tests/unit/test_row_set.py index c66341b84..d2b1cc8d8 100644 --- a/tests/unit/test_row_set.py +++ b/tests/unit/test_row_set.py @@ -167,6 +167,11 @@ def test_add_row_range_from_keys(self): ) self.assertEqual(row_set.row_ranges[0].end_key, b"row_key9") + def test_add_row_range_by_prefix_from_keys(self): + row_set = self._make_one() + row_set.add_row_range_by_prefix_from_keys("row") + self.assertEqual(row_set.row_ranges[0].end_key, b"rox") + def test__update_message_request(self): row_set = self._make_one() table_name = "table_name" From 7e3827697641e99666b0ead96096f7583c200f2d Mon Sep 17 00:00:00 2001 From: HemangChothani Date: Fri, 8 May 2020 16:02:37 +0530 Subject: [PATCH 2/3] docs(bigtable): nit --- google/cloud/bigtable/row_set.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/bigtable/row_set.py b/google/cloud/bigtable/row_set.py index d9b1fe235..357447554 100644 --- a/google/cloud/bigtable/row_set.py +++ b/google/cloud/bigtable/row_set.py @@ -119,7 +119,7 @@ def add_row_range_by_prefix_from_keys(self, row_key_prefix): :end-before: [END bigtable_row_range_by_prefix_from_keys] :type row_key_prefix: str - :param row_key_prefix: Add all rows that start with this row key prefix. + :param row_key_prefix: To retrieve all rows that start with this row key prefix. Prefix cannot be zero length.""" end_key = row_key_prefix[:-1] + chr(ord(row_key_prefix[-1]) + 1) From 909ba28d04fc544d6a6e38a4b47f4e96a9eb8d29 Mon Sep 17 00:00:00 2001 From: HemangChothani Date: Wed, 20 May 2020 14:46:57 +0530 Subject: [PATCH 3/3] feat(bigtable): separate the snippet test --- docs/snippets_table.py | 34 +++++++++++++++++++++----------- google/cloud/bigtable/row_set.py | 6 +++--- tests/system.py | 2 +- tests/unit/test_row_set.py | 4 ++-- 4 files changed, 28 insertions(+), 18 deletions(-) diff --git a/docs/snippets_table.py b/docs/snippets_table.py index b465f82cb..a1d8d37fc 100644 --- a/docs/snippets_table.py +++ b/docs/snippets_table.py @@ -620,7 +620,7 @@ def test_bigtable_create_update_delete_column_family(): assert column_family_id not in column_families -def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by_prefix_from_keys(): +def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): row_keys = [ b"row_key_1", b"row_key_2", @@ -631,7 +631,6 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by b"row_key_7", b"row_key_8", b"row_key_9", - b"sample_row_key_1", ] rows = [] @@ -692,8 +691,26 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by expected_row_keys = [b"row_key_3", b"row_key_4", b"row_key_5", b"row_key_6"] found_row_keys = [row.row_key for row in read_rows] assert found_row_keys == expected_row_keys + table.truncate(timeout=200) + + +def test_bigtable_add_row_range_with_prefix(): + row_keys = [ + b"row_key_1", + b"row_key_2", + b"row_key_3", + b"sample_row_key_1", + b"sample_row_key_2", + ] + + rows = [] + for row_key in row_keys: + row = Config.TABLE.row(row_key) + row.set_cell(COLUMN_FAMILY_ID, COL_NAME1, CELL_VAL1) + rows.append(row) + Config.TABLE.mutate_rows(rows) - # [START bigtable_row_range_by_prefix_from_keys] + # [START bigtable_add_row_range_with_prefix] from google.cloud.bigtable import Client from google.cloud.bigtable.row_set import RowSet @@ -702,24 +719,17 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by table = instance.table(TABLE_ID) row_set = RowSet() - row_set.add_row_range_by_prefix_from_keys("row") - # [END bigtable_row_range_by_prefix_from_keys] + row_set.add_row_range_with_prefix("row") + # [END bigtable_add_row_range_with_prefix] read_rows = table.read_rows(row_set=row_set) expected_row_keys = [ b"row_key_1", b"row_key_2", b"row_key_3", - b"row_key_4", - b"row_key_5", - b"row_key_6", - b"row_key_7", - b"row_key_8", - b"row_key_9", ] found_row_keys = [row.row_key for row in read_rows] assert found_row_keys == expected_row_keys - table.truncate(timeout=200) diff --git a/google/cloud/bigtable/row_set.py b/google/cloud/bigtable/row_set.py index 357447554..e71b44875 100644 --- a/google/cloud/bigtable/row_set.py +++ b/google/cloud/bigtable/row_set.py @@ -109,14 +109,14 @@ def add_row_range_from_keys( row_range = RowRange(start_key, end_key, start_inclusive, end_inclusive) self.row_ranges.append(row_range) - def add_row_range_by_prefix_from_keys(self, row_key_prefix): + def add_row_range_with_prefix(self, row_key_prefix): """Add row range to row_ranges list that start with the row_key_prefix from the row keys For example: .. literalinclude:: snippets_table.py - :start-after: [START bigtable_row_range_by_prefix_from_keys] - :end-before: [END bigtable_row_range_by_prefix_from_keys] + :start-after: [START bigtable_add_row_range_with_prefix] + :end-before: [END bigtable_add_row_range_with_prefix] :type row_key_prefix: str :param row_key_prefix: To retrieve all rows that start with this row key prefix. diff --git a/tests/system.py b/tests/system.py index 736197254..8c5d8e0e4 100644 --- a/tests/system.py +++ b/tests/system.py @@ -1032,7 +1032,7 @@ def test_add_row_range_by_prefix_from_keys(self): self._table.mutate_rows(rows) row_set = RowSet() - row_set.add_row_range_by_prefix_from_keys("row") + row_set.add_row_range_with_prefix("row") read_rows = self._table.yield_rows(row_set=row_set) diff --git a/tests/unit/test_row_set.py b/tests/unit/test_row_set.py index d2b1cc8d8..a855099a1 100644 --- a/tests/unit/test_row_set.py +++ b/tests/unit/test_row_set.py @@ -167,9 +167,9 @@ def test_add_row_range_from_keys(self): ) self.assertEqual(row_set.row_ranges[0].end_key, b"row_key9") - def test_add_row_range_by_prefix_from_keys(self): + def test_add_row_range_with_prefix(self): row_set = self._make_one() - row_set.add_row_range_by_prefix_from_keys("row") + row_set.add_row_range_with_prefix("row") self.assertEqual(row_set.row_ranges[0].end_key, b"rox") def test__update_message_request(self):