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