Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright 2017 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15"""Helpers for wrapping low-level gRPC methods with common functionality.
17This is used by gapic clients to provide common error mapping, retry, timeout,
18compression, pagination, and long-running operations to gRPC methods.
19"""
21import enum
22import functools
23from typing import List, Tuple
25from google.api_core import grpc_helpers
26from google.api_core.gapic_v1 import client_info
27from google.api_core.timeout import TimeToDeadlineTimeout
29USE_DEFAULT_METADATA = object()
32class _MethodDefault(enum.Enum):
33 # Uses enum so that pytype/mypy knows that this is the only possible value.
34 # https://stackoverflow.com/a/60605919/101923
35 _DEFAULT_VALUE = object()
38DEFAULT = _MethodDefault._DEFAULT_VALUE
39"""Sentinel value indicating that a retry, timeout, or compression argument was unspecified,
40so the default should be used."""
43def _is_not_none_or_false(value):
44 return value is not None and value is not False
47def _apply_decorators(func, decorators):
48 """Apply a list of decorators to a given function.
50 ``decorators`` may contain items that are ``None`` or ``False`` which will
51 be ignored.
52 """
53 filtered_decorators = filter(_is_not_none_or_false, reversed(decorators))
55 for decorator in filtered_decorators:
56 func = decorator(func)
58 return func
61def _deduplicate_metadata_tokens(*headers: str | None) -> str:
62 """
63 Given one or more metadata payload strings, create a combined
64 string with deduplicated tokens, while preserving token order.
66 Inputs are expected to contain a set of metadata tokens separated by spaces
67 Example: `gl-python/3.14.0 grpc/1.76.0 gax/2.29.0 gapic/3.8.0 pb/6.33.4`
69 Args:
70 *headers: one or more metadata payload strings
72 Returns:
73 a single combined payload string
74 """
75 # Split all non-empty headers into individual tokens
76 token_list = " ".join(filter(None, headers)).split()
77 # Deduplicate while preserving order
78 return " ".join(dict.fromkeys(token_list))
81def _extract_metrics_header(metadata) -> Tuple[str, List[Tuple[str, str]]]:
82 """Extract x-google-api-client header from metadata list.
84 Args:
85 metadata (Sequence[Tuple[str, str]]): The metadata to extract from.
87 Returns:
88 A tuple containing:
89 - a string representing the header value.
90 - A sequence of remaining metadata tuples.
91 """
92 if not metadata:
93 return "", []
95 key_to_find = client_info.METRICS_METADATA_KEY
97 metric_str = _deduplicate_metadata_tokens(
98 " ".join([v for k, v in metadata if k == key_to_find])
99 )
100 if not metric_str:
101 return "", list(metadata)
103 arbitrary_metadata = [item for item in metadata if item[0] != key_to_find]
104 return metric_str, arbitrary_metadata
107class _GapicCallable(object):
108 """Callable that applies retry, timeout, and metadata logic.
110 Args:
111 target (Callable): The low-level RPC method.
112 retry (google.api_core.retry.Retry): The default retry for the
113 callable. If ``None``, this callable will not retry by default
114 timeout (google.api_core.timeout.Timeout): The default timeout for the
115 callable (i.e. duration of time within which an RPC must terminate
116 after its start, not to be confused with deadline). If ``None``,
117 this callable will not specify a timeout argument to the low-level
118 RPC method.
119 compression (grpc.Compression): The default compression for the callable.
120 If ``None``, this callable will not specify a compression argument
121 to the low-level RPC method.
122 metadata (Sequence[Tuple[str, str]]): Additional metadata that is
123 provided to the RPC method on every invocation. This is merged with
124 any metadata specified during invocation. If ``None``, no
125 additional metadata will be passed to the RPC method.
126 """
128 def __init__(
129 self,
130 target,
131 retry,
132 timeout,
133 compression,
134 metadata=None,
135 ):
136 self._target = target
137 self._retry = retry
138 self._timeout = timeout
139 self._compression = compression
140 # Pre-extract the x-goog-api-client header from the initialized metadata.
141 self._x_goog_api_client, remaining = _extract_metrics_header(metadata)
142 self._static_metadata = tuple(remaining)
143 if self._x_goog_api_client:
144 self._default_metadata = (
145 (client_info.METRICS_METADATA_KEY, self._x_goog_api_client),
146 *self._static_metadata,
147 )
148 else:
149 self._default_metadata = self._static_metadata
151 def __call__(
152 self, *args, timeout=DEFAULT, retry=DEFAULT, compression=DEFAULT, **kwargs
153 ):
154 """Invoke the low-level RPC with retry, timeout, compression, and metadata."""
156 if retry is DEFAULT:
157 retry = self._retry
159 if timeout is DEFAULT:
160 timeout = self._timeout
162 if compression is DEFAULT:
163 compression = self._compression
165 if isinstance(timeout, (int, float)):
166 timeout = TimeToDeadlineTimeout(timeout=timeout)
168 # Apply all applicable decorators.
169 wrapped_func = _apply_decorators(self._target, [retry, timeout])
171 if user_metadata := kwargs.get("metadata"):
172 # Add the user agent metadata to the call.
173 final_metadata = list(self._static_metadata)
174 user_x_goog, remaining = _extract_metrics_header(user_metadata)
176 merged_header = _deduplicate_metadata_tokens(
177 self._x_goog_api_client, user_x_goog
178 )
179 if merged_header:
180 final_metadata.append((client_info.METRICS_METADATA_KEY, merged_header))
181 final_metadata.extend(remaining)
182 kwargs["metadata"] = final_metadata
183 elif self._default_metadata:
184 kwargs["metadata"] = self._default_metadata
186 if self._compression is not None:
187 kwargs["compression"] = compression
189 return wrapped_func(*args, **kwargs)
192def wrap_method(
193 func,
194 default_retry=None,
195 default_timeout=None,
196 default_compression=None,
197 client_info=client_info.DEFAULT_CLIENT_INFO,
198 *,
199 with_call=False,
200):
201 """Wrap an RPC method with common behavior.
203 This applies common error wrapping, retry, timeout, and compression behavior to a function.
204 The wrapped function will take optional ``retry``, ``timeout``, and ``compression``
205 arguments.
207 For example::
209 import google.api_core.gapic_v1.method
210 from google.api_core import retry
211 from google.api_core import timeout
212 from grpc import Compression
214 # The original RPC method.
215 def get_topic(name, timeout=None):
216 request = publisher_v2.GetTopicRequest(name=name)
217 return publisher_stub.GetTopic(request, timeout=timeout)
219 default_retry = retry.Retry(deadline=60)
220 default_timeout = timeout.Timeout(deadline=60)
221 default_compression = Compression.NoCompression
222 wrapped_get_topic = google.api_core.gapic_v1.method.wrap_method(
223 get_topic, default_retry)
225 # Execute get_topic with default retry and timeout:
226 response = wrapped_get_topic()
228 # Execute get_topic without doing any retying but with the default
229 # timeout:
230 response = wrapped_get_topic(retry=None)
232 # Execute get_topic but only retry on 5xx errors:
233 my_retry = retry.Retry(retry.if_exception_type(
234 exceptions.InternalServerError))
235 response = wrapped_get_topic(retry=my_retry)
237 The way this works is by late-wrapping the given function with the retry
238 and timeout decorators. Essentially, when ``wrapped_get_topic()`` is
239 called:
241 * ``get_topic()`` is first wrapped with the ``timeout`` into
242 ``get_topic_with_timeout``.
243 * ``get_topic_with_timeout`` is wrapped with the ``retry`` into
244 ``get_topic_with_timeout_and_retry()``.
245 * The final ``get_topic_with_timeout_and_retry`` is called passing through
246 the ``args`` and ``kwargs``.
248 The callstack is therefore::
250 method.__call__() ->
251 Retry.__call__() ->
252 Timeout.__call__() ->
253 wrap_errors() ->
254 get_topic()
256 Note that if ``timeout`` or ``retry`` is ``None``, then they are not
257 applied to the function. For example,
258 ``wrapped_get_topic(timeout=None, retry=None)`` is more or less
259 equivalent to just calling ``get_topic`` but with error re-mapping.
261 Args:
262 func (Callable[Any]): The function to wrap. It should accept an
263 optional ``timeout`` argument. If ``metadata`` is not ``None``, it
264 should accept a ``metadata`` argument.
265 default_retry (Optional[google.api_core.Retry]): The default retry
266 strategy. If ``None``, the method will not retry by default.
267 default_timeout (Optional[google.api_core.Timeout]): The default
268 timeout strategy. Can also be specified as an int or float. If
269 ``None``, the method will not have timeout specified by default.
270 default_compression (Optional[grpc.Compression]): The default
271 grpc.Compression. If ``None``, the method will not have
272 compression specified by default.
273 client_info
274 (Optional[google.api_core.gapic_v1.client_info.ClientInfo]):
275 Client information used to create a user-agent string that's
276 passed as gRPC metadata to the method. If unspecified, then
277 a sane default will be used. If ``None``, then no user agent
278 metadata will be provided to the RPC method.
279 with_call (bool): If True, wrapped grpc.UnaryUnaryMulticallables will
280 return a tuple of (response, grpc.Call) instead of just the response.
281 This is useful for extracting trailing metadata from unary calls.
282 Defaults to False.
284 Returns:
285 Callable: A new callable that takes optional ``retry``, ``timeout``,
286 and ``compression``
287 arguments and applies the common error mapping, retry, timeout, compression,
288 and metadata behavior to the low-level RPC method.
289 """
290 if with_call:
291 try:
292 func = func.with_call
293 except AttributeError as exc:
294 raise ValueError(
295 "with_call=True is only supported for unary calls."
296 ) from exc
297 func = grpc_helpers.wrap_errors(func)
298 if client_info is not None:
299 user_agent_metadata = [client_info.to_grpc_metadata()]
300 else:
301 user_agent_metadata = None
303 return functools.wraps(func)(
304 _GapicCallable(
305 func,
306 default_retry,
307 default_timeout,
308 default_compression,
309 metadata=user_agent_metadata,
310 )
311 )