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 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 gRPC methods.
18"""
19
20import functools
21
22from google.api_core import grpc_helpers_async
23from google.api_core.gapic_v1 import client_info
24from google.api_core.gapic_v1.method import _GapicCallable
25from google.api_core.gapic_v1.method import DEFAULT # noqa: F401
26from google.api_core.gapic_v1.method import USE_DEFAULT_METADATA # noqa: F401
27
28_DEFAULT_ASYNC_TRANSPORT_KIND = "grpc_asyncio"
29
30
31def wrap_method(
32 func,
33 default_retry=None,
34 default_timeout=None,
35 default_compression=None,
36 client_info=client_info.DEFAULT_CLIENT_INFO,
37 kind=_DEFAULT_ASYNC_TRANSPORT_KIND,
38):
39 """Wrap an async RPC method with common behavior.
40
41 Returns:
42 Callable: A new callable that takes optional ``retry``, ``timeout``,
43 and ``compression`` arguments and applies the common error mapping,
44 retry, timeout, metadata, and compression behavior to the low-level RPC method.
45 """
46 if kind == _DEFAULT_ASYNC_TRANSPORT_KIND:
47 func = grpc_helpers_async.wrap_errors(func)
48
49 metadata = [client_info.to_grpc_metadata()] if client_info is not None else None
50
51 return functools.wraps(func)(
52 _GapicCallable(
53 func,
54 default_retry,
55 default_timeout,
56 default_compression,
57 metadata=metadata,
58 )
59 )