set up channel interceptions

Author: daniel-sanche

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

@@ -19,6 +19,7 @@
     cast,
     Any,
     AsyncIterable,
+    Callable,
     Optional,
     Set,
     Sequence,
@@ -104,6 +105,7 @@
 else:
     from typing import Iterable  # noqa: F401
     from grpc import insecure_channel
+    from grpc import intercept_channel
     from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport as TransportType  # type: ignore
     from google.cloud.bigtable_v2.services.bigtable import BigtableClient as GapicClient  # type: ignore
     from google.cloud.bigtable.data._sync_autogen.mutations_batcher import _MB_SIZE
@@ -206,7 +208,7 @@ def __init__(
                 credentials = google.auth.credentials.AnonymousCredentials()
             if project is None:
                 project = _DEFAULT_BIGTABLE_EMULATOR_CLIENT
-
+        self._metrics_interceptor = MetricInterceptorType()
         # initialize client
         ClientWithProject.__init__(
             self,
@@ -234,7 +236,6 @@ def __init__(
         self._executor: concurrent.futures.ThreadPoolExecutor | None = (
             concurrent.futures.ThreadPoolExecutor() if not CrossSync.is_async else None
         )
-        self._interceptor = MetricInterceptorType()
         if self._emulator_host is None:
             # attempt to start background channel refresh tasks
             try:
@@ -263,12 +264,23 @@ def _build_grpc_channel(self, *args, **kwargs) -> AsyncSwappableChannel:
         Returns:
           a custom wrapped swappable channel
         """
+        create_channel_fn: Callable[[], Channel]
         if self._emulator_host is not None:
             # emulators use insecure channel
             create_channel_fn = partial(insecure_channel, self._emulator_host)
-        else:
+        elif CrossSync.is_async:
             create_channel_fn = partial(TransportType.create_channel, *args, **kwargs)
-        return AsyncSwappableChannel(create_channel_fn)
+        else:
+            # attach sync interceptors in create_channel_fn
+            create_channel_fn = lambda: intercept_channel(
+                TransportType.create_channel(*args, **kwargs), self._metrics_interceptor
+            )
+        new_channel = AsyncSwappableChannel(create_channel_fn)
+        if CrossSync.is_async:
+            # attach async interceptors
+            new_channel._unary_unary_interceptors.append(self._metrics_interceptor)
+            new_channel._unary_stream_interceptors.append(self._metrics_interceptor)
+        return new_channel
 
     @staticmethod
     def _client_version() -> str:
@@ -922,24 +934,12 @@ def __init__(
         )
 
         self._metrics = BigtableClientSideMetricsController(
-            client._interceptor,
+            client._metrics_interceptor,
             project_id=self.client.project,
             instance_id=instance_id,
             table_id=table_id,
             app_profile_id=app_profile_id,
         )
-        # TODO: simplify interceptors
-        if CrossSync.is_async:
-            client.transport.grpc_channel._unary_unary_interceptors.append(
-                self._metrics.interceptor
-            )
-            client.transport.grpc_channel._unary_stream_interceptors.append(
-                self._metrics.interceptor
-            )
-        else:
-            client.transport.grpc_channel = intercept_channel(
-                self._metrics.interceptor, client.transport.grpc_channel
-            )
 
         try:
             self._register_instance_future = CrossSync.create_task(

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

@@ -11,6 +11,7 @@
 # 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
 
 import time
 from typing import Any, Callable

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

@@ -17,7 +17,7 @@
 # This file is automatically generated by CrossSync. Do not edit manually.
 
 from __future__ import annotations
-from typing import cast, Any, Optional, Set, Sequence, TYPE_CHECKING
+from typing import cast, Any, Callable, Optional, Set, Sequence, TYPE_CHECKING
 import abc
 import time
 import warnings
@@ -76,6 +76,7 @@
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from typing import Iterable
 from grpc import insecure_channel
+from grpc import intercept_channel
 from google.cloud.bigtable_v2.services.bigtable.transports import (
     BigtableGrpcTransport as TransportType,
 )
@@ -147,6 +148,7 @@ def __init__(
                 credentials = google.auth.credentials.AnonymousCredentials()
             if project is None:
                 project = _DEFAULT_BIGTABLE_EMULATOR_CLIENT
+        self._metrics_interceptor = MetricInterceptorType()
         ClientWithProject.__init__(
             self,
             credentials=credentials,
@@ -172,7 +174,6 @@ def __init__(
             if not CrossSync._Sync_Impl.is_async
             else None
         )
-        self._interceptor = MetricInterceptorType()
         if self._emulator_host is None:
             try:
                 self._start_background_channel_refresh()
@@ -196,11 +197,15 @@ def _build_grpc_channel(self, *args, **kwargs) -> SwappableChannel:
           - **kwargs: keyword arguments passed by the gapic layer to create a new channel with
         Returns:
           a custom wrapped swappable channel"""
+        create_channel_fn: Callable[[], Channel]
         if self._emulator_host is not None:
             create_channel_fn = partial(insecure_channel, self._emulator_host)
         else:
-            create_channel_fn = partial(TransportType.create_channel, *args, **kwargs)
-        return SwappableChannel(create_channel_fn)
+            create_channel_fn = lambda: intercept_channel(
+                TransportType.create_channel(*args, **kwargs), self._metrics_interceptor
+            )
+        new_channel = SwappableChannel(create_channel_fn)
+        return new_channel
 
     @staticmethod
     def _client_version() -> str:
@@ -723,15 +728,12 @@ def __init__(
             default_retryable_errors or ()
         )
         self._metrics = BigtableClientSideMetricsController(
-            client._interceptor,
+            client._metrics_interceptor,
             project_id=self.client.project,
             instance_id=instance_id,
             table_id=table_id,
             app_profile_id=app_profile_id,
         )
-        client.transport.grpc_channel = intercept_channel(
-            self._metrics.interceptor, client.transport.grpc_channel
-        )
         try:
             self._register_instance_future = CrossSync._Sync_Impl.create_task(
                 self.client._register_instance,

google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 
-
 # This file is automatically generated by CrossSync. Do not edit manually.
 
+from __future__ import annotations
 from functools import wraps
 from google.cloud.bigtable.data._metrics.data_model import (
     OPERATION_INTERCEPTOR_METADATA_KEY,

tests/system/data/test_system_async.py

@@ -292,23 +292,32 @@ async def test_channel_refresh(self, table_id, instance_id, temp_rows):
             async with client.get_table(instance_id, table_id) as table:
                 rows = await table.read_rows({})
                 channel_wrapper = client.transport.grpc_channel
-                first_channel = client.transport.grpc_channel._channel
+                first_channel = channel_wrapper._channel
                 assert len(rows) == 2
                 await CrossSync.sleep(2)
                 rows_after_refresh = await table.read_rows({})
                 assert len(rows_after_refresh) == 2
                 assert client.transport.grpc_channel is channel_wrapper
-                assert client.transport.grpc_channel._channel is not first_channel
-                # ensure gapic's logging interceptor is still active
+                updated_channel = channel_wrapper._channel
+                assert channel_wrapper._channel is not first_channel
+                # ensure interceptors are kept (gapic's logging interceptor, and metric interceptor)
                 if CrossSync.is_async:
-                    interceptors = (
-                        client.transport.grpc_channel._channel._unary_unary_interceptors
+                    unary_interceptors = (
+                        updated_channel._unary_unary_interceptors
                     )
-                    assert GapicInterceptor in [type(i) for i in interceptors]
+                    assert len(unary_interceptors) == 2
+                    assert GapicInterceptor in [type(i) for i in unary_interceptors]
+                    assert client._metrics_interceptor in unary_interceptors
+                    stream_interceptors = (
+                        updated_channel._unary_stream_interceptors
+                    )
+                    assert len(stream_interceptors) == 1
+                    assert client._metrics_interceptor in stream_interceptors
                 else:
                     assert isinstance(
                         client.transport._logged_channel._interceptor, GapicInterceptor
                     )
+                    assert updated_channel._interceptor == client._metrics_interceptor
         finally:
             await client.close()
 

tests/system/data/test_system_autogen.py

@@ -235,16 +235,18 @@ def test_channel_refresh(self, table_id, instance_id, temp_rows):
             with client.get_table(instance_id, table_id) as table:
                 rows = table.read_rows({})
                 channel_wrapper = client.transport.grpc_channel
-                first_channel = client.transport.grpc_channel._channel
+                first_channel = channel_wrapper._channel
                 assert len(rows) == 2
                 CrossSync._Sync_Impl.sleep(2)
                 rows_after_refresh = table.read_rows({})
                 assert len(rows_after_refresh) == 2
                 assert client.transport.grpc_channel is channel_wrapper
-                assert client.transport.grpc_channel._channel is not first_channel
+                updated_channel = channel_wrapper._channel
+                assert channel_wrapper._channel is not first_channel
                 assert isinstance(
                     client.transport._logged_channel._interceptor, GapicInterceptor
                 )
+                assert updated_channel._interceptor == client._metrics_interceptor
         finally:
             client.close()
 

tests/unit/data/_async/test_client.py

@@ -54,18 +54,22 @@
     from google.cloud.bigtable.data._async._swappable_channel import (
         AsyncSwappableChannel,
     )
+    from google.cloud.bigtable.data._async.metrics_interceptor import AsyncBigtableMetricsInterceptor
 
     CrossSync.add_mapping("grpc_helpers", grpc_helpers_async)
     CrossSync.add_mapping("SwappableChannel", AsyncSwappableChannel)
+    CrossSync.add_mapping("MetricsInterceptor", AsyncBigtableMetricsInterceptor)
 else:
     from google.api_core import grpc_helpers
     from google.cloud.bigtable.data._sync_autogen.client import Table  # noqa: F401
     from google.cloud.bigtable.data._sync_autogen._swappable_channel import (
         SwappableChannel,
     )
+    from google.cloud.bigtable.data._sync_autogen.metrics_interceptor import BigtableMetricsInterceptor
 
     CrossSync.add_mapping("grpc_helpers", grpc_helpers)
     CrossSync.add_mapping("SwappableChannel", SwappableChannel)
+    CrossSync.add_mapping("MetricsInterceptor", BigtableMetricsInterceptor)
 
 __CROSS_SYNC_OUTPUT__ = "tests.unit.data._sync_autogen.test_client"
 
@@ -113,6 +117,7 @@ async def test_ctor(self):
         assert not client._active_instances
         assert client._channel_refresh_task is not None
         assert client.transport._credentials == expected_credentials
+        assert isinstance(client._metrics_interceptor, CrossSync.MetricsInterceptor)
         await client.close()
 
     @CrossSync.pytest

tests/unit/data/_sync_autogen/test_client.py

@@ -46,9 +46,13 @@
 )
 from google.api_core import grpc_helpers
 from google.cloud.bigtable.data._sync_autogen._swappable_channel import SwappableChannel
+from google.cloud.bigtable.data._sync_autogen.metrics_interceptor import (
+    BigtableMetricsInterceptor,
+)
 
 CrossSync._Sync_Impl.add_mapping("grpc_helpers", grpc_helpers)
 CrossSync._Sync_Impl.add_mapping("SwappableChannel", SwappableChannel)
+CrossSync._Sync_Impl.add_mapping("MetricsInterceptor", BigtableMetricsInterceptor)
 
 
 @CrossSync._Sync_Impl.add_mapping_decorator("TestBigtableDataClient")
@@ -84,6 +88,9 @@ def test_ctor(self):
         assert not client._active_instances
         assert client._channel_refresh_task is not None
         assert client.transport._credentials == expected_credentials
+        assert isinstance(
+            client._metrics_interceptor, CrossSync._Sync_Impl.MetricsInterceptor
+        )
         client.close()
 
     def test_ctor_super_inits(self):