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