feat(bigtable): client side metrics handlers (#16760)

Migrate googleapis/python-bigtable#1189 to the
monorepo

This PR builds off of
googleapis/python-bigtable#1187 to add handlers
to the client-side metrics system, which can subscribe to the metrics
stream, and export the results into different collection systems


We add two handlers to the system:
- `GoogleCloudMetricsHandler`: sends metrics to a private OpenTelemetry
meter, and then periodically exports them to GCP. Built on top of
`OpenTelemetryMetricsHandler`
- `OpenTelemetryMetricsHandler`: sends metrics to the root
MeterProvider, so the user can access the exported metrics for their own
systems. This will be off by default, but can be added alongside
`GoogleCloudMetricsHandler` if needed

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: daniel-sanche

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,16 @@
     _WarmedInstanceKey,
 )
 from google.cloud.bigtable.data._metrics import (
+    ActiveOperationMetric,
     BigtableClientSideMetricsController,
     OperationType,
     tracked_retry,
 )
+from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
+from google.cloud.bigtable.data._metrics.handlers.gcp_exporter import (
+    BigtableMetricsExporter,
+    GoogleCloudMetricsHandler,
+)
 from google.cloud.bigtable.data.exceptions import (
     FailedQueryShardError,
     ShardedReadRowsExceptionGroup,
@@ -159,6 +166,8 @@
 
 __CROSS_SYNC_OUTPUT__ = "google.cloud.bigtable.data._sync_autogen.client"
 
+_LOGGER = logging.getLogger(__name__)
+
 
 @CrossSync.convert_class(
     sync_name="BigtableDataClient",
@@ -262,6 +271,27 @@ def __init__(
                 "is the default."
             )
         self._is_closed = CrossSync.Event()
+        handlers: list[MetricsHandler] = []
+        if self._emulator_host is None:
+            try:
+                # create a metrics exporter using the same client configuration
+                exporter = BigtableMetricsExporter(
+                    credentials=self._credentials,
+                    client_options=client_options,
+                )
+                handlers.append(
+                    GoogleCloudMetricsHandler(
+                        exporter=exporter,
+                        client_version=self._client_version(),
+                    )
+                )
+            except Exception as e:
+                _LOGGER.warning(
+                    "Failed to initialize Google Cloud Metrics Exporter: %s. "
+                    "Client-side metrics will be disabled.",
+                    e,
+                )
+        self._metrics = BigtableClientSideMetricsController(handlers=handlers)
         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 +424,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 +1140,6 @@ def __init__(
             default_retryable_errors or ()
         )
 
-        self._metrics = BigtableClientSideMetricsController()
-
         try:
             self._register_instance_future = CrossSync.create_task(
                 self.client._register_instance,
@@ -1124,6 +1153,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 +1233,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 +1337,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 +1552,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 +1684,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 +1758,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 +1817,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 +1875,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 +1896,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,
         )