diff --git a/api_core/google/api_core/gapic_v1/method.py b/api_core/google/api_core/gapic_v1/method.py index e9a0bfa753a6..164b712cda27 100644 --- a/api_core/google/api_core/gapic_v1/method.py +++ b/api_core/google/api_core/gapic_v1/method.py @@ -94,16 +94,16 @@ class _GapicCallable(object): timeout (google.api_core.timeout.Timeout): The default timeout for the callable. If ``None``, this callable will not specify a timeout argument to the low-level RPC method by default. - user_agent_metadata (Tuple[str, str]): The user agent metadata key and - value to provide to the RPC method. If ``None``, no additional - metadata will be passed to the RPC method. + metadata (Optional[Sequence[Tuple[str, str]]]): Additional metadata to + provide to the RPC method. If ``None``, no additional metadata will + be passed to the RPC method. """ - def __init__(self, target, retry, timeout, user_agent_metadata=None): + def __init__(self, target, retry, timeout, metadata=None): self._target = target self._retry = retry self._timeout = timeout - self._user_agent_metadata = user_agent_metadata + self._metadata = metadata def __call__(self, *args, **kwargs): """Invoke the low-level RPC with retry, timeout, and metadata.""" @@ -125,10 +125,10 @@ def __call__(self, *args, **kwargs): # Apply all applicable decorators. wrapped_func = _apply_decorators(self._target, [retry, timeout_]) - # Add the user agent metadata to the call. - if self._user_agent_metadata is not None: + # Add the metadata to the call. + if self._metadata is not None: metadata = kwargs.get('metadata', []) - metadata.append(self._user_agent_metadata) + metadata.extend(self._metadata) kwargs['metadata'] = metadata return wrapped_func(*args, **kwargs) @@ -136,7 +136,7 @@ def __call__(self, *args, **kwargs): def wrap_method( func, default_retry=None, default_timeout=None, - client_info=client_info.DEFAULT_CLIENT_INFO): + client_info=client_info.DEFAULT_CLIENT_INFO, metadata=None): """Wrap an RPC method with common behavior. This applies common error wrapping, retry, and timeout behavior a function. @@ -210,6 +210,8 @@ 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. + metadata (Optional[Sequence[Tuple[str, str]]]): Additional metadata to + provide to the RPC method. Returns: Callable: A new callable that takes optional ``retry`` and ``timeout`` @@ -219,11 +221,12 @@ def get_topic(name, timeout=None): func = grpc_helpers.wrap_errors(func) if client_info is not None: - user_agent_metadata = client_info.to_grpc_metadata() - else: - user_agent_metadata = None + if metadata is None: + metadata = [client_info.to_grpc_metadata()] + else: + metadata.append(client_info.to_grpc_metadata()) return general_helpers.wraps(func)( _GapicCallable( func, default_retry, default_timeout, - user_agent_metadata=user_agent_metadata)) + metadata=metadata)) diff --git a/api_core/tests/unit/gapic/test_method.py b/api_core/tests/unit/gapic/test_method.py index 403463322043..41d1e0526e71 100644 --- a/api_core/tests/unit/gapic/test_method.py +++ b/api_core/tests/unit/gapic/test_method.py @@ -79,6 +79,23 @@ def test_wrap_method_with_custom_client_info(): assert client_info.to_grpc_metadata() in metadata +def test_wrap_method_with_custom_metadata(): + method = mock.Mock(spec=['__call__']) + + routing_header = (b'x-google-request-params', b'name=foo&book.read=True') + + wrapped_method = google.api_core.gapic_v1.method.wrap_method( + method, metadata=[routing_header]) + + wrapped_method(1, 2, meep='moop') + + method.assert_called_once_with(1, 2, meep='moop', metadata=mock.ANY) + + # Check that the custom metadata was specified. + metadata = method.call_args[1]['metadata'] + assert routing_header in metadata + + @mock.patch('time.sleep') def test_wrap_method_with_default_retry_and_timeout(unusued_sleep): method = mock.Mock(