diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 82e345ed20..49a005afba 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -1874,13 +1874,18 @@ definitions: - "$ref": "#/definitions/CustomBackoffStrategy" max_retries: title: Max Retry Count - description: The maximum number of time to retry a retryable request before giving up and failing. - type: integer + description: The maximum number of times to retry a retryable request before giving up and failing. Can be a hardcoded integer or a string interpolated from the connector config. + anyOf: + - type: integer + - type: string + interpolation_context: + - config default: 5 examples: - 5 - 0 - 10 + - "{{ config['max_retries_on_throttle'] }}" response_filters: title: Response Filters description: List of response filters to iterate on when deciding how to handle an error. When using an array of multiple filters, the filters will be applied sequentially and the response will be selected if it matches any of the filter's predicate. diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 8bd7b11463..4415855277 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2045,10 +2045,10 @@ class DefaultErrorHandler(BaseModel): description="List of backoff strategies to use to determine how long to wait before retrying a retryable request.", title="Backoff Strategies", ) - max_retries: Optional[int] = Field( + max_retries: Optional[Union[int, str]] = Field( 5, - description="The maximum number of time to retry a retryable request before giving up and failing.", - examples=[5, 0, 10], + description="The maximum number of times to retry a retryable request before giving up and failing. Can be a hardcoded integer or a string interpolated from the connector config.", + examples=[5, 0, 10, "{{ config['max_retries_on_throttle'] }}"], title="Max Retry Count", ) response_filters: Optional[List[HttpResponseFilter]] = Field( diff --git a/airbyte_cdk/sources/declarative/requesters/error_handlers/default_error_handler.py b/airbyte_cdk/sources/declarative/requesters/error_handlers/default_error_handler.py index b70ceaaeb2..1ab2f5cf2a 100644 --- a/airbyte_cdk/sources/declarative/requesters/error_handlers/default_error_handler.py +++ b/airbyte_cdk/sources/declarative/requesters/error_handlers/default_error_handler.py @@ -2,11 +2,12 @@ # Copyright (c) 2023 Airbyte, Inc., all rights reserved. # -from dataclasses import InitVar, dataclass, field +from dataclasses import InitVar, dataclass from typing import Any, List, Mapping, MutableMapping, Optional, Union import requests +from airbyte_cdk.sources.declarative.interpolation import InterpolatedString from airbyte_cdk.sources.declarative.requesters.error_handlers.default_http_response_filter import ( DefaultHttpResponseFilter, ) @@ -88,7 +89,10 @@ class DefaultErrorHandler(ErrorHandler): Attributes: response_filters (Optional[List[HttpResponseFilter]]): response filters to iterate on - max_retries (Optional[int]): maximum retry attempts + max_retries (Optional[Union[int, str]]): maximum retry attempts. Either a hardcoded int or + a string that interpolates from the connector config (e.g. + `"{{ config['max_retries_on_throttle'] }}"`). The string variant is evaluated once at + construction time and replaced with the resolved int. backoff_strategies (Optional[List[BackoffStrategy]]): list of backoff strategies to use to determine how long to wait before retrying """ @@ -96,10 +100,12 @@ class DefaultErrorHandler(ErrorHandler): parameters: InitVar[Mapping[str, Any]] config: Config response_filters: Optional[List[HttpResponseFilter]] = None - max_retries: Optional[int] = 5 + # The base class declares max_retries as Optional[int]. We widen the input type to + # also accept a Jinja-interpolatable string (e.g. "{{ config['max_retries_on_throttle'] }}"), + # which is resolved to an int in __post_init__ so the post-construction invariant matches + # the base class contract. + max_retries: Optional[Union[int, str]] = 5 # type: ignore[assignment] max_time: int = 60 * 10 - _max_retries: int = field(init=False, repr=False, default=5) - _max_time: int = field(init=False, repr=False, default=60 * 10) backoff_strategies: Optional[List[BackoffStrategy]] = None def __post_init__(self, parameters: Mapping[str, Any]) -> None: @@ -108,6 +114,18 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: self._last_request_to_attempt_count: MutableMapping[requests.PreparedRequest, int] = {} + if isinstance(self.max_retries, str): + evaluated = InterpolatedString( + string=self.max_retries, default="5", parameters=parameters + ).eval(config=self.config) + try: + self.max_retries = int(evaluated) + except (TypeError, ValueError) as exc: + raise ValueError( + f"DefaultErrorHandler.max_retries did not evaluate to an integer " + f"(got {evaluated!r})" + ) from exc + def interpret_response( self, response_or_exception: Optional[Union[requests.Response, Exception]] ) -> ErrorResolution: 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 bf4e1e3214..526e5760fc 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 @@ -308,3 +308,63 @@ def test_predicate_takes_precedent_over_default_mapped_error(): assert actual_error_resolution.response_action == ResponseAction.FAIL assert actual_error_resolution.failure_type == FailureType.system_error assert actual_error_resolution.error_message == DEFAULT_ERROR_MAPPING.get(404).error_message + + +def test_max_retries_default_when_unspecified(): + error_handler = DefaultErrorHandler(config={}, parameters={}) + assert error_handler.max_retries == 5 + + +def test_max_retries_with_literal_int(): + error_handler = DefaultErrorHandler(config={}, parameters={}, max_retries=10) + assert error_handler.max_retries == 10 + + +def test_max_retries_with_zero(): + error_handler = DefaultErrorHandler(config={}, parameters={}, max_retries=0) + assert error_handler.max_retries == 0 + + +def test_max_retries_interpolated_from_config(): + error_handler = DefaultErrorHandler( + config={"max_retries_on_throttle": 1}, + parameters={}, + max_retries="{{ config['max_retries_on_throttle'] }}", + ) + assert error_handler.max_retries == 1 + + +def test_max_retries_interpolated_from_config_with_jinja_default(): + error_handler = DefaultErrorHandler( + config={}, + parameters={}, + max_retries="{{ config.get('max_retries_on_throttle', 7) }}", + ) + assert error_handler.max_retries == 7 + + +def test_max_retries_interpolated_string_resolving_to_zero(): + error_handler = DefaultErrorHandler( + config={"max_retries_on_throttle": 0}, + parameters={}, + max_retries="{{ config['max_retries_on_throttle'] }}", + ) + assert error_handler.max_retries == 0 + + +def test_max_retries_interpolated_string_with_numeric_string(): + error_handler = DefaultErrorHandler( + config={"max_retries_on_throttle": "3"}, + parameters={}, + max_retries="{{ config['max_retries_on_throttle'] }}", + ) + assert error_handler.max_retries == 3 + + +def test_max_retries_raises_when_interpolation_does_not_resolve_to_int(): + with pytest.raises(ValueError, match="did not evaluate to an integer"): + DefaultErrorHandler( + config={"max_retries_on_throttle": "not-a-number"}, + parameters={}, + max_retries="{{ config['max_retries_on_throttle'] }}", + )