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.
15
16This is used by gapic clients to provide common error mapping, retry, timeout,
17compression, pagination, and long-running operations to methods.
18"""
19
20import functools
21import inspect
22
23from google.api_core import grpc_helpers_async
24from google.api_core.gapic_v1 import client_info
25from google.api_core.gapic_v1.method import ( # noqa: F401
26 DEFAULT,
27 USE_DEFAULT_METADATA,
28 _GapicCallable,
29)
30
31_TRANSPORT_KIND_GRPC_ASYNC = "grpc_asyncio"
32_TRANSPORT_KIND_REST_ASYNC = "rest_asyncio"
33_DEFAULT_ASYNC_TRANSPORT_KIND = _TRANSPORT_KIND_GRPC_ASYNC
34
35
36class _AsyncGapicCallable(_GapicCallable):
37 """Async callable object that wraps an async RPC method with retry, timeout, metadata, and tracing."""
38
39 _SUPPORTED_TRACING_KINDS = (
40 _TRANSPORT_KIND_GRPC_ASYNC,
41 _TRANSPORT_KIND_REST_ASYNC,
42 )
43
44 async def __call__(
45 self, *args, timeout=DEFAULT, retry=DEFAULT, compression=DEFAULT, **kwargs
46 ):
47 """Invoke the low-level async RPC with retry, timeout, compression, and metadata."""
48 wrapped_func = self._prepare_call(timeout, retry, compression, kwargs)
49 with self._trace_span():
50 res = wrapped_func(*args, **kwargs)
51 if inspect.isawaitable(res):
52 return await res
53 return res
54
55
56def wrap_method(
57 func,
58 default_retry=None,
59 default_timeout=None,
60 default_compression=None,
61 client_info=client_info.DEFAULT_CLIENT_INFO,
62 kind=_DEFAULT_ASYNC_TRANSPORT_KIND,
63 *,
64 client_options=None,
65 method_name=None,
66 is_streaming=False,
67 default_host=None,
68):
69 """Wrap an async RPC method with common behavior.
70
71 Args:
72 func (Callable): The low-level async RPC method.
73 default_retry (Optional[google.api_core.retry_async.AsyncRetry]): The default
74 retry strategy. If ``None``, the method will not retry by default.
75 default_timeout (Optional[Union[google.api_core.timeout.Timeout, float]]): The
76 default timeout strategy. Can also be specified as an int or float. If
77 ``None``, the method will not have a timeout specified by default.
78 default_compression (Optional[grpc.Compression]): The default
79 grpc.Compression. If ``None``, the method will not have
80 compression specified by default.
81 client_info (Optional[google.api_core.gapic_v1.client_info.ClientInfo]):
82 Client information used to create a user-agent string that's
83 passed as gRPC metadata to the method. If unspecified, then
84 a sane default will be used. If ``None``, then no user agent
85 metadata will be provided to the RPC method.
86 kind (str): The transport kind for the RPC method. Defaults to "grpc_asyncio".
87 Allowed values for OpenTelemetry method tracing are "grpc_asyncio" and "rest_asyncio".
88 client_options
89 (Optional[google.api_core.client_options.ClientOptions]):
90 Client options used to configure client-level behavior, such as
91 custom OpenTelemetry tracer providers. Defaults to None.
92 method_name (Optional[str]): Optional explicit full RPC method name
93 (e.g. "/google.cloud.secretmanager.v1.SecretManagerService/AccessSecretVersion").
94 Used to identify the RPC for observability.
95 is_streaming (bool): Whether the RPC method is streaming. Defaults to False.
96 Streaming methods are currently gated and do not generate Tier 3 spans.
97 default_host (Optional[str]): Fallback transport host when ``api_endpoint``
98 is not set on ``client_options``. Defaults to None.
99
100 Returns:
101 Callable: A new callable that takes optional ``retry``, ``timeout``,
102 and ``compression`` arguments and applies the common error mapping,
103 retry, timeout, metadata, and compression behavior to the low-level RPC method.
104 """
105 if kind == _DEFAULT_ASYNC_TRANSPORT_KIND:
106 func = grpc_helpers_async.wrap_errors(func)
107
108 metadata = [client_info.to_grpc_metadata()] if client_info is not None else None
109
110 return functools.wraps(func)(
111 _AsyncGapicCallable(
112 func,
113 default_retry,
114 default_timeout,
115 default_compression,
116 metadata=metadata,
117 client_options=client_options,
118 method_name=method_name,
119 is_streaming=is_streaming,
120 client_info=client_info,
121 kind=kind,
122 default_host=default_host,
123 )
124 )