diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 6fbda8f006..45ecd98dc1 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -482,6 +482,7 @@ definitions: anyOf: - type: number title: Number of seconds + minimum: 0 - type: string title: Interpolated Value interpolation_context: @@ -490,6 +491,13 @@ definitions: - 30 - 30.5 - "{{ config['backoff_time'] }}" + jitter_range_in_seconds: + title: Jitter Range + description: Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between backoff_time_in_seconds and backoff_time_in_seconds + (jitter_range_in_seconds * 2), so jitter only increases the base backoff. + type: number + minimum: 0 + examples: + - 15 $parameters: type: object additionalProperties: true @@ -2033,6 +2041,13 @@ definitions: - 5 - 5.5 - "10" + jitter_range_in_seconds: + title: Jitter Range + description: Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between computed_backoff and computed_backoff + (jitter_range_in_seconds * 2), so jitter only increases the computed backoff. + type: number + minimum: 0 + examples: + - 2 $parameters: type: object additionalProperties: true diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 931cef7f1b..f14c04da40 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -104,6 +104,13 @@ class ConstantBackoffStrategy(BaseModel): examples=[30, 30.5, "{{ config['backoff_time'] }}"], title="Backoff Time", ) + jitter_range_in_seconds: Optional[float] = Field( + None, + description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between backoff_time_in_seconds and backoff_time_in_seconds + (jitter_range_in_seconds * 2), so jitter only increases the base backoff.", + examples=[15], + ge=0, + title="Jitter Range", + ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -512,6 +519,13 @@ class ExponentialBackoffStrategy(BaseModel): examples=[5, 5.5, "10"], title="Factor", ) + jitter_range_in_seconds: Optional[float] = Field( + None, + description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between computed_backoff and computed_backoff + (jitter_range_in_seconds * 2), so jitter only increases the computed backoff.", + examples=[2], + ge=0, + title="Jitter Range", + ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -3077,6 +3091,7 @@ class AsyncRetriever(BaseModel): failed_retry_wait_time_in_seconds: Optional[Union[int, str]] = Field( None, description="Time in seconds to wait before retrying a failed async job. Only applies to jobs that ran on the API side and reported a FAILED status (e.g. report generation failed due to a cooldown). Creation failures (HTTP errors when starting a job, such as 429s) and TIMED_OUT jobs are retried immediately and are not affected by this setting. When set, the orchestrator defers retry of real failed jobs until the wait time has elapsed, without blocking other jobs.", + ge=1, ) download_target_requester: Optional[Union[HttpRequester, CustomRequester]] = Field( None, 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 896d82a0f6..8bcb9e2dc7 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1752,12 +1752,19 @@ def create_concurrent_cursor_from_perpartition_cursor( def create_constant_backoff_strategy( model: ConstantBackoffStrategyModel, config: Config, **kwargs: Any ) -> ConstantBackoffStrategy: + ModelToComponentFactory._validate_jitter_range(model.jitter_range_in_seconds) return ConstantBackoffStrategy( backoff_time_in_seconds=model.backoff_time_in_seconds, + jitter_range_in_seconds=model.jitter_range_in_seconds, config=config, parameters=model.parameters or {}, ) + @staticmethod + def _validate_jitter_range(jitter_range_in_seconds: Optional[float]) -> None: + if jitter_range_in_seconds is not None and jitter_range_in_seconds < 0: + raise ValueError("jitter_range_in_seconds must be greater than or equal to 0") + def create_cursor_pagination( self, model: CursorPaginationModel, config: Config, decoder: Decoder, **kwargs: Any ) -> CursorPaginationStrategy: @@ -2432,8 +2439,12 @@ def create_response_to_file_extractor( def create_exponential_backoff_strategy( model: ExponentialBackoffStrategyModel, config: Config ) -> ExponentialBackoffStrategy: + ModelToComponentFactory._validate_jitter_range(model.jitter_range_in_seconds) return ExponentialBackoffStrategy( - factor=model.factor or 5, parameters=model.parameters or {}, config=config + factor=model.factor or 5, + jitter_range_in_seconds=model.jitter_range_in_seconds, + parameters=model.parameters or {}, + config=config, ) @staticmethod diff --git a/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/constant_backoff_strategy.py b/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/constant_backoff_strategy.py index 26c7c76736..931183b2ef 100644 --- a/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/constant_backoff_strategy.py +++ b/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/constant_backoff_strategy.py @@ -2,8 +2,9 @@ # Copyright (c) 2023 Airbyte, Inc., all rights reserved. # +import random from dataclasses import InitVar, dataclass -from typing import Any, Mapping, Optional, Union +from typing import Any, Mapping, Optional, Union, cast import requests @@ -24,6 +25,7 @@ class ConstantBackoffStrategy(BackoffStrategy): backoff_time_in_seconds: Union[float, InterpolatedString, str] parameters: InitVar[Mapping[str, Any]] config: Config + jitter_range_in_seconds: Optional[float] = None def __post_init__(self, parameters: Mapping[str, Any]) -> None: if not isinstance(self.backoff_time_in_seconds, InterpolatedString): @@ -42,4 +44,10 @@ def backoff_time( response_or_exception: Optional[Union[requests.Response, requests.RequestException]], attempt_count: int, ) -> Optional[float]: - return self.backoff_time_in_seconds.eval(self.config) # type: ignore # backoff_time_in_seconds is always cast to an interpolated string + backoff_time = float( + cast(InterpolatedString, self.backoff_time_in_seconds).eval(self.config) + ) + if self.jitter_range_in_seconds is None: + return backoff_time + + return random.uniform(backoff_time, backoff_time + (self.jitter_range_in_seconds * 2)) diff --git a/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/exponential_backoff_strategy.py b/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/exponential_backoff_strategy.py index cdd1fe6506..dd96ff18b7 100644 --- a/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/exponential_backoff_strategy.py +++ b/airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/exponential_backoff_strategy.py @@ -2,6 +2,7 @@ # Copyright (c) 2023 Airbyte, Inc., all rights reserved. # +import random from dataclasses import InitVar, dataclass from typing import Any, Mapping, Optional, Union @@ -24,6 +25,7 @@ class ExponentialBackoffStrategy(BackoffStrategy): parameters: InitVar[Mapping[str, Any]] config: Config factor: Union[float, InterpolatedString, str] = 5 + jitter_range_in_seconds: Optional[float] = None def __post_init__(self, parameters: Mapping[str, Any]) -> None: if not isinstance(self.factor, InterpolatedString): @@ -35,11 +37,15 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: @property def _retry_factor(self) -> float: - return self._factor.eval(self.config) # type: ignore # factor is always cast to an interpolated string + return float(self._factor.eval(self.config)) def backoff_time( self, response_or_exception: Optional[Union[requests.Response, requests.RequestException]], attempt_count: int, ) -> Optional[float]: - return self._retry_factor * 2**attempt_count # type: ignore # factor is always cast to an interpolated string + backoff_time = float(self._retry_factor * 2**attempt_count) + if self.jitter_range_in_seconds is None: + return backoff_time + + return random.uniform(backoff_time, backoff_time + (self.jitter_range_in_seconds * 2)) 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 9fd1cb58b0..5b39b069e5 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 @@ -89,9 +89,15 @@ from airbyte_cdk.sources.declarative.models import ( SubstreamPartitionRouter as SubstreamPartitionRouterModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + ConstantBackoffStrategy as ConstantBackoffStrategyModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( CustomRequester as CustomRequesterModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + ExponentialBackoffStrategy as ExponentialBackoffStrategyModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( OffsetIncrement as OffsetIncrementModel, ) @@ -1754,6 +1760,123 @@ def test_create_requester(test_name, error_handler, expected_backoff_strategy_ty ) +@pytest.mark.parametrize( + "backoff_strategy_yaml, expected_backoff_strategy_type, expected_jitter_range", + [ + pytest.param( + """ + error_handler: + backoff_strategies: + - type: "ConstantBackoffStrategy" + backoff_time_in_seconds: 60 + jitter_range_in_seconds: 7 + """, + ConstantBackoffStrategy, + 7, + id="constant_backoff_strategy", + ), + pytest.param( + """ + error_handler: + backoff_strategies: + - type: "ExponentialBackoffStrategy" + factor: 5 + jitter_range_in_seconds: 15 + """, + ExponentialBackoffStrategy, + 15, + id="exponential_backoff_strategy", + ), + ], +) +def test_create_requester_with_backoff_jitter( + backoff_strategy_yaml, expected_backoff_strategy_type, expected_jitter_range +): + content = f""" +requester: + type: HttpRequester + path: "/v3/marketing/lists" + url_base: "https://api.sendgrid.com" + {backoff_strategy_yaml} + """ + 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) + assert len(requester.error_handler.backoff_strategies) == 1 + backoff_strategy = requester.error_handler.backoff_strategies[0] + assert isinstance(backoff_strategy, expected_backoff_strategy_type) + assert backoff_strategy.jitter_range_in_seconds == expected_jitter_range + + +@pytest.mark.parametrize( + "backoff_strategy_model, backoff_strategy_arguments", + [ + pytest.param( + ConstantBackoffStrategyModel, + { + "type": "ConstantBackoffStrategy", + "backoff_time_in_seconds": 60, + }, + id="constant_backoff_strategy", + ), + pytest.param( + ExponentialBackoffStrategyModel, + { + "type": "ExponentialBackoffStrategy", + "factor": 5, + }, + id="exponential_backoff_strategy", + ), + ], +) +def test_backoff_jitter_schema_validation(backoff_strategy_model, backoff_strategy_arguments): + backoff_strategy_model(**backoff_strategy_arguments, jitter_range_in_seconds=0) + + with pytest.raises(ValidationError, match="jitter_range_in_seconds"): + backoff_strategy_model( + **backoff_strategy_arguments, + jitter_range_in_seconds="{{ config['backoff_jitter'] }}", + ) + + with pytest.raises(ValidationError, match="jitter_range_in_seconds"): + backoff_strategy_model(**backoff_strategy_arguments, jitter_range_in_seconds=-1) + + +@pytest.mark.parametrize( + "backoff_strategy_model", + [ + pytest.param( + ConstantBackoffStrategyModel( + type="ConstantBackoffStrategy", backoff_time_in_seconds=60 + ), + id="constant_backoff_strategy", + ), + pytest.param( + ExponentialBackoffStrategyModel(type="ExponentialBackoffStrategy", factor=5), + id="exponential_backoff_strategy", + ), + ], +) +def test_create_backoff_strategy_with_negative_jitter_raises_error(backoff_strategy_model): + # Verify factory validation catches negative jitter even if Pydantic validation is bypassed. + backoff_strategy_model.__dict__["jitter_range_in_seconds"] = -1 + + with pytest.raises(ValueError, match="jitter_range_in_seconds"): + factory._create_component_from_model(backoff_strategy_model, config=input_config) + + def test_create_request_with_legacy_session_authenticator(): content = """ requester: diff --git a/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_constant_backoff.py b/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_constant_backoff.py index 04d54951e7..7fae93fa15 100644 --- a/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_constant_backoff.py +++ b/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_constant_backoff.py @@ -16,29 +16,67 @@ @pytest.mark.parametrize( - "test_name, attempt_count, backofftime, expected_backoff_time", + "test_name, attempt_count, backofftime, jitter_range, expected_backoff_time", [ - ("test_constant_backoff_first_attempt", 1, BACKOFF_TIME, BACKOFF_TIME), - ("test_constant_backoff_first_attempt_float", 1, 6.7, 6.7), - ("test_constant_backoff_attempt_round_float", 1.0, 6.7, 6.7), - ("test_constant_backoff_attempt_round_float", 1.5, 6.7, 6.7), - ("test_constant_backoff_first_attempt_round_float", 1, 10.0, BACKOFF_TIME), - ("test_constant_backoff_second_attempt_round_float", 2, 10.0, BACKOFF_TIME), + ("test_constant_backoff_first_attempt", 1, BACKOFF_TIME, None, BACKOFF_TIME), + ("test_constant_backoff_first_attempt_float", 1, 6.7, None, 6.7), + ("test_constant_backoff_attempt_round_float", 1.0, 6.7, None, 6.7), + ("test_constant_backoff_attempt_round_float", 1.5, 6.7, None, 6.7), + ("test_constant_backoff_first_attempt_round_float", 1, 10.0, None, BACKOFF_TIME), + ("test_constant_backoff_second_attempt_round_float", 2, 10.0, None, BACKOFF_TIME), + ("test_constant_backoff_zero_jitter", 2, BACKOFF_TIME, 0, BACKOFF_TIME), ( "test_constant_backoff_from_parameters", 1, "{{ parameters['backoff'] }}", + None, PARAMETERS_BACKOFF_TIME, ), - ("test_constant_backoff_from_config", 1, "{{ config['backoff'] }}", CONFIG_BACKOFF_TIME), + ( + "test_constant_backoff_from_config", + 1, + "{{ config['backoff'] }}", + None, + CONFIG_BACKOFF_TIME, + ), ], ) -def test_constant_backoff(test_name, attempt_count, backofftime, expected_backoff_time): +def test_constant_backoff( + test_name, attempt_count, backofftime, jitter_range, expected_backoff_time +): response_mock = MagicMock() backoff_strategy = ConstantBackoffStrategy( parameters={"backoff": PARAMETERS_BACKOFF_TIME}, backoff_time_in_seconds=backofftime, - config={"backoff": CONFIG_BACKOFF_TIME}, + jitter_range_in_seconds=jitter_range, + config={"backoff": CONFIG_BACKOFF_TIME, "jitter": 0}, ) backoff = backoff_strategy.backoff_time(response_mock, attempt_count=attempt_count) assert backoff == expected_backoff_time + + +@pytest.mark.parametrize( + "backofftime, jitter_range, expected_lower_bound, expected_upper_bound", + [ + pytest.param(60, 15, 60, 90, id="base_backoff_floor"), + pytest.param(10, 30, 10, 70, id="large_jitter"), + pytest.param(0, 5, 0, 10, id="zero_base"), + ], +) +def test_constant_backoff_with_jitter_bounds( + backofftime, jitter_range, expected_lower_bound, expected_upper_bound +): + response_mock = MagicMock() + backoff_strategy = ConstantBackoffStrategy( + parameters={}, + backoff_time_in_seconds=backofftime, + jitter_range_in_seconds=jitter_range, + config={}, + ) + + backoff_times = [ + backoff_strategy.backoff_time(response_mock, attempt_count=1) for _ in range(2000) + ] + + assert all(expected_lower_bound <= backoff <= expected_upper_bound for backoff in backoff_times) + assert len(set(backoff_times)) > 1 diff --git a/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_exponential_backoff.py b/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_exponential_backoff.py index 166afeab66..67a561a829 100644 --- a/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_exponential_backoff.py +++ b/unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_exponential_backoff.py @@ -15,18 +15,22 @@ @pytest.mark.parametrize( - "test_name, attempt_count, factor, expected_backoff_time", + "test_name, attempt_count, factor, jitter_range, expected_backoff_time", [ - ("test_exponential_backoff_first_attempt", 1, 5, 10), - ("test_exponential_backoff_second_attempt", 2, 5, 20), - ("test_exponential_backoff_from_parameters", 2, "{{parameters['backoff']}}", 20), - ("test_exponential_backoff_from_config", 2, "{{config['backoff']}}", 20), + ("test_exponential_backoff_first_attempt", 1, 5, None, 10), + ("test_exponential_backoff_second_attempt", 2, 5, None, 20), + ("test_exponential_backoff_zero_jitter", 2, 5, 0, 20), + ("test_exponential_backoff_from_parameters", 2, "{{parameters['backoff']}}", None, 20), + ("test_exponential_backoff_from_config", 2, "{{config['backoff']}}", None, 20), ], ) -def test_exponential_backoff(test_name, attempt_count, factor, expected_backoff_time): +def test_exponential_backoff(test_name, attempt_count, factor, jitter_range, expected_backoff_time): response_mock = MagicMock() backoff_strategy = ExponentialBackoffStrategy( - factor=factor, parameters=parameters, config=config + factor=factor, + jitter_range_in_seconds=jitter_range, + parameters=parameters, + config={**config, "jitter": 0}, ) backoff = backoff_strategy.backoff_time(response_mock, attempt_count=attempt_count) assert backoff == expected_backoff_time @@ -37,3 +41,30 @@ def test_exponential_backoff_default(): backoff_strategy = ExponentialBackoffStrategy(parameters=parameters, config=config) backoff = backoff_strategy.backoff_time(response_mock, attempt_count=3) assert backoff == 40 + + +@pytest.mark.parametrize( + "attempt_count, factor, jitter_range, expected_lower_bound, expected_upper_bound", + [ + pytest.param(2, 5, 5, 20, 30, id="base_backoff_floor"), + pytest.param(1, 2, 10, 4, 24, id="large_jitter"), + ], +) +def test_exponential_backoff_with_jitter_bounds( + attempt_count, factor, jitter_range, expected_lower_bound, expected_upper_bound +): + response_mock = MagicMock() + backoff_strategy = ExponentialBackoffStrategy( + factor=factor, + jitter_range_in_seconds=jitter_range, + parameters=parameters, + config=config, + ) + + backoff_times = [ + backoff_strategy.backoff_time(response_mock, attempt_count=attempt_count) + for _ in range(2000) + ] + + assert all(expected_lower_bound <= backoff <= expected_upper_bound for backoff in backoff_times) + assert len(set(backoff_times)) > 1