From 4796ac85c877d75ed596cde7628dae31918ef726 Mon Sep 17 00:00:00 2001
From: HemangChothani <50404902+HemangChothani@users.noreply.github.com>
Date: Tue, 4 Aug 2020 22:44:38 +0530
Subject: [PATCH] feat: add 'Rowset.add_row_range_with_prefix' (#30)

Co-authored-by: Tres Seaver <tseaver@palladion.com>
---
 docs/snippets_table.py           | 38 ++++++++++++++++++++++++++++++++
 google/cloud/bigtable/row_set.py | 18 +++++++++++++++
 tests/system.py                  | 32 +++++++++++++++++++++++++++
 tests/unit/test_row_set.py       |  5 +++++
 4 files changed, 93 insertions(+)

diff --git a/docs/snippets_table.py b/docs/snippets_table.py
index 702cf31b1..a1d8d37fc 100644
--- a/docs/snippets_table.py
+++ b/docs/snippets_table.py
@@ -691,7 +691,45 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys():
     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_add_row_range_with_prefix]
+    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_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",
+    ]
+    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 aa3e5eb92..5de7dabff 100644
--- a/google/cloud/bigtable/row_set.py
+++ b/google/cloud/bigtable/row_set.py
@@ -112,6 +112,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_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_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.
+                            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 c41c90a6a..7dae9862e 100644
--- a/tests/system.py
+++ b/tests/system.py
@@ -1083,6 +1083,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_with_prefix("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):
         self._maybe_emulator_skip(
             "Maximum gRPC received message size for emulator is 4194304 bytes."
diff --git a/tests/unit/test_row_set.py b/tests/unit/test_row_set.py
index c66341b84..a855099a1 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_with_prefix(self):
+        row_set = self._make_one()
+        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):
         row_set = self._make_one()
         table_name = "table_name"