feat: retry and retry_async support streaming rpcs

google/api_core/__init__.py

@@ -20,3 +20,6 @@
 from google.api_core import version as api_core_version
 
 __version__ = api_core_version.__version__
+
+# for backwards compatibility, expose async unary retries as google.api_core.retry_async
+from .retry import retry_unary_async as retry_async  # noqa: F401

google/api_core/retry/__init__.py

@@ -0,0 +1,45 @@
+# 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.
+
+"""Retry implementation for Google API client libraries."""
+
+from .retry_base import exponential_sleep_generator
+from .retry_base import if_exception_type
+from .retry_base import if_transient_error
+from .retry_base import _build_retry_error
+from .retry_base import RetryFailureReason
+from .retry_unary import Retry
+from .retry_unary import retry_target
+from .retry_unary_async import AsyncRetry
+from .retry_unary_async import retry_target as retry_target_async
+from .retry_streaming import StreamingRetry
+from .retry_streaming import retry_target_stream
+from .retry_streaming_async import AsyncStreamingRetry
+from .retry_streaming_async import retry_target_stream as retry_target_stream_async
+
+__all__ = (
+    "exponential_sleep_generator",
+    "if_exception_type",
+    "if_transient_error",
+    "_build_retry_error",
+    "RetryFailureReason",
+    "Retry",
+    "AsyncRetry",
+    "StreamingRetry",
+    "AsyncStreamingRetry",
+    "retry_target",
+    "retry_target_async",
+    "retry_target_stream",
+    "retry_target_stream_async",
+)

google/api_core/retry/retry_base.py

@@ -0,0 +1,346 @@
+# 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.
+
+"""Shared classes and functions for retrying requests.
+
+:class:`_BaseRetry` is the base class for :class:`Retry`,
+:class:`AsyncRetry`, :class:`StreamingRetry`, and :class:`AsyncStreamingRetry`.
+"""
+
+from __future__ import annotations
+
+import logging
+import random
+import time
+
+from enum import Enum
+from typing import Any, Callable, TYPE_CHECKING
+
+import requests.exceptions
+
+from google.api_core import exceptions
+from google.auth import exceptions as auth_exceptions
+
+if TYPE_CHECKING:
+    import sys
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
+_DEFAULT_INITIAL_DELAY = 1.0  # seconds
+_DEFAULT_MAXIMUM_DELAY = 60.0  # seconds
+_DEFAULT_DELAY_MULTIPLIER = 2.0
+_DEFAULT_DEADLINE = 60.0 * 2.0  # seconds
+
+_LOGGER = logging.getLogger("google.api_core.retry")
+
+
+def if_exception_type(
+    *exception_types: type[Exception],
+) -> Callable[[Exception], bool]:
+    """Creates a predicate to check if the exception is of a given type.
+
+    Args:
+        exception_types (Sequence[:func:`type`]): The exception types to check
+            for.
+
+    Returns:
+        Callable[Exception]: A predicate that returns True if the provided
+            exception is of the given type(s).
+    """
+
+    def if_exception_type_predicate(exception: Exception) -> bool:
+        """Bound predicate for checking an exception type."""
+        return isinstance(exception, exception_types)
+
+    return if_exception_type_predicate
+
+
+# pylint: disable=invalid-name
+# Pylint sees this as a constant, but it is also an alias that should be
+# considered a function.
+if_transient_error = if_exception_type(
+    exceptions.InternalServerError,
+    exceptions.TooManyRequests,
+    exceptions.ServiceUnavailable,
+    requests.exceptions.ConnectionError,
+    requests.exceptions.ChunkedEncodingError,
+    auth_exceptions.TransportError,
+)
+"""A predicate that checks if an exception is a transient API error.
+
+The following server errors are considered transient:
+
+- :class:`google.api_core.exceptions.InternalServerError` - HTTP 500, gRPC
+    ``INTERNAL(13)`` and its subclasses.
+- :class:`google.api_core.exceptions.TooManyRequests` - HTTP 429
+- :class:`google.api_core.exceptions.ServiceUnavailable` - HTTP 503
+- :class:`requests.exceptions.ConnectionError`
+- :class:`requests.exceptions.ChunkedEncodingError` - The server declared
+    chunked encoding but sent an invalid chunk.
+- :class:`google.auth.exceptions.TransportError` - Used to indicate an
+    error occurred during an HTTP request.
+"""
+# pylint: enable=invalid-name
+
+
+def exponential_sleep_generator(
+    initial: float, maximum: float, multiplier: float = _DEFAULT_DELAY_MULTIPLIER
+):
+    """Generates sleep intervals based on the exponential back-off algorithm.
+
+    This implements the `Truncated Exponential Back-off`_ algorithm.
+
+    .. _Truncated Exponential Back-off:
+        https://cloud.google.com/storage/docs/exponential-backoff
+
+    Args:
+        initial (float): The minimum amount of time to delay. This must
+            be greater than 0.
+        maximum (float): The maximum amount of time to delay.
+        multiplier (float): The multiplier applied to the delay.
+
+    Yields:
+        float: successive sleep intervals.
+    """
+    delay = min(initial, maximum)
+    while True:
+        yield random.uniform(0.0, delay)
+        delay = min(delay * multiplier, maximum)
+
+
+class RetryFailureReason(Enum):
+    """
+    The cause of a failed retry, used when building exceptions
+    """
+
+    TIMEOUT = 0
+    NON_RETRYABLE_ERROR = 1
+
+
+def _build_retry_error(
+    exc_list: list[Exception],
+    reason: RetryFailureReason,
+    timeout_val: float | None,
+    **kwargs: Any,
+) -> tuple[Exception, Exception | None]:
+    """
+    Default exception_factory implementation. Builds an exception after the retry fails
+
+    Args:
+      - exc_list: list of exceptions that occurred during the retry
+      - reason: reason for the retry failure.
+            Can be TIMEOUT or NON_RETRYABLE_ERROR
+      - timeout_val: the original timeout value for the retry, for use in the exception message
+
+    Returns:
+      - tuple: a tuple of the exception to be raised, and the cause exception if any
+    """
+    if reason == RetryFailureReason.TIMEOUT:
+        # return RetryError with the most recent exception as the cause
+        src_exc = exc_list[-1] if exc_list else None
+        timeout_val_str = f"of {timeout_val:0.1f}s " if timeout_val is not None else ""
+        return (
+            exceptions.RetryError(
+                f"Timeout {timeout_val_str}exceeded",
+                src_exc,
+            ),
+            src_exc,
+        )
+    elif exc_list:
+        # return most recent exception encountered
+        return exc_list[-1], None
+    else:
+        # no exceptions were given in exc_list. Raise generic RetryError
+        return exceptions.RetryError("Unknown error", None), None
+
+
+def _retry_error_helper(
+    exc: Exception,
+    deadline: float | None,
+    next_sleep: float,
+    error_list: list[Exception],
+    predicate_fn: Callable[[Exception], bool],
+    on_error_fn: Callable[[Exception], None] | None,
+    exc_factory_fn: Callable[
+        [list[Exception], RetryFailureReason],
+        tuple[Exception, Exception | None],
+    ],
+):
+    """
+    Shared logic for handling an error for all retry implementations
+
+    - Raises an error on timeout or non-retryable error
+    - Calls on_error_fn if provided
+    - Logs the error
+
+    Args:
+       - exc: the exception that was raised
+       - deadline: the deadline for the retry, calculated as a diff from time.monotonic()
+       - next_sleep: the calculated next sleep interval
+       - error_list: the list of exceptions that have been raised so far
+       - predicate_fn: the predicate that was used to determine if the exception should be retried
+       - on_error_fn: the callback that was called when the exception was raised
+       - exc_factory_fn: the callback that was called to build the exception to be raised on terminal failure
+    """
+    error_list.append(exc)
+    if not predicate_fn(exc):
+        final_exc, source_exc = exc_factory_fn(
+            error_list,
+            RetryFailureReason.NON_RETRYABLE_ERROR,
+        )
+        raise final_exc from source_exc
+    if on_error_fn is not None:
+        on_error_fn(exc)
+    if deadline is not None and time.monotonic() + next_sleep > deadline:
+        final_exc, source_exc = exc_factory_fn(
+            error_list,
+            RetryFailureReason.TIMEOUT,
+        )
+        raise final_exc from source_exc
+    _LOGGER.debug(
+        "Retrying due to {}, sleeping {:.1f}s ...".format(error_list[-1], next_sleep)
+    )
+
+
+class _BaseRetry(object):
+    """
+    Base class for retry configuration objects. This class is intended to capture retry
+    and backoff configuration that is common to both synchronous and asynchronous retries,
+    for both unary and streaming RPCs. It is not intended to be instantiated directly,
+    but rather to be subclassed by the various retry configuration classes.
+    """
+
+    def __init__(
+        self,
+        predicate: Callable[[Exception], bool] = if_transient_error,
+        initial: float = _DEFAULT_INITIAL_DELAY,
+        maximum: float = _DEFAULT_MAXIMUM_DELAY,
+        multiplier: float = _DEFAULT_DELAY_MULTIPLIER,
+        timeout: float = _DEFAULT_DEADLINE,
+        on_error: Callable[[Exception], Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self._predicate = predicate
+        self._initial = initial
+        self._multiplier = multiplier
+        self._maximum = maximum
+        self._timeout = kwargs.get("deadline", timeout)
+        self._deadline = self._timeout
+        self._on_error = on_error
+
+    def __call__(self, *args, **kwargs) -> Any:
+        raise NotImplementedError("Not implemented in base class")
+
+    @property
+    def deadline(self) -> float | None:
+        """
+        DEPRECATED: use ``timeout`` instead.  Refer to the ``Retry`` class
+        documentation for details.
+        """
+        return self._timeout
+
+    @property
+    def timeout(self) -> float | None:
+        return self._timeout
+
+    def _replace(
+        self,
+        predicate: Callable[[Exception], bool] | None = None,
+        initial: float | None = None,
+        maximum: float | None = None,
+        multiplier: float | None = None,
+        timeout: float | None = None,
+        on_error: Callable[[Exception], Any] | None = None,
+    ) -> Self:
+        return type(self)(
+            predicate=predicate or self._predicate,
+            initial=initial or self._initial,
+            maximum=maximum or self._maximum,
+            multiplier=multiplier or self._multiplier,
+            timeout=timeout or self._timeout,
+            on_error=on_error or self._on_error,
+        )
+
+    def with_deadline(self, deadline: float | None) -> Self:
+        """Return a copy of this retry with the given timeout.
+
+        DEPRECATED: use :meth:`with_timeout` instead. Refer to the ``Retry`` class
+        documentation for details.
+
+        Args:
+            deadline (float): How long to keep retrying in seconds.
+
+        Returns:
+            Retry: A new retry instance with the given timeout.
+        """
+        return self._replace(timeout=deadline)
+
+    def with_timeout(self, timeout: float) -> Self:
+        """Return a copy of this retry with the given timeout.
+
+        Args:
+            timeout (float): How long to keep retrying, in seconds.
+
+        Returns:
+            Retry: A new retry instance with the given timeout.
+        """
+        return self._replace(timeout=timeout)
+
+    def with_predicate(self, predicate: Callable[[Exception], bool]) -> Self:
+        """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 self._replace(predicate=predicate)
+
+    def with_delay(
+        self,
+        initial: float | None = None,
+        maximum: float | None = None,
+        multiplier: float | None = None,
+    ) -> Self:
+        """Return a copy of this retry with the given delay options.
+
+        Args:
+            initial (float): The minimum amount of time to delay. This must
+                be greater than 0.
+            maximum (float): The maximum amount of time to delay.
+            multiplier (float): The multiplier applied to the delay.
+
+        Returns:
+            Retry: A new retry instance with the given predicate.
+        """
+        return self._replace(initial=initial, maximum=maximum, multiplier=multiplier)
+
+    def __str__(self) -> str:
+        return (
+            "<{} predicate={}, initial={:.1f}, maximum={:.1f}, "
+            "multiplier={:.1f}, timeout={}, on_error={}>".format(
+                type(self).__name__,
+                self._predicate,
+                self._initial,
+                self._maximum,
+                self._multiplier,
+                self._timeout,  # timeout can be None, thus no {:.1f}
+                self._on_error,
+            )
+        )