feat: client side metrics handlers

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

@@ -17,6 +17,7 @@
 
 import abc
 import concurrent.futures
+import logging
 import os
 import random
 import time
@@ -64,10 +65,18 @@
     _WarmedInstanceKey,
 )
 from google.cloud.bigtable.data._metrics import (
+    ActiveOperationMetric,
     BigtableClientSideMetricsController,
     OperationType,
     tracked_retry,
 )
+from google.cloud.bigtable.data._metrics.handlers.gcp_exporter import (
+    BigtableMetricsExporter,
+    GoogleCloudMetricsHandler,
+)
+from google.cloud.bigtable.data._metrics.handlers.opentelemetry import (
+    OpenTelemetryMetricsHandler,
+)
 from google.cloud.bigtable.data.exceptions import (
     FailedQueryShardError,
     ShardedReadRowsExceptionGroup,
@@ -103,6 +112,8 @@
     SampleRowKeysRequest,
 )
 
+_LOGGER = logging.getLogger(__name__)
+
 if CrossSync.is_async:
     from grpc.aio import insecure_channel
 
@@ -262,6 +273,34 @@ def __init__(
                 "is the default."
             )
         self._is_closed = CrossSync.Event()
+        if os.getenv("BIGTABLE_EMULATOR_HOST"):
+            self._metrics_handler = OpenTelemetryMetricsHandler(
+                client_version=self._client_version(),
+            )
+        else:
+            try:
+                # create a metrics exporter using the same client configuration
+                exporter = BigtableMetricsExporter(
+                    project_id=self.project,
+                    credentials=credentials,
+                    client_options=client_options,
+                )
+                self._metrics_handler = GoogleCloudMetricsHandler(
+                    exporter=exporter,
+                    client_version=self._client_version(),
+                )
+            except Exception as e:
+                _LOGGER.warning(
+                    "Failed to initialize Google Cloud Metrics Exporter: %s. "
+                    "Falling back to local OpenTelemetry metrics handler.",
+                    e,
+                )
+                self._metrics_handler = OpenTelemetryMetricsHandler(
+                    client_version=self._client_version(),
+                )
+        self._metrics = BigtableClientSideMetricsController(
+            handlers=[self._metrics_handler]
+        )
         self.transport = cast(TransportType, self._gapic_client.transport)
         # keep track of active instances to for warmup on channel refresh
         self._active_instances: Set[_WarmedInstanceKey] = set()
@@ -394,6 +433,7 @@ async def close(self, timeout: float | None = 2.0):
         if self._executor:
             self._executor.shutdown(wait=False)
         self._channel_refresh_task = None
+        self._metrics.close()
 
     @CrossSync.convert
     async def _ping_and_warm_instances(
@@ -1109,8 +1149,6 @@ def __init__(
             default_retryable_errors or ()
         )
 
-        self._metrics = BigtableClientSideMetricsController()
-
         try:
             self._register_instance_future = CrossSync.create_task(
                 self.client._register_instance,
@@ -1124,6 +1162,21 @@ def __init__(
                 f"{self.__class__.__name__} must be created within an async event loop context."
             ) from e
 
+    def _create_operation(
+        self, op_type: OperationType, **kwargs
+    ) -> ActiveOperationMetric:
+        table_id = getattr(self, "table_id", None) or getattr(
+            self, "materialized_view_id", None
+        )
+        return self.client._metrics.create_operation(
+            op_type,
+            project_id=self.client.project,
+            instance_id=self.instance_id,
+            table_id=table_id,
+            app_profile_id=self.app_profile_id,
+            **kwargs,
+        )
+
     @property
     @abc.abstractmethod
     def _request_path(self) -> dict[str, str]:
@@ -1189,9 +1242,7 @@ async def read_rows_stream(
             self,
             operation_timeout=operation_timeout,
             attempt_timeout=attempt_timeout,
-            metric=self._metrics.create_operation(
-                OperationType.READ_ROWS, is_streaming=True
-            ),
+            metric=self._create_operation(OperationType.READ_ROWS, is_streaming=True),
             retryable_exceptions=retryable_excs,
         )
         return row_merger.start_operation()
@@ -1295,9 +1346,7 @@ async def read_row(
             self,
             operation_timeout=operation_timeout,
             attempt_timeout=attempt_timeout,
-            metric=self._metrics.create_operation(
-                OperationType.READ_ROWS, is_streaming=False
-            ),
+            metric=self._create_operation(OperationType.READ_ROWS, is_streaming=False),
             retryable_exceptions=retryable_excs,
         )
         results_generator = row_merger.start_operation()
@@ -1512,9 +1561,7 @@ async def sample_row_keys(
         retryable_excs = _get_retryable_errors(retryable_errors, self)
         predicate = retries.if_exception_type(*retryable_excs)
 
-        with self._metrics.create_operation(
-            OperationType.SAMPLE_ROW_KEYS
-        ) as operation_metric:
+        with self._create_operation(OperationType.SAMPLE_ROW_KEYS) as operation_metric:
 
             @CrossSync.convert
             async def execute_rpc():
@@ -1646,9 +1693,7 @@ async def mutate_row(
             # mutations should not be retried
             predicate = retries.if_exception_type()
 
-        with self._metrics.create_operation(
-            OperationType.MUTATE_ROW
-        ) as operation_metric:
+        with self._create_operation(OperationType.MUTATE_ROW) as operation_metric:
             target = partial(
                 self.client._gapic_client.mutate_row,
                 request=MutateRowRequest(
@@ -1722,7 +1767,7 @@ async def bulk_mutate_rows(
             mutation_entries,
             operation_timeout,
             attempt_timeout,
-            metric=self._metrics.create_operation(OperationType.BULK_MUTATE_ROWS),
+            metric=self._create_operation(OperationType.BULK_MUTATE_ROWS),
             retryable_exceptions=retryable_excs,
         )
         await operation.start()
@@ -1781,7 +1826,7 @@ async def check_and_mutate_row(
             false_case_mutations = [false_case_mutations]
         false_case_list = [m._to_pb() for m in false_case_mutations or []]
 
-        with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE):
+        with self._create_operation(OperationType.CHECK_AND_MUTATE):
             result = await self.client._gapic_client.check_and_mutate_row(
                 request=CheckAndMutateRowRequest(
                     true_mutations=true_case_list,
@@ -1839,7 +1884,7 @@ async def read_modify_write_row(
         if not rules:
             raise ValueError("rules must contain at least one item")
 
-        with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE):
+        with self._create_operation(OperationType.READ_MODIFY_WRITE):
             result = await self.client._gapic_client.read_modify_write_row(
                 request=ReadModifyWriteRowRequest(
                     rules=[rule._to_pb() for rule in rules],
@@ -1860,7 +1905,6 @@ async def close(self):
         """
         Called to close the Table instance and release any resources held by it.
         """
-        self._metrics.close()
         if self._register_instance_future:
             self._register_instance_future.cancel()
         self.client._remove_instance_registration(

packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py

@@ -379,7 +379,7 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]):
         # flush new entries
         in_process_requests: list[CrossSync.Future[list[FailedMutationEntryError]]] = []
         async for batch, metric in self._flow_control.add_to_flow_with_metrics(
-            new_entries, self._target._metrics
+            new_entries, self._target.client._metrics
         ):
             batch_task = CrossSync.create_task(
                 self._execute_mutate_rows,

packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/__init__.py

@@ -19,13 +19,21 @@
     OperationState,
     OperationType,
 )
+from google.cloud.bigtable.data._metrics.handlers.gcp_exporter import (
+    GoogleCloudMetricsHandler,
+)
+from google.cloud.bigtable.data._metrics.handlers.opentelemetry import (
+    OpenTelemetryMetricsHandler,
+)
 from google.cloud.bigtable.data._metrics.metrics_controller import (
     BigtableClientSideMetricsController,
 )
 from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
 
 __all__ = (
     "BigtableClientSideMetricsController",
+    "OpenTelemetryMetricsHandler",
+    "GoogleCloudMetricsHandler",
     "OperationType",
     "OperationState",
     "ActiveOperationMetric",

packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/data_model.py

@@ -118,6 +118,10 @@ class CompletedOperationMetric:
     cluster_id: str
     zone: str
     is_streaming: bool
+    project_id: str | None = None
+    instance_id: str | None = None
+    table_id: str | None = None
+    app_profile_id: str | None = None
     first_response_latency_ns: int | None = None
     flow_throttling_time_ns: int = 0
 
@@ -160,6 +164,10 @@ class ActiveOperationMetric:
     active_attempt: ActiveAttemptMetric | None = None
     cluster_id: str | None = None
     zone: str | None = None
+    project_id: str | None = None
+    instance_id: str | None = None
+    table_id: str | None = None
+    app_profile_id: str | None = None
     completed_attempts: list[CompletedAttemptMetric] = field(default_factory=list)
     is_streaming: bool = False  # only True for read_rows operations
     handlers: list[MetricsHandler] = field(default_factory=list)
@@ -375,6 +383,10 @@ def end_with_status(self, status: StatusCode | BaseException) -> None:
             cluster_id=self.cluster_id or DEFAULT_CLUSTER_ID,
             zone=self.zone or DEFAULT_ZONE,
             is_streaming=self.is_streaming,
+            project_id=self.project_id,
+            instance_id=self.instance_id,
+            table_id=self.table_id,
+            app_profile_id=self.app_profile_id,
             first_response_latency_ns=self.first_response_latency_ns,
             flow_throttling_time_ns=self.flow_throttling_time_ns,
         )