fix(declarative): fail check fast as transient_error when rate limited

airbyte_cdk/sources/declarative/checks/check_stream.py

@@ -7,11 +7,13 @@
 from dataclasses import InitVar, dataclass
 from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
 
+from airbyte_cdk.models import FailureType
 from airbyte_cdk.sources import Source
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.core import Stream
 from airbyte_cdk.sources.streams.http.availability_strategy import HttpAvailabilityStrategy
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 CHECK_STREAM_NAMES_CONFIG_KEY = "__airbyte_check_stream_names"
 
@@ -131,6 +133,10 @@ def _check_stream_availability(
                 message = f"Stream {stream_name} is not available: {reason}"
                 logger.warning(message)
                 return stream_is_available, message
+        except AirbyteTracedException as error:
+            if error.failure_type == FailureType.transient_error:
+                raise
+            return self._log_error(logger, f"checking availability of stream {stream_name}", error)
         except Exception as error:
             return self._log_error(logger, f"checking availability of stream {stream_name}", error)
         return True, None
@@ -194,6 +200,12 @@ def _check_generated_streams_availability(
                     message = f"Dynamic Stream {stream.name} is not available: {reason}"
                     logger.warning(message)
                     return False, message
+            except AirbyteTracedException as error:
+                if error.failure_type == FailureType.transient_error:
+                    raise
+                return self._log_error(
+                    logger, f"checking availability of dynamic stream {stream.name}", error
+                )
             except Exception as error:
                 return self._log_error(
                     logger, f"checking availability of dynamic stream {stream.name}", error

airbyte_cdk/sources/declarative/requesters/http_requester.py

@@ -70,6 +70,7 @@ class HttpRequester(Requester):
     read_timeout_in_seconds: Optional[float] = None
     use_tcp_keepalive: bool = False
     _exit_on_rate_limit: bool = False
+    _fail_fast_on_rate_limit: bool = False
     stream_response: bool = False
     decoder: Decoder = field(default_factory=lambda: JsonDecoder(parameters={}))
 
@@ -134,6 +135,14 @@ def exit_on_rate_limit(self) -> bool:
     def exit_on_rate_limit(self, value: bool) -> None:
         self._exit_on_rate_limit = value
 
+    @property
+    def fail_fast_on_rate_limit(self) -> bool:
+        return self._fail_fast_on_rate_limit
+
+    @fail_fast_on_rate_limit.setter
+    def fail_fast_on_rate_limit(self, value: bool) -> None:
+        self._fail_fast_on_rate_limit = value
+
     def get_authenticator(self) -> DeclarativeAuthenticator:
         return self._authenticator
 
@@ -489,6 +498,7 @@ def send_request(
             dedupe_query_params=True,
             log_formatter=log_formatter,
             exit_on_rate_limit=self._exit_on_rate_limit,
+            fail_fast_on_rate_limit=self._fail_fast_on_rate_limit,
         )
 
         return response

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -2,10 +2,11 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+from contextlib import contextmanager
 from logging import Logger
-from typing import Any, Callable, Iterable, List, Mapping, Optional, Union
+from typing import Any, Callable, Iterable, Iterator, List, Mapping, Optional, Union
 
-from airbyte_cdk.models import AirbyteStream, SyncMode
+from airbyte_cdk.models import AirbyteStream, FailureType, SyncMode
 from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import (
     ConcurrentPerPartitionCursor,
 )
@@ -121,32 +122,63 @@ def get_partition_router(self) -> PartitionRouter | None:
             return None
         return stream_slicer._partition_router
 
+    @contextmanager
+    def _fail_fast_on_rate_limit(self) -> Iterator[None]:
+        if not isinstance(self._stream_partition_generator, StreamSlicerPartitionGenerator):
+            yield
+            return
+
+        retriever = self._stream_partition_generator._partition_factory._retriever
+        target: Any
+        if hasattr(retriever, "fail_fast_on_rate_limit"):
+            target = retriever
+        else:
+            target = getattr(retriever, "requester", None)
+
+        if target is None or not hasattr(target, "fail_fast_on_rate_limit"):
+            yield
+            return
+
+        previous_value = target.fail_fast_on_rate_limit
+        target.fail_fast_on_rate_limit = True
+        try:
+            yield
+        finally:
+            target.fail_fast_on_rate_limit = previous_value
+
     def check_availability(self) -> StreamAvailability:
         """
         Check stream availability by attempting to read the first record of the stream.
         """
-        try:
-            partition = next(iter(self.generate_partitions()))
-        except StopIteration:
-            # NOTE: The following comment was copied from legacy stuff and I don't know how relevant it is:
-            # If stream_slices has no `next()` item (Note - this is different from stream_slices returning [None]!)
-            # This can happen when a substream's `stream_slices` method does a `for record in parent_records: yield <something>`
-            # without accounting for the case in which the parent stream is empty.
-            return StreamAvailability.unavailable(
-                f"Cannot attempt to connect to stream {self.name} - no stream slices were found"
-            )
-        except AirbyteTracedException as error:
-            return StreamAvailability.unavailable(
-                error.message or error.internal_message or "<no error message>"
-            )
-
-        try:
-            next(iter(partition.read()))
-            return StreamAvailability.available()
-        except StopIteration:
-            self._logger.info(f"Successfully connected to stream {self.name}, but got 0 records.")
-            return StreamAvailability.available()
-        except AirbyteTracedException as error:
-            return StreamAvailability.unavailable(
-                error.message or error.internal_message or "<no error message>"
-            )
+        with self._fail_fast_on_rate_limit():
+            try:
+                partition = next(iter(self.generate_partitions()))
+            except StopIteration:
+                # NOTE: The following comment was copied from legacy stuff and I don't know how relevant it is:
+                # If stream_slices has no `next()` item (Note - this is different from stream_slices returning [None]!)
+                # This can happen when a substream's `stream_slices` method does a `for record in parent_records: yield <something>`
+                # without accounting for the case in which the parent stream is empty.
+                return StreamAvailability.unavailable(
+                    f"Cannot attempt to connect to stream {self.name} - no stream slices were found"
+                )
+            except AirbyteTracedException as error:
+                if error.failure_type == FailureType.transient_error:
+                    raise
+                return StreamAvailability.unavailable(
+                    error.message or error.internal_message or "<no error message>"
+                )
+
+            try:
+                next(iter(partition.read()))
+                return StreamAvailability.available()
+            except StopIteration:
+                self._logger.info(
+                    f"Successfully connected to stream {self.name}, but got 0 records."
+                )
+                return StreamAvailability.available()
+            except AirbyteTracedException as error:
+                if error.failure_type == FailureType.transient_error:
+                    raise
+                return StreamAvailability.unavailable(
+                    error.message or error.internal_message or "<no error message>"
+                )

airbyte_cdk/sources/streams/http/http_client.py

@@ -333,6 +333,7 @@ def _send_with_retry(
         request_kwargs: Mapping[str, Any],
         log_formatter: Optional[Callable[[requests.Response], Any]] = None,
         exit_on_rate_limit: Optional[bool] = False,
+        fail_fast_on_rate_limit: bool = False,
     ) -> requests.Response:
         """
         Sends a request with retry logic.
@@ -363,6 +364,7 @@ def _send_with_retry(
                 request_kwargs,
                 log_formatter=log_formatter,
                 exit_on_rate_limit=exit_on_rate_limit,
+                fail_fast_on_rate_limit=fail_fast_on_rate_limit,
             )  # type: ignore # mypy can't infer that backoff_handler wraps _send
 
             return response
@@ -450,6 +452,7 @@ def _send(
         request_kwargs: Mapping[str, Any],
         log_formatter: Optional[Callable[[requests.Response], Any]] = None,
         exit_on_rate_limit: Optional[bool] = False,
+        fail_fast_on_rate_limit: bool = False,
     ) -> requests.Response:
         if request not in self._request_attempt_count:
             self._request_attempt_count[request] = 1
@@ -548,6 +551,7 @@ def _send(
             request=request,
             error_resolution=error_resolution,
             exit_on_rate_limit=exit_on_rate_limit,
+            fail_fast_on_rate_limit=fail_fast_on_rate_limit,
         )
 
         return response  # type: ignore # will either return a valid response of type requests.Response or raise an exception
@@ -613,6 +617,7 @@ def _handle_error_resolution(
         request: requests.PreparedRequest,
         error_resolution: ErrorResolution,
         exit_on_rate_limit: Optional[bool] = False,
+        fail_fast_on_rate_limit: bool = False,
     ) -> None:
         if error_resolution.response_action not in self._ACTIONS_TO_RETRY_ON:
             self._evict_key(request)
@@ -793,6 +798,29 @@ def _handle_error_resolution(
                 and self._can_retry_on_another_token(request)
             )
 
+            if (
+                error_resolution.response_action == ResponseAction.RATE_LIMITED
+                and fail_fast_on_rate_limit
+                and not rotate_instead_of_waiting
+            ):
+                error_message = error_resolution.error_message or (
+                    f"Request to {request.url} failed with failure type "
+                    f"{error_resolution.failure_type}, response action "
+                    f"{error_resolution.response_action}."
+                )
+                raise AirbyteTracedException(
+                    message="API rate limit exceeded. Wait for the rate limit to reset and try again.",
+                    internal_message=(
+                        "Rate limited while exit_on_rate_limit is set, not waiting: "
+                        f"{error_message}"
+                    ),
+                    failure_type=FailureType.transient_error,
+                    exception=(
+                        requests.HTTPError(response=response) if response is not None else exc
+                    ),
+                    stream_descriptor=StreamDescriptor(name=self._name),
+                )
+
             if rotate_instead_of_waiting:
                 # Says that a wait was skipped without the number, which is no longer computed,
                 # and names the cap explicitly: a connector that configured one gets no other
@@ -871,6 +899,7 @@ def send_request(
         dedupe_query_params: bool = False,
         log_formatter: Optional[Callable[[requests.Response], Any]] = None,
         exit_on_rate_limit: Optional[bool] = False,
+        fail_fast_on_rate_limit: bool = False,
     ) -> Tuple[requests.PreparedRequest, requests.Response]:
         """
         Prepares and sends request and return request and response objects.
@@ -900,6 +929,7 @@ def send_request(
             request_kwargs=request_kwargs,
             log_formatter=log_formatter,
             exit_on_rate_limit=exit_on_rate_limit,
+            fail_fast_on_rate_limit=fail_fast_on_rate_limit,
         )
 
         return request, response