api_core: Add metadata option to wrap_method

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,18 +125,18 @@ 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)
 
 
 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))

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(