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 retry_target_stream
 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[[], 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,18 @@ 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.
+            To avoid duplicate values, retryable streams should typically be
+            wrapped in additional filter logic before use. For more details, see
+            ``google.api_core.retry_streaming.retry_target_stream``.
         deadline (float): DEPRECATED: use `timeout` instead. For backward
             compatibility, if specified it will override the ``timeout`` parameter.
     """
@@ -313,7 +326,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 +336,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 +361,8 @@ def retry_wrapped_func(*args, **kwargs):
             sleep_generator = exponential_sleep_generator(
                 self._initial, self._maximum, multiplier=self._multiplier
             )
-            return retry_target(
+            retry_func = retry_target_stream 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 retry_target_stream
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -74,13 +75,13 @@ 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`.
 
     Args:
-        target(Callable): The function to call and retry. This must be a
+        target(Callable[[], 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.
@@ -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.
@@ -175,6 +176,15 @@ class AsyncRetry:
         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.
+            To avoid duplicate values, retryable streams should typically be
+            wrapped in additional filter logic before use. For more details, see
+            ``google.api_core.retry_streaming_async.retry_target_stream``.
         deadline (float): DEPRECATED use ``timeout`` instead. If set it will
         override ``timeout`` parameter.
     """
@@ -187,6 +197,7 @@ def __init__(
         multiplier=_DEFAULT_DELAY_MULTIPLIER,
         timeout=_DEFAULT_TIMEOUT,
         on_error=None,
+        is_stream=False,
         **kwargs
     ):
         self._predicate = predicate
@@ -196,12 +207,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 +236,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 retry_target_stream(
+                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,221 @@
+# 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.
+
+"""
+Generator wrapper for retryable streaming RPCs.
+This function will be used when initilizing a retry with
+``Retry(is_stream=True)``.
+
+When ``is_stream=False``, the target is treated as a callable,
+and will retry when the callable returns an error. When ``is_stream=True``,
+the target will be treated as a callable that retruns an iterable. Instead
+of just wrapping the initial call in retry logic, the entire iterable is
+wrapped, with each yield passing through the retryable generator. If any yield
+in the stream raises a retryable exception, the entire stream will be
+retried.
+
+Important Note: when a stream is encounters a retryable error, it will
+silently construct a fresh iterator instance in the background
+and continue yielding (likely duplicate) values as if no error occurred.
+This is the most general way to retry a stream, but it often is not the
+desired behavior. Example: iter([1, 2, 1/0]) -> [1, 2, 1, 2, ...]
+
+There are two ways to build more advanced retry logic for streams:
+
+1. Wrap the target
+    Use a ``target`` that maintains state between retries, and creates a
+    different generator on each retry call. For example, you can wrap a
+    network call in a function that modifies the request based on what has
+    already been returned:
+
+    ```
+    def attempt_with_modified_request(target, request, seen_items=[]):
+        # remove seen items from request on each attempt
+        new_request = modify_request(request, seen_items)
+        new_generator = target(new_request)
+        for item in new_generator:
+            yield item
+            seen_items.append(item)
+
+    retry_wrapped = Retry(is_stream=True)(attempt_with_modified_request, target, request, [])
+    ```
+
+    2. Wrap the retry generator
+        Alternatively, you can wrap the retryable generator itself before
+        passing it to the end-user to add a filter on the stream. For
+        example, you can keep track of the items that were successfully yielded
+        in previous retry attempts, and only yield new items when the
+        new attempt surpasses the previous ones:
+
+        ``
+        def retryable_with_filter(target):
+            stream_idx = 0
+            # reset stream_idx when the stream is retried
+            def on_error(e):
+                nonlocal stream_idx
+                stream_idx = 0
+            # build retryable
+            retryable_gen = Retry(is_stream=True, on_error=on_error, ...)(target)
+            # keep track of what has been yielded out of filter
+            yielded_items = []
+            for item in retryable_gen:
+                if stream_idx >= len(yielded_items):
+                    yield item
+                    yielded_items.append(item)
+                elif item != previous_stream[stream_idx]:
+                    raise ValueError("Stream differs from last attempt")"
+                stream_idx += 1
+
+        filter_retry_wrapped = retryable_with_filter(target)
+        ```
+"""
+
+from typing import (
+    Callable,
+    Optional,
+    List,
+    Tuple,
+    Iterable,
+    Generator,
+    TypeVar,
+    Any,
+    Union,
+    cast,
+)
+
+import logging
+import time
+from functools import partial
+
+from google.api_core import exceptions
+
+_LOGGER = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def _build_timeout_error(
+    exc_list: List[Exception], is_timeout: bool, timeout_val: float
+) -> Tuple[Exception, Optional[Exception]]:
+    """
+    Default exception_factory implementation. Builds an exception after the retry fails
+
+    Args:
+      - exc_list (list[Exception]): list of exceptions that occurred during the retry
+      - is_timeout (bool): whether the failure is due to the timeout value being exceeded,
+          or due to a non-retryable exception
+      - timeout_val (float): the original timeout value for the retry, for use in the exception message
+
+    Returns:
+      - tuple[Exception, Exception|None]: a tuple of the exception to be raised, and the cause exception if any
+    """
+    src_exc = exc_list[-1] if exc_list else None
+    if is_timeout:
+        return (
+            exceptions.RetryError(
+                "Timeout of {:.1f}s exceeded".format(timeout_val),
+                src_exc,
+            ),
+            src_exc,
+        )
+    else:
+        return exc_list[-1], None
+
+
+def retry_target_stream(
+    target: Callable[[], Union[Iterable[T], Generator[T, Any, None]]],
+    predicate: Callable[[Exception], bool],
+    sleep_generator: Iterable[float],
+    timeout: Optional[float] = None,
+    on_error: Optional[Callable[[Exception], None]] = None,
+    exception_factory: Optional[
+        Callable[[List[Exception], bool, float], Tuple[Exception, Optional[Exception]]]
+    ] = None,
+    **kwargs,
+) -> Generator[T, Any, None]:
+    """Create a generator wrapper that retries the wrapped stream if it fails.
+
+    This is the lowest-level retry helper. Generally, you'll use the
+    higher-level retry helper :class:`Retry`.
+
+    Args:
+        target: The generator function to call and 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.
+            Note: timeout is only checked before initiating a retry, so the target may
+            run past the timeout value as long as it is healthy.
+        on_error: A function to call while processing a
+            retryable exception.  Any error raised by this function will *not*
+            be caught.
+        exception_factory: A function that is called when the retryable reaches
+            a terminal failure state, used to construct an exception to be raised.
+            It it given a list of all exceptions encountered, a boolean indicating
+            whether the failure was due to a timeout, and the original timeout value
+            as arguments. It should return a tuple of the exception to be raised,
+            along with the cause exception if any.
+            If not provided, a default implementation will raise a RetryError
+            on timeout, or the last exception encountered otherwise.
+
+    Returns:
+        Generator: A retryable generator that wraps the target generator function.
+
+    Raises:
+        ValueError: If the sleep generator stops yielding values.
+        Exception: a custom exception specified by the exception_factory if provided.
+            If no exception_factory is provided:
+                google.api_core.RetryError: If the deadline is exceeded while retrying.
+                Exception: If the target raises an error that isn't retryable.
+    """
+
+    timeout = kwargs.get("deadline", timeout)
+    deadline: Optional[float] = time.monotonic() + timeout if timeout else None
+    error_list: List[Exception] = []
+    exc_factory = partial(
+        exception_factory or _build_timeout_error, timeout_val=timeout
+    )
+
+    for sleep in sleep_generator:
+        # Start a new retry loop
+        try:
+            # create and yeild from a new instance of the generator from input generator function
+            subgenerator = target()
+            return (yield from subgenerator)
+        # handle exceptions raised by the subgenerator
+        except Exception as exc:
+            error_list.append(exc)
+            if not predicate(exc):
+                final_exc, source_exc = exc_factory(
+                    exc_list=error_list, is_timeout=False
+                )
+                raise final_exc from source_exc
+            if on_error is not None:
+                on_error(exc)
+        finally:
+            if subgenerator is not None and getattr(subgenerator, "close", None):
+                cast(Generator, subgenerator).close()
+
+        if deadline is not None and time.monotonic() + sleep > deadline:
+            final_exc, source_exc = exc_factory(exc_list=error_list, is_timeout=True)
+            raise final_exc from source_exc
+        _LOGGER.debug(
+            "Retrying due to {}, sleeping {:.1f}s ...".format(error_list[-1], sleep)
+        )
+        time.sleep(sleep)
+
+    raise ValueError("Sleep generator stopped yielding sleep values.")