Skip to content

Commit 4500dc7

Browse files
Jon Wayne Parrottlukesneeringer
authored andcommitted
Make gapic's wrapper accept generic metadata (#4339)
1 parent fd1d18f commit 4500dc7

File tree

2 files changed

+26
-10
lines changed

2 files changed

+26
-10
lines changed

google/api_core/gapic_v1/method.py

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -94,16 +94,17 @@ class _GapicCallable(object):
9494
timeout (google.api_core.timeout.Timeout): The default timeout
9595
for the callable. If ``None``, this callable will not specify
9696
a timeout argument to the low-level RPC method by default.
97-
user_agent_metadata (Tuple[str, str]): The user agent metadata key and
98-
value to provide to the RPC method. If ``None``, no additional
99-
metadata will be passed to the RPC method.
97+
metadata (Sequence[Tuple[str, str]]): Additional metadata that is
98+
provided to the RPC method on every invocation. This is merged with
99+
any metadata specified during invocation. If ``None``, no
100+
additional metadata will be passed to the RPC method.
100101
"""
101102

102-
def __init__(self, target, retry, timeout, user_agent_metadata=None):
103+
def __init__(self, target, retry, timeout, metadata=None):
103104
self._target = target
104105
self._retry = retry
105106
self._timeout = timeout
106-
self._user_agent_metadata = user_agent_metadata
107+
self._metadata = metadata
107108

108109
def __call__(self, *args, **kwargs):
109110
"""Invoke the low-level RPC with retry, timeout, and metadata."""
@@ -126,9 +127,9 @@ def __call__(self, *args, **kwargs):
126127
wrapped_func = _apply_decorators(self._target, [retry, timeout_])
127128

128129
# Add the user agent metadata to the call.
129-
if self._user_agent_metadata is not None:
130-
metadata = kwargs.get('metadata', [])
131-
metadata.append(self._user_agent_metadata)
130+
if self._metadata is not None:
131+
metadata = list(kwargs.get('metadata', []))
132+
metadata.extend(self._metadata)
132133
kwargs['metadata'] = metadata
133134

134135
return wrapped_func(*args, **kwargs)
@@ -219,11 +220,11 @@ def get_topic(name, timeout=None):
219220
func = grpc_helpers.wrap_errors(func)
220221

221222
if client_info is not None:
222-
user_agent_metadata = client_info.to_grpc_metadata()
223+
user_agent_metadata = [client_info.to_grpc_metadata()]
223224
else:
224225
user_agent_metadata = None
225226

226227
return general_helpers.wraps(func)(
227228
_GapicCallable(
228229
func, default_retry, default_timeout,
229-
user_agent_metadata=user_agent_metadata))
230+
metadata=user_agent_metadata))

tests/unit/gapic/test_method.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,21 @@ def test_wrap_method_with_custom_client_info():
7979
assert client_info.to_grpc_metadata() in metadata
8080

8181

82+
def test_invoke_wrapped_method_with_metadata():
83+
method = mock.Mock(spec=['__call__'])
84+
85+
wrapped_method = google.api_core.gapic_v1.method.wrap_method(method)
86+
87+
wrapped_method(mock.sentinel.request, metadata=[('a', 'b')])
88+
89+
method.assert_called_once_with(mock.sentinel.request, metadata=mock.ANY)
90+
metadata = method.call_args[1]['metadata']
91+
# Metadata should have two items: the client info metadata and our custom
92+
# metadata.
93+
assert len(metadata) == 2
94+
assert ('a', 'b') in metadata
95+
96+
8297
@mock.patch('time.sleep')
8398
def test_wrap_method_with_default_retry_and_timeout(unusued_sleep):
8499
method = mock.Mock(

0 commit comments

Comments
 (0)