added unit tests for interceptor

Author: daniel-sanche

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

@@ -54,7 +54,7 @@ def wrapper(self, continuation, client_call_details, request):
         operation: "ActiveOperationMetric" = self.operation_map.get(key)
         if operation:
             # start a new attempt if not started
-            if operation.state != OperationState.ACTIVE_ATTEMPT:
+            if operation.state == OperationState.CREATED or operation.state == OperationState.BETWEEN_ATTEMPTS:
                 operation.start_attempt()
             # wrap continuation in logic to process the operation
             return func(self, operation, continuation, client_call_details, request)
@@ -65,6 +65,26 @@ def wrapper(self, continuation, client_call_details, request):
     return wrapper
 
 
+def _end_attempt(operation, exc, metadata):
+    """Helper to add metadata and exception to an operation"""
+    if metadata is not None:
+        operation.add_response_metadata(metadata)
+    if exc is not None:
+        # end attempt. If it succeeded, let higher levels decide when to end operation
+        operation.end_attempt_with_status(exc)
+
+
+@CrossSync.convert
+async def _get_metadata(source):
+    """Helper to extract metadata from a call or RpcError"""
+    try:
+        return (await source.trailing_metadata() or []) + (
+            await source.initial_metadata() or []
+        )
+    except Exception:
+        # ignore errors while fetching metadata
+        return None
+
 @CrossSync.convert_class(sync_name="BigtableMetricsInterceptor")
 class AsyncBigtableMetricsInterceptor(
     UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor, MetricsHandler
@@ -105,29 +125,23 @@ async def intercept_unary_unary(
         self, operation, continuation, client_call_details, request
     ):
         encountered_exc: Exception | None = None
-        call = None
+        metadata = None
         try:
             call = await continuation(client_call_details, request)
+            metadata = await _get_metadata(call)
             return call
-        except Exception as e:
-            encountered_exc = e
-            raise
+        except Exception as rpc_error:
+            metadata = await _get_metadata(rpc_error)
+            encountered_exc = rpc_error
+            raise rpc_error
         finally:
-            if call is not None:
-                metadata = (
-                    await call.trailing_metadata() + await call.initial_metadata()
-                )
-                operation.add_response_metadata(metadata)
-                if encountered_exc is not None:
-                    # end attempt. If it succeeded, let higher levels decide when to end operation
-                    operation.end_attempt_with_status(encountered_exc)
+            _end_attempt(operation, encountered_exc, metadata)
 
     @CrossSync.convert
     @_with_operation_from_metadata
     async def intercept_unary_stream(
         self, operation, continuation, client_call_details, request
     ):
-        # TODO: benchmark
         async def response_wrapper(call):
             has_first_response = operation.first_response_latency is not None
             encountered_exc = None
@@ -140,17 +154,17 @@ async def response_wrapper(call):
                         )
                         has_first_response = True
                     yield response
-
             except Exception as e:
+                # handle errors while processing stream
                 encountered_exc = e
                 raise
             finally:
-                metadata = (
-                    await call.trailing_metadata() + await call.initial_metadata()
-                )
-                operation.add_response_metadata(metadata)
-                if encountered_exc is not None:
-                    # end attempt. If it succeeded, let higher levels decide when to end operation
-                    operation.end_attempt_with_status(encountered_exc)
-
-        return response_wrapper(await continuation(client_call_details, request))
+                if call is not None:
+                    _end_attempt(operation, encountered_exc, await _get_metadata(call))
+
+        try:
+            return response_wrapper(await continuation(client_call_details, request))
+        except Exception as rpc_error:
+            # handle errors while intializing stream
+            _end_attempt(operation, rpc_error, await _get_metadata(rpc_error))
+            raise rpc_error

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

@@ -15,6 +15,7 @@
 # This file is automatically generated by CrossSync. Do not edit manually.
 
 from __future__ import annotations
+import time
 from functools import wraps
 from google.cloud.bigtable.data._metrics.data_model import (
     OPERATION_INTERCEPTOR_METADATA_KEY,
@@ -42,7 +43,10 @@ def wrapper(self, continuation, client_call_details, request):
         )
         operation: "ActiveOperationMetric" = self.operation_map.get(key)
         if operation:
-            if operation.state != OperationState.ACTIVE_ATTEMPT:
+            if (
+                operation.state == OperationState.CREATED
+                or operation.state == OperationState.BETWEEN_ATTEMPTS
+            ):
                 operation.start_attempt()
             return func(self, operation, continuation, client_call_details, request)
         else:
@@ -51,6 +55,22 @@ def wrapper(self, continuation, client_call_details, request):
     return wrapper
 
 
+def _end_attempt(operation, exc, metadata):
+    """Helper to add metadata and exception to an operation"""
+    if metadata is not None:
+        operation.add_response_metadata(metadata)
+    if exc is not None:
+        operation.end_attempt_with_status(exc)
+
+
+def _get_metadata(source):
+    """Helper to extract metadata from a call or RpcError"""
+    try:
+        return (source.trailing_metadata() or []) + (source.initial_metadata() or [])
+    except Exception:
+        return None
+
+
 class BigtableMetricsInterceptor(
     UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor, MetricsHandler
 ):
@@ -76,7 +96,8 @@ def register_operation(self, operation):
         operation.handlers.append(self)
 
     def on_operation_complete(self, op):
-        del self.operation_map[op.uuid]
+        if op.uuid in self.operation_map:
+            del self.operation_map[op.uuid]
 
     def on_operation_cancelled(self, op):
         self.on_operation_complete(op)
@@ -86,36 +107,42 @@ def intercept_unary_unary(
         self, operation, continuation, client_call_details, request
     ):
         encountered_exc: Exception | None = None
-        call = None
+        metadata = None
         try:
             call = continuation(client_call_details, request)
+            metadata = _get_metadata(call)
             return call
-        except Exception as e:
-            encountered_exc = e
-            raise
+        except Exception as rpc_error:
+            metadata = _get_metadata(rpc_error)
+            encountered_exc = rpc_error
+            raise rpc_error
         finally:
-            if call is not None:
-                metadata = call.trailing_metadata() + call.initial_metadata()
-                operation.add_response_metadata(metadata)
-                if encountered_exc is not None:
-                    operation.end_attempt_with_status(encountered_exc)
+            _end_attempt(operation, encountered_exc, metadata)
 
     @_with_operation_from_metadata
     def intercept_unary_stream(
         self, operation, continuation, client_call_details, request
     ):
         def response_wrapper(call):
+            has_first_response = operation.first_response_latency is not None
             encountered_exc = None
             try:
                 for response in call:
+                    if not has_first_response:
+                        operation.first_response_latency_ns = (
+                            time.monotonic_ns() - operation.start_time_ns
+                        )
+                        has_first_response = True
                     yield response
             except Exception as e:
                 encountered_exc = e
                 raise
             finally:
-                metadata = call.trailing_metadata() + call.initial_metadata()
-                operation.add_response_metadata(metadata)
-                if encountered_exc is not None:
-                    operation.end_attempt_with_status(encountered_exc)
+                if call is not None:
+                    _end_attempt(operation, encountered_exc, _get_metadata(call))
 
-        return response_wrapper(continuation(client_call_details, request))
+        try:
+            return response_wrapper(continuation(client_call_details, request))
+        except Exception as rpc_error:
+            _end_attempt(operation, rpc_error, _get_metadata(rpc_error))
+            raise rpc_error