feat: add basic interceptor to client (#1206)

Author: daniel-sanche

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

@@ -19,6 +19,7 @@
 cast,
 Any,
 AsyncIterable,
+Callable,
 Optional,
 Set,
 Sequence,
@@ -99,18 +100,24 @@
  )
 fromgoogle.cloud.bigtable.data._async.mutations_batcherimport_MB_SIZE
 fromgoogle.cloud.bigtable.data._async._swappable_channelimport (
-AsyncSwappableChannel,
+AsyncSwappableChannelasSwappableChannelType,
+ )
+fromgoogle.cloud.bigtable.data._async.metrics_interceptorimport (
+AsyncBigtableMetricsInterceptorasMetricsInterceptorType,
  )
 else:
 fromtypingimportIterable# noqa: F401
 fromgrpcimportinsecure_channel
+fromgrpcimportintercept_channel
 fromgoogle.cloud.bigtable_v2.services.bigtable.transportsimportBigtableGrpcTransportasTransportType# type: ignore
 fromgoogle.cloud.bigtable_v2.services.bigtableimportBigtableClientasGapicClient# type: ignore
 fromgoogle.cloud.bigtable.data._sync_autogen.mutations_batcherimport_MB_SIZE
 fromgoogle.cloud.bigtable.data._sync_autogen._swappable_channelimport ( # noqa: F401
-SwappableChannel,
+SwappableChannelasSwappableChannelType,
+ )
+fromgoogle.cloud.bigtable.data._sync_autogen.metrics_interceptorimport ( # noqa: F401
+BigtableMetricsInterceptorasMetricsInterceptorType,
  )
-
 
 ifTYPE_CHECKING:
 fromgoogle.cloud.bigtable.data._helpersimportRowKeySamples
@@ -205,7 +212,7 @@ def __init__(
 credentials=google.auth.credentials.AnonymousCredentials()
 ifprojectisNone:
 project=_DEFAULT_BIGTABLE_EMULATOR_CLIENT
-
+self._metrics_interceptor=MetricsInterceptorType()
 # initialize client
 ClientWithProject.__init__(
 self,
@@ -259,12 +266,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
@@ -275,12 +281,30 @@ def _build_grpc_channel(self, *args, **kwargs) -> AsyncSwappableChannel:
  Returns:
  a custom wrapped swappable channel
  """
+create_channel_fn: Callable[[], Channel]
 ifself._emulator_hostisnotNone:
-# emulators use insecure channel
+# Emulators use insecure channels
 create_channel_fn=partial(insecure_channel, self._emulator_host)
-else:
+elifCrossSync.is_async:
+# For async client, use the default create_channel.
 create_channel_fn=partial(TransportType.create_channel, *args, **kwargs)
-returnAsyncSwappableChannel(create_channel_fn)
+else:
+# For sync client, wrap create_channel with interceptors.
+defsync_create_channel_fn():
+returnintercept_channel(
+TransportType.create_channel(*args, **kwargs),
+self._metrics_interceptor,
+ )
+
+create_channel_fn=sync_create_channel_fn
+
+# Instantiate SwappableChannelType with the determined creation function.
+new_channel=SwappableChannelType(create_channel_fn)
+ifCrossSync.is_async:
+# Attach async interceptors to the channel instance itself.
+new_channel._unary_unary_interceptors.append(self._metrics_interceptor)
+new_channel._unary_stream_interceptors.append(self._metrics_interceptor)
+returnnew_channel
 
 @property
 defuniverse_domain(self) ->str:
@@ -402,7 +426,7 @@ def _invalidate_channel_stubs(self):
 self.transport._stubs= {}
 self.transport._prep_wrapped_messages(self.client_info)
 
-@CrossSync.convert(replace_symbols={"AsyncSwappableChannel": "SwappableChannel"})
+@CrossSync.convert
 asyncdef_manage_channel(
 self,
 refresh_interval_min: float=60*35,
@@ -427,10 +451,10 @@ async def _manage_channel(
  grace_period: time to allow previous channel to serve existing
  requests before closing, in seconds
  """
-ifnotisinstance(self.transport.grpc_channel, AsyncSwappableChannel):
+ifnotisinstance(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,78 @@
+# 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__ importannotations
+
+fromgoogle.cloud.bigtable.data._cross_syncimportCrossSync
+
+ifCrossSync.is_async:
+fromgrpc.aioimportUnaryUnaryClientInterceptor
+fromgrpc.aioimportUnaryStreamClientInterceptor
+else:
+fromgrpcimportUnaryUnaryClientInterceptor
+fromgrpcimportUnaryStreamClientInterceptor
+
+
+__CROSS_SYNC_OUTPUT__="google.cloud.bigtable.data._sync_autogen.metrics_interceptor"
+
+
+@CrossSync.convert_class(sync_name="BigtableMetricsInterceptor")
+classAsyncBigtableMetricsInterceptor(
+UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor
+):
+"""
+ An async gRPC interceptor to add client metadata and print server metadata.
+ """
+
+@CrossSync.convert
+asyncdefintercept_unary_unary(self, continuation, client_call_details, request):
+"""
+ Interceptor for unary rpcs:
+ - MutateRow
+ - CheckAndMutateRow
+ - ReadModifyWriteRow
+ """
+try:
+call=awaitcontinuation(client_call_details, request)
+returncall
+exceptExceptionasrpc_error:
+raiserpc_error
+
+@CrossSync.convert
+asyncdefintercept_unary_stream(self, continuation, client_call_details, request):
+"""
+ Interceptor for streaming rpcs:
+ - ReadRows
+ - MutateRows
+ - SampleRowKeys
+ """
+try:
+returnself._streaming_generator_wrapper(
+awaitcontinuation(client_call_details, request)
+ )
+exceptExceptionasrpc_error:
+# handle errors while intializing stream
+raiserpc_error
+
+@staticmethod
+@CrossSync.convert
+asyncdef_streaming_generator_wrapper(call):
+"""
+ Wrapped generator to be returned by intercept_unary_stream.
+ """
+try:
+asyncforresponseincall:
+yieldresponse
+exceptExceptionase:
+# handle errors while processing stream
+raisee

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__ importannotations
-fromtypingimportcast, Any, Optional, Set, Sequence, TYPE_CHECKING
+fromtypingimportcast, Any, Callable, Optional, Set, Sequence, TYPE_CHECKING
 importabc
 importtime
 importwarnings
@@ -77,12 +77,18 @@
 fromgoogle.cloud.bigtable.data._cross_syncimportCrossSync
 fromtypingimportIterable
 fromgrpcimportinsecure_channel
+fromgrpcimportintercept_channel
 fromgoogle.cloud.bigtable_v2.services.bigtable.transportsimport (
 BigtableGrpcTransportasTransportType,
 )
 fromgoogle.cloud.bigtable_v2.services.bigtableimportBigtableClientasGapicClient
 fromgoogle.cloud.bigtable.data._sync_autogen.mutations_batcherimport_MB_SIZE
-fromgoogle.cloud.bigtable.data._sync_autogen._swappable_channelimportSwappableChannel
+fromgoogle.cloud.bigtable.data._sync_autogen._swappable_channelimport (
+SwappableChannelasSwappableChannelType,
+)
+fromgoogle.cloud.bigtable.data._sync_autogen.metrics_interceptorimport (
+BigtableMetricsInterceptorasMetricsInterceptorType,
+)
 
 ifTYPE_CHECKING:
 fromgoogle.cloud.bigtable.data._helpersimportRowKeySamples
@@ -145,6 +151,7 @@ def __init__(
 credentials=google.auth.credentials.AnonymousCredentials()
 ifprojectisNone:
 project=_DEFAULT_BIGTABLE_EMULATOR_CLIENT
+self._metrics_interceptor=MetricsInterceptorType()
 ClientWithProject.__init__(
 self,
 credentials=credentials,
@@ -188,7 +195,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
@@ -201,11 +208,20 @@ 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]
 ifself._emulator_hostisnotNone:
 create_channel_fn=partial(insecure_channel, self._emulator_host)
 else:
-create_channel_fn=partial(TransportType.create_channel, *args, **kwargs)
-returnSwappableChannel(create_channel_fn)
+
+defsync_create_channel_fn():
+returnintercept_channel(
+TransportType.create_channel(*args, **kwargs),
+self._metrics_interceptor,
+ )
+
+create_channel_fn=sync_create_channel_fn
+new_channel=SwappableChannelType(create_channel_fn)
+returnnew_channel
 
 @property
 defuniverse_domain(self) ->str:
@@ -326,10 +342,10 @@ 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"""
-ifnotisinstance(self.transport.grpc_channel, SwappableChannel):
+ifnotisinstance(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
  )

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

@@ -0,0 +1,59 @@
+# 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__ importannotations
+fromgrpcimportUnaryUnaryClientInterceptor
+fromgrpcimportUnaryStreamClientInterceptor
+
+
+classBigtableMetricsInterceptor(
+UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor
+):
+"""
+ An async gRPC interceptor to add client metadata and print server metadata.
+ """
+
+defintercept_unary_unary(self, continuation, client_call_details, request):
+"""Interceptor for unary rpcs:
+ - MutateRow
+ - CheckAndMutateRow
+ - ReadModifyWriteRow"""
+try:
+call=continuation(client_call_details, request)
+returncall
+exceptExceptionasrpc_error:
+raiserpc_error
+
+defintercept_unary_stream(self, continuation, client_call_details, request):
+"""Interceptor for streaming rpcs:
+ - ReadRows
+ - MutateRows
+ - SampleRowKeys"""
+try:
+returnself._streaming_generator_wrapper(
+continuation(client_call_details, request)
+ )
+exceptExceptionasrpc_error:
+raiserpc_error
+
+@staticmethod
+def_streaming_generator_wrapper(call):
+"""Wrapped generator to be returned by intercept_unary_stream."""
+try:
+forresponseincall:
+yieldresponse
+exceptExceptionase:
+raisee

tests/system/data/test_system_async.py

@@ -285,23 +285,28 @@ async def test_channel_refresh(self, table_id, instance_id, temp_rows):
 asyncwithclient.get_table(instance_id, table_id) astable:
 rows=awaittable.read_rows({})
 channel_wrapper=client.transport.grpc_channel
-first_channel=client.transport.grpc_channel._channel
+first_channel=channel_wrapper._channel
 assertlen(rows) ==2
 awaitCrossSync.sleep(2)
 rows_after_refresh=awaittable.read_rows({})
 assertlen(rows_after_refresh) ==2
 assertclient.transport.grpc_channelischannel_wrapper
-assertclient.transport.grpc_channel._channelisnotfirst_channel
-# ensure gapic's logging interceptor is still active
+updated_channel=channel_wrapper._channel
+assertupdated_channelisnotfirst_channel
+# ensure interceptors are kept (gapic's logging interceptor, and metric interceptor)
 ifCrossSync.is_async:
-interceptors= (
-client.transport.grpc_channel._channel._unary_unary_interceptors
- )
-assertGapicInterceptorin [type(i) foriininterceptors]
+unary_interceptors=updated_channel._unary_unary_interceptors
+assertlen(unary_interceptors) ==2
+assertGapicInterceptorin [type(i) foriinunary_interceptors]
+assertclient._metrics_interceptorinunary_interceptors
+stream_interceptors=updated_channel._unary_stream_interceptors
+assertlen(stream_interceptors) ==1
+assertclient._metrics_interceptorinstream_interceptors
 else:
 assertisinstance(
 client.transport._logged_channel._interceptor, GapicInterceptor
  )
+assertupdated_channel._interceptor==client._metrics_interceptor
 finally:
 awaitclient.close()
 

tests/system/data/test_system_autogen.py

@@ -237,16 +237,18 @@ def test_channel_refresh(self, table_id, instance_id, temp_rows):
 withclient.get_table(instance_id, table_id) astable:
 rows=table.read_rows({})
 channel_wrapper=client.transport.grpc_channel
-first_channel=client.transport.grpc_channel._channel
+first_channel=channel_wrapper._channel
 assertlen(rows) ==2
 CrossSync._Sync_Impl.sleep(2)
 rows_after_refresh=table.read_rows({})
 assertlen(rows_after_refresh) ==2
 assertclient.transport.grpc_channelischannel_wrapper
-assertclient.transport.grpc_channel._channelisnotfirst_channel
+updated_channel=channel_wrapper._channel
+assertupdated_channelisnotfirst_channel
 assertisinstance(
 client.transport._logged_channel._interceptor, GapicInterceptor
  )
+assertupdated_channel._interceptor==client._metrics_interceptor
 finally:
 client.close()