????
Your IP : 216.73.216.210
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""AsyncIO helpers for :mod:`grpc` supporting 3.7+.
Please combine more detailed docstring in grpc_helpers.py to use following
functions. This module is implementing the same surface with AsyncIO semantics.
"""
import asyncio
import functools
from typing import AsyncGenerator, Generic, Iterator, Optional, TypeVar
import grpc
from grpc import aio
from google.api_core import exceptions, grpc_helpers
# denotes the proto response type for grpc calls
P = TypeVar("P")
# NOTE(lidiz) Alternatively, we can hack "__getattribute__" to perform
# automatic patching for us. But that means the overhead of creating an
# extra Python function spreads to every single send and receive.
class _WrappedCall(aio.Call):
def __init__(self):
self._call = None
def with_call(self, call):
"""Supplies the call object separately to keep __init__ clean."""
self._call = call
return self
async def initial_metadata(self):
return await self._call.initial_metadata()
async def trailing_metadata(self):
return await self._call.trailing_metadata()
async def code(self):
return await self._call.code()
async def details(self):
return await self._call.details()
def cancelled(self):
return self._call.cancelled()
def done(self):
return self._call.done()
def time_remaining(self):
return self._call.time_remaining()
def cancel(self):
return self._call.cancel()
def add_done_callback(self, callback):
self._call.add_done_callback(callback)
async def wait_for_connection(self):
try:
await self._call.wait_for_connection()
except grpc.RpcError as rpc_error:
raise exceptions.from_grpc_error(rpc_error) from rpc_error
class _WrappedUnaryResponseMixin(Generic[P], _WrappedCall):
def __await__(self) -> Iterator[P]:
try:
response = yield from self._call.__await__()
return response
except grpc.RpcError as rpc_error:
raise exceptions.from_grpc_error(rpc_error) from rpc_error
class _WrappedStreamResponseMixin(Generic[P], _WrappedCall):
def __init__(self):
self._wrapped_async_generator = None
async def read(self) -> P:
try:
return await self._call.read()
except grpc.RpcError as rpc_error:
raise exceptions.from_grpc_error(rpc_error) from rpc_error
async def _wrapped_aiter(self) -> AsyncGenerator[P, None]:
try:
# NOTE(lidiz) coverage doesn't understand the exception raised from
# __anext__ method. It is covered by test case:
# test_wrap_stream_errors_aiter_non_rpc_error
async for response in self._call: # pragma: no branch
yield response
except grpc.RpcError as rpc_error:
raise exceptions.from_grpc_error(rpc_error) from rpc_error
def __aiter__(self) -> AsyncGenerator[P, None]:
if not self._wrapped_async_generator:
self._wrapped_async_generator = self._wrapped_aiter()
return self._wrapped_async_generator
class _WrappedStreamRequestMixin(_WrappedCall):
async def write(self, request):
try:
await self._call.write(request)
except grpc.RpcError as rpc_error:
raise exceptions.from_grpc_error(rpc_error) from rpc_error
async def done_writing(self):
try:
await self._call.done_writing()
except grpc.RpcError as rpc_error:
raise exceptions.from_grpc_error(rpc_error) from rpc_error
# NOTE(lidiz) Implementing each individual class separately, so we don't
# expose any API that should not be seen. E.g., __aiter__ in unary-unary
# RPC, or __await__ in stream-stream RPC.
class _WrappedUnaryUnaryCall(_WrappedUnaryResponseMixin[P], aio.UnaryUnaryCall):
"""Wrapped UnaryUnaryCall to map exceptions."""
class _WrappedUnaryStreamCall(_WrappedStreamResponseMixin[P], aio.UnaryStreamCall):
"""Wrapped UnaryStreamCall to map exceptions."""
class _WrappedStreamUnaryCall(
_WrappedUnaryResponseMixin[P], _WrappedStreamRequestMixin, aio.StreamUnaryCall
):
"""Wrapped StreamUnaryCall to map exceptions."""
class _WrappedStreamStreamCall(
_WrappedStreamRequestMixin, _WrappedStreamResponseMixin[P], aio.StreamStreamCall
):
"""Wrapped StreamStreamCall to map exceptions."""
# public type alias denoting the return type of async streaming gapic calls
GrpcAsyncStream = _WrappedStreamResponseMixin[P]
# public type alias denoting the return type of unary gapic calls
AwaitableGrpcCall = _WrappedUnaryResponseMixin[P]
def _wrap_unary_errors(callable_):
"""Map errors for Unary-Unary async callables."""
grpc_helpers._patch_callable_name(callable_)
@functools.wraps(callable_)
def error_remapped_callable(*args, **kwargs):
call = callable_(*args, **kwargs)
return _WrappedUnaryUnaryCall().with_call(call)
return error_remapped_callable
def _wrap_stream_errors(callable_):
"""Map errors for streaming RPC async callables."""
grpc_helpers._patch_callable_name(callable_)
@functools.wraps(callable_)
async def error_remapped_callable(*args, **kwargs):
call = callable_(*args, **kwargs)
if isinstance(call, aio.UnaryStreamCall):
call = _WrappedUnaryStreamCall().with_call(call)
elif isinstance(call, aio.StreamUnaryCall):
call = _WrappedStreamUnaryCall().with_call(call)
elif isinstance(call, aio.StreamStreamCall):
call = _WrappedStreamStreamCall().with_call(call)
else:
raise TypeError("Unexpected type of call %s" % type(call))
await call.wait_for_connection()
return call
return error_remapped_callable
def wrap_errors(callable_):
"""Wrap a gRPC async callable and map :class:`grpc.RpcErrors` to
friendly error classes.
Errors raised by the gRPC callable are mapped to the appropriate
:class:`google.api_core.exceptions.GoogleAPICallError` subclasses. The
original `grpc.RpcError` (which is usually also a `grpc.Call`) is
available from the ``response`` property on the mapped exception. This
is useful for extracting metadata from the original error.
Args:
callable_ (Callable): A gRPC callable.
Returns: Callable: The wrapped gRPC callable.
"""
if isinstance(callable_, aio.UnaryUnaryMultiCallable):
return _wrap_unary_errors(callable_)
else:
return _wrap_stream_errors(callable_)
def create_channel(
target,
credentials=None,
scopes=None,
ssl_credentials=None,
credentials_file=None,
quota_project_id=None,
default_scopes=None,
default_host=None,
compression=None,
attempt_direct_path: Optional[bool] = False,
**kwargs
):
"""Create an AsyncIO secure channel with credentials.
Args:
target (str): The target service address in the format 'hostname:port'.
credentials (google.auth.credentials.Credentials): The credentials. If
not specified, then this function will attempt to ascertain the
credentials from the environment using :func:`google.auth.default`.
scopes (Sequence[str]): A optional list of scopes needed for this
service. These are only used when credentials are not specified and
are passed to :func:`google.auth.default`.
ssl_credentials (grpc.ChannelCredentials): Optional SSL channel
credentials. This can be used to specify different certificates.
credentials_file (str): A file with credentials that can be loaded with
:func:`google.auth.load_credentials_from_file`. This argument is
mutually exclusive with credentials.
quota_project_id (str): An optional project to use for billing and quota.
default_scopes (Sequence[str]): Default scopes passed by a Google client
library. Use 'scopes' for user-defined scopes.
default_host (str): The default endpoint. e.g., "pubsub.googleapis.com".
compression (grpc.Compression): An optional value indicating the
compression method to be used over the lifetime of the channel.
attempt_direct_path (Optional[bool]): If set, Direct Path will be attempted
when the request is made. Direct Path is only available within a Google
Compute Engine (GCE) environment and provides a proxyless connection
which increases the available throughput, reduces latency, and increases
reliability. Note:
- This argument should only be set in a GCE environment and for Services
that are known to support Direct Path.
- If this argument is set outside of GCE, then this request will fail
unless the back-end service happens to have configured fall-back to DNS.
- If the request causes a `ServiceUnavailable` response, it is recommended
that the client repeat the request with `attempt_direct_path` set to
`False` as the Service may not support Direct Path.
- Using `ssl_credentials` with `attempt_direct_path` set to `True` will
result in `ValueError` as this combination is not yet supported.
kwargs: Additional key-word args passed to :func:`aio.secure_channel`.
Returns:
aio.Channel: The created channel.
Raises:
google.api_core.DuplicateCredentialArgs: If both a credentials object and credentials_file are passed.
ValueError: If `ssl_credentials` is set and `attempt_direct_path` is set to `True`.
"""
# If `ssl_credentials` is set and `attempt_direct_path` is set to `True`,
# raise ValueError as this is not yet supported.
# See https://github.com/googleapis/python-api-core/issues/590
if ssl_credentials and attempt_direct_path:
raise ValueError("Using ssl_credentials with Direct Path is not supported")
composite_credentials = grpc_helpers._create_composite_credentials(
credentials=credentials,
credentials_file=credentials_file,
scopes=scopes,
default_scopes=default_scopes,
ssl_credentials=ssl_credentials,
quota_project_id=quota_project_id,
default_host=default_host,
)
if attempt_direct_path:
target = grpc_helpers._modify_target_for_direct_path(target)
return aio.secure_channel(
target, composite_credentials, compression=compression, **kwargs
)
class FakeUnaryUnaryCall(_WrappedUnaryUnaryCall):
"""Fake implementation for unary-unary RPCs.
It is a dummy object for response message. Supply the intended response
upon the initialization, and the coroutine will return the exact response
message.
"""
def __init__(self, response=object()):
self.response = response
self._future = asyncio.get_event_loop().create_future()
self._future.set_result(self.response)
def __await__(self):
response = yield from self._future.__await__()
return response
class FakeStreamUnaryCall(_WrappedStreamUnaryCall):
"""Fake implementation for stream-unary RPCs.
It is a dummy object for response message. Supply the intended response
upon the initialization, and the coroutine will return the exact response
message.
"""
def __init__(self, response=object()):
self.response = response
self._future = asyncio.get_event_loop().create_future()
self._future.set_result(self.response)
def __await__(self):
response = yield from self._future.__await__()
return response
async def wait_for_connection(self):
pass