added metadata capture for failed rpcs

Author: daniel-sanche

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

@@ -15,6 +15,7 @@
 
 import time
 from functools import wraps
+from grpc import RpcError
 from google.cloud.bigtable.data._metrics.data_model import (
     OPERATION_INTERCEPTOR_METADATA_KEY,
 )
@@ -105,16 +106,24 @@ 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 call.trailing_metadata() or []) + (await call.initial_metadata() or [])
             return call
+        except RpcError as rpc_error:
+            # attempt extracting metadata from error
+            try:
+                metadata = (await rpc_error.trailing_metadata() or []) + (await rpc_error.initial_metadata() or [])
+            except Exception:
+                pass
+            encountered_exc = rpc_error
+            raise rpc_error
         except Exception as e:
             encountered_exc = e
             raise
         finally:
-            if call is not None:
-                metadata = (await call.trailing_metadata() or []) + (await call.initial_metadata() or [])
+            if metadata is not None:
                 operation.add_response_metadata(metadata)
             if encountered_exc is not None:
                 # end attempt. If it succeeded, let higher levels decide when to end operation
@@ -138,6 +147,7 @@ async def response_wrapper(call):
                         has_first_response = True
                     yield response
 
+
             except Exception as e:
                 encountered_exc = e
                 raise
@@ -149,4 +159,17 @@ async def response_wrapper(call):
                         # 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))
+        try:
+            return response_wrapper(await continuation(client_call_details, request))
+        except RpcError as rpc_error:
+            # attempt extracting metadata from error
+            try:
+                metadata = (await rpc_error.trailing_metadata() or []) + (await rpc_error.initial_metadata() or [])
+                operation.add_response_metadata(metadata)
+            except Exception:
+                pass
+            operation.end_attempt_with_status(rpc_error)
+            raise rpc_error
+        except Exception as e:
+            operation.end_attempt_with_status(e)
+            raise

tests/unit/data/_async/test_metrics_interceptor.py

@@ -14,6 +14,7 @@
 
 import pytest
 import asyncio
+from grpc import RpcError
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
 
@@ -158,7 +159,61 @@ async def test_unary_unary_interceptor_success(self):
 
     @CrossSync.pytest
     async def test_unary_unary_interceptor_failure(self):
-        """Test that interceptor handles failed unary-unary calls"""
+        """test a failed RpcError with metadata"""
+        from google.cloud.bigtable.data._metrics.data_model import (
+            OPERATION_INTERCEPTOR_METADATA_KEY,
+        )
+
+        instance = self._make_one()
+        op = mock.Mock()
+        op.uuid = "test-uuid"
+        op.state = 1  # ACTIVE_ATTEMPT
+        instance.operation_map[op.uuid] = op
+        exc = RpcError("test")
+        exc.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")])
+        exc.initial_metadata = CrossSync.Mock(return_value=[("c", "d")])
+        continuation = CrossSync.Mock(side_effect=exc)
+        call = continuation.return_value
+        details = mock.Mock()
+        details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
+        request = mock.Mock()
+        with pytest.raises(RpcError) as e:
+            await instance.intercept_unary_unary(continuation, details, request)
+        assert e.value == exc
+        continuation.assert_called_once_with(details, request)
+        op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")])
+        op.end_attempt_with_status.assert_called_once_with(exc)
+
+    @CrossSync.pytest
+    async def test_unary_unary_interceptor_failure_no_metadata(self):
+        """test with RpcError without without metadata attached"""
+        from google.cloud.bigtable.data._metrics.data_model import (
+            OPERATION_INTERCEPTOR_METADATA_KEY,
+        )
+
+        instance = self._make_one()
+        op = mock.Mock()
+        op.uuid = "test-uuid"
+        op.state = 1  # ACTIVE_ATTEMPT
+        instance.operation_map[op.uuid] = op
+        exc = RpcError("test")
+        continuation = CrossSync.Mock(side_effect=exc)
+        call = continuation.return_value
+        call.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")])
+        call.initial_metadata = CrossSync.Mock(return_value=[("c", "d")])
+        details = mock.Mock()
+        details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
+        request = mock.Mock()
+        with pytest.raises(RpcError) as e:
+            await instance.intercept_unary_unary(continuation, details, request)
+        assert e.value == exc
+        continuation.assert_called_once_with(details, request)
+        op.add_response_metadata.assert_not_called()
+        op.end_attempt_with_status.assert_called_once_with(exc)
+
+    @CrossSync.pytest
+    async def test_unary_unary_interceptor_failure_generic(self):
+        """test generic exception"""
         from google.cloud.bigtable.data._metrics.data_model import (
             OPERATION_INTERCEPTOR_METADATA_KEY,
         )
@@ -180,9 +235,10 @@ async def test_unary_unary_interceptor_failure(self):
             await instance.intercept_unary_unary(continuation, details, request)
         assert e.value == exc
         continuation.assert_called_once_with(details, request)
-        op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")])
+        op.add_response_metadata.assert_not_called()
         op.end_attempt_with_status.assert_called_once_with(exc)
 
+
     @CrossSync.pytest
     async def test_unary_stream_interceptor_op_not_found(self):
         """Test that interceptor calls continuation if op is not found"""
@@ -263,7 +319,67 @@ async def mock_generator():
 
     @CrossSync.pytest
     async def test_unary_stream_interceptor_failure_start_stream(self):
-        """Test that interceptor handles failures at start of stream"""
+        """Test that interceptor handles failures at start of stream with RpcError with metadata"""
+        from google.cloud.bigtable.data._metrics.data_model import (
+            OPERATION_INTERCEPTOR_METADATA_KEY,
+        )
+
+        instance = self._make_one()
+        op = mock.Mock()
+        op.uuid = "test-uuid"
+        op.state = 1  # ACTIVE_ATTEMPT
+        op.start_time_ns = 0
+        op.first_response_latency = None
+        instance.operation_map[op.uuid] = op
+        exc = RpcError("test")
+        exc.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")])
+        exc.initial_metadata = CrossSync.Mock(return_value=[("c", "d")])
+
+        continuation = CrossSync.Mock()
+        continuation.side_effect = exc
+        details = mock.Mock()
+        details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
+        request = mock.Mock()
+        with pytest.raises(RpcError) as e:
+            await instance.intercept_unary_stream(continuation, details, request)
+        assert e.value == exc
+        continuation.assert_called_once_with(details, request)
+        assert op.first_response_latency_ns is not None
+        op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")])
+        op.end_attempt_with_status.assert_called_once_with(exc)
+
+    @CrossSync.pytest
+    async def test_unary_stream_interceptor_failure_start_stream_no_metadata(self):
+        """Test that interceptor handles failures at start of stream with RpcError with no metadata"""
+        from google.cloud.bigtable.data._metrics.data_model import (
+            OPERATION_INTERCEPTOR_METADATA_KEY,
+        )
+
+        instance = self._make_one()
+        op = mock.Mock()
+        op.uuid = "test-uuid"
+        op.state = 1  # ACTIVE_ATTEMPT
+        op.start_time_ns = 0
+        op.first_response_latency = None
+        instance.operation_map[op.uuid] = op
+        exc = RpcError("test")
+
+        continuation = CrossSync.Mock()
+        continuation.side_effect = exc
+        details = mock.Mock()
+        details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
+        request = mock.Mock()
+        with pytest.raises(RpcError) as e:
+            await instance.intercept_unary_stream(continuation, details, request)
+        assert e.value == exc
+        continuation.assert_called_once_with(details, request)
+        assert op.first_response_latency_ns is not None
+        op.add_response_metadata.assert_not_called()
+        op.end_attempt_with_status.assert_called_once_with(exc)
+
+    @CrossSync.pytest
+    async def test_unary_stream_interceptor_failure_start_stream_generic(self):
+        """Test that interceptor handles failures at start of stream with generic exception"""
         from google.cloud.bigtable.data._metrics.data_model import (
             OPERATION_INTERCEPTOR_METADATA_KEY,
         )
@@ -287,5 +403,5 @@ async def test_unary_stream_interceptor_failure_start_stream(self):
         assert e.value == exc
         continuation.assert_called_once_with(details, request)
         assert op.first_response_latency_ns is not None
-        op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")])
+        op.add_response_metadata.assert_not_called()
         op.end_attempt_with_status.assert_called_once_with(exc)