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 2020 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"""AsyncIO helpers for wrapping gRPC and REST methods with common functionality.
16This is used by gapic clients to provide common error mapping, retry, timeout,
17compression, pagination, and long-running operations to methods.
18"""
20import asyncio
21import contextlib
22import functools
23import inspect
25from google.api_core import _observability, grpc_helpers_async
26from google.api_core.gapic_v1 import client_info
27from google.api_core.gapic_v1.client_info import METRICS_METADATA_KEY
29# Retain _GapicCallable import for backward compatibility with external packages
30from google.api_core.gapic_v1.method import ( # noqa: F401
31 DEFAULT,
32 USE_DEFAULT_METADATA,
33 _apply_decorators,
34 _deduplicate_metadata_tokens,
35 _extract_error_attributes,
36 _extract_metrics_header,
37 _extract_rpc_identity,
38 _extract_status_code,
39 _GapicCallable,
40)
41from google.api_core.timeout import TimeToDeadlineTimeout
43_DEFAULT_ASYNC_TRANSPORT_KIND = "grpc_asyncio"
46class _AsyncGapicCallable(object):
47 """Async callable object that wraps an async RPC method with retry, timeout, metadata, and tracing.
49 Args:
50 target (Callable): The low-level async RPC method.
51 retry (Optional[google.api_core.retry_async.AsyncRetry]): The default retry for the
52 callable. If ``None``, this callable will not retry by default.
53 timeout (Optional[Union[google.api_core.timeout.Timeout, float]]): The default timeout for the
54 callable. If ``None``, this callable will not specify a timeout argument to the
55 low-level RPC method.
56 compression (Optional[grpc.Compression]): The default compression for the callable.
57 If ``None``, this callable will not specify a compression argument to the low-level
58 RPC method.
59 metadata (Optional[Sequence[Tuple[str, str]]]): Additional metadata that is
60 provided to the RPC method on every invocation. This is merged with
61 any metadata specified during invocation. If ``None``, no
62 additional metadata will be passed to the RPC method.
63 client_options (Optional[google.api_core.client_options.ClientOptions]):
64 Client options used to configure client-level behavior, such as
65 custom OpenTelemetry tracer providers. Defaults to None.
66 method_name (Optional[str]): The optional explicit full RPC method name
67 (e.g. "/google.cloud.secretmanager.v1.SecretManagerService/AccessSecretVersion").
68 is_streaming (bool): Whether the RPC method is streaming. Defaults to False.
69 Note: Streaming methods do not currently generate Tier 3 observability spans.
70 client_info (Optional[google.api_core.gapic_v1.client_info.ClientInfo]):
71 Client information used for metadata headers. Defaults to None.
72 kind (str): The transport kind for the RPC method. Defaults to "grpc_asyncio".
73 Allowed values for OpenTelemetry method tracing are "grpc_asyncio" and "rest_asyncio".
74 """
76 def __init__(
77 self,
78 target,
79 retry,
80 timeout,
81 compression,
82 metadata=None,
83 client_options=None,
84 method_name=None,
85 is_streaming=False,
86 client_info=None,
87 kind=_DEFAULT_ASYNC_TRANSPORT_KIND,
88 ):
89 self._target = target
90 self._retry = retry
91 self._timeout = timeout
92 self._compression = compression
94 # Pre-extract the x-goog-api-client header from the initialized metadata.
95 self._x_goog_api_client, remaining = _extract_metrics_header(metadata)
96 self._static_metadata = tuple(remaining)
97 if self._x_goog_api_client:
98 self._default_metadata = (
99 (METRICS_METADATA_KEY, self._x_goog_api_client),
100 *self._static_metadata,
101 )
102 else:
103 self._default_metadata = self._static_metadata
105 # Configure the OpenTelemetry span factory once at initialization.
106 self._start_span_fn = None
107 if (
108 not is_streaming
109 and kind in ("grpc_asyncio", "rest_asyncio")
110 and method_name is not None
111 and _observability.is_otel_capabilities_enabled(client_options)
112 ):
113 try:
114 from opentelemetry import trace
116 tracer_provider = None
117 if isinstance(client_options, dict):
118 tracer_provider = client_options.get("tracer_provider")
119 elif client_options is not None:
120 tracer_provider = getattr(client_options, "tracer_provider", None)
121 if tracer_provider is not None:
122 tracer = tracer_provider.get_tracer("google.api_core")
123 else:
124 tracer = trace.get_tracer("google.api_core")
126 span_name, _, _ = _extract_rpc_identity(method_name)
127 is_rest = kind in ("rest", "rest_asyncio")
128 span_attributes = {
129 "rpc.system.name": "http" if is_rest else "grpc",
130 "rpc.method": span_name,
131 }
132 self._start_span_fn = functools.partial(
133 tracer.start_as_current_span,
134 span_name,
135 kind=trace.SpanKind.CLIENT,
136 attributes=span_attributes,
137 )
138 except (ImportError, AttributeError, TypeError):
139 # Gracefully disable tracing if OpenTelemetry or custom provider fails
140 self._start_span_fn = None
142 async def __call__(
143 self, *args, timeout=DEFAULT, retry=DEFAULT, compression=DEFAULT, **kwargs
144 ):
145 """Invoke the low-level async RPC with retry, timeout, compression, and metadata."""
146 if retry is DEFAULT:
147 retry = self._retry
149 if timeout is DEFAULT:
150 timeout = self._timeout
152 if compression is DEFAULT:
153 compression = self._compression
155 if isinstance(timeout, (int, float)):
156 timeout = TimeToDeadlineTimeout(timeout=timeout)
158 # Apply all applicable decorators.
159 wrapped_func = _apply_decorators(self._target, [retry, timeout])
161 if user_metadata := kwargs.get("metadata"):
162 # Add the user agent metadata to the call.
163 final_metadata = list(self._static_metadata)
164 user_x_goog, remaining = _extract_metrics_header(user_metadata)
166 merged_header = _deduplicate_metadata_tokens(
167 self._x_goog_api_client, user_x_goog
168 )
169 if merged_header:
170 final_metadata.append((METRICS_METADATA_KEY, merged_header))
171 final_metadata.extend(remaining)
172 kwargs["metadata"] = final_metadata
173 elif self._default_metadata:
174 kwargs["metadata"] = self._default_metadata
176 if compression is not None:
177 kwargs["compression"] = compression
179 span_cm = contextlib.nullcontext()
180 if self._start_span_fn is not None:
181 try:
182 span_cm = self._start_span_fn()
183 except (
184 Exception
185 ): # Fail-open: proceed without span if tracing initialization fails
186 span_cm = contextlib.nullcontext()
188 with span_cm as span:
189 try:
190 res = wrapped_func(*args, **kwargs)
191 if inspect.isawaitable(res):
192 result = await res
193 else:
194 result = res
195 if span is not None and hasattr(span, "set_attribute"):
196 span.set_attribute("rpc.response.status_code", "OK")
197 return result
198 except (Exception, asyncio.CancelledError) as exc:
199 if span is not None and hasattr(span, "set_attribute"):
200 span.set_attribute(
201 "rpc.response.status_code", _extract_status_code(exc)
202 )
203 for k, v in _extract_error_attributes(exc).items():
204 span.set_attribute(k, v)
205 raise
208def wrap_method(
209 func,
210 default_retry=None,
211 default_timeout=None,
212 default_compression=None,
213 client_info=client_info.DEFAULT_CLIENT_INFO,
214 kind=_DEFAULT_ASYNC_TRANSPORT_KIND,
215 *,
216 client_options=None,
217 method_name=None,
218 is_streaming=False,
219):
220 """Wrap an async RPC method with common behavior.
222 Returns:
223 Callable: A new callable that takes optional ``retry``, ``timeout``,
224 and ``compression`` arguments and applies the common error mapping,
225 retry, timeout, metadata, and compression behavior to the low-level RPC method.
226 """
227 if kind == _DEFAULT_ASYNC_TRANSPORT_KIND:
228 func = grpc_helpers_async.wrap_errors(func)
230 metadata = [client_info.to_grpc_metadata()] if client_info is not None else None
232 return functools.wraps(func)(
233 _AsyncGapicCallable(
234 func,
235 default_retry,
236 default_timeout,
237 default_compression,
238 metadata=metadata,
239 client_options=client_options,
240 method_name=method_name,
241 is_streaming=is_streaming,
242 client_info=client_info,
243 kind=kind,
244 )
245 )