From e78eaffd72f32288ea96a27b644c14cfc55aaaf5 Mon Sep 17 00:00:00 2001 From: Aldo Gonzalez <168454423+aldogonzalez8@users.noreply.github.com> Date: Wed, 8 Jan 2025 16:53:09 -0600 Subject: [PATCH] fix(airbyte-cdk): unable to create custom retriever (#198) --- .../parsers/model_to_component_factory.py | 15 +++++--- .../test_model_to_component_factory.py | 36 +++++++++++++++++++ 2 files changed, 46 insertions(+), 5 deletions(-) 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 7cb04c2a6..bd164abc2 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1574,7 +1574,12 @@ def create_exponential_backoff_strategy( ) def create_http_requester( - self, model: HttpRequesterModel, decoder: Decoder, config: Config, *, name: str + self, + model: HttpRequesterModel, + config: Config, + decoder: Decoder = JsonDecoder(parameters={}), + *, + name: str, ) -> HttpRequester: authenticator = ( self._create_component_from_model( @@ -1990,9 +1995,9 @@ def create_record_selector( config: Config, *, name: str, - transformations: List[RecordTransformation], - decoder: Optional[Decoder] = None, - client_side_incremental_sync: Optional[Dict[str, Any]] = None, + transformations: List[RecordTransformation] | None = None, + decoder: Decoder | None = None, + client_side_incremental_sync: Dict[str, Any] | None = None, **kwargs: Any, ) -> RecordSelector: assert model.schema_normalization is not None # for mypy @@ -2022,7 +2027,7 @@ def create_record_selector( name=name, config=config, record_filter=record_filter, - transformations=transformations, + transformations=transformations or [], schema_normalization=schema_normalization, parameters=model.parameters or {}, ) 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 c50cfd521..c50e9e6e9 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 @@ -2634,6 +2634,42 @@ def test_create_custom_schema_loader(): assert isinstance(component, MyCustomSchemaLoader) +class MyCustomRetriever(SimpleRetriever): + pass + + +def test_create_custom_retriever(): + stream_model = { + "type": "DeclarativeStream", + "retriever": { + "type": "CustomRetriever", + "class_name": "unit_tests.sources.declarative.parsers.test_model_to_component_factory.MyCustomRetriever", + "record_selector": { + "type": "RecordSelector", + "extractor": { + "type": "DpathExtractor", + "field_path": [], + }, + "$parameters": {"name": ""}, + }, + "requester": { + "type": "HttpRequester", + "name": "list", + "url_base": "orange.com", + "path": "/v1/api", + "$parameters": {"name": ""}, + }, + }, + } + + stream = factory.create_component( + model_type=DeclarativeStreamModel, component_definition=stream_model, config=input_config + ) + + assert isinstance(stream, DeclarativeStream) + assert isinstance(stream.retriever, MyCustomRetriever) + + @freezegun.freeze_time("2021-01-01 00:00:00") @pytest.mark.parametrize( "config, manifest, expected",