feat: retry and retry_async support streaming rpcs

google/api_core/retry.py

@@ -66,6 +66,7 @@ def check_if_exists():
 
 from google.api_core import datetime_helpers
 from google.api_core import exceptions
+from google.api_core.retry_streaming import RetryableGenerator
 from google.auth import exceptions as auth_exceptions
 
 _LOGGER = logging.getLogger(__name__)
@@ -154,7 +155,7 @@ def retry_target(
     higher-level retry helper :class:`Retry`.
 
     Args:
-        target(Callable): The function to call and retry. This must be a
+        target(Callable[None, Any]): 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.
@@ -301,6 +302,15 @@ class Retry(object):
         maximum (float): The maximum amount of time to delay in seconds.
         multiplier (float): The multiplier applied to the delay.
         timeout (float): How long to keep retrying, in seconds.
+        on_error (Callable[Exception]): A function to call while processing
+            a retryable exception. Any error raised by this function will
+            *not* be caught.
+        is_stream (bool): Indicates whether the input function
+            should be treated as an stream function (i.e. a Generator,
+            or function that returns an Iterable). If True, the iterable
+            will be wrapped with retry logic, and any failed outputs will
+            restart the stream. If False, only the input function call itself
+            will be retried. Defaults to False.
         deadline (float): DEPRECATED: use `timeout` instead. For backward
             compatibility, if specified it will override the ``timeout`` parameter.
     """
@@ -313,7 +323,8 @@ def __init__(
         multiplier=_DEFAULT_DELAY_MULTIPLIER,
         timeout=_DEFAULT_DEADLINE,
         on_error=None,
-        **kwargs
+        is_stream=False,
+        **kwargs,
     ):
         self._predicate = predicate
         self._initial = initial
@@ -322,6 +333,7 @@ def __init__(
         self._timeout = kwargs.get("deadline", timeout)
         self._deadline = self._timeout
         self._on_error = on_error
+        self._is_stream = is_stream
 
     def __call__(self, func, on_error=None):
         """Wrap a callable with retry behavior.
@@ -346,7 +358,8 @@ def retry_wrapped_func(*args, **kwargs):
             sleep_generator = exponential_sleep_generator(
                 self._initial, self._maximum, multiplier=self._multiplier
             )
-            return retry_target(
+            retry_func = RetryableGenerator if self._is_stream else retry_target
+            return retry_func(
                 target,
                 self._predicate,
                 sleep_generator,

google/api_core/retry_async.py

@@ -61,6 +61,7 @@ async def check_if_exists():
 from google.api_core.retry import exponential_sleep_generator
 from google.api_core.retry import if_exception_type  # noqa: F401
 from google.api_core.retry import if_transient_error
+from google.api_core.retry_streaming_async import AsyncRetryableGenerator
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -74,7 +75,7 @@ async def check_if_exists():
 async def retry_target(
     target, predicate, sleep_generator, timeout=None, on_error=None, **kwargs
 ):
-    """Call a function and retry if it fails.
+    """Await a coroutine and retry if it fails.
 
     This is the lowest-level retry helper. Generally, you'll use the
     higher-level retry helper :class:`Retry`.
@@ -156,7 +157,7 @@ async def retry_target(
 
 
 class AsyncRetry:
-    """Exponential retry decorator for async functions.
+    """Exponential retry decorator for async coroutines.
 
     This class is a decorator used to add exponential back-off retry behavior
     to an RPC call.
@@ -172,9 +173,17 @@ class AsyncRetry:
         maximum (float): The maximum amount of time to delay in seconds.
         multiplier (float): The multiplier applied to the delay.
         timeout (float): How long to keep retrying in seconds.
+            When ``is_stream``, only time spent waiting on the
+            target or sleeping between retries is counted towards the timeout.
         on_error (Callable[Exception]): A function to call while processing
             a retryable exception. Any error raised by this function will
             *not* be caught.
+        is_stream (bool): Indicates whether the input function
+            should be treated as an stream function (i.e. an AsyncGenerator,
+            or function or coroutine that returns an AsyncIterable).
+            If True, the iterable will be wrapped with retry logic, and any
+            failed outputs will restart the stream. If False, only the input
+            function call itself will be retried. Defaults to False.
         deadline (float): DEPRECATED use ``timeout`` instead. If set it will
         override ``timeout`` parameter.
     """
@@ -187,6 +196,7 @@ def __init__(
         multiplier=_DEFAULT_DELAY_MULTIPLIER,
         timeout=_DEFAULT_TIMEOUT,
         on_error=None,
+        is_stream=False,
         **kwargs
     ):
         self._predicate = predicate
@@ -196,12 +206,13 @@ def __init__(
         self._timeout = kwargs.get("deadline", timeout)
         self._deadline = self._timeout
         self._on_error = on_error
+        self._is_stream = is_stream
 
     def __call__(self, func, on_error=None):
         """Wrap a callable with retry behavior.
 
         Args:
-            func (Callable): The callable to add retry behavior to.
+            func (Callable): The callable or stream 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.
@@ -224,11 +235,29 @@ async def retry_wrapped_func(*args, **kwargs):
                 target,
                 self._predicate,
                 sleep_generator,
-                self._timeout,
+                timeout=self._timeout,
+                on_error=on_error,
+            )
+
+        @functools.wraps(func)
+        def retry_wrapped_stream(*args, **kwargs):
+            """A wrapper that iterates over target stream with retry."""
+            target = functools.partial(func, *args, **kwargs)
+            sleep_generator = exponential_sleep_generator(
+                self._initial, self._maximum, multiplier=self._multiplier
+            )
+            return AsyncRetryableGenerator(
+                target,
+                self._predicate,
+                sleep_generator,
+                timeout=self._timeout,
                 on_error=on_error,
             )
 
-        return retry_wrapped_func
+        if self._is_stream:
+            return retry_wrapped_stream
+        else:
+            return retry_wrapped_func
 
     def _replace(
         self,

google/api_core/retry_streaming.py

@@ -0,0 +1,217 @@
+# Copyright 2023 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 retries for streaming APIs."""
+
+from typing import Callable, Optional, Iterable, Iterator, Generator, TypeVar, Any, cast
+
+import datetime
+import logging
+import time
+
+from google.api_core import datetime_helpers
+from google.api_core import exceptions
+
+_LOGGER = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+class RetryableGenerator(Generator[T, Any, None]):
+    """
+    Helper class for retrying Iterator and Generator-based
+    streaming APIs.
+    """
+
+    def __init__(
+        self,
+        target: Callable[[], Iterable[T]],
+        predicate: Callable[[Exception], bool],
+        sleep_generator: Iterable[float],
+        timeout: Optional[float] = None,
+        on_error: Optional[Callable[[Exception], None]] = None,
+    ):
+        """
+        Args:
+            target: The function to call to produce iterables for each retry.
+                This must be a nullary function - apply arguments with
+                `functools.partial`.
+            predicate: 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: An infinite iterator that determines
+                how long to sleep between retries.
+            timeout: How long to keep retrying the target, in seconds.
+            on_error: A function to call while processing a
+                retryable exception.  Any error raised by this function will *not*
+                be caught.
+        """
+        self.target_fn = target
+        self.active_target: Iterator[T] = self.target_fn().__iter__()
+        self.predicate = predicate
+        self.sleep_generator = iter(sleep_generator)
+        self.on_error = on_error
+        self.timeout = timeout
+        if self.timeout is not None:
+            self.deadline = datetime_helpers.utcnow() + datetime.timedelta(
+                seconds=self.timeout
+            )
+        else:
+            self.deadline = None
+
+    def __iter__(self) -> Generator[T, Any, None]:
+        """
+        Implement the iterator protocol.
+        """
+        return self
+
+    def _handle_exception(self, exc) -> None:
+        """
+        When an exception is raised while iterating over the active_target,
+        check if it is retryable. If so, create a new active_target and
+        continue iterating. If not, raise the exception.
+        """
+        if not self.predicate(exc):
+            raise exc
+        else:
+            # run on_error callback if provided
+            if self.on_error:
+                self.on_error(exc)
+            try:
+                next_sleep = next(self.sleep_generator)
+            except StopIteration:
+                raise ValueError("Sleep generator stopped yielding sleep values")
+            # if deadline is exceeded, raise RetryError
+            if self.deadline is not None:
+                next_attempt = datetime_helpers.utcnow() + datetime.timedelta(
+                    seconds=next_sleep
+                )
+                self._check_timeout(next_attempt, exc)
+            # sleep before retrying
+            _LOGGER.debug(
+                "Retrying due to {}, sleeping {:.1f}s ...".format(exc, next_sleep)
+            )
+            time.sleep(next_sleep)
+            self.active_target = self.target_fn().__iter__()
+
+    def _check_timeout(
+        self, current_time: float, source_exception: Optional[Exception] = None
+    ) -> None:
+        """
+        Helper function to check if the timeout has been exceeded, and raise a RetryError if so.
+
+        Args:
+          - current_time: the timestamp to check against the deadline
+          - source_exception: the exception that triggered the timeout check, if any
+        Raises:
+          - RetryError if the deadline has been exceeded
+        """
+        if (
+            self.deadline is not None
+            and self.timeout is not None
+            and self.deadline < current_time
+        ):
+            raise exceptions.RetryError(
+                "Timeout of {:.1f}s exceeded".format(self.timeout),
+                source_exception,
+            ) from source_exception
+
+    def __next__(self) -> T:
+        """
+        Implement the iterator protocol.
+
+        Returns:
+          - the next value of the active_target iterator
+        """
+        # check for expired timeouts before attempting to iterate
+        self._check_timeout(datetime_helpers.utcnow())
+        try:
+            return next(self.active_target)
+        except Exception as exc:
+            self._handle_exception(exc)
+        # if retryable exception was handled, try again with new active_target
+        return self.__next__()
+
+    def close(self) -> None:
+        """
+        Close the active_target if supported. (e.g. target is a generator)
+
+        Raises:
+          - AttributeError if the active_target does not have a close() method
+        """
+        if getattr(self.active_target, "close", None):
+            casted_target = cast(Generator, self.active_target)
+            return casted_target.close()
+        else:
+            raise AttributeError(
+                "close() not implemented for {}".format(self.active_target)
+            )
+
+    def send(self, *args, **kwargs) -> T:
+        """
+        Call send on the active_target if supported. (e.g. target is a generator)
+
+        If an exception is raised, a retry may be attempted before returning
+        a result.
+
+        Args:
+          - *args: arguments to pass to the wrapped generator's send method
+          - **kwargs: keyword arguments to pass to the wrapped generator's send method
+        Returns:
+          - the next value of the active_target iterator after calling send
+        Raises:
+          - AttributeError if the active_target does not have a send() method
+        """
+        # check for expired timeouts before attempting to iterate
+        self._check_timeout(datetime_helpers.utcnow())
+        if getattr(self.active_target, "send", None):
+            casted_target = cast(Generator, self.active_target)
+            try:
+                return casted_target.send(*args, **kwargs)
+            except Exception as exc:
+                self._handle_exception(exc)
+            # if exception was retryable, use new target for return value
+            return self.__next__()
+        else:
+            raise AttributeError(
+                "send() not implemented for {}".format(self.active_target)
+            )
+
+    def throw(self, *args, **kwargs) -> T:
+        """
+        Call throw on the active_target if supported. (e.g. target is a generator)
+
+        If an exception is raised, a retry may be attempted before returning
+        a result.
+
+        Args:
+          - *args: arguments to pass to the wrapped generator's throw method
+          - **kwargs: keyword arguments to pass to the wrapped generator's throw method
+        Returns:
+          - the next vale of the active_target iterator after calling throw
+        Raises:
+          - AttributeError if the active_target does not have a throw() method
+        """
+        if getattr(self.active_target, "throw", None):
+            casted_target = cast(Generator, self.active_target)
+            try:
+                return casted_target.throw(*args, **kwargs)
+            except Exception as exc:
+                self._handle_exception(exc)
+            # if retryable exception was handled, return next from new active_target
+            return self.__next__()
+        else:
+            raise AttributeError(
+                "throw() not implemented for {}".format(self.active_target)
+            )