feat: support with_call for wrapped rpcs (#550)
diff --git a/google/api_core/gapic_v1/method.py b/google/api_core/gapic_v1/method.py
index e6df133..0f14ea9 100644
--- a/google/api_core/gapic_v1/method.py
+++ b/google/api_core/gapic_v1/method.py
@@ -137,6 +137,8 @@
default_timeout=None,
default_compression=None,
client_info=client_info.DEFAULT_CLIENT_INFO,
+ *,
+ with_call=False,
):
"""Wrap an RPC method with common behavior.
@@ -216,6 +218,10 @@
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``,
@@ -223,6 +229,13 @@
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
index 0623a5b..d966f47 100644
--- a/tests/unit/gapic/test_method.py
+++ b/tests/unit/gapic/test_method.py
@@ -201,3 +201,24 @@
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)
