diff --git a/CHANGELOG.md b/CHANGELOG.md index 492584a7..cf7e3d20 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ [1]: https://pypi.org/project/google-api-core/#history +## [2.14.0](https://github.com/googleapis/python-api-core/compare/v2.13.1...v2.14.0) (2023-11-09) + + +### Features + +* Support with_call for wrapped rpcs ([#550](https://github.com/googleapis/python-api-core/issues/550)) ([01a57a7](https://github.com/googleapis/python-api-core/commit/01a57a745f4c8345c9c93412c27dd416b49f5953)) + ## [2.13.1](https://github.com/googleapis/python-api-core/compare/v2.13.0...v2.13.1) (2023-11-09) diff --git a/google/api_core/gapic_v1/method.py b/google/api_core/gapic_v1/method.py index e6df1332..0f14ea9c 100644 --- 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,6 +218,10 @@ 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``, @@ -223,6 +229,13 @@ def get_topic(name, timeout=None): 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/google/api_core/version.py b/google/api_core/version.py index 2330d0c2..ba8b4e8a 100644 --- a/google/api_core/version.py +++ b/google/api_core/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2.13.1" +__version__ = "2.14.0" diff --git a/tests/unit/gapic/test_method.py b/tests/unit/gapic/test_method.py index 0623a5bc..d966f478 100644 --- 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)