diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index aa3bbbc5e0..071f5eeaf4 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2258,6 +2258,22 @@ definitions: - skip - emit_parent default: skip + truncation_indicator_path: + title: Truncation Indicator Path + description: Path within each record to a field indicating that the embedded nested list is truncated (e.g. a `has_more` flag on the list object). When the field evaluates to a truthy value and `truncated_list_retriever` is configured, the retriever is used to fetch the complete list instead of expanding the embedded items. When the field is truthy and no retriever is configured, the embedded items are expanded as normal and a WARNING is logged once per stream so the truncation is visible instead of silent. Glob characters (`*`, `?`, `[`) are not supported in this path, nor in `expand_records_from_field` when a retriever is configured; this is enforced on the interpolated values. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it. + type: array + items: + type: string + interpolation_context: + - config + examples: + - ["data", "object", "lines", "has_more"] + truncated_list_retriever: + title: Truncated List Retriever + description: Retriever used to fetch the complete list of items when the field at `truncation_indicator_path` is truthy on a record. The record being expanded is exposed to the retriever's interpolation context as `stream_slice['parent_record']`. One fetch is issued per truncated record, so enable `use_cache` on the requester when the same list can be fetched repeatedly. Configure a `paginator`, since without one only the first page of the complete list is read. If the retriever returns no records, the embedded items are expanded as a fallback; if it returns fewer records than the `total_count` field next to the indicator, a WARNING is logged once per stream. Request failures surface through the retriever's `error_handler` and fail the stream like any other request. `$parameters` of the enclosing stream propagate into this retriever's components (including `request_parameters` on its requester); move request-shaping parameters to the outer requester's `request_parameters` when adopting this field. `partition_router` and `pagination_reset` are not supported. In Connector Builder test reads, its requests appear as auxiliary requests and the test-read page limit applies to each fetch independently, so the fetched list may be shorter than `total_count`; the incomplete-fetch warning is not emitted in test reads when a `paginator` is configured (without one the retriever is not capped, so the warning still applies). Requires `truncation_indicator_path`. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it. + anyOf: + - "$ref": "#/definitions/SimpleRetriever" + - "$ref": "#/definitions/CustomRetriever" $parameters: type: object additionalProperties: true diff --git a/airbyte_cdk/sources/declarative/expanders/record_expander.py b/airbyte_cdk/sources/declarative/expanders/record_expander.py index 6b489243d1..68c71af693 100644 --- a/airbyte_cdk/sources/declarative/expanders/record_expander.py +++ b/airbyte_cdk/sources/declarative/expanders/record_expander.py @@ -3,14 +3,42 @@ # import copy +import logging +import threading from dataclasses import InitVar, dataclass from enum import Enum -from typing import Any, Iterable, Mapping, MutableMapping, Sequence +from typing import TYPE_CHECKING, Any, Iterable, Mapping, MutableMapping, Optional, Sequence import dpath +from airbyte_cdk.models import ( + AirbyteControlMessage, + AirbyteLogMessage, + AirbyteMessage, + AirbyteStateMessage, + AirbyteTraceMessage, + Level, +) +from airbyte_cdk.models import Type as MessageType from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString -from airbyte_cdk.sources.types import Config +from airbyte_cdk.sources.message import MessageRepository +from airbyte_cdk.sources.types import Config, Record, StreamSlice + +if TYPE_CHECKING: + from airbyte_cdk.sources.declarative.retrievers import Retriever + +logger = logging.getLogger("airbyte") + +# dpath treats these characters as glob metacharacters (fnmatch semantics) inside a path segment. +_GLOB_METACHARACTERS = ("*", "?", "[") + +# Protocol payloads a retriever may yield unwrapped (i.e. not inside an `AirbyteMessage` envelope). +_BARE_PROTOCOL_MESSAGES = ( + AirbyteControlMessage, + AirbyteLogMessage, + AirbyteStateMessage, + AirbyteTraceMessage, +) class OnNoRecords(Enum): @@ -62,6 +90,34 @@ class RecordExpander: parent record in an "original_record" field. Defaults to False. on_no_records: Behavior when expansion produces no records. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged. + truncation_indicator_path: Path within each record to a field indicating that the + embedded nested list is truncated (e.g. a `has_more` flag on the list object). + When the indicator is truthy and no `truncated_list_retriever` is configured, the + embedded items are expanded as normal and a WARNING is logged (once per stream + instance) describing the expansion path and the embedded item count, so that the + data loss is visible instead of silent. Glob metacharacters (`*`, `?`, `[`) are + rejected in this path, and in `expand_records_from_field` when a retriever is + configured; the check runs on the interpolated values. + truncated_list_retriever: Retriever used to fetch the complete list of items when + the field at `truncation_indicator_path` is truthy. The record being expanded is + exposed to the retriever's interpolation context as `stream_slice['parent_record']`. + One fetch is issued per truncated parent record; enable `use_cache` on its requester + when the same list can be fetched repeatedly. Without a `paginator` only the first + page is read. If the retriever returns no records, the embedded items are expanded + as a fallback; if it returns fewer records than the `total_count` field next to the + indicator, a WARNING is logged once per stream instance. Request failures surface + through the retriever's error handler and fail the stream like any other request. + `$parameters` of the enclosing stream propagate into this retriever's components. + In Connector Builder test reads the page limit applies to each fetch independently, + so the fetched list may be shorter than `total_count`; no incomplete-fetch warning + is emitted there. + message_repository: Optional repository through which the truncation warnings are emitted + as Airbyte LOG messages so they are visible in the Connector Builder. When it is not + set, the warnings go to the `airbyte` logger instead. + suppress_incomplete_fetch_warning: Skip the incomplete-fetch WARNING. Set by the factory + for Connector Builder test reads when the page limit caps the retriever's paginator, + so a shortfall against `total_count` is expected. The truncated-without-retriever + warning is not affected. config: The user-provided configuration as specified by the source's spec. """ @@ -70,12 +126,48 @@ class RecordExpander: parameters: InitVar[Mapping[str, Any]] remain_original_record: bool = False on_no_records: OnNoRecords = OnNoRecords.skip + truncation_indicator_path: Optional[Sequence[str]] = None + truncated_list_retriever: Optional["Retriever"] = None + message_repository: Optional[MessageRepository] = None + suppress_incomplete_fetch_warning: bool = False def __post_init__(self, parameters: Mapping[str, Any]) -> None: self._expand_path: list[InterpolatedString] = [ InterpolatedString.create(path, parameters=parameters) for path in self.expand_records_from_field ] + if self.truncated_list_retriever and not self.truncation_indicator_path: + raise ValueError( + "`truncation_indicator_path` is required when `truncated_list_retriever` is configured." + ) + self._truncation_indicator_path: list[InterpolatedString] = [ + InterpolatedString.create(path, parameters=parameters) + for path in (self.truncation_indicator_path or []) + ] + # The paths only interpolate `config`, so their evaluated values are validated up front. + self._reject_globs(self._evaluated_indicator_path(), "truncation_indicator_path") + if self.truncated_list_retriever: + self._reject_globs(self._evaluated_expand_path(), "expand_records_from_field") + self._warning_lock = threading.Lock() + self._warned_truncation_without_retriever = False + self._warned_incomplete_fetch = False + + @staticmethod + def _reject_globs(path: Sequence[Any], field_name: str) -> None: + if any( + isinstance(segment, str) and any(char in segment for char in _GLOB_METACHARACTERS) + for segment in path + ): + raise ValueError( + f"Glob characters {_GLOB_METACHARACTERS} are not supported in `{field_name}` when " + "truncation handling is configured: the path must identify a single field." + ) + + def _evaluated_indicator_path(self) -> list[Any]: + return [segment.eval(self.config) for segment in self._truncation_indicator_path] + + def _evaluated_expand_path(self) -> list[Any]: + return [segment.eval(self.config) for segment in self._expand_path] def expand_record(self, record: MutableMapping[Any, Any]) -> Iterable[MutableMapping[Any, Any]]: """Expand a record by extracting items from a nested array field.""" @@ -89,37 +181,153 @@ def expand_record(self, record: MutableMapping[Any, Any]) -> Iterable[MutableMap return parent_record = record - expand_path = [path.eval(self.config) for path in self._expand_path] - expanded_any = False + + expand_path = self._evaluated_expand_path() + truncated = bool(self._truncation_indicator_path) and self._is_truncated(parent_record) + + if truncated and self.truncated_list_retriever: + # Streamed, so the shortfall is only known once the retriever is exhausted. If the + # consumer stops early the fetch was cut short by it, not by the API, and no warning + # would be accurate anyway. + fetched_count = 0 + for fetched in self._fetch_complete_list(parent_record): + fetched_count += 1 + yield fetched + self._warn_if_fetch_incomplete(parent_record, expand_path, fetched_count) + if fetched_count > 0: + return try: extracted_values = dpath.values(parent_record, expand_path) except KeyError: extracted_values = [] - for extracted in extracted_values: - if not isinstance(extracted, list): - continue - items = extracted + embedded_lists = [ + extracted for extracted in extracted_values if isinstance(extracted, list) + ] + embedded_count = sum(len(items) for items in embedded_lists) + + if truncated and not self.truncated_list_retriever: + self._warn_truncated_without_retriever(parent_record, expand_path, embedded_count) + + for items in embedded_lists: for item in items: if isinstance(item, dict): expanded_record = dict(item) self._apply_parent_context(parent_record, expanded_record) yield expanded_record - expanded_any = True + elif self.remain_original_record: + yield { + "value": item, + "original_record": copy.deepcopy(parent_record), + } else: - if self.remain_original_record: - yield { - "value": item, - "original_record": copy.deepcopy(parent_record), - } - else: - yield item - expanded_any = True - - if not expanded_any and self.on_no_records == OnNoRecords.emit_parent: + yield item + + if embedded_count == 0 and self.on_no_records == OnNoRecords.emit_parent: yield parent_record + def _warn_truncated_without_retriever( + self, parent_record: Mapping[str, Any], expand_path: list[Any], embedded_count: int + ) -> None: + with self._warning_lock: + if self._warned_truncation_without_retriever: + return + self._warned_truncation_without_retriever = True + + indicator_path = self._evaluated_indicator_path() + total_count = self._get_sibling_total_count(parent_record, indicator_path) + total_fragment = f" of {total_count} total" if total_count is not None else "" + self._emit_warning( + f"The nested list at {expand_path} is marked as truncated (the field at {indicator_path} " + f"is truthy) but no `truncated_list_retriever` is configured, so only the " + f"{embedded_count} embedded item(s){total_fragment} were expanded and the remaining " + "items are not emitted. Configure `truncated_list_retriever` to fetch the complete list " + "if the API provides an endpoint for it. This warning is emitted once per stream; other " + "records may be truncated as well." + ) + + def _warn_if_fetch_incomplete( + self, parent_record: Mapping[str, Any], expand_path: list[Any], fetched_count: int + ) -> None: + if self.suppress_incomplete_fetch_warning: + return + indicator_path = self._evaluated_indicator_path() + total_count = self._get_sibling_total_count(parent_record, indicator_path) + if total_count is None or fetched_count >= total_count: + return + with self._warning_lock: + if self._warned_incomplete_fetch: + return + self._warned_incomplete_fetch = True + fallback_fragment = ( + " The embedded items were expanded as a fallback." if fetched_count == 0 else "" + ) + self._emit_warning( + f"The `truncated_list_retriever` for the nested list at {expand_path} returned " + f"{fetched_count} record(s) but the `total_count` field next to {indicator_path} reports " + f"{total_count}, so the fetched list is still incomplete.{fallback_fragment} Check that " + "the retriever has a `paginator` configured and that its request matches the " + "complete-list endpoint. This warning is emitted once per stream; other records may be " + "affected as well." + ) + + def _emit_warning(self, message: str) -> None: + if self.message_repository: + self.message_repository.emit_message( + AirbyteMessage( + type=MessageType.LOG, + log=AirbyteLogMessage(level=Level.WARN, message=message), + ) + ) + else: + logger.warning(message) + + def _get_sibling_total_count( + self, parent_record: Mapping[str, Any], indicator_path: list[Any] + ) -> Optional[int]: + """Best-effort lookup of a `total_count` field next to the truncation indicator.""" + if not indicator_path: + return None + try: + total = dpath.get(dict(parent_record), [*indicator_path[:-1], "total_count"]) + except (KeyError, ValueError): + return None + return total if isinstance(total, int) and not isinstance(total, bool) else None + + def _is_truncated(self, parent_record: MutableMapping[Any, Any]) -> bool: + indicator_path = self._evaluated_indicator_path() + try: + return bool(dpath.get(parent_record, indicator_path)) + except (KeyError, ValueError): + return False + + def _fetch_complete_list(self, parent_record: Mapping[str, Any]) -> Iterable[Any]: + if not self.truncated_list_retriever: + return + stream_slice = StreamSlice(partition={"parent_record": parent_record}, cursor_slice={}) + for item in self.truncated_list_retriever.read_records( + records_schema={}, stream_slice=stream_slice + ): + if isinstance(item, AirbyteMessage): + if item.type != MessageType.RECORD or item.record is None: + continue + data: Any = item.record.data + elif isinstance(item, Record): + data = item.data + elif isinstance(item, _BARE_PROTOCOL_MESSAGES): + continue + else: + data = item + if isinstance(data, Mapping): + expanded_record = dict(data) + self._apply_parent_context(parent_record, expanded_record) + yield expanded_record + elif self.remain_original_record: + yield {"value": data, "original_record": copy.deepcopy(parent_record)} + else: + yield data + def _apply_parent_context( self, parent_record: Mapping[str, Any], child_record: MutableMapping[str, Any] ) -> None: diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 21a92cc3be..41265c9db9 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2305,6 +2305,17 @@ class RecordExpander(BaseModel): description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', title="On No Records", ) + truncation_indicator_path: Optional[List[str]] = Field( + None, + description="Path within each record to a field indicating that the embedded nested list is truncated (e.g. a `has_more` flag on the list object). When the field evaluates to a truthy value and `truncated_list_retriever` is configured, the retriever is used to fetch the complete list instead of expanding the embedded items. When the field is truthy and no retriever is configured, the embedded items are expanded as normal and a WARNING is logged once per stream so the truncation is visible instead of silent. Glob characters (`*`, `?`, `[`) are not supported in this path, nor in `expand_records_from_field` when a retriever is configured; this is enforced on the interpolated values. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.", + examples=[["data", "object", "lines", "has_more"]], + title="Truncation Indicator Path", + ) + truncated_list_retriever: Optional[Union[SimpleRetriever, CustomRetriever]] = Field( + None, + description="Retriever used to fetch the complete list of items when the field at `truncation_indicator_path` is truthy on a record. The record being expanded is exposed to the retriever's interpolation context as `stream_slice['parent_record']`. One fetch is issued per truncated record, so enable `use_cache` on the requester when the same list can be fetched repeatedly. Configure a `paginator`, since without one only the first page of the complete list is read. If the retriever returns no records, the embedded items are expanded as a fallback; if it returns fewer records than the `total_count` field next to the indicator, a WARNING is logged once per stream. Request failures surface through the retriever's `error_handler` and fail the stream like any other request. `$parameters` of the enclosing stream propagate into this retriever's components (including `request_parameters` on its requester); move request-shaping parameters to the outer requester's `request_parameters` when adopting this field. `partition_router` and `pagination_reset` are not supported. In Connector Builder test reads, its requests appear as auxiliary requests and the test-read page limit applies to each fetch independently, so the fetched list may be shorter than `total_count`; the incomplete-fetch warning is not emitted in test reads when a `paginator` is configured (without one the retriever is not capped, so the warning still applies). Requires `truncation_indicator_path`. This field is ignored by CDK versions that predate it, so pin the connector to a CDK version that supports it.", + title="Truncated List Retriever", + ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -3485,6 +3496,7 @@ class DynamicDeclarativeStream(BaseModel): FileUploader.update_forward_refs() DeclarativeStream.update_forward_refs() SessionTokenAuthenticator.update_forward_refs() +RecordExpander.update_forward_refs() HttpRequester.update_forward_refs() DynamicSchemaLoader.update_forward_refs() ParentStreamConfig.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..bec394f85e 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2506,6 +2506,40 @@ def create_record_expander( config: Config, **kwargs: Any, ) -> RecordExpander: + truncated_list_retriever = None + suppress_incomplete_fetch_warning = False + if model.truncated_list_retriever: + retriever_model = model.truncated_list_retriever + name = "record_expander_truncated_list" + # `CustomRetriever` allows extra fields, so read from the dumped model to cover both types. + retriever_fields = retriever_model.dict() + for unsupported_option in ("partition_router", "pagination_reset"): + if retriever_fields.get(unsupported_option): + raise ValueError( + f"`{unsupported_option}` is not supported on `truncated_list_retriever`." + ) + log_formatter = lambda response: format_http_message( + response, + f"Record expander '{name}' request", + "Request performed in order to fetch the complete nested list of a truncated record.", + name, + is_auxiliary=True, + ) + # `name`/`primary_key`/`transformations` are also what a `CustomRetriever` forwards to its + # nested `requester`/`record_selector`, so they are passed for both retriever types. + truncated_list_retriever = self._create_component_from_model( + model=retriever_model, + config=config, + name=name, + primary_key=None, + transformations=[], + log_formatter=log_formatter, + ) + # Only a capped paginator makes a shortfall expected; `NoPagination` and retrievers + # without a paginator are never capped. + suppress_incomplete_fetch_warning = isinstance( + truncated_list_retriever, SimpleRetriever + ) and isinstance(truncated_list_retriever.paginator, PaginatorTestReadDecorator) return RecordExpander( expand_records_from_field=model.expand_records_from_field, config=config, @@ -2514,6 +2548,10 @@ def create_record_expander( on_no_records=OnNoRecords(model.on_no_records.value) if model.on_no_records else OnNoRecords.skip, + truncation_indicator_path=model.truncation_indicator_path, + truncated_list_retriever=truncated_list_retriever, + message_repository=self._message_repository, + suppress_incomplete_fetch_warning=suppress_incomplete_fetch_warning, ) @staticmethod diff --git a/unit_tests/sources/declarative/expanders/__init__.py b/unit_tests/sources/declarative/expanders/__init__.py new file mode 100644 index 0000000000..89adeb6e68 --- /dev/null +++ b/unit_tests/sources/declarative/expanders/__init__.py @@ -0,0 +1,3 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# diff --git a/unit_tests/sources/declarative/expanders/test_record_expander.py b/unit_tests/sources/declarative/expanders/test_record_expander.py new file mode 100644 index 0000000000..23731dda9e --- /dev/null +++ b/unit_tests/sources/declarative/expanders/test_record_expander.py @@ -0,0 +1,574 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import threading +from unittest.mock import MagicMock + +import pytest + +from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteRecordMessage, Level, Type +from airbyte_cdk.sources.declarative.expanders.record_expander import RecordExpander +from airbyte_cdk.sources.message import InMemoryMessageRepository +from airbyte_cdk.sources.types import Record, StreamSlice + +config = {} +parameters = {} + + +def _make_retriever(records): + retriever = MagicMock() + retriever.read_records.return_value = iter(records) + return retriever + + +def _event(lines, has_more, total_count): + return { + "id": "evt_1", + "data": { + "object": { + "id": "in_1", + "lines": { + "data": lines, + "has_more": has_more, + "total_count": total_count, + "url": "/v1/invoices/in_1/lines", + }, + } + }, + } + + +def test_truncated_list_is_fetched_via_retriever(): + embedded = [{"id": f"il_{i}"} for i in range(10)] + complete = [{"id": f"il_{i}"} for i in range(15)] + retriever = _make_retriever([Record(data=item, stream_name="test") for item in complete]) + expander = RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + remain_original_record=True, + truncation_indicator_path=["data", "object", "lines", "has_more"], + truncated_list_retriever=retriever, + ) + + parent = _event(embedded, has_more=True, total_count=15) + records = list(expander.expand_record(parent)) + + assert len(records) == 15 + assert [record["id"] for record in records] == [f"il_{i}" for i in range(15)] + assert all(record["original_record"] == parent for record in records) + stream_slice = retriever.read_records.call_args.kwargs["stream_slice"] + assert stream_slice == StreamSlice(partition={"parent_record": parent}, cursor_slice={}) + + +def test_no_retriever_call_when_not_truncated(): + embedded = [{"id": f"il_{i}"} for i in range(10)] + retriever = _make_retriever([]) + expander = RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + truncated_list_retriever=retriever, + ) + + records = list(expander.expand_record(_event(embedded, has_more=False, total_count=10))) + + assert len(records) == 10 + retriever.read_records.assert_not_called() + + +def test_no_retriever_call_when_indicator_missing(): + embedded = [{"id": "il_0"}] + retriever = _make_retriever([]) + expander = RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + truncated_list_retriever=retriever, + ) + + parent = _event(embedded, has_more=False, total_count=1) + del parent["data"]["object"]["lines"]["has_more"] + records = list(expander.expand_record(parent)) + + assert len(records) == 1 + retriever.read_records.assert_not_called() + + +def test_falls_back_to_embedded_items_when_retriever_returns_nothing(): + embedded = [{"id": f"il_{i}"} for i in range(10)] + retriever = _make_retriever([]) + expander = RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + truncated_list_retriever=retriever, + ) + + records = list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + + assert len(records) == 10 + retriever.read_records.assert_called_once() + + +def test_retriever_requires_truncation_indicator_path(): + with pytest.raises(ValueError): + RecordExpander( + expand_records_from_field=["lines", "data"], + config=config, + parameters=parameters, + truncated_list_retriever=_make_retriever([]), + ) + + +def test_wildcard_rejected_with_truncation_handling(): + with pytest.raises(ValueError): + RecordExpander( + expand_records_from_field=["sections", "*", "items"], + config=config, + parameters=parameters, + truncation_indicator_path=["lines", "has_more"], + truncated_list_retriever=_make_retriever([]), + ) + + +@pytest.mark.parametrize("segment", ["*", "has_*", "has_more?", "[hl]ines", "**"]) +def test_glob_metacharacters_rejected_in_truncation_indicator_path(segment): + with pytest.raises(ValueError): + RecordExpander( + expand_records_from_field=["lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=[segment, "has_more"], + ) + + +@pytest.mark.parametrize("segment", ["sect?ons", "[s]ections", "sec*"]) +def test_glob_metacharacters_rejected_in_expand_path_with_retriever(segment): + with pytest.raises(ValueError): + RecordExpander( + expand_records_from_field=[segment, "items"], + config=config, + parameters=parameters, + truncation_indicator_path=["lines", "has_more"], + truncated_list_retriever=_make_retriever([]), + ) + + +def test_glob_metacharacters_allowed_in_expand_path_without_retriever(): + expander = RecordExpander( + expand_records_from_field=["sections", "*", "items"], + config=config, + parameters=parameters, + truncation_indicator_path=["has_more"], + ) + records = list( + expander.expand_record({"sections": [{"items": [{"id": 1}]}], "has_more": False}) + ) + assert records == [{"id": 1}] + + +def test_interpolated_glob_rejected(): + with pytest.raises(ValueError): + RecordExpander( + expand_records_from_field=["lines", "data"], + config={"indicator": "has_*"}, + parameters=parameters, + truncation_indicator_path=["lines", "{{ config['indicator'] }}"], + ) + + +def test_interpolated_indicator_path_is_used(): + expander = RecordExpander( + expand_records_from_field=["lines", "data"], + config={"indicator": "has_more"}, + parameters=parameters, + truncation_indicator_path=["lines", "{{ config['indicator'] }}"], + truncated_list_retriever=_make_retriever([{"id": 2}]), + ) + records = list(expander.expand_record({"lines": {"data": [{"id": 1}], "has_more": True}})) + assert records == [{"id": 2}] + + +def test_indicator_path_through_non_mapping_is_not_truncated(): + expander = RecordExpander( + expand_records_from_field=["lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["lines", "has_more", "nested"], + ) + records = list(expander.expand_record({"lines": {"data": [{"id": 1}], "has_more": True}})) + assert records == [{"id": 1}] + + +def test_fetched_scalar_items_are_yielded_like_embedded_items(): + retriever = _make_retriever(["a", "b"]) + expander = RecordExpander( + expand_records_from_field=["lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["lines", "has_more"], + truncated_list_retriever=retriever, + ) + records = list(expander.expand_record({"lines": {"data": ["a"], "has_more": True}})) + assert records == ["a", "b"] + + +def test_fetched_scalar_items_are_wrapped_when_remain_original_record(): + retriever = _make_retriever(["a", "b"]) + parent = {"lines": {"data": ["a"], "has_more": True}} + expander = RecordExpander( + expand_records_from_field=["lines", "data"], + config=config, + parameters=parameters, + remain_original_record=True, + truncation_indicator_path=["lines", "has_more"], + truncated_list_retriever=retriever, + ) + records = list(expander.expand_record(parent)) + assert records == [ + {"value": "a", "original_record": parent}, + {"value": "b", "original_record": parent}, + ] + + +def test_retriever_errors_propagate(): + retriever = MagicMock() + retriever.read_records.side_effect = RuntimeError("boom") + expander = RecordExpander( + expand_records_from_field=["lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["lines", "has_more"], + truncated_list_retriever=retriever, + ) + with pytest.raises(RuntimeError): + list(expander.expand_record({"lines": {"data": [{"id": 1}], "has_more": True}})) + + +def _retriever_expander(retriever, message_repository=None): + return RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + truncated_list_retriever=retriever, + message_repository=message_repository, + ) + + +def test_protocol_messages_from_retriever_are_not_treated_as_child_records(): + embedded = [{"id": f"il_{i}"} for i in range(10)] + log_message = AirbyteMessage( + type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="custom retriever log") + ) + record_message = AirbyteMessage( + type=Type.RECORD, + record=AirbyteRecordMessage(stream="lines", data={"id": "il_1"}, emitted_at=0), + ) + bare_log_message = AirbyteLogMessage(level=Level.DEBUG, message="request/response log") + retriever = _make_retriever( + [ + {"id": "il_0"}, + log_message, + bare_log_message, + record_message, + Record(data={"id": "il_2"}, stream_name="t"), + ] + ) + expander = _retriever_expander(retriever) + + records = list(expander.expand_record(_event(embedded, has_more=True, total_count=3))) + + assert records == [{"id": "il_0"}, {"id": "il_1"}, {"id": "il_2"}] + + +def test_warns_once_when_retriever_fetches_fewer_than_total_count(caplog): + embedded = [{"id": f"il_{i}"} for i in range(10)] + first_page_only = [{"id": f"il_{i}"} for i in range(12)] + retriever = MagicMock() + retriever.read_records.side_effect = lambda **_: iter(first_page_only) + expander = _retriever_expander(retriever) + + with caplog.at_level("WARNING", logger="airbyte"): + first = list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + second = list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + + assert len(first) == len(second) == 12 + warnings = [r for r in caplog.records if r.levelname == "WARNING"] + assert len(warnings) == 1 + message = warnings[0].getMessage() + assert "`truncated_list_retriever`" in message + assert "returned 12 record(s)" in message + assert "reports 15" in message + assert "['data', 'object', 'lines', 'data']" in message + assert "il_" not in message + + +@pytest.mark.parametrize( + "suppress_incomplete_fetch_warning, expected_warning_count", + [ + pytest.param(True, 0, id="test_read_page_cap_suppresses_warning"), + pytest.param(False, 1, id="no_page_cap_still_warns"), + ], +) +def test_incomplete_fetch_warning_under_test_read_page_cap( + caplog, suppress_incomplete_fetch_warning, expected_warning_count +): + embedded = [{"id": f"il_{i}"} for i in range(10)] + capped_pages = [{"id": f"il_{i}"} for i in range(200)] + expander = _retriever_expander(_make_retriever(capped_pages)) + expander.suppress_incomplete_fetch_warning = suppress_incomplete_fetch_warning + + with caplog.at_level("WARNING", logger="airbyte"): + records = list(expander.expand_record(_event(embedded, has_more=True, total_count=250))) + + assert len(records) == 200 + messages = [r.getMessage() for r in caplog.records if r.levelname == "WARNING"] + assert len(messages) == expected_warning_count + for message in messages: + assert "returned 200 record(s)" in message + assert "reports 250" in message + + +def test_test_read_page_cap_flag_does_not_suppress_no_retriever_warning(caplog): + embedded = [{"id": f"il_{i}"} for i in range(10)] + expander = RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + suppress_incomplete_fetch_warning=True, + ) + + with caplog.at_level("WARNING", logger="airbyte"): + records = list(expander.expand_record(_event(embedded, has_more=True, total_count=250))) + + assert len(records) == 10 + messages = [r.getMessage() for r in caplog.records if r.levelname == "WARNING"] + assert len(messages) == 1 + assert "no `truncated_list_retriever` is configured" in messages[0] + + +def test_no_retriever_warning_is_emitted_even_when_consumer_stops_after_first_child(caplog): + embedded = [{"id": f"il_{i}"} for i in range(10)] + expander = RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + ) + + with caplog.at_level("WARNING", logger="airbyte"): + first = next(iter(expander.expand_record(_event(embedded, True, 15)))) + + assert first == {"id": "il_0"} + messages = [r.getMessage() for r in caplog.records if r.levelname == "WARNING"] + assert len(messages) == 1 + assert "no `truncated_list_retriever` is configured" in messages[0] + assert "10 embedded item(s) of 15 total" in messages[0] + + +def test_fetched_children_are_streamed_not_buffered(): + embedded = [{"id": f"il_{i}"} for i in range(10)] + consumed = [] + + def stream(): + for i in range(3): + consumed.append(i) + yield {"id": f"il_{i}"} + + retriever = MagicMock() + retriever.read_records.return_value = stream() + expander = _retriever_expander(retriever) + + records = expander.expand_record(_event(embedded, has_more=True, total_count=3)) + assert next(records) == {"id": "il_0"} + assert consumed == [0] + + +@pytest.mark.parametrize("total_count", [15, None, True, "15"]) +def test_no_incomplete_fetch_warning_when_count_matches_or_total_unknown(caplog, total_count): + embedded = [{"id": f"il_{i}"} for i in range(10)] + complete = [{"id": f"il_{i}"} for i in range(15)] + expander = _retriever_expander(_make_retriever(complete)) + + with caplog.at_level("WARNING", logger="airbyte"): + records = list( + expander.expand_record(_event(embedded, has_more=True, total_count=total_count)) + ) + + assert len(records) == 15 + assert not [r for r in caplog.records if r.levelname == "WARNING"] + + +def test_warns_when_retriever_returns_nothing_but_total_count_is_positive(caplog): + embedded = [{"id": f"il_{i}"} for i in range(10)] + expander = _retriever_expander(_make_retriever([])) + + with caplog.at_level("WARNING", logger="airbyte"): + records = list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + + assert [record["id"] for record in records] == [f"il_{i}" for i in range(10)] + warnings = [r for r in caplog.records if r.levelname == "WARNING"] + assert len(warnings) == 1 + message = warnings[0].getMessage() + assert "returned 0 record(s)" in message + assert "reports 15" in message + assert "embedded items were expanded as a fallback" in message + + +def test_no_warning_when_retriever_returns_nothing_and_total_count_unknown(caplog): + embedded = [{"id": f"il_{i}"} for i in range(10)] + expander = _retriever_expander(_make_retriever([])) + + with caplog.at_level("WARNING", logger="airbyte"): + records = list(expander.expand_record(_event(embedded, has_more=True, total_count=None))) + + assert len(records) == 10 + assert not [r for r in caplog.records if r.levelname == "WARNING"] + + +def test_warning_is_emitted_once_across_concurrent_expansions(): + repository = InMemoryMessageRepository() + expander = RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + message_repository=repository, + ) + embedded = [{"id": f"il_{i}"} for i in range(10)] + barrier = threading.Barrier(8) + + def expand(): + barrier.wait() + list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + + threads = [threading.Thread(target=expand) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert len(list(repository.consume_queue())) == 1 + + +def test_warnings_go_to_message_repository_instead_of_logger_when_configured(caplog): + repository = InMemoryMessageRepository() + expander = _retriever_expander(_make_retriever([]), message_repository=repository) + embedded = [{"id": f"il_{i}"} for i in range(10)] + + with caplog.at_level("WARNING", logger="airbyte"): + list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + + assert len(list(repository.consume_queue())) == 1 + assert not [r for r in caplog.records if r.levelname == "WARNING"] + + +def test_warnings_are_emitted_through_message_repository(): + repository = InMemoryMessageRepository() + expander = RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + message_repository=repository, + ) + embedded = [{"id": f"il_{i}"} for i in range(10)] + + list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + + messages = list(repository.consume_queue()) + assert len(messages) == 1 + assert messages[0].type == Type.LOG + assert messages[0].log.level == Level.WARN + assert "10 embedded item(s) of 15 total" in messages[0].log.message + + +def _indicator_only_expander(): + return RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + ) + + +def test_warns_when_truncated_and_no_retriever_configured(caplog): + embedded = [{"id": f"il_{i}"} for i in range(10)] + expander = _indicator_only_expander() + + with caplog.at_level("WARNING", logger="airbyte"): + records = list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + + assert len(records) == 10 + warnings = [r for r in caplog.records if r.levelname == "WARNING"] + assert len(warnings) == 1 + message = warnings[0].getMessage() + assert "['data', 'object', 'lines', 'data']" in message + assert "['data', 'object', 'lines', 'has_more']" in message + assert "10 embedded item(s)" in message + assert "of 15 total" in message + + +def test_warning_emitted_once_per_stream_instance(caplog): + embedded = [{"id": f"il_{i}"} for i in range(10)] + expander = _indicator_only_expander() + + with caplog.at_level("WARNING", logger="airbyte"): + list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + list(expander.expand_record(_event(embedded, has_more=True, total_count=20))) + list(expander.expand_record(_event(embedded, has_more=True, total_count=25))) + + warnings = [r for r in caplog.records if r.levelname == "WARNING"] + assert len(warnings) == 1 + + +def test_warning_omits_total_when_no_total_count_sibling(caplog): + embedded = [{"id": "il_0"}] + expander = _indicator_only_expander() + parent = _event(embedded, has_more=True, total_count=5) + del parent["data"]["object"]["lines"]["total_count"] + + with caplog.at_level("WARNING", logger="airbyte"): + records = list(expander.expand_record(parent)) + + assert len(records) == 1 + warnings = [r for r in caplog.records if r.levelname == "WARNING"] + assert len(warnings) == 1 + assert "total" not in warnings[0].getMessage().split("embedded item(s)")[1].split(" were")[0] + + +def test_no_warning_when_indicator_falsy(caplog): + embedded = [{"id": f"il_{i}"} for i in range(10)] + expander = _indicator_only_expander() + + with caplog.at_level("WARNING", logger="airbyte"): + records = list(expander.expand_record(_event(embedded, has_more=False, total_count=10))) + + assert len(records) == 10 + assert not [r for r in caplog.records if r.levelname == "WARNING"] + + +def test_no_warning_when_retriever_configured(caplog): + embedded = [{"id": f"il_{i}"} for i in range(10)] + complete = [{"id": f"il_{i}"} for i in range(15)] + retriever = _make_retriever([Record(data=item, stream_name="test") for item in complete]) + expander = RecordExpander( + expand_records_from_field=["data", "object", "lines", "data"], + config=config, + parameters=parameters, + truncation_indicator_path=["data", "object", "lines", "has_more"], + truncated_list_retriever=retriever, + ) + + with caplog.at_level("WARNING", logger="airbyte"): + records = list(expander.expand_record(_event(embedded, has_more=True, total_count=15))) + + assert len(records) == 15 + assert not [r for r in caplog.records if r.levelname == "WARNING"] diff --git a/unit_tests/sources/declarative/expanders/test_record_expander_http.py b/unit_tests/sources/declarative/expanders/test_record_expander_http.py new file mode 100644 index 0000000000..99d468a1b2 --- /dev/null +++ b/unit_tests/sources/declarative/expanders/test_record_expander_http.py @@ -0,0 +1,217 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import json +from copy import deepcopy +from typing import Any, Dict, List +from unittest.mock import MagicMock + +from airbyte_cdk.models import ConfiguredAirbyteCatalog, ConfiguredAirbyteStream, Type +from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( + ConcurrentDeclarativeSource, +) +from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse + +_CONFIG: Dict[str, Any] = {} +_URL_BASE = "https://api.test.com/v1" + +_MANIFEST: Dict[str, Any] = { + "version": "6.0.0", + "type": "DeclarativeSource", + "check": {"type": "CheckStream", "stream_names": ["invoice_line_items"]}, + "streams": [ + { + "type": "DeclarativeStream", + "name": "invoice_line_items", + "primary_key": ["id"], + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": { + "$schema": "http://json-schema.org/schema#", + "type": "object", + "properties": {"id": {"type": "string"}}, + }, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": _URL_BASE, + "path": "events", + "http_method": "GET", + "request_parameters": {"types[]": "invoice.created"}, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": { + "type": "DpathExtractor", + "field_path": ["data"], + "record_expander": { + "type": "RecordExpander", + "expand_records_from_field": ["data", "object", "lines", "data"], + "truncation_indicator_path": ["data", "object", "lines", "has_more"], + "truncated_list_retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": _URL_BASE, + "path": "invoices/{{ stream_slice['parent_record']['data']['object']['id'] }}/lines", + "http_method": "GET", + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": ["data"]}, + }, + "paginator": { + "type": "DefaultPaginator", + "page_token_option": { + "type": "RequestOption", + "inject_into": "request_parameter", + "field_name": "starting_after", + }, + "pagination_strategy": { + "type": "CursorPagination", + "cursor_value": '{{ response["data"][-1]["id"] }}', + "stop_condition": '{{ not response.get("has_more", False) }}', + }, + }, + }, + }, + }, + }, + }, + } + ], + "spec": { + "type": "Spec", + "connection_specification": { + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": {}, + "additionalProperties": True, + }, + }, +} + + +def _invoice_event(invoice_id: str, embedded: List[str], has_more: bool, total_count: int): + return { + "id": f"evt_{invoice_id}", + "type": "invoice.created", + "data": { + "object": { + "id": invoice_id, + "object": "invoice", + "lines": { + "object": "list", + "data": [{"id": line_id, "object": "line_item"} for line_id in embedded], + "has_more": has_more, + "total_count": total_count, + "url": f"/v1/invoices/{invoice_id}/lines", + }, + } + }, + } + + +def _lines_page(line_ids: List[str], has_more: bool): + return HttpResponse( + body=json.dumps( + { + "object": "list", + "data": [{"id": line_id, "object": "line_item"} for line_id in line_ids], + "has_more": has_more, + } + ) + ) + + +def _read_records(manifest: Dict[str, Any]) -> List[Dict[str, Any]]: + source = ConcurrentDeclarativeSource( + source_config=manifest, config=_CONFIG, catalog=None, state=None + ) + catalog = ConfiguredAirbyteCatalog( + streams=[ + ConfiguredAirbyteStream( + stream=stream, sync_mode="full_refresh", destination_sync_mode="overwrite" + ) + for stream in source.discover(logger=source.logger, config=_CONFIG).streams + ] + ) + return [ + message.record.data + for message in source.read(logger=MagicMock(), config=_CONFIG, catalog=catalog, state=None) + if message.type == Type.RECORD + ] + + +def test_truncated_nested_list_is_fetched_over_http_with_pagination(): + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url=f"{_URL_BASE}/events", query_params={"types[]": "invoice.created"}), + HttpResponse( + body=json.dumps( + { + "object": "list", + "data": [ + _invoice_event("in_1", ["il_1", "il_2"], has_more=True, total_count=5), + _invoice_event("in_2", ["il_6"], has_more=False, total_count=1), + ], + "has_more": False, + } + ) + ), + ) + first_page_request = HttpRequest(url=f"{_URL_BASE}/invoices/in_1/lines") + second_page_request = HttpRequest( + url=f"{_URL_BASE}/invoices/in_1/lines", query_params={"starting_after": "il_3"} + ) + http_mocker.get(first_page_request, _lines_page(["il_1", "il_2", "il_3"], has_more=True)) + http_mocker.get(second_page_request, _lines_page(["il_4", "il_5"], has_more=False)) + + records = _read_records(_MANIFEST) + + assert [record["id"] for record in records] == [ + "il_1", + "il_2", + "il_3", + "il_4", + "il_5", + "il_6", + ] + http_mocker.assert_number_of_calls(first_page_request, 1) + http_mocker.assert_number_of_calls(second_page_request, 1) + + +def test_stream_parameters_propagate_into_truncated_list_requester(): + """Request-shaping options declared in the stream's `$parameters` also reach the nested requester, + which is why the docs recommend declaring them on the outer requester's `request_parameters`.""" + with HttpMocker() as http_mocker: + manifest = deepcopy(_MANIFEST) + stream = manifest["streams"][0] + del stream["retriever"]["requester"]["request_parameters"] + stream["$parameters"] = {"request_parameters": {"types[]": "invoice.created"}} + + http_mocker.get( + HttpRequest(url=f"{_URL_BASE}/events", query_params={"types[]": "invoice.created"}), + HttpResponse( + body=json.dumps( + { + "object": "list", + "data": [_invoice_event("in_1", ["il_1"], has_more=True, total_count=2)], + "has_more": False, + } + ) + ), + ) + http_mocker.get( + HttpRequest( + url=f"{_URL_BASE}/invoices/in_1/lines", query_params={"types[]": "invoice.created"} + ), + _lines_page(["il_1", "il_2"], has_more=False), + ) + + records = _read_records(manifest) + + assert [record["id"] for record in records] == ["il_1", "il_2"] 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..39adbde700 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 @@ -156,7 +156,10 @@ WaitUntilTimeFromHeaderBackoffStrategy, ) from airbyte_cdk.sources.declarative.requesters.http_job_repository import AsyncHttpJobRepository -from airbyte_cdk.sources.declarative.requesters.paginators import DefaultPaginator +from airbyte_cdk.sources.declarative.requesters.paginators import ( + DefaultPaginator, + PaginatorTestReadDecorator, +) from airbyte_cdk.sources.declarative.requesters.paginators.strategies import ( CursorPaginationStrategy, OffsetIncrement, @@ -225,6 +228,7 @@ from airbyte_cdk.utils import AirbyteTracedException from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now, ab_datetime_parse from unit_tests.sources.declarative.parsers.testing_components import ( + TestingCustomRetriever, TestingCustomSubstreamPartitionRouter, TestingSomeComponent, ) @@ -1986,6 +1990,254 @@ def test_client_side_incremental_with_partition_router(): ) +def test_create_record_expander_with_truncated_list_retriever(): + content = """ + selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data"] + record_expander: + type: RecordExpander + expand_records_from_field: ["data", "object", "lines", "data"] + remain_original_record: true + truncation_indicator_path: ["data", "object", "lines", "has_more"] + truncated_list_retriever: + type: SimpleRetriever + requester: + type: HttpRequester + url_base: "https://api.stripe.com/v1/" + path: "invoices/{{ stream_slice['parent_record']['data']['object']['id'] }}/lines" + http_method: "GET" + record_selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + expander = selector.extractor.record_expander + assert expander is not None + assert expander.truncation_indicator_path == ["data", "object", "lines", "has_more"] + assert isinstance(expander.truncated_list_retriever, SimpleRetriever) + assert expander.truncated_list_retriever.name == "record_expander_truncated_list" + assert expander.message_repository is factory._message_repository + assert expander.suppress_incomplete_fetch_warning is False + + +# `$parameters.url_base` is how a `CustomRetriever` shares the requester's `url_base` with a nested paginator. +_CURSOR_PAGINATOR_YAML = """ + paginator: + type: DefaultPaginator + $parameters: + url_base: "https://api.test.com/" + page_token_option: + type: RequestOption + inject_into: request_parameter + field_name: starting_after + pagination_strategy: + type: CursorPagination + cursor_value: "{{ last_record['id'] }}" + stop_condition: "{{ not response['has_more'] }}" +""" + + +@pytest.mark.parametrize( + "paginator_yaml, expected_suppressed", + [ + pytest.param( + _CURSOR_PAGINATOR_YAML, True, id="capped_default_paginator_suppresses_warning" + ), + pytest.param("", False, id="no_pagination_is_not_capped_so_warning_stays"), + ], +) +@pytest.mark.parametrize( + "retriever_type", + [ + "type: SimpleRetriever", + "type: CustomRetriever\n class_name: unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever", + ], + ids=["simple_retriever", "custom_retriever"], +) +def test_create_record_expander_suppresses_incomplete_fetch_warning_only_for_capped_paginator( + retriever_type, paginator_yaml, expected_suppressed +): + content = _record_expander_selector( + f""" + {retriever_type} + requester: + type: HttpRequester + url_base: "https://api.test.com/" + path: "invoices/{{{{ stream_slice['parent_record']['id'] }}}}/lines" + http_method: "GET" + record_selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data"] + {paginator_yaml} + """ + ) + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + test_read_factory = ModelToComponentFactory(limit_pages_fetched_per_slice=2) + + selector = test_read_factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + expander = selector.extractor.record_expander + assert expander.suppress_incomplete_fetch_warning is expected_suppressed + assert ( + isinstance(expander.truncated_list_retriever.paginator, PaginatorTestReadDecorator) + is expected_suppressed + ) + + +def _record_expander_selector(retriever_yaml: str) -> str: + return f""" + selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data"] + record_expander: + type: RecordExpander + expand_records_from_field: ["lines", "data"] + truncation_indicator_path: ["lines", "has_more"] + truncated_list_retriever: +{retriever_yaml} + """ + + +def _create_record_expander(content: str): + 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, + ) + return selector.extractor.record_expander + + +def test_create_record_expander_with_custom_truncated_list_retriever(): + expander = _create_record_expander( + _record_expander_selector( + """ + type: CustomRetriever + class_name: unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever + requester: + type: HttpRequester + url_base: "https://api.test.com/" + path: "invoices/{{ stream_slice['parent_record']['id'] }}/lines" + http_method: "GET" + record_selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data"] + """ + ) + ) + retriever = expander.truncated_list_retriever + assert isinstance(retriever, TestingCustomRetriever) + assert retriever.name == "record_expander_truncated_list" + assert isinstance(retriever.requester, HttpRequester) + assert isinstance(retriever.record_selector, RecordSelector) + assert expander.suppress_incomplete_fetch_warning is False + + request = requests.PreparedRequest() + request.headers = {} + request.url = "https://api.test.com/invoices/in_1/lines" + response = requests.Response() + response.request = request + response.status_code = 200 + assert retriever.log_formatter is not None + assert retriever.log_formatter(response)["http"]["is_auxiliary"] is True + + +@pytest.mark.parametrize( + "unsupported_option", + [ + pytest.param( + """ + partition_router: + type: ListPartitionRouter + values: ["a", "b"] + cursor_field: "partition" + """, + id="partition_router", + ), + pytest.param( + """ + pagination_reset: + type: PaginationReset + action: RESET + """, + id="pagination_reset", + ), + ], +) +@pytest.mark.parametrize( + "retriever_type", + [ + "type: SimpleRetriever", + "type: CustomRetriever\n class_name: unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever", + ], + ids=["simple_retriever", "custom_retriever"], +) +def test_create_record_expander_rejects_unsupported_retriever_options( + unsupported_option, retriever_type +): + content = _record_expander_selector( + f""" + {retriever_type} + requester: + type: HttpRequester + url_base: "https://api.test.com/" + path: "lines" + http_method: "GET" + record_selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data"] + """ + + unsupported_option + ) + with pytest.raises(ValueError, match="not supported on `truncated_list_retriever`"): + _create_record_expander(content) + + @pytest.mark.parametrize( "test_name, record_selector, expected_runtime_selector", [