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.
15"""AsyncIO helpers for :mod:`grpc`.
17Please combine more detailed docstring in grpc_helpers.py to use following
18functions. This module is implementing the same surface with AsyncIO semantics.
19"""
21import asyncio
22import functools
23import warnings
24from typing import AsyncGenerator, Generic, Iterator, Optional, Sequence, TypeVar
26import grpc
27from grpc import aio
29from google.api_core import exceptions, general_helpers, grpc_helpers
31# denotes the proto response type for grpc calls
32P = TypeVar("P")
34# NOTE(lidiz) Alternatively, we can hack "__getattribute__" to perform
35# automatic patching for us. But that means the overhead of creating an
36# extra Python function spreads to every single send and receive.
39class _WrappedCall(aio.Call):
40 def __init__(self):
41 self._call = None
43 def with_call(self, call):
44 """Supplies the call object separately to keep __init__ clean."""
45 self._call = call
46 return self
48 async def initial_metadata(self):
49 return await self._call.initial_metadata()
51 async def trailing_metadata(self):
52 return await self._call.trailing_metadata()
54 async def code(self):
55 return await self._call.code()
57 async def details(self):
58 return await self._call.details()
60 def cancelled(self):
61 return self._call.cancelled()
63 def done(self):
64 return self._call.done()
66 def time_remaining(self):
67 return self._call.time_remaining()
69 def cancel(self):
70 return self._call.cancel()
72 def add_done_callback(self, callback):
73 self._call.add_done_callback(callback)
75 async def wait_for_connection(self):
76 try:
77 await self._call.wait_for_connection()
78 except grpc.RpcError as rpc_error:
79 raise exceptions.from_grpc_error(rpc_error) from rpc_error
82class _WrappedUnaryResponseMixin(Generic[P], _WrappedCall):
83 def __await__(self) -> Iterator[P]:
84 try:
85 response = yield from self._call.__await__()
86 return response
87 except grpc.RpcError as rpc_error:
88 raise exceptions.from_grpc_error(rpc_error) from rpc_error
91class _WrappedStreamResponseMixin(Generic[P], _WrappedCall):
92 def __init__(self):
93 self._wrapped_async_generator = None
95 async def read(self) -> P:
96 try:
97 return await self._call.read()
98 except grpc.RpcError as rpc_error:
99 raise exceptions.from_grpc_error(rpc_error) from rpc_error
101 async def _wrapped_aiter(self) -> AsyncGenerator[P, None]:
102 try:
103 # NOTE(lidiz) coverage doesn't understand the exception raised from
104 # __anext__ method. It is covered by test case:
105 # test_wrap_stream_errors_aiter_non_rpc_error
106 async for response in self._call: # pragma: no branch
107 yield response
108 except grpc.RpcError as rpc_error:
109 raise exceptions.from_grpc_error(rpc_error) from rpc_error
111 def __aiter__(self) -> AsyncGenerator[P, None]:
112 if not self._wrapped_async_generator:
113 self._wrapped_async_generator = self._wrapped_aiter()
114 return self._wrapped_async_generator
117class _WrappedStreamRequestMixin(_WrappedCall):
118 async def write(self, request):
119 try:
120 await self._call.write(request)
121 except grpc.RpcError as rpc_error:
122 raise exceptions.from_grpc_error(rpc_error) from rpc_error
124 async def done_writing(self):
125 try:
126 await self._call.done_writing()
127 except grpc.RpcError as rpc_error:
128 raise exceptions.from_grpc_error(rpc_error) from rpc_error
131# NOTE(lidiz) Implementing each individual class separately, so we don't
132# expose any API that should not be seen. E.g., __aiter__ in unary-unary
133# RPC, or __await__ in stream-stream RPC.
134class _WrappedUnaryUnaryCall(_WrappedUnaryResponseMixin[P], aio.UnaryUnaryCall):
135 """Wrapped UnaryUnaryCall to map exceptions."""
138class _WrappedUnaryStreamCall(_WrappedStreamResponseMixin[P], aio.UnaryStreamCall):
139 """Wrapped UnaryStreamCall to map exceptions."""
142class _WrappedStreamUnaryCall(
143 _WrappedUnaryResponseMixin[P], _WrappedStreamRequestMixin, aio.StreamUnaryCall
144):
145 """Wrapped StreamUnaryCall to map exceptions."""
148class _WrappedStreamStreamCall(
149 _WrappedStreamRequestMixin, _WrappedStreamResponseMixin[P], aio.StreamStreamCall
150):
151 """Wrapped StreamStreamCall to map exceptions."""
154# public type alias denoting the return type of async streaming gapic calls
155GrpcAsyncStream = _WrappedStreamResponseMixin
156# public type alias denoting the return type of unary gapic calls
157AwaitableGrpcCall = _WrappedUnaryResponseMixin
160def _wrap_unary_errors(callable_):
161 """Map errors for Unary-Unary async callables."""
163 @functools.wraps(callable_)
164 def error_remapped_callable(*args, **kwargs):
165 call = callable_(*args, **kwargs)
166 return _WrappedUnaryUnaryCall().with_call(call)
168 return error_remapped_callable
171def _wrap_stream_errors(callable_, wrapper_type):
172 """Map errors for streaming RPC async callables."""
174 @functools.wraps(callable_)
175 async def error_remapped_callable(*args, **kwargs):
176 call = callable_(*args, **kwargs)
177 call = wrapper_type().with_call(call)
178 await call.wait_for_connection()
179 return call
181 return error_remapped_callable
184def wrap_errors(callable_):
185 """Wrap a gRPC async callable and map :class:`grpc.RpcErrors` to
186 friendly error classes.
188 Errors raised by the gRPC callable are mapped to the appropriate
189 :class:`google.api_core.exceptions.GoogleAPICallError` subclasses. The
190 original `grpc.RpcError` (which is usually also a `grpc.Call`) is
191 available from the ``response`` property on the mapped exception. This
192 is useful for extracting metadata from the original error.
194 Args:
195 callable_ (Callable): A gRPC callable.
197 Returns: Callable: The wrapped gRPC callable.
198 """
199 grpc_helpers._patch_callable_name(callable_)
201 if isinstance(callable_, aio.UnaryStreamMultiCallable):
202 return _wrap_stream_errors(callable_, _WrappedUnaryStreamCall)
203 elif isinstance(callable_, aio.StreamUnaryMultiCallable):
204 return _wrap_stream_errors(callable_, _WrappedStreamUnaryCall)
205 elif isinstance(callable_, aio.StreamStreamMultiCallable):
206 return _wrap_stream_errors(callable_, _WrappedStreamStreamCall)
207 else:
208 return _wrap_unary_errors(callable_)
211def create_channel(
212 target,
213 credentials=None,
214 scopes=None,
215 ssl_credentials=None,
216 credentials_file=None,
217 quota_project_id=None,
218 default_scopes=None,
219 default_host=None,
220 compression=None,
221 attempt_direct_path: Optional[bool] = False,
222 attempt_direct_path_xds_over_interconnect: Optional[bool] = None,
223 **kwargs,
224):
225 """Create an AsyncIO secure channel with credentials.
227 Args:
228 target (str): The target service address in the format 'hostname:port'.
229 credentials (google.auth.credentials.Credentials): The credentials. If
230 not specified, then this function will attempt to ascertain the
231 credentials from the environment using :func:`google.auth.default`.
232 scopes (Sequence[str]): A optional list of scopes needed for this
233 service. These are only used when credentials are not specified and
234 are passed to :func:`google.auth.default`.
235 ssl_credentials (grpc.ChannelCredentials): Optional SSL channel
236 credentials. This can be used to specify different certificates.
237 credentials_file (str): Deprecated. A file with credentials that can be loaded with
238 :func:`google.auth.load_credentials_from_file`. This argument is
239 mutually exclusive with credentials. This argument will be
240 removed in the next major version of `google-api-core`.
242 .. warning::
243 Important: If you accept a credential configuration (credential JSON/File/Stream)
244 from an external source for authentication to Google Cloud Platform, you must
245 validate it before providing it to any Google API or client library. Providing an
246 unvalidated credential configuration to Google APIs or libraries can compromise
247 the security of your systems and data. For more information, refer to
248 `Validate credential configurations from external sources`_.
250 .. _Validate credential configurations from external sources:
252 https://cloud.google.com/docs/authentication/external/externally-sourced-credentials
253 quota_project_id (str): An optional project to use for billing and quota.
254 default_scopes (Sequence[str]): Default scopes passed by a Google client
255 library. Use 'scopes' for user-defined scopes.
256 default_host (str): The default endpoint. e.g., "pubsub.googleapis.com".
257 compression (grpc.Compression): An optional value indicating the
258 compression method to be used over the lifetime of the channel.
259 attempt_direct_path (Optional[bool]): If set, Direct Path will be attempted
260 when the request is made. Direct Path is only available within a Google
261 Compute Engine (GCE) environment and provides a proxyless connection
262 which increases the available throughput, reduces latency, and increases
263 reliability. Note:
265 - This argument should only be set in a GCE environment and for Services
266 that are known to support Direct Path.
267 - If this argument is set outside of GCE, then this request will fail
268 unless the back-end service happens to have configured fall-back to DNS.
269 - If the request causes a `ServiceUnavailable` response, it is recommended
270 that the client repeat the request with `attempt_direct_path` set to
271 `False` as the Service may not support Direct Path.
272 - Using `ssl_credentials` with `attempt_direct_path` set to `True` will
273 result in `ValueError` as this combination is not yet supported.
274 attempt_direct_path_xds_over_interconnect (Optional[bool]): If set,
275 DirectPath over Cloud Interconnect will be attempted using standard
276 TLS credentials and ``?force-xds`` C2P target resolution. If None
277 (default), the setting is resolved from the
278 ``GOOGLE_CLOUD_ENABLE_DIRECT_PATH_XDS_OVER_INTERCONNECT``
279 environment variable.
281 kwargs: Additional key-word args passed to :func:`aio.secure_channel`.
283 Returns:
284 aio.Channel: The created channel.
286 Raises:
287 google.api_core.DuplicateCredentialArgs: If both a credentials object
288 and credentials_file are passed.
289 ValueError: If `ssl_credentials` is set and `attempt_direct_path` is
290 set to `True` without `attempt_direct_path_xds_over_interconnect`.
291 """
293 if credentials_file is not None:
294 warnings.warn(general_helpers._CREDENTIALS_FILE_WARNING, DeprecationWarning)
296 target, composite_credentials, kwargs = (
297 grpc_helpers._setup_direct_path_and_credentials(
298 target=target,
299 credentials=credentials,
300 scopes=scopes,
301 ssl_credentials=ssl_credentials,
302 credentials_file=credentials_file,
303 quota_project_id=quota_project_id,
304 default_scopes=default_scopes,
305 default_host=default_host,
306 attempt_direct_path=attempt_direct_path,
307 attempt_direct_path_xds_over_interconnect=attempt_direct_path_xds_over_interconnect,
308 **kwargs,
309 )
310 )
312 return aio.secure_channel(
313 target, composite_credentials, compression=compression, **kwargs
314 )
317def apply_channel_interceptors(
318 channel: aio.Channel,
319 interceptors: Optional[Sequence[aio.ClientInterceptor]] = None,
320) -> aio.Channel:
321 """Applies client interceptors to a gRPC AsyncIO channel.
323 In grpc.aio, channels maintain internal interceptor lists
324 (_unary_unary_interceptors, etc.). To preserve the public API contract for
325 callers who supply their own pre-instantiated ``channel`` object or a custom
326 channel factory callable (which does not accept ``interceptors``), we attach
327 interceptors post-instantiation directly to the channel's interceptor lists.
329 Args:
330 channel (aio.Channel): The async gRPC channel to intercept.
331 interceptors (Optional[Sequence[aio.ClientInterceptor]]):
332 Additional interceptors to apply to the channel.
334 Returns:
335 aio.Channel: The channel with interceptors attached, or the original channel
336 if no interceptors were provided.
337 """
338 if not interceptors:
339 return channel
341 mapping = (
342 ("intercept_unary_unary", "_unary_unary_interceptors"),
343 ("intercept_unary_stream", "_unary_stream_interceptors"),
344 ("intercept_stream_unary", "_stream_unary_interceptors"),
345 ("intercept_stream_stream", "_stream_stream_interceptors"),
346 )
347 for interceptor in interceptors:
348 matched = False
349 for method_name, attr_name in mapping:
350 if hasattr(interceptor, method_name) and hasattr(channel, attr_name):
351 target_list = getattr(channel, attr_name)
352 if isinstance(target_list, list):
353 if interceptor not in target_list:
354 target_list.append(interceptor)
355 matched = True
356 elif hasattr(target_list, "append"):
357 target_list.append(interceptor)
358 matched = True
359 if not matched and hasattr(channel, "_unary_unary_interceptors"):
360 unary_interceptors = channel._unary_unary_interceptors
361 if isinstance(unary_interceptors, list):
362 if interceptor not in unary_interceptors:
363 unary_interceptors.append(interceptor)
364 elif hasattr(unary_interceptors, "append"):
365 unary_interceptors.append(interceptor)
367 return channel
370class FakeUnaryUnaryCall(_WrappedUnaryUnaryCall):
371 """Fake implementation for unary-unary RPCs.
373 It is a dummy object for response message. Supply the intended response
374 upon the initialization, and the coroutine will return the exact response
375 message.
376 """
378 def __init__(self, response=object()):
379 self.response = response
380 self._future = asyncio.get_event_loop().create_future()
381 self._future.set_result(self.response)
383 def __await__(self):
384 response = yield from self._future.__await__()
385 return response
388class FakeStreamUnaryCall(_WrappedStreamUnaryCall):
389 """Fake implementation for stream-unary RPCs.
391 It is a dummy object for response message. Supply the intended response
392 upon the initialization, and the coroutine will return the exact response
393 message.
394 """
396 def __init__(self, response=object()):
397 self.response = response
398 self._future = asyncio.get_event_loop().create_future()
399 self._future.set_result(self.response)
401 def __await__(self):
402 response = yield from self._future.__await__()
403 return response
405 async def wait_for_connection(self):
406 pass