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 else:
220 # Fallback per OpenTelemetry Semantic Conventions: every failed span should record
221 # a low-cardinality error.type. Use canonical status code name or exception class name.
222 status_code = _extract_status_code(target_exc)
223 attrs["error.type"] = status_code or target_exc.__class__.__name__
224 metadata = getattr(source, "metadata", None)
225 if metadata:
226 for k, v in metadata.items():
227 attrs[f"gcp.errors.metadata.{k}"] = str(v)
228
229 # 5. Extract human-readable error description for cross-language PRD parity
230 message = getattr(target_exc, "message", None)
231 if not message and hasattr(target_exc, "details") and callable(target_exc.details):
232 message = target_exc.details()
233 if not message and isinstance(target_exc, Exception):
234 message = str(target_exc)
235 if message:
236 attrs["status.message"] = str(message)
237
238 return attrs
239
240
241class _GapicCallable(object):
242 """Callable that applies retry, timeout, and metadata logic.
243
244 Args:
245 target (Callable): The low-level RPC method.
246 retry (google.api_core.retry.Retry): The default retry for the
247 callable. If ``None``, this callable will not retry by default
248 timeout (google.api_core.timeout.Timeout): The default timeout for the
249 callable (i.e. duration of time within which an RPC must terminate
250 after its start, not to be confused with deadline). If ``None``,
251 this callable will not specify a timeout argument to the low-level
252 RPC method.
253 compression (grpc.Compression): The default compression for the callable.
254 If ``None``, this callable will not specify a compression argument
255 to the low-level RPC method.
256 metadata (Sequence[Tuple[str, str]]): Additional metadata that is
257 provided to the RPC method on every invocation. This is merged with
258 any metadata specified during invocation. If ``None``, no
259 additional metadata will be passed to the RPC method.
260 client_options
261 (Optional[google.api_core.client_options.ClientOptions]):
262 Client options used to configure client-level behavior, such as
263 custom OpenTelemetry tracer providers. Defaults to None.
264 method_name (Optional[str]): The optional explicit full RPC method name
265 (e.g. "/google.cloud.secretmanager.v1.SecretManagerService/AccessSecretVersion").
266 is_streaming (bool): Whether the RPC method is streaming. Defaults to False.
267 Note: Streaming methods do not currently generate Tier 3 observability spans.
268 client_info (Optional[google.api_core.gapic_v1.client_info.ClientInfo]):
269 Client information used for metadata headers. Defaults to None.
270 kind (str): The transport kind for the RPC method. Defaults to "grpc".
271 Allowed values for OpenTelemetry method tracing are "grpc" and "rest".
272 """
273
274 def __init__(
275 self,
276 target,
277 retry,
278 timeout,
279 compression,
280 metadata=None,
281 client_options=None,
282 method_name=None,
283 is_streaming=False,
284 client_info=None,
285 kind="grpc",
286 ):
287 self._target = target
288 self._retry = retry
289 self._timeout = timeout
290 self._compression = compression
291
292 # Pre-extract the x-goog-api-client header from the initialized metadata.
293 self._x_goog_api_client, remaining = _extract_metrics_header(metadata)
294 self._static_metadata = tuple(remaining)
295 if self._x_goog_api_client:
296 self._default_metadata = (
297 (METRICS_METADATA_KEY, self._x_goog_api_client),
298 *self._static_metadata,
299 )
300 else:
301 self._default_metadata = self._static_metadata
302
303 # Configure the OpenTelemetry span factory once at initialization.
304 # For now, method tracing is gated to non-streaming gRPC calls where an explicit method_name is provided.
305 self._start_span_fn = None
306 if (
307 not is_streaming
308 and kind in ("grpc", "rest")
309 and method_name is not None
310 and _observability.is_otel_capabilities_enabled(client_options)
311 ):
312 try:
313 from opentelemetry import trace
314
315 tracer_provider = None
316 if isinstance(client_options, dict):
317 tracer_provider = client_options.get("tracer_provider")
318 elif client_options is not None:
319 tracer_provider = getattr(client_options, "tracer_provider", None)
320 if tracer_provider is not None:
321 tracer = tracer_provider.get_tracer("google.api_core")
322 else:
323 tracer = trace.get_tracer("google.api_core")
324
325 span_name, _, _ = _extract_rpc_identity(method_name)
326 is_rest = kind in ("rest", "rest_asyncio")
327 span_attributes = {
328 "rpc.system.name": "http" if is_rest else "grpc",
329 "rpc.method": span_name,
330 }
331 self._start_span_fn = functools.partial(
332 tracer.start_as_current_span,
333 span_name,
334 kind=trace.SpanKind.CLIENT,
335 attributes=span_attributes,
336 )
337 except (ImportError, AttributeError, TypeError):
338 # Gracefully disable tracing if OpenTelemetry or custom provider fails
339 self._start_span_fn = None
340
341 def __call__(
342 self, *args, timeout=DEFAULT, retry=DEFAULT, compression=DEFAULT, **kwargs
343 ):
344 """Invoke the low-level RPC with retry, timeout, compression, and metadata."""
345
346 if retry is DEFAULT:
347 retry = self._retry
348
349 if timeout is DEFAULT:
350 timeout = self._timeout
351
352 if compression is DEFAULT:
353 compression = self._compression
354
355 if isinstance(timeout, (int, float)):
356 timeout = TimeToDeadlineTimeout(timeout=timeout)
357
358 # Apply all applicable decorators.
359 wrapped_func = _apply_decorators(self._target, [retry, timeout])
360
361 if user_metadata := kwargs.get("metadata"):
362 # Add the user agent metadata to the call.
363 final_metadata = list(self._static_metadata)
364 user_x_goog, remaining = _extract_metrics_header(user_metadata)
365
366 merged_header = _deduplicate_metadata_tokens(
367 self._x_goog_api_client, user_x_goog
368 )
369 if merged_header:
370 final_metadata.append((METRICS_METADATA_KEY, merged_header))
371 final_metadata.extend(remaining)
372 kwargs["metadata"] = final_metadata
373 elif self._default_metadata:
374 kwargs["metadata"] = self._default_metadata
375
376 if compression is not None:
377 kwargs["compression"] = compression
378
379 span_cm = contextlib.nullcontext()
380 if self._start_span_fn is not None:
381 try:
382 span_cm = self._start_span_fn()
383 except (
384 Exception
385 ): # Fail-open: proceed without span if tracing initialization fails
386 span_cm = contextlib.nullcontext()
387
388 with span_cm as span:
389 try:
390 result = wrapped_func(*args, **kwargs)
391 if span is not None and hasattr(span, "set_attribute"):
392 span.set_attribute("rpc.response.status_code", "OK")
393 return result
394 except Exception as exc:
395 if span is not None and hasattr(span, "set_attribute"):
396 span.set_attribute(
397 "rpc.response.status_code", _extract_status_code(exc)
398 )
399 for k, v in _extract_error_attributes(exc).items():
400 span.set_attribute(k, v)
401 raise
402
403
404def wrap_method(
405 func,
406 default_retry=None,
407 default_timeout=None,
408 default_compression=None,
409 client_info=client_info.DEFAULT_CLIENT_INFO,
410 *,
411 with_call=False,
412 client_options=None,
413 method_name=None,
414 is_streaming=False,
415 kind="grpc",
416):
417 """Wrap an RPC method with common behavior.
418
419 This applies common error wrapping, retry, timeout, and compression behavior to a function.
420 The wrapped function will take optional ``retry``, ``timeout``, and ``compression``
421 arguments.
422
423 For example::
424
425 import google.api_core.gapic_v1.method
426 from google.api_core import retry
427 from google.api_core import timeout
428 from grpc import Compression
429
430 # The original RPC method.
431 def get_topic(name, timeout=None):
432 request = publisher_v2.GetTopicRequest(name=name)
433 return publisher_stub.GetTopic(request, timeout=timeout)
434
435 default_retry = retry.Retry(deadline=60)
436 default_timeout = timeout.Timeout(deadline=60)
437 default_compression = Compression.NoCompression
438 wrapped_get_topic = google.api_core.gapic_v1.method.wrap_method(
439 get_topic, default_retry)
440
441 # Execute get_topic with default retry and timeout:
442 response = wrapped_get_topic()
443
444 # Execute get_topic without doing any retying but with the default
445 # timeout:
446 response = wrapped_get_topic(retry=None)
447
448 # Execute get_topic but only retry on 5xx errors:
449 my_retry = retry.Retry(retry.if_exception_type(
450 exceptions.InternalServerError))
451 response = wrapped_get_topic(retry=my_retry)
452
453 The way this works is by late-wrapping the given function with the retry
454 and timeout decorators. Essentially, when ``wrapped_get_topic()`` is
455 called:
456
457 * ``get_topic()`` is first wrapped with the ``timeout`` into
458 ``get_topic_with_timeout``.
459 * ``get_topic_with_timeout`` is wrapped with the ``retry`` into
460 ``get_topic_with_timeout_and_retry()``.
461 * The final ``get_topic_with_timeout_and_retry`` is called passing through
462 the ``args`` and ``kwargs``.
463
464 The callstack is therefore::
465
466 method.__call__() ->
467 Retry.__call__() ->
468 Timeout.__call__() ->
469 wrap_errors() ->
470 get_topic()
471
472 Note that if ``timeout`` or ``retry`` is ``None``, then they are not
473 applied to the function. For example,
474 ``wrapped_get_topic(timeout=None, retry=None)`` is more or less
475 equivalent to just calling ``get_topic`` but with error re-mapping.
476
477 Args:
478 func (Callable[Any]): The function to wrap. It should accept an
479 optional ``timeout`` argument. If ``metadata`` is not ``None``, it
480 should accept a ``metadata`` argument.
481 default_retry (Optional[google.api_core.Retry]): The default retry
482 strategy. If ``None``, the method will not retry by default.
483 default_timeout (Optional[google.api_core.Timeout]): The default
484 timeout strategy. Can also be specified as an int or float. If
485 ``None``, the method will not have timeout specified by default.
486 default_compression (Optional[grpc.Compression]): The default
487 grpc.Compression. If ``None``, the method will not have
488 compression specified by default.
489 client_info
490 (Optional[google.api_core.gapic_v1.client_info.ClientInfo]):
491 Client information used to create a user-agent string that's
492 passed as gRPC metadata to the method. If unspecified, then
493 a sane default will be used. If ``None``, then no user agent
494 metadata will be provided to the RPC method.
495 with_call (bool): If True, wrapped grpc.UnaryUnaryMulticallables will
496 return a tuple of (response, grpc.Call) instead of just the response.
497 This is useful for extracting trailing metadata from unary calls.
498 Defaults to False.
499 client_options
500 (Optional[google.api_core.client_options.ClientOptions]):
501 Client options used to configure client-level behavior, such as
502 custom OpenTelemetry tracer providers. Defaults to None.
503 method_name (Optional[str]): Optional explicit full RPC method name
504 (e.g. "/google.cloud.secretmanager.v1.SecretManagerService/AccessSecretVersion").
505 Used to identify the RPC for observability.
506 is_streaming (bool): Whether the RPC method is streaming. Defaults to False.
507 Streaming methods are currently gated and do not generate Tier 3 spans.
508 kind (str): The transport kind for the RPC method. Defaults to "grpc".
509 Allowed values for OpenTelemetry method tracing are "grpc" and "rest".
510
511 Returns:
512 Callable: A new callable that takes optional ``retry``, ``timeout``,
513 and ``compression``
514 arguments and applies the common error mapping, retry, timeout, compression,
515 and metadata behavior to the low-level RPC method.
516 """
517 if with_call:
518 try:
519 func = func.with_call
520 except AttributeError as exc:
521 raise ValueError(
522 "with_call=True is only supported for unary calls."
523 ) from exc
524 func = grpc_helpers.wrap_errors(func)
525 if client_info is not None:
526 user_agent_metadata = [client_info.to_grpc_metadata()]
527 else:
528 user_agent_metadata = None
529
530 return functools.wraps(func)(
531 _GapicCallable(
532 func,
533 default_retry,
534 default_timeout,
535 default_compression,
536 metadata=user_agent_metadata,
537 client_options=client_options,
538 method_name=method_name,
539 is_streaming=is_streaming,
540 client_info=client_info,
541 kind=kind,
542 )
543 )