fix(airbyte-cdk): unable to create custom retriever (#198)

airbytehq · Jan 8, 2025 · e78eaff · e78eaff
1 parent 3344441
commit e78eaff
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 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
@@ -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
@@ -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",