capture status for unary failed attempts

Author: daniel-sanche

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

@@ -17,6 +17,8 @@
 
 import time
 from functools import wraps
+from grpc import StatusCode
+
 from google.cloud.bigtable.data._metrics.data_model import (
     OPERATION_INTERCEPTOR_METADATA_KEY,
 )
@@ -147,18 +149,23 @@ def on_operation_cancelled(self, op):
     async def intercept_unary_unary(
         self, operation, continuation, client_call_details, request
     ):
-        encountered_exc: Exception | None = None
+        encountered_status: Exception | StatusCode | None = None
         metadata = None
         try:
             call = await continuation(client_call_details, request)
             metadata = await _get_metadata(call)
+            if CrossSync.is_async:
+                encountered_status = await call.code()
+            elif isinstance(call, Exception):
+                # sync unary calls return exception objects without raising
+                encountered_status = call
             return call
         except Exception as rpc_error:
             metadata = await _get_metadata(rpc_error)
-            encountered_exc = rpc_error
+            encountered_status = rpc_error
             raise rpc_error
         finally:
-            _end_attempt(operation, encountered_exc, metadata)
+            _end_attempt(operation, encountered_status, metadata)
 
     @CrossSync.convert
     @_with_operation_from_metadata

tests/system/data/test_metrics_async.py

@@ -1561,6 +1561,45 @@ async def test_mutate_row_failure_unauthorized(
             and attempt.gfe_latency_ns < operation.duration_ns
         )
 
+    @CrossSync.pytest
+    async def test_mutate_row_failure_unauthorized_with_retries(
+        self, handler, authorized_view, cluster_config
+    ):
+        """
+        retry unauthorized request multiple times before timing out
+        """
+        from google.cloud.bigtable.data.mutations import SetCell
+
+        row_key = b"row_key_1"
+        mutation = SetCell("unauthorized", b"q", b"v")
+
+        with pytest.raises(GoogleAPICallError) as e:
+            await authorized_view.mutate_row(row_key, [mutation], retryable_errors=[PermissionDenied], operation_timeout=30)
+        assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED"
+        # validate counts
+        assert len(handler.completed_operations) == 1
+        assert len(handler.completed_attempts) > 1
+        assert len(handler.cancelled_operations) == 0
+        # validate operation
+        operation = handler.completed_operations[0]
+        assert isinstance(operation, CompletedOperationMetric)
+        assert operation.final_status.name == "DEADLINE_EXCEEDED"
+        assert operation.op_type.value == "MutateRow"
+        assert operation.is_streaming is False
+        assert len(operation.completed_attempts) > 1
+        assert operation.cluster_id == next(iter(cluster_config.keys()))
+        assert (
+            operation.zone
+            == cluster_config[operation.cluster_id].location.split("/")[-1]
+        )
+        # validate attempts
+        for attempt in handler.completed_attempts:
+            assert attempt.end_status.name == "PERMISSION_DENIED"
+            assert (
+                attempt.gfe_latency_ns >= 0
+                and attempt.gfe_latency_ns < operation.duration_ns
+            )
+
     @CrossSync.pytest
     async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config):
         await table.sample_row_keys()