added try; generated sync

Author: daniel-sanche

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

@@ -44,18 +44,19 @@ def _with_operation_from_metadata(func):
     @wraps(func)
     def wrapper(self, continuation, client_call_details, request):
         found_operation_id: str | None = None
-        new_metadata = client_call_details.metadata
-        if client_call_details.metadata:
-            # find operation key from metadata
-            temp_metadata = []
-            for k, v in client_call_details.metadata:
-                if k == OPERATION_INTERCEPTOR_METADATA_KEY:
-                    found_operation_id = v
-                else:
-                    temp_metadata.append((k, v))
-            new_metadata = temp_metadata
-        # update client_call_details to drop the operation key metadata
-        client_call_details.metadata = new_metadata
+        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)
         if operation:

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

@@ -166,7 +166,7 @@ class ActiveOperationMetric:
     @property
     def interceptor_metadata(self) -> tuple[str, str]:
         """
-        returns a tuple to attach to the grpc metadata. 
+        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
         """

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

@@ -33,15 +33,19 @@ def _with_operation_from_metadata(func):
 
     @wraps(func)
     def wrapper(self, continuation, client_call_details, request):
-        key = next(
-            (
-                m[1]
-                for m in client_call_details.metadata
-                if m[0] == OPERATION_INTERCEPTOR_METADATA_KEY
-            ),
-            None,
-        )
-        operation: "ActiveOperationMetric" = self.operation_map.get(key)
+        found_operation_id: str | None = None
+        try:
+            new_metadata: list[tuple[str, str]] = []
+            if client_call_details.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))
+            client_call_details.metadata = new_metadata
+        except Exception:
+            pass
+        operation: "ActiveOperationMetric" = self.operation_map.get(found_operation_id)
         if operation:
             if (
                 operation.state == OperationState.CREATED

tests/unit/data/_async/test_metrics_interceptor.py

@@ -145,7 +145,10 @@ async def test_strip_operation_id_metadata(self):
         instance.operation_map[op.uuid] = op
         continuation = CrossSync.Mock()
         details = ClientCallDetails()
-        details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid), ("other_key", "other_value")]
+        details.metadata = [
+            (OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid),
+            ("other_key", "other_value"),
+        ]
         await instance.intercept_unary_unary(continuation, details, mock.Mock())
         assert details.metadata == [("other_key", "other_value")]
         assert continuation.call_count == 1

tests/unit/data/_metrics/test_data_model.py

@@ -577,6 +577,7 @@ def test_end_with_status_with_default_cluster_zone(self):
             called_with = h.on_operation_complete.call_args[0][0]
             assert called_with.cluster_id == DEFAULT_CLUSTER_ID
             assert called_with.zone == DEFAULT_ZONE
+
     def test_end_with_success(self):
         """
         end with success should be a pass-through helper for end_with_status

tests/unit/data/_sync_autogen/test_metrics_interceptor.py

@@ -17,6 +17,7 @@
 
 import pytest
 from grpc import RpcError
+from grpc import ClientCallDetails
 from google.cloud.bigtable.data._metrics.data_model import OperationState
 from google.cloud.bigtable.data._cross_sync import CrossSync
 
@@ -103,11 +104,33 @@ def test_on_operation_cancelled(self):
         op.cancel()
         assert op.uuid not in instance.operation_map
 
+    def test_strip_operation_id_metadata(self):
+        """After operation id is detected in metadata, the field should be stripped out before calling continuation"""
+        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 = OperationState.ACTIVE_ATTEMPT
+        instance.operation_map[op.uuid] = op
+        continuation = CrossSync._Sync_Impl.Mock()
+        details = ClientCallDetails()
+        details.metadata = [
+            (OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid),
+            ("other_key", "other_value"),
+        ]
+        instance.intercept_unary_unary(continuation, details, mock.Mock())
+        assert details.metadata == [("other_key", "other_value")]
+        assert continuation.call_count == 1
+        assert continuation.call_args[0][0].metadata == [("other_key", "other_value")]
+
     def test_unary_unary_interceptor_op_not_found(self):
         """Test that interceptor call cuntinuation if op is not found"""
         instance = self._make_one()
         continuation = CrossSync._Sync_Impl.Mock()
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = []
         request = mock.Mock()
         instance.intercept_unary_unary(continuation, details, request)
@@ -128,7 +151,7 @@ def test_unary_unary_interceptor_success(self):
         call = continuation.return_value
         call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")])
         call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")])
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         result = instance.intercept_unary_unary(continuation, details, request)
@@ -152,7 +175,7 @@ def test_unary_unary_interceptor_failure(self):
         exc.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")])
         exc.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")])
         continuation = CrossSync._Sync_Impl.Mock(side_effect=exc)
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         with pytest.raises(RpcError) as e:
@@ -178,7 +201,7 @@ def test_unary_unary_interceptor_failure_no_metadata(self):
         call = continuation.return_value
         call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")])
         call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")])
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         with pytest.raises(RpcError) as e:
@@ -204,7 +227,7 @@ def test_unary_unary_interceptor_failure_generic(self):
         call = continuation.return_value
         call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")])
         call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")])
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         with pytest.raises(ValueError) as e:
@@ -218,7 +241,7 @@ def test_unary_stream_interceptor_op_not_found(self):
         """Test that interceptor calls continuation if op is not found"""
         instance = self._make_one()
         continuation = CrossSync._Sync_Impl.Mock()
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = []
         request = mock.Mock()
         instance.intercept_unary_stream(continuation, details, request)
@@ -243,7 +266,7 @@ def test_unary_stream_interceptor_success(self):
         call = continuation.return_value
         call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")])
         call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")])
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         wrapper = instance.intercept_unary_stream(continuation, details, request)
@@ -274,7 +297,7 @@ def test_unary_stream_interceptor_failure_mid_stream(self):
         call = continuation.return_value
         call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")])
         call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")])
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         wrapper = instance.intercept_unary_stream(continuation, details, request)
@@ -304,7 +327,7 @@ def test_unary_stream_interceptor_failure_start_stream(self):
         exc.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")])
         continuation = CrossSync._Sync_Impl.Mock()
         continuation.side_effect = exc
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         with pytest.raises(RpcError) as e:
@@ -331,7 +354,7 @@ def test_unary_stream_interceptor_failure_start_stream_no_metadata(self):
         exc = RpcError("test")
         continuation = CrossSync._Sync_Impl.Mock()
         continuation.side_effect = exc
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         with pytest.raises(RpcError) as e:
@@ -358,7 +381,7 @@ def test_unary_stream_interceptor_failure_start_stream_generic(self):
         exc = ValueError("test")
         continuation = CrossSync._Sync_Impl.Mock()
         continuation.side_effect = exc
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         with pytest.raises(ValueError) as e:
@@ -387,7 +410,7 @@ def test_unary_unary_interceptor_start_operation(self, initial_state):
         call = continuation.return_value
         call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[])
         call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[])
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         instance.intercept_unary_unary(continuation, details, request)
@@ -413,7 +436,7 @@ def test_unary_stream_interceptor_start_operation(self, initial_state):
         call = continuation.return_value
         call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[])
         call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[])
-        details = mock.Mock()
+        details = ClientCallDetails()
         details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)]
         request = mock.Mock()
         instance.intercept_unary_stream(continuation, details, request)