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 97a1f4ce2d..09be2f8bb3 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1932,10 +1932,16 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) -> for key, arg in kwargs.items(): model_args[key] = arg + # A custom component declaring a `decoder` field shares it with its nested subcomponents (requester, record selector, + # paginator) like `create_simple_retriever` does, so the decoder is built first. A decoder passed by the parent still wins. + nested_component_kwargs = dict(kwargs) + fields_in_build_order = sorted(model_args, key=lambda model_field: model_field != "decoder") + # Pydantic is unable to parse a custom component's fields that are subcomponents into models because their fields and types are not # defined in the schema. The fields and types are defined within the Python class implementation. Pydantic can only parse down to # the custom component and this code performs a second parse to convert the sub-fields first into models, then declarative components - for model_field, model_value in model_args.items(): + for model_field in fields_in_build_order: + model_value = model_args[model_field] # If a custom component field doesn't have a type set, we try to use the type hints to infer the type if ( isinstance(model_value, dict) @@ -1954,7 +1960,7 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) -> model_field, model_value, config, - **kwargs, + **nested_component_kwargs, ) elif isinstance(model_value, list): vals = [] @@ -1972,13 +1978,20 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) -> model_field, v, config, - **kwargs, + **nested_component_kwargs, ) ) else: vals.append(v) model_args[model_field] = vals + if ( + model_field == "decoder" + and "decoder" in component_fields + and isinstance(model_args[model_field], Decoder) + ): + nested_component_kwargs.setdefault("decoder", model_args[model_field]) + kwargs = { class_field: model_args[class_field] for class_field in component_fields.keys() @@ -2742,12 +2755,14 @@ def create_http_requester( self, model: HttpRequesterModel, config: Config, - decoder: Decoder = JsonDecoder(parameters={}), query_properties_key: Optional[str] = None, use_cache: Optional[bool] = None, *, name: str, + decoder: Optional[Decoder] = None, ) -> HttpRequester: + if decoder is None: + decoder = JsonDecoder(parameters={}) authenticator = ( self._create_component_from_model( model=model.authenticator, @@ -2804,7 +2819,7 @@ def create_http_requester( message_repository=self._message_repository, use_cache=should_use_cache, decoder=decoder, - stream_response=decoder.is_stream_response() if decoder else False, + stream_response=decoder.is_stream_response(), ) @staticmethod 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 56e1ad76f0..5c0f22de1f 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 @@ -85,6 +85,7 @@ from airbyte_cdk.sources.declarative.models import ( CustomRecordExtractor as CustomRecordExtractorModel, ) +from airbyte_cdk.sources.declarative.models import CustomRetriever as CustomRetrieverModel from airbyte_cdk.sources.declarative.models import CustomSchemaLoader as CustomSchemaLoaderModel from airbyte_cdk.sources.declarative.models import DatetimeBasedCursor as DatetimeBasedCursorModel from airbyte_cdk.sources.declarative.models import DeclarativeStream as DeclarativeStreamModel @@ -279,6 +280,7 @@ from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now, ab_datetime_parse from unit_tests.sources.declarative.parsers.testing_components import ( TestingCustomRetriever, + TestingCustomRetrieverWithDecoder, TestingCustomSubstreamPartitionRouter, TestingSomeComponent, ) @@ -10096,3 +10098,157 @@ def test_create_file_uploader_with_a_combined_download_target_extractor(): ) assert isinstance(file_uploader.download_target_extractor, CombinedExtractor) + + +_CUSTOM_RETRIEVER_WITH_DECODER = ( + "unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetrieverWithDecoder" +) +_CUSTOM_RETRIEVER_WITHOUT_DECODER = ( + "unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever" +) + + +def _create_stream_with_custom_retriever( + retriever_definition: Mapping[str, Any], component_factory: ModelToComponentFactory = factory +) -> DefaultStream: + stream_manifest = transformer.propagate_types_and_parameters( + "", + { + "type": "DeclarativeStream", + "primary_key": "id", + "schema_loader": {"type": "InlineSchemaLoader", "schema": {"type": "object"}}, + "retriever": { + "requester": { + "type": "HttpRequester", + "url": "https://api.test.com/items", + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + **retriever_definition, + }, + "$parameters": {"name": "items"}, + }, + {}, + ) + return component_factory.create_component( + model_type=DeclarativeStreamModel, component_definition=stream_manifest, config=input_config + ) + + +def test_custom_retriever_decoder_is_passed_to_nested_requester_and_record_selector(): + stream = _create_stream_with_custom_retriever( + { + "type": "CustomRetriever", + "class_name": _CUSTOM_RETRIEVER_WITH_DECODER, + "decoder": {"type": "JsonlDecoder"}, + } + ) + + retriever = get_retriever(stream) + assert isinstance(retriever, TestingCustomRetrieverWithDecoder) + assert isinstance(retriever.decoder, CompositeRawDecoder) + assert retriever.decoder.is_stream_response() + assert retriever.requester.decoder is retriever.decoder + assert retriever.requester.stream_response is True + assert retriever.record_selector.extractor.decoder is retriever.decoder + + +def test_custom_retriever_decoder_is_passed_to_nested_paginator(): + stream = _create_stream_with_custom_retriever( + { + "type": "CustomRetriever", + "class_name": _CUSTOM_RETRIEVER_WITH_DECODER, + "decoder": {"type": "GzipDecoder", "decoder": {"type": "JsonDecoder"}}, + "paginator": { + "type": "DefaultPaginator", + "pagination_strategy": { + "type": "CursorPagination", + "cursor_value": "{{ response.next }}", + }, + "$parameters": {"url_base": "https://api.test.com"}, + }, + } + ) + + retriever = get_retriever(stream) + assert retriever.requester.stream_response is True + assert retriever.requester.decoder is retriever.decoder + assert retriever.record_selector.extractor.decoder is retriever.decoder + assert retriever.paginator.decoder.decoder is retriever.decoder + + +def test_custom_retriever_decoder_is_buffered_in_connector_builder(): + stream = _create_stream_with_custom_retriever( + { + "type": "CustomRetriever", + "class_name": _CUSTOM_RETRIEVER_WITH_DECODER, + "decoder": {"type": "JsonlDecoder"}, + }, + component_factory=ModelToComponentFactory(emit_connector_builder_messages=True), + ) + + retriever = get_retriever(stream) + assert retriever.requester.decoder is retriever.decoder + assert retriever.requester.stream_response is False + assert retriever.record_selector.extractor.decoder is retriever.decoder + + +def test_custom_retriever_without_decoder_keeps_json_decoder_defaults(): + stream = _create_stream_with_custom_retriever( + {"type": "CustomRetriever", "class_name": _CUSTOM_RETRIEVER_WITH_DECODER} + ) + + retriever = get_retriever(stream) + assert retriever.decoder is None + assert isinstance(retriever.requester.decoder, JsonDecoder) + assert retriever.requester.stream_response is False + assert isinstance(retriever.record_selector.extractor.decoder, JsonDecoder) + + +def test_custom_retriever_without_decoder_field_does_not_pass_manifest_decoder_to_subcomponents(): + stream = _create_stream_with_custom_retriever( + { + "type": "CustomRetriever", + "class_name": _CUSTOM_RETRIEVER_WITHOUT_DECODER, + "decoder": {"type": "JsonlDecoder"}, + } + ) + + retriever = get_retriever(stream) + assert isinstance(retriever.requester.decoder, JsonDecoder) + assert retriever.requester.stream_response is False + assert isinstance(retriever.record_selector.extractor.decoder, JsonDecoder) + + +def test_custom_retriever_decoder_passed_by_parent_takes_precedence_over_manifest_decoder(): + parent_decoder = JsonDecoder(parameters={}) + + retriever = factory.create_component( + model_type=CustomRetrieverModel, + component_definition={ + "type": "CustomRetriever", + "class_name": _CUSTOM_RETRIEVER_WITH_DECODER, + "decoder": {"type": "JsonlDecoder"}, + "requester": { + "type": "HttpRequester", + "url": "https://api.test.com/items", + "$parameters": {"name": "items"}, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + "$parameters": {"name": "items"}, + }, + }, + config=input_config, + name="items", + primary_key="id", + decoder=parent_decoder, + ) + + assert retriever.decoder is parent_decoder + assert retriever.requester.decoder is parent_decoder + assert retriever.requester.stream_response is False + assert retriever.record_selector.extractor.decoder is parent_decoder diff --git a/unit_tests/sources/declarative/parsers/testing_components.py b/unit_tests/sources/declarative/parsers/testing_components.py index 9d5ecd9642..157b4309ef 100644 --- a/unit_tests/sources/declarative/parsers/testing_components.py +++ b/unit_tests/sources/declarative/parsers/testing_components.py @@ -6,6 +6,7 @@ from typing import Any, ClassVar, List, Mapping, Optional from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator +from airbyte_cdk.sources.declarative.decoders import Decoder from airbyte_cdk.sources.declarative.extractors import DpathExtractor from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration from airbyte_cdk.sources.declarative.partition_routers import SubstreamPartitionRouter @@ -141,3 +142,14 @@ class TestingRequesterWithDefaultBudget(TestingRequester): def __post_init__(self, parameters: Mapping[str, Any]) -> None: self.api_budget = None super().__post_init__(parameters) + + +@dataclass +class TestingCustomRetrieverWithDecoder(SimpleRetriever): + """ + A test class based on SimpleRetriever that declares a `decoder` field, like source-google-ads' GoogleAdsRetriever. + """ + + __test__: ClassVar[bool] = False # Tell Pytest this is not a Pytest class, despite its name + + decoder: Optional[Decoder] = None