allow for specific parameters to be passed to custom components

Author: maxime.c

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -752,7 +752,7 @@ def _init_mappings(self) -> None:
             OAuthAuthenticatorModel: self.create_oauth_authenticator,
             OffsetIncrementModel: self.create_offset_increment,
             PageIncrementModel: self.create_page_increment,
-            ParentStreamConfigModel: self.create_parent_stream_config,
+            ParentStreamConfigModel: self._create_message_repository_substream_wrapper,
             PredicateValidatorModel: self.create_predicate_validator,
             PropertiesFromEndpointModel: self.create_properties_from_endpoint,
             PropertyChunkingModel: self.create_property_chunking,
@@ -1748,7 +1748,7 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) ->
 
             if self._is_component(model_value):
                 model_args[model_field] = self._create_nested_component(
-                    model, model_field, model_value, config
+                    model, model_field, model_value, config, **kwargs,
                 )
             elif isinstance(model_value, list):
                 vals = []
@@ -1760,7 +1760,7 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) ->
                         if derived_type:
                             v["type"] = derived_type
                     if self._is_component(v):
-                        vals.append(self._create_nested_component(model, model_field, v, config))
+                        vals.append(self._create_nested_component(model, model_field, v, config, **kwargs,))
                     else:
                         vals.append(v)
                 model_args[model_field] = vals
@@ -1850,7 +1850,7 @@ def _extract_missing_parameters(error: TypeError) -> List[str]:
             return []
 
     def _create_nested_component(
-        self, model: Any, model_field: str, model_value: Any, config: Config
+        self, model: Any, model_field: str, model_value: Any, config: Config, **kwargs: Any
     ) -> Any:
         type_name = model_value.get("type", None)
         if not type_name:
@@ -1875,8 +1875,11 @@ def _create_nested_component(
                     for kwarg in constructor_kwargs
                     if kwarg in model_parameters
                 }
+                matching_kwargs = {
+                    kwarg: kwargs[kwarg] for kwarg in constructor_kwargs if kwarg in kwargs
+                }
                 return self._create_component_from_model(
-                    model=parsed_model, config=config, **matching_parameters
+                    model=parsed_model, config=config, **(matching_parameters | matching_kwargs)
                 )
             except TypeError as error:
                 missing_parameters = self._extract_missing_parameters(error)
@@ -2871,7 +2874,7 @@ def create_page_increment(
         )
 
     def create_parent_stream_config(
-        self, model: ParentStreamConfigModel, config: Config, **kwargs: Any
+        self, model: ParentStreamConfigModel, config: Config, stream_name: str, **kwargs: Any
     ) -> ParentStreamConfig:
         declarative_stream = self._create_component_from_model(
             model.stream,
@@ -3695,11 +3698,11 @@ def create_substream_partition_router(
         )
 
     def _create_message_repository_substream_wrapper(
-        self, model: ParentStreamConfigModel, config: Config, **kwargs: Any
+        self, model: ParentStreamConfigModel, config: Config, *, stream_name: str, **kwargs: Any
     ) -> Any:
         # getting the parent state
         child_state = self._connector_state_manager.get_stream_state(
-            kwargs["stream_name"], None
+            stream_name, None
         )
 
         # This flag will be used exclusively for StateDelegatingStream when a parent stream is created
@@ -3731,8 +3734,8 @@ def _create_message_repository_substream_wrapper(
             ),
         )
 
-        return substream_factory._create_component_from_model(
-            model=model, config=config, has_parent_state=has_parent_state, **kwargs
+        return substream_factory.create_parent_stream_config(
+            model=model, config=config, stream_name=stream_name, **kwargs
         )
 
     def _instantiate_parent_stream_state_manager(

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -163,6 +163,7 @@
 from airbyte_cdk.sources.declarative.transformations import AddFields, RemoveFields
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
 from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource
+from airbyte_cdk.sources.message.repository import StateFilteringMessageRepository
 from airbyte_cdk.sources.streams.call_rate import MovingWindowCallRatePolicy
 from airbyte_cdk.sources.streams.concurrent.clamping import (
     ClampingEndProvider,
@@ -944,6 +945,58 @@ def test_stream_with_incremental_and_retriever_with_partition_router():
     assert list_stream_slicer._cursor_field.string == "a_key"
 
 
+def test_stream_with_custom_retriever_and_transformations():
+    content = """
+a_stream:
+  type: DeclarativeStream
+  primary_key: id
+  schema_loader:
+    type: InlineSchemaLoader
+    schema:
+      $schema: "http://json-schema.org/draft-07/schema"
+      type: object
+      properties:
+        id:
+          type: string
+  retriever:
+    type: CustomRetriever
+    class_name: unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever
+    name: "{{ parameters['name'] }}"
+    decoder:
+      type: JsonDecoder
+    requester:
+      type: HttpRequester
+      name: "{{ parameters['name'] }}"
+      url_base: "https://api.sendgrid.com/v3/"
+      http_method: "GET"
+    record_selector:
+      type: RecordSelector
+      extractor:
+        type: DpathExtractor
+        field_path: ["records"]
+  transformations:
+    - type: AddFields
+      fields:
+        - path: ["extra"]
+          value: "{{ response.to_add }}"
+  $parameters:
+   name: a_stream
+"""
+
+    parsed_manifest = YamlDeclarativeSource._parse(content)
+    resolved_manifest = resolver.preprocess_manifest(parsed_manifest)
+    stream_manifest = transformer.propagate_types_and_parameters(
+        "", resolved_manifest["a_stream"], {}
+    )
+
+    stream = factory.create_component(
+        model_type=DeclarativeStreamModel, component_definition=stream_manifest, config=input_config
+    )
+
+    assert isinstance(stream, DefaultStream)
+    assert get_retriever(stream).record_selector.transformations
+
+
 @pytest.mark.parametrize(
     "use_legacy_state",
     [
@@ -2053,11 +2106,12 @@ def test_custom_components_do_not_contain_extra_fields():
     }
 
     custom_substream_partition_router = factory.create_component(
-        CustomPartitionRouterModel, custom_substream_partition_router_manifest, input_config
+        CustomPartitionRouterModel, custom_substream_partition_router_manifest, input_config, stream_name="child_stream_name",
     )
     assert isinstance(custom_substream_partition_router, TestingCustomSubstreamPartitionRouter)
 
     assert len(custom_substream_partition_router.parent_stream_configs) == 1
+    assert isinstance(custom_substream_partition_router.parent_stream_configs[0].stream.cursor._message_repository, StateFilteringMessageRepository)
     assert custom_substream_partition_router.parent_stream_configs[0].parent_key.eval({}) == "id"
     assert (
         custom_substream_partition_router.parent_stream_configs[0].partition_field.eval({})
@@ -2120,7 +2174,7 @@ def test_parse_custom_component_fields_if_subcomponent():
     }
 
     custom_substream_partition_router = factory.create_component(
-        CustomPartitionRouterModel, custom_substream_partition_router_manifest, input_config
+        CustomPartitionRouterModel, custom_substream_partition_router_manifest, input_config, stream_name="child_stream_name"
     )
     assert isinstance(custom_substream_partition_router, TestingCustomSubstreamPartitionRouter)
     assert custom_substream_partition_router.custom_field == "here"

unit_tests/sources/declarative/parsers/testing_components.py

@@ -13,6 +13,7 @@
     DefaultPaginator,
     PaginationStrategy,
 )
+from airbyte_cdk.sources.declarative.retrievers import SimpleRetriever
 
 
 @dataclass
@@ -43,3 +44,8 @@ class TestingCustomSubstreamPartitionRouter(SubstreamPartitionRouter):
 
     custom_field: str
     custom_pagination_strategy: PaginationStrategy
+
+
+@dataclass
+class TestingCustomRetriever(SimpleRetriever):
+    pass