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 a02be6fa0e..8e0ec9b9bf 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -1911,6 +1911,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: @@ -2183,6 +2184,66 @@ 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.' + type: string + enum: + - union + - first_match + examples: + - union + - first_match + 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 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" + - "$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. 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. 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: + 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 + skip_empty_records: true + extractors: + - type: DpathExtractor + field_path: ["data", "boards", "*", "items_page", "items", "*"] + - type: DpathExtractor + field_path: ["data", "next_items_page", "items"] DpathExtractor: title: Dpath Extractor description: Record extractor that searches a decoded response over a path defined as an array of fields. @@ -3978,6 +4039,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. @@ -4540,17 +4602,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..bb57c421c3 --- /dev/null +++ b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py @@ -0,0 +1,248 @@ +# +# 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, Iterable, Iterator, List, Mapping, Union + +import requests + +from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor + +logger = logging.getLogger("airbyte") + + +class CombineMode(Enum): + """ + How a `CombinedExtractor` combines the output of its sub-extractors. + """ + + union = "union" + first_match = "first_match" + + +class _NoRecord: + """Sentinel telling "this extractor yielded nothing" apart from a falsy or `None` record.""" + + +_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. + + 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 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 + 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. + + 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"] + ``` + + ## 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 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. + + ## 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. + + 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 + 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. + + ## 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, 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. + + 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: + 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) + else: + yield from self._extract_union(response) + + def _extract_union(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: + 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 index in range(len(self.extractors)): + records: Iterator[Mapping[str, Any]] = iter(self._records_of(index, response)) + # 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(dropped_placeholders, [first_record], records) + return + + 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: + 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) " + "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__, + ) diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index a62c3db2fe..cd97715a2d 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -508,6 +508,11 @@ class Config: ) +class CombineMode(Enum): + union = "union" + first_match = "first_match" + + class ResponseToFileExtractor(BaseModel): type: Literal["ResponseToFileExtractor"] preserve_na_values: Optional[bool] = Field( @@ -2589,6 +2594,26 @@ class DpathExtractor(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +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 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. 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. 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") + + class ZipfileDecoder(BaseModel): class Config: extra = Extra.allow @@ -2603,7 +2628,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.", @@ -2848,9 +2873,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, @@ -3336,15 +3363,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( ..., @@ -3554,6 +3583,7 @@ class DynamicDeclarativeStream(BaseModel): ComplexFieldType.update_forward_refs() GzipDecoder.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 8940f33fa3..b34d7a8914 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -113,6 +113,8 @@ RecordExpander, ) from airbyte_cdk.sources.declarative.extractors import ( + CombinedExtractor, + CombineMode, DpathExtractor, RecordFilter, RecordSelector, @@ -172,6 +174,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, ) @@ -702,6 +707,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] @@ -799,6 +805,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, @@ -2446,7 +2453,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]: @@ -2490,6 +2499,81 @@ 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. + + 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 ( + 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.", + internal_message=( + f"CombinedExtractor was configured with {type(inner_decoder).__name__}, which " + 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, + ) + + def create_combined_extractor( + self, + model: CombinedExtractorModel, + config: Config, + 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 + ] + 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 {}, + ) + def create_dpath_extractor( self, model: DpathExtractorModel, @@ -3123,12 +3207,92 @@ 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` 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 + 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` 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_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, 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,18 +3308,14 @@ 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 # 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. @@ -3177,17 +3337,15 @@ 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 # 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. @@ -4356,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/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/.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/extractors/test_combined_extractor.py b/unit_tests/sources/declarative/extractors/test_combined_extractor.py new file mode 100644 index 0000000000..0048111896 --- /dev/null +++ b/unit_tests/sources/declarative/extractors/test_combined_extractor.py @@ -0,0 +1,578 @@ +# +# Copyright (c) 2025 Airbyte, Inc., all rights reserved. +# +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 +from airbyte_cdk.sources.declarative.extractors import ( + CombinedExtractor, + CombineMode, + DpathExtractor, +) +from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor + +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 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"}], +} + + +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.status_code = 200 + response.raw = urllib3.HTTPResponse(body=io.BytesIO(raw), preload_content=False, status=200) + 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(RecordExtractor): + """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))) == [] + + +@pytest.mark.parametrize( + "mode", + [CombineMode.union, CombineMode.first_match], +) +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 _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. + + 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 = _streaming_decoder() + extractor = CombinedExtractor( + extractors=[ + dpath("data", decoder=streaming_decoder), + dpath("data", decoder=streaming_decoder), + ], + mode=CombineMode.union, + parameters=parameters, + ) + + 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 _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 _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 ea3b0cf151..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 @@ -54,7 +54,17 @@ 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.decoders.composite_raw_decoder import ( + CompositeRawDecoder, + JsonParser, +) +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, @@ -112,9 +122,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, ) @@ -124,6 +149,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, @@ -7816,3 +7844,805 @@ 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}] + + +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: + 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"], {}) + + +class _BufferedCompositeRawDecoderSubclass(CompositeRawDecoder): + """Stands in for a `CustomDecoder` that subclasses `CompositeRawDecoder` and buffers.""" + + +_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) + + +@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 = """ + 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 _retriever_manifest_with_paginator( + extractor_definition: str, pagination_strategy: str +) -> Mapping[str, Any]: + content = f""" + retriever: + type: SimpleRetriever + requester: + type: HttpRequester + url_base: "https://api.test.com" + path: "/records" + http_method: GET + paginator: + type: DefaultPaginator + pagination_strategy: +{pagination_strategy} + record_selector: + type: RecordSelector + extractor: +{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"]""" +_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"]""" + + +def _build_retriever(retriever_manifest: Mapping[str, Any]) -> SimpleRetriever: + return factory.create_component( + model_type=SimpleRetrieverModel, + component_definition=retriever_manifest, + config=input_config, + name="test_stream", + primary_key="id", + stream_slicer=None, + 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) + # 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 + + # `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_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) + + +_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 = { + "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 c8aaca355b..9d4e438940 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -6299,3 +6299,171 @@ 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": "first_match", + "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__) + ) + + +@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, + ]