improved from_context

Author: daniel-sanche

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

@@ -45,7 +45,7 @@ def _with_active_operation(func):
 
     @wraps(func)
     def wrapper(self, continuation, client_call_details, request):
-        operation: "ActiveOperationMetric" | None = ActiveOperationMetric.from_context()
+        operation: ActiveOperationMetric | None = ActiveOperationMetric.from_context()
 
         if operation:
             # start a new attempt if not started

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

@@ -178,13 +178,25 @@ class ActiveOperationMetric:
 
     @classmethod
     def from_context(cls) -> ActiveOperationMetric | None:
-        """
-        Reads the active operation saved to contextvars
+        """Retrieves the active operation from the current execution context.
+
+        Because execution within a context is sequential, this guarantees
+        retrieval of the single, unique operation, isolated from other
+        concurrent RPCs.
+
+        Note:
+            This is intended to be called by gRPC interceptors at the start
+            of an RPC.
 
-        This is meant to be called by grpc interceptor to grab a reference to
-        the active attempt reference at the start of an rpc
+        Returns:
+            ActiveOperationMetric: The current active operation.
+            None: If no operation is set, or if the current operation is
+            already in the `COMPLETED` state.
         """
-        return cls._active_operation_context.get(None)
+        op = cls._active_operation_context.get(None)
+        if op and op.state == OperationState.COMPLETED:
+            return None
+        return op
 
     @property
     def state(self) -> OperationState:

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

@@ -32,7 +32,7 @@ def _with_active_operation(func):
 
     @wraps(func)
     def wrapper(self, continuation, client_call_details, request):
-        operation: "ActiveOperationMetric" | None = ActiveOperationMetric.from_context()
+        operation: ActiveOperationMetric | None = ActiveOperationMetric.from_context()
         if operation:
             if (
                 operation.state == OperationState.CREATED