feat: support row_range in sample_row_keys method (googleapis#17330)

Thank you for opening a Pull Request! Before submitting your PR, there
are a few things you can do to make sure it goes smoothly:
- [x] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/google-cloud-python/issues)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [x] Ensure the tests and linter pass
- [x] Code coverage does not decrease (if any source code was changed)
- [x] Appropriate docs were updated (if necessary)

Fixes googleapis#17329

Author: anvitTawar

packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py

@@ -84,7 +84,7 @@
 from google.cloud.bigtable.data.execute_query.values import ExecuteQueryValueType
 from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry
 from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule
-from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
+from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery, RowRange
 from google.cloud.bigtable.data.row import Row
 from google.cloud.bigtable.data.row_filters import (
     CellsRowLimitFilter,
@@ -1412,6 +1412,7 @@ async def row_exists(
     async def sample_row_keys(
         self,
         *,
+        row_range: RowRange | None = None,
         operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT,
         attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT,
         retryable_errors: Sequence[type[Exception]]
@@ -1429,6 +1430,8 @@ async def sample_row_keys(
         row_keys, along with offset positions in the table
 
         Args:
+            row_range: the range of rows to sample. If not provided, samples the
+                entire table.
             operation_timeout: the time budget for the entire operation, in seconds.
                 Failed requests will be retried within the budget.i
                 Defaults to the Table's default_operation_timeout
@@ -1466,7 +1469,9 @@ async def sample_row_keys(
             async def execute_rpc():
                 results = await self.client._gapic_client.sample_row_keys(
                     request=SampleRowKeysRequest(
-                        app_profile_id=self.app_profile_id, **self._request_path
+                        app_profile_id=self.app_profile_id,
+                        row_range=row_range._to_pb() if row_range is not None else None,
+                        **self._request_path,
                     ),
                     timeout=next(attempt_timeout_gen),
                     retry=None,

packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py

@@ -85,7 +85,7 @@
 from google.cloud.bigtable.data.execute_query.values import ExecuteQueryValueType
 from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry
 from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule
-from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
+from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery, RowRange
 from google.cloud.bigtable.data.row import Row
 from google.cloud.bigtable.data.row_filters import (
     CellsRowLimitFilter,
@@ -1160,6 +1160,7 @@ def row_exists(
     def sample_row_keys(
         self,
         *,
+        row_range: RowRange | None = None,
         operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT,
         attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT,
         retryable_errors: Sequence[type[Exception]]
@@ -1176,6 +1177,8 @@ def sample_row_keys(
         row_keys, along with offset positions in the table
 
         Args:
+            row_range: the range of rows to sample. If not provided, samples the
+                entire table.
             operation_timeout: the time budget for the entire operation, in seconds.
                 Failed requests will be retried within the budget.i
                 Defaults to the Table's default_operation_timeout
@@ -1208,7 +1211,9 @@ def sample_row_keys(
             def execute_rpc():
                 results = self.client._gapic_client.sample_row_keys(
                     request=SampleRowKeysRequest(
-                        app_profile_id=self.app_profile_id, **self._request_path
+                        app_profile_id=self.app_profile_id,
+                        row_range=row_range._to_pb() if row_range is not None else None,
+                        **self._request_path,
                     ),
                     timeout=next(attempt_timeout_gen),
                     retry=None,

packages/google-cloud-bigtable/test_proxy/handlers/client_handler_data_async.py

@@ -250,7 +250,11 @@ async def SampleRowKeys(self, request, **kwargs):
         kwargs["operation_timeout"] = (
             kwargs.get("operation_timeout", self.per_operation_timeout) or 20
         )
-        result = CrossSync.rm_aio(await table.sample_row_keys(**kwargs))
+        row_range = None
+        if "row_range" in request:
+            from google.cloud.bigtable.data.read_rows_query import RowRange
+            row_range = RowRange._from_dict(request["row_range"])
+        result = CrossSync.rm_aio(await table.sample_row_keys(row_range=row_range, **kwargs))
         return result
 
     @error_safe

packages/google-cloud-bigtable/test_proxy/handlers/client_handler_data_sync_autogen.py

@@ -187,7 +187,12 @@ async def SampleRowKeys(self, request, **kwargs):
         kwargs["operation_timeout"] = (
             kwargs.get("operation_timeout", self.per_operation_timeout) or 20
         )
-        result = table.sample_row_keys(**kwargs)
+        row_range = None
+        if "row_range" in request:
+            from google.cloud.bigtable.data.read_rows_query import RowRange
+
+            row_range = RowRange._from_dict(request["row_range"])
+        result = table.sample_row_keys(row_range=row_range, **kwargs)
         return result
 
     @error_safe

packages/google-cloud-bigtable/tests/system/data/test_system_async.py

@@ -346,6 +346,40 @@ async def test_sample_row_keys(
         assert results[-1][0] == b""
         assert isinstance(results[-1][1], int)
 
+    @pytest.mark.skipif(
+        bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't use splits"
+    )
+    @pytest.mark.usefixtures("client")
+    @pytest.mark.usefixtures("target")
+    @CrossSync.Retry(
+        predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
+    )
+    @CrossSync.pytest
+    async def test_sample_row_keys_w_row_range(
+        self, client, target, column_split_config
+    ):
+        """
+        Sample keys with row range should return samples within the range,
+        with the last key matching the end of the range.
+        """
+        if len(column_split_config) < 4:
+            pytest.skip("Not enough splits in column_split_config for this test")
+
+        from google.cloud.bigtable.data import RowRange
+
+        start_key = column_split_config[1]
+        end_key = column_split_config[3]
+        row_range = RowRange(start_key=start_key, end_key=end_key)
+
+        results = await target.sample_row_keys(row_range=row_range)
+        assert len(results) == 2
+
+        assert results[0][0] == column_split_config[2]
+        assert results[1][0] == column_split_config[3]
+
+        for _, offset in results:
+            assert isinstance(offset, int)
+
     @pytest.mark.usefixtures("client")
     @pytest.mark.usefixtures("target")
     @CrossSync.pytest

packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py

@@ -270,6 +270,31 @@ def test_sample_row_keys(self, client, target, temp_rows, column_split_config):
         assert results[-1][0] == b""
         assert isinstance(results[-1][1], int)
 
+    @pytest.mark.skipif(
+        bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't use splits"
+    )
+    @pytest.mark.usefixtures("client")
+    @pytest.mark.usefixtures("target")
+    @CrossSync._Sync_Impl.Retry(
+        predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
+    )
+    def test_sample_row_keys_w_row_range(self, client, target, column_split_config):
+        """Sample keys with row range should return samples within the range,
+        with the last key matching the end of the range."""
+        if len(column_split_config) < 4:
+            pytest.skip("Not enough splits in column_split_config for this test")
+        from google.cloud.bigtable.data import RowRange
+
+        start_key = column_split_config[1]
+        end_key = column_split_config[3]
+        row_range = RowRange(start_key=start_key, end_key=end_key)
+        results = target.sample_row_keys(row_range=row_range)
+        assert len(results) == 2
+        assert results[0][0] == column_split_config[2]
+        assert results[1][0] == column_split_config[3]
+        for _, offset in results:
+            assert isinstance(offset, int)
+
     @pytest.mark.usefixtures("client")
     @pytest.mark.usefixtures("target")
     def test_bulk_mutations_set_cell(self, client, target, temp_rows):

packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py

@@ -2451,6 +2451,32 @@ async def test_sample_row_keys(self):
                     assert result[1] == samples[1]
                     assert result[2] == samples[2]
 
+    @CrossSync.pytest
+    async def test_sample_row_keys_w_row_range(self):
+        """
+        Test that method returns the expected key samples when row_range is provided
+        """
+        samples = [
+            (b"a_key1", 100),
+            (b"b", 200),
+        ]
+        from google.cloud.bigtable.data import RowRange
+
+        row_range = RowRange(start_key=b"a", end_key=b"b")
+        async with self._make_client() as client:
+            async with client.get_table("instance", "table") as table:
+                with mock.patch.object(
+                    table.client._gapic_client, "sample_row_keys", CrossSync.Mock()
+                ) as sample_row_keys:
+                    sample_row_keys.return_value = self._make_gapic_stream(samples)
+                    result = await table.sample_row_keys(row_range=row_range)
+                    assert len(result) == 2
+                    assert result[0] == samples[0]
+                    assert result[1] == samples[1]
+                    sample_row_keys.assert_called_once()
+                    called_request = sample_row_keys.call_args[1]["request"]
+                    assert called_request.row_range == row_range._to_pb()
+
     @CrossSync.pytest
     async def test_sample_row_keys_bad_timeout(self):
         """

packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py

@@ -2030,6 +2030,28 @@ def test_sample_row_keys(self):
                     assert result[1] == samples[1]
                     assert result[2] == samples[2]
 
+    def test_sample_row_keys_w_row_range(self):
+        """Test that method returns the expected key samples when row_range is provided"""
+        samples = [(b"a_key1", 100), (b"b", 200)]
+        from google.cloud.bigtable.data import RowRange
+
+        row_range = RowRange(start_key=b"a", end_key=b"b")
+        with self._make_client() as client:
+            with client.get_table("instance", "table") as table:
+                with mock.patch.object(
+                    table.client._gapic_client,
+                    "sample_row_keys",
+                    CrossSync._Sync_Impl.Mock(),
+                ) as sample_row_keys:
+                    sample_row_keys.return_value = self._make_gapic_stream(samples)
+                    result = table.sample_row_keys(row_range=row_range)
+                    assert len(result) == 2
+                    assert result[0] == samples[0]
+                    assert result[1] == samples[1]
+                    sample_row_keys.assert_called_once()
+                    called_request = sample_row_keys.call_args[1]["request"]
+                    assert called_request.row_range == row_range._to_pb()
+
     def test_sample_row_keys_bad_timeout(self):
         """should raise error if timeout is negative"""
         with self._make_client() as client: