feat: add pooled grpc transport

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "gapic-generator-fork"]
+	path = gapic-generator-fork
+	url = [email protected]:googleapis/gapic-generator-python.git

google/cloud/bigtable/client.py

@@ -15,12 +15,28 @@
 
 from __future__ import annotations
 
-from typing import Any, AsyncIterable, TYPE_CHECKING
+from typing import cast, Any, Optional, AsyncIterable, Set, TYPE_CHECKING
 
-from google.cloud.client import ClientWithProject
+import asyncio
+import grpc
+import time
+import warnings
+import sys
+
+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,
+)
+from google.cloud.client import _ClientProjectMixin
+from google.api_core.exceptions import GoogleAPICallError
 
 
 import google.auth.credentials
+import google.auth._default
+from google.api_core import client_options as client_options_lib
+
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.mutations import Mutation, BulkMutationsEntry
@@ -32,7 +48,7 @@
     from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule
 
 
-class BigtableDataClient(ClientWithProject):
+class BigtableDataClient(BigtableAsyncClient, _ClientProjectMixin):
     def __init__(
         self,
         *,
@@ -47,6 +63,8 @@ def __init__(
         """
         Create a client instance for the Bigtable Data API
 
+        Client must be created within an async run loop context
+
         Args:
             project: the project which the client acts on behalf of.
                 If not passed, falls back to the default inferred
@@ -62,29 +80,227 @@ def __init__(
                 Client options used to set user options
                 on the client. API Endpoint should be set through client_options.
             metadata: a list of metadata headers to be attached to all calls with this client
+        Raises:
+          - RuntimeError if called outside of an async run loop context
+          - ValueError if pool_size is less than 1
         """
-        raise NotImplementedError
+        # set up transport in registry
+        transport_str = f"pooled_grpc_asyncio_{pool_size}"
+        transport = PooledBigtableGrpcAsyncIOTransport.with_fixed_size(pool_size)
+        BigtableClientMeta._transport_registry[transport_str] = transport
+        # set up client info headers for veneer library
+        client_info = DEFAULT_CLIENT_INFO
+        client_info.client_library_version = client_info.gapic_version
+        # parse client options
+        if type(client_options) is dict:
+            client_options = client_options_lib.from_dict(client_options)
+        client_options = cast(
+            Optional[client_options_lib.ClientOptions], client_options
+        )
+        mixin_args = {"project": project, "credentials": credentials}
+        # support google-api-core <=1.5.0, which does not have credentials
+        if "credentials" not in _ClientProjectMixin.__init__.__code__.co_varnames:
+            mixin_args.pop("credentials")
+        # initialize client
+        _ClientProjectMixin.__init__(self, **mixin_args)
+        # raises RuntimeError if called outside of an async run loop context
+        BigtableAsyncClient.__init__(
+            self,
+            transport=transport_str,
+            credentials=credentials,
+            client_options=client_options,
+            client_info=client_info,
+        )
+        self.metadata = metadata or []
+        # keep track of active instances to for warmup on channel refresh
+        self._active_instances: Set[str] = set()
+        # 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:
+            warnings.warn(
+                f"{self.__class__.__name__} should be started in an "
+                "asyncio event loop. Channel refresh will not be started",
+                RuntimeWarning,
+            )
+
+    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:
+            # raise RuntimeError if there is no event loop
+            asyncio.get_running_loop()
+            for channel_idx in range(len(self.transport.channel_pool)):
+                refresh_task = asyncio.create_task(self._manage_channel(channel_idx))
+                if sys.version_info >= (3, 8):
+                    refresh_task.set_name(
+                        f"{self.__class__.__name__} channel refresh {channel_idx}"
+                    )
+                self._channel_refresh_tasks.append(refresh_task)
+
+    @property
+    def transport(self) -> PooledBigtableGrpcAsyncIOTransport:
+        """Returns the transport used by the client instance.
+        Returns:
+            BigtableTransport: The transport used by the client instance.
+        """
+        return cast(PooledBigtableGrpcAsyncIOTransport, self._client.transport)
+
+    async def close(self, timeout: float = 2.0):
+        """
+        Cancel all background tasks
+        """
+        for task in self._channel_refresh_tasks:
+            task.cancel()
+        group = asyncio.gather(*self._channel_refresh_tasks, return_exceptions=True)
+        await asyncio.wait_for(group, timeout=timeout)
+        await self.transport.close()
+        self._channel_refresh_tasks = []
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """
+        Cleanly close context manager on exit
+        """
+        await self.close()
+
+    async def _ping_and_warm_instances(
+        self, channel: grpc.aio.Channel
+    ) -> list[GoogleAPICallError | None]:
+        """
+        Prepares the backend for requests on a channel
+
+        Pings each Bigtable instance registered in `_active_instances` on the client
+
+        Args:
+            channel: grpc channel to ping
+        Returns:
+            - squence of results or exceptions from the ping requests
+        """
+        ping_rpc = channel.unary_unary(
+            "/google.bigtable.v2.Bigtable/PingAndWarmChannel"
+        )
+        tasks = [ping_rpc({"name": n}) for n in self._active_instances]
+        return await asyncio.gather(*tasks, return_exceptions=True)
+
+    async def _manage_channel(
+        self,
+        channel_idx: int,
+        refresh_interval: float = 60 * 45,
+        grace_period: float = 60 * 15,
+    ) -> None:
+        """
+        Background coroutine that periodically refreshes and warms a grpc channel
+
+        The backend will automatically close channels after 60 minutes, so
+        `refresh_interval` + `grace_period` should be < 60 minutes
+
+        Runs continuously until the client is closed
+
+        Args:
+            channel_idx: index of the channel in the transport's channel pool
+            refresh_interval: interval before initiating refresh process in seconds
+            grace_period: time to allow previous channel to serve existing
+                requests before closing, in seconds
+        """
+        first_refresh = self._channel_init_time + refresh_interval
+        next_sleep = max(first_refresh - time.time(), 0)
+        if next_sleep > 0:
+            # warm the current channel immediately
+            channel = self.transport.channel_pool[channel_idx]
+            await self._ping_and_warm_instances(channel)
+        # continuously refresh the channel every `refresh_interval` seconds
+        while True:
+            await asyncio.sleep(next_sleep)
+            # prepare new channel for use
+            new_channel = self.transport.create_channel(
+                self.transport._host,
+                credentials=self.transport._credentials,
+                scopes=self.transport._scopes,
+                ssl_credentials=self.transport._ssl_channel_credentials,
+                quota_project_id=self.transport._quota_project_id,
+                options=[
+                    ("grpc.max_send_message_length", -1),
+                    ("grpc.max_receive_message_length", -1),
+                ],
+            )
+            await self._ping_and_warm_instances(new_channel)
+            # cycle channel out of use, with long grace window before closure
+            start_timestamp = time.time()
+            await self.transport.replace_channel(channel_idx, grace_period, new_channel)
+            # subtract the time spent waiting for the channel to be replaced
+            next_sleep = refresh_interval - (time.time() - start_timestamp)
+
+    async def register_instance(self, instance_id: str):
+        """
+        Registers an instance with the client, and warms the channel pool
+        for the instance
+        The client will periodically refresh grpc channel pool used to make
+        requests, and new channels will be warmed for each registered instance
+        Channels will not be refreshed unless at least one instance is registered
+        """
+        instance_name = self.instance_path(self.project, instance_id)
+        if instance_name not in self._active_instances:
+            self._active_instances.add(instance_name)
+            if self._channel_refresh_tasks:
+                # refresh tasks already running
+                # call ping and warm on all existing channels
+                for channel in self.transport.channel_pool:
+                    await self._ping_and_warm_instances(channel)
+            else:
+                # refresh tasks aren't active. start them as background tasks
+                self.start_background_channel_refresh()
+
+    async def remove_instance_registration(self, instance_id: str) -> bool:
+        """
+        Removes an instance from the client's registered instances, to prevent
+        warming new channels for the instance
+
+        If instance_id is not registered, returns False
+
+        Args:
+            instance_id: id of the instance to remove
+        Returns:
+            - True if instance was removed
+        """
+        instance_name = self.instance_path(self.project, instance_id)
+        try:
+            self._active_instances.remove(instance_name)
+            return True
+        except KeyError:
+            return False
 
     def get_table(
-        self, instance_id: str, table_id: str, app_profile_id: str | None = None
+        self,
+        instance_id: str,
+        table_id: str,
+        app_profile_id: str | None = None,
+        metadata: list[tuple[str, str]] | None = None,
     ) -> Table:
         """
-        Return a Table instance to make API requests for a specific table.
+        Returns a table instance for making data API requests
 
         Args:
-            instance_id: The ID of the instance that owns the table.
+            instance_id: The Bigtable instance ID to associate with this client
+                instance_id is combined with the client's project to fully
+                specify the instance
             table_id: The ID of the table.
             app_profile_id: (Optional) The app profile to associate with requests.
                 https://cloud.google.com/bigtable/docs/app-profiles
+            metadata: a list of metadata headers to be attached to all calls with this client
         """
-        raise NotImplementedError
+        return Table(self, instance_id, table_id, app_profile_id, metadata)
 
 
 class Table:
     """
     Main Data API surface
 
-    Table object maintains instance_id, table_id, and app_profile_id context, and passes them with
+    Table object maintains table_id, and app_profile_id context, and passes them with
     each call
     """
 
@@ -94,8 +310,40 @@ def __init__(
         instance_id: str,
         table_id: str,
         app_profile_id: str | None = None,
+        metadata: list[tuple[str, str]] | None = None,
     ):
-        raise NotImplementedError
+        """
+        Initialize a Table instance
+
+        Must be created within an async run loop context
+
+        Args:
+            instance_id: The Bigtable instance ID to associate with this client
+                instance_id is combined with the client's project to fully
+                specify the instance
+            table_id: The ID of the table.
+            app_profile_id: (Optional) The app profile to associate with requests.
+                https://cloud.google.com/bigtable/docs/app-profiles
+            metadata: a list of metadata headers to be attached to all calls with this client
+        Raises:
+          - RuntimeError if called outside of an async run loop context
+        """
+        self.client = client
+        self.instance = instance_id
+        self.table_id = table_id
+        self.app_profile_id = app_profile_id
+        self.metadata = metadata or []
+        # raises RuntimeError if called outside of an async run loop context
+        try:
+            self._register_instance_task = asyncio.create_task(
+                self.client.register_instance(instance_id)
+            )
+        except RuntimeError:
+            warnings.warn(
+                "Table should be created in an asyncio event loop."
+                " Instance will not be registered with client for refresh",
+                RuntimeWarning,
+            )
 
     async def read_rows_stream(
         self,

google/cloud/bigtable_v2/services/bigtable/async_client.py

@@ -807,8 +807,8 @@ async def ping_and_warm(
 
         Args:
             request (Optional[Union[google.cloud.bigtable_v2.types.PingAndWarmRequest, dict]]):
-                The request object. Request message for client
-                connection keep-alive and warming.
+                The request object. Request message for client connection
+                keep-alive and warming.
             name (:class:`str`):
                 Required. The unique name of the instance to check
                 permissions for as well as respond. Values are of the
@@ -1027,8 +1027,9 @@ def generate_initial_change_stream_partitions(
 
         Args:
             request (Optional[Union[google.cloud.bigtable_v2.types.GenerateInitialChangeStreamPartitionsRequest, dict]]):
-                The request object. NOTE: This API is intended to be
-                used by Apache Beam BigtableIO. Request message for
+                The request object. NOTE: This API is intended to be used
+                by Apache Beam BigtableIO. Request
+                message for
                 Bigtable.GenerateInitialChangeStreamPartitions.
             table_name (:class:`str`):
                 Required. The unique name of the table from which to get
@@ -1126,9 +1127,9 @@ def read_change_stream(
 
         Args:
             request (Optional[Union[google.cloud.bigtable_v2.types.ReadChangeStreamRequest, dict]]):
-                The request object. NOTE: This API is intended to be
-                used by Apache Beam BigtableIO. Request message for
-                Bigtable.ReadChangeStream.
+                The request object. NOTE: This API is intended to be used
+                by Apache Beam BigtableIO. Request
+                message for Bigtable.ReadChangeStream.
             table_name (:class:`str`):
                 Required. The unique name of the table from which to
                 read a change stream. Values are of the form

google/cloud/bigtable_v2/services/bigtable/client.py

@@ -53,6 +53,7 @@
 from .transports.base import BigtableTransport, DEFAULT_CLIENT_INFO
 from .transports.grpc import BigtableGrpcTransport
 from .transports.grpc_asyncio import BigtableGrpcAsyncIOTransport
+from .transports.pooled_grpc_asyncio import PooledBigtableGrpcAsyncIOTransport
 from .transports.rest import BigtableRestTransport
 
 
@@ -67,6 +68,7 @@ class BigtableClientMeta(type):
     _transport_registry = OrderedDict()  # type: Dict[str, Type[BigtableTransport]]
     _transport_registry["grpc"] = BigtableGrpcTransport
     _transport_registry["grpc_asyncio"] = BigtableGrpcAsyncIOTransport
+    _transport_registry["pooled_grpc_asyncio"] = PooledBigtableGrpcAsyncIOTransport
     _transport_registry["rest"] = BigtableRestTransport
 
     def get_transport_class(
@@ -380,6 +382,9 @@ def __init__(
             transport (Union[str, BigtableTransport]): The
                 transport to use. If set to None, a transport is chosen
                 automatically.
+                NOTE: "rest" transport functionality is currently in a
+                beta state (preview). We welcome your feedback via an
+                issue in this library's source repository.
             client_options (Optional[Union[google.api_core.client_options.ClientOptions, dict]]): Custom options for the
                 client. It won't take effect if a ``transport`` instance is provided.
                 (1) The ``api_endpoint`` property can be used to override the