diff --git a/airbyte_cdk/sources/declarative/checks/check_stream.py b/airbyte_cdk/sources/declarative/checks/check_stream.py index 618511f8d7..85bed21114 100644 --- a/airbyte_cdk/sources/declarative/checks/check_stream.py +++ b/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 diff --git a/airbyte_cdk/sources/declarative/requesters/http_requester.py b/airbyte_cdk/sources/declarative/requesters/http_requester.py index bd1468bc7b..31b15853e1 100644 --- a/airbyte_cdk/sources/declarative/requesters/http_requester.py +++ b/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 diff --git a/airbyte_cdk/sources/streams/concurrent/default_stream.py b/airbyte_cdk/sources/streams/concurrent/default_stream.py index 6cc6e44d44..e88fdd40d0 100644 --- a/airbyte_cdk/sources/streams/concurrent/default_stream.py +++ b/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 ` - # 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 "" - ) - - 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 "" - ) + 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 ` + # 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 "" + ) + + 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 "" + ) diff --git a/airbyte_cdk/sources/streams/http/http_client.py b/airbyte_cdk/sources/streams/http/http_client.py index 8da499aa32..92fb308480 100644 --- a/airbyte_cdk/sources/streams/http/http_client.py +++ b/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 diff --git a/unit_tests/sources/declarative/checks/test_check_stream.py b/unit_tests/sources/declarative/checks/test_check_stream.py index 8bc9571d1a..7613ab1612 100644 --- a/unit_tests/sources/declarative/checks/test_check_stream.py +++ b/unit_tests/sources/declarative/checks/test_check_stream.py @@ -6,20 +6,22 @@ import logging from copy import deepcopy from typing import Any, Iterable, Mapping, Optional -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import requests from jsonschema.exceptions import ValidationError -from airbyte_cdk.models import Status +from airbyte_cdk.models import FailureType, Status from airbyte_cdk.sources.declarative.checks.check_stream import CheckStream from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( ConcurrentDeclarativeSource, ) +from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream from airbyte_cdk.sources.streams.core import Stream from airbyte_cdk.sources.streams.http import HttpStream from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse +from airbyte_cdk.utils.traced_exception import AirbyteTracedException logger = logging.getLogger("test") config = dict() @@ -796,6 +798,126 @@ def test_check_stream1( assert connection_status.status == expected_result +@pytest.mark.parametrize( + "failure_type, should_raise", + [ + (FailureType.transient_error, True), + (FailureType.config_error, False), + ], +) +def test_check_connection_propagates_transient_availability_errors(failure_type, should_raise): + stream = MagicMock(spec=AbstractStream) + stream.name = "mock_stream" + error = AirbyteTracedException(message="request failed", failure_type=failure_type) + stream.check_availability.side_effect = error + source = MagicMock() + source.streams.return_value = [stream] + check_stream = CheckStream(stream_names=["mock_stream"], parameters={}) + + if should_raise: + with pytest.raises(AirbyteTracedException) as raised_error: + check_stream.check_connection(source, logger, config) + assert raised_error.value is error + else: + available, reason = check_stream.check_connection(source, logger, config) + assert available is False + assert "request failed" in reason + + +def test_check_generated_stream_availability_propagates_transient_error(): + stream = MagicMock(spec=AbstractStream) + stream.name = "generated_stream" + error = AirbyteTracedException( + message="request failed", failure_type=FailureType.transient_error + ) + stream.check_availability.side_effect = error + check_stream = CheckStream(stream_names=[], parameters={}) + + with pytest.raises(AirbyteTracedException) as raised_error: + check_stream._check_generated_streams_availability( + [{"name": "generated_stream"}], {"generated_stream": stream}, logger, None + ) + + assert raised_error.value is error + + +@pytest.mark.parametrize("status_code", [200, 429]) +def test_concurrent_declarative_check_fails_fast_when_rate_limited(status_code): + manifest = { + "version": "6.7.0", + "type": "DeclarativeSource", + "check": {"type": "CheckStream", "stream_names": ["stream"]}, + "streams": [ + { + "type": "DeclarativeStream", + "name": "stream", + "primary_key": ["id"], + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": { + "type": "object", + "properties": {"id": {"type": "integer"}}, + }, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "/items", + "http_method": "GET", + "error_handler": { + "type": "DefaultErrorHandler", + "max_retries": 1, + "backoff_strategies": [ + {"type": "WaitTimeFromHeader", "header": "Retry-After"} + ], + }, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "paginator": {"type": "NoPagination"}, + }, + } + ], + } + request = HttpRequest(url="https://api.test.com/items") + response = HttpResponse( + body=json.dumps([{"id": 1}] if status_code == 200 else {"error": "rate limited"}), + status_code=status_code, + headers={"Retry-After": "900"} if status_code == 429 else {}, + ) + + with HttpMocker() as http_mocker: + http_mocker.get(request, response) + source = ConcurrentDeclarativeSource( + source_config=manifest, + config={}, + catalog=None, + state=None, + ) + + with patch("time.sleep") as mocked_sleep: + if status_code == 429: + with pytest.raises(AirbyteTracedException) as error: + source.check(logger, {}) + + assert error.value.failure_type == FailureType.transient_error + assert error.value.message == ( + "API rate limit exceeded. Wait for the rate limit to reset and try again." + ) + http_mocker.assert_number_of_calls(request, 1) + else: + connection_status = source.check(logger, {}) + + assert connection_status.status == Status.SUCCEEDED + http_mocker.assert_number_of_calls(request, 1) + + mocked_sleep.assert_not_called() + + def test_check_empty_static_stream_override_falls_back_to_manifest_streams_and_checks_dynamic_streams(): manifest = { **deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT), @@ -840,11 +962,12 @@ def test_check_empty_static_stream_override_falls_back_to_manifest_streams_and_c state=None, ) - connection_status = source.check(logger, check_config) + with pytest.raises(AirbyteTracedException) as error: + source.check(logger, check_config) + assert error.value.failure_type == FailureType.transient_error http_mocker.assert_number_of_calls(static_stream_request, 6) http_mocker.assert_number_of_calls(item_request_2, 0) - assert connection_status.status == Status.FAILED def test_check_stream_missing_fields(): diff --git a/unit_tests/sources/declarative/requesters/test_http_requester.py b/unit_tests/sources/declarative/requesters/test_http_requester.py index 88eabad54f..12c916595f 100644 --- a/unit_tests/sources/declarative/requesters/test_http_requester.py +++ b/unit_tests/sources/declarative/requesters/test_http_requester.py @@ -228,6 +228,22 @@ def test_basic_send_request(): assert sent_request.body is None +def test_fail_fast_on_rate_limit_is_forwarded_to_http_client(mocker): + requester = create_requester() + assert requester.fail_fast_on_rate_limit is False + requester.fail_fast_on_rate_limit = True + send_request = mocker.patch.object( + requester._http_client, + "send_request", + return_value=(PreparedRequest(), requests.Response()), + ) + + requester.send_request() + + send_request.assert_called_once() + assert send_request.call_args.kwargs["fail_fast_on_rate_limit"] is True + + @pytest.mark.parametrize( "provider_data, provider_json, param_data, param_json, authenticator_data, authenticator_json, expected_exception, expected_body", [ diff --git a/unit_tests/sources/streams/concurrent/test_default_stream.py b/unit_tests/sources/streams/concurrent/test_default_stream.py index e294928160..3aebe1f89d 100644 --- a/unit_tests/sources/streams/concurrent/test_default_stream.py +++ b/unit_tests/sources/streams/concurrent/test_default_stream.py @@ -2,11 +2,15 @@ # Copyright (c) 2023 Airbyte, Inc., all rights reserved. # import unittest +from types import SimpleNamespace from unittest.mock import Mock import pytest -from airbyte_cdk.models import AirbyteStream, SyncMode +from airbyte_cdk.models import AirbyteStream, FailureType, SyncMode +from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import ( + StreamSlicerPartitionGenerator, +) from airbyte_cdk.sources.message import InMemoryMessageRepository from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, CursorField, FinalStateCursor from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream @@ -367,6 +371,102 @@ def test_given_AirbyteTracedException_when_reading_records_when_get_availability assert availability.is_available == False + def _stream_with_retriever(self, retriever): + partition_generator = Mock(spec=StreamSlicerPartitionGenerator) + partition_generator._partition_factory = SimpleNamespace(_retriever=retriever) + partition_generator.generate.return_value = [self._partition] + return DefaultStream( + partition_generator, + self._name, + self._json_schema, + self._primary_key, + self._cursor_field, + self._logger, + self._cursor, + ), partition_generator + + def test_check_availability_temporarily_enables_fail_fast_on_rate_limit(self): + for use_requester in [False, True]: + for previous_value in [False, True]: + with self.subTest(use_requester=use_requester, previous_value=previous_value): + rate_limited_component = SimpleNamespace( + fail_fast_on_rate_limit=previous_value, exit_on_rate_limit=False + ) + retriever = ( + SimpleNamespace(requester=rate_limited_component) + if use_requester + else rate_limited_component + ) + stream, partition_generator = self._stream_with_retriever(retriever) + + def generate(): + assert rate_limited_component.fail_fast_on_rate_limit is True + assert rate_limited_component.exit_on_rate_limit is False + return [self._partition] + + def read(): + assert rate_limited_component.fail_fast_on_rate_limit is True + assert rate_limited_component.exit_on_rate_limit is False + return [Mock(spec=Record)] + + partition_generator.generate.side_effect = generate + self._partition.read.side_effect = read + + availability = stream.check_availability() + + assert availability.is_available is True + assert rate_limited_component.fail_fast_on_rate_limit is previous_value + assert rate_limited_component.exit_on_rate_limit is False + + def test_check_availability_propagates_transient_errors_and_restores_flag(self): + for failure_type, should_raise in [ + (FailureType.transient_error, True), + (FailureType.config_error, False), + ]: + for raise_during_generation in [False, True]: + for previous_value in [False, True]: + for use_requester in [False, True]: + with self.subTest( + failure_type=failure_type, + raise_during_generation=raise_during_generation, + previous_value=previous_value, + use_requester=use_requester, + ): + fail_fast_component = SimpleNamespace( + fail_fast_on_rate_limit=previous_value, exit_on_rate_limit=False + ) + retriever = ( + SimpleNamespace(requester=fail_fast_component) + if use_requester + else fail_fast_component + ) + stream, partition_generator = self._stream_with_retriever(retriever) + error = AirbyteTracedException( + message="request failed", failure_type=failure_type + ) + + def raise_error(): + assert fail_fast_component.fail_fast_on_rate_limit is True + assert fail_fast_component.exit_on_rate_limit is False + raise error + + if raise_during_generation: + partition_generator.generate.side_effect = raise_error + else: + self._partition.read.side_effect = raise_error + + if should_raise: + with pytest.raises(AirbyteTracedException) as raised_error: + stream.check_availability() + assert raised_error.value is error + else: + availability = stream.check_availability() + assert availability.is_available is False + assert availability.reason == "request failed" + + assert fail_fast_component.fail_fast_on_rate_limit is previous_value + assert fail_fast_component.exit_on_rate_limit is False + def test_given_unknown_error_when_reading_record_when_get_availability_then_raise(self) -> None: """ I'm not sure why we handle AirbyteTracedException but not other exceptions but this is to keep feature compatibility with HttpAvailabilityStrategy diff --git a/unit_tests/sources/streams/http/test_http_client.py b/unit_tests/sources/streams/http/test_http_client.py index 31d4ff9a1c..38dd009d6c 100644 --- a/unit_tests/sources/streams/http/test_http_client.py +++ b/unit_tests/sources/streams/http/test_http_client.py @@ -742,6 +742,85 @@ def test_backoff_strategy_endless( assert mocked_send.call_count == expected_call_count +def test_fail_fast_on_rate_limit_without_evaluating_backoff_strategy_or_sleeping(requests_mock): + url = "https://test_base_url.com/v1/endpoint" + requests_mock.get(url, [{"status_code": 429}, {"status_code": 200}]) + backoff_strategy = CustomBackoffStrategy(backoff_time_value=900) + backoff_strategy.backoff_time = MagicMock(wraps=backoff_strategy.backoff_time) + http_client = HttpClient( + name="test", + logger=MagicMock(), + error_handler=HttpStatusErrorHandler( + logger=MagicMock(), + max_retries=1, + max_time=timedelta(seconds=1800), + error_mapping={ + 429: ErrorResolution( + ResponseAction.RATE_LIMITED, FailureType.transient_error, "rate limit" + ) + }, + ), + backoff_strategy=backoff_strategy, + ) + + with patch("time.sleep") as mocked_sleep: + with pytest.raises(AirbyteTracedException) as error: + http_client.send_request( + http_method="get", + url=url, + request_kwargs={}, + fail_fast_on_rate_limit=True, + ) + + assert error.value.failure_type == FailureType.transient_error + assert error.value.message == ( + "API rate limit exceeded. Wait for the rate limit to reset and try again." + ) + assert error.value.internal_message == ( + "Rate limited while exit_on_rate_limit is set, not waiting: rate limit" + ) + assert isinstance(error.value._exception, requests.HTTPError) + assert error.value._stream_descriptor.name == "test" + backoff_strategy.backoff_time.assert_not_called() + mocked_sleep.assert_not_called() + assert requests_mock.call_count == 1 + + +def test_exit_on_rate_limit_alone_retries_rate_limit(requests_mock): + url = "https://test_base_url.com/v1/endpoint" + requests_mock.get(url, [{"status_code": 429}, {"status_code": 200}]) + backoff_strategy = CustomBackoffStrategy(backoff_time_value=900) + backoff_strategy.backoff_time = MagicMock(wraps=backoff_strategy.backoff_time) + http_client = HttpClient( + name="test", + logger=MagicMock(), + error_handler=HttpStatusErrorHandler( + logger=MagicMock(), + max_retries=1, + max_time=timedelta(seconds=1800), + error_mapping={ + 429: ErrorResolution( + ResponseAction.RATE_LIMITED, FailureType.transient_error, "rate limit" + ) + }, + ), + backoff_strategy=backoff_strategy, + ) + + with patch("time.sleep") as mocked_sleep: + _, response = http_client.send_request( + http_method="get", + url=url, + request_kwargs={}, + exit_on_rate_limit=True, + ) + + assert response.status_code == 200 + assert backoff_strategy.backoff_time.call_count == 1 + assert mocked_sleep.call_args_list[0].args == (901,) + assert requests_mock.call_count == 2 + + def test_given_different_headers_then_response_is_not_cached(requests_mock): http_client = HttpClient(name="test", logger=MagicMock(), use_cache=True) first_request_headers = {"header_key": "first"}