WIP: AsyncIO integration

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
         )

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)

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

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)
 

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)

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
         )
     )

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))