diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 7068105cb..a77bc34ea 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -391,6 +391,7 @@ definitions: anyOf: - "$ref": "#/definitions/CompositeErrorHandler" - "$ref": "#/definitions/DefaultErrorHandler" + - "$ref": "#/definitions/CustomErrorHandler" $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 49b48a232..915f942d0 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2149,10 +2149,12 @@ class ConfigAddFields(BaseModel): class CompositeErrorHandler(BaseModel): type: Literal["CompositeErrorHandler"] - error_handlers: List[Union[CompositeErrorHandler, DefaultErrorHandler]] = Field( - ..., - description="List of error handlers to iterate on to determine how to handle a failed response.", - title="Error Handlers", + error_handlers: List[Union[CompositeErrorHandler, DefaultErrorHandler, CustomErrorHandler]] = ( + Field( + ..., + description="List of error handlers to iterate on to determine how to handle a failed response.", + title="Error Handlers", + ) ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") 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 5c4d2c8d5..903eeb14c 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 @@ -2027,6 +2027,44 @@ def test_create_composite_error_handler(): assert error_handler_1.response_filters[0].action == ResponseAction.RETRY +def test_create_composite_error_handler_with_custom_error_handler(): + content = """ + error_handler: + type: "CompositeErrorHandler" + error_handlers: + - type: "CustomErrorHandler" + class_name: "unit_tests.sources.declarative.parsers.testing_components.TestingCustomErrorHandler" + - response_filters: + - http_codes: [ 429 ] + action: RETRY + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + error_handler_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["error_handler"], {} + ) + + error_handler = factory.create_component( + model_type=CompositeErrorHandlerModel, + component_definition=error_handler_manifest, + config=input_config, + ) + + assert isinstance(error_handler, CompositeErrorHandler) + assert len(error_handler.error_handlers) == 2 + + # First error handler should be a custom error handler + error_handler_0 = error_handler.error_handlers[0] + assert error_handler_0.__class__.__name__ == "TestingCustomErrorHandler" + + # Second error handler should be a default error handler + error_handler_1 = error_handler.error_handlers[1] + assert isinstance(error_handler_1, DefaultErrorHandler) + assert isinstance(error_handler_1.response_filters[0], HttpResponseFilter) + assert error_handler_1.response_filters[0].http_codes == {429} + assert error_handler_1.response_filters[0].action == ResponseAction.RETRY + + # This might be a better test for the manifest transformer but also worth testing end-to-end here as well def test_config_with_defaults(): content = """ diff --git a/unit_tests/sources/declarative/parsers/testing_components.py b/unit_tests/sources/declarative/parsers/testing_components.py index 88316b521..6269acb8e 100644 --- a/unit_tests/sources/declarative/parsers/testing_components.py +++ b/unit_tests/sources/declarative/parsers/testing_components.py @@ -90,6 +90,15 @@ def migrate(self, stream_state: Mapping[str, Any]) -> Mapping[str, Any]: return stream_state +@dataclass +class TestingCustomErrorHandler(DefaultErrorHandler): + """ + A test class based on DefaultErrorHandler used for testing manifests that use custom error handlers. + """ + + __test__: ClassVar[bool] = False # Tell Pytest this is not a Pytest class, despite its name + + @dataclass class TestingRequester(HttpRequester): request_parameters: Optional[RequestInput] = None