googleapis · daniel-sanche · Apr 6, 2023 · Apr 6, 2023 · Apr 6, 2023 · Apr 6, 2023
diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py
@@ -20,8 +20,6 @@
     Any,
     Optional,
     Set,
-    Callable,
-    Coroutine,
     TYPE_CHECKING,
 )
 
@@ -31,15 +29,20 @@
 import warnings
 import sys
 import random
+import os
+from itertools import chain
 
 from google.cloud.bigtable_v2.services.bigtable.client import BigtableClientMeta
 from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient
 from google.cloud.bigtable_v2.services.bigtable.async_client import DEFAULT_CLIENT_INFO
 from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import (
     PooledBigtableGrpcAsyncIOTransport,
+    PooledChannel,
 )
+from google.cloud.bigtable_v2.types.bigtable import PingAndWarmRequest
 from google.cloud.client import ClientWithProject
 from google.api_core.exceptions import GoogleAPICallError
+from google.cloud.environment_vars import BIGTABLE_EMULATOR  # type: ignore
 from google.api_core import retry_async as retries
 from google.api_core import exceptions as core_exceptions
 from google.cloud.bigtable._read_rows import _ReadRowsOperation
@@ -50,10 +53,14 @@
 from google.cloud.bigtable.row import Row
 from google.cloud.bigtable.read_rows_query import ReadRowsQuery
 from google.cloud.bigtable.iterators import ReadRowsIterator
+from google.cloud.bigtable.exceptions import FailedQueryShardError
+from google.cloud.bigtable.exceptions import ShardedReadRowsExceptionGroup
+
 from google.cloud.bigtable.mutations import Mutation, RowMutationEntry
 from google.cloud.bigtable._mutate_rows import _MutateRowsOperation
 from google.cloud.bigtable._helpers import _make_metadata
 from google.cloud.bigtable._helpers import _convert_retry_deadline
+from google.cloud.bigtable._helpers import _attempt_timeout_generator
 
 from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule
 from google.cloud.bigtable.row_filters import RowFilter
@@ -65,6 +72,9 @@
     from google.cloud.bigtable.mutations_batcher import MutationsBatcher
     from google.cloud.bigtable import RowKeySamples
 
+# used by read_rows_sharded to limit how many requests are attempted in parallel
+CONCURRENCY_LIMIT = 10
+
 
 class BigtableDataClient(ClientWithProject):
     def __init__(
@@ -107,6 +117,7 @@ def __init__(
         # set up client info headers for veneer library
         client_info = DEFAULT_CLIENT_INFO
         client_info.client_library_version = client_info.gapic_version
+        client_info.user_agent = f"bigtable-python/{client_info.client_library_version}"
         # parse client options
         if type(client_options) is dict:
             client_options = client_options_lib.from_dict(client_options)
@@ -137,23 +148,41 @@ def __init__(
         # attempt to start background tasks
         self._channel_init_time = time.time()
         self._channel_refresh_tasks: list[asyncio.Task[None]] = []
-        try:
-            self.start_background_channel_refresh()
-        except RuntimeError:
+        self._emulator_host = os.getenv(BIGTABLE_EMULATOR)
+        if self._emulator_host is not None:
+            # connect to an emulator host
             warnings.warn(
-                f"{self.__class__.__name__} should be started in an "
-                "asyncio event loop. Channel refresh will not be started",
+                "Connecting to Bigtable emulator at {}".format(self._emulator_host),
                 RuntimeWarning,
                 stacklevel=2,
             )
+            self.transport._grpc_channel = PooledChannel(
+                pool_size=pool_size,
+                host=self._emulator_host,
+                insecure=True,
+            )
+            # refresh cached stubs to use emulator pool
+            self.transport._stubs = {}
+            self.transport._prep_wrapped_messages(client_info)
+        else:
+            # attempt to start background channel refresh tasks
+            try:
+                self.start_background_channel_refresh()
+            except RuntimeError:
+                warnings.warn(
+                    f"{self.__class__.__name__} should be started in an "
+                    "asyncio event loop. Channel refresh will not be started",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
 
     def start_background_channel_refresh(self) -> None:
         """
         Starts a background task to ping and warm each channel in the pool
         Raises:
           - RuntimeError if not called in an asyncio event loop
         """
-        if not self._channel_refresh_tasks:
+        if not self._channel_refresh_tasks and not self._emulator_host:
             # raise RuntimeError if there is no event loop
             asyncio.get_running_loop()
             for channel_idx in range(self.transport.pool_size):
@@ -190,10 +219,13 @@ async def _ping_and_warm_instances(
             - sequence of results or exceptions from the ping requests
         """
         ping_rpc = channel.unary_unary(
-            "/google.bigtable.v2.Bigtable/PingAndWarmChannel"
+            "/google.bigtable.v2.Bigtable/PingAndWarm",
+            request_serializer=PingAndWarmRequest.serialize,
         )
         tasks = [ping_rpc({"name": n}) for n in self._active_instances]
-        return await asyncio.gather(*tasks, return_exceptions=True)
+        result = await asyncio.gather(*tasks, return_exceptions=True)
+        # return None in place of empty successful responses
+        return [r or None for r in result]
 
     async def _manage_channel(
         self,
@@ -534,20 +566,59 @@ async def read_rows_sharded(
         self,
         query_list: list[ReadRowsQuery] | list[dict[str, Any]],
         *,
-        limit: int | None,
-        operation_timeout: int | float | None = 60,
+        operation_timeout: int | float | None = None,
         per_request_timeout: int | float | None = None,
-    ) -> ReadRowsIterator:
+    ) -> list[Row]:
         """
-        Runs a sharded query in parallel
+        Runs a sharded query in parallel, then return the results in a single list.
+        Results will be returned in the order of the input queries.
 
-        Each query in query list will be run concurrently, with results yielded as they are ready
-        yielded results may be out of order
+        This function is intended to be run on the results on a query.shard() call:
+
+        ```
+        table_shard_keys = await table.sample_row_keys()
+        query = ReadRowsQuery(...)
+        shard_queries = query.shard(table_shard_keys)
+        results = await table.read_rows_sharded(shard_queries)
+        ```
 
         Args:
             - query_list: a list of queries to run in parallel
+        Raises:
+            - ShardedReadRowsExceptionGroup: if any of the queries failed
+            - ValueError: if the query_list is empty
         """
-        raise NotImplementedError
+        if not query_list:
+            raise ValueError("query_list must contain at least one query")
+        routine_list = [
+            self.read_rows(
+                query,
+                operation_timeout=operation_timeout,
+                per_request_timeout=per_request_timeout,
+            )
+            for query in query_list
+        ]
+        # submit requests in batches to limit concurrency
+        batched_routines = [
+            routine_list[i : i + CONCURRENCY_LIMIT]
+            for i in range(0, len(routine_list), CONCURRENCY_LIMIT)
+        ]
+        # run batches and collect results
+        results_list = []
+        for batch in batched_routines:
+            batch_result = await asyncio.gather(*batch, return_exceptions=True)
+            results_list.extend(batch_result)
+        # collect exceptions
+        exception_list = [
+            FailedQueryShardError(idx, query_list[idx], e)
+            for idx, e in enumerate(results_list)
+            if isinstance(e, Exception)
+        ]
+        if exception_list:
+            # if any sub-request failed, raise an exception instead of returning results
+            raise ShardedReadRowsExceptionGroup(exception_list, len(query_list))
+        combined_list = list(chain.from_iterable(results_list))
+        return combined_list
 
     async def row_exists(
         self,
@@ -577,32 +648,81 @@ async def row_exists(
         )
         return len(results) > 0
 
-    async def sample_keys(
+    async def sample_row_keys(
         self,
         *,
-        operation_timeout: int | float | None = 60,
-        per_sample_timeout: int | float | None = 10,
-        per_request_timeout: int | float | None = None,
+        operation_timeout: float | None = None,
+        per_request_timeout: float | None = None,
     ) -> RowKeySamples:
         """
         Return a set of RowKeySamples that delimit contiguous sections of the table of
         approximately equal size
 
         RowKeySamples output can be used with ReadRowsQuery.shard() to create a sharded query that
         can be parallelized across multiple backend nodes read_rows and read_rows_stream
-        requests will call sample_keys internally for this purpose when sharding is enabled
+        requests will call sample_row_keys internally for this purpose when sharding is enabled
 
         RowKeySamples is simply a type alias for list[tuple[bytes, int]]; a list of
             row_keys, along with offset positions in the table
 
         Returns:
             - a set of RowKeySamples the delimit contiguous sections of the table
         Raises:
-            - DeadlineExceeded: raised after operation timeout
-                will be chained with a RetryExceptionGroup containing all GoogleAPIError
-                exceptions from any retries that failed
+            - GoogleAPICallError: if the sample_row_keys request fails
         """
-        raise NotImplementedError
+        # prepare timeouts
+        operation_timeout = operation_timeout or self.default_operation_timeout
+        per_request_timeout = per_request_timeout or self.default_per_request_timeout
+
+        if operation_timeout <= 0:
+            raise ValueError("operation_timeout must be greater than 0")
+        if per_request_timeout is not None and per_request_timeout <= 0:
+            raise ValueError("per_request_timeout must be greater than 0")
+        if per_request_timeout is not None and per_request_timeout > operation_timeout:
+            raise ValueError(
+                "per_request_timeout must not be greater than operation_timeout"
+            )
+        attempt_timeout_gen = _attempt_timeout_generator(
+            per_request_timeout, operation_timeout
+        )
+        # prepare retryable
+        predicate = retries.if_exception_type(
+            core_exceptions.DeadlineExceeded,
+            core_exceptions.ServiceUnavailable,
+        )
+        transient_errors = []
+
+        def on_error_fn(exc):
+            # add errors to list if retryable
+            if predicate(exc):
+                transient_errors.append(exc)
+
+        retry = retries.AsyncRetry(
+            predicate=predicate,
+            timeout=operation_timeout,
+            initial=0.01,
+            multiplier=2,
+            maximum=60,
+            on_error=on_error_fn,
+            is_stream=False,
+        )
+
+        # prepare request
+        metadata = _make_metadata(self.table_name, self.app_profile_id)
+
+        async def execute_rpc():
+            results = await self.client._gapic_client.sample_row_keys(
+                table_name=self.table_name,
+                app_profile_id=self.app_profile_id,
+                timeout=next(attempt_timeout_gen),
+                metadata=metadata,
+            )
+            return [(s.row_key, s.offset_bytes) async for s in results]
+
+        wrapped_fn = _convert_retry_deadline(
+            retry(execute_rpc), operation_timeout, transient_errors
+        )
+        return await wrapped_fn()
 
     def mutations_batcher(self, **kwargs) -> MutationsBatcher:
         """
@@ -712,10 +832,6 @@ async def bulk_mutate_rows(
         *,
         operation_timeout: float | None = 60,
         per_request_timeout: float | None = None,
-        on_success: Callable[
-            [int, RowMutationEntry], None | Coroutine[None, None, None]
-        ]
-        | None = None,
     ):
         """
         Applies mutations for multiple rows in a single batched request.
@@ -741,9 +857,6 @@ async def bulk_mutate_rows(
                 in seconds. If it takes longer than this time to complete, the request
                 will be cancelled with a DeadlineExceeded exception, and a retry will
                 be attempted if within operation_timeout budget
-            - on_success: a callback function that will be called when each mutation
-                entry is confirmed to be applied successfully. Will be passed the
-                index and the entry itself.
         Raises:
             - MutationsExceptionGroup if one or more mutations fails
                 Contains details about any failed entries in .exceptions
@@ -896,16 +1009,17 @@ async def close(self):
         """
         Called to close the Table instance and release any resources held by it.
         """
+        self._register_instance_task.cancel()
         await self.client._remove_instance_registration(self.instance_id, self)
 
     async def __aenter__(self):
         """
         Implement async context manager protocol
 
-        Register this instance with the client, so that
+        Ensure registration task has time to run, so that
         grpc channels will be warmed for the specified instance
         """
-        await self.client._register_instance(self.instance_id, self)
+        await self._register_instance_task
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):

diff --git a/google/cloud/bigtable/exceptions.py b/google/cloud/bigtable/exceptions.py
@@ -16,14 +16,15 @@
 
 import sys
 
-from typing import TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 from google.api_core import exceptions as core_exceptions
 
 is_311_plus = sys.version_info >= (3, 11)
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.mutations import RowMutationEntry
+    from google.cloud.bigtable.read_rows_query import ReadRowsQuery
 
 
 class IdleTimeout(core_exceptions.DeadlineExceeded):
@@ -137,3 +138,39 @@ def __init__(self, excs: list[Exception]):
 
     def __new__(cls, excs: list[Exception]):
         return super().__new__(cls, cls._format_message(excs), excs)
+
+
+class ShardedReadRowsExceptionGroup(BigtableExceptionGroup):
+    """
+    Represents one or more exceptions that occur during a sharded read rows operation
+    """
+
+    @staticmethod
+    def _format_message(excs: list[FailedQueryShardError], total_queries: int):
+        query_str = "query" if total_queries == 1 else "queries"
+        plural_str = "" if len(excs) == 1 else "s"
+        return f"{len(excs)} sub-exception{plural_str} (from {total_queries} {query_str} attempted)"
+
+    def __init__(self, excs: list[FailedQueryShardError], total_queries: int):
+        super().__init__(self._format_message(excs, total_queries), excs)
+
+    def __new__(cls, excs: list[FailedQueryShardError], total_queries: int):
+        return super().__new__(cls, cls._format_message(excs, total_queries), excs)
+
+
+class FailedQueryShardError(Exception):
+    """
+    Represents an individual failed query in a sharded read rows operation
+    """
+
+    def __init__(
+        self,
+        failed_index: int,
+        failed_query: "ReadRowsQuery" | dict[str, Any],
+        cause: Exception,
+    ):
+        message = f"Failed query at index {failed_index} with cause: {cause!r}"
+        super().__init__(message)
+        self.index = failed_index
+        self.query = failed_query
+        self.__cause__ = cause