From 788f1db15a28c8913372c06e6fa712d5a22a85db Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Tue, 15 Sep 2026 19:14:25 +0300 Subject: [PATCH 01/10] feat: add CombinedExtractor for combining multiple record extractors A DpathExtractor can only describe one path into a response. Three connectors need more than that and hand-roll it in a custom components.py today, which keeps them out of the Connector Builder and out of manifest-only packaging: - source-google-analytics-data-api ships its own class literally named CombinedExtractor (components.py:17) that zip-merges the i-th record of each sub-extractor - the `zip_merge` mode here. - source-monday's MondayIncrementalItemsExtractor (components.py:75) falls back to a second field path only when the first one produced nothing - the `first_match` mode here. - source-github's in-progress GraphQL extractors read records from several sibling paths of one document - the `union` mode here. Their deep traversal stays out of scope; a separate paginator PR covers it. `zip_merge` truncates at the shortest sub-extractor instead of padding, matching GA4's existing zip() so that connector can drop its custom class without a data change. `first_match` decides "produced at least one record" by peeking a single element and chaining it back, so no sub-extractor is materialized and the peeked record is still emitted. All sub-extractors read the same requests.Response. That is safe for the buffered decoders (JsonDecoder, XmlDecoder, ZipfileDecoder) and not for the streaming ones (CompositeRawDecoder with stream_response=True, i.e. CsvDecoder, JsonlDecoder, JsonItemsDecoder, GzipDecoder outside the Connector Builder). The limitation is documented in the docstring and in the schema, and pinned by a test. The generated pydantic model was updated by hand on purpose: a local `poe assemble` on arm64 rewrites ~340 unrelated lines relative to main (conint/ confloat where main uses Field(..., ge=N), plus unrelated class reordering), so only the delta this schema change implies was applied. A maintainer should run /poe build to regenerate the models authoritatively before merge. Co-Authored-By: Claude Opus 5 (1M context) --- .../declarative_component_schema.yaml | 59 ++++ .../declarative/extractors/__init__.py | 6 + .../extractors/combined_extractor.py | 142 ++++++++ .../models/declarative_component_schema.py | 40 ++- .../parsers/model_to_component_factory.py | 23 ++ .../extractors/test_combined_extractor.py | 329 ++++++++++++++++++ .../test_model_to_component_factory.py | 105 +++++- 7 files changed, 696 insertions(+), 8 deletions(-) create mode 100644 airbyte_cdk/sources/declarative/extractors/combined_extractor.py create mode 100644 unit_tests/sources/declarative/extractors/test_combined_extractor.py diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index aa3bbbc5e0..25f04bace2 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -1900,6 +1900,7 @@ definitions: anyOf: - "$ref": "#/definitions/DpathExtractor" - "$ref": "#/definitions/CustomRecordExtractor" + - "$ref": "#/definitions/CombinedExtractor" file_extractor: description: Responsible for fetching the content of the file. If not defined, the assumption is that the whole response body is the file content anyOf: @@ -2172,6 +2173,60 @@ definitions: $parameters: type: object additionalProperties: true + CombinedExtractor: + title: Combined Extractor + description: Record extractor that combines the output of several sub-extractors. Use it when a single response holds records under more than one path, when the records can be found under one of several alternative paths, or when a record has to be assembled from several parallel lists. + type: object + required: + - type + - extractors + properties: + type: + type: string + enum: [CombinedExtractor] + extractors: + title: Extractors + description: The record extractors to combine. At least one is required. Each sub-extractor is given the same HTTP response, so they must not rely on a streaming decoder (CsvDecoder, JsonlDecoder, JsonItemsDecoder, GzipDecoder), whose response body can only be read once. + type: array + items: + anyOf: + - "$ref": "#/definitions/DpathExtractor" + - "$ref": "#/definitions/CustomRecordExtractor" + - "$ref": "#/definitions/CombinedExtractor" + mode: + title: Combine Mode + description: 'How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor.' + type: string + enum: + - union + - first_match + - zip_merge + default: union + $parameters: + type: object + additionalProperties: true + examples: + - type: CombinedExtractor + mode: union + extractors: + - type: DpathExtractor + field_path: ["data", "repository", "issues", "nodes"] + - type: DpathExtractor + field_path: ["data", "repository", "pullRequests", "nodes"] + - type: CombinedExtractor + mode: first_match + extractors: + - type: DpathExtractor + field_path: ["data", "boards", "*", "items_page", "items"] + - type: DpathExtractor + field_path: ["data", "next_items_page", "items"] + - type: CombinedExtractor + mode: zip_merge + extractors: + - type: DpathExtractor + field_path: ["reports", "*", "dimensions"] + - type: DpathExtractor + field_path: ["reports", "*", "metrics"] DpathExtractor: title: Dpath Extractor description: Record extractor that searches a decoded response over a path defined as an array of fields. @@ -3949,6 +4004,7 @@ definitions: anyOf: - "$ref": "#/definitions/DpathExtractor" - "$ref": "#/definitions/CustomRecordExtractor" + - "$ref": "#/definitions/CombinedExtractor" record_filter: title: Record Filter description: Responsible for filtering records to be emitted by the Source. @@ -4377,17 +4433,20 @@ definitions: anyOf: - "$ref": "#/definitions/DpathExtractor" - "$ref": "#/definitions/CustomRecordExtractor" + - "$ref": "#/definitions/CombinedExtractor" download_target_extractor: description: Responsible for fetching the final result `urls` provided by the completed / finished / ready async job. anyOf: - "$ref": "#/definitions/DpathExtractor" - "$ref": "#/definitions/CustomRecordExtractor" + - "$ref": "#/definitions/CombinedExtractor" download_extractor: description: Responsible for fetching the records from provided urls. anyOf: - "$ref": "#/definitions/DpathExtractor" - "$ref": "#/definitions/CustomRecordExtractor" - "$ref": "#/definitions/ResponseToFileExtractor" + - "$ref": "#/definitions/CombinedExtractor" creation_requester: description: Requester component that describes how to prepare HTTP requests to send to the source API to create the async server-side job. anyOf: diff --git a/airbyte_cdk/sources/declarative/extractors/__init__.py b/airbyte_cdk/sources/declarative/extractors/__init__.py index 8f1d18d12f..5fb4d98a75 100644 --- a/airbyte_cdk/sources/declarative/extractors/__init__.py +++ b/airbyte_cdk/sources/declarative/extractors/__init__.py @@ -2,6 +2,10 @@ # Copyright (c) 2023 Airbyte, Inc., all rights reserved. # +from airbyte_cdk.sources.declarative.extractors.combined_extractor import ( + CombinedExtractor, + CombineMode, +) from airbyte_cdk.sources.declarative.extractors.dpath_extractor import DpathExtractor from airbyte_cdk.sources.declarative.extractors.http_selector import HttpSelector from airbyte_cdk.sources.declarative.extractors.record_filter import RecordFilter @@ -14,6 +18,8 @@ __all__ = [ "TypeTransformer", "HttpSelector", + "CombinedExtractor", + "CombineMode", "DpathExtractor", "RecordFilter", "RecordSelector", diff --git a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py new file mode 100644 index 0000000000..91360c40f5 --- /dev/null +++ b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py @@ -0,0 +1,142 @@ +# +# Copyright (c) 2025 Airbyte, Inc., all rights reserved. +# + +from dataclasses import InitVar, dataclass +from enum import Enum +from itertools import chain +from typing import Any, Dict, Iterable, Iterator, List, Mapping, Union + +import requests + +from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor + + +class CombineMode(Enum): + """ + How a `CombinedExtractor` combines the output of its sub-extractors. + """ + + union = "union" + first_match = "first_match" + zip_merge = "zip_merge" + + +class _NoRecord: + """Sentinel telling "this extractor yielded nothing" apart from a falsy or `None` record.""" + + +_NO_RECORD = _NoRecord() + + +@dataclass +class CombinedExtractor(RecordExtractor): + """Combines the output of several record extractors into a single stream of records. + + A single `DpathExtractor` can only describe one path into the response. Several connectors need + more than that and reach for a custom `components.py` today, which keeps them out of the + Connector Builder. This component covers the three shapes those connectors implement: + + - `union` (default) — yield every record of every sub-extractor, in sub-extractor order. Use it + when one response carries records under several paths, for example a GraphQL document that + returns `issues.nodes` and `pullRequests.nodes` side by side. + - `first_match` — yield the records of the first sub-extractor that produces at least one + record, and skip the remaining sub-extractors. If no sub-extractor produces a record, nothing + is yielded. Use it when an API answers with one of several alternative shapes. The winning + sub-extractor is only peeked at, never restarted: the peeked record is chained back in front + of the remaining ones, so large responses are still streamed lazily and no sub-extractor is + ever materialized into a list. + - `zip_merge` — merge the i-th record of every sub-extractor into a single dictionary, with + later sub-extractors overwriting keys set by earlier ones. Iteration stops at the SHORTEST + sub-extractor (plain `zip` semantics, not `zip_longest`): records past the end of the + shortest sequence are dropped rather than padded with empty values. This matches the + hand-rolled `CombinedExtractor` of `source-google-analytics-data-api`, which this component + is meant to replace. + + Examples of instantiating this component: + ``` + extractor: + type: CombinedExtractor + mode: union + extractors: + - type: DpathExtractor + field_path: ["data", "repository", "issues", "nodes"] + - type: DpathExtractor + field_path: ["data", "repository", "pullRequests", "nodes"] + ``` + + ``` + extractor: + type: CombinedExtractor + mode: first_match + extractors: + - type: DpathExtractor + field_path: ["data", "boards", "*", "items_page", "items"] + - type: DpathExtractor + field_path: ["data", "next_items_page", "items"] + ``` + + Known limitation - the sub-extractors share one `requests.Response`: + Every sub-extractor is handed the same response object, so the response body must be + readable more than once. That holds for the decoders that buffer the whole body, which is + the common case: `JsonDecoder` and any `CompositeRawDecoder` built with + `stream_response=False` read `response.content`, `XmlDecoder` reads `response.text` and + `ZipfileDecoder` reads `response.content` - `requests` caches all of those, so every + sub-extractor sees the full body. + + It does NOT hold for streaming decoders, i.e. a `CompositeRawDecoder` with + `stream_response=True`, which is what `CsvDecoder`, `JsonlDecoder`, `JsonItemsDecoder` and + `GzipDecoder` resolve to outside the Connector Builder, nor for `IterableDecoder`. Those + read `response.raw` and close it afterwards, so the first sub-extractor drains the stream + and the later ones see an exhausted (and closed) stream. Note that the Connector Builder + forces those same decoders to `stream_response=False`, so a manifest that combines + extractors over a streaming decoder can look fine in the Builder and still lose records in + production. Do not use `CombinedExtractor` with a streaming decoder. + + Attributes: + extractors (List[RecordExtractor]): The sub-extractors to combine. At least one is required. + mode (CombineMode): How the sub-extractor outputs are combined. Defaults to `union`. + """ + + extractors: List[RecordExtractor] + parameters: InitVar[Mapping[str, Any]] + mode: CombineMode = CombineMode.union + + def __post_init__(self, parameters: Mapping[str, Any]) -> None: + if not self.extractors: + raise ValueError( + "CombinedExtractor requires at least one extractor in its `extractors` field." + ) + if not isinstance(self.mode, CombineMode): + self.mode = CombineMode(self.mode) + + def extract_records(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: + if self.mode == CombineMode.first_match: + yield from self._extract_first_match(response) + elif self.mode == CombineMode.zip_merge: + yield from self._extract_zip_merge(response) + else: + yield from self._extract_union(response) + + def _extract_union(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: + for extractor in self.extractors: + yield from extractor.extract_records(response) + + def _extract_first_match(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: + for extractor in self.extractors: + records: Iterator[Mapping[str, Any]] = iter(extractor.extract_records(response)) + first_record: Union[Mapping[str, Any], _NoRecord] = next(records, _NO_RECORD) + if isinstance(first_record, _NoRecord): + continue + # Chain the peeked record back instead of restarting the extractor: the rest of the + # records stay lazy and the response is never read twice. + yield from chain([first_record], records) + return + + def _extract_zip_merge(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: + # `zip` stops at the shortest sub-extractor, matching the behavior this component replaces. + for records in zip(*(extractor.extract_records(response) for extractor in self.extractors)): + merged: Dict[str, Any] = {} + for record in records: + merged.update(record) + yield merged diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 21a92cc3be..0030367773 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2517,6 +2517,27 @@ class DpathExtractor(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class CombineMode(Enum): + union = "union" + first_match = "first_match" + zip_merge = "zip_merge" + + +class CombinedExtractor(BaseModel): + type: Literal["CombinedExtractor"] + extractors: List[Union[DpathExtractor, CustomRecordExtractor, CombinedExtractor]] = Field( + ..., + description="The record extractors to combine. At least one is required. Each sub-extractor is given the same HTTP response, so they must not rely on a streaming decoder (CsvDecoder, JsonlDecoder, JsonItemsDecoder, GzipDecoder), whose response body can only be read once.", + title="Extractors", + ) + mode: Optional[CombineMode] = Field( + CombineMode.union, + description='How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor.', + title="Combine Mode", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class ZipfileDecoder(BaseModel): class Config: extra = Extra.allow @@ -2531,7 +2552,7 @@ class Config: class RecordSelector(BaseModel): type: Literal["RecordSelector"] - extractor: Union[DpathExtractor, CustomRecordExtractor] + extractor: Union[DpathExtractor, CustomRecordExtractor, CombinedExtractor] record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( None, description="Responsible for filtering records to be emitted by the Source.", @@ -2776,9 +2797,11 @@ class FileUploader(BaseModel): ..., description="Requester component that describes how to prepare HTTP requests to send to the source API.", ) - download_target_extractor: Union[DpathExtractor, CustomRecordExtractor] = Field( - ..., - description="Responsible for fetching the url where the file is located. This is applied on each records and not on the HTTP response", + download_target_extractor: Union[DpathExtractor, CustomRecordExtractor, CombinedExtractor] = ( + Field( + ..., + description="Responsible for fetching the url where the file is located. This is applied on each records and not on the HTTP response", + ) ) file_extractor: Optional[Union[DpathExtractor, CustomRecordExtractor]] = Field( None, @@ -3260,15 +3283,17 @@ class AsyncRetriever(BaseModel): status_mapping: AsyncJobStatusMap = Field( ..., description="Async Job Status to Airbyte CDK Async Job Status mapping." ) - status_extractor: Union[DpathExtractor, CustomRecordExtractor] = Field( + status_extractor: Union[DpathExtractor, CustomRecordExtractor, CombinedExtractor] = Field( ..., description="Responsible for fetching the actual status of the async job." ) - download_target_extractor: Optional[Union[DpathExtractor, CustomRecordExtractor]] = Field( + download_target_extractor: Optional[ + Union[DpathExtractor, CustomRecordExtractor, CombinedExtractor] + ] = Field( None, description="Responsible for fetching the final result `urls` provided by the completed / finished / ready async job.", ) download_extractor: Optional[ - Union[DpathExtractor, CustomRecordExtractor, ResponseToFileExtractor] + Union[DpathExtractor, CustomRecordExtractor, ResponseToFileExtractor, CombinedExtractor] ] = Field(None, description="Responsible for fetching the records from provided urls.") creation_requester: Union[HttpRequester, CustomRequester] = Field( ..., @@ -3477,6 +3502,7 @@ class DynamicDeclarativeStream(BaseModel): ComplexFieldType.update_forward_refs() GzipDecoder.update_forward_refs() +CombinedExtractor.update_forward_refs() CompositeErrorHandler.update_forward_refs() DeclarativeSource1.update_forward_refs() DeclarativeSource2.update_forward_refs() 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 aa546c73ab..b192883238 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -112,6 +112,8 @@ RecordExpander, ) from airbyte_cdk.sources.declarative.extractors import ( + CombinedExtractor, + CombineMode, DpathExtractor, RecordFilter, RecordSelector, @@ -168,6 +170,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( CheckStream as CheckStreamModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + CombinedExtractor as CombinedExtractorModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ComplexFieldType as ComplexFieldTypeModel, ) @@ -784,6 +789,7 @@ def _init_mappings(self) -> None: DeclarativeStreamModel: self.create_default_stream, DefaultErrorHandlerModel: self.create_default_error_handler, DefaultPaginatorModel: self.create_default_paginator, + CombinedExtractorModel: self.create_combined_extractor, DpathExtractorModel: self.create_dpath_extractor, DpathValidatorModel: self.create_dpath_validator, ResponseToFileExtractorModel: self.create_response_to_file_extractor, @@ -2472,6 +2478,23 @@ def create_default_paginator( return PaginatorTestReadDecorator(paginator, self._limit_pages_fetched_per_slice) return paginator + def create_combined_extractor( + self, + model: CombinedExtractorModel, + config: Config, + decoder: Optional[Decoder] = None, + **kwargs: Any, + ) -> CombinedExtractor: + extractors = [ + self._create_component_from_model(model=sub_extractor, config=config, decoder=decoder) + for sub_extractor in model.extractors + ] + return CombinedExtractor( + extractors=extractors, + mode=CombineMode(model.mode.value) if model.mode else CombineMode.union, + parameters=model.parameters or {}, + ) + def create_dpath_extractor( self, model: DpathExtractorModel, diff --git a/unit_tests/sources/declarative/extractors/test_combined_extractor.py b/unit_tests/sources/declarative/extractors/test_combined_extractor.py new file mode 100644 index 0000000000..2a3760e7f4 --- /dev/null +++ b/unit_tests/sources/declarative/extractors/test_combined_extractor.py @@ -0,0 +1,329 @@ +# +# Copyright (c) 2025 Airbyte, Inc., all rights reserved. +# +import io +import json +from dataclasses import dataclass, field +from typing import Any, Dict, Iterable, List, Mapping, Union + +import pytest +import requests + +from airbyte_cdk.sources.declarative.decoders import CompositeRawDecoder, JsonDecoder +from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import JsonLineParser +from airbyte_cdk.sources.declarative.extractors import ( + CombinedExtractor, + CombineMode, + DpathExtractor, +) + +config = {"field": "issues"} +parameters = {"parameters_field": "issues"} + +decoder_json = JsonDecoder(parameters={}) + +# A GraphQL document that carries records under two sibling paths, which is the shape that +# source-github and source-monday work around with a custom components.py today. +GRAPHQL_BODY = { + "data": { + "repository": { + "issues": {"nodes": [{"id": "I_1", "title": "first issue"}]}, + "pullRequests": { + "nodes": [{"id": "PR_1", "title": "first pr"}, {"id": "PR_2", "title": "second pr"}] + }, + "discussions": {"nodes": []}, + } + } +} + +# The GA4 report shape: dimensions and metrics come back as two parallel lists that have to be +# merged position by position. +REPORT_BODY = { + "rows": [ + {"dimensions": {"date": "20240101"}, "metrics": {"activeUsers": 1}}, + {"dimensions": {"date": "20240102"}, "metrics": {"activeUsers": 2}}, + ] +} + + +def create_response(body: Union[Dict, List, bytes]) -> requests.Response: + response = requests.Response() + response.raw = io.BytesIO(body if isinstance(body, bytes) else json.dumps(body).encode("utf-8")) + return response + + +def dpath(*field_path: str, decoder=decoder_json) -> DpathExtractor: + return DpathExtractor( + field_path=list(field_path), config=config, decoder=decoder, parameters=parameters + ) + + +@dataclass +class _CountingExtractor: + """Records how many times it is asked for records, so laziness can be asserted.""" + + records: List[Mapping[str, Any]] + calls: List[str] = field(default_factory=list) + + def extract_records(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: + self.calls.append("extract_records") + yield from self.records + + +def test_union_yields_every_record_in_sub_extractor_order(): + extractor = CombinedExtractor( + extractors=[ + dpath("data", "repository", "pullRequests", "nodes"), + dpath("data", "repository", "issues", "nodes"), + ], + mode=CombineMode.union, + parameters=parameters, + ) + + records = list(extractor.extract_records(create_response(GRAPHQL_BODY))) + + assert records == [ + {"id": "PR_1", "title": "first pr"}, + {"id": "PR_2", "title": "second pr"}, + {"id": "I_1", "title": "first issue"}, + ] + + +def test_union_is_the_default_mode(): + extractor = CombinedExtractor( + extractors=[ + dpath("data", "repository", "issues", "nodes"), + dpath("data", "repository", "pullRequests", "nodes"), + ], + parameters=parameters, + ) + + assert extractor.mode == CombineMode.union + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [ + {"id": "I_1", "title": "first issue"}, + {"id": "PR_1", "title": "first pr"}, + {"id": "PR_2", "title": "second pr"}, + ] + + +def test_union_skips_sub_extractors_that_yield_nothing(): + extractor = CombinedExtractor( + extractors=[ + dpath("data", "repository", "discussions", "nodes"), + dpath("data", "repository", "issues", "nodes"), + dpath("data", "repository", "does_not_exist"), + ], + mode=CombineMode.union, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [ + {"id": "I_1", "title": "first issue"} + ] + + +def test_first_match_falls_through_an_empty_extractor(): + extractor = CombinedExtractor( + extractors=[ + dpath("data", "repository", "discussions", "nodes"), + dpath("data", "repository", "pullRequests", "nodes"), + dpath("data", "repository", "issues", "nodes"), + ], + mode=CombineMode.first_match, + parameters=parameters, + ) + + records = list(extractor.extract_records(create_response(GRAPHQL_BODY))) + + # The second extractor wins, and the third one is not used at all. + assert records == [ + {"id": "PR_1", "title": "first pr"}, + {"id": "PR_2", "title": "second pr"}, + ] + + +def test_first_match_does_not_drop_the_peeked_record(): + """The winner is peeked to find out whether it produced anything; that record must still be emitted.""" + winner = _CountingExtractor(records=[{"id": 1}, {"id": 2}, {"id": 3}]) + extractor = CombinedExtractor( + extractors=[_CountingExtractor(records=[]), winner], + mode=CombineMode.first_match, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [ + {"id": 1}, + {"id": 2}, + {"id": 3}, + ] + assert winner.calls == ["extract_records"] + + +def test_first_match_does_not_touch_extractors_after_the_winner(): + loser = _CountingExtractor(records=[{"id": 99}]) + extractor = CombinedExtractor( + extractors=[_CountingExtractor(records=[{"id": 1}]), loser], + mode=CombineMode.first_match, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [{"id": 1}] + assert loser.calls == [] + + +def test_first_match_yields_nothing_when_every_extractor_is_empty(): + extractor = CombinedExtractor( + extractors=[ + dpath("data", "repository", "discussions", "nodes"), + dpath("data", "repository", "does_not_exist"), + ], + mode=CombineMode.first_match, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [] + + +def test_zip_merge_merges_the_ith_record_of_every_extractor(): + extractor = CombinedExtractor( + extractors=[dpath("rows", "*", "dimensions"), dpath("rows", "*", "metrics")], + mode=CombineMode.zip_merge, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(REPORT_BODY))) == [ + {"date": "20240101", "activeUsers": 1}, + {"date": "20240102", "activeUsers": 2}, + ] + + +def test_zip_merge_later_extractors_win_on_key_collision(): + extractor = CombinedExtractor( + extractors=[ + _CountingExtractor(records=[{"id": 1, "source": "first", "only_first": True}]), + _CountingExtractor(records=[{"id": 1, "source": "second"}]), + ], + mode=CombineMode.zip_merge, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(REPORT_BODY))) == [ + {"id": 1, "source": "second", "only_first": True} + ] + + +def test_zip_merge_stops_at_the_shortest_sub_extractor(): + """`zip`, not `zip_longest`: the trailing records of the longer extractor are dropped, not padded.""" + extractor = CombinedExtractor( + extractors=[ + _CountingExtractor(records=[{"a": 1}, {"a": 2}, {"a": 3}]), + _CountingExtractor(records=[{"b": 1}, {"b": 2}]), + ], + mode=CombineMode.zip_merge, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(REPORT_BODY))) == [ + {"a": 1, "b": 1}, + {"a": 2, "b": 2}, + ] + + +def test_zip_merge_yields_nothing_when_one_extractor_is_empty(): + extractor = CombinedExtractor( + extractors=[dpath("rows", "*", "dimensions"), dpath("rows", "*", "does_not_exist")], + mode=CombineMode.zip_merge, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(REPORT_BODY))) == [] + + +@pytest.mark.parametrize( + "mode", + [CombineMode.union, CombineMode.first_match, CombineMode.zip_merge], +) +def test_a_single_sub_extractor_behaves_like_that_extractor_alone(mode: CombineMode): + field_path = ["data", "repository", "pullRequests", "nodes"] + direct = DpathExtractor( + field_path=field_path, config=config, decoder=decoder_json, parameters=parameters + ) + combined = CombinedExtractor( + extractors=[ + DpathExtractor( + field_path=field_path, config=config, decoder=decoder_json, parameters=parameters + ) + ], + mode=mode, + parameters=parameters, + ) + + expected = list(direct.extract_records(create_response(GRAPHQL_BODY))) + assert list(combined.extract_records(create_response(GRAPHQL_BODY))) == expected + assert expected == [ + {"id": "PR_1", "title": "first pr"}, + {"id": "PR_2", "title": "second pr"}, + ] + + +def test_an_empty_extractors_list_is_a_configuration_error(): + with pytest.raises(ValueError) as exc_info: + CombinedExtractor(extractors=[], parameters=parameters) + + assert "extractors" in str(exc_info.value) + + +def test_mode_can_be_given_as_a_string(): + extractor = CombinedExtractor( + extractors=[dpath("data", "repository", "issues", "nodes")], + mode="first_match", + parameters=parameters, + ) + + assert extractor.mode == CombineMode.first_match + + +def test_an_unknown_mode_is_rejected(): + with pytest.raises(ValueError): + CombinedExtractor( + extractors=[dpath("data", "repository", "issues", "nodes")], + mode="concatenate", + parameters=parameters, + ) + + +def test_buffered_decoders_can_be_read_by_every_sub_extractor(): + """`JsonDecoder` reads `response.content`, which `requests` caches, so sharing the response is safe.""" + extractor = CombinedExtractor( + extractors=[ + dpath("data", "repository", "issues", "nodes"), + dpath("data", "repository", "pullRequests", "nodes"), + dpath("data", "repository", "issues", "nodes"), + ], + mode=CombineMode.union, + parameters=parameters, + ) + + assert len(list(extractor.extract_records(create_response(GRAPHQL_BODY)))) == 4 + + +def test_streaming_decoder_is_a_known_limitation(): + """Pins the documented limitation: a streaming decoder can only be read by the first sub-extractor. + + `CompositeRawDecoder(stream_response=True)` - what `CsvDecoder`, `JsonlDecoder`, + `JsonItemsDecoder` and `GzipDecoder` resolve to outside the Connector Builder - consumes and + then closes `response.raw`, so the second sub-extractor cannot read the body again. + """ + streaming_decoder = CompositeRawDecoder(parser=JsonLineParser(), stream_response=True) + extractor = CombinedExtractor( + extractors=[ + dpath("data", decoder=streaming_decoder), + dpath("data", decoder=streaming_decoder), + ], + mode=CombineMode.union, + parameters=parameters, + ) + response = create_response(b'{"data": [{"id": 1}, {"id": 2}]}') + + with pytest.raises(ValueError, match="closed file"): + list(extractor.extract_records(response)) 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 21c99adc71..5f8ed7d4c8 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 @@ -54,7 +54,13 @@ from airbyte_cdk.sources.declarative.concurrency_level import ConcurrencyLevel from airbyte_cdk.sources.declarative.datetime.min_max_datetime import MinMaxDatetime from airbyte_cdk.sources.declarative.decoders import JsonDecoder, PaginationDecoderDecorator -from airbyte_cdk.sources.declarative.extractors import DpathExtractor, RecordFilter, RecordSelector +from airbyte_cdk.sources.declarative.extractors import ( + CombinedExtractor, + CombineMode, + DpathExtractor, + RecordFilter, + RecordSelector, +) from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor from airbyte_cdk.sources.declarative.extractors.record_filter import ( ClientSideIncrementalRecordFilterDecorator, @@ -6555,3 +6561,100 @@ def test_incremental_dependency_without_incremental_sync_warns( for record in caplog.records ) assert warning_emitted == expected_warning + + +def test_create_combined_extractor_through_the_record_selector(): + """A manifest-level test: it exercises the schema, the generated model and the factory registration.""" + content = """ + selector: + type: RecordSelector + extractor: + type: CombinedExtractor + mode: first_match + extractors: + - type: DpathExtractor + field_path: ["data", "boards", "*", "items_page", "items"] + - type: CombinedExtractor + mode: union + extractors: + - type: DpathExtractor + field_path: ["data", "next_items_page", "items"] + - type: DpathExtractor + field_path: ["data", "{{ parameters['name'] }}"] + $parameters: + name: "lists" + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + assert isinstance(selector, RecordSelector) + extractor = selector.extractor + assert isinstance(extractor, CombinedExtractor) + assert extractor.mode == CombineMode.first_match + assert len(extractor.extractors) == 2 + assert isinstance(extractor.extractors[0], DpathExtractor) + assert [fp.eval(input_config) for fp in extractor.extractors[0]._field_path] == [ + "data", + "boards", + "*", + "items_page", + "items", + ] + + nested = extractor.extractors[1] + assert isinstance(nested, CombinedExtractor) + assert nested.mode == CombineMode.union + assert [fp.eval(input_config) for fp in nested.extractors[1]._field_path] == ["data", "lists"] + + response = requests.Response() + response._content = json.dumps( + {"data": {"next_items_page": {"items": [{"id": 1}]}, "lists": [{"id": 2}]}} + ).encode("utf-8") + assert list(extractor.extract_records(response)) == [{"id": 1}, {"id": 2}] + + +def test_create_combined_extractor_defaults_to_union(): + content = """ + selector: + type: RecordSelector + extractor: + type: CombinedExtractor + extractors: + - type: DpathExtractor + field_path: ["a"] + - type: DpathExtractor + field_path: ["b"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + assert isinstance(selector.extractor, CombinedExtractor) + assert selector.extractor.mode == CombineMode.union + + response = requests.Response() + response._content = json.dumps({"a": [{"id": 1}], "b": [{"id": 2}]}).encode("utf-8") + assert list(selector.extractor.extract_records(response)) == [{"id": 1}, {"id": 2}] From 5c7510c37244cdcb1a695d8cebc7dee205253327 Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Tue, 15 Sep 2026 20:13:04 +0300 Subject: [PATCH 02/10] fix(CombinedExtractor): reject streaming decoders, enforce minItems, address review The streaming-decoder hazard was documented as a loud `ValueError`, but that is only what `io.BytesIO` does. In production `requests` puts a `urllib3.HTTPResponse` in `response.raw`, whose read on a closed body returns `b""`, so a `CombinedExtractor` over `JsonlDecoder` or `GzipDecoder` silently drops records: `union` emits only the first sub-extractor's output and `first_match` emits nothing when the first path misses. The Connector Builder downgrades those decoders to `stream_response=False`, so the manifest test-reads green and loses data once published, with the cursor advancing regardless. - `create_combined_extractor` now raises an `AirbyteTracedException` (`config_error`) when the stream's decoder streams the response, naming the decoders and the buffered alternatives. It also rejects the downgraded `CompositeRawDecoder` instances the Builder substitutes, so the Builder rejects exactly what production rejects, and it fires for nested `CombinedExtractor`s, which inherit the retriever's decoder. - `test_streaming_decoder_is_a_known_limitation` is replaced by two tests built on a real `urllib3.HTTPResponse`, pinning the silent loss in `union` and the silent wrong answer in `first_match` against a buffered control, plus parametrized factory tests for the parse-time rejection. `minItems: 1` is added to `extractors`. The stated blocker did not exist: datamodel-codegen silently drops `minItems` on this self-referencing construct at the pinned 0.26.3, so pydantic never sees it, while manifest validation is jsonschema and does enforce it. The `__post_init__` guard stays as the backstop. The hand-written generated enum is renamed `CombineMode` -> `Mode` to match what the pinned generator emits (it derives an inline property enum's name from the property name, not from `title`), so a `/poe build` regen is not a surprise diff. The runtime enum stays `CombineMode`. Also from the review: - Three manifest-level tests through `ConcurrentDeclarativeSource` so the 59 new schema lines are covered by jsonschema validation rather than `parse_obj`. - `zip_merge` logs a warning when the sub-extractors have different lengths, and raises an error naming the mode and sub-extractor index when one yields a non-object instead of a bare `dictionary update sequence` built-in error. - Docstring and schema descriptions rewritten: the limitation is silent, not loud, and is now rejected; the per-sub-extractor decode cost and the wrong paginator record count are documented. The docstring uses headings so pdoc no longer renders the warning as a code block. - `AsyncHttpJobRepository.status_extractor` / `download_target_extractor` and the three factory `extractor_model` parameters widened to what they receive. - `CombinedExtractor` and `CombineMode` exported from `airbyte_cdk`. - Coverage for the four widened unions outside `RecordSelector` and for `CombinedExtractor` behind an `OffsetIncrement` paginator. Backward compatibility re-verified: 531 monorepo manifests validate identically against the base and head schemas, and 5,818 extractor payloads resolve identically through the base and head model modules. Co-Authored-By: Claude Opus 5 (1M context) --- airbyte_cdk/__init__.py | 9 +- .../declarative_component_schema.yaml | 9 +- .../extractors/combined_extractor.py | 104 ++++-- .../models/declarative_component_schema.py | 13 +- .../parsers/model_to_component_factory.py | 54 ++- .../requesters/http_job_repository.py | 9 +- .../extractors/test_combined_extractor.py | 118 ++++++- .../test_model_to_component_factory.py | 329 ++++++++++++++++++ .../test_concurrent_declarative_source.py | 96 +++++ 9 files changed, 686 insertions(+), 55 deletions(-) diff --git a/airbyte_cdk/__init__.py b/airbyte_cdk/__init__.py index d971690c86..4f32754ee7 100644 --- a/airbyte_cdk/__init__.py +++ b/airbyte_cdk/__init__.py @@ -103,7 +103,12 @@ from .sources.declarative.datetime.min_max_datetime import MinMaxDatetime from .sources.declarative.decoders import Decoder, JsonDecoder from .sources.declarative.exceptions import ReadException -from .sources.declarative.extractors import DpathExtractor, RecordSelector +from .sources.declarative.extractors import ( + CombinedExtractor, + CombineMode, + DpathExtractor, + RecordSelector, +) from .sources.declarative.extractors.record_extractor import RecordExtractor from .sources.declarative.extractors.record_filter import RecordFilter from .sources.declarative.interpolation import InterpolatedBoolean, InterpolatedString @@ -234,6 +239,8 @@ "BasicHttpAuthenticator", "BearerAuthenticator", "CartesianProductStreamSlicer", + "CombineMode", + "CombinedExtractor", "CursorPaginationStrategy", "DatetimeBasedCursor", "DeclarativeAuthenticator", diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 25f04bace2..55718cfd24 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2186,8 +2186,9 @@ definitions: enum: [CombinedExtractor] extractors: title: Extractors - description: The record extractors to combine. At least one is required. Each sub-extractor is given the same HTTP response, so they must not rely on a streaming decoder (CsvDecoder, JsonlDecoder, JsonItemsDecoder, GzipDecoder), whose response body can only be read once. + description: The record extractors to combine. At least one is required. Each sub-extractor is given the same HTTP response and decodes it independently, so the response is parsed once per sub-extractor. Streaming decoders (CsvDecoder, JsonlDecoder, JsonItemsDecoder, GzipDecoder, IterableDecoder) can only read the response body once and are rejected. type: array + minItems: 1 items: anyOf: - "$ref": "#/definitions/DpathExtractor" @@ -2195,13 +2196,17 @@ definitions: - "$ref": "#/definitions/CombinedExtractor" mode: title: Combine Mode - description: 'How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor.' + description: 'How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that an OffsetIncrement or PageIncrement paginator counts the combined records, which under "first_match" and "zip_merge" is not the number of records the API returned for the page.' type: string enum: - union - first_match - zip_merge default: union + examples: + - union + - first_match + - zip_merge $parameters: type: object additionalProperties: true diff --git a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py index 91360c40f5..28f34b9bb8 100644 --- a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py +++ b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py @@ -2,15 +2,18 @@ # Copyright (c) 2025 Airbyte, Inc., all rights reserved. # +import logging from dataclasses import InitVar, dataclass from enum import Enum from itertools import chain -from typing import Any, Dict, Iterable, Iterator, List, Mapping, Union +from typing import Any, Dict, Iterable, Iterator, List, Mapping, Sequence, Union import requests from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor +logger = logging.getLogger("airbyte") + class CombineMode(Enum): """ @@ -49,9 +52,9 @@ class CombinedExtractor(RecordExtractor): - `zip_merge` — merge the i-th record of every sub-extractor into a single dictionary, with later sub-extractors overwriting keys set by earlier ones. Iteration stops at the SHORTEST sub-extractor (plain `zip` semantics, not `zip_longest`): records past the end of the - shortest sequence are dropped rather than padded with empty values. This matches the - hand-rolled `CombinedExtractor` of `source-google-analytics-data-api`, which this component - is meant to replace. + shortest sequence are dropped rather than padded with empty values, and a warning is logged + when that happens. This matches the hand-rolled `CombinedExtractor` of + `source-google-analytics-data-api`, which this component is meant to replace. Examples of instantiating this component: ``` @@ -76,22 +79,40 @@ class CombinedExtractor(RecordExtractor): field_path: ["data", "next_items_page", "items"] ``` - Known limitation - the sub-extractors share one `requests.Response`: - Every sub-extractor is handed the same response object, so the response body must be - readable more than once. That holds for the decoders that buffer the whole body, which is - the common case: `JsonDecoder` and any `CompositeRawDecoder` built with - `stream_response=False` read `response.content`, `XmlDecoder` reads `response.text` and - `ZipfileDecoder` reads `response.content` - `requests` caches all of those, so every - sub-extractor sees the full body. - - It does NOT hold for streaming decoders, i.e. a `CompositeRawDecoder` with - `stream_response=True`, which is what `CsvDecoder`, `JsonlDecoder`, `JsonItemsDecoder` and - `GzipDecoder` resolve to outside the Connector Builder, nor for `IterableDecoder`. Those - read `response.raw` and close it afterwards, so the first sub-extractor drains the stream - and the later ones see an exhausted (and closed) stream. Note that the Connector Builder - forces those same decoders to `stream_response=False`, so a manifest that combines - extractors over a streaming decoder can look fine in the Builder and still lose records in - production. Do not use `CombinedExtractor` with a streaming decoder. + ## Streaming decoders are rejected + + Every sub-extractor is handed the same `requests.Response`, so the response body must be + readable more than once. That holds for the decoders that buffer the whole body, which is the + common case: `JsonDecoder` and any `CompositeRawDecoder` built with `stream_response=False` + read `response.content`, `XmlDecoder` reads `response.text` and `ZipfileDecoder` reads + `response.content` — `requests` caches all of those, so every sub-extractor sees the full body. + + It does NOT hold for streaming decoders, i.e. a `CompositeRawDecoder` with + `stream_response=True`, which is what `CsvDecoder`, `JsonlDecoder`, `JsonItemsDecoder` and + `GzipDecoder` resolve to outside the Connector Builder, nor for `IterableDecoder`. Those read + `response.raw` and close it afterwards, so the first sub-extractor drains the stream and the + later ones read a dead body. That loss is SILENT rather than loud: `requests` puts a + `urllib3.HTTPResponse` in `response.raw`, and reading a closed `urllib3.HTTPResponse` returns + an empty body instead of raising, so `union` emits only the first sub-extractor's records and + `first_match` emits nothing at all when the first path misses. + + Because of that, `ModelToComponentFactory.create_combined_extractor` refuses to build this + component over a streaming decoder and raises a configuration error naming the decoder. The + Connector Builder forces those same decoders to `stream_response=False`, so the rejection is + applied there too rather than letting a manifest test-read correctly and lose records once + published. + + ## Cost + + Each sub-extractor decodes the response independently — `CompositeRawDecoder` re-parses the + cached body on every `decode()` call — so a response is parsed once per sub-extractor. Under + `zip_merge` all sub-extractor outputs are also alive at the same time. Combining three + extractors over a large response costs roughly three times the decode time of a single one. + + Note that an `OffsetIncrement` or `PageIncrement` paginator builds its own copy of the + extractor to count the records of a page, which doubles that cost, and the count it obtains is + the combined count: the winner's record count under `first_match` and the shortest + sub-extractor's count under `zip_merge`, neither of which is the API's page size. Attributes: extractors (List[RecordExtractor]): The sub-extractors to combine. At least one is required. @@ -134,9 +155,40 @@ def _extract_first_match(self, response: requests.Response) -> Iterable[Mapping[ return def _extract_zip_merge(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: - # `zip` stops at the shortest sub-extractor, matching the behavior this component replaces. - for records in zip(*(extractor.extract_records(response) for extractor in self.extractors)): - merged: Dict[str, Any] = {} - for record in records: - merged.update(record) - yield merged + # Iteration stops at the shortest sub-extractor, matching the behavior this component + # replaces. Unlike plain `zip`, the round in which a sub-extractor runs out is completed so + # that an uneven number of records can be reported instead of silently truncated. + iterators = [iter(extractor.extract_records(response)) for extractor in self.extractors] + emitted = 0 + while True: + round_records = [next(iterator, _NO_RECORD) for iterator in iterators] + exhausted = [ + index for index, record in enumerate(round_records) if isinstance(record, _NoRecord) + ] + if exhausted: + if len(exhausted) != len(iterators): + logger.warning( + "CombinedExtractor in `zip_merge` mode discarded the records of the " + "sub-extractors that outlasted the shortest one: sub-extractors " + "%s ran out after %s record(s) while the others still had records. " + "`zip_merge` stops at the shortest sub-extractor.", + exhausted, + emitted, + ) + return + yield self._merge(round_records) + emitted += 1 + + def _merge(self, records: Sequence[Any]) -> Mapping[str, Any]: + merged: Dict[str, Any] = {} + for index, record in enumerate(records): + if not isinstance(record, Mapping): + raise ValueError( + f"CombinedExtractor in `zip_merge` mode can only merge records that are " + f"objects, but sub-extractor {index} " + f"({type(self.extractors[index]).__name__}) yielded a " + f"{type(record).__name__}. Point that sub-extractor at a path that holds " + f"objects, or use the `union` mode." + ) + merged.update(record) + return merged diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 0030367773..98d47b7254 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2517,7 +2517,7 @@ class DpathExtractor(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class CombineMode(Enum): +class Mode(Enum): union = "union" first_match = "first_match" zip_merge = "zip_merge" @@ -2527,12 +2527,13 @@ class CombinedExtractor(BaseModel): type: Literal["CombinedExtractor"] extractors: List[Union[DpathExtractor, CustomRecordExtractor, CombinedExtractor]] = Field( ..., - description="The record extractors to combine. At least one is required. Each sub-extractor is given the same HTTP response, so they must not rely on a streaming decoder (CsvDecoder, JsonlDecoder, JsonItemsDecoder, GzipDecoder), whose response body can only be read once.", + description="The record extractors to combine. At least one is required. Each sub-extractor is given the same HTTP response and decodes it independently, so the response is parsed once per sub-extractor. Streaming decoders (CsvDecoder, JsonlDecoder, JsonItemsDecoder, GzipDecoder, IterableDecoder) can only read the response body once and are rejected.", title="Extractors", ) - mode: Optional[CombineMode] = Field( - CombineMode.union, - description='How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor.', + mode: Optional[Mode] = Field( + Mode.union, + description='How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that an OffsetIncrement or PageIncrement paginator counts the combined records, which under "first_match" and "zip_merge" is not the number of records the API returned for the page.', + examples=["union", "first_match", "zip_merge"], title="Combine Mode", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -3502,8 +3503,8 @@ class DynamicDeclarativeStream(BaseModel): ComplexFieldType.update_forward_refs() GzipDecoder.update_forward_refs() -CombinedExtractor.update_forward_refs() CompositeErrorHandler.update_forward_refs() +CombinedExtractor.update_forward_refs() DeclarativeSource1.update_forward_refs() DeclarativeSource2.update_forward_refs() SelectiveAuthenticator.update_forward_refs() 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 b192883238..c5826c11d9 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -692,6 +692,7 @@ from airbyte_cdk.sources.streams.http.error_handlers.response_models import ResponseAction from airbyte_cdk.sources.types import Config from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer +from airbyte_cdk.utils.traced_exception import AirbyteTracedException ComponentDefinition = Mapping[str, Any] @@ -2434,7 +2435,9 @@ def create_default_paginator( config: Config, *, url_base: str, - extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None, + extractor_model: Optional[ + Union[CustomRecordExtractorModel, DpathExtractorModel, CombinedExtractorModel] + ] = None, decoder: Optional[Decoder] = None, cursor_used_for_stop_condition: Optional[Cursor] = None, ) -> Union[DefaultPaginator, PaginatorTestReadDecorator]: @@ -2478,6 +2481,46 @@ def create_default_paginator( return PaginatorTestReadDecorator(paginator, self._limit_pages_fetched_per_slice) return paginator + def _reject_combined_extractor_over_streaming_decoder(self, decoder: Optional[Decoder]) -> None: + """Refuse to build a `CombinedExtractor` whose response body can only be read once. + + Every sub-extractor is handed the same `requests.Response`. A streaming decoder consumes + and closes `response.raw`, so every sub-extractor after the first reads a closed + `urllib3.HTTPResponse`, which returns an empty body instead of raising: `union` emits only + the first sub-extractor's records and `first_match` emits nothing when the first path + misses. Rejecting the manifest is the only way to make that loud. + + The Connector Builder downgrades those same decoders to `stream_response=False` so a test + read can be replayed, which would let a manifest look correct in the Builder and lose + records once published. The downgraded decoders are therefore rejected in the Builder too. + """ + if decoder is None: + return + inner_decoder = ( + decoder.decoder if isinstance(decoder, PaginationDecoderDecorator) else decoder + ) + streams_in_production = inner_decoder.is_stream_response() or ( + # The four decoders the Builder downgrades are the only bare `CompositeRawDecoder`s + # this factory builds; every buffered decoder is a distinct class. + self._emit_connector_builder_messages and isinstance(inner_decoder, CompositeRawDecoder) + ) + if not streams_in_production: + return + raise AirbyteTracedException( + message=( + "CombinedExtractor is not supported with a streaming decoder (CsvDecoder, " + "JsonlDecoder, JsonItemsDecoder, GzipDecoder, IterableDecoder), which reads the " + "response body only once. Use JsonDecoder, XmlDecoder or ZipfileDecoder, or " + "declare a single extractor." + ), + internal_message=( + f"CombinedExtractor was configured with {type(inner_decoder).__name__}, which " + f"streams the response. Sub-extractors after the first would read a closed " + f"response and silently return no records." + ), + failure_type=FailureType.config_error, + ) + def create_combined_extractor( self, model: CombinedExtractorModel, @@ -2485,6 +2528,7 @@ def create_combined_extractor( decoder: Optional[Decoder] = None, **kwargs: Any, ) -> CombinedExtractor: + self._reject_combined_extractor_over_streaming_decoder(decoder) extractors = [ self._create_component_from_model(model=sub_extractor, config=config, decoder=decoder) for sub_extractor in model.extractors @@ -3095,7 +3139,9 @@ def create_offset_increment( model: OffsetIncrementModel, config: Config, decoder: Decoder, - extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None, + extractor_model: Optional[ + Union[CustomRecordExtractorModel, DpathExtractorModel, CombinedExtractorModel] + ] = None, **kwargs: Any, ) -> OffsetIncrement: if isinstance(decoder, PaginationDecoderDecorator): @@ -3144,7 +3190,9 @@ def create_page_increment( model: PageIncrementModel, config: Config, decoder: Optional[Decoder] = None, - extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None, + extractor_model: Optional[ + Union[CustomRecordExtractorModel, DpathExtractorModel, CombinedExtractorModel] + ] = None, **kwargs: Any, ) -> PageIncrement: # Like OffsetIncrement, we instantiate a separate extractor with identical behavior to the diff --git a/airbyte_cdk/sources/declarative/requesters/http_job_repository.py b/airbyte_cdk/sources/declarative/requesters/http_job_repository.py index d837ed9024..efdcb461e8 100644 --- a/airbyte_cdk/sources/declarative/requesters/http_job_repository.py +++ b/airbyte_cdk/sources/declarative/requesters/http_job_repository.py @@ -14,10 +14,7 @@ from airbyte_cdk.sources.declarative.async_job.job import AsyncJob from airbyte_cdk.sources.declarative.async_job.repository import AsyncJobRepository from airbyte_cdk.sources.declarative.async_job.status import AsyncJobStatus -from airbyte_cdk.sources.declarative.extractors.dpath_extractor import ( - DpathExtractor, - RecordExtractor, -) +from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor from airbyte_cdk.sources.declarative.extractors.response_to_file_extractor import ( ResponseToFileExtractor, ) @@ -41,9 +38,9 @@ class AsyncHttpJobRepository(AsyncJobRepository): download_retriever: SimpleRetriever abort_requester: Optional[Requester] delete_requester: Optional[Requester] - status_extractor: DpathExtractor + status_extractor: RecordExtractor status_mapping: Mapping[str, AsyncJobStatus] - download_target_extractor: Optional[DpathExtractor] + download_target_extractor: Optional[RecordExtractor] # timeout for the job to be completed, passed from `polling_job_timeout` job_timeout: Optional[timedelta] = None diff --git a/unit_tests/sources/declarative/extractors/test_combined_extractor.py b/unit_tests/sources/declarative/extractors/test_combined_extractor.py index 2a3760e7f4..21cab56510 100644 --- a/unit_tests/sources/declarative/extractors/test_combined_extractor.py +++ b/unit_tests/sources/declarative/extractors/test_combined_extractor.py @@ -3,11 +3,13 @@ # import io import json +import logging from dataclasses import dataclass, field from typing import Any, Dict, Iterable, List, Mapping, Union import pytest import requests +import urllib3 from airbyte_cdk.sources.declarative.decoders import CompositeRawDecoder, JsonDecoder from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import JsonLineParser @@ -16,6 +18,7 @@ CombineMode, DpathExtractor, ) +from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor config = {"field": "issues"} parameters = {"parameters_field": "issues"} @@ -47,8 +50,15 @@ def create_response(body: Union[Dict, List, bytes]) -> requests.Response: + """A response whose `raw` is a real `urllib3.HTTPResponse`, which is what `requests` builds. + + This matters for the streaming-decoder tests: `io.BytesIO` raises on a read after close, while + `urllib3.HTTPResponse` returns an empty body, so only the latter reproduces production. + """ + raw = body if isinstance(body, bytes) else json.dumps(body).encode("utf-8") response = requests.Response() - response.raw = io.BytesIO(body if isinstance(body, bytes) else json.dumps(body).encode("utf-8")) + response.status_code = 200 + response.raw = urllib3.HTTPResponse(body=io.BytesIO(raw), preload_content=False, status=200) return response @@ -59,7 +69,7 @@ def dpath(*field_path: str, decoder=decoder_json) -> DpathExtractor: @dataclass -class _CountingExtractor: +class _CountingExtractor(RecordExtractor): """Records how many times it is asked for records, so laziness can be asserted.""" records: List[Mapping[str, Any]] @@ -307,14 +317,21 @@ def test_buffered_decoders_can_be_read_by_every_sub_extractor(): assert len(list(extractor.extract_records(create_response(GRAPHQL_BODY)))) == 4 -def test_streaming_decoder_is_a_known_limitation(): - """Pins the documented limitation: a streaming decoder can only be read by the first sub-extractor. +def _streaming_decoder() -> CompositeRawDecoder: + """What `CsvDecoder`, `JsonlDecoder`, `JsonItemsDecoder` and `GzipDecoder` resolve to in production.""" + return CompositeRawDecoder(parser=JsonLineParser(), stream_response=True) + + +def test_union_over_a_streaming_decoder_silently_drops_the_later_sub_extractors(): + """Pins the real production behavior, which is silent data loss rather than an error. - `CompositeRawDecoder(stream_response=True)` - what `CsvDecoder`, `JsonlDecoder`, - `JsonItemsDecoder` and `GzipDecoder` resolve to outside the Connector Builder - consumes and - then closes `response.raw`, so the second sub-extractor cannot read the body again. + A streaming decoder consumes and closes `response.raw`. Reading a closed + `urllib3.HTTPResponse` - what `requests` puts there - returns an empty body instead of raising, + so the second sub-extractor yields nothing and no exception reaches the sync. This is why + `ModelToComponentFactory.create_combined_extractor` rejects the combination at parse time; see + `test_combined_extractor_over_a_streaming_decoder_is_rejected` in the factory tests. """ - streaming_decoder = CompositeRawDecoder(parser=JsonLineParser(), stream_response=True) + streaming_decoder = _streaming_decoder() extractor = CombinedExtractor( extractors=[ dpath("data", decoder=streaming_decoder), @@ -323,7 +340,86 @@ def test_streaming_decoder_is_a_known_limitation(): mode=CombineMode.union, parameters=parameters, ) - response = create_response(b'{"data": [{"id": 1}, {"id": 2}]}') - with pytest.raises(ValueError, match="closed file"): - list(extractor.extract_records(response)) + records = list(extractor.extract_records(create_response(b'{"data": [{"id": 1}, {"id": 2}]}'))) + + # No exception, and half of the expected four records are gone. + assert records == [{"id": 1}, {"id": 2}] + + +def test_first_match_over_a_streaming_decoder_silently_returns_the_wrong_answer(): + """The worst case: `first_match` is built for a first path that misses, and that path drains the body.""" + streaming_decoder = _streaming_decoder() + body = b'{"data": [{"id": 1}, {"id": 2}]}' + + streaming = CombinedExtractor( + extractors=[ + dpath("missing", decoder=streaming_decoder), + dpath("data", decoder=streaming_decoder), + ], + mode=CombineMode.first_match, + parameters=parameters, + ) + buffered = CombinedExtractor( + extractors=[dpath("missing"), dpath("data")], + mode=CombineMode.first_match, + parameters=parameters, + ) + + assert list(streaming.extract_records(create_response(body))) == [] + assert list(buffered.extract_records(create_response(body))) == [{"id": 1}, {"id": 2}] + + +def test_zip_merge_warns_when_a_sub_extractor_runs_out_early(caplog): + extractor = CombinedExtractor( + extractors=[ + _CountingExtractor(records=[{"a": 1}, {"a": 2}, {"a": 3}]), + _CountingExtractor(records=[{"b": 1}, {"b": 2}]), + ], + mode=CombineMode.zip_merge, + parameters=parameters, + ) + + with caplog.at_level(logging.WARNING, logger="airbyte"): + assert list(extractor.extract_records(create_response(REPORT_BODY))) == [ + {"a": 1, "b": 1}, + {"a": 2, "b": 2}, + ] + + assert any( + "zip_merge" in record.message and record.levelno == logging.WARNING + for record in caplog.records + ) + + +def test_zip_merge_does_not_warn_when_every_sub_extractor_has_the_same_length(caplog): + extractor = CombinedExtractor( + extractors=[dpath("rows", "*", "dimensions"), dpath("rows", "*", "metrics")], + mode=CombineMode.zip_merge, + parameters=parameters, + ) + + with caplog.at_level(logging.WARNING, logger="airbyte"): + assert len(list(extractor.extract_records(create_response(REPORT_BODY)))) == 2 + + assert [record.message for record in caplog.records] == [] + + +def test_zip_merge_names_the_sub_extractor_that_yielded_a_non_object(): + extractor = CombinedExtractor( + extractors=[ + _CountingExtractor(records=[{"a": 1}]), + _CountingExtractor(records=["not-an-object"]), + ], + mode=CombineMode.zip_merge, + parameters=parameters, + ) + + with pytest.raises(ValueError) as exc_info: + list(extractor.extract_records(create_response(REPORT_BODY))) + + message = str(exc_info.value) + assert "CombinedExtractor" in message + assert "zip_merge" in message + assert "sub-extractor 1" in message + assert "_CountingExtractor" in message 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 5f8ed7d4c8..763f35a74e 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 @@ -114,9 +114,24 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ExponentialBackoffStrategy as ExponentialBackoffStrategyModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + FileUploader as FileUploaderModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( GzipDecoder as GzipDecoderModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + IterableDecoder as IterableDecoderModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + JsonDecoder as JsonDecoderModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + JsonItemsDecoder as JsonItemsDecoderModel, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + JsonlDecoder as JsonlDecoderModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( OffsetIncrement as OffsetIncrementModel, ) @@ -126,6 +141,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( SelectiveAuthenticator, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + XmlDecoder as XmlDecoderModel, +) from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import ( ENV_VAR_ALLOW_CUSTOM_CODE, INJECTED_MANIFEST, @@ -6658,3 +6676,314 @@ def test_create_combined_extractor_defaults_to_union(): response = requests.Response() response._content = json.dumps({"a": [{"id": 1}], "b": [{"id": 2}]}).encode("utf-8") assert list(selector.extractor.extract_records(response)) == [{"id": 1}, {"id": 2}] + + +def _combined_extractor_selector_definition() -> Mapping[str, Any]: + content = """ + selector: + type: RecordSelector + extractor: + type: CombinedExtractor + extractors: + - type: DpathExtractor + field_path: ["a"] + - type: DpathExtractor + field_path: ["b"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + return transformer.propagate_types_and_parameters("", resolved_manifest["selector"], {}) + + +_STREAMING_DECODERS = { + "CsvDecoder": (CsvDecoderModel, {"type": "CsvDecoder"}), + "JsonlDecoder": (JsonlDecoderModel, {"type": "JsonlDecoder"}), + "JsonItemsDecoder": ( + JsonItemsDecoderModel, + {"type": "JsonItemsDecoder", "items_path": "items.item"}, + ), + "GzipDecoder": ( + GzipDecoderModel, + {"type": "GzipDecoder", "decoder": {"type": "JsonlDecoder"}}, + ), + "IterableDecoder": (IterableDecoderModel, {"type": "IterableDecoder"}), +} + + +@pytest.mark.parametrize("decoder_type", list(_STREAMING_DECODERS)) +@pytest.mark.parametrize("emit_connector_builder_messages", [False, True]) +def test_combined_extractor_over_a_streaming_decoder_is_rejected( + decoder_type: str, emit_connector_builder_messages: bool +): + """A streaming decoder lets every sub-extractor after the first read a closed response body. + + That loss is silent in production - see + `test_union_over_a_streaming_decoder_silently_drops_the_later_sub_extractors` - so the manifest + is rejected while it is being parsed. The Connector Builder downgrades these decoders to + `stream_response=False`, so it is rejected there too rather than letting the Builder test-read + a manifest that loses records once published. + """ + local_factory = ModelToComponentFactory( + emit_connector_builder_messages=emit_connector_builder_messages + ) + decoder_model, decoder_definition = _STREAMING_DECODERS[decoder_type] + decoder = local_factory.create_component( + model_type=decoder_model, + component_definition=decoder_definition, + config=input_config, + ) + + with pytest.raises(AirbyteTracedException) as exc_info: + local_factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=_combined_extractor_selector_definition(), + decoder=decoder, + transformations=[], + config=input_config, + ) + + assert exc_info.value.failure_type == FailureType.config_error + assert "CombinedExtractor" in exc_info.value.message + assert decoder_type in exc_info.value.internal_message or "CompositeRawDecoder" in ( + exc_info.value.internal_message + ) + + +@pytest.mark.parametrize("decoder_type", ["JsonDecoder", "XmlDecoder"]) +def test_combined_extractor_over_a_buffered_decoder_is_accepted(decoder_type: str): + decoder = factory.create_component( + model_type=JsonDecoderModel if decoder_type == "JsonDecoder" else XmlDecoderModel, + component_definition={"type": decoder_type}, + config=input_config, + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=_combined_extractor_selector_definition(), + decoder=decoder, + transformations=[], + config=input_config, + ) + + assert isinstance(selector.extractor, CombinedExtractor) + + +def test_combined_extractor_is_rejected_when_nested_under_another_combined_extractor(): + """Sub-extractors inherit the retriever's decoder, so the guard has to fire at every level.""" + content = """ + selector: + type: RecordSelector + extractor: + type: CombinedExtractor + extractors: + - type: CombinedExtractor + extractors: + - type: DpathExtractor + field_path: ["a"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + streaming_decoder = factory.create_component( + model_type=JsonlDecoderModel, + component_definition={"type": "JsonlDecoder"}, + config=input_config, + ) + + with pytest.raises(AirbyteTracedException): + factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=streaming_decoder, + transformations=[], + config=input_config, + ) + + +def test_combined_extractor_is_rejected_through_a_simple_retriever_with_a_streaming_decoder(): + """The decoder is never `None` on a manifest path, so the guard is reachable end to end.""" + content = """ + retriever: + type: SimpleRetriever + decoder: + type: JsonlDecoder + requester: + type: HttpRequester + url_base: "https://api.test.com" + path: "/records" + http_method: GET + record_selector: + type: RecordSelector + extractor: + type: CombinedExtractor + extractors: + - type: DpathExtractor + field_path: ["a"] + - type: DpathExtractor + field_path: ["b"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + retriever_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["retriever"], {} + ) + + with pytest.raises(AirbyteTracedException) as exc_info: + factory.create_component( + model_type=SimpleRetrieverModel, + component_definition=retriever_manifest, + config=input_config, + name="test_stream", + primary_key="id", + stream_slicer=None, + transformations=[], + ) + + assert exc_info.value.failure_type == FailureType.config_error + + +def test_combined_extractor_is_accepted_by_an_offset_increment_paginator(): + """`OffsetIncrement` builds a second copy of the extractor behind a `PaginationDecoderDecorator`.""" + content = """ + retriever: + type: SimpleRetriever + requester: + type: HttpRequester + url_base: "https://api.test.com" + path: "/records" + http_method: GET + paginator: + type: DefaultPaginator + pagination_strategy: + type: OffsetIncrement + page_size: 10 + record_selector: + type: RecordSelector + extractor: + type: CombinedExtractor + extractors: + - type: DpathExtractor + field_path: ["a"] + - type: DpathExtractor + field_path: ["b"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + retriever_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["retriever"], {} + ) + + retriever = factory.create_component( + model_type=SimpleRetrieverModel, + component_definition=retriever_manifest, + config=input_config, + name="test_stream", + primary_key="id", + stream_slicer=None, + transformations=[], + ) + + assert isinstance(retriever.record_selector.extractor, CombinedExtractor) + pagination_strategy = retriever.paginator.pagination_strategy + assert isinstance(pagination_strategy, OffsetIncrement) + # The paginator counts records with its own copy, so the response is traversed twice per page. + assert isinstance(pagination_strategy.extractor, CombinedExtractor) + assert pagination_strategy.extractor is not retriever.record_selector.extractor + + +def test_create_async_retriever_with_combined_extractors(): + """Covers the three `AsyncRetriever` unions widened by `CombinedExtractor`.""" + combined = { + "type": "CombinedExtractor", + "mode": "first_match", + "extractors": [ + {"type": "DpathExtractor", "field_path": ["primary"]}, + {"type": "DpathExtractor", "field_path": ["fallback"]}, + ], + } + definition = { + "type": "AsyncRetriever", + "status_mapping": { + "failed": ["failed"], + "running": ["pending"], + "timeout": ["timeout"], + "completed": ["ready"], + }, + "status_extractor": combined, + "download_target_extractor": combined, + "download_extractor": combined, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": ["data"]}, + }, + "polling_requester": { + "type": "HttpRequester", + "path": "/exports/{{ creation_response['id'] }}", + "url_base": "https://api.test.com", + "http_method": "GET", + }, + "creation_requester": { + "type": "HttpRequester", + "path": "/exports", + "url_base": "https://api.test.com", + "http_method": "POST", + }, + "download_requester": { + "type": "HttpRequester", + "path": "{{download_target}}", + "url_base": "", + "http_method": "GET", + }, + } + + component = factory.create_component( + model_type=AsyncRetrieverModel, + component_definition=definition, + name="test_stream", + primary_key="id", + stream_slicer=None, + transformations=[], + config={}, + ) + + job_repository = component.stream_slicer.job_orchestrator_factory( + [StreamSlice(partition={}, cursor_slice={})] + )._job_repository + assert isinstance(job_repository.status_extractor, CombinedExtractor) + assert isinstance(job_repository.download_target_extractor, CombinedExtractor) + assert isinstance( + job_repository.download_retriever.record_selector.extractor, CombinedExtractor + ) + + +def test_create_file_uploader_with_a_combined_download_target_extractor(): + definition = { + "type": "FileUploader", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "/download", + "http_method": "GET", + }, + "download_target_extractor": { + "type": "CombinedExtractor", + "mode": "first_match", + "extractors": [ + {"type": "DpathExtractor", "field_path": ["url"]}, + {"type": "DpathExtractor", "field_path": ["download_url"]}, + ], + }, + } + + file_uploader = factory.create_component( + model_type=FileUploaderModel, + component_definition=definition, + config=input_config, + ) + + assert isinstance(file_uploader.download_target_extractor, CombinedExtractor) diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index e2fef671c9..267836021d 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -6113,3 +6113,99 @@ def capturing_resolve(resolver_self, *args, **kwargs): "HttpComponentsResolver's requester should have api_budget set during dynamic stream " "discovery, but it was None. This means discovery HTTP requests are not rate-limited." ) + + +def _validation_messages(error: ValidationError) -> Iterable[str]: + """Flattens a jsonschema error tree; the top-level message of an `anyOf` miss carries no detail.""" + yield error.message + for sub_error in error.context or []: + yield from _validation_messages(sub_error) + + +def _combined_extractor_manifest(extractor: Mapping[str, Any]) -> Dict[str, Any]: + return { + "version": "6.0.0", + "type": "DeclarativeSource", + "check": {"type": "CheckStream", "stream_names": ["lists"]}, + "streams": [ + { + "type": "DeclarativeStream", + "name": "lists", + "primary_key": [], + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": { + "$schema": "http://json-schema.org/schema#", + "type": "object", + "properties": {}, + }, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "/lists", + "http_method": "GET", + }, + "record_selector": {"type": "RecordSelector", "extractor": extractor}, + }, + } + ], + } + + +def test_combined_extractor_manifest_passes_schema_validation(): + """Covers the CombinedExtractor JSON-Schema definition, which the factory tests never reach.""" + manifest = _combined_extractor_manifest( + { + "type": "CombinedExtractor", + "mode": "zip_merge", + "extractors": [ + {"type": "DpathExtractor", "field_path": ["rows", "*", "dimensions"]}, + { + "type": "CombinedExtractor", + "extractors": [ + {"type": "DpathExtractor", "field_path": ["rows", "*", "metrics"]} + ], + }, + ], + } + ) + + source = ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=create_catalog("lists"), state=None + ) + + assert len(source.streams(config={})) == 1 + + +def test_combined_extractor_with_an_unknown_mode_fails_schema_validation(): + manifest = _combined_extractor_manifest( + { + "type": "CombinedExtractor", + "mode": "concatenate", + "extractors": [{"type": "DpathExtractor", "field_path": ["rows"]}], + } + ) + + with pytest.raises(ValidationError): + ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=create_catalog("lists"), state=None + ) + + +def test_combined_extractor_with_an_empty_extractors_list_fails_schema_validation(): + """`minItems: 1` turns what used to be a runtime `ValueError` into a field-level schema error.""" + manifest = _combined_extractor_manifest({"type": "CombinedExtractor", "extractors": []}) + + with pytest.raises(ValidationError) as exc_info: + ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=create_catalog("lists"), state=None + ) + + # `_validate_source` re-raises a generic message; the field-level error is in the cause tree. + assert any( + "should be non-empty" in message + for message in _validation_messages(exc_info.value.__cause__) + ) From 28da58ef18e964e123a670bbcb9f2e87a93455db Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Tue, 15 Sep 2026 23:04:26 +0300 Subject: [PATCH 03/10] fix(CombinedExtractor): reject union under OffsetIncrement, rename mode enum - Reject a `union` CombinedExtractor (nested ones included) under an `OffsetIncrement` paginator with a config error: the offset is advanced by the summed record count of all sub-extractors, which overshoots the API page size and silently skips records. - Correct the pagination caveat in the `mode` schema description and in the component docstring: `union` is the harmful mode, `first_match` and `zip_merge` report the winning/shortest sub-extractor's count. - Narrow the Connector Builder half of the streaming-decoder guard to the four decoders the Builder actually downgrades (exact `CompositeRawDecoder` class with a streaming parser), so a buffered subclass is no longer rejected. - Move the `mode` enum into its own `CombineMode` schema definition so codegen emits `CombineMode` instead of the very generic top-level `Mode`; the generated model matches `datamodel_code_generator==0.26.3` output. - Keep both new user-facing messages to the violated condition and move the remediation into `internal_message`, per the error-message guidelines. - Ignore the `test_response.csv` fixture that `test_response_to_file_extractor.py` writes under `unit_tests/sources/declarative/extractors/`, mirroring the `decoders/` rule. Co-Authored-By: Claude Opus 5 (1M context) --- .../declarative_component_schema.yaml | 24 ++- .../extractors/combined_extractor.py | 23 +- .../models/declarative_component_schema.py | 19 +- .../parsers/model_to_component_factory.py | 86 +++++++- .../sources/declarative/extractors/.gitignore | 1 + .../test_model_to_component_factory.py | 197 ++++++++++++++++-- 6 files changed, 301 insertions(+), 49 deletions(-) create mode 100644 unit_tests/sources/declarative/extractors/.gitignore diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 55718cfd24..ebbbfee542 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2173,6 +2173,18 @@ definitions: $parameters: type: object additionalProperties: true + CombineMode: + title: Combine Mode + description: 'How the records of the sub-extractors of a CombinedExtractor are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones.' + type: string + enum: + - union + - first_match + - zip_merge + examples: + - union + - first_match + - zip_merge CombinedExtractor: title: Combined Extractor description: Record extractor that combines the output of several sub-extractors. Use it when a single response holds records under more than one path, when the records can be found under one of several alternative paths, or when a record has to be assembled from several parallel lists. @@ -2196,17 +2208,9 @@ definitions: - "$ref": "#/definitions/CombinedExtractor" mode: title: Combine Mode - description: 'How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that an OffsetIncrement or PageIncrement paginator counts the combined records, which under "first_match" and "zip_merge" is not the number of records the API returned for the page.' - type: string - enum: - - union - - first_match - - zip_merge + description: 'How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that a paginator which counts the records of a page counts the combined records: under "union" that is the sum over all sub-extractors, which overshoots the API page size, so "union" is rejected with an OffsetIncrement paginator because the offset would skip records. Under "first_match" the count is the count of the winning sub-extractor and under "zip_merge" the count of the shortest one, which are usually the number of records the API returned for the page.' + "$ref": "#/definitions/CombineMode" default: union - examples: - - union - - first_match - - zip_merge $parameters: type: object additionalProperties: true diff --git a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py index 28f34b9bb8..b4ad6741fd 100644 --- a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py +++ b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py @@ -97,7 +97,8 @@ class CombinedExtractor(RecordExtractor): `first_match` emits nothing at all when the first path misses. Because of that, `ModelToComponentFactory.create_combined_extractor` refuses to build this - component over a streaming decoder and raises a configuration error naming the decoder. The + component over a streaming decoder and raises a configuration error, naming the decoder in the + internal message. The Connector Builder forces those same decoders to `stream_response=False`, so the rejection is applied there too rather than letting a manifest test-read correctly and lose records once published. @@ -110,9 +111,23 @@ class CombinedExtractor(RecordExtractor): extractors over a large response costs roughly three times the decode time of a single one. Note that an `OffsetIncrement` or `PageIncrement` paginator builds its own copy of the - extractor to count the records of a page, which doubles that cost, and the count it obtains is - the combined count: the winner's record count under `first_match` and the shortest - sub-extractor's count under `zip_merge`, neither of which is the API's page size. + extractor to count the records of a page, which doubles that cost. + + ## Record-counting paginators + + The count those paginators obtain is the combined count, not the API's page size: + + - `union` returns the SUM over all sub-extractors. An `OffsetIncrement` advances the offset by + that sum, so the next page starts past the records that were never read — two sub-extractors + returning two records each with `page_size: 2` request offsets 0, 4, 8 instead of 0, 2, 4 and + silently drop two thirds of the records. `ModelToComponentFactory.create_offset_increment` + therefore rejects a `union` `CombinedExtractor` (nested ones included) with a configuration + error. A `PageIncrement` only compares the count against `page_size` to decide whether to + stop, so it loses nothing, but it issues one extra request whenever the summed count of the + last page happens to equal `page_size`. + - `first_match` returns the count of the winning sub-extractor and `zip_merge` the count of the + shortest one. Neither inflates the count, so both work with either paginator, as long as the + path the count comes from is the one the API paginates over. Attributes: extractors (List[RecordExtractor]): The sub-extractors to combine. At least one is required. diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 98d47b7254..07d9255d9a 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -507,6 +507,12 @@ class Config: ) +class CombineMode(Enum): + union = "union" + first_match = "first_match" + zip_merge = "zip_merge" + + class ResponseToFileExtractor(BaseModel): type: Literal["ResponseToFileExtractor"] preserve_na_values: Optional[bool] = Field( @@ -2517,12 +2523,6 @@ class DpathExtractor(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class Mode(Enum): - union = "union" - first_match = "first_match" - zip_merge = "zip_merge" - - class CombinedExtractor(BaseModel): type: Literal["CombinedExtractor"] extractors: List[Union[DpathExtractor, CustomRecordExtractor, CombinedExtractor]] = Field( @@ -2530,10 +2530,9 @@ class CombinedExtractor(BaseModel): description="The record extractors to combine. At least one is required. Each sub-extractor is given the same HTTP response and decodes it independently, so the response is parsed once per sub-extractor. Streaming decoders (CsvDecoder, JsonlDecoder, JsonItemsDecoder, GzipDecoder, IterableDecoder) can only read the response body once and are rejected.", title="Extractors", ) - mode: Optional[Mode] = Field( - Mode.union, - description='How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that an OffsetIncrement or PageIncrement paginator counts the combined records, which under "first_match" and "zip_merge" is not the number of records the API returned for the page.', - examples=["union", "first_match", "zip_merge"], + mode: Optional[CombineMode] = Field( + CombineMode.union, + description='How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that a paginator which counts the records of a page counts the combined records: under "union" that is the sum over all sub-extractors, which overshoots the API page size, so "union" is rejected with an OffsetIncrement paginator because the offset would skip records. Under "first_match" the count is the count of the winning sub-extractor and under "zip_merge" the count of the shortest one, which are usually the number of records the API returned for the page.', title="Combine Mode", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") 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 c5826c11d9..bc140c2cc6 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2481,6 +2481,25 @@ def create_default_paginator( return PaginatorTestReadDecorator(paginator, self._limit_pages_fetched_per_slice) return paginator + @staticmethod + def _is_decoder_downgraded_by_connector_builder(decoder: Decoder) -> bool: + """Is this the buffered stand-in the Connector Builder builds for a streaming decoder? + + `create_csv_decoder`, `create_jsonl_decoder`, `create_json_items_decoder` and + `create_gzip_decoder` build a `CompositeRawDecoder` with `stream_response=False` when + `_emit_connector_builder_messages` is set, so a Builder test read can be replayed. Those + four instances are the only ones that stream in production while reporting + `is_stream_response() == False` in the Builder. + + The match is on the exact class and on the parser, not on `isinstance`: a `CustomDecoder` + subclassing `CompositeRawDecoder` with `stream_response=False`, or a buffered + `CompositeRawDecoder` built around a `JsonParser`, reads the body from `response.content` + in production too and must not be rejected. + """ + return type(decoder) is CompositeRawDecoder and isinstance( + decoder.parser, (CsvParser, JsonLineParser, JsonItemsParser, GzipParser) + ) + def _reject_combined_extractor_over_streaming_decoder(self, decoder: Optional[Decoder]) -> None: """Refuse to build a `CombinedExtractor` whose response body can only be read once. @@ -2500,23 +2519,20 @@ def _reject_combined_extractor_over_streaming_decoder(self, decoder: Optional[De decoder.decoder if isinstance(decoder, PaginationDecoderDecorator) else decoder ) streams_in_production = inner_decoder.is_stream_response() or ( - # The four decoders the Builder downgrades are the only bare `CompositeRawDecoder`s - # this factory builds; every buffered decoder is a distinct class. - self._emit_connector_builder_messages and isinstance(inner_decoder, CompositeRawDecoder) + self._emit_connector_builder_messages + and self._is_decoder_downgraded_by_connector_builder(inner_decoder) ) if not streams_in_production: return raise AirbyteTracedException( - message=( - "CombinedExtractor is not supported with a streaming decoder (CsvDecoder, " - "JsonlDecoder, JsonItemsDecoder, GzipDecoder, IterableDecoder), which reads the " - "response body only once. Use JsonDecoder, XmlDecoder or ZipfileDecoder, or " - "declare a single extractor." - ), + message="CombinedExtractor is not supported with a streaming decoder.", internal_message=( f"CombinedExtractor was configured with {type(inner_decoder).__name__}, which " - f"streams the response. Sub-extractors after the first would read a closed " - f"response and silently return no records." + f"streams the response and can only be read once. Sub-extractors after the first " + f"would read a closed response and silently return no records. The streaming " + f"decoders are CsvDecoder, JsonlDecoder, JsonItemsDecoder, GzipDecoder and " + f"IterableDecoder; use JsonDecoder, XmlDecoder or ZipfileDecoder, or declare a " + f"single extractor." ), failure_type=FailureType.config_error, ) @@ -3134,6 +3150,52 @@ def _get_refresh_token_error_information( # returning default values we think cover most cases return (400,), "error", ("invalid_grant", "invalid_permissions") + @staticmethod + def _reject_union_combined_extractor_for_offset_increment( + extractor_model: Optional[BaseModel], + ) -> None: + """Refuse an `OffsetIncrement` paginator driven by a `union` `CombinedExtractor`. + + `OffsetIncrement.next_page_token` advances the offset by the number of records its + extractor returns for the page. Under `union` that number is the sum over all + sub-extractors, so the offset overshoots the API page size and every page after the first + starts past the records that were never read: with two sub-extractors returning two records + each and `page_size: 2`, the requested offsets are 0, 4, 8 instead of 0, 2, 4 and two + thirds of the records are silently dropped. + + `first_match` (the winning sub-extractor's count) and `zip_merge` (the shortest + sub-extractor's count) do not inflate the count and are left alone. A `union` nested + anywhere in the tree inflates the count of the node above it, so the whole tree is walked. + """ + if not isinstance(extractor_model, CombinedExtractorModel): + return + if not ModelToComponentFactory._combined_extractor_tree_contains_union(extractor_model): + return + raise AirbyteTracedException( + message=( + 'CombinedExtractor mode "union" is not supported with an OffsetIncrement paginator.' + ), + internal_message=( + "OffsetIncrement counts the records of its extractor to advance the offset. A " + "`union` CombinedExtractor returns the sum of its sub-extractors' records, which " + "overshoots the page the API returned, so records would be skipped. Use the " + "`first_match` or `zip_merge` mode, a CursorPagination or PageIncrement " + "paginator, or a single extractor." + ), + failure_type=FailureType.config_error, + ) + + @staticmethod + def _combined_extractor_tree_contains_union(model: CombinedExtractorModel) -> bool: + mode = CombineMode(model.mode.value) if model.mode else CombineMode.union + if mode == CombineMode.union: + return True + return any( + isinstance(sub_extractor, CombinedExtractorModel) + and ModelToComponentFactory._combined_extractor_tree_contains_union(sub_extractor) + for sub_extractor in model.extractors + ) + def create_offset_increment( self, model: OffsetIncrementModel, @@ -3157,6 +3219,8 @@ def create_offset_increment( self._UNSUPPORTED_DECODER_ERROR.format(decoder_type=type(inner_decoder)) ) + self._reject_union_combined_extractor_for_offset_increment(extractor_model) + # Ideally we would instantiate the runtime extractor from highest most level (in this case the SimpleRetriever) # so that it can be shared by OffSetIncrement and RecordSelector. However, due to how we instantiate the # decoder with various decorators here, but not in create_record_selector, it is simpler to retain existing diff --git a/unit_tests/sources/declarative/extractors/.gitignore b/unit_tests/sources/declarative/extractors/.gitignore new file mode 100644 index 0000000000..7b1b301200 --- /dev/null +++ b/unit_tests/sources/declarative/extractors/.gitignore @@ -0,0 +1 @@ +test_response.csv 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 763f35a74e..882ce3689a 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 @@ -54,6 +54,10 @@ from airbyte_cdk.sources.declarative.concurrency_level import ConcurrencyLevel from airbyte_cdk.sources.declarative.datetime.min_max_datetime import MinMaxDatetime from airbyte_cdk.sources.declarative.decoders import JsonDecoder, PaginationDecoderDecorator +from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import ( + CompositeRawDecoder, + JsonParser, +) from airbyte_cdk.sources.declarative.extractors import ( CombinedExtractor, CombineMode, @@ -6695,6 +6699,10 @@ def _combined_extractor_selector_definition() -> Mapping[str, Any]: return transformer.propagate_types_and_parameters("", resolved_manifest["selector"], {}) +class _BufferedCompositeRawDecoderSubclass(CompositeRawDecoder): + """Stands in for a `CustomDecoder` that subclasses `CompositeRawDecoder` and buffers.""" + + _STREAMING_DECODERS = { "CsvDecoder": (CsvDecoderModel, {"type": "CsvDecoder"}), "JsonlDecoder": (JsonlDecoderModel, {"type": "JsonlDecoder"}), @@ -6770,6 +6778,42 @@ def test_combined_extractor_over_a_buffered_decoder_is_accepted(decoder_type: st assert isinstance(selector.extractor, CombinedExtractor) +@pytest.mark.parametrize( + "decoder", + [ + pytest.param( + CompositeRawDecoder(parser=JsonParser(), stream_response=False), + id="buffered_composite_raw_decoder", + ), + pytest.param( + _BufferedCompositeRawDecoderSubclass(parser=JsonParser(), stream_response=False), + id="buffered_composite_raw_decoder_subclass", + ), + ], +) +def test_combined_extractor_over_a_buffered_composite_raw_decoder_is_accepted_in_the_builder( + decoder: CompositeRawDecoder, +): + """The Builder rejection only covers the decoders the Builder itself downgrades. + + A `CompositeRawDecoder` that reads `response.content` in production too - a custom decoder + subclassing it, or one built around a buffering parser - can be read by every sub-extractor and + must not be rejected just because the Builder is emitting messages. + """ + local_factory = ModelToComponentFactory(emit_connector_builder_messages=True) + + selector = local_factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=_combined_extractor_selector_definition(), + decoder=decoder, + transformations=[], + config=input_config, + ) + + assert isinstance(selector.extractor, CombinedExtractor) + + def test_combined_extractor_is_rejected_when_nested_under_another_combined_extractor(): """Sub-extractors inherit the retriever's decoder, so the guard has to fire at every level.""" content = """ @@ -6847,9 +6891,10 @@ def test_combined_extractor_is_rejected_through_a_simple_retriever_with_a_stream assert exc_info.value.failure_type == FailureType.config_error -def test_combined_extractor_is_accepted_by_an_offset_increment_paginator(): - """`OffsetIncrement` builds a second copy of the extractor behind a `PaginationDecoderDecorator`.""" - content = """ +def _retriever_manifest_with_paginator( + extractor_definition: str, pagination_strategy: str +) -> Mapping[str, Any]: + content = f""" retriever: type: SimpleRetriever requester: @@ -6860,25 +6905,50 @@ def test_combined_extractor_is_accepted_by_an_offset_increment_paginator(): paginator: type: DefaultPaginator pagination_strategy: - type: OffsetIncrement - page_size: 10 +{pagination_strategy} record_selector: type: RecordSelector extractor: - type: CombinedExtractor +{extractor_definition} + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + return transformer.propagate_types_and_parameters("", resolved_manifest["retriever"], {}) + + +_OFFSET_INCREMENT = """ type: OffsetIncrement + page_size: 2""" +_PAGE_INCREMENT = """ type: PageIncrement + page_size: 2""" +_UNION_EXTRACTOR_DEFAULT_MODE = """ type: CombinedExtractor extractors: - type: DpathExtractor field_path: ["a"] - type: DpathExtractor - field_path: ["b"] - """ - parsed_manifest = YamlDeclarativeSource._parse(content) - resolved_manifest = resolver.preprocess_manifest(parsed_manifest) - retriever_manifest = transformer.propagate_types_and_parameters( - "", resolved_manifest["retriever"], {} - ) + field_path: ["b"]""" +_UNION_EXTRACTOR_EXPLICIT_MODE = """ type: CombinedExtractor + mode: union + extractors: + - type: DpathExtractor + field_path: ["a"] + - type: DpathExtractor + field_path: ["b"]""" +_NESTED_UNION_EXTRACTOR = """ type: CombinedExtractor + mode: first_match + extractors: + - type: CombinedExtractor + mode: union + extractors: + - type: DpathExtractor + field_path: ["a"] + - type: DpathExtractor + field_path: ["b"] + - type: DpathExtractor + field_path: ["c"]""" - retriever = factory.create_component( + +def _build_retriever(retriever_manifest: Mapping[str, Any]) -> SimpleRetriever: + return factory.create_component( model_type=SimpleRetrieverModel, component_definition=retriever_manifest, config=input_config, @@ -6888,6 +6958,55 @@ def test_combined_extractor_is_accepted_by_an_offset_increment_paginator(): transformations=[], ) + +def _json_response(body: Mapping[str, Any]) -> requests.Response: + response = requests.Response() + response.status_code = 200 + response._content = json.dumps(body).encode("utf-8") + return response + + +@pytest.mark.parametrize( + "extractor_definition", + [ + pytest.param(_UNION_EXTRACTOR_DEFAULT_MODE, id="union_by_default"), + pytest.param(_UNION_EXTRACTOR_EXPLICIT_MODE, id="union_declared"), + pytest.param(_NESTED_UNION_EXTRACTOR, id="union_nested_under_first_match"), + ], +) +def test_union_combined_extractor_is_rejected_by_an_offset_increment_paginator( + extractor_definition: str, +): + """`OffsetIncrement` advances the offset by the record count of its extractor. + + Under `union` that count is the sum over all sub-extractors, so the offset overshoots the page + the API returned and the records in between are never requested. A `union` nested anywhere in + the tree inflates the count of the node above it, so it is rejected as well. + """ + with pytest.raises(AirbyteTracedException) as exc_info: + _build_retriever( + _retriever_manifest_with_paginator(extractor_definition, _OFFSET_INCREMENT) + ) + + assert exc_info.value.failure_type == FailureType.config_error + assert "union" in exc_info.value.message + assert "OffsetIncrement" in exc_info.value.message + + +def test_first_match_combined_extractor_counts_only_the_winning_sub_extractor_for_the_offset(): + """`first_match` is accepted because the count it reports is the winner's, not a sum.""" + extractor_definition = """ type: CombinedExtractor + mode: first_match + extractors: + - type: DpathExtractor + field_path: ["a"] + - type: DpathExtractor + field_path: ["b"]""" + + retriever = _build_retriever( + _retriever_manifest_with_paginator(extractor_definition, _OFFSET_INCREMENT) + ) + assert isinstance(retriever.record_selector.extractor, CombinedExtractor) pagination_strategy = retriever.paginator.pagination_strategy assert isinstance(pagination_strategy, OffsetIncrement) @@ -6895,6 +7014,56 @@ def test_combined_extractor_is_accepted_by_an_offset_increment_paginator(): assert isinstance(pagination_strategy.extractor, CombinedExtractor) assert pagination_strategy.extractor is not retriever.record_selector.extractor + # `a` misses, so the winner is `b` with its two records: the offset advances by 2, not by 4. + response = _json_response({"a": [], "b": [{"id": 1}, {"id": 2}]}) + assert ( + pagination_strategy.next_page_token( + response=response, last_page_size=2, last_record=None, last_page_token_value=0 + ) + == 2 + ) + + +def test_zip_merge_combined_extractor_counts_the_shortest_sub_extractor_for_the_offset(): + extractor_definition = """ type: CombinedExtractor + mode: zip_merge + extractors: + - type: DpathExtractor + field_path: ["a"] + - type: DpathExtractor + field_path: ["b"]""" + + retriever = _build_retriever( + _retriever_manifest_with_paginator(extractor_definition, _OFFSET_INCREMENT) + ) + + pagination_strategy = retriever.paginator.pagination_strategy + assert isinstance(pagination_strategy, OffsetIncrement) + + # Two merged records, not the four records the two sub-extractors yield together. + response = _json_response({"a": [{"x": 1}, {"x": 2}], "b": [{"y": 1}, {"y": 2}]}) + assert ( + pagination_strategy.next_page_token( + response=response, last_page_size=2, last_record=None, last_page_token_value=0 + ) + == 2 + ) + + +def test_union_combined_extractor_is_accepted_by_a_page_increment_paginator(): + """`PageIncrement` only compares the count against `page_size`, so an inflated count is not lossy. + + It can cost one extra request when the summed count of the last page equals `page_size`, which + is documented on the component rather than rejected. + """ + retriever = _build_retriever( + _retriever_manifest_with_paginator(_UNION_EXTRACTOR_DEFAULT_MODE, _PAGE_INCREMENT) + ) + + pagination_strategy = retriever.paginator.pagination_strategy + assert isinstance(pagination_strategy, PageIncrement) + assert isinstance(pagination_strategy.extractor, CombinedExtractor) + def test_create_async_retriever_with_combined_extractors(): """Covers the three `AsyncRetriever` unions widened by `CombinedExtractor`.""" From 502336821d8d46104f44caa98d2ccd0cc8e685fe Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Wed, 16 Sep 2026 00:03:10 +0300 Subject: [PATCH 04/10] feat(CombinedExtractor): add `skip_empty_records` so `first_match` can fall through a page of nulls `source-monday`'s `MondayIncrementalItemsExtractor` was described as a motivating connector for this component but was not a drop-in adopter, for one reason: it drops falsy entries before deciding whether the primary path produced anything (`components.py:106-124`). A GraphQL partial response puts `null` in the record list and explains itself in a sibling `errors` field, and monday has to fall through to `field_path_pagination` when a page is all nulls. `first_match` counted a `None` as a record, so the primary path won with nothing usable and the fallback was never tried. `skip_empty_records` drops falsy records - `None`, `{}`, `[]`, `""` - before any mode sees them, which both keeps them out of the record stream and takes them out of the match decision. It is off by default, so existing manifests are unaffected. The flag applies to all three modes. Under `zip_merge` it shifts the alignment of the merged records, because dropping the i-th record of one sub-extractor pairs its next record with the i-th record of the others; the schema description and a test say so rather than leaving it to be discovered. One divergence from monday remains and is deliberate: monday's warning names `body.get("errors")`, which assumes a GraphQL error envelope. The component logs the dropped count and the sub-extractor that produced them instead. The generated model was hand-applied and then verified against the pinned generator: the `CombinedExtractor` block is byte-identical to `poe assemble` output once `ruff format` is applied. Co-Authored-By: Claude Opus 5 (1M context) --- .../declarative_component_schema.yaml | 6 + .../extractors/combined_extractor.py | 62 ++++++- .../models/declarative_component_schema.py | 5 + .../parsers/model_to_component_factory.py | 1 + .../extractors/test_combined_extractor.py | 157 ++++++++++++++++++ .../test_model_to_component_factory.py | 74 +++++++++ 6 files changed, 300 insertions(+), 5 deletions(-) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index ebbbfee542..34b9d2eb19 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2211,6 +2211,11 @@ definitions: description: 'How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that a paginator which counts the records of a page counts the combined records: under "union" that is the sum over all sub-extractors, which overshoots the API page size, so "union" is rejected with an OffsetIncrement paginator because the offset would skip records. Under "first_match" the count is the count of the winning sub-extractor and under "zip_merge" the count of the shortest one, which are usually the number of records the API returned for the page.' "$ref": "#/definitions/CombineMode" default: union + skip_empty_records: + title: Skip Empty Records + description: 'Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried, and the record count a paginator obtains is the count after the empty records were dropped. Under "zip_merge" it shifts the alignment of the merged records, because dropping the i-th record of one sub-extractor pairs its next record with the i-th record of the others.' + type: boolean + default: false $parameters: type: object additionalProperties: true @@ -2224,6 +2229,7 @@ definitions: field_path: ["data", "repository", "pullRequests", "nodes"] - type: CombinedExtractor mode: first_match + skip_empty_records: true extractors: - type: DpathExtractor field_path: ["data", "boards", "*", "items_page", "items"] diff --git a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py index b4ad6741fd..e981dec913 100644 --- a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py +++ b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py @@ -103,6 +103,25 @@ class CombinedExtractor(RecordExtractor): applied there too rather than letting a manifest test-read correctly and lose records once published. + ## Empty records + + `skip_empty_records` drops falsy records — `None`, `{}`, `[]`, `""` — before any of the modes + sees them. It is off by default, so a sub-extractor's output is passed through verbatim. + + Turn it on when an API can answer with nulls in the middle of a record list. A GraphQL API + that returns a partial response puts `null` in the `data` array for the fields it failed to + resolve and describes the failure in a sibling `errors` field; without this flag those nulls + reach the record stream and fail schema validation downstream. + + Under `first_match` the flag also changes which sub-extractor wins: a sub-extractor whose + records are all empty no longer counts as a match, so the next sub-extractor is tried. That is + the behaviour of `source-monday`'s `MondayIncrementalItemsExtractor`, where a page of nulls + must fall through to the pagination path rather than lock in the primary path. + + Under `zip_merge` the flag shifts alignment: dropping the i-th record of one sub-extractor + pairs its i+1-th record with the i-th record of the others. Only enable it there if the + sub-extractors drop records in lockstep. + ## Cost Each sub-extractor decodes the response independently — `CompositeRawDecoder` re-parses the @@ -132,11 +151,14 @@ class CombinedExtractor(RecordExtractor): Attributes: extractors (List[RecordExtractor]): The sub-extractors to combine. At least one is required. mode (CombineMode): How the sub-extractor outputs are combined. Defaults to `union`. + skip_empty_records (bool): Whether falsy records are dropped before they are combined. + Defaults to `False`. """ extractors: List[RecordExtractor] parameters: InitVar[Mapping[str, Any]] mode: CombineMode = CombineMode.union + skip_empty_records: bool = False def __post_init__(self, parameters: Mapping[str, Any]) -> None: if not self.extractors: @@ -155,12 +177,12 @@ def extract_records(self, response: requests.Response) -> Iterable[Mapping[str, yield from self._extract_union(response) def _extract_union(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: - for extractor in self.extractors: - yield from extractor.extract_records(response) + for index in range(len(self.extractors)): + yield from self._records_of(index, response) def _extract_first_match(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: - for extractor in self.extractors: - records: Iterator[Mapping[str, Any]] = iter(extractor.extract_records(response)) + for index in range(len(self.extractors)): + records: Iterator[Mapping[str, Any]] = iter(self._records_of(index, response)) first_record: Union[Mapping[str, Any], _NoRecord] = next(records, _NO_RECORD) if isinstance(first_record, _NoRecord): continue @@ -173,7 +195,9 @@ def _extract_zip_merge(self, response: requests.Response) -> Iterable[Mapping[st # Iteration stops at the shortest sub-extractor, matching the behavior this component # replaces. Unlike plain `zip`, the round in which a sub-extractor runs out is completed so # that an uneven number of records can be reported instead of silently truncated. - iterators = [iter(extractor.extract_records(response)) for extractor in self.extractors] + iterators = [ + iter(self._records_of(index, response)) for index in range(len(self.extractors)) + ] emitted = 0 while True: round_records = [next(iterator, _NO_RECORD) for iterator in iterators] @@ -194,6 +218,34 @@ def _extract_zip_merge(self, response: requests.Response) -> Iterable[Mapping[st yield self._merge(round_records) emitted += 1 + def _records_of( + self, extractor_index: int, response: requests.Response + ) -> Iterable[Mapping[str, Any]]: + records = self.extractors[extractor_index].extract_records(response) + if not self.skip_empty_records: + return records + return self._without_empty_records(records, extractor_index) + + def _without_empty_records( + self, records: Iterable[Mapping[str, Any]], extractor_index: int + ) -> Iterator[Mapping[str, Any]]: + dropped = 0 + for record in records: + if not record: + dropped += 1 + continue + yield record + if dropped: + logger.warning( + "CombinedExtractor dropped %s empty record(s) yielded by sub-extractor %s (%s) " + "because `skip_empty_records` is enabled. An API that returns empty records " + "alongside real ones usually reports the reason in an error field of the " + "response body.", + dropped, + extractor_index, + type(self.extractors[extractor_index]).__name__, + ) + def _merge(self, records: Sequence[Any]) -> Mapping[str, Any]: merged: Dict[str, Any] = {} for index, record in enumerate(records): diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 07d9255d9a..bf57750946 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2535,6 +2535,11 @@ class CombinedExtractor(BaseModel): description='How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that a paginator which counts the records of a page counts the combined records: under "union" that is the sum over all sub-extractors, which overshoots the API page size, so "union" is rejected with an OffsetIncrement paginator because the offset would skip records. Under "first_match" the count is the count of the winning sub-extractor and under "zip_merge" the count of the shortest one, which are usually the number of records the API returned for the page.', title="Combine Mode", ) + skip_empty_records: Optional[bool] = Field( + False, + description='Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried, and the record count a paginator obtains is the count after the empty records were dropped. Under "zip_merge" it shifts the alignment of the merged records, because dropping the i-th record of one sub-extractor pairs its next record with the i-th record of the others.', + title="Skip Empty Records", + ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") 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 bc140c2cc6..2afd392fbf 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2552,6 +2552,7 @@ def create_combined_extractor( return CombinedExtractor( extractors=extractors, mode=CombineMode(model.mode.value) if model.mode else CombineMode.union, + skip_empty_records=bool(model.skip_empty_records), parameters=model.parameters or {}, ) diff --git a/unit_tests/sources/declarative/extractors/test_combined_extractor.py b/unit_tests/sources/declarative/extractors/test_combined_extractor.py index 21cab56510..9bab3eb7ef 100644 --- a/unit_tests/sources/declarative/extractors/test_combined_extractor.py +++ b/unit_tests/sources/declarative/extractors/test_combined_extractor.py @@ -39,6 +39,27 @@ } } +# The source-monday shape: a GraphQL partial response puts nulls where it failed to resolve an +# item and explains itself in a sibling `errors` field. The primary path therefore holds records +# that are all null, and the connector has to fall through to the pagination path. +MONDAY_NULL_PAGE_BODY = { + "data": { + "boards": [{"items_page": {"items": [None, None]}}], + "next_items_page": {"items": [{"id": "item_1"}]}, + }, + "errors": [{"message": "Item not found"}], +} + +# The same shape, but the primary path resolved one of its two items. The primary path wins and +# the pagination path is never read. +MONDAY_PARTIAL_PAGE_BODY = { + "data": { + "boards": [{"items_page": {"items": [{"id": "item_1"}, None]}}], + "next_items_page": {"items": [{"id": "item_9"}]}, + }, + "errors": [{"message": "Item not found"}], +} + # The GA4 report shape: dimensions and metrics come back as two parallel lists that have to be # merged position by position. REPORT_BODY = { @@ -423,3 +444,139 @@ def test_zip_merge_names_the_sub_extractor_that_yielded_a_non_object(): assert "zip_merge" in message assert "sub-extractor 1" in message assert "_CountingExtractor" in message + + +def _monday_extractor(skip_empty_records: bool) -> CombinedExtractor: + return CombinedExtractor( + extractors=[ + dpath("data", "boards", "*", "items_page", "items", "*"), + dpath("data", "next_items_page", "items", "*"), + ], + mode=CombineMode.first_match, + skip_empty_records=skip_empty_records, + parameters=parameters, + ) + + +def test_skip_empty_records_is_off_by_default(): + extractor = CombinedExtractor( + extractors=[_CountingExtractor(records=[{"id": 1}, None, {}])], + parameters=parameters, + ) + + assert extractor.skip_empty_records is False + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [{"id": 1}, None, {}] + + +def test_skip_empty_records_drops_falsy_records_under_union(): + extractor = CombinedExtractor( + extractors=[ + _CountingExtractor(records=[{"id": 1}, None, {}]), + _CountingExtractor(records=[[], "", {"id": 2}]), + ], + mode=CombineMode.union, + skip_empty_records=True, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [{"id": 1}, {"id": 2}] + + +def test_first_match_locks_in_a_page_of_nulls_without_skip_empty_records(): + """Pins the behaviour `skip_empty_records` exists to change. + + Without the flag a null counts as a record, so the primary path wins with nothing usable and + the pagination path is never tried. + """ + records = list( + _monday_extractor(skip_empty_records=False).extract_records( + create_response(MONDAY_NULL_PAGE_BODY) + ) + ) + + assert records == [None, None] + + +def test_first_match_falls_through_a_page_of_nulls_with_skip_empty_records(): + records = list( + _monday_extractor(skip_empty_records=True).extract_records( + create_response(MONDAY_NULL_PAGE_BODY) + ) + ) + + assert records == [{"id": "item_1"}] + + +def test_first_match_keeps_the_winner_when_some_of_its_records_survive(): + records = list( + _monday_extractor(skip_empty_records=True).extract_records( + create_response(MONDAY_PARTIAL_PAGE_BODY) + ) + ) + + assert records == [{"id": "item_1"}] + + +def test_skip_empty_records_does_not_touch_extractors_after_the_winner(): + loser = _CountingExtractor(records=[{"id": 99}]) + extractor = CombinedExtractor( + extractors=[_CountingExtractor(records=[None, {"id": 1}]), loser], + mode=CombineMode.first_match, + skip_empty_records=True, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [{"id": 1}] + assert loser.calls == [] + + +def test_skip_empty_records_warns_and_names_the_sub_extractor(caplog): + extractor = CombinedExtractor( + extractors=[_CountingExtractor(records=[{"id": 1}, None, None])], + mode=CombineMode.union, + skip_empty_records=True, + parameters=parameters, + ) + + with caplog.at_level(logging.WARNING, logger="airbyte"): + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [{"id": 1}] + + warnings = [record.message for record in caplog.records if record.levelno == logging.WARNING] + assert len(warnings) == 1 + assert "dropped 2 empty record(s)" in warnings[0] + assert "sub-extractor 0 (_CountingExtractor)" in warnings[0] + + +def test_skip_empty_records_does_not_warn_when_nothing_is_dropped(caplog): + extractor = CombinedExtractor( + extractors=[dpath("data", "repository", "issues", "nodes")], + mode=CombineMode.union, + skip_empty_records=True, + parameters=parameters, + ) + + with caplog.at_level(logging.WARNING, logger="airbyte"): + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [ + {"id": "I_1", "title": "first issue"} + ] + + assert [record.message for record in caplog.records if record.levelno == logging.WARNING] == [] + + +def test_skip_empty_records_shifts_zip_merge_alignment(): + """Documents the caveat the schema description warns about. + + Dropping the first sub-extractor's null pairs its next record with the other sub-extractor's + first record, so the merge is off by one rather than skipping a row. + """ + extractor = CombinedExtractor( + extractors=[ + _CountingExtractor(records=[None, {"a": 2}]), + _CountingExtractor(records=[{"b": 1}, {"b": 2}]), + ], + mode=CombineMode.zip_merge, + skip_empty_records=True, + parameters=parameters, + ) + + assert list(extractor.extract_records(create_response(REPORT_BODY))) == [{"a": 2, "b": 1}] 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 882ce3689a..23c6a5974e 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 @@ -6682,6 +6682,80 @@ def test_create_combined_extractor_defaults_to_union(): assert list(selector.extractor.extract_records(response)) == [{"id": 1}, {"id": 2}] +def test_create_combined_extractor_defaults_to_not_skipping_empty_records(): + content = """ + selector: + type: RecordSelector + extractor: + type: CombinedExtractor + extractors: + - type: DpathExtractor + field_path: ["a"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + assert selector.extractor.skip_empty_records is False + + +def test_create_combined_extractor_with_skip_empty_records(): + """The source-monday shape: a page of nulls must fall through to the pagination path.""" + content = """ + selector: + type: RecordSelector + extractor: + type: CombinedExtractor + mode: first_match + skip_empty_records: true + extractors: + - type: DpathExtractor + field_path: ["data", "boards", "*", "items_page", "items", "*"] + - type: DpathExtractor + field_path: ["data", "next_items_page", "items", "*"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + assert isinstance(selector.extractor, CombinedExtractor) + assert selector.extractor.skip_empty_records is True + + response = requests.Response() + response._content = json.dumps( + { + "data": { + "boards": [{"items_page": {"items": [None, None]}}], + "next_items_page": {"items": [{"id": "item_1"}]}, + }, + "errors": [{"message": "Item not found"}], + } + ).encode("utf-8") + assert list(selector.extractor.extract_records(response)) == [{"id": "item_1"}] + + def _combined_extractor_selector_definition() -> Mapping[str, Any]: content = """ selector: From bbfd864cdb5080e36d0b061c20ad954992a52fb8 Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Wed, 16 Sep 2026 17:28:26 +0300 Subject: [PATCH 05/10] feat(CombinedExtractor): drop the `zip_merge` mode The architecture review of this PR found one adopter for `zip_merge` (source-google-analytics-data-api), and that connector keeps its `components.py` in every scenario because both of its inner extractors are a custom `KeyValueExtractor`. `union` and `first_match` each have two adopters, so they stay; `zip_merge` was fitted to a single connector that it does not free. Removes the mode from the enum, the schema, the generated models, the component and its tests, along with the shortest-wins truncation warning, the Mapping-only merge check and the `skip_empty_records` alignment caveat that only applied to it. The mode can come back when a second columnar API turns up, ideally alongside a key/value-list-to-object transformation that would free GA4's inner extractor too. Co-Authored-By: Claude Opus 5 (1M context) --- .../declarative_component_schema.yaml | 15 +- .../extractors/combined_extractor.py | 77 ++-------- .../models/declarative_component_schema.py | 5 +- .../parsers/model_to_component_factory.py | 10 +- .../extractors/test_combined_extractor.py | 140 +----------------- .../test_model_to_component_factory.py | 26 ---- .../test_concurrent_declarative_source.py | 2 +- 7 files changed, 23 insertions(+), 252 deletions(-) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 34b9d2eb19..e8f4a98ce5 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2175,16 +2175,14 @@ definitions: additionalProperties: true CombineMode: title: Combine Mode - description: 'How the records of the sub-extractors of a CombinedExtractor are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones.' + description: 'How the records of the sub-extractors of a CombinedExtractor are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record.' type: string enum: - union - first_match - - zip_merge examples: - union - first_match - - zip_merge CombinedExtractor: title: Combined Extractor description: Record extractor that combines the output of several sub-extractors. Use it when a single response holds records under more than one path, when the records can be found under one of several alternative paths, or when a record has to be assembled from several parallel lists. @@ -2208,12 +2206,12 @@ definitions: - "$ref": "#/definitions/CombinedExtractor" mode: title: Combine Mode - description: 'How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that a paginator which counts the records of a page counts the combined records: under "union" that is the sum over all sub-extractors, which overshoots the API page size, so "union" is rejected with an OffsetIncrement paginator because the offset would skip records. Under "first_match" the count is the count of the winning sub-extractor and under "zip_merge" the count of the shortest one, which are usually the number of records the API returned for the page.' + description: 'How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. Note that a paginator which counts the records of a page counts the combined records: under "union" that is the sum over all sub-extractors, which overshoots the API page size, so "union" is rejected with an OffsetIncrement paginator because the offset would skip records. Under "first_match" the count is the count of the winning sub-extractor, which is usually the number of records the API returned for the page.' "$ref": "#/definitions/CombineMode" default: union skip_empty_records: title: Skip Empty Records - description: 'Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried, and the record count a paginator obtains is the count after the empty records were dropped. Under "zip_merge" it shifts the alignment of the merged records, because dropping the i-th record of one sub-extractor pairs its next record with the i-th record of the others.' + description: 'Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried, and the record count a paginator obtains is the count after the empty records were dropped.' type: boolean default: false $parameters: @@ -2235,13 +2233,6 @@ definitions: field_path: ["data", "boards", "*", "items_page", "items"] - type: DpathExtractor field_path: ["data", "next_items_page", "items"] - - type: CombinedExtractor - mode: zip_merge - extractors: - - type: DpathExtractor - field_path: ["reports", "*", "dimensions"] - - type: DpathExtractor - field_path: ["reports", "*", "metrics"] DpathExtractor: title: Dpath Extractor description: Record extractor that searches a decoded response over a path defined as an array of fields. diff --git a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py index e981dec913..7f366456fc 100644 --- a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py +++ b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py @@ -6,7 +6,7 @@ from dataclasses import InitVar, dataclass from enum import Enum from itertools import chain -from typing import Any, Dict, Iterable, Iterator, List, Mapping, Sequence, Union +from typing import Any, Iterable, Iterator, List, Mapping, Union import requests @@ -22,7 +22,6 @@ class CombineMode(Enum): union = "union" first_match = "first_match" - zip_merge = "zip_merge" class _NoRecord: @@ -38,7 +37,7 @@ class CombinedExtractor(RecordExtractor): A single `DpathExtractor` can only describe one path into the response. Several connectors need more than that and reach for a custom `components.py` today, which keeps them out of the - Connector Builder. This component covers the three shapes those connectors implement: + Connector Builder. This component covers the two shapes those connectors implement: - `union` (default) — yield every record of every sub-extractor, in sub-extractor order. Use it when one response carries records under several paths, for example a GraphQL document that @@ -49,12 +48,6 @@ class CombinedExtractor(RecordExtractor): sub-extractor is only peeked at, never restarted: the peeked record is chained back in front of the remaining ones, so large responses are still streamed lazily and no sub-extractor is ever materialized into a list. - - `zip_merge` — merge the i-th record of every sub-extractor into a single dictionary, with - later sub-extractors overwriting keys set by earlier ones. Iteration stops at the SHORTEST - sub-extractor (plain `zip` semantics, not `zip_longest`): records past the end of the - shortest sequence are dropped rather than padded with empty values, and a warning is logged - when that happens. This matches the hand-rolled `CombinedExtractor` of - `source-google-analytics-data-api`, which this component is meant to replace. Examples of instantiating this component: ``` @@ -98,10 +91,9 @@ class CombinedExtractor(RecordExtractor): Because of that, `ModelToComponentFactory.create_combined_extractor` refuses to build this component over a streaming decoder and raises a configuration error, naming the decoder in the - internal message. The - Connector Builder forces those same decoders to `stream_response=False`, so the rejection is - applied there too rather than letting a manifest test-read correctly and lose records once - published. + internal message. The Connector Builder forces those same decoders to `stream_response=False`, + so the rejection is applied there too rather than letting a manifest test-read correctly and + lose records once published. ## Empty records @@ -118,16 +110,12 @@ class CombinedExtractor(RecordExtractor): the behaviour of `source-monday`'s `MondayIncrementalItemsExtractor`, where a page of nulls must fall through to the pagination path rather than lock in the primary path. - Under `zip_merge` the flag shifts alignment: dropping the i-th record of one sub-extractor - pairs its i+1-th record with the i-th record of the others. Only enable it there if the - sub-extractors drop records in lockstep. - ## Cost Each sub-extractor decodes the response independently — `CompositeRawDecoder` re-parses the - cached body on every `decode()` call — so a response is parsed once per sub-extractor. Under - `zip_merge` all sub-extractor outputs are also alive at the same time. Combining three - extractors over a large response costs roughly three times the decode time of a single one. + cached body on every `decode()` call — so a response is parsed once per sub-extractor. + Combining three extractors over a large response costs roughly three times the decode time of + a single one. Note that an `OffsetIncrement` or `PageIncrement` paginator builds its own copy of the extractor to count the records of a page, which doubles that cost. @@ -144,9 +132,9 @@ class CombinedExtractor(RecordExtractor): error. A `PageIncrement` only compares the count against `page_size` to decide whether to stop, so it loses nothing, but it issues one extra request whenever the summed count of the last page happens to equal `page_size`. - - `first_match` returns the count of the winning sub-extractor and `zip_merge` the count of the - shortest one. Neither inflates the count, so both work with either paginator, as long as the - path the count comes from is the one the API paginates over. + - `first_match` returns the count of the winning sub-extractor, which does not inflate the + count, so it works with either paginator as long as the path the count comes from is the one + the API paginates over. Attributes: extractors (List[RecordExtractor]): The sub-extractors to combine. At least one is required. @@ -171,8 +159,6 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: def extract_records(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: if self.mode == CombineMode.first_match: yield from self._extract_first_match(response) - elif self.mode == CombineMode.zip_merge: - yield from self._extract_zip_merge(response) else: yield from self._extract_union(response) @@ -191,33 +177,6 @@ def _extract_first_match(self, response: requests.Response) -> Iterable[Mapping[ yield from chain([first_record], records) return - def _extract_zip_merge(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: - # Iteration stops at the shortest sub-extractor, matching the behavior this component - # replaces. Unlike plain `zip`, the round in which a sub-extractor runs out is completed so - # that an uneven number of records can be reported instead of silently truncated. - iterators = [ - iter(self._records_of(index, response)) for index in range(len(self.extractors)) - ] - emitted = 0 - while True: - round_records = [next(iterator, _NO_RECORD) for iterator in iterators] - exhausted = [ - index for index, record in enumerate(round_records) if isinstance(record, _NoRecord) - ] - if exhausted: - if len(exhausted) != len(iterators): - logger.warning( - "CombinedExtractor in `zip_merge` mode discarded the records of the " - "sub-extractors that outlasted the shortest one: sub-extractors " - "%s ran out after %s record(s) while the others still had records. " - "`zip_merge` stops at the shortest sub-extractor.", - exhausted, - emitted, - ) - return - yield self._merge(round_records) - emitted += 1 - def _records_of( self, extractor_index: int, response: requests.Response ) -> Iterable[Mapping[str, Any]]: @@ -245,17 +204,3 @@ def _without_empty_records( extractor_index, type(self.extractors[extractor_index]).__name__, ) - - def _merge(self, records: Sequence[Any]) -> Mapping[str, Any]: - merged: Dict[str, Any] = {} - for index, record in enumerate(records): - if not isinstance(record, Mapping): - raise ValueError( - f"CombinedExtractor in `zip_merge` mode can only merge records that are " - f"objects, but sub-extractor {index} " - f"({type(self.extractors[index]).__name__}) yielded a " - f"{type(record).__name__}. Point that sub-extractor at a path that holds " - f"objects, or use the `union` mode." - ) - merged.update(record) - return merged diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index bf57750946..508114dbd3 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -510,7 +510,6 @@ class Config: class CombineMode(Enum): union = "union" first_match = "first_match" - zip_merge = "zip_merge" class ResponseToFileExtractor(BaseModel): @@ -2532,12 +2531,12 @@ class CombinedExtractor(BaseModel): ) mode: Optional[CombineMode] = Field( CombineMode.union, - description='How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. "zip_merge" merges the i-th record of every sub-extractor into a single record, with later sub-extractors overwriting the fields set by earlier ones, and stops at the shortest sub-extractor, discarding the trailing records of the longer ones. Note that a paginator which counts the records of a page counts the combined records: under "union" that is the sum over all sub-extractors, which overshoots the API page size, so "union" is rejected with an OffsetIncrement paginator because the offset would skip records. Under "first_match" the count is the count of the winning sub-extractor and under "zip_merge" the count of the shortest one, which are usually the number of records the API returned for the page.', + description='How the records of the sub-extractors are combined. "union" (default) yields every record of every sub-extractor, in the order the extractors are declared. "first_match" yields the records of the first sub-extractor that produces at least one record and skips the remaining ones; nothing is yielded if none of them produces a record. Note that a paginator which counts the records of a page counts the combined records: under "union" that is the sum over all sub-extractors, which overshoots the API page size, so "union" is rejected with an OffsetIncrement paginator because the offset would skip records. Under "first_match" the count is the count of the winning sub-extractor, which is usually the number of records the API returned for the page.', title="Combine Mode", ) skip_empty_records: Optional[bool] = Field( False, - description='Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried, and the record count a paginator obtains is the count after the empty records were dropped. Under "zip_merge" it shifts the alignment of the merged records, because dropping the i-th record of one sub-extractor pairs its next record with the i-th record of the others.', + description='Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried, and the record count a paginator obtains is the count after the empty records were dropped.', title="Skip Empty Records", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") 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 2afd392fbf..fadfefbb3f 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -3164,9 +3164,9 @@ def _reject_union_combined_extractor_for_offset_increment( each and `page_size: 2`, the requested offsets are 0, 4, 8 instead of 0, 2, 4 and two thirds of the records are silently dropped. - `first_match` (the winning sub-extractor's count) and `zip_merge` (the shortest - sub-extractor's count) do not inflate the count and are left alone. A `union` nested - anywhere in the tree inflates the count of the node above it, so the whole tree is walked. + `first_match` returns the winning sub-extractor's count, which does not inflate the count, + and is left alone. A `union` nested anywhere in the tree inflates the count of the node + above it, so the whole tree is walked. """ if not isinstance(extractor_model, CombinedExtractorModel): return @@ -3180,8 +3180,8 @@ def _reject_union_combined_extractor_for_offset_increment( "OffsetIncrement counts the records of its extractor to advance the offset. A " "`union` CombinedExtractor returns the sum of its sub-extractors' records, which " "overshoots the page the API returned, so records would be skipped. Use the " - "`first_match` or `zip_merge` mode, a CursorPagination or PageIncrement " - "paginator, or a single extractor." + "`first_match` mode, a CursorPagination or PageIncrement paginator, or a single " + "extractor." ), failure_type=FailureType.config_error, ) diff --git a/unit_tests/sources/declarative/extractors/test_combined_extractor.py b/unit_tests/sources/declarative/extractors/test_combined_extractor.py index 9bab3eb7ef..8eb0225e0a 100644 --- a/unit_tests/sources/declarative/extractors/test_combined_extractor.py +++ b/unit_tests/sources/declarative/extractors/test_combined_extractor.py @@ -60,15 +60,6 @@ "errors": [{"message": "Item not found"}], } -# The GA4 report shape: dimensions and metrics come back as two parallel lists that have to be -# merged position by position. -REPORT_BODY = { - "rows": [ - {"dimensions": {"date": "20240101"}, "metrics": {"activeUsers": 1}}, - {"dimensions": {"date": "20240102"}, "metrics": {"activeUsers": 2}}, - ] -} - def create_response(body: Union[Dict, List, bytes]) -> requests.Response: """A response whose `raw` is a real `urllib3.HTTPResponse`, which is what `requests` builds. @@ -215,64 +206,9 @@ def test_first_match_yields_nothing_when_every_extractor_is_empty(): assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [] -def test_zip_merge_merges_the_ith_record_of_every_extractor(): - extractor = CombinedExtractor( - extractors=[dpath("rows", "*", "dimensions"), dpath("rows", "*", "metrics")], - mode=CombineMode.zip_merge, - parameters=parameters, - ) - - assert list(extractor.extract_records(create_response(REPORT_BODY))) == [ - {"date": "20240101", "activeUsers": 1}, - {"date": "20240102", "activeUsers": 2}, - ] - - -def test_zip_merge_later_extractors_win_on_key_collision(): - extractor = CombinedExtractor( - extractors=[ - _CountingExtractor(records=[{"id": 1, "source": "first", "only_first": True}]), - _CountingExtractor(records=[{"id": 1, "source": "second"}]), - ], - mode=CombineMode.zip_merge, - parameters=parameters, - ) - - assert list(extractor.extract_records(create_response(REPORT_BODY))) == [ - {"id": 1, "source": "second", "only_first": True} - ] - - -def test_zip_merge_stops_at_the_shortest_sub_extractor(): - """`zip`, not `zip_longest`: the trailing records of the longer extractor are dropped, not padded.""" - extractor = CombinedExtractor( - extractors=[ - _CountingExtractor(records=[{"a": 1}, {"a": 2}, {"a": 3}]), - _CountingExtractor(records=[{"b": 1}, {"b": 2}]), - ], - mode=CombineMode.zip_merge, - parameters=parameters, - ) - - assert list(extractor.extract_records(create_response(REPORT_BODY))) == [ - {"a": 1, "b": 1}, - {"a": 2, "b": 2}, - ] - - -def test_zip_merge_yields_nothing_when_one_extractor_is_empty(): - extractor = CombinedExtractor( - extractors=[dpath("rows", "*", "dimensions"), dpath("rows", "*", "does_not_exist")], - mode=CombineMode.zip_merge, - parameters=parameters, - ) - - assert list(extractor.extract_records(create_response(REPORT_BODY))) == [] - - @pytest.mark.parametrize( "mode", - [CombineMode.union, CombineMode.first_match, CombineMode.zip_merge], + [CombineMode.union, CombineMode.first_match], ) def test_a_single_sub_extractor_behaves_like_that_extractor_alone(mode: CombineMode): field_path = ["data", "repository", "pullRequests", "nodes"] @@ -391,61 +327,6 @@ def test_first_match_over_a_streaming_decoder_silently_returns_the_wrong_answer( assert list(buffered.extract_records(create_response(body))) == [{"id": 1}, {"id": 2}] -def test_zip_merge_warns_when_a_sub_extractor_runs_out_early(caplog): - extractor = CombinedExtractor( - extractors=[ - _CountingExtractor(records=[{"a": 1}, {"a": 2}, {"a": 3}]), - _CountingExtractor(records=[{"b": 1}, {"b": 2}]), - ], - mode=CombineMode.zip_merge, - parameters=parameters, - ) - - with caplog.at_level(logging.WARNING, logger="airbyte"): - assert list(extractor.extract_records(create_response(REPORT_BODY))) == [ - {"a": 1, "b": 1}, - {"a": 2, "b": 2}, - ] - - assert any( - "zip_merge" in record.message and record.levelno == logging.WARNING - for record in caplog.records - ) - - -def test_zip_merge_does_not_warn_when_every_sub_extractor_has_the_same_length(caplog): - extractor = CombinedExtractor( - extractors=[dpath("rows", "*", "dimensions"), dpath("rows", "*", "metrics")], - mode=CombineMode.zip_merge, - parameters=parameters, - ) - - with caplog.at_level(logging.WARNING, logger="airbyte"): - assert len(list(extractor.extract_records(create_response(REPORT_BODY)))) == 2 - - assert [record.message for record in caplog.records] == [] - - -def test_zip_merge_names_the_sub_extractor_that_yielded_a_non_object(): - extractor = CombinedExtractor( - extractors=[ - _CountingExtractor(records=[{"a": 1}]), - _CountingExtractor(records=["not-an-object"]), - ], - mode=CombineMode.zip_merge, - parameters=parameters, - ) - - with pytest.raises(ValueError) as exc_info: - list(extractor.extract_records(create_response(REPORT_BODY))) - - message = str(exc_info.value) - assert "CombinedExtractor" in message - assert "zip_merge" in message - assert "sub-extractor 1" in message - assert "_CountingExtractor" in message - - def _monday_extractor(skip_empty_records: bool) -> CombinedExtractor: return CombinedExtractor( extractors=[ @@ -561,22 +442,3 @@ def test_skip_empty_records_does_not_warn_when_nothing_is_dropped(caplog): ] assert [record.message for record in caplog.records if record.levelno == logging.WARNING] == [] - - -def test_skip_empty_records_shifts_zip_merge_alignment(): - """Documents the caveat the schema description warns about. - - Dropping the first sub-extractor's null pairs its next record with the other sub-extractor's - first record, so the merge is off by one rather than skipping a row. - """ - extractor = CombinedExtractor( - extractors=[ - _CountingExtractor(records=[None, {"a": 2}]), - _CountingExtractor(records=[{"b": 1}, {"b": 2}]), - ], - mode=CombineMode.zip_merge, - skip_empty_records=True, - parameters=parameters, - ) - - assert list(extractor.extract_records(create_response(REPORT_BODY))) == [{"a": 2, "b": 1}] 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 23c6a5974e..32ad5bc560 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 @@ -7098,32 +7098,6 @@ def test_first_match_combined_extractor_counts_only_the_winning_sub_extractor_fo ) -def test_zip_merge_combined_extractor_counts_the_shortest_sub_extractor_for_the_offset(): - extractor_definition = """ type: CombinedExtractor - mode: zip_merge - extractors: - - type: DpathExtractor - field_path: ["a"] - - type: DpathExtractor - field_path: ["b"]""" - - retriever = _build_retriever( - _retriever_manifest_with_paginator(extractor_definition, _OFFSET_INCREMENT) - ) - - pagination_strategy = retriever.paginator.pagination_strategy - assert isinstance(pagination_strategy, OffsetIncrement) - - # Two merged records, not the four records the two sub-extractors yield together. - response = _json_response({"a": [{"x": 1}, {"x": 2}], "b": [{"y": 1}, {"y": 2}]}) - assert ( - pagination_strategy.next_page_token( - response=response, last_page_size=2, last_record=None, last_page_token_value=0 - ) - == 2 - ) - - def test_union_combined_extractor_is_accepted_by_a_page_increment_paginator(): """`PageIncrement` only compares the count against `page_size`, so an inflated count is not lossy. diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index 267836021d..1f46fd8da1 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -6160,7 +6160,7 @@ def test_combined_extractor_manifest_passes_schema_validation(): manifest = _combined_extractor_manifest( { "type": "CombinedExtractor", - "mode": "zip_merge", + "mode": "first_match", "extractors": [ {"type": "DpathExtractor", "field_path": ["rows", "*", "dimensions"]}, { From d72f9a92059013732152a4e11bda061362a682d8 Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Wed, 16 Sep 2026 17:34:00 +0300 Subject: [PATCH 06/10] feat(RecordExpander): add `parent_fields` to copy named parent values onto expanded items `RecordExpander` can already carry parent context onto each expanded item, but only by way of `remain_original_record: true`, which deep-copies the entire parent record once per item. When the parent is large and the nested list is long - a GitHub pull request node with ten reviews, say - that copies the whole pull request ten times to carry one URL. `parent_fields` names the values to copy instead: ```yaml record_expander: type: RecordExpander expand_records_from_field: ["reviews", "nodes"] parent_fields: - type: ParentFieldPath parent_path: ["url"] record_path: ["pull_request_url"] ``` Both paths are lists, so a value nested in the parent can be written to a nested position on the item; intermediate objects are created as needed. An existing value at `record_path` is overwritten, and a `parent_path` the parent does not have copies `None` - which is what `parent.get(field)` does in the connector classes this replaces. Glob metacharacters are rejected in both paths, matching the existing rule on `truncation_indicator_path`. The copy happens in `_apply_parent_context`, the single point both the embedded-items path and the `truncated_list_retriever` path go through, so fetched items are stamped the same way as embedded ones. `parent_fields` and `remain_original_record` are independent and may both be set. Making one suppress the other would be surprising; the reason to prefer `parent_fields` is that it does not need the deep copy, which the docstring and the schema description both say. Scalar list items previously fell through unwrapped unless `remain_original_record` was set, in which case they were wrapped as `{"value": ..., "original_record": ...}`. They are now wrapped as `{"value": ...}` plus the copied fields whenever either option asks for parent context, and still yielded bare when neither does. The generated model was hand-applied and verified against the pinned generator (`datamodel_code_generator==0.26.3` with the flag set from `bin/generate_component_manifest_files.py`): the `ParentFieldPath` class and the `parent_fields` field are byte-identical to its output after `ruff format`. Co-Authored-By: Claude Opus 5 (1M context) --- .../declarative_component_schema.yaml | 45 ++++ .../sources/declarative/expanders/__init__.py | 8 +- .../declarative/expanders/record_expander.py | 95 +++++++- .../models/declarative_component_schema.py | 22 ++ .../parsers/model_to_component_factory.py | 24 ++ .../expanders/test_record_expander.py | 220 +++++++++++++++++- .../test_model_to_component_factory.py | 94 ++++++++ 7 files changed, 498 insertions(+), 10 deletions(-) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 071f5eeaf4..62e8209660 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2250,6 +2250,12 @@ definitions: description: If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false. type: boolean default: false + parent_fields: + title: Parent Fields + description: 'Named values to copy from the record being expanded onto each expanded item. Use this instead of `remain_original_record` when only a few parent fields are needed: it copies the named values rather than deep-copying the whole parent once per item, which matters when the parent record is large and the nested list is long. An existing value at `record_path` is overwritten, and a `parent_path` the parent does not have copies null. Independent of `remain_original_record`; both may be set. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.' + type: array + items: + "$ref": "#/definitions/ParentFieldPath" on_no_records: title: On No Records description: Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged. @@ -2277,6 +2283,45 @@ definitions: $parameters: type: object additionalProperties: true + ParentFieldPath: + title: Parent Field Path + description: One value to copy from the record being expanded onto each expanded item. Both paths are lists, so a value nested inside the parent can be copied into a nested position on the item; intermediate objects on `record_path` are created as needed. Glob characters (`*`, `?`, `[`) are not supported in either path, since each must identify a single field; this is enforced on the interpolated values. + type: object + required: + - type + - parent_path + - record_path + properties: + type: + type: string + enum: [ParentFieldPath] + parent_path: + title: Parent Path + description: Path to the value on the record being expanded. + type: array + items: + type: string + interpolation_context: + - config + examples: + - ["url"] + - ["id"] + - ["repository", "name"] + record_path: + title: Record Path + description: Path on the expanded item to write the value to. An existing value is overwritten. + type: array + items: + type: string + interpolation_context: + - config + examples: + - ["pull_request_url"] + - ["comment_id"] + - ["parent", "id"] + $parameters: + type: object + additionalProperties: true ExponentialBackoffStrategy: title: Exponential Backoff description: Backoff strategy with an exponential backoff interval. The interval is defined as factor * 2^attempt_count. diff --git a/airbyte_cdk/sources/declarative/expanders/__init__.py b/airbyte_cdk/sources/declarative/expanders/__init__.py index 1b10ee82d7..23f59bf1bf 100644 --- a/airbyte_cdk/sources/declarative/expanders/__init__.py +++ b/airbyte_cdk/sources/declarative/expanders/__init__.py @@ -2,6 +2,10 @@ # Copyright (c) 2023 Airbyte, Inc., all rights reserved. # -from airbyte_cdk.sources.declarative.expanders.record_expander import OnNoRecords, RecordExpander +from airbyte_cdk.sources.declarative.expanders.record_expander import ( + OnNoRecords, + ParentFieldPath, + RecordExpander, +) -__all__ = ["OnNoRecords", "RecordExpander"] +__all__ = ["OnNoRecords", "ParentFieldPath", "RecordExpander"] diff --git a/airbyte_cdk/sources/declarative/expanders/record_expander.py b/airbyte_cdk/sources/declarative/expanders/record_expander.py index 68c71af693..99125d5ae9 100644 --- a/airbyte_cdk/sources/declarative/expanders/record_expander.py +++ b/airbyte_cdk/sources/declarative/expanders/record_expander.py @@ -50,6 +50,57 @@ class OnNoRecords(Enum): emit_parent = "emit_parent" +@dataclass +class ParentFieldPath: + """One field to copy from the record being expanded onto each expanded item. + + `parent_path` locates the value on the parent, `record_path` says where to put it on the + child. Both are lists, so a value nested inside the parent can be copied into a nested + position on the child. Glob metacharacters are rejected: each path must identify a single + field. + """ + + parent_path: Sequence[str] + record_path: Sequence[str] + config: Config + parameters: InitVar[Mapping[str, Any]] + + def __post_init__(self, parameters: Mapping[str, Any]) -> None: + if not self.parent_path: + raise ValueError("`parent_path` cannot be empty.") + if not self.record_path: + raise ValueError("`record_path` cannot be empty.") + self._parent_path: list[InterpolatedString] = [ + InterpolatedString.create(path, parameters=parameters) for path in self.parent_path + ] + self._record_path: list[InterpolatedString] = [ + InterpolatedString.create(path, parameters=parameters) for path in self.record_path + ] + RecordExpander._reject_globs(self.evaluated_parent_path(), "parent_path") + RecordExpander._reject_globs(self.evaluated_record_path(), "record_path") + + def evaluated_parent_path(self) -> list[Any]: + return [segment.eval(self.config) for segment in self._parent_path] + + def evaluated_record_path(self) -> list[Any]: + return [segment.eval(self.config) for segment in self._record_path] + + def copy_onto( + self, parent_record: Mapping[str, Any], child_record: MutableMapping[str, Any] + ) -> None: + """Copy the parent value onto the child, overwriting whatever was there. + + A `parent_path` that the parent does not have copies `None`, which is what the custom + connector classes this replaces do (`parent.get(field)`). Distinguishing "absent" from + "present and null" would need a third option and no connector needs one. + """ + try: + value = dpath.get(dict(parent_record), self.evaluated_parent_path()) + except (KeyError, ValueError): + value = None + dpath.new(child_record, self.evaluated_record_path(), value) + + @dataclass class RecordExpander: """Expands records by extracting items from a nested array field. @@ -59,6 +110,11 @@ class RecordExpander: to embed the full parent record under `original_record` in each expanded item when you need downstream transformations to access parent context. + When only a few parent fields are needed, prefer `parent_fields` over + `remain_original_record`: it copies the named values onto each expanded item instead of + deep-copying the whole parent once per item, which matters when the parent is large and the + nested list is long. + The expand_records_from_field path supports wildcards (*) for matching multiple arrays. When wildcards are used, items from all matched arrays are extracted and emitted. @@ -72,6 +128,17 @@ class RecordExpander: remain_original_record: true ``` + ``` + record_expander: + type: RecordExpander + expand_records_from_field: + - "reviews" + - "nodes" + parent_fields: + - parent_path: ["url"] + record_path: ["pull_request_url"] + ``` + ``` record_expander: type: RecordExpander @@ -88,6 +155,13 @@ class RecordExpander: Supports wildcards (*). remain_original_record: If True, each expanded record will include the original parent record in an "original_record" field. Defaults to False. + parent_fields: Named values to copy from the record being expanded onto each expanded + item. Each entry has a `parent_path` and a `record_path`; an existing value at + `record_path` is overwritten, and a `parent_path` the parent does not have copies + `None`. Independent of `remain_original_record` - both may be set, though copying + named fields is the cheaper way to carry parent context. Applies to items fetched + through `truncated_list_retriever` as well as to embedded ones. Glob metacharacters + are rejected in both paths. on_no_records: Behavior when expansion produces no records. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged. truncation_indicator_path: Path within each record to a field indicating that the @@ -125,6 +199,7 @@ class RecordExpander: config: Config parameters: InitVar[Mapping[str, Any]] remain_original_record: bool = False + parent_fields: Optional[Sequence[ParentFieldPath]] = None on_no_records: OnNoRecords = OnNoRecords.skip truncation_indicator_path: Optional[Sequence[str]] = None truncated_list_retriever: Optional["Retriever"] = None @@ -216,11 +291,10 @@ def expand_record(self, record: MutableMapping[Any, Any]) -> Iterable[MutableMap expanded_record = dict(item) self._apply_parent_context(parent_record, expanded_record) yield expanded_record - elif self.remain_original_record: - yield { - "value": item, - "original_record": copy.deepcopy(parent_record), - } + elif self._carries_parent_context(): + wrapped: MutableMapping[str, Any] = {"value": item} + self._apply_parent_context(parent_record, wrapped) + yield wrapped else: yield item @@ -323,8 +397,10 @@ def _fetch_complete_list(self, parent_record: Mapping[str, Any]) -> Iterable[Any expanded_record = dict(data) self._apply_parent_context(parent_record, expanded_record) yield expanded_record - elif self.remain_original_record: - yield {"value": data, "original_record": copy.deepcopy(parent_record)} + elif self._carries_parent_context(): + wrapped: MutableMapping[str, Any] = {"value": data} + self._apply_parent_context(parent_record, wrapped) + yield wrapped else: yield data @@ -334,3 +410,8 @@ def _apply_parent_context( """Apply parent context to a child record.""" if self.remain_original_record: child_record["original_record"] = copy.deepcopy(parent_record) + for parent_field in self.parent_fields or []: + parent_field.copy_onto(parent_record, child_record) + + def _carries_parent_context(self) -> bool: + return bool(self.remain_original_record or self.parent_fields) diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 41265c9db9..4345e40df2 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -522,6 +522,23 @@ class OnNoRecords(Enum): emit_parent = "emit_parent" +class ParentFieldPath(BaseModel): + type: Literal["ParentFieldPath"] + parent_path: List[str] = Field( + ..., + description="Path to the value on the record being expanded.", + examples=[["url"], ["id"], ["repository", "name"]], + title="Parent Path", + ) + record_path: List[str] = Field( + ..., + description="Path on the expanded item to write the value to. An existing value is overwritten.", + examples=[["pull_request_url"], ["comment_id"], ["parent", "id"]], + title="Record Path", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class ExponentialBackoffStrategy(BaseModel): type: Literal["ExponentialBackoffStrategy"] factor: Optional[Union[float, str]] = Field( @@ -2300,6 +2317,11 @@ class RecordExpander(BaseModel): description='If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false.', title="Remain Original Record", ) + parent_fields: Optional[List[ParentFieldPath]] = Field( + None, + description="Named values to copy from the record being expanded onto each expanded item. Use this instead of `remain_original_record` when only a few parent fields are needed: it copies the named values rather than deep-copying the whole parent once per item, which matters when the parent record is large and the nested list is long. An existing value at `record_path` is overwritten, and a `parent_path` the parent does not have copies null. Independent of `remain_original_record`; both may be set. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.", + title="Parent Fields", + ) on_no_records: Optional[OnNoRecords] = Field( OnNoRecords.skip, description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', 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 bec394f85e..6dab0a84e9 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -109,6 +109,7 @@ ) from airbyte_cdk.sources.declarative.expanders.record_expander import ( OnNoRecords, + ParentFieldPath, RecordExpander, ) from airbyte_cdk.sources.declarative.extractors import ( @@ -390,6 +391,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ParametrizedComponentsResolver as ParametrizedComponentsResolverModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + ParentFieldPath as ParentFieldPathModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ParentStreamConfig as ParentStreamConfigModel, ) @@ -824,6 +828,7 @@ def _init_mappings(self) -> None: PropertiesFromEndpointModel: self.create_properties_from_endpoint, PropertyChunkingModel: self.create_property_chunking, QueryPropertiesModel: self.create_query_properties, + ParentFieldPathModel: self.create_parent_field_path, RecordExpanderModel: self.create_record_expander, RecordFilterModel: self.create_record_filter, RecordSelectorModel: self.create_record_selector, @@ -2545,6 +2550,12 @@ def create_record_expander( config=config, parameters=model.parameters or {}, remain_original_record=model.remain_original_record or False, + parent_fields=[ + self._create_component_from_model(model=parent_field, config=config) + for parent_field in model.parent_fields + ] + if model.parent_fields + else None, on_no_records=OnNoRecords(model.on_no_records.value) if model.on_no_records else OnNoRecords.skip, @@ -2554,6 +2565,19 @@ def create_record_expander( suppress_incomplete_fetch_warning=suppress_incomplete_fetch_warning, ) + @staticmethod + def create_parent_field_path( + model: ParentFieldPathModel, + config: Config, + **kwargs: Any, + ) -> ParentFieldPath: + return ParentFieldPath( + parent_path=model.parent_path, + record_path=model.record_path, + config=config, + parameters=model.parameters or {}, + ) + @staticmethod def create_response_to_file_extractor( model: ResponseToFileExtractorModel, diff --git a/unit_tests/sources/declarative/expanders/test_record_expander.py b/unit_tests/sources/declarative/expanders/test_record_expander.py index 23731dda9e..42fc37109a 100644 --- a/unit_tests/sources/declarative/expanders/test_record_expander.py +++ b/unit_tests/sources/declarative/expanders/test_record_expander.py @@ -8,7 +8,10 @@ import pytest from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteRecordMessage, Level, Type -from airbyte_cdk.sources.declarative.expanders.record_expander import RecordExpander +from airbyte_cdk.sources.declarative.expanders.record_expander import ( + ParentFieldPath, + RecordExpander, +) from airbyte_cdk.sources.message import InMemoryMessageRepository from airbyte_cdk.sources.types import Record, StreamSlice @@ -572,3 +575,218 @@ def test_no_warning_when_retriever_configured(caplog): assert len(records) == 15 assert not [r for r in caplog.records if r.levelname == "WARNING"] + + +def _parent_field(parent_path, record_path): + return ParentFieldPath( + parent_path=parent_path, + record_path=record_path, + config=config, + parameters=parameters, + ) + + +def _reviews_parent(): + return { + "number": 7, + "url": "https://github.com/airbytehq/airbyte/pull/7", + "reviews": {"nodes": [{"id": "PRR_1"}, {"id": "PRR_2"}]}, + } + + +def test_parent_fields_copies_the_named_value_onto_every_item(): + expander = RecordExpander( + expand_records_from_field=["reviews", "nodes"], + parent_fields=[_parent_field(["url"], ["pull_request_url"])], + config=config, + parameters=parameters, + ) + + assert list(expander.expand_record(_reviews_parent())) == [ + {"id": "PRR_1", "pull_request_url": "https://github.com/airbytehq/airbyte/pull/7"}, + {"id": "PRR_2", "pull_request_url": "https://github.com/airbytehq/airbyte/pull/7"}, + ] + + +def test_parent_fields_does_not_embed_the_parent(): + """The reason to use `parent_fields`: no `original_record`, so no deep copy per item.""" + expander = RecordExpander( + expand_records_from_field=["reviews", "nodes"], + parent_fields=[_parent_field(["url"], ["pull_request_url"])], + config=config, + parameters=parameters, + ) + + for record in expander.expand_record(_reviews_parent()): + assert "original_record" not in record + + +def test_parent_fields_and_remain_original_record_are_independent(): + expander = RecordExpander( + expand_records_from_field=["reviews", "nodes"], + remain_original_record=True, + parent_fields=[_parent_field(["number"], ["pull_request_number"])], + config=config, + parameters=parameters, + ) + + records = list(expander.expand_record(_reviews_parent())) + + assert [record["pull_request_number"] for record in records] == [7, 7] + assert all(record["original_record"]["number"] == 7 for record in records) + + +def test_parent_fields_reads_and_writes_nested_paths(): + expander = RecordExpander( + expand_records_from_field=["items"], + parent_fields=[_parent_field(["repository", "name"], ["parent", "repository_name"])], + config=config, + parameters=parameters, + ) + + parent = {"repository": {"name": "airbyte"}, "items": [{"id": 1}]} + + assert list(expander.expand_record(parent)) == [ + {"id": 1, "parent": {"repository_name": "airbyte"}} + ] + + +def test_parent_fields_overwrites_an_existing_value_on_the_item(): + expander = RecordExpander( + expand_records_from_field=["items"], + parent_fields=[_parent_field(["id"], ["id"])], + config=config, + parameters=parameters, + ) + + assert list(expander.expand_record({"id": "parent", "items": [{"id": "child"}]})) == [ + {"id": "parent"} + ] + + +def test_parent_fields_copies_none_when_the_parent_lacks_the_field(): + """Matches `parent.get(field)` in the connector classes this replaces.""" + expander = RecordExpander( + expand_records_from_field=["items"], + parent_fields=[_parent_field(["missing"], ["copied"])], + config=config, + parameters=parameters, + ) + + assert list(expander.expand_record({"items": [{"id": 1}]})) == [{"id": 1, "copied": None}] + + +def test_parent_fields_copies_a_present_null_the_same_way(): + expander = RecordExpander( + expand_records_from_field=["items"], + parent_fields=[_parent_field(["maybe"], ["copied"])], + config=config, + parameters=parameters, + ) + + assert list(expander.expand_record({"maybe": None, "items": [{"id": 1}]})) == [ + {"id": 1, "copied": None} + ] + + +def test_several_parent_fields_are_applied_in_order(): + expander = RecordExpander( + expand_records_from_field=["items"], + parent_fields=[ + _parent_field(["a"], ["copied"]), + _parent_field(["b"], ["copied"]), + ], + config=config, + parameters=parameters, + ) + + assert list(expander.expand_record({"a": 1, "b": 2, "items": [{}]})) == [{"copied": 2}] + + +def test_parent_fields_does_not_mutate_the_parent_record(): + expander = RecordExpander( + expand_records_from_field=["reviews", "nodes"], + parent_fields=[_parent_field(["url"], ["pull_request_url"])], + config=config, + parameters=parameters, + ) + parent = _reviews_parent() + + list(expander.expand_record(parent)) + + assert parent == _reviews_parent() + + +def test_parent_fields_wraps_scalar_items(): + expander = RecordExpander( + expand_records_from_field=["items"], + parent_fields=[_parent_field(["id"], ["parent_id"])], + config=config, + parameters=parameters, + ) + + assert list(expander.expand_record({"id": 9, "items": ["a", "b"]})) == [ + {"value": "a", "parent_id": 9}, + {"value": "b", "parent_id": 9}, + ] + + +def test_scalar_items_stay_bare_without_parent_context(): + expander = RecordExpander( + expand_records_from_field=["items"], + config=config, + parameters=parameters, + ) + + assert list(expander.expand_record({"items": ["a", "b"]})) == ["a", "b"] + + +def test_parent_fields_apply_to_items_fetched_by_the_truncated_list_retriever(): + expander = RecordExpander( + expand_records_from_field=["reviews", "nodes"], + truncation_indicator_path=["reviews", "has_more"], + truncated_list_retriever=_make_retriever([{"id": "PRR_9"}]), + parent_fields=[_parent_field(["url"], ["pull_request_url"])], + config=config, + parameters=parameters, + ) + parent = { + "url": "https://github.com/airbytehq/airbyte/pull/7", + "reviews": {"nodes": [{"id": "PRR_1"}], "has_more": True}, + } + + assert list(expander.expand_record(parent)) == [ + {"id": "PRR_9", "pull_request_url": "https://github.com/airbytehq/airbyte/pull/7"} + ] + + +@pytest.mark.parametrize("path", [["*"], ["a", "?"], ["a[0]"]]) +def test_globs_are_rejected_in_parent_path(path): + with pytest.raises(ValueError, match="Glob characters"): + _parent_field(path, ["copied"]) + + +@pytest.mark.parametrize("path", [["*"], ["a", "?"], ["a[0]"]]) +def test_globs_are_rejected_in_record_path(path): + with pytest.raises(ValueError, match="Glob characters"): + _parent_field(["id"], path) + + +@pytest.mark.parametrize("parent_path,record_path", [([], ["a"]), (["a"], [])]) +def test_empty_paths_are_rejected(parent_path, record_path): + with pytest.raises(ValueError, match="cannot be empty"): + _parent_field(parent_path, record_path) + + +def test_parent_field_paths_interpolate_config(): + parent_field = ParentFieldPath( + parent_path=["{{ config['from'] }}"], + record_path=["{{ config['to'] }}"], + config={"from": "url", "to": "pull_request_url"}, + parameters=parameters, + ) + child = {} + + parent_field.copy_onto({"url": "https://example.com/7"}, child) + + assert child == {"pull_request_url": "https://example.com/7"} 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 39adbde700..101f3105e0 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 @@ -1990,6 +1990,100 @@ def test_client_side_incremental_with_partition_router(): ) +def test_create_record_expander_with_parent_fields(): + """The source-github `reviews` shape: children stamped with a field from the parent node.""" + content = """ + selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data", "repository", "pullRequests", "nodes"] + record_expander: + type: RecordExpander + expand_records_from_field: ["reviews", "nodes"] + parent_fields: + - type: ParentFieldPath + parent_path: ["url"] + record_path: ["pull_request_url"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + expander = selector.extractor.record_expander + assert expander is not None + assert expander.remain_original_record is False + assert len(expander.parent_fields) == 1 + + response = requests.Response() + response._content = json.dumps( + { + "data": { + "repository": { + "pullRequests": { + "nodes": [ + { + "url": "https://github.com/airbytehq/airbyte/pull/7", + "reviews": {"nodes": [{"id": "PRR_1"}, {"id": "PRR_2"}]}, + }, + { + "url": "https://github.com/airbytehq/airbyte/pull/8", + "reviews": {"nodes": [{"id": "PRR_3"}]}, + }, + ] + } + } + } + } + ).encode("utf-8") + + assert list(selector.extractor.extract_records(response)) == [ + {"id": "PRR_1", "pull_request_url": "https://github.com/airbytehq/airbyte/pull/7"}, + {"id": "PRR_2", "pull_request_url": "https://github.com/airbytehq/airbyte/pull/7"}, + {"id": "PRR_3", "pull_request_url": "https://github.com/airbytehq/airbyte/pull/8"}, + ] + + +def test_create_record_expander_without_parent_fields_leaves_them_unset(): + content = """ + selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data"] + record_expander: + type: RecordExpander + expand_records_from_field: ["items"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + assert selector.extractor.record_expander.parent_fields is None + + def test_create_record_expander_with_truncated_list_retriever(): content = """ selector: From 120c1efd2652d46e50429a768edd9ecbbf501401 Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Wed, 16 Sep 2026 18:22:27 +0300 Subject: [PATCH 07/10] feat(RecordExpander): add `merge_parent` to shallow-merge the parent onto each expanded item With `merge_parent: true` each expanded item becomes the parent record shallow-merged with the item, the item's own keys winning on collision, and the value at `expand_records_from_field` removed from the parent's copy. Only that value is removed: on a multi-segment path the top-level key stays with its other fields, and the parent record itself is never mutated. This is the `{**record, **item}` flatten that source-mailchimp (`components.py:16`) and source-orb (`components.py:16-45`) implement in custom Python today. The order inside `_apply_parent_context` is fixed: the merge first, then `parent_fields` copies, then `original_record` when `remain_original_record` is set, so a named copy can overwrite a merged value. The three options are independent and may be combined. Scalar list items are wrapped as `{"value": item}` before the merge, and items fetched through `truncated_list_retriever` are merged the same way as embedded ones. The stripped parent is built once per parent record rather than once per item. The generated model change was hand-applied in the same style as `parent_fields`. Co-Authored-By: Claude Fable 5.1 --- .../declarative_component_schema.yaml | 5 + .../declarative/expanders/record_expander.py | 112 +++++++++-- .../models/declarative_component_schema.py | 5 + .../parsers/model_to_component_factory.py | 1 + .../expanders/test_record_expander.py | 181 ++++++++++++++++++ .../test_model_to_component_factory.py | 87 +++++++++ 6 files changed, 371 insertions(+), 20 deletions(-) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 62e8209660..4ef4215e5d 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2256,6 +2256,11 @@ definitions: type: array items: "$ref": "#/definitions/ParentFieldPath" + merge_parent: + title: Merge Parent + description: If true, each expanded item is the parent record shallow-merged with the item, the item's own keys winning on collision, and the expanded list removed from the parent's copy. Only the value at `expand_records_from_field` is removed, so for a multi-segment path the top-level key stays with its other fields. The merge happens first, then `parent_fields` are copied, then `original_record` is embedded when `remain_original_record` is set; all three may be combined. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it. + type: boolean + default: false on_no_records: title: On No Records description: Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged. diff --git a/airbyte_cdk/sources/declarative/expanders/record_expander.py b/airbyte_cdk/sources/declarative/expanders/record_expander.py index 99125d5ae9..4dfd6b0caa 100644 --- a/airbyte_cdk/sources/declarative/expanders/record_expander.py +++ b/airbyte_cdk/sources/declarative/expanders/record_expander.py @@ -115,6 +115,10 @@ class RecordExpander: deep-copying the whole parent once per item, which matters when the parent is large and the nested list is long. + Set `merge_parent: true` to flatten the parent into each item instead: the parent's fields, + minus the expanded list, are shallow-merged underneath the item's own fields, so the item wins + on any key both have. + The expand_records_from_field path supports wildcards (*) for matching multiple arrays. When wildcards are used, items from all matched arrays are extracted and emitted. @@ -139,6 +143,14 @@ class RecordExpander: record_path: ["pull_request_url"] ``` + ``` + record_expander: + type: RecordExpander + expand_records_from_field: + - "activity" + merge_parent: true + ``` + ``` record_expander: type: RecordExpander @@ -162,6 +174,13 @@ class RecordExpander: named fields is the cheaper way to carry parent context. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. Glob metacharacters are rejected in both paths. + merge_parent: If True, each expanded item is the parent record shallow-merged with the + item, the item's own keys winning on collision, and the expanded list removed from + the parent's copy. Only the value at `expand_records_from_field` is removed: for a + multi-segment path the top-level key stays with its other fields. The merge happens + before `parent_fields` are copied and before `original_record` is embedded, so both + can overwrite merged values. Applies to items fetched through + `truncated_list_retriever` as well as to embedded ones. Defaults to False. on_no_records: Behavior when expansion produces no records. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged. truncation_indicator_path: Path within each record to a field indicating that the @@ -200,6 +219,7 @@ class RecordExpander: parameters: InitVar[Mapping[str, Any]] remain_original_record: bool = False parent_fields: Optional[Sequence[ParentFieldPath]] = None + merge_parent: bool = False on_no_records: OnNoRecords = OnNoRecords.skip truncation_indicator_path: Optional[Sequence[str]] = None truncated_list_retriever: Optional["Retriever"] = None @@ -259,13 +279,18 @@ def expand_record(self, record: MutableMapping[Any, Any]) -> Iterable[MutableMap expand_path = self._evaluated_expand_path() truncated = bool(self._truncation_indicator_path) and self._is_truncated(parent_record) + # Built once per parent, not once per item: the walk that finds the expanded list costs + # as much as extracting it does. + merge_base = ( + self._without_expanded_list(parent_record, expand_path) if self.merge_parent else None + ) if truncated and self.truncated_list_retriever: # Streamed, so the shortfall is only known once the retriever is exhausted. If the # consumer stops early the fetch was cut short by it, not by the API, and no warning # would be accurate anyway. fetched_count = 0 - for fetched in self._fetch_complete_list(parent_record): + for fetched in self._fetch_complete_list(parent_record, merge_base): fetched_count += 1 yield fetched self._warn_if_fetch_incomplete(parent_record, expand_path, fetched_count) @@ -288,13 +313,9 @@ def expand_record(self, record: MutableMapping[Any, Any]) -> Iterable[MutableMap for items in embedded_lists: for item in items: if isinstance(item, dict): - expanded_record = dict(item) - self._apply_parent_context(parent_record, expanded_record) - yield expanded_record + yield self._apply_parent_context(parent_record, dict(item), merge_base) elif self._carries_parent_context(): - wrapped: MutableMapping[str, Any] = {"value": item} - self._apply_parent_context(parent_record, wrapped) - yield wrapped + yield self._apply_parent_context(parent_record, {"value": item}, merge_base) else: yield item @@ -376,7 +397,9 @@ def _is_truncated(self, parent_record: MutableMapping[Any, Any]) -> bool: except (KeyError, ValueError): return False - def _fetch_complete_list(self, parent_record: Mapping[str, Any]) -> Iterable[Any]: + def _fetch_complete_list( + self, parent_record: Mapping[str, Any], merge_base: Optional[Mapping[str, Any]] + ) -> Iterable[Any]: if not self.truncated_list_retriever: return stream_slice = StreamSlice(partition={"parent_record": parent_record}, cursor_slice={}) @@ -394,24 +417,73 @@ def _fetch_complete_list(self, parent_record: Mapping[str, Any]) -> Iterable[Any else: data = item if isinstance(data, Mapping): - expanded_record = dict(data) - self._apply_parent_context(parent_record, expanded_record) - yield expanded_record + yield self._apply_parent_context(parent_record, dict(data), merge_base) elif self._carries_parent_context(): - wrapped: MutableMapping[str, Any] = {"value": data} - self._apply_parent_context(parent_record, wrapped) - yield wrapped + yield self._apply_parent_context(parent_record, {"value": data}, merge_base) else: yield data def _apply_parent_context( - self, parent_record: Mapping[str, Any], child_record: MutableMapping[str, Any] - ) -> None: - """Apply parent context to a child record.""" - if self.remain_original_record: - child_record["original_record"] = copy.deepcopy(parent_record) + self, + parent_record: Mapping[str, Any], + child_record: MutableMapping[str, Any], + merge_base: Optional[Mapping[str, Any]], + ) -> MutableMapping[str, Any]: + """Return the child record carrying the configured parent context. + + The order is fixed: the parent is merged underneath the child first, then `parent_fields` + are copied on top, then `original_record` is embedded. A named copy can therefore + overwrite a merged value. + """ + if merge_base is not None: + child_record = {**merge_base, **child_record} for parent_field in self.parent_fields or []: parent_field.copy_onto(parent_record, child_record) + if self.remain_original_record: + child_record["original_record"] = copy.deepcopy(parent_record) + return child_record def _carries_parent_context(self) -> bool: - return bool(self.remain_original_record or self.parent_fields) + return bool(self.remain_original_record or self.parent_fields or self.merge_parent) + + @classmethod + def _without_expanded_list( + cls, parent_record: Mapping[str, Any], expand_path: list[Any] + ) -> Mapping[str, Any]: + """Copy of the parent with only the value at `expand_path` removed. + + The containers on the way to the list are copied and the rest of the parent is shared, + so siblings of the list are kept and the parent itself is never mutated. Every match of + the path is removed, which covers glob paths that select several lists. + """ + matched_paths = [ + path + for path, _ in dpath.segments.walk(parent_record) # type: ignore[no-untyped-call] + if dpath.segments.match(path, expand_path) + ] + stripped: Any = parent_record + # Deepest and right-most first, so removing a list element never shifts the index of a + # match still to be removed, and a match nested inside another match goes first. + for path in sorted(matched_paths, key=cls._path_sort_key, reverse=True): + stripped = cls._without_path(stripped, path) + return dict(stripped) + + @staticmethod + def _path_sort_key(path: Sequence[Any]) -> list[tuple[int, Any]]: + # Integers and strings are not mutually comparable, so tag each segment by kind. + return [(0, segment) if isinstance(segment, int) else (1, str(segment)) for segment in path] + + @classmethod + def _without_path(cls, container: Any, path: Sequence[Any]) -> Any: + head, rest = path[0], path[1:] + if isinstance(container, Mapping): + copied: Any = dict(container) + elif isinstance(container, list): + copied = list(container) + else: + return container + if rest: + copied[head] = cls._without_path(copied[head], rest) + else: + del copied[head] + return copied diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 4345e40df2..8aab6fea73 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2322,6 +2322,11 @@ class RecordExpander(BaseModel): description="Named values to copy from the record being expanded onto each expanded item. Use this instead of `remain_original_record` when only a few parent fields are needed: it copies the named values rather than deep-copying the whole parent once per item, which matters when the parent record is large and the nested list is long. An existing value at `record_path` is overwritten, and a `parent_path` the parent does not have copies null. Independent of `remain_original_record`; both may be set. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.", title="Parent Fields", ) + merge_parent: Optional[bool] = Field( + False, + description="If true, each expanded item is the parent record shallow-merged with the item, the item's own keys winning on collision, and the expanded list removed from the parent's copy. Only the value at `expand_records_from_field` is removed, so for a multi-segment path the top-level key stays with its other fields. The merge happens first, then `parent_fields` are copied, then `original_record` is embedded when `remain_original_record` is set; all three may be combined. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.", + title="Merge Parent", + ) on_no_records: Optional[OnNoRecords] = Field( OnNoRecords.skip, description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', 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 6dab0a84e9..f5242ef700 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2556,6 +2556,7 @@ def create_record_expander( ] if model.parent_fields else None, + merge_parent=model.merge_parent or False, on_no_records=OnNoRecords(model.on_no_records.value) if model.on_no_records else OnNoRecords.skip, diff --git a/unit_tests/sources/declarative/expanders/test_record_expander.py b/unit_tests/sources/declarative/expanders/test_record_expander.py index 42fc37109a..0a7bdbe1e4 100644 --- a/unit_tests/sources/declarative/expanders/test_record_expander.py +++ b/unit_tests/sources/declarative/expanders/test_record_expander.py @@ -790,3 +790,184 @@ def test_parent_field_paths_interpolate_config(): parent_field.copy_onto({"url": "https://example.com/7"}, child) assert child == {"pull_request_url": "https://example.com/7"} + + +def _mailchimp_parent(): + return { + "email_id": "e1", + "list_id": "l1", + "activity": [ + {"action": "open", "timestamp": "t1"}, + {"action": "click", "timestamp": "t2"}, + ], + } + + +def _merge_expander(expand_records_from_field, **kwargs): + return RecordExpander( + expand_records_from_field=expand_records_from_field, + merge_parent=True, + config=config, + parameters=parameters, + **kwargs, + ) + + +def test_merge_parent_flattens_the_parent_into_every_item(): + """The source-mailchimp `email_activity` shape: `{**record, **activity_item}` per item.""" + expander = _merge_expander(["activity"]) + + records = list(expander.expand_record(_mailchimp_parent())) + + assert records == [ + {"email_id": "e1", "list_id": "l1", "action": "open", "timestamp": "t1"}, + {"email_id": "e1", "list_id": "l1", "action": "click", "timestamp": "t2"}, + ] + assert all("activity" not in record for record in records) + + +def test_merge_parent_lets_the_item_win_on_collision(): + expander = _merge_expander(["items"]) + + assert list(expander.expand_record({"id": "parent", "items": [{"id": "child"}]})) == [ + {"id": "child"} + ] + + +def test_merge_parent_removes_only_the_expanded_list_on_a_multi_segment_path(): + expander = _merge_expander(["reviews", "nodes"]) + parent = { + "number": 7, + "reviews": {"nodes": [{"id": "PRR_1"}], "totalCount": 1}, + } + + assert list(expander.expand_record(parent)) == [ + {"number": 7, "reviews": {"totalCount": 1}, "id": "PRR_1"} + ] + + +def test_merge_parent_removes_every_list_matched_by_a_glob_path(): + expander = _merge_expander(["sections", "*", "items"]) + parent = { + "id": 1, + "sections": { + "a": {"items": [{"n": 1}], "title": "A"}, + "b": {"items": [{"n": 2}], "title": "B"}, + }, + } + + assert list(expander.expand_record(parent)) == [ + {"id": 1, "sections": {"a": {"title": "A"}, "b": {"title": "B"}}, "n": 1}, + {"id": 1, "sections": {"a": {"title": "A"}, "b": {"title": "B"}}, "n": 2}, + ] + + +def test_merge_parent_removes_lists_nested_in_a_list_of_sections(): + expander = _merge_expander(["sections", "*", "items"]) + parent = {"sections": [{"items": [{"n": 1}], "k": "s0"}, {"items": [{"n": 2}], "k": "s1"}]} + + assert list(expander.expand_record(parent)) == [ + {"sections": [{"k": "s0"}, {"k": "s1"}], "n": 1}, + {"sections": [{"k": "s0"}, {"k": "s1"}], "n": 2}, + ] + + +def test_parent_fields_applied_after_merge_parent_overwrite_a_merged_value(): + expander = _merge_expander( + ["items"], + parent_fields=[_parent_field(["id"], ["id"])], + ) + + assert list(expander.expand_record({"id": "parent", "items": [{"id": "child"}]})) == [ + {"id": "parent"} + ] + + +def test_merge_parent_and_remain_original_record_are_independent(): + expander = _merge_expander(["activity"], remain_original_record=True) + + records = list(expander.expand_record(_mailchimp_parent())) + + assert records == [ + { + "email_id": "e1", + "list_id": "l1", + "action": "open", + "timestamp": "t1", + "original_record": _mailchimp_parent(), + }, + { + "email_id": "e1", + "list_id": "l1", + "action": "click", + "timestamp": "t2", + "original_record": _mailchimp_parent(), + }, + ] + + +def test_merge_parent_wraps_scalar_items(): + expander = _merge_expander(["items"]) + + assert list(expander.expand_record({"id": 9, "items": ["a", "b"]})) == [ + {"id": 9, "value": "a"}, + {"id": 9, "value": "b"}, + ] + + +def test_merge_parent_does_not_mutate_the_parent_record(): + expander = _merge_expander(["reviews", "nodes"]) + parent = _reviews_parent() + + records = list(expander.expand_record(parent)) + + assert parent == _reviews_parent() + assert records[0]["reviews"] == {} + + +def test_merge_parent_shares_untouched_nested_values_rather_than_deep_copying_them(): + expander = _merge_expander(["items"]) + metadata = {"k": "v"} + + records = list(expander.expand_record({"metadata": metadata, "items": [{"id": 1}]})) + + assert records[0]["metadata"] is metadata + + +def test_merge_parent_applies_to_items_fetched_by_the_truncated_list_retriever(): + expander = _merge_expander( + ["reviews", "nodes"], + truncation_indicator_path=["reviews", "has_more"], + truncated_list_retriever=_make_retriever([{"id": "PRR_9"}, "scalar"]), + ) + parent = { + "url": "https://github.com/airbytehq/airbyte/pull/7", + "reviews": {"nodes": [{"id": "PRR_1"}], "has_more": True}, + } + + assert list(expander.expand_record(parent)) == [ + { + "url": "https://github.com/airbytehq/airbyte/pull/7", + "reviews": {"has_more": True}, + "id": "PRR_9", + }, + { + "url": "https://github.com/airbytehq/airbyte/pull/7", + "reviews": {"has_more": True}, + "value": "scalar", + }, + ] + + +def test_merge_parent_unset_leaves_items_unchanged(): + expander = RecordExpander( + expand_records_from_field=["activity"], + config=config, + parameters=parameters, + ) + + assert expander.merge_parent is False + assert list(expander.expand_record(_mailchimp_parent())) == [ + {"action": "open", "timestamp": "t1"}, + {"action": "click", "timestamp": "t2"}, + ] 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 101f3105e0..dbd8b7f4d4 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 @@ -2084,6 +2084,93 @@ def test_create_record_expander_without_parent_fields_leaves_them_unset(): assert selector.extractor.record_expander.parent_fields is None +def test_create_record_expander_with_merge_parent(): + """The source-mailchimp `email_activity` shape: the parent flattened into each activity item.""" + content = """ + selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["emails"] + record_expander: + type: RecordExpander + expand_records_from_field: ["activity"] + merge_parent: true + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + expander = selector.extractor.record_expander + assert expander is not None + assert expander.merge_parent is True + assert expander.remain_original_record is False + assert expander.parent_fields is None + + response = requests.Response() + response._content = json.dumps( + { + "emails": [ + { + "email_id": "e1", + "list_id": "l1", + "activity": [ + {"action": "open", "timestamp": "t1"}, + {"action": "click", "timestamp": "t2"}, + ], + }, + {"email_id": "e2", "list_id": "l1", "activity": [{"action": "bounce"}]}, + ] + } + ).encode("utf-8") + + assert list(selector.extractor.extract_records(response)) == [ + {"email_id": "e1", "list_id": "l1", "action": "open", "timestamp": "t1"}, + {"email_id": "e1", "list_id": "l1", "action": "click", "timestamp": "t2"}, + {"email_id": "e2", "list_id": "l1", "action": "bounce"}, + ] + + +def test_create_record_expander_without_merge_parent_defaults_to_false(): + content = """ + selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data"] + record_expander: + type: RecordExpander + expand_records_from_field: ["items"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + assert selector.extractor.record_expander.merge_parent is False + + def test_create_record_expander_with_truncated_list_retriever(): content = """ selector: From 606c32f3f8aa0b80f30d82cea8017b8151219e95 Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Wed, 16 Sep 2026 18:53:41 +0300 Subject: [PATCH 08/10] fix(RecordExpander): isolate copied parent context from the parent and sibling items `merge_parent` shallow-merged the parent, so every nested container of the parent was aliased into every expanded item and into the parent itself. The same held for a `parent_fields` entry whose `parent_path` pointed at a dict or a list, because `dpath.get` returns the parent's object by reference. A downstream `AddFields`/`RemoveFields` with a multi-segment path then wrote into records already yielded, and into the record being expanded. `merge_base` is now deep-copied per item, and a container copied by `ParentFieldPath.copy_onto` is deep-copied as well; scalars pass through untouched. `merge_base` still excludes the expanded list and is still built once per parent record, so the copy stays cheaper than `remain_original_record`. Also: - register `RecordExpander.parent_fields` in `DEFAULT_MODEL_TYPES`, so a manifest may omit `type: ParentFieldPath` as it may for `AddFields.fields` - drop the "when truncation handling is configured" clause from the shared glob-rejection message, which `parent_fields` paths also reach - document `dpath.new` behaviour on `record_path` and the per-item copying in the schema descriptions and docstrings Co-Authored-By: Claude Opus 5 (1M context) --- .../declarative_component_schema.yaml | 6 +-- .../declarative/expanders/record_expander.py | 49 ++++++++++++----- .../models/declarative_component_schema.py | 6 +-- .../parsers/manifest_component_transformer.py | 2 + .../expanders/test_record_expander.py | 39 ++++++++++++-- .../test_model_to_component_factory.py | 52 +++++++++++++++++++ 6 files changed, 131 insertions(+), 23 deletions(-) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 4ef4215e5d..7f3ccd845f 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2252,13 +2252,13 @@ definitions: default: false parent_fields: title: Parent Fields - description: 'Named values to copy from the record being expanded onto each expanded item. Use this instead of `remain_original_record` when only a few parent fields are needed: it copies the named values rather than deep-copying the whole parent once per item, which matters when the parent record is large and the nested list is long. An existing value at `record_path` is overwritten, and a `parent_path` the parent does not have copies null. Independent of `remain_original_record`; both may be set. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.' + description: 'Named values to copy from the record being expanded onto each expanded item. Use this instead of `remain_original_record` when only a few parent fields are needed: it copies the named values rather than deep-copying the whole parent once per item, which matters when the parent record is large and the nested list is long. An existing value at `record_path` is overwritten, and a `parent_path` the parent does not have copies null. A copied object or array is deep-copied, so a transformation that writes inside it affects only that item. Independent of `remain_original_record`; both may be set. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.' type: array items: "$ref": "#/definitions/ParentFieldPath" merge_parent: title: Merge Parent - description: If true, each expanded item is the parent record shallow-merged with the item, the item's own keys winning on collision, and the expanded list removed from the parent's copy. Only the value at `expand_records_from_field` is removed, so for a multi-segment path the top-level key stays with its other fields. The merge happens first, then `parent_fields` are copied, then `original_record` is embedded when `remain_original_record` is set; all three may be combined. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it. + description: If true, each expanded item is the parent record shallow-merged with the item, the item's own keys winning on collision, and the expanded list removed from the parent's copy. Only the value at `expand_records_from_field` is removed, so for a multi-segment path the top-level key stays with its other fields. Each item receives its own deep copy of the merged parent fields, so a transformation that writes into a nested value affects only that item. The merge happens first, then `parent_fields` are copied, then `original_record` is embedded when `remain_original_record` is set; all three may be combined. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it. type: boolean default: false on_no_records: @@ -2314,7 +2314,7 @@ definitions: - ["repository", "name"] record_path: title: Record Path - description: Path on the expanded item to write the value to. An existing value is overwritten. + description: Path on the expanded item to write the value to. An existing value is overwritten. Intermediate objects are created as needed, so a path may not pass through a value the item already holds as a scalar, and a numeric path segment creates an array rather than an object - the same behavior as `AddFields`. type: array items: type: string diff --git a/airbyte_cdk/sources/declarative/expanders/record_expander.py b/airbyte_cdk/sources/declarative/expanders/record_expander.py index 4dfd6b0caa..2f3de73cf6 100644 --- a/airbyte_cdk/sources/declarative/expanders/record_expander.py +++ b/airbyte_cdk/sources/declarative/expanders/record_expander.py @@ -41,6 +41,13 @@ ) +def _detached(value: Any) -> Any: + """Copy of `value` sharing nothing with the record it came from; scalars are returned as-is.""" + if isinstance(value, (Mapping, list, set, tuple)): + return copy.deepcopy(value) + return value + + class OnNoRecords(Enum): """ Behavior when record expansion produces no records. @@ -93,12 +100,16 @@ def copy_onto( A `parent_path` that the parent does not have copies `None`, which is what the custom connector classes this replaces do (`parent.get(field)`). Distinguishing "absent" from "present and null" would need a third option and no connector needs one. + + A container value is deep-copied, so that a downstream transformation writing inside it + cannot reach the parent record or the items expanded from it alongside this one. The copy + costs proportionally to the named value, not to the whole parent. """ try: value = dpath.get(dict(parent_record), self.evaluated_parent_path()) except (KeyError, ValueError): value = None - dpath.new(child_record, self.evaluated_record_path(), value) + dpath.new(child_record, self.evaluated_record_path(), _detached(value)) @dataclass @@ -170,17 +181,20 @@ class RecordExpander: parent_fields: Named values to copy from the record being expanded onto each expanded item. Each entry has a `parent_path` and a `record_path`; an existing value at `record_path` is overwritten, and a `parent_path` the parent does not have copies - `None`. Independent of `remain_original_record` - both may be set, though copying - named fields is the cheaper way to carry parent context. Applies to items fetched - through `truncated_list_retriever` as well as to embedded ones. Glob metacharacters - are rejected in both paths. + `None`. A copied container is deep-copied, so writing into it downstream does not + reach the parent or the sibling items. Independent of `remain_original_record` - + both may be set, though copying named fields is the cheaper way to carry parent + context. Applies to items fetched through `truncated_list_retriever` as well as to + embedded ones. Glob metacharacters are rejected in both paths. merge_parent: If True, each expanded item is the parent record shallow-merged with the item, the item's own keys winning on collision, and the expanded list removed from the parent's copy. Only the value at `expand_records_from_field` is removed: for a - multi-segment path the top-level key stays with its other fields. The merge happens - before `parent_fields` are copied and before `original_record` is embedded, so both - can overwrite merged values. Applies to items fetched through - `truncated_list_retriever` as well as to embedded ones. Defaults to False. + multi-segment path the top-level key stays with its other fields. Each item gets its + own deep copy of the merged parent, so a downstream transformation writing into a + nested value affects only that item. The merge happens before `parent_fields` are + copied and before `original_record` is embedded, so both can overwrite merged + values. Applies to items fetched through `truncated_list_retriever` as well as to + embedded ones. Defaults to False. on_no_records: Behavior when expansion produces no records. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged. truncation_indicator_path: Path within each record to a field indicating that the @@ -254,8 +268,8 @@ def _reject_globs(path: Sequence[Any], field_name: str) -> None: for segment in path ): raise ValueError( - f"Glob characters {_GLOB_METACHARACTERS} are not supported in `{field_name}` when " - "truncation handling is configured: the path must identify a single field." + f"Glob characters {_GLOB_METACHARACTERS} are not supported in `{field_name}`: " + "the path must identify a single field." ) def _evaluated_indicator_path(self) -> list[Any]: @@ -436,7 +450,12 @@ def _apply_parent_context( overwrite a merged value. """ if merge_base is not None: - child_record = {**merge_base, **child_record} + # Deep-copied per item: a shallow merge would alias every nested container of the + # parent into every item, so a downstream transformation writing into one of them + # would reach the parent record and the items already yielded. The expanded list is + # already stripped from `merge_base`, so the copy costs less than + # `remain_original_record` does. + child_record = {**copy.deepcopy(merge_base), **child_record} for parent_field in self.parent_fields or []: parent_field.copy_onto(parent_record, child_record) if self.remain_original_record: @@ -453,8 +472,10 @@ def _without_expanded_list( """Copy of the parent with only the value at `expand_path` removed. The containers on the way to the list are copied and the rest of the parent is shared, - so siblings of the list are kept and the parent itself is never mutated. Every match of - the path is removed, which covers glob paths that select several lists. + so siblings of the list are kept and the parent itself is not mutated here. This base is + built once per parent record and then deep-copied per item in `_apply_parent_context`, + which is what keeps the parent and the sibling items isolated from one another. Every + match of the path is removed, which covers glob paths that select several lists. """ matched_paths = [ path diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 8aab6fea73..6c944d8679 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -532,7 +532,7 @@ class ParentFieldPath(BaseModel): ) record_path: List[str] = Field( ..., - description="Path on the expanded item to write the value to. An existing value is overwritten.", + description="Path on the expanded item to write the value to. An existing value is overwritten. Intermediate objects are created as needed, so a path may not pass through a value the item already holds as a scalar, and a numeric path segment creates an array rather than an object - the same behavior as `AddFields`.", examples=[["pull_request_url"], ["comment_id"], ["parent", "id"]], title="Record Path", ) @@ -2319,12 +2319,12 @@ class RecordExpander(BaseModel): ) parent_fields: Optional[List[ParentFieldPath]] = Field( None, - description="Named values to copy from the record being expanded onto each expanded item. Use this instead of `remain_original_record` when only a few parent fields are needed: it copies the named values rather than deep-copying the whole parent once per item, which matters when the parent record is large and the nested list is long. An existing value at `record_path` is overwritten, and a `parent_path` the parent does not have copies null. Independent of `remain_original_record`; both may be set. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.", + description="Named values to copy from the record being expanded onto each expanded item. Use this instead of `remain_original_record` when only a few parent fields are needed: it copies the named values rather than deep-copying the whole parent once per item, which matters when the parent record is large and the nested list is long. An existing value at `record_path` is overwritten, and a `parent_path` the parent does not have copies null. A copied object or array is deep-copied, so a transformation that writes inside it affects only that item. Independent of `remain_original_record`; both may be set. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.", title="Parent Fields", ) merge_parent: Optional[bool] = Field( False, - description="If true, each expanded item is the parent record shallow-merged with the item, the item's own keys winning on collision, and the expanded list removed from the parent's copy. Only the value at `expand_records_from_field` is removed, so for a multi-segment path the top-level key stays with its other fields. The merge happens first, then `parent_fields` are copied, then `original_record` is embedded when `remain_original_record` is set; all three may be combined. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.", + description="If true, each expanded item is the parent record shallow-merged with the item, the item's own keys winning on collision, and the expanded list removed from the parent's copy. Only the value at `expand_records_from_field` is removed, so for a multi-segment path the top-level key stays with its other fields. Each item receives its own deep copy of the merged parent fields, so a transformation that writes into a nested value affects only that item. The merge happens first, then `parent_fields` are copied, then `original_record` is embedded when `remain_original_record` is set; all three may be combined. Applies to items fetched through `truncated_list_retriever` as well as to embedded ones. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.", title="Merge Parent", ) on_no_records: Optional[OnNoRecords] = Field( diff --git a/airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py b/airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py index 00b1a18ffb..e11a0af515 100644 --- a/airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py +++ b/airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py @@ -50,6 +50,8 @@ # ParentStreamConfig "ParentStreamConfig.request_option": "RequestOption", "ParentStreamConfig.stream": "DeclarativeStream", + # RecordExpander + "RecordExpander.parent_fields": "ParentFieldPath", # RecordSelector "RecordSelector.extractor": "DpathExtractor", "RecordSelector.record_filter": "RecordFilter", diff --git a/unit_tests/sources/declarative/expanders/test_record_expander.py b/unit_tests/sources/declarative/expanders/test_record_expander.py index 0a7bdbe1e4..26269e9f63 100644 --- a/unit_tests/sources/declarative/expanders/test_record_expander.py +++ b/unit_tests/sources/declarative/expanders/test_record_expander.py @@ -772,6 +772,14 @@ def test_globs_are_rejected_in_record_path(path): _parent_field(["id"], path) +def test_glob_rejection_message_does_not_mention_truncation_handling(): + with pytest.raises(ValueError) as error: + _parent_field(["*"], ["copied"]) + + assert "truncation" not in str(error.value) + assert "`parent_path`" in str(error.value) + + @pytest.mark.parametrize("parent_path,record_path", [([], ["a"]), (["a"], [])]) def test_empty_paths_are_rejected(parent_path, record_path): with pytest.raises(ValueError, match="cannot be empty"): @@ -925,13 +933,38 @@ def test_merge_parent_does_not_mutate_the_parent_record(): assert records[0]["reviews"] == {} -def test_merge_parent_shares_untouched_nested_values_rather_than_deep_copying_them(): +def test_merge_parent_gives_each_item_its_own_copy_of_nested_values(): expander = _merge_expander(["items"]) metadata = {"k": "v"} + parent = {"metadata": metadata, "items": [{"id": 1}, {"id": 2}]} + + records = list(expander.expand_record(parent)) + + assert records[0]["metadata"] is not metadata + assert records[0]["metadata"] is not records[1]["metadata"] + # A downstream transformation writing into one item cannot reach the parent or its siblings. + records[0]["metadata"]["written"] = True + assert records[1]["metadata"] == {"k": "v"} + assert metadata == {"k": "v"} + - records = list(expander.expand_record({"metadata": metadata, "items": [{"id": 1}]})) +def test_parent_fields_gives_each_item_its_own_copy_of_a_copied_container(): + expander = RecordExpander( + expand_records_from_field=["reviews", "nodes"], + parent_fields=[_parent_field(["repository"], ["repository"])], + config=config, + parameters=parameters, + ) + repository = {"name": "airbyte"} + parent = {"repository": repository, "reviews": {"nodes": [{"id": 1}, {"id": 2}]}} + + records = list(expander.expand_record(parent)) - assert records[0]["metadata"] is metadata + assert records[0]["repository"] is not repository + assert records[0]["repository"] is not records[1]["repository"] + records[0]["repository"]["review_id"] = 1 + assert records[1]["repository"] == {"name": "airbyte"} + assert repository == {"name": "airbyte"} def test_merge_parent_applies_to_items_fetched_by_the_truncated_list_retriever(): 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 dbd8b7f4d4..a1f612e479 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 @@ -2055,6 +2055,58 @@ def test_create_record_expander_with_parent_fields(): ] +def test_create_record_expander_with_parent_fields_without_explicit_type(): + """`type: ParentFieldPath` is optional: the transformer fills it in, as it does for AddFields.""" + content = """ + selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["nodes"] + record_expander: + type: RecordExpander + expand_records_from_field: ["reviews", "nodes"] + parent_fields: + - parent_path: ["url"] + record_path: ["pull_request_url"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + assert ( + selector_manifest["extractor"]["record_expander"]["parent_fields"][0]["type"] + == "ParentFieldPath" + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + response = requests.Response() + response._content = json.dumps( + { + "nodes": [ + { + "url": "https://github.com/airbytehq/airbyte/pull/7", + "reviews": {"nodes": [{"id": "PRR_1"}]}, + } + ] + } + ).encode("utf-8") + + assert list(selector.extractor.extract_records(response)) == [ + {"id": "PRR_1", "pull_request_url": "https://github.com/airbytehq/airbyte/pull/7"}, + ] + + def test_create_record_expander_without_parent_fields_leaves_them_unset(): content = """ selector: From 92daa4ebb93611df582a2188fcd1c751b18e4166 Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Mon, 28 Sep 2026 18:42:15 +0300 Subject: [PATCH 09/10] fix(RecordExpander): make parent_fields lookup linear and harden merge_parent stripping - Look up `parent_path` by walking only the containers on the path instead of calling `dpath.get`, which folds over the whole parent (the expanded list included) for every item and every field. The expansion was quadratic in the list length and slower than `remain_original_record`. Paths are evaluated once at construction. - Raise a config-error `AirbyteTracedException` naming the `record_path` when it runs through a value the item holds as a scalar or a list, instead of a raw `PathNotFound`/`TypeError`. - `merge_parent` strips only the list matches of a glob path; non-list matches were never expanded and are kept. - Pin the strip order, `$parameters` propagation into `parent_fields`, and the merge-then-copy order with tests. Co-Authored-By: Claude Opus 5.5 --- .../declarative/expanders/record_expander.py | 71 ++++++++-- .../expanders/test_record_expander.py | 132 +++++++++++++++++- .../test_model_to_component_factory.py | 49 +++++++ 3 files changed, 240 insertions(+), 12 deletions(-) diff --git a/airbyte_cdk/sources/declarative/expanders/record_expander.py b/airbyte_cdk/sources/declarative/expanders/record_expander.py index 2f3de73cf6..19693e2529 100644 --- a/airbyte_cdk/sources/declarative/expanders/record_expander.py +++ b/airbyte_cdk/sources/declarative/expanders/record_expander.py @@ -17,12 +17,14 @@ AirbyteMessage, AirbyteStateMessage, AirbyteTraceMessage, + FailureType, Level, ) from airbyte_cdk.models import Type as MessageType from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString from airbyte_cdk.sources.message import MessageRepository from airbyte_cdk.sources.types import Config, Record, StreamSlice +from airbyte_cdk.utils.traced_exception import AirbyteTracedException if TYPE_CHECKING: from airbyte_cdk.sources.declarative.retrievers import Retriever @@ -32,6 +34,9 @@ # dpath treats these characters as glob metacharacters (fnmatch semantics) inside a path segment. _GLOB_METACHARACTERS = ("*", "?", "[") +# Returned by `_lookup` for a path the record does not have, since `None` is a valid value. +_MISSING = object() + # Protocol payloads a retriever may yield unwrapped (i.e. not inside an `AirbyteMessage` envelope). _BARE_PROTOCOL_MESSAGES = ( AirbyteControlMessage, @@ -41,6 +46,34 @@ ) +def _lookup(record: Any, path: Sequence[Any]) -> Any: + """Value at a glob-free `path`, or `_MISSING`, reading only the containers on the path. + + Matches the way `dpath.get` resolves such a path: a mapping key is compared as a string, and + a segment reaching a list is taken as an index, negative ones counting from the end. + `dpath.get` itself walks the whole record, so on a parent holding a long list each lookup + costs as much as that list. + """ + value = record + for segment in path: + if isinstance(value, Mapping): + key = str(segment) + if key not in value: + return _MISSING + value = value[key] + elif isinstance(value, list): + try: + index = int(segment) + except (TypeError, ValueError): + return _MISSING + if not -len(value) <= index < len(value): + return _MISSING + value = value[index] + else: + return _MISSING + return value + + def _detached(value: Any) -> Any: """Copy of `value` sharing nothing with the record it came from; scalars are returned as-is.""" if isinstance(value, (Mapping, list, set, tuple)): @@ -83,8 +116,11 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: self._record_path: list[InterpolatedString] = [ InterpolatedString.create(path, parameters=parameters) for path in self.record_path ] - RecordExpander._reject_globs(self.evaluated_parent_path(), "parent_path") - RecordExpander._reject_globs(self.evaluated_record_path(), "record_path") + # The paths only interpolate `config`, so they are evaluated once instead of per item. + self._evaluated_parent_path = self.evaluated_parent_path() + self._evaluated_record_path = self.evaluated_record_path() + RecordExpander._reject_globs(self._evaluated_parent_path, "parent_path") + RecordExpander._reject_globs(self._evaluated_record_path, "record_path") def evaluated_parent_path(self) -> list[Any]: return [segment.eval(self.config) for segment in self._parent_path] @@ -102,14 +138,26 @@ def copy_onto( "present and null" would need a third option and no connector needs one. A container value is deep-copied, so that a downstream transformation writing inside it - cannot reach the parent record or the items expanded from it alongside this one. The copy - costs proportionally to the named value, not to the whole parent. + cannot reach the parent record or the items expanded from it alongside this one. Both the + lookup and the copy cost proportionally to the named value, not to the whole parent. + + A `record_path` running through a value the item holds as a scalar or a list cannot be + written and fails the sync as a config error naming the path. """ + value = _lookup(parent_record, self._evaluated_parent_path) + copied = None if value is _MISSING else _detached(value) try: - value = dpath.get(dict(parent_record), self.evaluated_parent_path()) - except (KeyError, ValueError): - value = None - dpath.new(child_record, self.evaluated_record_path(), _detached(value)) + dpath.new(child_record, self._evaluated_record_path, copied) + except (dpath.exceptions.PathNotFound, TypeError) as error: + message = ( + f"Cannot write parent field to `record_path` {self._evaluated_record_path}: " + "the expanded item already holds a non-object value on that path." + ) + raise AirbyteTracedException( + message=message, + internal_message=f"{message} {error!r}", + failure_type=FailureType.config_error, + ) from error @dataclass @@ -475,12 +523,13 @@ def _without_expanded_list( so siblings of the list are kept and the parent itself is not mutated here. This base is built once per parent record and then deep-copied per item in `_apply_parent_context`, which is what keeps the parent and the sibling items isolated from one another. Every - match of the path is removed, which covers glob paths that select several lists. + list matched by the path is removed, which covers glob paths that select several lists; + a match that is not a list is kept, since nothing is expanded from it. """ matched_paths = [ path - for path, _ in dpath.segments.walk(parent_record) # type: ignore[no-untyped-call] - if dpath.segments.match(path, expand_path) + for path, value in dpath.segments.walk(parent_record) # type: ignore[no-untyped-call] + if isinstance(value, list) and dpath.segments.match(path, expand_path) ] stripped: Any = parent_record # Deepest and right-most first, so removing a list element never shifts the index of a diff --git a/unit_tests/sources/declarative/expanders/test_record_expander.py b/unit_tests/sources/declarative/expanders/test_record_expander.py index 26269e9f63..d73a0f1b9e 100644 --- a/unit_tests/sources/declarative/expanders/test_record_expander.py +++ b/unit_tests/sources/declarative/expanders/test_record_expander.py @@ -5,15 +5,24 @@ import threading from unittest.mock import MagicMock +import dpath import pytest -from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteRecordMessage, Level, Type +from airbyte_cdk.models import ( + AirbyteLogMessage, + AirbyteMessage, + AirbyteRecordMessage, + FailureType, + Level, + Type, +) from airbyte_cdk.sources.declarative.expanders.record_expander import ( ParentFieldPath, RecordExpander, ) from airbyte_cdk.sources.message import InMemoryMessageRepository from airbyte_cdk.sources.types import Record, StreamSlice +from airbyte_cdk.utils.traced_exception import AirbyteTracedException config = {} parameters = {} @@ -800,6 +809,82 @@ def test_parent_field_paths_interpolate_config(): assert child == {"pull_request_url": "https://example.com/7"} +@pytest.mark.parametrize( + "parent,parent_path", + [ + ({"a": {"b": 1}}, ["a", "b"]), + ({"a": {"b": None}}, ["a", "b"]), + ({"a": {"b": {"c": [1, 2]}}}, ["a", "b"]), + ({"a": [{"b": 1}, {"b": 2}]}, ["a", "1", "b"]), + ({"a": [{"b": 1}, {"b": 2}]}, ["a", "01", "b"]), + ({"a": [{"b": 1}]}, ["a", "1", "b"]), + ({"a": [{"b": 1}, {"b": 2}]}, ["a", "-1", "b"]), + ({"a": [{"b": 1}]}, ["a", "-2", "b"]), + ({"a": [{"b": 1}]}, ["a", "x"]), + ({"a": {"0": "zero"}}, ["a", "0"]), + ({"a": "text"}, ["a", "0"]), + ({"a": 1}, ["a", "b"]), + ({}, ["a"]), + ], +) +def test_parent_field_lookup_matches_dpath_get(parent, parent_path): + """The direct lookup replacing `dpath.get` resolves every path the way `dpath.get` did.""" + try: + expected = dpath.get(parent, parent_path) + except (KeyError, ValueError): + expected = None + child = {} + + _parent_field(parent_path, ["copied"]).copy_onto(parent, child) + + assert child == {"copied": expected} + + +def test_parent_field_lookup_does_not_walk_the_parent(monkeypatch): + """Walking the whole parent per item made `parent_fields` quadratic in the list length.""" + + def fail(*args, **kwargs): + raise AssertionError("the parent was walked") + + monkeypatch.setattr(dpath, "get", fail) + expander = RecordExpander( + expand_records_from_field=["items"], + parent_fields=[_parent_field(["meta", "url"], ["url"])], + config=config, + parameters=parameters, + ) + parent = {"meta": {"url": "u"}, "items": [{"n": n} for n in range(3)]} + + assert list(expander.expand_record(parent)) == [ + {"n": 0, "url": "u"}, + {"n": 1, "url": "u"}, + {"n": 2, "url": "u"}, + ] + + +def test_parent_field_paths_interpolate_parameters(): + parent_field = ParentFieldPath( + parent_path=["{{ parameters['from'] }}"], + record_path=["{{ parameters['to'] }}"], + config=config, + parameters={"from": "url", "to": "pull_request_url"}, + ) + child = {} + + parent_field.copy_onto({"url": "https://example.com/7"}, child) + + assert child == {"pull_request_url": "https://example.com/7"} + + +@pytest.mark.parametrize("blocking_value", [1, [1]], ids=["scalar", "list"]) +def test_record_path_through_a_non_object_value_raises_a_config_error(blocking_value): + with pytest.raises(AirbyteTracedException) as error: + _parent_field(["id"], ["q", "x"]).copy_onto({"id": 1}, {"q": blocking_value}) + + assert error.value.failure_type == FailureType.config_error + assert "['q', 'x']" in error.value.message + + def _mailchimp_parent(): return { "email_id": "e1", @@ -880,6 +965,51 @@ def test_merge_parent_removes_lists_nested_in_a_list_of_sections(): ] +def test_merge_parent_removes_several_lists_from_the_same_list(): + """Removal runs right-most first; left to right, the second index would have shifted.""" + expander = _merge_expander(["sec", "*"]) + parent = {"id": 1, "sec": [[{"n": 1}], [{"n": 2}], [{"n": 3}]]} + + assert list(expander.expand_record(parent)) == [ + {"id": 1, "sec": [], "n": 1}, + {"id": 1, "sec": [], "n": 2}, + {"id": 1, "sec": [], "n": 3}, + ] + + +def test_merge_parent_removes_a_list_matched_inside_another_matched_list(): + """Removal runs deepest first, so the outer list is still there when the inner one goes.""" + expander = _merge_expander(["**"]) + parent = {"outer": [{"n": 1}, [{"n": 2}]]} + merge_base = expander._without_expanded_list(parent, ["**"]) + + assert merge_base == {} + assert parent == {"outer": [{"n": 1}, [{"n": 2}]]} + + +def test_merge_parent_keeps_glob_matches_that_are_not_lists(): + expander = _merge_expander(["sec", "*"]) + parent = {"sec": [[{"n": 1}], "s", [{"n": 2}]]} + + assert list(expander.expand_record(parent)) == [ + {"sec": ["s"], "n": 1}, + {"sec": ["s"], "n": 2}, + ] + + +def test_parent_fields_write_into_a_merged_nested_value(): + """Merge first, then the named copy: the copy lands inside the merged object, not over it.""" + expander = _merge_expander( + ["items"], + parent_fields=[_parent_field(["id"], ["meta", "parent_id"])], + ) + parent = {"id": 7, "meta": {"a": 1}, "items": [{"n": 1}]} + + assert list(expander.expand_record(parent)) == [ + {"id": 7, "meta": {"a": 1, "parent_id": 7}, "n": 1}, + ] + + def test_parent_fields_applied_after_merge_parent_overwrite_a_merged_value(): expander = _merge_expander( ["items"], 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 71bf059192..af5be7167c 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 @@ -2118,6 +2118,55 @@ def test_create_record_expander_with_parent_fields_without_explicit_type(): ] +def test_create_record_expander_propagates_parameters_into_parent_fields(): + content = """ + selector: + type: RecordSelector + $parameters: + parent_key: "url" + child_key: "pull_request_url" + extractor: + type: DpathExtractor + field_path: ["nodes"] + record_expander: + type: RecordExpander + expand_records_from_field: ["reviews", "nodes"] + parent_fields: + - parent_path: ["{{ parameters['parent_key'] }}"] + record_path: ["{{ parameters['child_key'] }}"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + response = requests.Response() + response._content = json.dumps( + { + "nodes": [ + { + "url": "https://github.com/airbytehq/airbyte/pull/7", + "reviews": {"nodes": [{"id": "PRR_1"}]}, + } + ] + } + ).encode("utf-8") + + assert list(selector.extractor.extract_records(response)) == [ + {"id": "PRR_1", "pull_request_url": "https://github.com/airbytehq/airbyte/pull/7"}, + ] + + def test_create_record_expander_without_parent_fields_leaves_them_unset(): content = """ selector: From e78f1efbdad4a21cd5ac62e7b2ba506d51fb597e Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Mon, 28 Sep 2026 18:45:10 +0300 Subject: [PATCH 10/10] fix(CombinedExtractor): count dropped empty records for paginators, guard the async download paginator - `skip_empty_records` dropped nulls before an `OffsetIncrement` or `PageIncrement` counted the page, so one null on a full page made the page look short and silently ended pagination. The paginator's copy of the extractor now yields a placeholder for every empty record it drops, so the count is what the API returned. The `first_match` winner is still chosen by the surviving records, so the paginator counts the same sub-extractor the record stream reads, and only the record stream logs the drop. - The `union` + `OffsetIncrement` rejection never reached `AsyncRetriever.download_paginator`, which was built without an extractor model and fell back to the retriever's summed count. A `CombinedExtractor` download extractor is now handed to the download paginator, which applies both the rejection and the dropped-record count. Other download extractors keep the retriever's count, because a paginator copy would re-read a response that a streaming download decoder has already consumed. - The documented `first_match` example lacked the trailing `"*"` on its wildcard path, so it emitted each board's item list as one record. Co-Authored-By: Claude Opus 5.5 --- .../declarative_component_schema.yaml | 4 +- .../extractors/combined_extractor.py | 50 ++++- .../models/declarative_component_schema.py | 2 +- .../parsers/model_to_component_factory.py | 56 ++++-- .../extractors/test_combined_extractor.py | 134 +++++++++++++ .../test_model_to_component_factory.py | 183 +++++++++++++++++- .../test_concurrent_declarative_source.py | 72 +++++++ 7 files changed, 481 insertions(+), 20 deletions(-) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 608d509638..8e0ec9b9bf 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2222,7 +2222,7 @@ definitions: default: union skip_empty_records: title: Skip Empty Records - description: 'Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried, and the record count a paginator obtains is the count after the empty records were dropped.' + description: 'Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried. The dropped records still count toward the page size of an OffsetIncrement or PageIncrement paginator, so a full page holding a null does not end pagination early.' type: boolean default: false $parameters: @@ -2241,7 +2241,7 @@ definitions: skip_empty_records: true extractors: - type: DpathExtractor - field_path: ["data", "boards", "*", "items_page", "items"] + field_path: ["data", "boards", "*", "items_page", "items", "*"] - type: DpathExtractor field_path: ["data", "next_items_page", "items"] DpathExtractor: diff --git a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py index 7f366456fc..bb57c421c3 100644 --- a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py +++ b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py @@ -31,6 +31,18 @@ class _NoRecord: _NO_RECORD = _NoRecord() +class _DroppedEmptyRecord: + """Stands in for an empty record dropped by `skip_empty_records` when the extractor counts a page. + + A record-counting paginator only takes the length of what the extractor yields, so the + placeholder makes the dropped record count toward the page size without ever reaching the + record stream. + """ + + +_DROPPED_EMPTY_RECORD = _DroppedEmptyRecord() + + @dataclass class CombinedExtractor(RecordExtractor): """Combines the output of several record extractors into a single stream of records. @@ -67,7 +79,7 @@ class CombinedExtractor(RecordExtractor): mode: first_match extractors: - type: DpathExtractor - field_path: ["data", "boards", "*", "items_page", "items"] + field_path: ["data", "boards", "*", "items_page", "items", "*"] - type: DpathExtractor field_path: ["data", "next_items_page", "items"] ``` @@ -110,6 +122,9 @@ class CombinedExtractor(RecordExtractor): the behaviour of `source-monday`'s `MondayIncrementalItemsExtractor`, where a page of nulls must fall through to the pagination path rather than lock in the primary path. + The dropped records still count toward the size of the page, see "Record-counting paginators" + below. + ## Cost Each sub-extractor decodes the response independently — `CompositeRawDecoder` re-parses the @@ -136,17 +151,30 @@ class CombinedExtractor(RecordExtractor): count, so it works with either paginator as long as the path the count comes from is the one the API paginates over. + Empty records dropped by `skip_empty_records` are part of that count. The API returned them as + slots of the page, so a full page holding a null must still read as full: counting only the + surviving records would make both paginators stop at the first page with a null, and would + make an `OffsetIncrement` re-request records it already read. The factory builds the + paginator's copy of the extractor with `count_dropped_empty_records` set, which makes that copy + yield a placeholder for every empty record it drops. The winner under `first_match` is still + chosen by the surviving records alone, so the paginator counts the same sub-extractor the + record stream reads. + Attributes: extractors (List[RecordExtractor]): The sub-extractors to combine. At least one is required. mode (CombineMode): How the sub-extractor outputs are combined. Defaults to `union`. skip_empty_records (bool): Whether falsy records are dropped before they are combined. Defaults to `False`. + count_dropped_empty_records (bool): Set by the factory on the copy a record-counting + paginator uses, never on the one that emits records: yield a placeholder in place of + every dropped empty record so it counts toward the page size. Defaults to `False`. """ extractors: List[RecordExtractor] parameters: InitVar[Mapping[str, Any]] mode: CombineMode = CombineMode.union skip_empty_records: bool = False + count_dropped_empty_records: bool = False def __post_init__(self, parameters: Mapping[str, Any]) -> None: if not self.extractors: @@ -169,12 +197,21 @@ def _extract_union(self, response: requests.Response) -> Iterable[Mapping[str, A def _extract_first_match(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: for index in range(len(self.extractors)): records: Iterator[Mapping[str, Any]] = iter(self._records_of(index, response)) - first_record: Union[Mapping[str, Any], _NoRecord] = next(records, _NO_RECORD) + # Placeholders for dropped empty records do not make a sub-extractor win, but they are + # kept so the winner's count still includes the ones that came before its first record. + dropped_placeholders: List[Mapping[str, Any]] = [] + first_record: Union[Mapping[str, Any], _NoRecord] = _NO_RECORD + for record in records: + if isinstance(record, _DroppedEmptyRecord): + dropped_placeholders.append(record) + continue + first_record = record + break if isinstance(first_record, _NoRecord): continue # Chain the peeked record back instead of restarting the extractor: the rest of the # records stay lazy and the response is never read twice. - yield from chain([first_record], records) + yield from chain(dropped_placeholders, [first_record], records) return def _records_of( @@ -191,9 +228,14 @@ def _without_empty_records( dropped = 0 for record in records: if not record: - dropped += 1 + if self.count_dropped_empty_records: + yield _DROPPED_EMPTY_RECORD # type: ignore[misc] # only ever counted, never emitted + else: + dropped += 1 continue yield record + # The paginator's copy re-reads the page the record stream already read, so only the + # record stream reports the drop. if dropped: logger.warning( "CombinedExtractor dropped %s empty record(s) yielded by sub-extractor %s (%s) " diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index e17e5e0258..cd97715a2d 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2608,7 +2608,7 @@ class CombinedExtractor(BaseModel): ) skip_empty_records: Optional[bool] = Field( False, - description='Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried, and the record count a paginator obtains is the count after the empty records were dropped.', + description='Whether to drop empty records - null, {}, [], "" - yielded by a sub-extractor before the records are combined. Off by default, so the output of a sub-extractor is passed through as it is. Turn it on when the API can return nulls in the middle of a record list, which a GraphQL API does when it answers a partial response and reports the failure in a sibling error field. Under "first_match" this also changes which sub-extractor wins: a sub-extractor whose records are all empty no longer counts as a match, so the next one is tried. The dropped records still count toward the page size of an OffsetIncrement or PageIncrement paginator, so a full page holding a null does not end pagination early.', title="Skip Empty Records", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") 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 369a13f94f..b34d7a8914 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -3253,6 +3253,38 @@ def _combined_extractor_tree_contains_union(model: CombinedExtractorModel) -> bo for sub_extractor in model.extractors ) + def _create_record_counting_extractor( + self, + extractor_model: Optional[ + Union[CustomRecordExtractorModel, DpathExtractorModel, CombinedExtractorModel] + ], + config: Config, + decoder: Optional[Decoder], + ) -> Optional[RecordExtractor]: + """Build the copy of the extractor an `OffsetIncrement` or `PageIncrement` counts a page with. + + The paginator stops when the count falls below the page size, so the count has to be the + number of records the API returned. A `CombinedExtractor` with `skip_empty_records` drops + empty records before they are counted, which would make a full page holding a null look + short and silently end pagination, so every `CombinedExtractor` in the copy is told to count + the records it drops. + """ + if not extractor_model: + return None + extractor: RecordExtractor = self._create_component_from_model( + model=extractor_model, config=config, decoder=decoder + ) + self._count_dropped_empty_records(extractor) + return extractor + + @staticmethod + def _count_dropped_empty_records(extractor: RecordExtractor) -> None: + if not isinstance(extractor, CombinedExtractor): + return + extractor.count_dropped_empty_records = True + for sub_extractor in extractor.extractors: + ModelToComponentFactory._count_dropped_empty_records(sub_extractor) + def create_offset_increment( self, model: OffsetIncrementModel, @@ -3283,13 +3315,7 @@ def create_offset_increment( # decoder with various decorators here, but not in create_record_selector, it is simpler to retain existing # behavior by having two separate extractors with identical behavior since they use the same extractor model. # When we have more time to investigate we can look into reusing the same component. - extractor = ( - self._create_component_from_model( - model=extractor_model, config=config, decoder=decoder_to_use - ) - if extractor_model - else None - ) + extractor = self._create_record_counting_extractor(extractor_model, config, decoder_to_use) # Pydantic v1 Union type coercion can convert int to string depending on Union order. # If page_size is a string that represents an integer (not an interpolation), convert it back. @@ -3319,11 +3345,7 @@ def create_page_increment( # Like OffsetIncrement, we instantiate a separate extractor with identical behavior to the # RecordSelector's so the strategy can count the raw records in the response. This ensures # pagination is driven by the API's page size, not the post-filter record count. - extractor = ( - self._create_component_from_model(model=extractor_model, config=config, decoder=decoder) - if extractor_model - else None - ) + extractor = self._create_record_counting_extractor(extractor_model, config, decoder) # Pydantic v1 Union type coercion can convert int to string depending on Union order. # If page_size is a string that represents an integer (not an interpolation), convert it back. @@ -4492,6 +4514,16 @@ def _get_download_retriever( decoder=_decoder, config=config, url_base="", + # Only a `CombinedExtractor` is handed over, so the `union` rejection and the + # count of dropped empty records apply to downloads too. Other extractors keep + # the retriever's own count: a paginator's copy would re-read the response, + # which a streaming download decoder has already consumed, whereas a + # `CombinedExtractor` is only ever built over a buffered decoder. + extractor_model=( + model.download_extractor + if isinstance(model.download_extractor, CombinedExtractorModel) + else None + ), ) if model.download_paginator else NoPagination(parameters={}) diff --git a/unit_tests/sources/declarative/extractors/test_combined_extractor.py b/unit_tests/sources/declarative/extractors/test_combined_extractor.py index 8eb0225e0a..0048111896 100644 --- a/unit_tests/sources/declarative/extractors/test_combined_extractor.py +++ b/unit_tests/sources/declarative/extractors/test_combined_extractor.py @@ -442,3 +442,137 @@ def test_skip_empty_records_does_not_warn_when_nothing_is_dropped(caplog): ] assert [record.message for record in caplog.records if record.levelno == logging.WARNING] == [] + + +def _counting_copy(extractor: CombinedExtractor) -> CombinedExtractor: + """What the factory hands a record-counting paginator: the same tree, told to count its drops.""" + extractor.count_dropped_empty_records = True + for sub_extractor in extractor.extractors: + if isinstance(sub_extractor, CombinedExtractor): + _counting_copy(sub_extractor) + return extractor + + +def test_the_counting_copy_counts_dropped_empty_records_under_union(): + extractor = _counting_copy( + CombinedExtractor( + extractors=[ + _CountingExtractor(records=[{"id": 1}, None, {}]), + _CountingExtractor(records=[{"id": 2}]), + ], + mode=CombineMode.union, + skip_empty_records=True, + parameters=parameters, + ) + ) + + assert len(list(extractor.extract_records(create_response(GRAPHQL_BODY)))) == 4 + + +def test_the_counting_copy_counts_the_same_winner_the_record_stream_reads(): + """A page of nulls must not win the count while the record stream falls through past it. + + The winner is chosen by the surviving records, so the count is the fallback path's, dropped + records included, and not the two nulls of the primary path. + """ + body = { + "data": { + "boards": [{"items_page": {"items": [None, None]}}], + "next_items_page": {"items": [None, {"id": "item_1"}, {"id": "item_2"}]}, + } + } + + assert list( + _monday_extractor(skip_empty_records=True).extract_records(create_response(body)) + ) == [{"id": "item_1"}, {"id": "item_2"}] + assert ( + len( + list( + _counting_copy(_monday_extractor(skip_empty_records=True)).extract_records( + create_response(body) + ) + ) + ) + == 3 + ) + + +def test_the_counting_copy_counts_nothing_when_every_record_is_empty(): + extractor = _counting_copy( + CombinedExtractor( + extractors=[_CountingExtractor(records=[None]), _CountingExtractor(records=[{}])], + mode=CombineMode.first_match, + skip_empty_records=True, + parameters=parameters, + ) + ) + + assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [] + + +def test_the_counting_copy_counts_the_drops_of_a_nested_combined_extractor(): + extractor = _counting_copy( + CombinedExtractor( + extractors=[ + CombinedExtractor( + extractors=[_CountingExtractor(records=[None, None])], + skip_empty_records=True, + parameters=parameters, + ), + CombinedExtractor( + extractors=[_CountingExtractor(records=[{"id": 1}, None])], + skip_empty_records=True, + parameters=parameters, + ), + ], + mode=CombineMode.first_match, + parameters=parameters, + ) + ) + + # The first nested extractor yields no surviving record, so it does not win even though the + # outer extractor does not skip empty records itself. + assert len(list(extractor.extract_records(create_response(GRAPHQL_BODY)))) == 2 + + +def test_the_counting_copy_does_not_warn(caplog): + """The record stream already warned about the page the paginator re-reads.""" + extractor = _counting_copy( + CombinedExtractor( + extractors=[_CountingExtractor(records=[{"id": 1}, None])], + skip_empty_records=True, + parameters=parameters, + ) + ) + + with caplog.at_level(logging.WARNING, logger="airbyte"): + list(extractor.extract_records(create_response(GRAPHQL_BODY))) + + assert not [record for record in caplog.records if record.levelno == logging.WARNING] + + +def test_the_documented_first_match_example_yields_one_record_per_item(): + """The schema and docstring example; without the trailing `*` each board's list was one record.""" + extractor = CombinedExtractor( + extractors=[ + dpath("data", "boards", "*", "items_page", "items", "*"), + dpath("data", "next_items_page", "items"), + ], + mode=CombineMode.first_match, + skip_empty_records=True, + parameters=parameters, + ) + body = { + "data": { + "boards": [ + {"items_page": {"items": [{"id": 1}, {"id": 2}]}}, + {"items_page": {"items": [{"id": 3}]}}, + ], + } + } + + assert list(extractor.extract_records(create_response(body))) == [ + {"id": 1}, + {"id": 2}, + {"id": 3}, + ] 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 b7279559c3..b4bc4d7d13 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 @@ -10,7 +10,7 @@ # mypy: ignore-errors from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Any, Iterable, Mapping, Optional, Union +from typing import Any, Dict, Iterable, Mapping, Optional, Union import freezegun import pytest @@ -8374,6 +8374,187 @@ def test_union_combined_extractor_is_accepted_by_a_page_increment_paginator(): assert isinstance(pagination_strategy.extractor, CombinedExtractor) +_FIRST_MATCH_SKIPPING_EMPTY_RECORDS = """ type: CombinedExtractor + mode: first_match + skip_empty_records: true + extractors: + - type: DpathExtractor + field_path: ["a"] + - type: DpathExtractor + field_path: ["b"]""" + + +@pytest.mark.parametrize( + "pagination_strategy", + [ + pytest.param(_OFFSET_INCREMENT, id="offset_increment"), + pytest.param(_PAGE_INCREMENT, id="page_increment"), + ], +) +def test_a_record_counting_paginator_counts_the_empty_records_the_extractor_drops( + pagination_strategy: str, +): + """A null on a full page must not make the page look short and silently end pagination.""" + retriever = _build_retriever( + _retriever_manifest_with_paginator(_FIRST_MATCH_SKIPPING_EMPTY_RECORDS, pagination_strategy) + ) + + record_selector_extractor = retriever.record_selector.extractor + counting_extractor = retriever.paginator.pagination_strategy.extractor + assert isinstance(record_selector_extractor, CombinedExtractor) + assert isinstance(counting_extractor, CombinedExtractor) + assert record_selector_extractor.count_dropped_empty_records is False + assert counting_extractor.count_dropped_empty_records is True + + full_page_with_a_null = _json_response({"a": [None, {"id": 1}]}) + assert list(record_selector_extractor.extract_records(full_page_with_a_null)) == [{"id": 1}] + # `last_page_size` is what the retriever emitted; the strategy must count the page itself. + assert ( + retriever.paginator.pagination_strategy.next_page_token( + response=full_page_with_a_null, + last_page_size=1, + last_record={"id": 1}, + last_page_token_value=0, + ) + is not None + ) + + +def _async_retriever_definition( + download_extractor: Mapping[str, Any], download_paginator: Mapping[str, Any] +) -> Dict[str, Any]: + return { + "type": "AsyncRetriever", + "status_mapping": { + "failed": ["failed"], + "running": ["pending"], + "timeout": ["timeout"], + "completed": ["ready"], + }, + "status_extractor": {"type": "DpathExtractor", "field_path": ["status"]}, + "download_target_extractor": {"type": "DpathExtractor", "field_path": ["url"]}, + "download_extractor": download_extractor, + "download_paginator": download_paginator, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "polling_requester": { + "type": "HttpRequester", + "path": "/exports/{{ creation_response['id'] }}", + "url_base": "https://api.test.com", + "http_method": "GET", + }, + "creation_requester": { + "type": "HttpRequester", + "path": "/exports", + "url_base": "https://api.test.com", + "http_method": "POST", + }, + "download_requester": { + "type": "HttpRequester", + "path": "{{download_target}}", + "url_base": "", + "http_method": "GET", + }, + } + + +def _download_paginator(pagination_strategy_type: str) -> Dict[str, Any]: + return { + "type": "DefaultPaginator", + "pagination_strategy": {"type": pagination_strategy_type, "page_size": 2}, + "page_token_option": { + "type": "RequestOption", + "inject_into": "request_parameter", + "field_name": "offset", + }, + } + + +def _build_download_paginator(definition: Mapping[str, Any]) -> Any: + component = factory.create_component( + model_type=AsyncRetrieverModel, + component_definition=definition, + name="test_stream", + primary_key="id", + stream_slicer=None, + transformations=[], + config={}, + ) + job_repository = component.stream_slicer.job_orchestrator_factory( + [StreamSlice(partition={}, cursor_slice={})] + )._job_repository + return job_repository.download_retriever.paginator + + +def test_union_combined_download_extractor_is_rejected_by_an_offset_increment_download_paginator(): + """The download paginator is built apart from the record selector's, so it needs its own check. + + Without it the strategy falls back to the retriever's record count, which under `union` is the + sum over all sub-extractors, and the offsets skip the records in between. + """ + definition = _async_retriever_definition( + { + "type": "CombinedExtractor", + "mode": "union", + "extractors": [ + {"type": "DpathExtractor", "field_path": ["a"]}, + {"type": "DpathExtractor", "field_path": ["b"]}, + ], + }, + _download_paginator("OffsetIncrement"), + ) + + with pytest.raises(AirbyteTracedException) as exc_info: + _build_download_paginator(definition) + + assert exc_info.value.failure_type == FailureType.config_error + assert "union" in exc_info.value.message + assert "OffsetIncrement" in exc_info.value.message + + +def test_combined_download_extractor_skipping_empty_records_counts_them_for_the_download_paginator(): + definition = _async_retriever_definition( + { + "type": "CombinedExtractor", + "mode": "first_match", + "skip_empty_records": True, + "extractors": [ + {"type": "DpathExtractor", "field_path": ["a"]}, + {"type": "DpathExtractor", "field_path": ["b"]}, + ], + }, + _download_paginator("PageIncrement"), + ) + + pagination_strategy = _build_download_paginator(definition).pagination_strategy + assert isinstance(pagination_strategy, PageIncrement) + assert isinstance(pagination_strategy.extractor, CombinedExtractor) + assert pagination_strategy.extractor.count_dropped_empty_records is True + assert ( + pagination_strategy.next_page_token( + response=_json_response({"a": [None, {"id": 1}]}), + last_page_size=1, + last_record={"id": 1}, + last_page_token_value=0, + ) + is not None + ) + + +def test_download_paginator_keeps_the_retriever_count_for_other_download_extractors(): + """A paginator's copy would re-read a response a streaming download decoder already consumed.""" + definition = _async_retriever_definition( + {"type": "DpathExtractor", "field_path": ["a"]}, + _download_paginator("OffsetIncrement"), + ) + + pagination_strategy = _build_download_paginator(definition).pagination_strategy + assert isinstance(pagination_strategy, OffsetIncrement) + assert pagination_strategy.extractor is None + + def test_create_async_retriever_with_combined_extractors(): """Covers the three `AsyncRetriever` unions widened by `CombinedExtractor`.""" combined = { diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index 3b72d18ee8..9d4e438940 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -6395,3 +6395,75 @@ def test_combined_extractor_with_an_empty_extractors_list_fails_schema_validatio "should be non-empty" in message for message in _validation_messages(exc_info.value.__cause__) ) + + +@pytest.mark.parametrize( + "pagination_strategy, page_token_field, page_tokens", + [ + pytest.param( + {"type": "OffsetIncrement", "page_size": 3}, "offset", ["3", "6"], id="offset_increment" + ), + pytest.param( + {"type": "PageIncrement", "page_size": 3, "start_from_page": 0}, + "page", + ["1", "2"], + id="page_increment", + ), + ], +) +def test_combined_extractor_skipping_empty_records_reads_every_page( + pagination_strategy, page_token_field, page_tokens +): + """A null on a full page used to make the page look short, so the later pages were never read.""" + manifest = _combined_extractor_manifest( + { + "type": "CombinedExtractor", + "mode": "first_match", + "skip_empty_records": True, + "extractors": [ + {"type": "DpathExtractor", "field_path": ["data", "items"]}, + {"type": "DpathExtractor", "field_path": ["data", "fallback"]}, + ], + } + ) + manifest["streams"][0]["retriever"]["paginator"] = { + "type": "DefaultPaginator", + "pagination_strategy": pagination_strategy, + "page_token_option": { + "type": "RequestOption", + "inject_into": "request_parameter", + "field_name": page_token_field, + }, + } + catalog = create_catalog("lists") + source = ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=catalog, state=None + ) + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest("https://api.test.com/lists"), + HttpResponse(json.dumps({"data": {"items": [{"id": 1}, None, {"id": 3}]}})), + ) + http_mocker.get( + HttpRequest( + "https://api.test.com/lists", query_params={page_token_field: page_tokens[0]} + ), + HttpResponse(json.dumps({"data": {"items": [{"id": 4}, {"id": 5}, {"id": 6}]}})), + ) + http_mocker.get( + HttpRequest( + "https://api.test.com/lists", query_params={page_token_field: page_tokens[1]} + ), + HttpResponse(json.dumps({"data": {"items": [{"id": 7}]}})), + ) + messages = list(source.read(logger=source.logger, config={}, catalog=catalog, state=[])) + + assert [record.data["id"] for record in get_records_for_stream("lists", messages)] == [ + 1, + 3, + 4, + 5, + 6, + 7, + ]