From 788f1db15a28c8913372c06e6fa712d5a22a85db Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Tue, 15 Sep 2026 19:14:25 +0300 Subject: [PATCH 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 e78f1efbdad4a21cd5ac62e7b2ba506d51fb597e Mon Sep 17 00:00:00 2001 From: Anatolii Yatsuk Date: Mon, 28 Sep 2026 18:45:10 +0300 Subject: [PATCH 6/6] 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, + ]