replace custom metadata with contextvars

Author: daniel-sanche

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

@@ -188,7 +188,6 @@ async def _run_attempt(self):
                 ),
                 timeout=next(self.timeout_generator),
                 retry=None,
-                metadata=[self._operation_metric.interceptor_metadata],
             )
             async for result_list in result_generator:
                 for result in result_list.entries:

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

@@ -159,7 +159,6 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
             self.request,
             timeout=next(self.attempt_timeout_gen),
             retry=None,
-            metadata=[self._operation_metric.interceptor_metadata],
         )
         chunked_stream = self.chunk_stream(gapic_stream)
         return self.merge_rows(chunked_stream)

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

@@ -972,7 +972,6 @@ def __init__(
         )
 
         self._metrics = BigtableClientSideMetricsController(
-            client._metrics_interceptor,
             handlers=[],
             project_id=self.client.project,
             instance_id=instance_id,
@@ -1382,7 +1381,6 @@ async def execute_rpc():
                     ),
                     timeout=next(attempt_timeout_gen),
                     retry=None,
-                    metadata=[operation_metric.interceptor_metadata],
                 )
                 return [(s.row_key, s.offset_bytes) async for s in results]
 
@@ -1519,7 +1517,6 @@ async def mutate_row(
                 ),
                 timeout=attempt_timeout,
                 retry=None,
-                metadata=[operation_metric.interceptor_metadata],
             )
             return await CrossSync.retry_target(
                 target,
@@ -1657,7 +1654,6 @@ async def check_and_mutate_row(
                 ),
                 timeout=operation_timeout,
                 retry=None,
-                metadata=[op.interceptor_metadata],
             )
             return result.predicate_matched
 
@@ -1712,7 +1708,6 @@ async def read_modify_write_row(
                 ),
                 timeout=operation_timeout,
                 retry=None,
-                metadata=[op.interceptor_metadata],
             )
             # construct Row from result
             return Row._from_pb(result.row)

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

@@ -19,9 +19,6 @@
 from functools import wraps
 from grpc import StatusCode
 
-from google.cloud.bigtable.data._metrics.data_model import (
-    OPERATION_INTERCEPTOR_METADATA_KEY,
-)
 from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric
 from google.cloud.bigtable.data._metrics.data_model import OperationState
 from google.cloud.bigtable.data._metrics.data_model import OperationType
@@ -49,22 +46,8 @@ def _with_operation_from_metadata(func):
 
     @wraps(func)
     def wrapper(self, continuation, client_call_details, request):
-        found_operation_id: str | None = None
-        try:
-            new_metadata: list[tuple[str, str]] = []
-            if client_call_details.metadata:
-                # find operation key from metadata
-                for k, v in client_call_details.metadata:
-                    if k == OPERATION_INTERCEPTOR_METADATA_KEY:
-                        found_operation_id = v
-                    else:
-                        new_metadata.append((k, v))
-            # update client_call_details to drop the operation key metadata
-            client_call_details.metadata = new_metadata
-        except Exception:
-            pass
-
-        operation: "ActiveOperationMetric" = self.operation_map.get(found_operation_id)
+        operation: "ActiveOperationMetric" | None = ActiveOperationMetric.get_active()
+
         if operation:
             # start a new attempt if not started
             if (
@@ -108,31 +91,6 @@ class AsyncBigtableMetricsInterceptor(
     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
-
-        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
     @_with_operation_from_metadata
     async def intercept_unary_unary(

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

@@ -195,7 +195,6 @@ async def add_to_flow_with_metrics(
             try:
                 value = await inner_generator.__anext__()
             except CrossSync.StopIteration:
-                metric.cancel()
                 return
             metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time
             yield value, metric

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

@@ -13,12 +13,14 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import Tuple, cast, TYPE_CHECKING
+from typing import ClassVar, Tuple, cast, TYPE_CHECKING
+
 
 import time
 import re
 import logging
 import uuid
+import contextvars
 
 from enum import Enum
 from functools import lru_cache
@@ -54,8 +56,6 @@
 
 INVALID_STATE_ERROR = "Invalid state for {}: {}"
 
-OPERATION_INTERCEPTOR_METADATA_KEY = "x-goog-operation-key"
-
 
 class OperationType(Enum):
     """Enum for the type of operation being performed."""
@@ -169,14 +169,12 @@ class ActiveOperationMetric:
     # time waiting on flow control, in nanoseconds
     flow_throttling_time_ns: int = 0
 
-    @property
-    def interceptor_metadata(self) -> tuple[str, str]:
-        """
-        returns a tuple to attach to the grpc metadata.
 
-        This metadata field will be read by the BigtableMetricsInterceptor to associate a request with an operation
-        """
-        return OPERATION_INTERCEPTOR_METADATA_KEY, self.uuid
+    _active_operation_context: ClassVar[contextvars.ContextVar] = contextvars.ContextVar("active_operation_context")
+
+    @classmethod
+    def get_active(cls):
+        return cls._active_operation_context.get(None)
 
     @property
     def state(self) -> OperationState:
@@ -190,6 +188,9 @@ def state(self) -> OperationState:
         else:
             return OperationState.ACTIVE_ATTEMPT
 
+    def __post_init__(self):
+        self._active_operation_context.set(self)
+
     def start(self) -> None:
         """
         Optionally called to mark the start of the operation. If not called,
@@ -200,6 +201,7 @@ def start(self) -> None:
         if self.state != OperationState.CREATED:
             return self._handle_error(INVALID_STATE_ERROR.format("start", self.state))
         self.start_time_ns = time.monotonic_ns()
+        self._active_operation_context.set(self)
 
     def start_attempt(self) -> ActiveAttemptMetric | None:
         """
@@ -214,6 +216,7 @@ def start_attempt(self) -> ActiveAttemptMetric | None:
             return self._handle_error(
                 INVALID_STATE_ERROR.format("start_attempt", self.state)
             )
+        self._active_operation_context.set(self)
 
         try:
             # find backoff value before this attempt
@@ -365,13 +368,6 @@ def end_with_success(self):
         """
         return self.end_with_status(StatusCode.OK)
 
-    def cancel(self):
-        """
-        Called to cancel an operation without processing emitting it.
-        """
-        for handler in self.handlers:
-            handler.on_operation_cancelled(self)
-
     @staticmethod
     def _exc_to_status(exc: BaseException) -> StatusCode:
         """

google/cloud/bigtable/data/_metrics/handlers/_base.py

@@ -29,9 +29,6 @@ def __init__(self, **kwargs):
     def on_operation_complete(self, op: CompletedOperationMetric) -> None:
         pass
 
-    def on_operation_cancelled(self, op: ActiveOperationMetric) -> None:
-        pass
-
     def on_attempt_complete(
         self, attempt: CompletedAttemptMetric, op: ActiveOperationMetric
     ) -> None:

google/cloud/bigtable/data/_metrics/metrics_controller.py

@@ -38,7 +38,6 @@ class BigtableClientSideMetricsController:
 
     def __init__(
         self,
-        interceptor: AsyncBigtableMetricsInterceptor | BigtableMetricsInterceptor,
         handlers: list[MetricsHandler] | None = None,
         **kwargs,
     ):
@@ -50,12 +49,7 @@ def __init__(
           - handlers: A list of MetricsHandler objects to subscribe to metrics events.
           - **kwargs: Optional arguments to pass to the metrics handlers.
         """
-        self.interceptor = interceptor
         self.handlers: list[MetricsHandler] = handlers or []
-        if handlers is None:
-            # handlers not given. Use default handlers.
-            # TODO: add default handlers
-            pass
 
     def add_handler(self, handler: MetricsHandler) -> None:
         """
@@ -72,6 +66,4 @@ def create_operation(
         """
         Creates a new operation and registers it with the subscribed handlers.
         """
-        new_op = ActiveOperationMetric(op_type, **kwargs, handlers=self.handlers)
-        self.interceptor.register_operation(new_op)
-        return new_op
+        return ActiveOperationMetric(op_type, **kwargs, handlers=self.handlers)