use replaceable channel wrapper

Author: daniel-sanche

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

@@ -92,7 +92,11 @@
     from google.cloud.bigtable_v2.services.bigtable.transports import (
         BigtableGrpcAsyncIOTransport as TransportType,
     )
+    from google.cloud.bigtable_v2.services.bigtable.transports.grpc_asyncio import (
+        _LoggingClientAIOInterceptor
+    )
     from google.cloud.bigtable.data._async.mutations_batcher import _MB_SIZE
+    from google.cloud.bigtable.data._async.replaceable_channel import _AsyncReplaceableChannel
 else:
     from typing import Iterable  # noqa: F401
     from grpc import insecure_channel
@@ -182,7 +186,6 @@ def __init__(
         client_options = cast(
             Optional[client_options_lib.ClientOptions], client_options
         )
-        custom_channel = None
         self._emulator_host = os.getenv(BIGTABLE_EMULATOR)
         if self._emulator_host is not None:
             warnings.warn(
@@ -191,11 +194,11 @@ def __init__(
                 stacklevel=2,
             )
             # use insecure channel if emulator is set
-            custom_channel = insecure_channel(self._emulator_host)
             if credentials is None:
                 credentials = google.auth.credentials.AnonymousCredentials()
             if project is None:
                 project = _DEFAULT_BIGTABLE_EMULATOR_CLIENT
+
         # initialize client
         ClientWithProject.__init__(
             self,
@@ -208,7 +211,7 @@ def __init__(
             client_options=client_options,
             client_info=self.client_info,
             transport=lambda *args, **kwargs: TransportType(
-                *args, **kwargs, channel=custom_channel
+                *args, **kwargs, channel=self._build_grpc_channel
             ),
         )
         self._is_closed = CrossSync.Event()
@@ -235,6 +238,23 @@ def __init__(
                     stacklevel=2,
                 )
 
+    def _build_grpc_channel(self, *args, **kwargs):
+        if self._emulator_host is not None:
+            # emulators use insecure channel
+            return insecure_channel(self._emulator_host)
+        create_channel_fn = partial(TransportType.create_channel, *args, **kwargs)
+        # interceptors are handled differently between async and sync, because of differences in the grpc and gapic layers
+        if CrossSync.is_async:
+            # for async, add interceptors to the creation function
+            create_channel_fn = partial(create_channel_fn, interceptors=[_LoggingClientAIOInterceptor()])
+            return _AsyncReplaceableChannel(create_channel_fn)
+        else:
+            # for sync, chain interceptors using grpc.channel.intercept
+            # LoggingClientInterceptor not needed, since it is chained in the gapic layer
+            return TransportType.create_channel(*args, **kwargs)
+
+
+
     @staticmethod
     def _client_version() -> str:
         """
@@ -376,32 +396,11 @@ async def _manage_channel(
                 break
             start_timestamp = time.monotonic()
             # prepare new channel for use
-            # TODO: refactor to avoid using internal references: https://github.com/googleapis/python-bigtable/issues/1094
-            old_channel = self.transport.grpc_channel
-            new_channel = self.transport.create_channel()
-            if CrossSync.is_async:
-                new_channel._unary_unary_interceptors.append(
-                    self.transport._interceptor
-                )
-            else:
-                new_channel = intercept_channel(
-                    new_channel, self.transport._interceptor
-                )
+            new_channel = self.transport.grpc_channel.create_channel()
             await self._ping_and_warm_instances(channel=new_channel)
             # cycle channel out of use, with long grace window before closure
-            self.transport._grpc_channel = new_channel
-            self.transport._logged_channel = new_channel
-            # invalidate caches
-            self.transport._stubs = {}
-            self.transport._prep_wrapped_messages(self.client_info)
-            # give old_channel a chance to complete existing rpcs
-            if CrossSync.is_async:
-                await old_channel.close(grace_period)
-            else:
-                if grace_period:
-                    self._is_closed.wait(grace_period)  # type: ignore
-                old_channel.close()  # type: ignore
-            # subtract thed time spent waiting for the channel to be replaced
+            await self.transport.grpc_channel.replace_channel(new_channel, grace_period)
+            # subtract the time spent waiting for the channel to be replaced
             next_refresh = random.uniform(refresh_interval_min, refresh_interval_max)
             next_sleep = max(next_refresh - (time.monotonic() - start_timestamp), 0)
 

google/cloud/bigtable/data/_async/replaceable_channel.py

@@ -0,0 +1,87 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+#
+from __future__ import annotations
+
+from typing import Callable
+
+import grpc  # type: ignore
+from grpc.experimental import aio  # type: ignore
+
+from google.cloud.bigtable.data._cross_sync import CrossSync
+
+class _AsyncReplaceableChannel(aio.Channel):
+    """
+    A wrapper around a gRPC channel. All methods are passed
+    through to the underlying channel.
+    """
+
+    def __init__(self, channel_fn: Callable[[], aio.Channel]):
+        self._channel_fn = channel_fn
+        self._channel = channel_fn()
+
+    def unary_unary(self, *args, **kwargs):
+        return self._channel.unary_unary(*args, **kwargs)
+
+    def unary_stream(self, *args, **kwargs):
+        return self._channel.unary_stream(*args, **kwargs)
+
+    def stream_unary(self, *args, **kwargs):
+        return self._channel.stream_unary(*args, **kwargs)
+
+    def stream_stream(self, *args, **kwargs):
+        return self._channel.stream_stream(*args, **kwargs)
+
+    async def close(self, grace=None):
+        return await self._channel.close(grace=grace)
+
+    async def channel_ready(self):
+        return await self._channel.channel_ready()
+
+    async def __aenter__(self):
+        await self._channel.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return await self._channel.__aexit__(exc_type, exc_val, exc_tb)
+
+    def get_state(self, try_to_connect: bool = False) -> grpc.ChannelConnectivity:
+        return self._channel.get_state(try_to_connect=try_to_connect)
+
+    async def wait_for_state_change(self, last_observed_state):
+        return await self._channel.wait_for_state_change(last_observed_state)
+
+    @property
+    def wrapped_channel(self):
+        return self._channel
+
+    def create_channel(self) -> aio.Channel:
+        return self._channel_fn()
+
+    async def replace_channel(self, new_channel: aio.Channel, grace_period: float | None) -> aio.Channel:
+        old_channel = self._channel
+        self._channel = new_channel
+        # give old_channel a chance to complete existing rpcs
+        if CrossSync.is_async:
+            await old_channel.close(grace_period)
+        else:
+            if grace_period:
+                self._is_closed.wait(grace_period)  # type: ignore
+            old_channel.close()  # type: ignore
+        return old_channel
+
+    @property
+    def _unary_unary_interceptors(self):
+        # return empty list for compatibility with gapic layer
+        return []