diff --git a/google/api_core/exceptions.py b/google/api_core/exceptions.py index eed4ee40..d1459ab2 100644 --- a/google/api_core/exceptions.py +++ b/google/api_core/exceptions.py @@ -444,6 +444,10 @@ def from_grpc_status(status_code, message, **kwargs): return error +def _is_informative_grpc_error(rpc_exc): + return hasattr(rpc_exc, "code") and hasattr(rpc_exc, "details") + + def from_grpc_error(rpc_exc): """Create a :class:`GoogleAPICallError` from a :class:`grpc.RpcError`. @@ -454,7 +458,9 @@ def from_grpc_error(rpc_exc): GoogleAPICallError: An instance of the appropriate subclass of :class:`GoogleAPICallError`. """ - if isinstance(rpc_exc, grpc.Call): + # NOTE(lidiz) All gRPC error shares the parent class grpc.RpcError. + # However, check for grpc.RpcError breaks backward compatibility. + if isinstance(rpc_exc, grpc.Call) or _is_informative_grpc_error(rpc_exc): return from_grpc_status( rpc_exc.code(), rpc_exc.details(), errors=(rpc_exc,), response=rpc_exc ) diff --git a/google/api_core/future/async_future.py b/google/api_core/future/async_future.py new file mode 100644 index 00000000..e1d158d0 --- /dev/null +++ b/google/api_core/future/async_future.py @@ -0,0 +1,157 @@ +# 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 implementation of the abstract base Future class.""" + +import asyncio + +from google.api_core import exceptions +from google.api_core import retry +from google.api_core import retry_async +from google.api_core.future import base + + +class _OperationNotComplete(Exception): + """Private exception used for polling via retry.""" + pass + + +RETRY_PREDICATE = retry.if_exception_type( + _OperationNotComplete, + exceptions.TooManyRequests, + exceptions.InternalServerError, + exceptions.BadGateway, +) +DEFAULT_RETRY = retry_async.AsyncRetry(predicate=RETRY_PREDICATE) + + +class AsyncFuture(base.Future): + """A Future that polls peer service to self-update. + + The :meth:`done` method should be implemented by subclasses. The polling + behavior will repeatedly call ``done`` until it returns True. + + .. note: Privacy here is intended to prevent the final class from + overexposing, not to prevent subclasses from accessing methods. + + Args: + retry (google.api_core.retry.Retry): The retry configuration used + when polling. This can be used to control how often :meth:`done` + is polled. Regardless of the retry's ``deadline``, it will be + overridden by the ``timeout`` argument to :meth:`result`. + """ + + def __init__(self, retry=DEFAULT_RETRY): + super().__init__() + self._retry = retry + self._future = asyncio.get_event_loop().create_future() + self._background_task = None + + async def done(self, retry=DEFAULT_RETRY): + """Checks to see if the operation is complete. + + Args: + retry (google.api_core.retry.Retry): (Optional) How to retry the RPC. + + Returns: + bool: True if the operation is complete, False otherwise. + """ + # pylint: disable=redundant-returns-doc, missing-raises-doc + raise NotImplementedError() + + async def _done_or_raise(self): + """Check if the future is done and raise if it's not.""" + result = await self.done() + if not result: + raise _OperationNotComplete() + + async def running(self): + """True if the operation is currently running.""" + result = await self.done() + return not result + + async def _blocking_poll(self, timeout=None): + """Poll and await for the Future to be resolved. + + Args: + timeout (int): + How long (in seconds) to wait for the operation to complete. + If None, wait indefinitely. + """ + if self._future.done(): + return + + retry_ = self._retry.with_deadline(timeout) + + try: + await retry_(self._done_or_raise)() + except exceptions.RetryError: + raise asyncio.TimeoutError( + "Operation did not complete within the designated " "timeout." + ) + + async def result(self, timeout=None): + """Get the result of the operation. + + Args: + timeout (int): + How long (in seconds) to wait for the operation to complete. + If None, wait indefinitely. + + Returns: + google.protobuf.Message: The Operation's result. + + Raises: + google.api_core.GoogleAPICallError: If the operation errors or if + the timeout is reached before the operation completes. + """ + await self._blocking_poll(timeout=timeout) + return self._future.result() + + async def exception(self, timeout=None): + """Get the exception from the operation. + + Args: + timeout (int): How long to wait for the operation to complete. + If None, wait indefinitely. + + Returns: + Optional[google.api_core.GoogleAPICallError]: The operation's + error. + """ + await self._blocking_poll(timeout=timeout) + return self._future.exception() + + def add_done_callback(self, fn): + """Add a callback to be executed when the operation is complete. + + If the operation is completed, the callback will be scheduled onto the + event loop. Otherwise, the callback will be stored and invoked when the + future is done. + + Args: + fn (Callable[Future]): The callback to execute when the operation + is complete. + """ + if self._background_task is None: + self._background_task = asyncio.get_event_loop().create_task(self._blocking_poll()) + self._future.add_done_callback(fn) + + def set_result(self, result): + """Set the Future's result.""" + self._future.set_result(result) + + def set_exception(self, exception): + """Set the Future's exception.""" + self._future.set_exception(exception) diff --git a/google/api_core/gapic_v1/__init__.py b/google/api_core/gapic_v1/__init__.py index e7a7a686..a8121c74 100644 --- a/google/api_core/gapic_v1/__init__.py +++ b/google/api_core/gapic_v1/__init__.py @@ -12,9 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import sys + from google.api_core.gapic_v1 import client_info from google.api_core.gapic_v1 import config from google.api_core.gapic_v1 import method from google.api_core.gapic_v1 import routing_header -__all__ = ["client_info", "config", "method", "routing_header"] +if sys.version_info[0] >= 3 and sys.version_info[1] >= 6: + from google.api_core.gapic_v1 import method_async # noqa: F401 + from google.api_core.gapic_v1 import config_async # noqa: F401 + __all__ = ["client_info", "config", "config_async", "method", "method_async", "routing_header"] +else: + __all__ = ["client_info", "config", "method", "routing_header"] # pragma: NO COVER diff --git a/google/api_core/gapic_v1/config.py b/google/api_core/gapic_v1/config.py index 3a3eb15f..2a56cf1b 100644 --- a/google/api_core/gapic_v1/config.py +++ b/google/api_core/gapic_v1/config.py @@ -45,7 +45,7 @@ def _exception_class_for_grpc_status_name(name): return exceptions.exception_class_for_grpc_status(getattr(grpc.StatusCode, name)) -def _retry_from_retry_config(retry_params, retry_codes): +def _retry_from_retry_config(retry_params, retry_codes, retry_impl=retry.Retry): """Creates a Retry object given a gapic retry configuration. Args: @@ -70,7 +70,7 @@ def _retry_from_retry_config(retry_params, retry_codes): exception_classes = [ _exception_class_for_grpc_status_name(code) for code in retry_codes ] - return retry.Retry( + return retry_impl( retry.if_exception_type(*exception_classes), initial=(retry_params["initial_retry_delay_millis"] / _MILLIS_PER_SECOND), maximum=(retry_params["max_retry_delay_millis"] / _MILLIS_PER_SECOND), @@ -110,7 +110,7 @@ def _timeout_from_retry_config(retry_params): MethodConfig = collections.namedtuple("MethodConfig", ["retry", "timeout"]) -def parse_method_configs(interface_config): +def parse_method_configs(interface_config, retry_impl=retry.Retry): """Creates default retry and timeout objects for each method in a gapic interface config. @@ -120,6 +120,8 @@ def parse_method_configs(interface_config): an interface named ``google.example.v1.ExampleService`` you would pass in just that interface's configuration, for example ``gapic_config['interfaces']['google.example.v1.ExampleService']``. + retry_impl (Callable): The constructor that creates a retry decorator + that will be applied to the method based on method configs. Returns: Mapping[str, MethodConfig]: A mapping of RPC method names to their @@ -151,7 +153,7 @@ def parse_method_configs(interface_config): if retry_params_name is not None: retry_params = retry_params_map[retry_params_name] retry_ = _retry_from_retry_config( - retry_params, retry_codes_map[method_params["retry_codes_name"]] + retry_params, retry_codes_map[method_params["retry_codes_name"]], retry_impl ) timeout_ = _timeout_from_retry_config(retry_params) diff --git a/google/api_core/gapic_v1/config_async.py b/google/api_core/gapic_v1/config_async.py new file mode 100644 index 00000000..00e5e240 --- /dev/null +++ b/google/api_core/gapic_v1/config_async.py @@ -0,0 +1,42 @@ +# 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 loading gapic configuration data. + +The Google API generator creates supplementary configuration for each RPC +method to tell the client library how to deal with retries and timeouts. +""" + +from google.api_core import retry_async +from google.api_core.gapic_v1 import config +from google.api_core.gapic_v1.config import MethodConfig # noqa: F401 + + +def parse_method_configs(interface_config): + """Creates default retry and timeout objects for each method in a gapic + interface config with AsyncIO semantics. + + Args: + interface_config (Mapping): The interface config section of the full + gapic library config. For example, If the full configuration has + an interface named ``google.example.v1.ExampleService`` you would + pass in just that interface's configuration, for example + ``gapic_config['interfaces']['google.example.v1.ExampleService']``. + + Returns: + Mapping[str, MethodConfig]: A mapping of RPC method names to their + configuration. + """ + return config.parse_method_configs( + interface_config, + retry_impl=retry_async.AsyncRetry) diff --git a/google/api_core/gapic_v1/method.py b/google/api_core/gapic_v1/method.py index 49982c03..4c46bbb3 100644 --- a/google/api_core/gapic_v1/method.py +++ b/google/api_core/gapic_v1/method.py @@ -86,7 +86,7 @@ def _determine_timeout(default_timeout, specified_timeout, retry): return specified_timeout -class _GapicCallable(object): +class GapicCallable(object): """Callable that applies retry, timeout, and metadata logic. Args: @@ -236,7 +236,7 @@ def get_topic(name, timeout=None): user_agent_metadata = None return general_helpers.wraps(func)( - _GapicCallable( + GapicCallable( func, default_retry, default_timeout, metadata=user_agent_metadata ) ) diff --git a/google/api_core/gapic_v1/method_async.py b/google/api_core/gapic_v1/method_async.py new file mode 100644 index 00000000..e081b8e7 --- /dev/null +++ b/google/api_core/gapic_v1/method_async.py @@ -0,0 +1,48 @@ +# 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 wrapping gRPC methods with common functionality. + +This is used by gapic clients to provide common error mapping, retry, timeout, +pagination, and long-running operations to gRPC methods. +""" + +from google.api_core import general_helpers, grpc_helpers_async +from google.api_core.gapic_v1 import client_info +from google.api_core.gapic_v1.method import (DEFAULT, # noqa: F401 + GapicCallable, + USE_DEFAULT_METADATA) + + +def wrap_method( + func, + default_retry=None, + default_timeout=None, + client_info=client_info.DEFAULT_CLIENT_INFO, +): + """Wrap an async RPC method with common behavior. + + Returns: + Callable: A new callable that takes optional ``retry`` and ``timeout`` + arguments and applies the common error mapping, retry, timeout, + and metadata behavior to the low-level RPC method. + """ + func = grpc_helpers_async.wrap_errors(func) + + if client_info is not None: + user_agent_metadata = [client_info.to_grpc_metadata()] + else: + user_agent_metadata = None + + return general_helpers.wraps(func)(GapicCallable( + func, default_retry, default_timeout, metadata=user_agent_metadata)) diff --git a/google/api_core/grpc_helpers.py b/google/api_core/grpc_helpers.py index c47b09fd..fde6c337 100644 --- a/google/api_core/grpc_helpers.py +++ b/google/api_core/grpc_helpers.py @@ -170,13 +170,10 @@ def wrap_errors(callable_): return _wrap_unary_errors(callable_) -def create_channel( - target, credentials=None, scopes=None, ssl_credentials=None, **kwargs -): - """Create a secure channel with credentials. +def _create_composite_credentials(credentials=None, scopes=None, ssl_credentials=None): + """Create the composite credentials for secure channels. 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`. @@ -185,11 +182,9 @@ def create_channel( are passed to :func:`google.auth.default`. ssl_credentials (grpc.ChannelCredentials): Optional SSL channel credentials. This can be used to specify different certificates. - kwargs: Additional key-word args passed to - :func:`grpc_gcp.secure_channel` or :func:`grpc.secure_channel`. Returns: - grpc.Channel: The created channel. + grpc.ChannelCredentials: The composed channel credentials object. """ if credentials is None: credentials, _ = google.auth.default(scopes=scopes) @@ -212,10 +207,34 @@ def create_channel( ssl_credentials = grpc.ssl_channel_credentials() # Combine the ssl credentials and the authorization credentials. - composite_credentials = grpc.composite_channel_credentials( + return grpc.composite_channel_credentials( ssl_credentials, google_auth_credentials ) + +def create_channel(target, credentials=None, scopes=None, ssl_credentials=None, **kwargs): + """Create a 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. + kwargs: Additional key-word args passed to + :func:`grpc_gcp.secure_channel` or :func:`grpc.secure_channel`. + + Returns: + grpc.Channel: The created channel. + """ + composite_credentials = _create_composite_credentials( + credentials, scopes, ssl_credentials + ) + if HAS_GRPC_GCP: # If grpc_gcp module is available use grpc_gcp.secure_channel, # otherwise, use grpc.secure_channel to create grpc channel. diff --git a/google/api_core/grpc_helpers_async.py b/google/api_core/grpc_helpers_async.py new file mode 100644 index 00000000..883d3936 --- /dev/null +++ b/google/api_core/grpc_helpers_async.py @@ -0,0 +1,273 @@ +# 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.6+. + +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 + +import grpc +from grpc.experimental import aio + +from google.api_core import exceptions, grpc_helpers + + +# TODO(lidiz) Support gRPC GCP wrapper +HAS_GRPC_GCP = False + +# 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(_WrappedCall): + + def __await__(self): + 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(_WrappedCall): + + def __init__(self): + self._wrapped_async_generator = None + + async def read(self): + 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): + 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): + 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, aio.UnaryUnaryCall): + """Wrapped UnaryUnaryCall to map exceptions.""" + + +class _WrappedUnaryStreamCall(_WrappedStreamResponseMixin, aio.UnaryStreamCall): + """Wrapped UnaryStreamCall to map exceptions.""" + + +class _WrappedStreamUnaryCall(_WrappedUnaryResponseMixin, _WrappedStreamRequestMixin, aio.StreamUnaryCall): + """Wrapped StreamUnaryCall to map exceptions.""" + + +class _WrappedStreamStreamCall(_WrappedStreamRequestMixin, _WrappedStreamResponseMixin, aio.StreamStreamCall): + """Wrapped StreamStreamCall to map exceptions.""" + + +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. enable_asyncio + (Optional[bool]): Optional flag to indicate whether to use AsyncIO + API of gRPC Python or not. + + 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, **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. + enable_asyncio: Optional flag to indicate whether to use AsyncIO + version of gRPC Channel or not. + kwargs: Additional key-word args passed to :func:`aio.secure_channel`. + + Returns: + aio.Channel: The created channel. + """ + composite_credentials = grpc_helpers._create_composite_credentials( + credentials, scopes, ssl_credentials + ) + + return aio.secure_channel(target, composite_credentials, **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 diff --git a/google/api_core/operation_async.py b/google/api_core/operation_async.py new file mode 100644 index 00000000..1199122e --- /dev/null +++ b/google/api_core/operation_async.py @@ -0,0 +1,217 @@ +# Copyright 2016 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. + +"""Futures for long-running operations returned from Google Cloud APIs. + +These futures can be used to synchronously wait for the result of a +long-running operation using :meth:`Operation.result`: + + +.. code-block:: python + + operation = my_api_client.long_running_method() + result = operation.result() + +Or asynchronously using callbacks and :meth:`Operation.add_done_callback`: + +.. code-block:: python + + operation = my_api_client.long_running_method() + + def my_callback(future): + result = future.result() + + operation.add_done_callback(my_callback) + +""" + +import functools +import threading + +from google.api_core import exceptions +from google.api_core import protobuf_helpers +from google.api_core.future import async_future +from google.longrunning import operations_pb2 +from google.rpc import code_pb2 + + +class AsyncOperation(async_future.AsyncFuture): + """A Future for interacting with a Google API Long-Running Operation. + + Args: + operation (google.longrunning.operations_pb2.Operation): The + initial operation. + refresh (Callable[[], ~.api_core.operation.Operation]): A callable that + returns the latest state of the operation. + cancel (Callable[[], None]): A callable that tries to cancel + the operation. + result_type (func:`type`): The protobuf type for the operation's + result. + metadata_type (func:`type`): The protobuf type for the operation's + metadata. + retry (google.api_core.retry.Retry): The retry configuration used + when polling. This can be used to control how often :meth:`done` + is polled. Regardless of the retry's ``deadline``, it will be + overridden by the ``timeout`` argument to :meth:`result`. + """ + + def __init__( + self, + operation, + refresh, + cancel, + result_type, + metadata_type=None, + retry=async_future.DEFAULT_RETRY, + ): + super().__init__(retry=retry) + self._operation = operation + self._refresh = refresh + self._cancel = cancel + self._result_type = result_type + self._metadata_type = metadata_type + self._completion_lock = threading.Lock() + # Invoke this in case the operation came back already complete. + self._set_result_from_operation() + + @property + def operation(self): + """google.longrunning.Operation: The current long-running operation.""" + return self._operation + + @property + def metadata(self): + """google.protobuf.Message: the current operation metadata.""" + if not self._operation.HasField("metadata"): + return None + + return protobuf_helpers.from_any_pb( + self._metadata_type, self._operation.metadata + ) + + @classmethod + def deserialize(cls, payload): + """Deserialize a ``google.longrunning.Operation`` protocol buffer. + + Args: + payload (bytes): A serialized operation protocol buffer. + + Returns: + ~.operations_pb2.Operation: An Operation protobuf object. + """ + return operations_pb2.Operation.FromString(payload) + + def _set_result_from_operation(self): + """Set the result or exception from the operation if it is complete.""" + # This must be done in a lock to prevent the async_future thread + # and main thread from both executing the completion logic + # at the same time. + with self._completion_lock: + # If the operation isn't complete or if the result has already been + # set, do not call set_result/set_exception again. + # Note: self._result_set is set to True in set_result and + # set_exception, in case those methods are invoked directly. + if not self._operation.done or self._future.done(): + return + + if self._operation.HasField("response"): + response = protobuf_helpers.from_any_pb( + self._result_type, self._operation.response + ) + self.set_result(response) + elif self._operation.HasField("error"): + exception = exceptions.GoogleAPICallError( + self._operation.error.message, + errors=(self._operation.error,), + response=self._operation, + ) + self.set_exception(exception) + else: + exception = exceptions.GoogleAPICallError( + "Unexpected state: Long-running operation had neither " + "response nor error set." + ) + self.set_exception(exception) + + async def _refresh_and_update(self, retry=async_future.DEFAULT_RETRY): + """Refresh the operation and update the result if needed. + + Args: + retry (google.api_core.retry.Retry): (Optional) How to retry the RPC. + """ + # If the currently cached operation is done, no need to make another + # RPC as it will not change once done. + if not self._operation.done: + self._operation = await self._refresh(retry=retry) + self._set_result_from_operation() + + async def done(self, retry=async_future.DEFAULT_RETRY): + """Checks to see if the operation is complete. + + Args: + retry (google.api_core.retry.Retry): (Optional) How to retry the RPC. + + Returns: + bool: True if the operation is complete, False otherwise. + """ + await self._refresh_and_update(retry) + return self._operation.done + + async def cancel(self): + """Attempt to cancel the operation. + + Returns: + bool: True if the cancel RPC was made, False if the operation is + already complete. + """ + result = await self.done() + if result: + return False + else: + await self._cancel() + return True + + async def cancelled(self): + """True if the operation was cancelled.""" + await self._refresh_and_update() + return ( + self._operation.HasField("error") + and self._operation.error.code == code_pb2.CANCELLED + ) + + +def from_gapic(operation, operations_client, result_type, **kwargs): + """Create an operation future from a gapic client. + + This interacts with the long-running operations `service`_ (specific + to a given API) via a gapic client. + + .. _service: https://github.com/googleapis/googleapis/blob/\ + 050400df0fdb16f63b63e9dee53819044bffc857/\ + google/longrunning/operations.proto#L38 + + Args: + operation (google.longrunning.operations_pb2.Operation): The operation. + operations_client (google.api_core.operations_v1.OperationsClient): + The operations client. + result_type (:func:`type`): The protobuf result type. + kwargs: Keyword args passed into the :class:`Operation` constructor. + + Returns: + ~.api_core.operation.Operation: The operation future to track the given + operation. + """ + refresh = functools.partial(operations_client.get_operation, operation.name) + cancel = functools.partial(operations_client.cancel_operation, operation.name) + return AsyncOperation(operation, refresh, cancel, result_type, **kwargs) diff --git a/google/api_core/operations_v1/__init__.py b/google/api_core/operations_v1/__init__.py index f0549561..0f6884f9 100644 --- a/google/api_core/operations_v1/__init__.py +++ b/google/api_core/operations_v1/__init__.py @@ -14,6 +14,12 @@ """Package for interacting with the google.longrunning.operations meta-API.""" +import sys + from google.api_core.operations_v1.operations_client import OperationsClient -__all__ = ["OperationsClient"] +if sys.version_info[0] >= 3 and sys.version_info[1] >= 6: + from google.api_core.operations_v1.operations_async_client import OperationsAsyncClient # noqa: F401 + __all__ = ["OperationsClient", "OperationsAsyncClient"] +else: + __all__ = ["OperationsClient"] # pragma: NO COVER diff --git a/google/api_core/operations_v1/operations_async_client.py b/google/api_core/operations_v1/operations_async_client.py new file mode 100644 index 00000000..039bec1b --- /dev/null +++ b/google/api_core/operations_v1/operations_async_client.py @@ -0,0 +1,274 @@ +# 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. + +"""An async client for the google.longrunning.operations meta-API. + +.. _Google API Style Guide: + https://cloud.google.com/apis/design/design_pattern + s#long_running_operations +.. _google/longrunning/operations.proto: + https://github.com/googleapis/googleapis/blob/master/google/longrunning + /operations.proto +""" + +import functools + +from google.api_core import gapic_v1, page_iterator_async +from google.api_core.operations_v1 import operations_client_config +from google.longrunning import operations_pb2 + + +class OperationsAsyncClient: + """Async client for interacting with long-running operations. + + Args: + channel (aio.Channel): The gRPC AsyncIO channel associated with the + service that implements the ``google.longrunning.operations`` + interface. + client_config (dict): + A dictionary of call options for each method. If not specified + the default configuration is used. + """ + + def __init__(self, channel, client_config=operations_client_config.config): + # Create the gRPC client stub with gRPC AsyncIO channel. + self.operations_stub = operations_pb2.OperationsStub(channel) + + # Create all wrapped methods using the interface configuration. + # The interface config contains all of the default settings for retry + # and timeout for each RPC method. + interfaces = client_config["interfaces"] + interface_config = interfaces["google.longrunning.Operations"] + method_configs = gapic_v1.config_async.parse_method_configs(interface_config) + + self._get_operation = gapic_v1.method_async.wrap_method( + self.operations_stub.GetOperation, + default_retry=method_configs["GetOperation"].retry, + default_timeout=method_configs["GetOperation"].timeout, + ) + + self._list_operations = gapic_v1.method_async.wrap_method( + self.operations_stub.ListOperations, + default_retry=method_configs["ListOperations"].retry, + default_timeout=method_configs["ListOperations"].timeout, + ) + + self._cancel_operation = gapic_v1.method_async.wrap_method( + self.operations_stub.CancelOperation, + default_retry=method_configs["CancelOperation"].retry, + default_timeout=method_configs["CancelOperation"].timeout, + ) + + self._delete_operation = gapic_v1.method_async.wrap_method( + self.operations_stub.DeleteOperation, + default_retry=method_configs["DeleteOperation"].retry, + default_timeout=method_configs["DeleteOperation"].timeout, + ) + + async def get_operation( + self, name, retry=gapic_v1.method_async.DEFAULT, timeout=gapic_v1.method_async.DEFAULT + ): + """Gets the latest state of a long-running operation. + + Clients can use this method to poll the operation result at intervals + as recommended by the API service. + + Example: + >>> from google.api_core import operations_v1 + >>> api = operations_v1.OperationsClient() + >>> name = '' + >>> response = await api.get_operation(name) + + Args: + name (str): The name of the operation resource. + retry (google.api_core.retry.Retry): The retry strategy to use + when invoking the RPC. If unspecified, the default retry from + the client configuration will be used. If ``None``, then this + method will not retry the RPC at all. + timeout (float): The amount of time in seconds to wait for the RPC + to complete. Note that if ``retry`` is used, this timeout + applies to each individual attempt and the overall time it + takes for this method to complete may be longer. If + unspecified, the the default timeout in the client + configuration is used. If ``None``, then the RPC method will + not time out. + + Returns: + google.longrunning.operations_pb2.Operation: The state of the + operation. + + Raises: + google.api_core.exceptions.GoogleAPICallError: If an error occurred + while invoking the RPC, the appropriate ``GoogleAPICallError`` + subclass will be raised. + """ + request = operations_pb2.GetOperationRequest(name=name) + return await self._get_operation(request, retry=retry, timeout=timeout) + + async def list_operations( + self, + name, + filter_, + retry=gapic_v1.method_async.DEFAULT, + timeout=gapic_v1.method_async.DEFAULT, + ): + """ + Lists operations that match the specified filter in the request. + + Example: + >>> from google.api_core import operations_v1 + >>> api = operations_v1.OperationsClient() + >>> name = '' + >>> + >>> # Iterate over all results + >>> for operation in await api.list_operations(name): + >>> # process operation + >>> pass + >>> + >>> # Or iterate over results one page at a time + >>> iter = await api.list_operations(name) + >>> for page in iter.pages: + >>> for operation in page: + >>> # process operation + >>> pass + + Args: + name (str): The name of the operation collection. + filter_ (str): The standard list filter. + retry (google.api_core.retry.Retry): The retry strategy to use + when invoking the RPC. If unspecified, the default retry from + the client configuration will be used. If ``None``, then this + method will not retry the RPC at all. + timeout (float): The amount of time in seconds to wait for the RPC + to complete. Note that if ``retry`` is used, this timeout + applies to each individual attempt and the overall time it + takes for this method to complete may be longer. If + unspecified, the the default timeout in the client + configuration is used. If ``None``, then the RPC method will + not time out. + + Returns: + google.api_core.page_iterator.Iterator: An iterator that yields + :class:`google.longrunning.operations_pb2.Operation` instances. + + Raises: + google.api_core.exceptions.MethodNotImplemented: If the server + does not support this method. Services are not required to + implement this method. + google.api_core.exceptions.GoogleAPICallError: If an error occurred + while invoking the RPC, the appropriate ``GoogleAPICallError`` + subclass will be raised. + """ + # Create the request object. + request = operations_pb2.ListOperationsRequest(name=name, filter=filter_) + + # Create the method used to fetch pages + method = functools.partial(self._list_operations, retry=retry, timeout=timeout) + + iterator = page_iterator_async.AsyncGRPCIterator( + client=None, + method=method, + request=request, + items_field="operations", + request_token_field="page_token", + response_token_field="next_page_token", + ) + + return iterator + + async def cancel_operation( + self, name, retry=gapic_v1.method_async.DEFAULT, timeout=gapic_v1.method_async.DEFAULT + ): + """Starts asynchronous cancellation on a long-running operation. + + The server makes a best effort to cancel the operation, but success is + not guaranteed. Clients can use :meth:`get_operation` or service- + specific methods to check whether the cancellation succeeded or whether + the operation completed despite cancellation. On successful + cancellation, the operation is not deleted; instead, it becomes an + operation with an ``Operation.error`` value with a + ``google.rpc.Status.code`` of ``1``, corresponding to + ``Code.CANCELLED``. + + Example: + >>> from google.api_core import operations_v1 + >>> api = operations_v1.OperationsClient() + >>> name = '' + >>> api.cancel_operation(name) + + Args: + name (str): The name of the operation resource to be cancelled. + retry (google.api_core.retry.Retry): The retry strategy to use + when invoking the RPC. If unspecified, the default retry from + the client configuration will be used. If ``None``, then this + method will not retry the RPC at all. + timeout (float): The amount of time in seconds to wait for the RPC + to complete. Note that if ``retry`` is used, this timeout + applies to each individual attempt and the overall time it + takes for this method to complete may be longer. If + unspecified, the the default timeout in the client + configuration is used. If ``None``, then the RPC method will + not time out. + + Raises: + google.api_core.exceptions.MethodNotImplemented: If the server + does not support this method. Services are not required to + implement this method. + google.api_core.exceptions.GoogleAPICallError: If an error occurred + while invoking the RPC, the appropriate ``GoogleAPICallError`` + subclass will be raised. + """ + # Create the request object. + request = operations_pb2.CancelOperationRequest(name=name) + await self._cancel_operation(request, retry=retry, timeout=timeout) + + async def delete_operation( + self, name, retry=gapic_v1.method_async.DEFAULT, timeout=gapic_v1.method_async.DEFAULT + ): + """Deletes a long-running operation. + + This method indicates that the client is no longer interested in the + operation result. It does not cancel the operation. + + Example: + >>> from google.api_core import operations_v1 + >>> api = operations_v1.OperationsClient() + >>> name = '' + >>> api.delete_operation(name) + + Args: + name (str): The name of the operation resource to be deleted. + retry (google.api_core.retry.Retry): The retry strategy to use + when invoking the RPC. If unspecified, the default retry from + the client configuration will be used. If ``None``, then this + method will not retry the RPC at all. + timeout (float): The amount of time in seconds to wait for the RPC + to complete. Note that if ``retry`` is used, this timeout + applies to each individual attempt and the overall time it + takes for this method to complete may be longer. If + unspecified, the the default timeout in the client + configuration is used. If ``None``, then the RPC method will + not time out. + + Raises: + google.api_core.exceptions.MethodNotImplemented: If the server + does not support this method. Services are not required to + implement this method. + google.api_core.exceptions.GoogleAPICallError: If an error occurred + while invoking the RPC, the appropriate ``GoogleAPICallError`` + subclass will be raised. + """ + # Create the request object. + request = operations_pb2.DeleteOperationRequest(name=name) + await self._delete_operation(request, retry=retry, timeout=timeout) diff --git a/google/api_core/page_iterator_async.py b/google/api_core/page_iterator_async.py new file mode 100644 index 00000000..d9ee9d8c --- /dev/null +++ b/google/api_core/page_iterator_async.py @@ -0,0 +1,278 @@ +# 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 iterators for paging through paged API methods. + +These iterators simplify the process of paging through API responses +where the request takes a page token and the response is a list of results with +a token for the next page. See `list pagination`_ in the Google API Style Guide +for more details. + +.. _list pagination: + https://cloud.google.com/apis/design/design_patterns#list_pagination + +API clients that have methods that follow the list pagination pattern can +return an :class:`.AsyncIterator`: + + >>> results_iterator = await client.list_resources() + +Or you can walk your way through items and call off the search early if +you find what you're looking for (resulting in possibly fewer requests):: + + >>> async for resource in results_iterator: + ... print(resource.name) + ... if not resource.is_valid: + ... break + +At any point, you may check the number of items consumed by referencing the +``num_results`` property of the iterator:: + + >>> async for my_item in results_iterator: + ... if results_iterator.num_results >= 10: + ... break + +When iterating, not every new item will send a request to the server. +To iterate based on each page of items (where a page corresponds to +a request):: + + >>> async for page in results_iterator.pages: + ... print('=' * 20) + ... print(' Page number: {:d}'.format(iterator.page_number)) + ... print(' Items in page: {:d}'.format(page.num_items)) + ... print(' First item: {!r}'.format(next(page))) + ... print('Items remaining: {:d}'.format(page.remaining)) + ... print('Next page token: {}'.format(iterator.next_page_token)) + ==================== + Page number: 1 + Items in page: 1 + First item: + Items remaining: 0 + Next page token: eav1OzQB0OM8rLdGXOEsyQWSG + ==================== + Page number: 2 + Items in page: 19 + First item: + Items remaining: 18 + Next page token: None +""" + +import abc + +from google.api_core.page_iterator import Page + + +def _item_to_value_identity(iterator, item): + """An item to value transformer that returns the item un-changed.""" + # pylint: disable=unused-argument + # We are conforming to the interface defined by Iterator. + return item + + +class AsyncIterator(abc.ABC): + """A generic class for iterating through API list responses. + + Args: + client(google.cloud.client.Client): The API client. + item_to_value (Callable[google.api_core.page_iterator_async.AsyncIterator, Any]): + Callable to convert an item from the type in the raw API response + into the native object. Will be called with the iterator and a + single item. + page_token (str): A token identifying a page in a result set to start + fetching results from. + max_results (int): The maximum number of results to fetch. + """ + + def __init__( + self, + client, + item_to_value=_item_to_value_identity, + page_token=None, + max_results=None, + ): + self._started = False + self.client = client + """Optional[Any]: The client that created this iterator.""" + self.item_to_value = item_to_value + """Callable[Iterator, Any]: Callable to convert an item from the type + in the raw API response into the native object. Will be called with + the iterator and a + single item. + """ + self.max_results = max_results + """int: The maximum number of results to fetch.""" + + # The attributes below will change over the life of the iterator. + self.page_number = 0 + """int: The current page of results.""" + self.next_page_token = page_token + """str: The token for the next page of results. If this is set before + the iterator starts, it effectively offsets the iterator to a + specific starting point.""" + self.num_results = 0 + """int: The total number of results fetched so far.""" + + @property + def pages(self): + """Iterator of pages in the response. + + returns: + types.GeneratorType[google.api_core.page_iterator.Page]: A + generator of page instances. + + raises: + ValueError: If the iterator has already been started. + """ + if self._started: + raise ValueError("Iterator has already started", self) + self._started = True + return self._page_aiter(increment=True) + + async def _items_aiter(self): + """Iterator for each item returned.""" + async for page in self._page_aiter(increment=False): + for item in page: + self.num_results += 1 + yield item + + def __aiter__(self): + """Iterator for each item returned. + + Returns: + types.GeneratorType[Any]: A generator of items from the API. + + Raises: + ValueError: If the iterator has already been started. + """ + if self._started: + raise ValueError("Iterator has already started", self) + self._started = True + return self._items_aiter() + + async def _page_aiter(self, increment): + """Generator of pages of API responses. + + Args: + increment (bool): Flag indicating if the total number of results + should be incremented on each page. This is useful since a page + iterator will want to increment by results per page while an + items iterator will want to increment per item. + + Yields: + Page: each page of items from the API. + """ + page = await self._next_page() + while page is not None: + self.page_number += 1 + if increment: + self.num_results += page.num_items + yield page + page = await self._next_page() + + @abc.abstractmethod + async def _next_page(self): + """Get the next page in the iterator. + + This does nothing and is intended to be over-ridden by subclasses + to return the next :class:`Page`. + + Raises: + NotImplementedError: Always, this method is abstract. + """ + raise NotImplementedError + + +class AsyncGRPCIterator(AsyncIterator): + """A generic class for iterating through gRPC list responses. + + .. note:: The class does not take a ``page_token`` argument because it can + just be specified in the ``request``. + + Args: + client (google.cloud.client.Client): The API client. This unused by + this class, but kept to satisfy the :class:`Iterator` interface. + method (Callable[protobuf.Message]): A bound gRPC method that should + take a single message for the request. + request (protobuf.Message): The request message. + items_field (str): The field in the response message that has the + items for the page. + item_to_value (Callable[GRPCIterator, Any]): Callable to convert an + item from the type in the JSON response into a native object. Will + be called with the iterator and a single item. + request_token_field (str): The field in the request message used to + specify the page token. + response_token_field (str): The field in the response message that has + the token for the next page. + max_results (int): The maximum number of results to fetch. + + .. autoattribute:: pages + """ + + _DEFAULT_REQUEST_TOKEN_FIELD = "page_token" + _DEFAULT_RESPONSE_TOKEN_FIELD = "next_page_token" + + def __init__( + self, + client, + method, + request, + items_field, + item_to_value=_item_to_value_identity, + request_token_field=_DEFAULT_REQUEST_TOKEN_FIELD, + response_token_field=_DEFAULT_RESPONSE_TOKEN_FIELD, + max_results=None, + ): + super().__init__(client, item_to_value, max_results=max_results) + self._method = method + self._request = request + self._items_field = items_field + self._request_token_field = request_token_field + self._response_token_field = response_token_field + + async def _next_page(self): + """Get the next page in the iterator. + + Returns: + Page: The next page in the iterator or :data:`None` if + there are no pages left. + """ + if not self._has_next_page(): + return None + + if self.next_page_token is not None: + setattr(self._request, self._request_token_field, self.next_page_token) + + response = await self._method(self._request) + + self.next_page_token = getattr(response, self._response_token_field) + items = getattr(response, self._items_field) + page = Page(self, items, self.item_to_value, raw_page=response) + + return page + + def _has_next_page(self): + """Determines whether or not there are more pages with results. + + Returns: + bool: Whether the iterator has more pages. + """ + if self.page_number == 0: + return True + + if self.max_results is not None: + if self.num_results >= self.max_results: + return False + + # Note: intentionally a falsy check instead of a None check. The RPC + # can return an empty string indicating no more pages. + return True if self.next_page_token else False diff --git a/google/api_core/retry_async.py b/google/api_core/retry_async.py new file mode 100644 index 00000000..270f224b --- /dev/null +++ b/google/api_core/retry_async.py @@ -0,0 +1,287 @@ +# Copyright 2017 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. + +"""Helpers for retrying coroutine functions with exponential back-off. + +The :class:`AsyncRetry` decorator shares most functionality and behavior with +:class:`Retry`, but supports coroutine functions. Please refer to description +of :class:`Retry` for more details. + +By default, this decorator will retry transient +API errors (see :func:`if_transient_error`). For example: + +.. code-block:: python + + @retry.AsyncRetry() + async def call_flaky_rpc(): + return await client.flaky_rpc() + + # Will retry flaky_rpc() if it raises transient API errors. + result = await call_flaky_rpc() + +You can pass a custom predicate to retry on different exceptions, such as +waiting for an eventually consistent item to be available: + +.. code-block:: python + + @retry.AsyncRetry(predicate=if_exception_type(exceptions.NotFound)) + async def check_if_exists(): + return await client.does_thing_exist() + + is_available = await check_if_exists() + +Some client library methods apply retry automatically. These methods can accept +a ``retry`` parameter that allows you to configure the behavior: + +.. code-block:: python + + my_retry = retry.AsyncRetry(deadline=60) + result = await client.some_method(retry=my_retry) + +""" + +import asyncio +import datetime +import functools +import logging + +from google.api_core import datetime_helpers, exceptions +from google.api_core.retry import (exponential_sleep_generator, # noqa: F401 + if_exception_type, if_transient_error) + +_LOGGER = logging.getLogger(__name__) +_DEFAULT_INITIAL_DELAY = 1.0 # seconds +_DEFAULT_MAXIMUM_DELAY = 60.0 # seconds +_DEFAULT_DELAY_MULTIPLIER = 2.0 +_DEFAULT_DEADLINE = 60.0 * 2.0 # seconds + + +async def retry_target(target, predicate, sleep_generator, deadline, on_error=None): + """Call a function and retry if it fails. + + This is the lowest-level retry helper. Generally, you'll use the + higher-level retry helper :class:`Retry`. + + Args: + target(Callable): The function to call and retry. This must be a + nullary function - apply arguments with `functools.partial`. + predicate (Callable[Exception]): A callable used to determine if an + exception raised by the target should be considered retryable. + It should return True to retry or False otherwise. + sleep_generator (Iterable[float]): An infinite iterator that determines + how long to sleep between retries. + deadline (float): How long to keep retrying the target. The last sleep + period is shortened as necessary, so that the last retry runs at + ``deadline`` (and not considerably beyond it). + on_error (Callable[Exception]): A function to call while processing a + retryable exception. Any error raised by this function will *not* + be caught. + + Returns: + Any: the return value of the target function. + + Raises: + google.api_core.RetryError: If the deadline is exceeded while retrying. + ValueError: If the sleep generator stops yielding values. + Exception: If the target raises a method that isn't retryable. + """ + if deadline is not None: + deadline_datetime = datetime_helpers.utcnow() + datetime.timedelta( + seconds=deadline + ) + else: + deadline_datetime = None + + last_exc = None + + for sleep in sleep_generator: + try: + if deadline_datetime is None: + return await target() + else: + return await asyncio.wait_for( + target(), + timeout=(deadline_datetime - datetime_helpers.utcnow()).total_seconds() + ) + # pylint: disable=broad-except + # This function explicitly must deal with broad exceptions. + except Exception as exc: + if not predicate(exc) and not isinstance(exc, asyncio.TimeoutError): + raise + last_exc = exc + if on_error is not None: + on_error(exc) + + now = datetime_helpers.utcnow() + + if deadline_datetime is not None: + if deadline_datetime <= now: + raise exceptions.RetryError( + "Deadline of {:.1f}s exceeded while calling {}".format( + deadline, target + ), + last_exc, + ) from last_exc + else: + time_to_deadline = (deadline_datetime - now).total_seconds() + sleep = min(time_to_deadline, sleep) + + _LOGGER.debug( + "Retrying due to {}, sleeping {:.1f}s ...".format(last_exc, sleep) + ) + await asyncio.sleep(sleep) + + raise ValueError("Sleep generator stopped yielding sleep values.") + + +class AsyncRetry(object): + """Exponential retry decorator for async functions. + + This class is a decorator used to add exponential back-off retry behavior + to an RPC call. + + Although the default behavior is to retry transient API errors, a + different predicate can be provided to retry other exceptions. + + Args: + predicate (Callable[Exception]): A callable that should return ``True`` + if the given exception is retryable. + initial (float): The minimum a,out of time to delay in seconds. This + must be greater than 0. + maximum (float): The maximum amout of time to delay in seconds. + multiplier (float): The multiplier applied to the delay. + deadline (float): How long to keep retrying in seconds. The last sleep + period is shortened as necessary, so that the last retry runs at + ``deadline`` (and not considerably beyond it). + """ + + def __init__( + self, + predicate=if_transient_error, + initial=_DEFAULT_INITIAL_DELAY, + maximum=_DEFAULT_MAXIMUM_DELAY, + multiplier=_DEFAULT_DELAY_MULTIPLIER, + deadline=_DEFAULT_DEADLINE, + on_error=None, + ): + self._predicate = predicate + self._initial = initial + self._multiplier = multiplier + self._maximum = maximum + self._deadline = deadline + self._on_error = on_error + + def __call__(self, func, on_error=None): + """Wrap a callable with retry behavior. + + Args: + func (Callable): The callable to add retry behavior to. + on_error (Callable[Exception]): A function to call while processing + a retryable exception. Any error raised by this function will + *not* be caught. + + Returns: + Callable: A callable that will invoke ``func`` with retry + behavior. + """ + if self._on_error is not None: + on_error = self._on_error + + @functools.wraps(func) + async def retry_wrapped_func(*args, **kwargs): + """A wrapper that calls target function with retry.""" + target = functools.partial(func, *args, **kwargs) + sleep_generator = exponential_sleep_generator( + self._initial, self._maximum, multiplier=self._multiplier + ) + return await retry_target( + target, + self._predicate, + sleep_generator, + self._deadline, + on_error=on_error, + ) + + return retry_wrapped_func + + def with_deadline(self, deadline): + """Return a copy of this retry with the given deadline. + + Args: + deadline (float): How long to keep retrying. + + Returns: + Retry: A new retry instance with the given deadline. + """ + return AsyncRetry( + predicate=self._predicate, + initial=self._initial, + maximum=self._maximum, + multiplier=self._multiplier, + deadline=deadline, + on_error=self._on_error, + ) + + def with_predicate(self, predicate): + """Return a copy of this retry with the given predicate. + + Args: + predicate (Callable[Exception]): A callable that should return + ``True`` if the given exception is retryable. + + Returns: + Retry: A new retry instance with the given predicate. + """ + return AsyncRetry( + predicate=predicate, + initial=self._initial, + maximum=self._maximum, + multiplier=self._multiplier, + deadline=self._deadline, + on_error=self._on_error, + ) + + def with_delay(self, initial=None, maximum=None, multiplier=None): + """Return a copy of this retry with the given delay options. + + Args: + initial (float): The minimum amout of time to delay. This must + be greater than 0. + maximum (float): The maximum amout of time to delay. + multiplier (float): The multiplier applied to the delay. + + Returns: + Retry: A new retry instance with the given predicate. + """ + return AsyncRetry( + predicate=self._predicate, + initial=initial if initial is not None else self._initial, + maximum=maximum if maximum is not None else self._maximum, + multiplier=multiplier if maximum is not None else self._multiplier, + deadline=self._deadline, + on_error=self._on_error, + ) + + def __str__(self): + return ( + "".format( + self._predicate, + self._initial, + self._maximum, + self._multiplier, + self._deadline, + self._on_error, + ) + ) diff --git a/noxfile.py b/noxfile.py index dfb12575..5eadbc8e 100644 --- a/noxfile.py +++ b/noxfile.py @@ -15,10 +15,23 @@ from __future__ import absolute_import import os import shutil +import sys # https://github.com/google/importlab/issues/25 import nox # pytype: disable=import-error +_MINIMAL_ASYNCIO_SUPPORT_PYTHON_VERSION = [3, 6] + + +def _greater_or_equal_than_36(version_string): + tokens = version_string.split('.') + for i, token in enumerate(tokens): + try: + tokens[i] = int(token) + except ValueError: + pass + return tokens >= [3, 6] + def default(session): """Default unit test session. @@ -32,19 +45,49 @@ def default(session): session.install("mock", "pytest", "pytest-cov", "grpcio >= 1.0.2") session.install("-e", ".") - # Run py.test against the unit tests. - session.run( - "py.test", - "--quiet", - "--cov=google.api_core", - "--cov=tests.unit", - "--cov-append", - "--cov-config=.coveragerc", - "--cov-report=", - "--cov-fail-under=0", - os.path.join("tests", "unit"), - *session.posargs - ) + # Inject AsyncIO content, if version >= 3.6. + if _greater_or_equal_than_36(session.python): + session.install( + "asyncmock", + "git+https://github.com/pytest-dev/pytest-asyncio.git") + # TODO(lidiz) Remove the daily build install once new API got released. + session.install( + "--pre", "--upgrade", "--force-reinstall", + "--extra-index-url", "https://packages.grpc.io/archive/2020/04/3dca4a321326dfbf6e3656b6d5fc29cf5f4b6f95-76408568-daeb-49cd-b869-be2fd4066e8f/python", + "grpcio") + + session.run( + "python", + "-m", + "py.test", + "--quiet", + "--cov=google.api_core", + "--cov=tests.unit", + "--cov=tests.asyncio", + "--cov-append", + "--cov-config=.coveragerc", + "--cov-report=", + "--cov-fail-under=0", + os.path.join("tests", "unit"), + os.path.join("tests", "asyncio"), + *session.posargs + ) + else: + # Run py.test against the unit tests. + session.run( + "python", + "-m", + "py.test", + "--quiet", + "--cov=google.api_core", + "--cov=tests.unit", + "--cov-append", + "--cov-config=.coveragerc", + "--cov-report=", + "--cov-fail-under=0", + os.path.join("tests", "unit"), + *session.posargs + ) @nox.session(python=["2.7", "3.5", "3.6", "3.7", "3.8"]) diff --git a/tests/asyncio/__init__.py b/tests/asyncio/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/asyncio/future/__init__.py b/tests/asyncio/future/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/asyncio/future/test_async_future.py b/tests/asyncio/future/test_async_future.py new file mode 100644 index 00000000..3322cb05 --- /dev/null +++ b/tests/asyncio/future/test_async_future.py @@ -0,0 +1,229 @@ +# Copyright 2017, 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. + +import asyncio + +import mock +import pytest + +from google.api_core import exceptions +from google.api_core.future import async_future + + +class AsyncFuture(async_future.AsyncFuture): + async def done(self): + return False + + async def cancel(self): + return True + + async def cancelled(self): + return False + + async def running(self): + return True + + +@pytest.mark.asyncio +async def test_polling_future_constructor(): + future = AsyncFuture() + assert not await future.done() + assert not await future.cancelled() + assert await future.running() + assert await future.cancel() + + +@pytest.mark.asyncio +async def test_set_result(): + future = AsyncFuture() + callback = mock.Mock() + + future.set_result(1) + + assert await future.result() == 1 + callback_called = asyncio.Event() + + def callback(unused_future): + callback_called.set() + + future.add_done_callback(callback) + await callback_called.wait() + + +@pytest.mark.asyncio +async def test_set_exception(): + future = AsyncFuture() + exception = ValueError("meep") + + future.set_exception(exception) + + assert await future.exception() == exception + with pytest.raises(ValueError): + await future.result() + + callback_called = asyncio.Event() + + def callback(unused_future): + callback_called.set() + + future.add_done_callback(callback) + await callback_called.wait() + + +@pytest.mark.asyncio +async def test_invoke_callback_exception(): + future = AsyncFuture() + future.set_result(42) + + # This should not raise, despite the callback causing an exception. + callback_called = asyncio.Event() + + def callback(unused_future): + callback_called.set() + raise ValueError() + + future.add_done_callback(callback) + await callback_called.wait() + + +class AsyncFutureWithPoll(AsyncFuture): + def __init__(self): + super().__init__() + self.poll_count = 0 + self.event = asyncio.Event() + + async def done(self): + self.poll_count += 1 + await self.event.wait() + self.set_result(42) + return True + + +@pytest.mark.asyncio +async def test_result_with_polling(): + future = AsyncFutureWithPoll() + + future.event.set() + result = await future.result() + + assert result == 42 + assert future.poll_count == 1 + # Repeated calls should not cause additional polling + assert await future.result() == result + assert future.poll_count == 1 + + +class AsyncFutureTimeout(AsyncFutureWithPoll): + + async def done(self): + await asyncio.sleep(0.2) + return False + + +@pytest.mark.asyncio +async def test_result_timeout(): + future = AsyncFutureTimeout() + with pytest.raises(asyncio.TimeoutError): + await future.result(timeout=0.2) + + +@pytest.mark.asyncio +async def test_exception_timeout(): + future = AsyncFutureTimeout() + with pytest.raises(asyncio.TimeoutError): + await future.exception(timeout=0.2) + + +@pytest.mark.asyncio +async def test_result_timeout_with_retry(): + future = AsyncFutureTimeout() + with pytest.raises(asyncio.TimeoutError): + await future.exception(timeout=0.4) + + +class AsyncFutureTransient(AsyncFutureWithPoll): + def __init__(self, errors): + super().__init__() + self._errors = errors + + async def done(self): + if self._errors: + error, self._errors = self._errors[0], self._errors[1:] + raise error("testing") + self.poll_count += 1 + self.set_result(42) + return True + + +@mock.patch("asyncio.sleep", autospec=True) +@pytest.mark.asyncio +async def test_result_transient_error(unused_sleep): + future = AsyncFutureTransient( + ( + exceptions.TooManyRequests, + exceptions.InternalServerError, + exceptions.BadGateway, + ) + ) + result = await future.result() + assert result == 42 + assert future.poll_count == 1 + # Repeated calls should not cause additional polling + assert await future.result() == result + assert future.poll_count == 1 + + +@pytest.mark.asyncio +async def test_callback_concurrency(): + future = AsyncFutureWithPoll() + + callback_called = asyncio.Event() + + def callback(unused_future): + callback_called.set() + + future.add_done_callback(callback) + + # Give the thread a second to poll + await asyncio.sleep(1) + assert future.poll_count == 1 + + future.event.set() + await callback_called.wait() + + +@pytest.mark.asyncio +async def test_double_callback_concurrency(): + future = AsyncFutureWithPoll() + + callback_called = asyncio.Event() + + def callback(unused_future): + callback_called.set() + + callback_called2 = asyncio.Event() + + def callback2(unused_future): + callback_called2.set() + + future.add_done_callback(callback) + future.add_done_callback(callback2) + + # Give the thread a second to poll + await asyncio.sleep(1) + future.event.set() + + assert future.poll_count == 1 + await callback_called.wait() + await callback_called2.wait() diff --git a/tests/asyncio/gapic/__init__.py b/tests/asyncio/gapic/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/asyncio/gapic/test_config_async.py b/tests/asyncio/gapic/test_config_async.py new file mode 100644 index 00000000..a0db0e4f --- /dev/null +++ b/tests/asyncio/gapic/test_config_async.py @@ -0,0 +1,87 @@ +# Copyright 2017 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. + +from google.api_core import exceptions +from google.api_core.gapic_v1 import config_async + + +INTERFACE_CONFIG = { + "retry_codes": { + "idempotent": ["DEADLINE_EXCEEDED", "UNAVAILABLE"], + "other": ["FAILED_PRECONDITION"], + "non_idempotent": [], + }, + "retry_params": { + "default": { + "initial_retry_delay_millis": 1000, + "retry_delay_multiplier": 2.5, + "max_retry_delay_millis": 120000, + "initial_rpc_timeout_millis": 120000, + "rpc_timeout_multiplier": 1.0, + "max_rpc_timeout_millis": 120000, + "total_timeout_millis": 600000, + }, + "other": { + "initial_retry_delay_millis": 1000, + "retry_delay_multiplier": 1, + "max_retry_delay_millis": 1000, + "initial_rpc_timeout_millis": 1000, + "rpc_timeout_multiplier": 1, + "max_rpc_timeout_millis": 1000, + "total_timeout_millis": 1000, + }, + }, + "methods": { + "AnnotateVideo": { + "timeout_millis": 60000, + "retry_codes_name": "idempotent", + "retry_params_name": "default", + }, + "Other": { + "timeout_millis": 60000, + "retry_codes_name": "other", + "retry_params_name": "other", + }, + "Plain": {"timeout_millis": 30000}, + }, +} + + +def test_create_method_configs(): + method_configs = config_async.parse_method_configs(INTERFACE_CONFIG) + + retry, timeout = method_configs["AnnotateVideo"] + assert retry._predicate(exceptions.DeadlineExceeded(None)) + assert retry._predicate(exceptions.ServiceUnavailable(None)) + assert retry._initial == 1.0 + assert retry._multiplier == 2.5 + assert retry._maximum == 120.0 + assert retry._deadline == 600.0 + assert timeout._initial == 120.0 + assert timeout._multiplier == 1.0 + assert timeout._maximum == 120.0 + + retry, timeout = method_configs["Other"] + assert retry._predicate(exceptions.FailedPrecondition(None)) + assert retry._initial == 1.0 + assert retry._multiplier == 1.0 + assert retry._maximum == 1.0 + assert retry._deadline == 1.0 + assert timeout._initial == 1.0 + assert timeout._multiplier == 1.0 + assert timeout._maximum == 1.0 + + retry, timeout = method_configs["Plain"] + assert retry is None + assert timeout._timeout == 30.0 diff --git a/tests/asyncio/gapic/test_method_async.py b/tests/asyncio/gapic/test_method_async.py new file mode 100644 index 00000000..7318362b --- /dev/null +++ b/tests/asyncio/gapic/test_method_async.py @@ -0,0 +1,243 @@ +# Copyright 2017 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. + +import datetime + +from grpc.experimental import aio +import mock +import pytest + +from google.api_core import (exceptions, gapic_v1, grpc_helpers_async, + retry_async, timeout) + + +def _utcnow_monotonic(): + current_time = datetime.datetime.min + delta = datetime.timedelta(seconds=0.5) + while True: + yield current_time + current_time += delta + + +@pytest.mark.asyncio +async def test_wrap_method_basic(): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall(42) + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, return_value=fake_call) + + wrapped_method = gapic_v1.method_async.wrap_method(method) + + result = await wrapped_method(1, 2, meep="moop") + + assert result == 42 + method.assert_called_once_with(1, 2, meep="moop", metadata=mock.ANY) + + # Check that the default client info was specified in the metadata. + metadata = method.call_args[1]["metadata"] + assert len(metadata) == 1 + client_info = gapic_v1.client_info.DEFAULT_CLIENT_INFO + user_agent_metadata = client_info.to_grpc_metadata() + assert user_agent_metadata in metadata + + +@pytest.mark.asyncio +async def test_wrap_method_with_no_client_info(): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall() + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, return_value=fake_call) + + wrapped_method = gapic_v1.method_async.wrap_method( + method, client_info=None + ) + + await wrapped_method(1, 2, meep="moop") + + method.assert_called_once_with(1, 2, meep="moop") + + +@pytest.mark.asyncio +async def test_wrap_method_with_custom_client_info(): + client_info = gapic_v1.client_info.ClientInfo( + python_version=1, + grpc_version=2, + api_core_version=3, + gapic_version=4, + client_library_version=5, + ) + fake_call = grpc_helpers_async.FakeUnaryUnaryCall() + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, return_value=fake_call) + + wrapped_method = gapic_v1.method_async.wrap_method( + method, client_info=client_info + ) + + await wrapped_method(1, 2, meep="moop") + + method.assert_called_once_with(1, 2, meep="moop", metadata=mock.ANY) + + # Check that the custom client info was specified in the metadata. + metadata = method.call_args[1]["metadata"] + assert client_info.to_grpc_metadata() in metadata + + +@pytest.mark.asyncio +async def test_invoke_wrapped_method_with_metadata(): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall() + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, return_value=fake_call) + + wrapped_method = gapic_v1.method_async.wrap_method(method) + + await wrapped_method(mock.sentinel.request, metadata=[("a", "b")]) + + method.assert_called_once_with(mock.sentinel.request, metadata=mock.ANY) + metadata = method.call_args[1]["metadata"] + # Metadata should have two items: the client info metadata and our custom + # metadata. + assert len(metadata) == 2 + assert ("a", "b") in metadata + + +@pytest.mark.asyncio +async def test_invoke_wrapped_method_with_metadata_as_none(): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall() + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, return_value=fake_call) + + wrapped_method = gapic_v1.method_async.wrap_method(method) + + await wrapped_method(mock.sentinel.request, metadata=None) + + method.assert_called_once_with(mock.sentinel.request, metadata=mock.ANY) + metadata = method.call_args[1]["metadata"] + # Metadata should have just one items: the client info metadata. + assert len(metadata) == 1 + + +@mock.patch("asyncio.sleep") +@pytest.mark.asyncio +async def test_wrap_method_with_default_retry_and_timeout(unused_sleep): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall(42) + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, side_effect=[ + exceptions.InternalServerError(None), + fake_call, + ]) + + default_retry = retry_async.AsyncRetry() + default_timeout = timeout.ConstantTimeout(60) + wrapped_method = gapic_v1.method_async.wrap_method( + method, default_retry, default_timeout + ) + + result = await wrapped_method() + + assert result == 42 + assert method.call_count == 2 + method.assert_called_with(timeout=60, metadata=mock.ANY) + + +@mock.patch("asyncio.sleep") +@pytest.mark.asyncio +async def test_wrap_method_with_default_retry_and_timeout_using_sentinel(unused_sleep): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall(42) + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, side_effect=[ + exceptions.InternalServerError(None), + fake_call, + ]) + + default_retry = retry_async.AsyncRetry() + default_timeout = timeout.ConstantTimeout(60) + wrapped_method = gapic_v1.method_async.wrap_method( + method, default_retry, default_timeout + ) + + result = await wrapped_method( + retry=gapic_v1.method_async.DEFAULT, + timeout=gapic_v1.method_async.DEFAULT, + ) + + assert result == 42 + assert method.call_count == 2 + method.assert_called_with(timeout=60, metadata=mock.ANY) + + +@mock.patch("asyncio.sleep") +@pytest.mark.asyncio +async def test_wrap_method_with_overriding_retry_and_timeout(unused_sleep): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall(42) + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, side_effect=[ + exceptions.NotFound(None), + fake_call, + ]) + + default_retry = retry_async.AsyncRetry() + default_timeout = timeout.ConstantTimeout(60) + wrapped_method = gapic_v1.method_async.wrap_method( + method, default_retry, default_timeout + ) + + result = await wrapped_method( + retry=retry_async.AsyncRetry(retry_async.if_exception_type(exceptions.NotFound)), + timeout=timeout.ConstantTimeout(22), + ) + + assert result == 42 + assert method.call_count == 2 + method.assert_called_with(timeout=22, metadata=mock.ANY) + + +@mock.patch("asyncio.sleep") +@mock.patch( + "google.api_core.datetime_helpers.utcnow", + side_effect=_utcnow_monotonic(), + autospec=True, +) +@pytest.mark.asyncio +async def test_wrap_method_with_overriding_retry_deadline(utcnow, unused_sleep): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall(42) + method = mock.Mock( + spec=aio.UnaryUnaryMultiCallable, + side_effect=([exceptions.InternalServerError(None)] * 4) + [fake_call]) + + default_retry = retry_async.AsyncRetry() + default_timeout = timeout.ExponentialTimeout(deadline=60) + wrapped_method = gapic_v1.method_async.wrap_method( + method, default_retry, default_timeout + ) + + # Overriding only the retry's deadline should also override the timeout's + # deadline. + result = await wrapped_method(retry=default_retry.with_deadline(30)) + + assert result == 42 + timeout_args = [call[1]["timeout"] for call in method.call_args_list] + assert timeout_args == [5.0, 10.0, 20.0, 26.0, 25.0] + assert utcnow.call_count == ( + 1 + + 1 # Compute wait_for timeout in retry_async + + 5 # First to set the deadline. + + 5 # One for each min(timeout, maximum, (DEADLINE - NOW).seconds) + ) + + +@pytest.mark.asyncio +async def test_wrap_method_with_overriding_timeout_as_a_number(): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall(42) + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, return_value=fake_call) + default_retry = retry_async.AsyncRetry() + default_timeout = timeout.ConstantTimeout(60) + wrapped_method = gapic_v1.method_async.wrap_method( + method, default_retry, default_timeout + ) + + result = await wrapped_method(timeout=22) + + assert result == 42 + method.assert_called_once_with(timeout=22, metadata=mock.ANY) diff --git a/tests/asyncio/operations_v1/__init__.py b/tests/asyncio/operations_v1/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/asyncio/operations_v1/test_operations_async_client.py b/tests/asyncio/operations_v1/test_operations_async_client.py new file mode 100644 index 00000000..0f9363ff --- /dev/null +++ b/tests/asyncio/operations_v1/test_operations_async_client.py @@ -0,0 +1,93 @@ +# Copyright 2017 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. + +from grpc.experimental import aio +import mock +import pytest + +from google.api_core import (grpc_helpers_async, operations_v1, + page_iterator_async) +from google.longrunning import operations_pb2 +from google.protobuf import empty_pb2 + + +def _mock_grpc_objects(response): + fake_call = grpc_helpers_async.FakeUnaryUnaryCall(response) + method = mock.Mock(spec=aio.UnaryUnaryMultiCallable, return_value=fake_call) + mocked_channel = mock.Mock() + mocked_channel.unary_unary = mock.Mock(return_value=method) + return mocked_channel, method, fake_call + + +@pytest.mark.asyncio +async def test_get_operation(): + mocked_channel, method, fake_call = _mock_grpc_objects( + operations_pb2.Operation(name="meep")) + client = operations_v1.OperationsAsyncClient(mocked_channel) + + response = await client.get_operation("name") + assert method.call_count == 1 + assert tuple(method.call_args_list[0])[0][0].name == "name" + assert response == fake_call.response + + +@pytest.mark.asyncio +async def test_list_operations(): + operations = [ + operations_pb2.Operation(name="1"), + operations_pb2.Operation(name="2"), + ] + list_response = operations_pb2.ListOperationsResponse(operations=operations) + + mocked_channel, method, fake_call = _mock_grpc_objects(list_response) + client = operations_v1.OperationsAsyncClient(mocked_channel) + + pager = await client.list_operations("name", "filter") + + assert isinstance(pager, page_iterator_async.AsyncIterator) + responses = [] + async for response in pager: + responses.append(response) + + assert responses == operations + + assert method.call_count == 1 + request = tuple(method.call_args_list[0])[0][0] + assert isinstance(request, operations_pb2.ListOperationsRequest) + assert request.name == "name" + assert request.filter == "filter" + + +@pytest.mark.asyncio +async def test_delete_operation(): + mocked_channel, method, fake_call = _mock_grpc_objects( + empty_pb2.Empty()) + client = operations_v1.OperationsAsyncClient(mocked_channel) + + await client.delete_operation("name") + + assert method.call_count == 1 + assert tuple(method.call_args_list[0])[0][0].name == "name" + + +@pytest.mark.asyncio +async def test_cancel_operation(): + mocked_channel, method, fake_call = _mock_grpc_objects( + empty_pb2.Empty()) + client = operations_v1.OperationsAsyncClient(mocked_channel) + + await client.cancel_operation("name") + + assert method.call_count == 1 + assert tuple(method.call_args_list[0])[0][0].name == "name" diff --git a/tests/asyncio/test_grpc_helpers_async.py b/tests/asyncio/test_grpc_helpers_async.py new file mode 100644 index 00000000..00539521 --- /dev/null +++ b/tests/asyncio/test_grpc_helpers_async.py @@ -0,0 +1,372 @@ +# Copyright 2017 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. + +import grpc +from grpc.experimental import aio +import mock +import pytest + +from google.api_core import exceptions +from google.api_core import grpc_helpers_async +import google.auth.credentials + + +class RpcErrorImpl(grpc.RpcError, grpc.Call): + def __init__(self, code): + super(RpcErrorImpl, self).__init__() + self._code = code + + def code(self): + return self._code + + def details(self): + return None + + +@pytest.mark.asyncio +async def test_wrap_unary_errors(): + grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT) + callable_ = mock.AsyncMock(spec=["__call__"], side_effect=grpc_error) + + wrapped_callable = grpc_helpers_async._wrap_unary_errors(callable_) + + with pytest.raises(exceptions.InvalidArgument) as exc_info: + await wrapped_callable(1, 2, three="four") + + callable_.assert_called_once_with(1, 2, three="four") + assert exc_info.value.response == grpc_error + + +@pytest.mark.asyncio +async def test_common_methods_in_wrapped_call(): + mock_call = mock.Mock(aio.UnaryUnaryCall, autospec=True) + wrapped_call = grpc_helpers_async._WrappedUnaryUnaryCall().with_call(mock_call) + + await wrapped_call.initial_metadata() + assert mock_call.initial_metadata.call_count == 1 + + await wrapped_call.trailing_metadata() + assert mock_call.trailing_metadata.call_count == 1 + + await wrapped_call.code() + assert mock_call.code.call_count == 1 + + await wrapped_call.details() + assert mock_call.details.call_count == 1 + + wrapped_call.cancelled() + assert mock_call.cancelled.call_count == 1 + + wrapped_call.done() + assert mock_call.done.call_count == 1 + + wrapped_call.time_remaining() + assert mock_call.time_remaining.call_count == 1 + + wrapped_call.cancel() + assert mock_call.cancel.call_count == 1 + + callback = mock.sentinel.callback + wrapped_call.add_done_callback(callback) + mock_call.add_done_callback.assert_called_once_with(callback) + + await wrapped_call.wait_for_connection() + assert mock_call.wait_for_connection.call_count == 1 + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_unary_stream(): + mock_call = mock.Mock(aio.UnaryStreamCall, autospec=True) + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + + await wrapped_callable(1, 2, three="four") + multicallable.assert_called_once_with(1, 2, three="four") + assert mock_call.wait_for_connection.call_count == 1 + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_stream_unary(): + mock_call = mock.Mock(aio.StreamUnaryCall, autospec=True) + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + + await wrapped_callable(1, 2, three="four") + multicallable.assert_called_once_with(1, 2, three="four") + assert mock_call.wait_for_connection.call_count == 1 + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_stream_stream(): + mock_call = mock.Mock(aio.StreamStreamCall, autospec=True) + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + + await wrapped_callable(1, 2, three="four") + multicallable.assert_called_once_with(1, 2, three="four") + assert mock_call.wait_for_connection.call_count == 1 + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_type_error(): + mock_call = mock.Mock() + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + + with pytest.raises(TypeError): + await wrapped_callable() + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_raised(): + grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT) + mock_call = mock.Mock(aio.StreamStreamCall, autospec=True) + mock_call.wait_for_connection = mock.AsyncMock(side_effect=[grpc_error]) + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + + with pytest.raises(exceptions.InvalidArgument): + await wrapped_callable() + assert mock_call.wait_for_connection.call_count == 1 + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_read(): + grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT) + + mock_call = mock.Mock(aio.StreamStreamCall, autospec=True) + mock_call.read = mock.AsyncMock(side_effect=grpc_error) + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + + wrapped_call = await wrapped_callable(1, 2, three="four") + multicallable.assert_called_once_with(1, 2, three="four") + assert mock_call.wait_for_connection.call_count == 1 + + with pytest.raises(exceptions.InvalidArgument) as exc_info: + await wrapped_call.read() + assert exc_info.value.response == grpc_error + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_aiter(): + grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT) + + mock_call = mock.Mock(aio.StreamStreamCall, autospec=True) + mocked_aiter = mock.Mock(spec=['__anext__']) + mocked_aiter.__anext__ = mock.AsyncMock(side_effect=[mock.sentinel.response, grpc_error]) + mock_call.__aiter__ = mock.Mock(return_value=mocked_aiter) + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + wrapped_call = await wrapped_callable() + + with pytest.raises(exceptions.InvalidArgument) as exc_info: + async for response in wrapped_call: + assert response == mock.sentinel.response + assert exc_info.value.response == grpc_error + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_aiter_non_rpc_error(): + non_grpc_error = TypeError('Not a gRPC error') + + mock_call = mock.Mock(aio.StreamStreamCall, autospec=True) + mocked_aiter = mock.Mock(spec=['__anext__']) + mocked_aiter.__anext__ = mock.AsyncMock(side_effect=[mock.sentinel.response, non_grpc_error]) + mock_call.__aiter__ = mock.Mock(return_value=mocked_aiter) + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + wrapped_call = await wrapped_callable() + + with pytest.raises(TypeError) as exc_info: + async for response in wrapped_call: + assert response == mock.sentinel.response + assert exc_info.value == non_grpc_error + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_aiter_called_multiple_times(): + mock_call = mock.Mock(aio.StreamStreamCall, autospec=True) + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + wrapped_call = await wrapped_callable() + + assert wrapped_call.__aiter__() == wrapped_call.__aiter__() + + +@pytest.mark.asyncio +async def test_wrap_stream_errors_write(): + grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT) + + mock_call = mock.Mock(aio.StreamStreamCall, autospec=True) + mock_call.write = mock.AsyncMock(side_effect=[None, grpc_error]) + mock_call.done_writing = mock.AsyncMock(side_effect=[None, grpc_error]) + multicallable = mock.Mock(return_value=mock_call) + + wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable) + + wrapped_call = await wrapped_callable() + + await wrapped_call.write(mock.sentinel.request) + with pytest.raises(exceptions.InvalidArgument) as exc_info: + await wrapped_call.write(mock.sentinel.request) + assert mock_call.write.call_count == 2 + assert exc_info.value.response == grpc_error + + await wrapped_call.done_writing() + with pytest.raises(exceptions.InvalidArgument) as exc_info: + await wrapped_call.done_writing() + assert mock_call.done_writing.call_count == 2 + assert exc_info.value.response == grpc_error + + +@mock.patch("google.api_core.grpc_helpers_async._wrap_unary_errors") +def test_wrap_errors_non_streaming(wrap_unary_errors): + callable_ = mock.create_autospec(aio.UnaryUnaryMultiCallable) + + result = grpc_helpers_async.wrap_errors(callable_) + + assert result == wrap_unary_errors.return_value + wrap_unary_errors.assert_called_once_with(callable_) + + +@mock.patch("google.api_core.grpc_helpers_async._wrap_stream_errors") +def test_wrap_errors_streaming(wrap_stream_errors): + callable_ = mock.create_autospec(aio.UnaryStreamMultiCallable) + + result = grpc_helpers_async.wrap_errors(callable_) + + assert result == wrap_stream_errors.return_value + wrap_stream_errors.assert_called_once_with(callable_) + + +@mock.patch("grpc.composite_channel_credentials") +@mock.patch( + "google.auth.default", + return_value=(mock.sentinel.credentials, mock.sentinel.projet), +) +@mock.patch("grpc.experimental.aio.secure_channel") +def test_create_channel_implicit(grpc_secure_channel, default, composite_creds_call): + target = "example.com:443" + composite_creds = composite_creds_call.return_value + + channel = grpc_helpers_async.create_channel(target) + + assert channel is grpc_secure_channel.return_value + default.assert_called_once_with(scopes=None) + grpc_secure_channel.assert_called_once_with(target, composite_creds) + + +@mock.patch("grpc.composite_channel_credentials") +@mock.patch( + "google.auth.default", + return_value=(mock.sentinel.credentials, mock.sentinel.projet), +) +@mock.patch("grpc.experimental.aio.secure_channel") +def test_create_channel_implicit_with_ssl_creds( + grpc_secure_channel, default, composite_creds_call +): + target = "example.com:443" + + ssl_creds = grpc.ssl_channel_credentials() + + grpc_helpers_async.create_channel(target, ssl_credentials=ssl_creds) + + default.assert_called_once_with(scopes=None) + composite_creds_call.assert_called_once_with(ssl_creds, mock.ANY) + composite_creds = composite_creds_call.return_value + grpc_secure_channel.assert_called_once_with(target, composite_creds) + + +@mock.patch("grpc.composite_channel_credentials") +@mock.patch( + "google.auth.default", + return_value=(mock.sentinel.credentials, mock.sentinel.projet), +) +@mock.patch("grpc.experimental.aio.secure_channel") +def test_create_channel_implicit_with_scopes( + grpc_secure_channel, default, composite_creds_call +): + target = "example.com:443" + composite_creds = composite_creds_call.return_value + + channel = grpc_helpers_async.create_channel(target, scopes=["one", "two"]) + + assert channel is grpc_secure_channel.return_value + default.assert_called_once_with(scopes=["one", "two"]) + grpc_secure_channel.assert_called_once_with(target, composite_creds) + + +@mock.patch("grpc.composite_channel_credentials") +@mock.patch("google.auth.credentials.with_scopes_if_required") +@mock.patch("grpc.experimental.aio.secure_channel") +def test_create_channel_explicit(grpc_secure_channel, auth_creds, composite_creds_call): + target = "example.com:443" + composite_creds = composite_creds_call.return_value + + channel = grpc_helpers_async.create_channel(target, credentials=mock.sentinel.credentials) + + auth_creds.assert_called_once_with(mock.sentinel.credentials, None) + assert channel is grpc_secure_channel.return_value + grpc_secure_channel.assert_called_once_with(target, composite_creds) + + +@mock.patch("grpc.composite_channel_credentials") +@mock.patch("grpc.experimental.aio.secure_channel") +def test_create_channel_explicit_scoped(grpc_secure_channel, composite_creds_call): + target = "example.com:443" + scopes = ["1", "2"] + composite_creds = composite_creds_call.return_value + + credentials = mock.create_autospec(google.auth.credentials.Scoped, instance=True) + credentials.requires_scopes = True + + channel = grpc_helpers_async.create_channel( + target, credentials=credentials, scopes=scopes + ) + + credentials.with_scopes.assert_called_once_with(scopes) + assert channel is grpc_secure_channel.return_value + grpc_secure_channel.assert_called_once_with(target, composite_creds) + + +@pytest.mark.skipif(grpc_helpers_async.HAS_GRPC_GCP, reason="grpc_gcp module not available") +@mock.patch("grpc.experimental.aio.secure_channel") +def test_create_channel_without_grpc_gcp(grpc_secure_channel): + target = "example.com:443" + scopes = ["test_scope"] + + credentials = mock.create_autospec(google.auth.credentials.Scoped, instance=True) + credentials.requires_scopes = True + + grpc_helpers_async.create_channel(target, credentials=credentials, scopes=scopes) + grpc_secure_channel.assert_called() + credentials.with_scopes.assert_called_once_with(scopes) + + +@pytest.mark.asyncio +async def test_fake_stream_unary_call(): + fake_call = grpc_helpers_async.FakeStreamUnaryCall() + await fake_call.wait_for_connection() + response = await fake_call + assert fake_call.response == response diff --git a/tests/asyncio/test_operation_async.py b/tests/asyncio/test_operation_async.py new file mode 100644 index 00000000..419749f3 --- /dev/null +++ b/tests/asyncio/test_operation_async.py @@ -0,0 +1,193 @@ +# Copyright 2017, 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. + + +import mock +import pytest + +from google.api_core import exceptions +from google.api_core import operation_async +from google.api_core import operations_v1 +from google.api_core import retry_async +from google.longrunning import operations_pb2 +from google.protobuf import struct_pb2 +from google.rpc import code_pb2 +from google.rpc import status_pb2 + +TEST_OPERATION_NAME = "test/operation" + + +def make_operation_proto( + name=TEST_OPERATION_NAME, metadata=None, response=None, error=None, **kwargs +): + operation_proto = operations_pb2.Operation(name=name, **kwargs) + + if metadata is not None: + operation_proto.metadata.Pack(metadata) + + if response is not None: + operation_proto.response.Pack(response) + + if error is not None: + operation_proto.error.CopyFrom(error) + + return operation_proto + + +def make_operation_future(client_operations_responses=None): + if client_operations_responses is None: + client_operations_responses = [make_operation_proto()] + + refresh = mock.AsyncMock(spec=["__call__"], side_effect=client_operations_responses) + refresh.responses = client_operations_responses + cancel = mock.AsyncMock(spec=["__call__"]) + operation_future = operation_async.AsyncOperation( + client_operations_responses[0], + refresh, + cancel, + result_type=struct_pb2.Struct, + metadata_type=struct_pb2.Struct, + ) + + return operation_future, refresh, cancel + + +@pytest.mark.asyncio +async def test_constructor(): + future, refresh, _ = make_operation_future() + + assert future.operation == refresh.responses[0] + assert future.operation.done is False + assert future.operation.name == TEST_OPERATION_NAME + assert future.metadata is None + assert await future.running() + + +def test_metadata(): + expected_metadata = struct_pb2.Struct() + future, _, _ = make_operation_future( + [make_operation_proto(metadata=expected_metadata)] + ) + + assert future.metadata == expected_metadata + + +@pytest.mark.asyncio +async def test_cancellation(): + responses = [ + make_operation_proto(), + # Second response indicates that the operation was cancelled. + make_operation_proto( + done=True, error=status_pb2.Status(code=code_pb2.CANCELLED) + ), + ] + future, _, cancel = make_operation_future(responses) + + assert await future.cancel() + assert await future.cancelled() + cancel.assert_called_once_with() + + # Cancelling twice should have no effect. + assert not await future.cancel() + cancel.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_result(): + expected_result = struct_pb2.Struct() + responses = [ + make_operation_proto(), + # Second operation response includes the result. + make_operation_proto(done=True, response=expected_result), + ] + future, _, _ = make_operation_future(responses) + + result = await future.result() + + assert result == expected_result + assert await future.done() + + +@pytest.mark.asyncio +async def test_done_w_retry(): + RETRY_PREDICATE = retry_async.if_exception_type(exceptions.TooManyRequests) + test_retry = retry_async.AsyncRetry(predicate=RETRY_PREDICATE) + + expected_result = struct_pb2.Struct() + responses = [ + make_operation_proto(), + # Second operation response includes the result. + make_operation_proto(done=True, response=expected_result), + ] + future, refresh, _ = make_operation_future(responses) + + await future.done(retry=test_retry) + refresh.assert_called_once_with(retry=test_retry) + + +@pytest.mark.asyncio +async def test_exception(): + expected_exception = status_pb2.Status(message="meep") + responses = [ + make_operation_proto(), + # Second operation response includes the error. + make_operation_proto(done=True, error=expected_exception), + ] + future, _, _ = make_operation_future(responses) + + exception = await future.exception() + + assert expected_exception.message in "{!r}".format(exception) + + +@mock.patch("asyncio.sleep", autospec=True) +@pytest.mark.asyncio +async def test_unexpected_result(unused_sleep): + responses = [ + make_operation_proto(), + # Second operation response is done, but has not error or response. + make_operation_proto(done=True), + ] + future, _, _ = make_operation_future(responses) + + exception = await future.exception() + + assert "Unexpected state" in "{!r}".format(exception) + + +def test_from_gapic(): + operation_proto = make_operation_proto(done=True) + operations_client = mock.create_autospec( + operations_v1.OperationsClient, instance=True + ) + + future = operation_async.from_gapic( + operation_proto, + operations_client, + struct_pb2.Struct, + metadata_type=struct_pb2.Struct, + ) + + assert future._result_type == struct_pb2.Struct + assert future._metadata_type == struct_pb2.Struct + assert future.operation.name == TEST_OPERATION_NAME + assert future.done + + +def test_deserialize(): + op = make_operation_proto(name="foobarbaz") + serialized = op.SerializeToString() + deserialized_op = operation_async.AsyncOperation.deserialize(serialized) + assert op.name == deserialized_op.name + assert type(op) is type(deserialized_op) diff --git a/tests/asyncio/test_page_iterator_async.py b/tests/asyncio/test_page_iterator_async.py new file mode 100644 index 00000000..42fac2a2 --- /dev/null +++ b/tests/asyncio/test_page_iterator_async.py @@ -0,0 +1,261 @@ +# Copyright 2015 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. + +import inspect + +import mock +import pytest + +from google.api_core import page_iterator_async + + +class PageAsyncIteratorImpl(page_iterator_async.AsyncIterator): + + async def _next_page(self): + return mock.create_autospec(page_iterator_async.Page, instance=True) + + +class TestAsyncIterator: + + def test_constructor(self): + client = mock.sentinel.client + item_to_value = mock.sentinel.item_to_value + token = "ab13nceor03" + max_results = 1337 + + iterator = PageAsyncIteratorImpl( + client, item_to_value, page_token=token, max_results=max_results + ) + + assert not iterator._started + assert iterator.client is client + assert iterator.item_to_value == item_to_value + assert iterator.max_results == max_results + # Changing attributes. + assert iterator.page_number == 0 + assert iterator.next_page_token == token + assert iterator.num_results == 0 + + def test_pages_property_starts(self): + iterator = PageAsyncIteratorImpl(None, None) + + assert not iterator._started + + assert inspect.isasyncgen(iterator.pages) + + assert iterator._started + + def test_pages_property_restart(self): + iterator = PageAsyncIteratorImpl(None, None) + + assert iterator.pages + + # Make sure we cannot restart. + with pytest.raises(ValueError): + assert iterator.pages + + @pytest.mark.asyncio + async def test__page_aiter_increment(self): + iterator = PageAsyncIteratorImpl(None, None) + page = page_iterator_async.Page( + iterator, ("item",), page_iterator_async._item_to_value_identity) + iterator._next_page = mock.AsyncMock(side_effect=[page, None]) + + assert iterator.num_results == 0 + + page_aiter = iterator._page_aiter(increment=True) + await page_aiter.__anext__() + + assert iterator.num_results == 1 + + @pytest.mark.asyncio + async def test__page_aiter_no_increment(self): + iterator = PageAsyncIteratorImpl(None, None) + + assert iterator.num_results == 0 + + page_aiter = iterator._page_aiter(increment=False) + await page_aiter.__anext__() + + # results should still be 0 after fetching a page. + assert iterator.num_results == 0 + + @pytest.mark.asyncio + async def test__items_aiter(self): + # Items to be returned. + item1 = 17 + item2 = 100 + item3 = 211 + + # Make pages from mock responses + parent = mock.sentinel.parent + page1 = page_iterator_async.Page( + parent, (item1, item2), page_iterator_async._item_to_value_identity) + page2 = page_iterator_async.Page( + parent, (item3,), page_iterator_async._item_to_value_identity) + + iterator = PageAsyncIteratorImpl(None, None) + iterator._next_page = mock.AsyncMock(side_effect=[page1, page2, None]) + + items_aiter = iterator._items_aiter() + + assert inspect.isasyncgen(items_aiter) + + # Consume items and check the state of the iterator. + assert iterator.num_results == 0 + assert await items_aiter.__anext__() == item1 + assert iterator.num_results == 1 + + assert await items_aiter.__anext__() == item2 + assert iterator.num_results == 2 + + assert await items_aiter.__anext__() == item3 + assert iterator.num_results == 3 + + with pytest.raises(StopAsyncIteration): + await items_aiter.__anext__() + + @pytest.mark.asyncio + async def test___aiter__(self): + async_iterator = PageAsyncIteratorImpl(None, None) + async_iterator._next_page = mock.AsyncMock(side_effect=[(1, 2), (3,), None]) + + assert not async_iterator._started + + result = [] + async for item in async_iterator: + result.append(item) + + assert result == [1, 2, 3] + assert async_iterator._started + + def test___aiter__restart(self): + iterator = PageAsyncIteratorImpl(None, None) + + iterator.__aiter__() + + # Make sure we cannot restart. + with pytest.raises(ValueError): + iterator.__aiter__() + + def test___aiter___restart_after_page(self): + iterator = PageAsyncIteratorImpl(None, None) + + assert iterator.pages + + # Make sure we cannot restart after starting the page iterator + with pytest.raises(ValueError): + iterator.__aiter__() + + +class TestAsyncGRPCIterator(object): + + def test_constructor(self): + client = mock.sentinel.client + items_field = "items" + iterator = page_iterator_async.AsyncGRPCIterator( + client, mock.sentinel.method, mock.sentinel.request, items_field + ) + + assert not iterator._started + assert iterator.client is client + assert iterator.max_results is None + assert iterator.item_to_value is page_iterator_async._item_to_value_identity + assert iterator._method == mock.sentinel.method + assert iterator._request == mock.sentinel.request + assert iterator._items_field == items_field + assert ( + iterator._request_token_field + == page_iterator_async.AsyncGRPCIterator._DEFAULT_REQUEST_TOKEN_FIELD + ) + assert ( + iterator._response_token_field + == page_iterator_async.AsyncGRPCIterator._DEFAULT_RESPONSE_TOKEN_FIELD + ) + # Changing attributes. + assert iterator.page_number == 0 + assert iterator.next_page_token is None + assert iterator.num_results == 0 + + def test_constructor_options(self): + client = mock.sentinel.client + items_field = "items" + request_field = "request" + response_field = "response" + iterator = page_iterator_async.AsyncGRPCIterator( + client, + mock.sentinel.method, + mock.sentinel.request, + items_field, + item_to_value=mock.sentinel.item_to_value, + request_token_field=request_field, + response_token_field=response_field, + max_results=42, + ) + + assert iterator.client is client + assert iterator.max_results == 42 + assert iterator.item_to_value is mock.sentinel.item_to_value + assert iterator._method == mock.sentinel.method + assert iterator._request == mock.sentinel.request + assert iterator._items_field == items_field + assert iterator._request_token_field == request_field + assert iterator._response_token_field == response_field + + @pytest.mark.asyncio + async def test_iterate(self): + request = mock.Mock(spec=["page_token"], page_token=None) + response1 = mock.Mock(items=["a", "b"], next_page_token="1") + response2 = mock.Mock(items=["c"], next_page_token="2") + response3 = mock.Mock(items=["d"], next_page_token="") + method = mock.AsyncMock(side_effect=[response1, response2, response3]) + iterator = page_iterator_async.AsyncGRPCIterator( + mock.sentinel.client, method, request, "items" + ) + + assert iterator.num_results == 0 + + items = [] + async for item in iterator: + items.append(item) + + assert items == ["a", "b", "c", "d"] + + method.assert_called_with(request) + assert method.call_count == 3 + assert request.page_token == "2" + + @pytest.mark.asyncio + async def test_iterate_with_max_results(self): + request = mock.Mock(spec=["page_token"], page_token=None) + response1 = mock.Mock(items=["a", "b"], next_page_token="1") + response2 = mock.Mock(items=["c"], next_page_token="2") + response3 = mock.Mock(items=["d"], next_page_token="") + method = mock.AsyncMock(side_effect=[response1, response2, response3]) + iterator = page_iterator_async.AsyncGRPCIterator( + mock.sentinel.client, method, request, "items", max_results=3 + ) + + assert iterator.num_results == 0 + + items = [] + async for item in iterator: + items.append(item) + + assert items == ["a", "b", "c"] + assert iterator.num_results == 3 + + method.assert_called_with(request) + assert method.call_count == 2 + assert request.page_token == "1" diff --git a/tests/asyncio/test_retry_async.py b/tests/asyncio/test_retry_async.py new file mode 100644 index 00000000..ab1a50c3 --- /dev/null +++ b/tests/asyncio/test_retry_async.py @@ -0,0 +1,397 @@ +# Copyright 2017 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. + +import datetime +import re + +import mock +import pytest + +from google.api_core import exceptions +from google.api_core import retry_async + + +@mock.patch("asyncio.sleep", autospec=True) +@mock.patch( + "google.api_core.datetime_helpers.utcnow", + return_value=datetime.datetime.min, + autospec=True, +) +@pytest.mark.asyncio +async def test_retry_target_success(utcnow, sleep): + predicate = retry_async.if_exception_type(ValueError) + call_count = [0] + + async def target(): + call_count[0] += 1 + if call_count[0] < 3: + raise ValueError() + return 42 + + result = await retry_async.retry_target(target, predicate, range(10), None) + + assert result == 42 + assert call_count[0] == 3 + sleep.assert_has_calls([mock.call(0), mock.call(1)]) + + +@mock.patch("asyncio.sleep", autospec=True) +@mock.patch( + "google.api_core.datetime_helpers.utcnow", + return_value=datetime.datetime.min, + autospec=True, +) +@pytest.mark.asyncio +async def test_retry_target_w_on_error(utcnow, sleep): + predicate = retry_async.if_exception_type(ValueError) + call_count = {"target": 0} + to_raise = ValueError() + + async def target(): + call_count["target"] += 1 + if call_count["target"] < 3: + raise to_raise + return 42 + + on_error = mock.Mock() + + result = await retry_async.retry_target(target, predicate, range(10), None, on_error=on_error) + + assert result == 42 + assert call_count["target"] == 3 + + on_error.assert_has_calls([mock.call(to_raise), mock.call(to_raise)]) + sleep.assert_has_calls([mock.call(0), mock.call(1)]) + + +@mock.patch("asyncio.sleep", autospec=True) +@mock.patch( + "google.api_core.datetime_helpers.utcnow", + return_value=datetime.datetime.min, + autospec=True, +) +@pytest.mark.asyncio +async def test_retry_target_non_retryable_error(utcnow, sleep): + predicate = retry_async.if_exception_type(ValueError) + exception = TypeError() + target = mock.Mock(side_effect=exception) + + with pytest.raises(TypeError) as exc_info: + await retry_async.retry_target(target, predicate, range(10), None) + + assert exc_info.value == exception + sleep.assert_not_called() + + +@mock.patch("asyncio.sleep", autospec=True) +@mock.patch("google.api_core.datetime_helpers.utcnow", autospec=True) +@pytest.mark.asyncio +async def test_retry_target_deadline_exceeded(utcnow, sleep): + predicate = retry_async.if_exception_type(ValueError) + exception = ValueError("meep") + target = mock.Mock(side_effect=exception) + # Setup the timeline so that the first call takes 5 seconds but the second + # call takes 6, which puts the retry over the deadline. + utcnow.side_effect = [ + # The first call to utcnow establishes the start of the timeline. + datetime.datetime.min, + datetime.datetime.min + datetime.timedelta(seconds=5), + datetime.datetime.min + datetime.timedelta(seconds=11), + ] + + with pytest.raises(exceptions.RetryError) as exc_info: + await retry_async.retry_target(target, predicate, range(10), deadline=10) + + assert exc_info.value.cause == exception + assert exc_info.match("Deadline of 10.0s exceeded") + assert exc_info.match("last exception: meep") + assert target.call_count == 2 + + +@pytest.mark.asyncio +async def test_retry_target_bad_sleep_generator(): + with pytest.raises(ValueError, match="Sleep generator"): + await retry_async.retry_target(mock.sentinel.target, mock.sentinel.predicate, [], None) + + +class TestAsyncRetry: + + def test_constructor_defaults(self): + retry_ = retry_async.AsyncRetry() + assert retry_._predicate == retry_async.if_transient_error + assert retry_._initial == 1 + assert retry_._maximum == 60 + assert retry_._multiplier == 2 + assert retry_._deadline == 120 + assert retry_._on_error is None + + def test_constructor_options(self): + _some_function = mock.Mock() + + retry_ = retry_async.AsyncRetry( + predicate=mock.sentinel.predicate, + initial=1, + maximum=2, + multiplier=3, + deadline=4, + on_error=_some_function, + ) + assert retry_._predicate == mock.sentinel.predicate + assert retry_._initial == 1 + assert retry_._maximum == 2 + assert retry_._multiplier == 3 + assert retry_._deadline == 4 + assert retry_._on_error is _some_function + + def test_with_deadline(self): + retry_ = retry_async.AsyncRetry( + predicate=mock.sentinel.predicate, + initial=1, + maximum=2, + multiplier=3, + deadline=4, + on_error=mock.sentinel.on_error, + ) + new_retry = retry_.with_deadline(42) + assert retry_ is not new_retry + assert new_retry._deadline == 42 + + # the rest of the attributes should remain the same + assert new_retry._predicate is retry_._predicate + assert new_retry._initial == retry_._initial + assert new_retry._maximum == retry_._maximum + assert new_retry._multiplier == retry_._multiplier + assert new_retry._on_error is retry_._on_error + + def test_with_predicate(self): + retry_ = retry_async.AsyncRetry( + predicate=mock.sentinel.predicate, + initial=1, + maximum=2, + multiplier=3, + deadline=4, + on_error=mock.sentinel.on_error, + ) + new_retry = retry_.with_predicate(mock.sentinel.predicate) + assert retry_ is not new_retry + assert new_retry._predicate == mock.sentinel.predicate + + # the rest of the attributes should remain the same + assert new_retry._deadline == retry_._deadline + assert new_retry._initial == retry_._initial + assert new_retry._maximum == retry_._maximum + assert new_retry._multiplier == retry_._multiplier + assert new_retry._on_error is retry_._on_error + + def test_with_delay_noop(self): + retry_ = retry_async.AsyncRetry( + predicate=mock.sentinel.predicate, + initial=1, + maximum=2, + multiplier=3, + deadline=4, + on_error=mock.sentinel.on_error, + ) + new_retry = retry_.with_delay() + assert retry_ is not new_retry + assert new_retry._initial == retry_._initial + assert new_retry._maximum == retry_._maximum + assert new_retry._multiplier == retry_._multiplier + + def test_with_delay(self): + retry_ = retry_async.AsyncRetry( + predicate=mock.sentinel.predicate, + initial=1, + maximum=2, + multiplier=3, + deadline=4, + on_error=mock.sentinel.on_error, + ) + new_retry = retry_.with_delay(initial=1, maximum=2, multiplier=3) + assert retry_ is not new_retry + assert new_retry._initial == 1 + assert new_retry._maximum == 2 + assert new_retry._multiplier == 3 + + # the rest of the attributes should remain the same + assert new_retry._deadline == retry_._deadline + assert new_retry._predicate is retry_._predicate + assert new_retry._on_error is retry_._on_error + + def test___str__(self): + def if_exception_type(exc): + return bool(exc) # pragma: NO COVER + + # Explicitly set all attributes as changed Retry defaults should not + # cause this test to start failing. + retry_ = retry_async.AsyncRetry( + predicate=if_exception_type, + initial=1.0, + maximum=60.0, + multiplier=2.0, + deadline=120.0, + on_error=None, + ) + assert re.match( + ( + r", " + r"initial=1.0, maximum=60.0, multiplier=2.0, deadline=120.0, " + r"on_error=None>" + ), + str(retry_), + ) + + @mock.patch("asyncio.sleep", autospec=True) + @pytest.mark.asyncio + async def test___call___and_execute_success(self, sleep): + retry_ = retry_async.AsyncRetry() + target = mock.AsyncMock(spec=["__call__"], return_value=42) + # __name__ is needed by functools.partial. + target.__name__ = "target" + + decorated = retry_(target) + target.assert_not_called() + + result = await decorated("meep") + + assert result == 42 + target.assert_called_once_with("meep") + sleep.assert_not_called() + + # Make uniform return half of its maximum, which is the calculated sleep time. + @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0) + @mock.patch("asyncio.sleep", autospec=True) + @pytest.mark.asyncio + async def test___call___and_execute_retry(self, sleep, uniform): + + on_error = mock.Mock(spec=["__call__"], side_effect=[None]) + retry_ = retry_async.AsyncRetry(predicate=retry_async.if_exception_type(ValueError)) + + target = mock.AsyncMock(spec=["__call__"], side_effect=[ValueError(), 42]) + # __name__ is needed by functools.partial. + target.__name__ = "target" + + decorated = retry_(target, on_error=on_error) + target.assert_not_called() + + result = await decorated("meep") + + assert result == 42 + assert target.call_count == 2 + target.assert_has_calls([mock.call("meep"), mock.call("meep")]) + sleep.assert_called_once_with(retry_._initial) + assert on_error.call_count == 1 + + # Make uniform return half of its maximum, which is the calculated sleep time. + @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0) + @mock.patch("asyncio.sleep", autospec=True) + @pytest.mark.asyncio + async def test___call___and_execute_retry_hitting_deadline(self, sleep, uniform): + + on_error = mock.Mock(spec=["__call__"], side_effect=[None] * 10) + retry_ = retry_async.AsyncRetry( + predicate=retry_async.if_exception_type(ValueError), + initial=1.0, + maximum=1024.0, + multiplier=2.0, + deadline=9.9, + ) + + utcnow = datetime.datetime.utcnow() + utcnow_patcher = mock.patch( + "google.api_core.datetime_helpers.utcnow", return_value=utcnow + ) + + target = mock.AsyncMock(spec=["__call__"], side_effect=[ValueError()] * 10) + # __name__ is needed by functools.partial. + target.__name__ = "target" + + decorated = retry_(target, on_error=on_error) + target.assert_not_called() + + with utcnow_patcher as patched_utcnow: + # Make sure that calls to fake asyncio.sleep() also advance the mocked + # time clock. + def increase_time(sleep_delay): + patched_utcnow.return_value += datetime.timedelta(seconds=sleep_delay) + sleep.side_effect = increase_time + + with pytest.raises(exceptions.RetryError): + await decorated("meep") + + assert target.call_count == 5 + target.assert_has_calls([mock.call("meep")] * 5) + assert on_error.call_count == 5 + + # check the delays + assert sleep.call_count == 4 # once between each successive target calls + last_wait = sleep.call_args.args[0] + total_wait = sum(call_args.args[0] for call_args in sleep.call_args_list) + + assert last_wait == 2.9 # and not 8.0, because the last delay was shortened + assert total_wait == 9.9 # the same as the deadline + + @mock.patch("asyncio.sleep", autospec=True) + @pytest.mark.asyncio + async def test___init___without_retry_executed(self, sleep): + _some_function = mock.Mock() + + retry_ = retry_async.AsyncRetry( + predicate=retry_async.if_exception_type(ValueError), on_error=_some_function + ) + # check the proper creation of the class + assert retry_._on_error is _some_function + + target = mock.AsyncMock(spec=["__call__"], side_effect=[42]) + # __name__ is needed by functools.partial. + target.__name__ = "target" + + wrapped = retry_(target) + + result = await wrapped("meep") + + assert result == 42 + target.assert_called_once_with("meep") + sleep.assert_not_called() + _some_function.assert_not_called() + + # Make uniform return half of its maximum, which is the calculated sleep time. + @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0) + @mock.patch("asyncio.sleep", autospec=True) + @pytest.mark.asyncio + async def test___init___when_retry_is_executed(self, sleep, uniform): + _some_function = mock.Mock() + + retry_ = retry_async.AsyncRetry( + predicate=retry_async.if_exception_type(ValueError), on_error=_some_function + ) + # check the proper creation of the class + assert retry_._on_error is _some_function + + target = mock.AsyncMock( + spec=["__call__"], side_effect=[ValueError(), ValueError(), 42] + ) + # __name__ is needed by functools.partial. + target.__name__ = "target" + + wrapped = retry_(target) + target.assert_not_called() + + result = await wrapped("meep") + + assert result == 42 + assert target.call_count == 3 + assert _some_function.call_count == 2 + target.assert_has_calls([mock.call("meep"), mock.call("meep")]) + sleep.assert_any_call(retry_._initial)