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):
68 """Wrap an async RPC method with common behavior.
69
70 Args:
71 func (Callable): The low-level async RPC method.
72 default_retry (Optional[google.api_core.retry_async.AsyncRetry]): The default
73 retry strategy. If ``None``, the method will not retry by default.
74 default_timeout (Optional[Union[google.api_core.timeout.Timeout, float]]): The
75 default timeout strategy. Can also be specified as an int or float. If
76 ``None``, the method will not have a timeout specified by default.
77 default_compression (Optional[grpc.Compression]): The default
78 grpc.Compression. If ``None``, the method will not have
79 compression specified by default.
80 client_info (Optional[google.api_core.gapic_v1.client_info.ClientInfo]):
81 Client information used to create a user-agent string that's
82 passed as gRPC metadata to the method. If unspecified, then
83 a sane default will be used. If ``None``, then no user agent
84 metadata will be provided to the RPC method.
85 kind (str): The transport kind for the RPC method. Defaults to "grpc_asyncio".
86 Allowed values for OpenTelemetry method tracing are "grpc_asyncio" and "rest_asyncio".
87 client_options
88 (Optional[google.api_core.client_options.ClientOptions]):
89 Client options used to configure client-level behavior, such as
90 custom OpenTelemetry tracer providers. Defaults to None.
91 method_name (Optional[str]): Optional explicit full RPC method name
92 (e.g. "/google.cloud.secretmanager.v1.SecretManagerService/AccessSecretVersion").
93 Used to identify the RPC for observability.
94 is_streaming (bool): Whether the RPC method is streaming. Defaults to False.
95 Streaming methods are currently gated and do not generate Tier 3 spans.
96
97 Returns:
98 Callable: A new callable that takes optional ``retry``, ``timeout``,
99 and ``compression`` arguments and applies the common error mapping,
100 retry, timeout, metadata, and compression behavior to the low-level RPC method.
101 """
102 if kind == _DEFAULT_ASYNC_TRANSPORT_KIND:
103 func = grpc_helpers_async.wrap_errors(func)
104
105 metadata = [client_info.to_grpc_metadata()] if client_info is not None else None
106
107 return functools.wraps(func)(
108 _AsyncGapicCallable(
109 func,
110 default_retry,
111 default_timeout,
112 default_compression,
113 metadata=metadata,
114 client_options=client_options,
115 method_name=method_name,
116 is_streaming=is_streaming,
117 client_info=client_info,
118 kind=kind,
119 )
120 )