feat: added metrics interceptor

Author: daniel-sanche

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

@@ -97,18 +97,24 @@
     )
     from google.cloud.bigtable.data._async.mutations_batcher import _MB_SIZE
     from google.cloud.bigtable.data._async._swappable_channel import (
-        AsyncSwappableChannel,
+        AsyncSwappableChannel as SwappableChannelType,
+    )
+    from google.cloud.bigtable.data._async.metrics_interceptor import (
+        AsyncBigtableMetricsInterceptor as MetricInterceptorType,
     )
 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
     from google.cloud.bigtable.data._sync_autogen._swappable_channel import (  # noqa: F401
-        SwappableChannel,
+        SwappableChannel as SwappableChannelType,
+    )
+    from google.cloud.bigtable.data._sync_autogen.metrics_interceptor import (  # noqa: F401
+        BigtableMetricsInterceptor as MetricInterceptorType,
     )
-
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._helpers import RowKeySamples
@@ -203,7 +209,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,
@@ -257,12 +263,11 @@ def __init__(
                     stacklevel=2,
                 )
 
-    @CrossSync.convert(replace_symbols={"AsyncSwappableChannel": "SwappableChannel"})
-    def _build_grpc_channel(self, *args, **kwargs) -> AsyncSwappableChannel:
+    def _build_grpc_channel(self, *args, **kwargs) -> SwappableChannelType:
         """
         This method is called by the gapic transport to create a grpc channel.
 
-        The init arguments passed down are captured in a partial used by AsyncSwappableChannel
+        The init arguments passed down are captured in a partial used by SwappableChannel
         to create new channel instances in the future, as part of the channel refresh logic
 
         Emulators always use an inseucre channel
@@ -276,9 +281,22 @@ def _build_grpc_channel(self, *args, **kwargs) -> AsyncSwappableChannel:
         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
+            def create_channel_fn():
+                return intercept_channel(
+                    TransportType.create_channel(*args, **kwargs),
+                    self._metrics_interceptor,
+                )
+
+        new_channel = SwappableChannelType(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
 
     @property
     def universe_domain(self) -> str:
@@ -400,7 +418,6 @@ def _invalidate_channel_stubs(self):
         self.transport._stubs = {}
         self.transport._prep_wrapped_messages(self.client_info)
 
-    @CrossSync.convert(replace_symbols={"AsyncSwappableChannel": "SwappableChannel"})
     async def _manage_channel(
         self,
         refresh_interval_min: float = 60 * 35,
@@ -425,10 +442,10 @@ async def _manage_channel(
             grace_period: time to allow previous channel to serve existing
                 requests before closing, in seconds
         """
-        if not isinstance(self.transport.grpc_channel, AsyncSwappableChannel):
+        if not isinstance(self.transport.grpc_channel, SwappableChannelType):
             warnings.warn("Channel does not support auto-refresh.")
             return
-        super_channel: AsyncSwappableChannel = self.transport.grpc_channel
+        super_channel: SwappableChannelType = self.transport.grpc_channel
         first_refresh = self._channel_init_time + random.uniform(
             refresh_interval_min, refresh_interval_max
         )

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

@@ -0,0 +1,85 @@
+# Copyright 2025 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 google.cloud.bigtable.data._cross_sync import CrossSync
+
+if CrossSync.is_async:
+    from grpc.aio import UnaryUnaryClientInterceptor
+    from grpc.aio import UnaryStreamClientInterceptor
+else:
+    from grpc import UnaryUnaryClientInterceptor
+    from grpc import UnaryStreamClientInterceptor
+
+
+__CROSS_SYNC_OUTPUT__ = "google.cloud.bigtable.data._sync_autogen.metrics_interceptor"
+
+
+@CrossSync.convert_class(sync_name="BigtableMetricsInterceptor")
+class AsyncBigtableMetricsInterceptor(
+    UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor
+):
+    """
+    An async gRPC interceptor to add client metadata and print server metadata.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.operation_map = {}
+
+    def register_operation(self, operation):
+        """
+        Register an operation object to be tracked my the interceptor
+
+        When registered, the operation will receive metadata updates:
+        - start_attempt if attempt not started when rpc is being sent
+        - add_response_metadata after call is complete
+        - end_attempt_with_status if attempt receives an error
+
+        The interceptor will register itself as a handeler for the operation,
+        so it can unregister the operation when it is complete
+        """
+        self.operation_map[operation.uuid] = operation
+        operation.handlers.append(self)
+
+    def on_operation_complete(self, op):
+        if op.uuid in self.operation_map:
+            del self.operation_map[op.uuid]
+
+    def on_operation_cancelled(self, op):
+        self.on_operation_complete(op)
+
+    @CrossSync.convert
+    async def intercept_unary_unary(self, continuation, client_call_details, request):
+        try:
+            call = await continuation(client_call_details, request)
+            return call
+        except Exception as rpc_error:
+            raise rpc_error
+
+    @CrossSync.convert
+    async def intercept_unary_stream(self, continuation, client_call_details, request):
+        async def response_wrapper(call):
+            try:
+                async for response in call:
+                    yield response
+            except Exception as e:
+                # handle errors while processing stream
+                raise e
+
+        try:
+            return response_wrapper(await continuation(client_call_details, request))
+        except Exception as rpc_error:
+            # handle errors while intializing stream
+            raise rpc_error

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

@@ -75,12 +75,18 @@
 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,
 )
 from google.cloud.bigtable_v2.services.bigtable import BigtableClient as GapicClient
 from google.cloud.bigtable.data._sync_autogen.mutations_batcher import _MB_SIZE
-from google.cloud.bigtable.data._sync_autogen._swappable_channel import SwappableChannel
+from google.cloud.bigtable.data._sync_autogen._swappable_channel import (
+    SwappableChannel as SwappableChannelType,
+)
+from google.cloud.bigtable.data._sync_autogen.metrics_interceptor import (
+    BigtableMetricsInterceptor as MetricInterceptorType,
+)
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._helpers import RowKeySamples
@@ -143,6 +149,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,
@@ -186,7 +193,7 @@ def __init__(
                     stacklevel=2,
                 )
 
-    def _build_grpc_channel(self, *args, **kwargs) -> SwappableChannel:
+    def _build_grpc_channel(self, *args, **kwargs) -> SwappableChannelType:
         """This method is called by the gapic transport to create a grpc channel.
 
         The init arguments passed down are captured in a partial used by SwappableChannel
@@ -202,8 +209,15 @@ def _build_grpc_channel(self, *args, **kwargs) -> SwappableChannel:
         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)
+
+            def create_channel_fn():
+                return intercept_channel(
+                    TransportType.create_channel(*args, **kwargs),
+                    self._metrics_interceptor,
+                )
+
+        new_channel = SwappableChannelType(create_channel_fn)
+        return new_channel
 
     @property
     def universe_domain(self) -> str:
@@ -302,7 +316,7 @@ def _invalidate_channel_stubs(self):
         self.transport._stubs = {}
         self.transport._prep_wrapped_messages(self.client_info)
 
-    def _manage_channel(
+    async def _manage_channel(
         self,
         refresh_interval_min: float = 60 * 35,
         refresh_interval_max: float = 60 * 45,
@@ -324,25 +338,25 @@ def _manage_channel(
                 between `refresh_interval_min` and `refresh_interval_max`
             grace_period: time to allow previous channel to serve existing
                 requests before closing, in seconds"""
-        if not isinstance(self.transport.grpc_channel, SwappableChannel):
+        if not isinstance(self.transport.grpc_channel, SwappableChannelType):
             warnings.warn("Channel does not support auto-refresh.")
             return
-        super_channel: SwappableChannel = self.transport.grpc_channel
+        super_channel: SwappableChannelType = self.transport.grpc_channel
         first_refresh = self._channel_init_time + random.uniform(
             refresh_interval_min, refresh_interval_max
         )
         next_sleep = max(first_refresh - time.monotonic(), 0)
         if next_sleep > 0:
-            self._ping_and_warm_instances(channel=super_channel)
+            await self._ping_and_warm_instances(channel=super_channel)
         while not self._is_closed.is_set():
-            CrossSync._Sync_Impl.event_wait(
+            await CrossSync._Sync_Impl.event_wait(
                 self._is_closed, next_sleep, async_break_early=False
             )
             if self._is_closed.is_set():
                 break
             start_timestamp = time.monotonic()
             new_channel = super_channel.create_channel()
-            self._ping_and_warm_instances(channel=new_channel)
+            await self._ping_and_warm_instances(channel=new_channel)
             old_channel = super_channel.swap_channel(new_channel)
             self._invalidate_channel_stubs()
             if grace_period:

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

@@ -0,0 +1,71 @@
+# Copyright 2025 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
+
+# This file is automatically generated by CrossSync. Do not edit manually.
+
+from __future__ import annotations
+from grpc import UnaryUnaryClientInterceptor
+from grpc import UnaryStreamClientInterceptor
+
+
+class BigtableMetricsInterceptor(
+    UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor
+):
+    """
+    An async gRPC interceptor to add client metadata and print server metadata.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.operation_map = {}
+
+    def register_operation(self, operation):
+        """Register an operation object to be tracked my the interceptor
+
+        When registered, the operation will receive metadata updates:
+        - start_attempt if attempt not started when rpc is being sent
+        - add_response_metadata after call is complete
+        - end_attempt_with_status if attempt receives an error
+
+        The interceptor will register itself as a handeler for the operation,
+        so it can unregister the operation when it is complete"""
+        self.operation_map[operation.uuid] = operation
+        operation.handlers.append(self)
+
+    def on_operation_complete(self, op):
+        if op.uuid in self.operation_map:
+            del self.operation_map[op.uuid]
+
+    def on_operation_cancelled(self, op):
+        self.on_operation_complete(op)
+
+    def intercept_unary_unary(self, continuation, client_call_details, request):
+        try:
+            call = continuation(client_call_details, request)
+            return call
+        except Exception as rpc_error:
+            raise rpc_error
+
+    def intercept_unary_stream(self, continuation, client_call_details, request):
+        def response_wrapper(call):
+            try:
+                for response in call:
+                    yield response
+            except Exception as e:
+                raise e
+
+        try:
+            return response_wrapper(continuation(client_call_details, request))
+        except Exception as rpc_error:
+            raise rpc_error