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.
14
15"""Helpers for wrapping low-level gRPC methods with common functionality.
16
17This is used by gapic clients to provide common error mapping, retry, timeout,
18compression, pagination, and long-running operations to gRPC methods.
19"""
20
21import contextlib
22import enum
23import functools
24from typing import Any, List, Optional, Tuple
25
26from google.api_core import _observability, grpc_helpers
27from google.api_core.gapic_v1 import client_info
28from google.api_core.gapic_v1.client_info import METRICS_METADATA_KEY
29from google.api_core.timeout import TimeToDeadlineTimeout
30
31USE_DEFAULT_METADATA = object()
32
33
34class _MethodDefault(enum.Enum):
35 # Uses enum so that pytype/mypy knows that this is the only possible value.
36 # https://stackoverflow.com/a/60605919/101923
37 _DEFAULT_VALUE = object()
38
39
40DEFAULT = _MethodDefault._DEFAULT_VALUE
41"""Sentinel value indicating that a retry, timeout, or compression argument was unspecified,
42so the default should be used."""
43
44
45def _is_not_none_or_false(value):
46 return value is not None and value is not False
47
48
49def _apply_decorators(func, decorators):
50 """Apply a list of decorators to a given function.
51
52 ``decorators`` may contain items that are ``None`` or ``False`` which will
53 be ignored.
54 """
55 filtered_decorators = filter(_is_not_none_or_false, reversed(decorators))
56
57 for decorator in filtered_decorators:
58 func = decorator(func)
59
60 return func
61
62
63def _deduplicate_metadata_tokens(*headers: str | None) -> str:
64 """
65 Given one or more metadata payload strings, create a combined
66 string with deduplicated tokens, while preserving token order.
67
68 Inputs are expected to contain a set of metadata tokens separated by spaces
69 Example: `gl-python/3.14.0 grpc/1.76.0 gax/2.29.0 gapic/3.8.0 pb/6.33.4`
70
71 Args:
72 *headers: one or more metadata payload strings
73
74 Returns:
75 a single combined payload string
76 """
77 # Split all non-empty headers into individual tokens
78 token_list = " ".join(filter(None, headers)).split()
79 # Deduplicate while preserving order
80 return " ".join(dict.fromkeys(token_list))
81
82
83def _extract_metrics_header(metadata) -> Tuple[str, List[Tuple[str, str]]]:
84 """Extract x-google-api-client header from metadata list.
85
86 Args:
87 metadata (Sequence[Tuple[str, str]]): The metadata to extract from.
88
89 Returns:
90 A tuple containing:
91 - a string representing the header value.
92 - A sequence of remaining metadata tuples.
93 """
94 if not metadata:
95 return "", []
96
97 key_to_find = METRICS_METADATA_KEY
98
99 metric_str = _deduplicate_metadata_tokens(
100 " ".join([v for k, v in metadata if k == key_to_find])
101 )
102 if not metric_str:
103 return "", list(metadata)
104
105 arbitrary_metadata = [item for item in metadata if item[0] != key_to_find]
106 return metric_str, arbitrary_metadata
107
108
109def _extract_rpc_identity(
110 method_name: str,
111) -> Tuple[str, str, str]:
112 """Extract (full_rpc_name, service_name, rpc_method_name) from an explicit method name.
113
114 Args:
115 method_name: Explicit RPC name (e.g. "/google.cloud.secretmanager.v1.SecretManagerService/AccessSecretVersion").
116
117 Returns:
118 Tuple[str, str, str]: A 3-tuple of (full_rpc_name, service_name, rpc_method_name).
119 """
120 method_str = method_name.lstrip("/")
121 service, _, method = method_str.rpartition("/")
122 return method_str, service, method
123
124
125def _extract_status_code(exc: Optional[Exception]) -> str:
126 """Extract canonical status code name string from an exception.
127
128 Status code name strings are resolved by inspecting the following locations:
129 * Chained exceptions: Unwraps RetryError or __cause__ to the root exception.
130 * Enum & code attributes: Inspects .grpc_status_code on GoogleAPICallError or .code on gRPC errors.
131 * Integer status codes: Maps raw gRPC integer status codes to canonical enum names.
132 * Fallback: Defaults to the exception class name for standard Python errors.
133
134 Args:
135 exc (Optional[Exception]): The exception to extract the status code name from.
136
137 Returns:
138 str: The canonical status code name (e.g. "NOT_FOUND", "UNAVAILABLE") or class name.
139 """
140 if exc is None:
141 return ""
142
143 # 1. Unwrap chained exceptions: unwrap RetryError or __cause__ to the root failure
144 target = getattr(exc, "cause", None) or getattr(exc, "__cause__", None) or exc
145
146 # 2. Check enum & code attributes: .grpc_status_code enum or callable/non-callable .code
147 status = getattr(target, "grpc_status_code", None)
148 if status is None and hasattr(target, "code"):
149 try:
150 status = target.code() if callable(target.code) else target.code
151 except Exception:
152 status = None
153
154 name = getattr(status, "name", None)
155 if name:
156 return str(name)
157
158 # 3. Check integer status codes: map raw gRPC integer status codes to canonical enum names
159 if isinstance(status, int):
160 from google.api_core import exceptions
161
162 status = exceptions._INT_TO_GRPC_CODE.get(status, status)
163 return getattr(status, "name", str(status))
164
165 # 4. Fallback: default to the exception class name for standard Python errors
166 return target.__class__.__name__
167
168
169def _extract_error_attributes(exc: Optional[Exception]) -> dict[str, Any]:
170 """Extract gcp.errors.* and error.type attributes from an exception.
171
172 Error details and ErrorInfo structures are resolved by inspecting the following locations:
173 * Chained exceptions: Unwraps RetryError or __cause__ to the root exception.
174 * GoogleAPICallError attributes: Reads ErrorInfo from ._error_info or .error_info.
175 * Native gRPC trailing metadata: Parses google.rpc.Status binary details from trailing_metadata.
176 * Unified attribute extraction: Extracts domain, reason, and metadata from ErrorInfo or exception attributes.
177
178 Args:
179 exc (Optional[Exception]): An exception (such as GoogleAPICallError or grpc.RpcError) or ErrorInfo object.
180
181 Returns:
182 dict[str, Any]: Extracted error attributes (e.g. gcp.errors.domain, error.type, gcp.errors.metadata.*).
183 """
184 attrs: dict[str, Any] = {}
185 if exc is None:
186 return attrs
187
188 # 1. Unwrap chained exceptions: unwrap RetryError or __cause__ to the root failure
189 target_exc = getattr(exc, "cause", None) or getattr(exc, "__cause__", None) or exc
190
191 # 2. Check GoogleAPICallError ErrorInfo attributes
192 error_info = getattr(target_exc, "_error_info", None) or getattr(
193 target_exc, "error_info", None
194 )
195
196 # 3. Check native gRPC trailing metadata for binary google.rpc.Status details
197 if error_info is None:
198 rpc_call = (
199 target_exc
200 if hasattr(target_exc, "trailing_metadata")
201 else getattr(target_exc, "response", None)
202 )
203 if rpc_call is not None and hasattr(rpc_call, "trailing_metadata"):
204 try:
205 from google.api_core import exceptions
206
207 _, error_info = exceptions._parse_grpc_error_details(rpc_call)
208 except Exception:
209 pass
210
211 # 4. Unified attribute extraction: extract domain, reason, and metadata from ErrorInfo or exception attributes
212 source = error_info or target_exc
213 domain = getattr(source, "domain", None)
214 if domain:
215 attrs["gcp.errors.domain"] = domain
216 reason = getattr(source, "reason", None)
217 if reason:
218 attrs["error.type"] = reason
219 metadata = getattr(source, "metadata", None)
220 if metadata:
221 for k, v in metadata.items():
222 attrs[f"gcp.errors.metadata.{k}"] = str(v)
223
224 return attrs
225
226
227class _GapicCallable(object):
228 """Callable that applies retry, timeout, and metadata logic.
229
230 Args:
231 target (Callable): The low-level RPC method.
232 retry (google.api_core.retry.Retry): The default retry for the
233 callable. If ``None``, this callable will not retry by default
234 timeout (google.api_core.timeout.Timeout): The default timeout for the
235 callable (i.e. duration of time within which an RPC must terminate
236 after its start, not to be confused with deadline). If ``None``,
237 this callable will not specify a timeout argument to the low-level
238 RPC method.
239 compression (grpc.Compression): The default compression for the callable.
240 If ``None``, this callable will not specify a compression argument
241 to the low-level RPC method.
242 metadata (Sequence[Tuple[str, str]]): Additional metadata that is
243 provided to the RPC method on every invocation. This is merged with
244 any metadata specified during invocation. If ``None``, no
245 additional metadata will be passed to the RPC method.
246 client_options
247 (Optional[google.api_core.client_options.ClientOptions]):
248 Client options used to configure client-level behavior, such as
249 custom OpenTelemetry tracer providers. Defaults to None.
250 method_name (Optional[str]): The optional explicit full RPC method name
251 (e.g. "/google.cloud.secretmanager.v1.SecretManagerService/AccessSecretVersion").
252 is_streaming (bool): Whether the RPC method is streaming. Defaults to False.
253 Note: Streaming methods do not currently generate Tier 3 observability spans.
254 client_info (Optional[google.api_core.gapic_v1.client_info.ClientInfo]):
255 Client information used for metadata headers. Defaults to None.
256 kind (str): The transport kind for the RPC method. Defaults to "grpc".
257 Allowed values for OpenTelemetry method tracing are "grpc" and "grpc_asyncio".
258 """
259
260 def __init__(
261 self,
262 target,
263 retry,
264 timeout,
265 compression,
266 metadata=None,
267 client_options=None,
268 method_name=None,
269 is_streaming=False,
270 client_info=None,
271 kind="grpc",
272 ):
273 self._target = target
274 self._retry = retry
275 self._timeout = timeout
276 self._compression = compression
277
278 # Pre-extract the x-goog-api-client header from the initialized metadata.
279 self._x_goog_api_client, remaining = _extract_metrics_header(metadata)
280 self._static_metadata = tuple(remaining)
281 if self._x_goog_api_client:
282 self._default_metadata = (
283 (METRICS_METADATA_KEY, self._x_goog_api_client),
284 *self._static_metadata,
285 )
286 else:
287 self._default_metadata = self._static_metadata
288
289 # Configure the OpenTelemetry span factory once at initialization.
290 # For now, method tracing is gated to non-streaming gRPC calls where an explicit method_name is provided.
291 self._start_span_fn = None
292 if (
293 not is_streaming
294 and kind == "grpc"
295 and method_name is not None
296 and _observability.is_otel_capabilities_enabled(client_options)
297 ):
298 try:
299 from opentelemetry import trace
300
301 tracer_provider = None
302 if isinstance(client_options, dict):
303 tracer_provider = client_options.get("tracer_provider")
304 elif client_options is not None:
305 tracer_provider = getattr(client_options, "tracer_provider", None)
306 if tracer_provider is not None:
307 tracer = tracer_provider.get_tracer("google.api_core")
308 else:
309 tracer = trace.get_tracer("google.api_core")
310
311 span_name, _, _ = _extract_rpc_identity(method_name)
312 span_attributes = {
313 "rpc.system.name": "grpc",
314 "rpc.method": span_name,
315 }
316 self._start_span_fn = functools.partial(
317 tracer.start_as_current_span,
318 span_name,
319 kind=trace.SpanKind.CLIENT,
320 attributes=span_attributes,
321 )
322 except (ImportError, AttributeError, TypeError):
323 # Gracefully disable tracing if OpenTelemetry or custom provider fails
324 self._start_span_fn = None
325
326 def __call__(
327 self, *args, timeout=DEFAULT, retry=DEFAULT, compression=DEFAULT, **kwargs
328 ):
329 """Invoke the low-level RPC with retry, timeout, compression, and metadata."""
330
331 if retry is DEFAULT:
332 retry = self._retry
333
334 if timeout is DEFAULT:
335 timeout = self._timeout
336
337 if compression is DEFAULT:
338 compression = self._compression
339
340 if isinstance(timeout, (int, float)):
341 timeout = TimeToDeadlineTimeout(timeout=timeout)
342
343 # Apply all applicable decorators.
344 wrapped_func = _apply_decorators(self._target, [retry, timeout])
345
346 if user_metadata := kwargs.get("metadata"):
347 # Add the user agent metadata to the call.
348 final_metadata = list(self._static_metadata)
349 user_x_goog, remaining = _extract_metrics_header(user_metadata)
350
351 merged_header = _deduplicate_metadata_tokens(
352 self._x_goog_api_client, user_x_goog
353 )
354 if merged_header:
355 final_metadata.append((METRICS_METADATA_KEY, merged_header))
356 final_metadata.extend(remaining)
357 kwargs["metadata"] = final_metadata
358 elif self._default_metadata:
359 kwargs["metadata"] = self._default_metadata
360
361 if self._compression is not None:
362 kwargs["compression"] = compression
363
364 span_cm = contextlib.nullcontext()
365 if self._start_span_fn is not None:
366 try:
367 span_cm = self._start_span_fn()
368 except Exception:
369 span_cm = contextlib.nullcontext()
370
371 with span_cm as span:
372 try:
373 result = wrapped_func(*args, **kwargs)
374 if span is not None and hasattr(span, "set_attribute"):
375 span.set_attribute("rpc.response.status_code", "OK")
376 return result
377 except Exception as exc:
378 if span is not None and hasattr(span, "set_attribute"):
379 span.set_attribute(
380 "rpc.response.status_code", _extract_status_code(exc)
381 )
382 for k, v in _extract_error_attributes(exc).items():
383 span.set_attribute(k, v)
384 raise
385
386
387def wrap_method(
388 func,
389 default_retry=None,
390 default_timeout=None,
391 default_compression=None,
392 client_info=client_info.DEFAULT_CLIENT_INFO,
393 *,
394 with_call=False,
395 client_options=None,
396 method_name=None,
397 is_streaming=False,
398 kind="grpc",
399):
400 """Wrap an RPC method with common behavior.
401
402 This applies common error wrapping, retry, timeout, and compression behavior to a function.
403 The wrapped function will take optional ``retry``, ``timeout``, and ``compression``
404 arguments.
405
406 For example::
407
408 import google.api_core.gapic_v1.method
409 from google.api_core import retry
410 from google.api_core import timeout
411 from grpc import Compression
412
413 # The original RPC method.
414 def get_topic(name, timeout=None):
415 request = publisher_v2.GetTopicRequest(name=name)
416 return publisher_stub.GetTopic(request, timeout=timeout)
417
418 default_retry = retry.Retry(deadline=60)
419 default_timeout = timeout.Timeout(deadline=60)
420 default_compression = Compression.NoCompression
421 wrapped_get_topic = google.api_core.gapic_v1.method.wrap_method(
422 get_topic, default_retry)
423
424 # Execute get_topic with default retry and timeout:
425 response = wrapped_get_topic()
426
427 # Execute get_topic without doing any retying but with the default
428 # timeout:
429 response = wrapped_get_topic(retry=None)
430
431 # Execute get_topic but only retry on 5xx errors:
432 my_retry = retry.Retry(retry.if_exception_type(
433 exceptions.InternalServerError))
434 response = wrapped_get_topic(retry=my_retry)
435
436 The way this works is by late-wrapping the given function with the retry
437 and timeout decorators. Essentially, when ``wrapped_get_topic()`` is
438 called:
439
440 * ``get_topic()`` is first wrapped with the ``timeout`` into
441 ``get_topic_with_timeout``.
442 * ``get_topic_with_timeout`` is wrapped with the ``retry`` into
443 ``get_topic_with_timeout_and_retry()``.
444 * The final ``get_topic_with_timeout_and_retry`` is called passing through
445 the ``args`` and ``kwargs``.
446
447 The callstack is therefore::
448
449 method.__call__() ->
450 Retry.__call__() ->
451 Timeout.__call__() ->
452 wrap_errors() ->
453 get_topic()
454
455 Note that if ``timeout`` or ``retry`` is ``None``, then they are not
456 applied to the function. For example,
457 ``wrapped_get_topic(timeout=None, retry=None)`` is more or less
458 equivalent to just calling ``get_topic`` but with error re-mapping.
459
460 Args:
461 func (Callable[Any]): The function to wrap. It should accept an
462 optional ``timeout`` argument. If ``metadata`` is not ``None``, it
463 should accept a ``metadata`` argument.
464 default_retry (Optional[google.api_core.Retry]): The default retry
465 strategy. If ``None``, the method will not retry by default.
466 default_timeout (Optional[google.api_core.Timeout]): The default
467 timeout strategy. Can also be specified as an int or float. If
468 ``None``, the method will not have timeout specified by default.
469 default_compression (Optional[grpc.Compression]): The default
470 grpc.Compression. If ``None``, the method will not have
471 compression specified by default.
472 client_info
473 (Optional[google.api_core.gapic_v1.client_info.ClientInfo]):
474 Client information used to create a user-agent string that's
475 passed as gRPC metadata to the method. If unspecified, then
476 a sane default will be used. If ``None``, then no user agent
477 metadata will be provided to the RPC method.
478 with_call (bool): If True, wrapped grpc.UnaryUnaryMulticallables will
479 return a tuple of (response, grpc.Call) instead of just the response.
480 This is useful for extracting trailing metadata from unary calls.
481 Defaults to False.
482 client_options
483 (Optional[google.api_core.client_options.ClientOptions]):
484 Client options used to configure client-level behavior, such as
485 custom OpenTelemetry tracer providers. Defaults to None.
486 method_name (Optional[str]): Optional explicit full RPC method name
487 (e.g. "/google.cloud.secretmanager.v1.SecretManagerService/AccessSecretVersion").
488 Used to identify the RPC for observability.
489 is_streaming (bool): Whether the RPC method is streaming. Defaults to False.
490 Streaming methods are currently gated and do not generate Tier 3 spans.
491 kind (str): The transport kind for the RPC method. Defaults to "grpc".
492 Non-gRPC transports (e.g. "rest") are currently gated and do not generate
493 Tier 3 method spans.
494
495 Returns:
496 Callable: A new callable that takes optional ``retry``, ``timeout``,
497 and ``compression``
498 arguments and applies the common error mapping, retry, timeout, compression,
499 and metadata behavior to the low-level RPC method.
500 """
501 if with_call:
502 try:
503 func = func.with_call
504 except AttributeError as exc:
505 raise ValueError(
506 "with_call=True is only supported for unary calls."
507 ) from exc
508 func = grpc_helpers.wrap_errors(func)
509 if client_info is not None:
510 user_agent_metadata = [client_info.to_grpc_metadata()]
511 else:
512 user_agent_metadata = None
513
514 return functools.wraps(func)(
515 _GapicCallable(
516 func,
517 default_retry,
518 default_timeout,
519 default_compression,
520 metadata=user_agent_metadata,
521 client_options=client_options,
522 method_name=method_name,
523 is_streaming=is_streaming,
524 client_info=client_info,
525 kind=kind,
526 )
527 )