diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 6df783d1db..d8f2a29676 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2134,6 +2134,7 @@ definitions: - "$ref": "#/definitions/ConstantBackoffStrategy" - "$ref": "#/definitions/ExponentialBackoffStrategy" - "$ref": "#/definitions/WaitTimeFromHeader" + - "$ref": "#/definitions/WaitTimeFromResponse" - "$ref": "#/definitions/WaitUntilTimeFromHeader" - "$ref": "#/definitions/CustomBackoffStrategy" max_retries: @@ -5117,6 +5118,42 @@ definitions: $parameters: type: object additionalProperties: true + WaitTimeFromResponse: + title: Wait Time Extracted From Response + description: Compute the wait time, in relative seconds, from the JSON response body or the response headers. A non-JSON body is exposed as an empty `response`. If the expression evaluates to an empty, zero, negative or non-numeric result, or cannot be evaluated on the response, the next backoff strategy is used. + type: object + required: + - type + - wait_time_in_seconds + properties: + type: + type: string + enum: [WaitTimeFromResponse] + wait_time_in_seconds: + title: Wait Time in Seconds + description: The expression defining how long to wait before retrying, in seconds relative to now (convert an absolute timestamp with `timestamp(...) - now_utc().timestamp()`). An empty, zero, negative or non-numeric result, or an expression that cannot be evaluated on the response, defers to the next backoff strategy. Nothing bounds the wait unless `max_waiting_time_in_seconds` is set, so setting it is strongly recommended. + type: string + interpolation_context: + - config + - response + - headers + examples: + - "{{ response.retry_after_seconds }}" + - "{{ max(1, timestamp(response.error.message | regex_search('Retry after (\\S+Z)')) - now_utc().timestamp()) if (response is mapping and response.get('error') is mapping and ((response.error.message or '') | string | regex_search('Retry after (\\S+Z)'))) else '' }}" + max_waiting_time_in_seconds: + title: Max Waiting Time in Seconds + description: 'Stop the stream instead of waiting, when the computed wait is greater than or equal to this value. Can be a hardcoded number, or a string interpolated from the connector config so that the bound can be changed without a connector release. A value of 0 means never wait. Not evaluated when a rate-limited retry can rotate to another credential with quota; any other retryable error still consults this strategy.' + anyOf: + - type: number + - type: string + interpolation_context: + - config + examples: + - 3600 + - "{{ config['max_waiting_time'] * 60 }}" + $parameters: + type: object + additionalProperties: true WaitUntilTimeFromHeader: title: Wait Until Time Defined In Response Header description: Extract time at which we can retry the request from response header and wait for the difference between now and that time. diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 0317d52d65..b79a35b777 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -1455,6 +1455,26 @@ class WaitTimeFromHeader(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class WaitTimeFromResponse(BaseModel): + type: Literal["WaitTimeFromResponse"] + wait_time_in_seconds: str = Field( + ..., + description="The expression defining how long to wait before retrying, in seconds relative to now (convert an absolute timestamp with `timestamp(...) - now_utc().timestamp()`). An empty, zero, negative or non-numeric result, or an expression that cannot be evaluated on the response, defers to the next backoff strategy. Nothing bounds the wait unless `max_waiting_time_in_seconds` is set, so setting it is strongly recommended.", + examples=[ + "{{ response.retry_after_seconds }}", + "{{ max(1, timestamp(response.error.message | regex_search('Retry after (\\S+Z)')) - now_utc().timestamp()) if (response is mapping and response.get('error') is mapping and ((response.error.message or '') | string | regex_search('Retry after (\\S+Z)'))) else '' }}", + ], + title="Wait Time in Seconds", + ) + max_waiting_time_in_seconds: Optional[Union[float, str]] = Field( + None, + description="Stop the stream instead of waiting, when the computed wait is greater than or equal to this value. Can be a hardcoded number, or a string interpolated from the connector config so that the bound can be changed without a connector release. A value of 0 means never wait. Not evaluated when a rate-limited retry can rotate to another credential with quota; any other retryable error still consults this strategy.", + examples=[3600, "{{ config['max_waiting_time'] * 60 }}"], + title="Max Waiting Time in Seconds", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class WaitUntilTimeFromHeader(BaseModel): type: Literal["WaitUntilTimeFromHeader"] header: str = Field( @@ -2296,6 +2316,7 @@ class DefaultErrorHandler(BaseModel): ConstantBackoffStrategy, ExponentialBackoffStrategy, WaitTimeFromHeader, + WaitTimeFromResponse, WaitUntilTimeFromHeader, CustomBackoffStrategy, ] diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index 219e0f70ad..86d696af74 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -495,6 +495,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( WaitTimeFromHeader as WaitTimeFromHeaderModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + WaitTimeFromResponse as WaitTimeFromResponseModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( WaitUntilTimeFromHeader as WaitUntilTimeFromHeaderModel, ) @@ -538,6 +541,7 @@ ConstantBackoffStrategy, ExponentialBackoffStrategy, WaitTimeFromHeaderBackoffStrategy, + WaitTimeFromResponseBackoffStrategy, WaitUntilTimeFromHeaderBackoffStrategy, ) from airbyte_cdk.sources.declarative.requesters.http_job_repository import AsyncHttpJobRepository @@ -885,6 +889,7 @@ def _init_mappings(self) -> None: SubstreamPartitionRouterModel: self.create_substream_partition_router, ValidateAdheresToSchemaModel: self.create_validate_adheres_to_schema, WaitTimeFromHeaderModel: self.create_wait_time_from_header, + WaitTimeFromResponseModel: self.create_wait_time_from_response, WaitUntilTimeFromHeaderModel: self.create_wait_until_time_from_header, AsyncRetrieverModel: self.create_async_retriever, HttpComponentsResolverModel: self.create_http_components_resolver, @@ -5188,6 +5193,17 @@ def create_wait_time_from_header( max_waiting_time_in_seconds=model.max_waiting_time_in_seconds, ) + @staticmethod + def create_wait_time_from_response( + model: WaitTimeFromResponseModel, config: Config, **kwargs: Any + ) -> WaitTimeFromResponseBackoffStrategy: + return WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds=model.wait_time_in_seconds, + parameters=model.parameters or {}, + config=config, + max_waiting_time_in_seconds=model.max_waiting_time_in_seconds, + ) + @staticmethod def create_wait_until_time_from_header( model: WaitUntilTimeFromHeaderModel, config: Config, **kwargs: Any diff --git a/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/__init__.py b/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/__init__.py index 26ecafbde7..4900813817 100644 --- a/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/__init__.py +++ b/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/__init__.py @@ -11,6 +11,9 @@ from airbyte_cdk.sources.declarative.requesters.error_handlers.backoff_strategies.wait_time_from_header_backoff_strategy import ( WaitTimeFromHeaderBackoffStrategy, ) +from airbyte_cdk.sources.declarative.requesters.error_handlers.backoff_strategies.wait_time_from_response_backoff_strategy import ( + WaitTimeFromResponseBackoffStrategy, +) from airbyte_cdk.sources.declarative.requesters.error_handlers.backoff_strategies.wait_until_time_from_header_backoff_strategy import ( WaitUntilTimeFromHeaderBackoffStrategy, ) @@ -19,5 +22,6 @@ "ConstantBackoffStrategy", "ExponentialBackoffStrategy", "WaitTimeFromHeaderBackoffStrategy", + "WaitTimeFromResponseBackoffStrategy", "WaitUntilTimeFromHeaderBackoffStrategy", ] diff --git a/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_time_from_response_backoff_strategy.py b/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_time_from_response_backoff_strategy.py new file mode 100644 index 0000000000..4409da9395 --- /dev/null +++ b/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_time_from_response_backoff_strategy.py @@ -0,0 +1,158 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import logging +import math +from dataclasses import InitVar, dataclass +from typing import Any, Mapping, Optional, Union + +import requests +from jinja2.exceptions import UndefinedError + +from airbyte_cdk.models import FailureType +from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString +from airbyte_cdk.sources.declarative.interpolation.jinja import JinjaInterpolation +from airbyte_cdk.sources.declarative.requesters.error_handlers.backoff_strategies.max_waiting_time_helper import ( + evaluate_max_waiting_time, + interpolated_max_waiting_time, +) +from airbyte_cdk.sources.streams.http.error_handlers import BackoffStrategy +from airbyte_cdk.sources.types import Config +from airbyte_cdk.utils import AirbyteTracedException + +LOGGER = logging.getLogger("airbyte") + +_INTERPOLATION_CONTEXT = frozenset({"config", "response", "headers", "parameters"}) + + +@dataclass +class WaitTimeFromResponseBackoffStrategy(BackoffStrategy): + """ + Extract wait time, in relative seconds, from the JSON response body or the response headers. + + Attributes: + wait_time_in_seconds: Expression defining how long to wait before retrying. An empty, + zero, negative, non-finite or non-numeric result, or an expression that cannot be + evaluated on the response body, defers to the next backoff strategy. + max_waiting_time_in_seconds: Optional inclusive cap on the computed wait. + """ + + wait_time_in_seconds: Union[InterpolatedString, str] + parameters: InitVar[Mapping[str, Any]] + config: Config + max_waiting_time_in_seconds: Optional[Union[float, InterpolatedString, str]] = None + + def __post_init__(self, parameters: Mapping[str, Any]) -> None: + # Evaluated with `JinjaInterpolation` directly rather than `InterpolatedString`: the latter + # caches a template as plain text when its first evaluation returns the template unchanged, + # which is what `JinjaInterpolation` does when rendering raises `TypeError`. One response + # body of an unexpected shape would then turn the strategy off for the rest of the sync. + self._wait_time_template = ( + self.wait_time_in_seconds.string + if isinstance(self.wait_time_in_seconds, InterpolatedString) + else self.wait_time_in_seconds + ) + self._interpolation = JinjaInterpolation() + self._parameters = parameters + undeclared = ( + JinjaInterpolation._find_undeclared_variables(self._wait_time_template) + - _INTERPOLATION_CONTEXT + ) + if undeclared: + # A typo such as `resposne` would otherwise fail on every response and leave the + # strategy silently inert; it is a manifest bug, so it fails at startup. + raise AirbyteTracedException( + internal_message=( + f"WaitTimeFromResponse `wait_time_in_seconds` {self._wait_time_template!r} " + f"uses undeclared variables {sorted(undeclared)}; only " + f"{sorted(_INTERPOLATION_CONTEXT)} are available." + ), + message=( + "The connector's rate limit wait time expression is invalid. This is a " + "problem with the connector rather than with your configuration." + ), + failure_type=FailureType.system_error, + ) + self._max_waiting_time_in_seconds = interpolated_max_waiting_time( + self.max_waiting_time_in_seconds, parameters + ) + evaluate_max_waiting_time(self._max_waiting_time_in_seconds, self.config) + + def backoff_time( + self, + response_or_exception: Optional[Union[requests.Response, requests.RequestException]], + attempt_count: int, + ) -> Optional[float]: + response: Any = {} + headers: Mapping[str, Any] = {} + if isinstance(response_or_exception, requests.Response): + try: + response = response_or_exception.json() + except (ValueError, requests.RequestException, RecursionError): + # Unreadable body (not JSON, too deeply nested, or a broken stream): interpolate + # with an empty response so an expression reading only `headers` still works. + pass + headers = response_or_exception.headers + + try: + evaluated_wait_time = self._interpolation.eval( + self._wait_time_template, + self.config, + # The template as its own default re-raises an `UndefinedError` from the first + # render instead of returning None, so it is reported below. + self._wait_time_template, + parameters=self._parameters, + response=response, + headers=headers, + ) + if evaluated_wait_time == self._wait_time_template: + # `JinjaInterpolation._eval` returns the raw template when rendering raises + # `TypeError`, e.g. arithmetic on a string from the body. + raise TypeError("the expression could not be rendered") + except (UndefinedError, TypeError, ValueError, ArithmeticError) as error: + # The strategy runs on every retryable response, not only the one the expression was + # written for: a body of another shape (a list, or a string where an object was + # expected) must defer to the next strategy, not fail the stream. The body is not + # logged; the error type is enough to point at the expression. + LOGGER.warning( + "WaitTimeFromResponse: `wait_time_in_seconds` could not be evaluated on this " + f"response ({type(error).__name__}); deferring to the next backoff strategy." + ) + return None + if evaluated_wait_time is None or ( + isinstance(evaluated_wait_time, str) and not evaluated_wait_time.strip() + ): + return None + + try: + wait_time = float(evaluated_wait_time) + except (TypeError, ValueError, ArithmeticError): + # Mirrors the header strategies, which return None for a value they cannot read, so + # the next backoff strategy decides. The value is not logged: it comes from the + # response body and may hold anything. + LOGGER.warning( + "WaitTimeFromResponse: `wait_time_in_seconds` resolved to a non-numeric " + f"{type(evaluated_wait_time).__name__}; deferring to the next backoff strategy." + ) + return None + if not math.isfinite(wait_time) or wait_time < 0: + # A retry time that has already passed gives a negative wait, which is not a wait to + # take: the next strategy decides, as `WaitUntilTimeFromHeader` does for a past + # timestamp. Returned as-is, a negative value would reach `time.sleep` and fail the + # stream. A zero wait is returned as 0, but `DefaultErrorHandler` treats it as no + # answer and moves on to the next strategy, as it does for `WaitTimeFromHeader`. + return None + max_waiting_time = evaluate_max_waiting_time(self._max_waiting_time_in_seconds, self.config) + if max_waiting_time is not None and wait_time and wait_time >= max_waiting_time: + raise AirbyteTracedException( + internal_message=( + f"Rate limit wait time {wait_time}s is greater than or equal to the " + f"maximum of {max_waiting_time}s this stream is allowed to wait. " + f"Stopping the stream..." + ), + message="The rate limit wait time is longer than the connector is allowed to wait.", + failure_type=FailureType.transient_error, + ) + + return wait_time diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index e54f23e8e8..d9caa54459 100644 --- a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +++ b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py @@ -189,6 +189,7 @@ ConstantBackoffStrategy, ExponentialBackoffStrategy, WaitTimeFromHeaderBackoffStrategy, + WaitTimeFromResponseBackoffStrategy, WaitUntilTimeFromHeaderBackoffStrategy, ) from airbyte_cdk.sources.declarative.requesters.http_job_repository import AsyncHttpJobRepository @@ -2786,6 +2787,44 @@ def test_create_requester_with_backoff_jitter( assert backoff_strategy.jitter_range_in_seconds == expected_jitter_range +def test_create_requester_with_wait_time_from_response_strategy(): + content = """ +requester: + type: HttpRequester + path: "/v3/marketing/lists" + url_base: "https://api.sendgrid.com" + error_handler: + backoff_strategies: + - type: WaitTimeFromResponse + wait_time_in_seconds: "{{ response[parameters.wait_field] }}" + max_waiting_time_in_seconds: 30 + $parameters: + wait_field: retry_after_seconds + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + requester_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["requester"], {} + ) + + requester = factory.create_component( + model_type=HttpRequesterModel, + component_definition=requester_manifest, + config=input_config, + name="name", + decoder=None, + ) + + assert isinstance(requester.error_handler, DefaultErrorHandler) + backoff_strategy = requester.error_handler.backoff_strategies[0] + assert isinstance(backoff_strategy, WaitTimeFromResponseBackoffStrategy) + assert backoff_strategy.max_waiting_time_in_seconds == 30 + rate_limited = requests.Response() + rate_limited.status_code = 429 + rate_limited._content = b'{"retry_after_seconds": 12}' + assert backoff_strategy.backoff_time(rate_limited, 1) == 12 + + @pytest.mark.parametrize( "backoff_strategy_model, backoff_strategy_arguments", [ diff --git a/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_wait_time_from_response.py b/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_wait_time_from_response.py new file mode 100644 index 0000000000..51829a4d85 --- /dev/null +++ b/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_wait_time_from_response.py @@ -0,0 +1,367 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import json +import logging +from unittest.mock import patch + +import pytest +import requests +from freezegun import freeze_time + +from airbyte_cdk.models import FailureType +from airbyte_cdk.sources.declarative.requesters.error_handlers.backoff_strategies.wait_time_from_response_backoff_strategy import ( + WaitTimeFromResponseBackoffStrategy, +) +from airbyte_cdk.sources.declarative.requesters.error_handlers.default_error_handler import ( + DefaultErrorHandler, +) +from airbyte_cdk.sources.streams.http.http_client import HttpClient +from airbyte_cdk.utils import AirbyteTracedException + +GMAIL_WAIT_TIME_EXPRESSION = r"{{ max(1, timestamp(response.error.message | regex_search('Retry after (\\S+Z)')) - now_utc().timestamp()) if (response is mapping and response.get('error') is mapping and ((response.error.message or '') | string | regex_search('Retry after (\\S+Z)'))) else '' }}" + + +def test_wait_time_from_response_interpolates_json_body(): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.retry_in }}", parameters={}, config={} + ) + + assert strategy.backoff_time(_response({"retry_in": "12.5"}), 1) == 12.5 + + +def test_wait_time_from_response_interpolates_headers(): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ headers['Retry-After'] }}", parameters={}, config={} + ) + + assert strategy.backoff_time(_response({}, headers={"Retry-After": "15"}), 1) == 15.0 + + +@pytest.mark.parametrize("empty_value", ["", " "]) +def test_wait_time_from_response_returns_none_for_empty_result(empty_value): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.get('retry_in', '') }}", parameters={}, config={} + ) + + assert strategy.backoff_time(_response({"retry_in": empty_value}), 1) is None + + +def test_wait_time_from_response_returns_none_for_missing_key(): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.get('retry_in', '') }}", parameters={}, config={} + ) + + assert strategy.backoff_time(_response({}), 1) is None + + +def test_wait_time_from_response_returns_none_for_non_json_body(): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.get('retry_in', '') }}", parameters={}, config={} + ) + + assert strategy.backoff_time(_response_body(b"not JSON"), 1) is None + + +@pytest.mark.parametrize( + "response_or_exception", + [None, requests.RequestException()], + ids=["none", "exception"], +) +def test_wait_time_from_response_returns_none_for_non_response_input(response_or_exception): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.get('retry_in', '') }}", parameters={}, config={} + ) + + assert strategy.backoff_time(response_or_exception, 1) is None + + +@pytest.mark.parametrize("wait_time", [10, 11], ids=["equal_to_cap", "above_cap"]) +def test_wait_time_from_response_at_or_above_cap_raises_transient_error(wait_time): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.retry_in }}", + parameters={}, + config={}, + max_waiting_time_in_seconds=10, + ) + + with pytest.raises(AirbyteTracedException) as exception: + strategy.backoff_time(_response({"retry_in": wait_time}), 1) + + assert exception.value.failure_type == FailureType.transient_error + assert ( + exception.value.message + == "The rate limit wait time is longer than the connector is allowed to wait." + ) + assert "greater than or equal to the maximum of 10.0s" in exception.value.internal_message + + +def test_wait_time_from_response_below_cap_is_returned(): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.retry_in }}", + parameters={}, + config={}, + max_waiting_time_in_seconds=10, + ) + + assert strategy.backoff_time(_response({"retry_in": 9}), 1) == 9.0 + + +def test_wait_time_from_response_cap_is_interpolated_from_config(): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.retry_in }}", + parameters={}, + config={"max_waiting_time": 20}, + max_waiting_time_in_seconds="{{ config['max_waiting_time'] }}", + ) + + assert strategy.backoff_time(_response({"retry_in": 10}), 1) == 10.0 + + +def test_wait_time_from_response_invalid_cap_fails_at_construction(): + with pytest.raises(AirbyteTracedException) as exception: + WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="10", + parameters={}, + config={"max_waiting_time": "not a number"}, + max_waiting_time_in_seconds="{{ config['max_waiting_time'] }}", + ) + + assert exception.value.failure_type == FailureType.system_error + + +def test_wait_time_from_response_zero_wait_is_not_refused_by_zero_cap(): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.retry_in }}", + parameters={}, + config={}, + max_waiting_time_in_seconds=0, + ) + + assert strategy.backoff_time(_response({"retry_in": 0}), 1) == 0 + + +@freeze_time("2026-10-15T13:00:00Z") +@pytest.mark.parametrize( + "message, expected_wait", + [ + pytest.param( + "User-rate limit exceeded. Retry after 2026-10-15T13:15:00Z", + 900, + id="future_timestamp", + ), + pytest.param( + "User-rate limit exceeded. Retry after 2026-10-15T12:59:59Z", + 1, + id="past_timestamp", + ), + pytest.param("User-rate limit exceeded.", None, id="missing_timestamp"), + ], +) +def test_wait_time_from_response_gmail_expression(message, expected_wait): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds=GMAIL_WAIT_TIME_EXPRESSION, parameters={}, config={} + ) + + actual_wait = strategy.backoff_time(_response({"error": {"message": message}}), 1) + + if expected_wait is None: + assert actual_wait is None + else: + assert actual_wait == pytest.approx(expected_wait) + + +@pytest.mark.parametrize( + "retry_in", + [ + pytest.param(-0.5, id="negative_fraction"), + pytest.param(-30, id="negative"), + pytest.param("nan", id="not_a_number"), + pytest.param("inf", id="infinite"), + ], +) +def test_wait_time_from_response_returns_none_for_negative_or_non_finite_result(retry_in): + """A negative wait would reach `time.sleep` and fail the stream; the next strategy decides.""" + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.retry_in }}", parameters={}, config={} + ) + + assert strategy.backoff_time(_response({"retry_in": retry_in}), 1) is None + + +def test_wait_time_from_response_returns_none_for_non_numeric_result(caplog): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.retry_in }}", parameters={}, config={} + ) + + assert strategy.backoff_time(_response({"retry_in": "soon"}), 1) is None + assert any("non-numeric str" in message for message in caplog.messages) + assert all("soon" not in message for message in caplog.messages) + + +@freeze_time("2026-10-15T13:00:00Z") +def test_wait_time_from_response_past_timestamp_without_a_floor_defers(): + """`retry_at - now` is negative once the time has passed, without the `max(1, ...)` floor.""" + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ timestamp(response.retry_at) - now_utc().timestamp() }}", + parameters={}, + config={}, + ) + + assert strategy.backoff_time(_response({"retry_at": "2026-10-15T12:59:30Z"}), 1) is None + assert strategy.backoff_time(_response({"retry_at": "2026-10-15T13:00:30Z"}), 1) == 30 + + +@pytest.mark.parametrize( + "body", + [ + pytest.param([{"error": {"message": "Retry after 2026-10-15T13:15:00Z"}}], id="list_body"), + pytest.param({"error": "rate limited"}, id="error_is_a_string"), + pytest.param({"error": {}}, id="error_without_message"), + pytest.param({"error": {"message": None}}, id="message_is_null"), + pytest.param(True, id="scalar_body"), + ], +) +def test_wait_time_from_response_gmail_expression_defers_on_other_body_shapes(body): + """The documented example runs on every retryable response, not only Gmail's rate limit.""" + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds=GMAIL_WAIT_TIME_EXPRESSION, parameters={}, config={} + ) + + assert strategy.backoff_time(_response(body), 1) is None + + +def test_wait_time_from_response_returns_none_when_the_expression_fails_on_the_body(caplog): + """An expression that does not fit the body defers to the next strategy instead of failing. + + `.get` on a list raises `UndefinedError` from Jinja, which `InterpolatedString.eval` re-raises + when it evaluates the default (the expression itself) a second time. + """ + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.get('retry') }}", parameters={}, config={} + ) + + assert strategy.backoff_time(_response([{"retry": "soon"}]), 1) is None + assert any("could not be evaluated" in message for message in caplog.messages) + assert all("soon" not in message for message in caplog.messages) + + +def test_wait_time_from_response_a_failed_render_does_not_disable_the_strategy(caplog): + """A `TypeError` while rendering one body must not turn the strategy off for later bodies.""" + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.reset_in - 1 }}", parameters={}, config={} + ) + + assert strategy.backoff_time(_response({"reset_in": "abc"}), 1) is None + assert any("could not be evaluated" in message for message in caplog.messages) + assert strategy.backoff_time(_response({"reset_in": 30}), 1) == 29 + + +@pytest.mark.parametrize( + "expression, body", + [ + pytest.param("{{ 10 / response.n }}", b'{"n": 0}', id="division_by_zero"), + pytest.param("{{ response.n }}", b'{"n": 1' + b"0" * 400 + b"}", id="float_overflow"), + ], +) +def test_wait_time_from_response_returns_none_for_arithmetic_errors(expression, body): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds=expression, parameters={}, config={} + ) + + assert strategy.backoff_time(_response_body(body), 1) is None + + +def test_wait_time_from_response_unreadable_body_falls_back_to_headers(): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ headers['Retry-After'] }}", parameters={}, config={} + ) + + deeply_nested_body = b"[" * 100_000 + assert strategy.backoff_time(_response_body(deeply_nested_body, {"Retry-After": "5"}), 1) == 5 + + +def test_wait_time_from_response_undeclared_variable_fails_at_construction(): + with pytest.raises(AirbyteTracedException) as error: + WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ resposne.retry_in }}", parameters={}, config={} + ) + + assert error.value.failure_type == FailureType.system_error + assert "resposne" in error.value.internal_message + + +def test_wait_time_from_response_interpolates_parameters(): + strategy = WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response[parameters.wait_field] }}", + parameters={"wait_field": "retry_in"}, + config={}, + ) + + assert strategy.backoff_time(_response({"retry_in": 7}), 1) == 7 + + +def test_wait_time_from_response_wait_reaches_the_sleep_in_http_client(requests_mock): + requests_mock.get( + "https://api.example.com/data", + [{"status_code": 429, "json": {"retry_in": 7}}, {"status_code": 200, "json": {}}], + ) + http_client = _http_client( + WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.retry_in }}", parameters={}, config={} + ) + ) + + sleeps = [] + with patch("time.sleep", side_effect=sleeps.append): + _, response = http_client.send_request( + http_method="GET", url="https://api.example.com/data", request_kwargs={} + ) + + assert response.status_code == 200 + # `user_defined_backoff_handler` adds one second to cover fractions of a second; the `backoff` + # library's own wait is 0. + assert max(sleeps) == 8 + + +def test_wait_time_from_response_wait_above_cap_stops_the_request_in_http_client(requests_mock): + requests_mock.get("https://api.example.com/data", status_code=429, json={"retry_in": 3600}) + http_client = _http_client( + WaitTimeFromResponseBackoffStrategy( + wait_time_in_seconds="{{ response.retry_in }}", + parameters={}, + config={}, + max_waiting_time_in_seconds=60, + ) + ) + + with patch("time.sleep") as sleep, pytest.raises(AirbyteTracedException) as error: + http_client.send_request( + http_method="GET", url="https://api.example.com/data", request_kwargs={} + ) + + assert error.value.failure_type == FailureType.transient_error + sleep.assert_not_called() + + +def _http_client(backoff_strategy): + return HttpClient( + name="test", + logger=logging.getLogger("test"), + error_handler=DefaultErrorHandler(config={}, parameters={}), + backoff_strategy=[backoff_strategy], + ) + + +def _response(json_body, headers=None): + return _response_body(json.dumps(json_body).encode(), headers=headers) + + +def _response_body(body, headers=None): + response = requests.Response() + response.status_code = 429 + response.headers = headers or {} + response._content = body + response.encoding = "utf-8" + return response diff --git a/unit_tests/sources/declarative/requesters/error_handlers/test_default_error_handler.py b/unit_tests/sources/declarative/requesters/error_handlers/test_default_error_handler.py index 526e5760fc..09d25f5a47 100644 --- a/unit_tests/sources/declarative/requesters/error_handlers/test_default_error_handler.py +++ b/unit_tests/sources/declarative/requesters/error_handlers/test_default_error_handler.py @@ -13,6 +13,9 @@ from airbyte_cdk.sources.declarative.requesters.error_handlers.backoff_strategies.exponential_backoff_strategy import ( ExponentialBackoffStrategy, ) +from airbyte_cdk.sources.declarative.requesters.error_handlers.backoff_strategies.wait_time_from_response_backoff_strategy import ( + WaitTimeFromResponseBackoffStrategy, +) from airbyte_cdk.sources.declarative.requesters.error_handlers.default_error_handler import ( DefaultErrorHandler, HttpResponseFilter, @@ -29,6 +32,22 @@ SOME_BACKOFF_TIME = 60 +def test_default_error_handler_falls_back_when_response_backoff_is_empty(): + response_mock = create_response(429, json_body={"retry_in": ""}) + error_handler = DefaultErrorHandler( + config={}, + parameters={}, + backoff_strategies=[ + WaitTimeFromResponseBackoffStrategy( + "{{ response.get('retry_in', '') }}", config={}, parameters={} + ), + ExponentialBackoffStrategy(factor=5, config={}, parameters={}), + ], + ) + + assert error_handler.backoff_time(response_mock, attempt_count=2) == 20 + + @pytest.mark.parametrize( "test_name, http_status_code, expected_error_resolution", [