feat: support with_call for wrapped rpcs (#550)

daniel-sanche · web-flow · commit 01a57a745f4c · 2023-11-09T21:31:36.000Z
diff --git a/google/api_core/gapic_v1/method.py b/google/api_core/gapic_v1/method.py
@@ -137,6 +137,8 @@ def wrap_method(
     default_timeout=None,
     default_compression=None,
     client_info=client_info.DEFAULT_CLIENT_INFO,
+    *,
+    with_call=False,
 ):
     """Wrap an RPC method with common behavior.
 
@@ -216,13 +218,24 @@ def get_topic(name, timeout=None):
                 passed as gRPC metadata to the method. If unspecified, then
                 a sane default will be used. If ``None``, then no user agent
                 metadata will be provided to the RPC method.
+        with_call (bool): If True, wrapped grpc.UnaryUnaryMulticallables will
+            return a tuple of (response, grpc.Call) instead of just the response.
+            This is useful for extracting trailing metadata from unary calls.
+            Defaults to False.
 
     Returns:
         Callable: A new callable that takes optional ``retry``, ``timeout``,
             and ``compression``
             arguments and applies the common error mapping, retry, timeout, compression,
             and metadata behavior to the low-level RPC method.
     """
+    if with_call:
+        try:
+            func = func.with_call
+        except AttributeError as exc:
+            raise ValueError(
+                "with_call=True is only supported for unary calls."
+            ) from exc
     func = grpc_helpers.wrap_errors(func)
     if client_info is not None:
         user_agent_metadata = [client_info.to_grpc_metadata()]
diff --git a/tests/unit/gapic/test_method.py b/tests/unit/gapic/test_method.py
@@ -201,3 +201,24 @@ def test_wrap_method_with_overriding_timeout_as_a_number():
 
     assert result == 42
     method.assert_called_once_with(timeout=22, metadata=mock.ANY)
+
+
+def test_wrap_method_with_call():
+    method = mock.Mock()
+    mock_call = mock.Mock()
+    method.with_call.return_value = 42, mock_call
+
+    wrapped_method = google.api_core.gapic_v1.method.wrap_method(method, with_call=True)
+    result = wrapped_method()
+    assert len(result) == 2
+    assert result[0] == 42
+    assert result[1] == mock_call
+
+
+def test_wrap_method_with_call_not_supported():
+    """Raises an error if wrapped callable doesn't have with_call method."""
+    method = lambda: None  # noqa: E731
+
+    with pytest.raises(ValueError) as exc_info:
+        google.api_core.gapic_v1.method.wrap_method(method, with_call=True)
+    assert "with_call=True is only supported for unary calls" in str(exc_info.value)
