diff --git a/airbyte_cdk/__init__.py b/airbyte_cdk/__init__.py index d971690c86..4ee2097c73 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 ( + DpathExtractor, + NestedRecordExtractor, + ParentFieldPath, + RecordSelector, +) from .sources.declarative.extractors.record_extractor import RecordExtractor from .sources.declarative.extractors.record_filter import RecordFilter from .sources.declarative.interpolation import InterpolatedBoolean, InterpolatedString @@ -254,10 +259,12 @@ "JsonFileSchemaLoader", "LegacyToPerPartitionStateMigration", "MinMaxDatetime", + "NestedRecordExtractor", "NoAuth", "OffsetIncrement", "PageIncrement", "PaginationStrategy", + "ParentFieldPath", "ParentStreamConfig", "ReadException", "RecordExtractor", diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index aa3bbbc5e0..a55e287cd3 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2203,6 +2203,88 @@ definitions: $parameters: type: object additionalProperties: true + NestedRecordExtractor: + title: Nested Record Extractor + description: Record extractor that yields the records of a collection nested inside another record, while carrying fields down from the record that contains it. For every record the parent extractor yields, the nested collection is read from the child field path and each of its elements is emitted, with the configured parent fields copied onto it first. + type: object + required: + - type + - parent_extractor + - child_field_path + properties: + type: + type: string + enum: [NestedRecordExtractor] + parent_extractor: + title: Parent Extractor + description: Extractor producing the records that contain the nested collection. Nest another NestedRecordExtractor here to reach a collection more than one level down, or to copy fields from more than one ancestor. The type is required here, because this field accepts a NestedRecordExtractor as readily as a DpathExtractor. + anyOf: + - "$ref": "#/definitions/DpathExtractor" + - "$ref": "#/definitions/NestedRecordExtractor" + - "$ref": "#/definitions/CustomRecordExtractor" + child_field_path: + title: Child Field Path + description: List of potentially nested fields describing the full path of the nested collection on each parent record. The collection may be a list or a single object. A path that does not resolve yields no records. The "*" wildcard is not supported and is rejected rather than read as a field name; flatten across several collections in the parent extractor instead. Every other character addresses the field that spells it. + type: array + minItems: 1 + items: + type: string + interpolation_context: + - config + examples: + - ["reviews", "nodes"] + - ["comments"] + - ["{{ parameters.child_connection }}", "nodes"] + parent_fields: + title: Parent Fields + description: Fields to copy from the parent record onto every record extracted from the nested collection. A parent path that is absent on the parent, or that runs into a null, copies null; one that runs into a value holding no such field fails the sync rather than copying null silently. A key already present on the child record is overwritten. + type: array + items: + "$ref": "#/definitions/ParentFieldPath" + examples: + - - parent_path: ["url"] + record_path: ["pull_request_url"] + - - parent_path: ["id"] + record_path: ["pull_request_id"] + - parent_path: ["author", "login"] + record_path: ["pull_request_author"] + $parameters: + type: object + additionalProperties: true + ParentFieldPath: + title: Parent Field Path + description: A single copy instruction, reading a field from the parent record and writing it onto every record extracted from the nested collection. + type: object + required: + - parent_path + - record_path + properties: + parent_path: + title: Parent Path + description: List of potentially nested fields describing the full path of the field to read on the parent record. The "*" wildcard is not supported and is rejected rather than read as a field name. + type: array + minItems: 1 + items: + type: string + interpolation_context: + - config + examples: + - ["url"] + - ["id"] + - ["author", "login"] + record_path: + title: Record Path + description: List of potentially nested fields describing where to write the value on the extracted record. Intermediate objects are created as needed, and an intermediate that is present but is not an object is replaced by one. Every segment is an object key, so a numeric segment does not create a list. + type: array + minItems: 1 + items: + type: string + interpolation_context: + - config + examples: + - ["pull_request_url"] + - ["comment_id"] + - ["parent", "id"] ResponseToFileExtractor: title: CSV To File Extractor description: A record extractor designed for handling large responses that may exceed memory limits (to prevent OOM issues). It downloads a CSV file to disk, reads the data from disk, and deletes the file once it has been fully processed. @@ -3948,6 +4030,7 @@ definitions: extractor: anyOf: - "$ref": "#/definitions/DpathExtractor" + - "$ref": "#/definitions/NestedRecordExtractor" - "$ref": "#/definitions/CustomRecordExtractor" record_filter: title: Record Filter diff --git a/airbyte_cdk/sources/declarative/extractors/__init__.py b/airbyte_cdk/sources/declarative/extractors/__init__.py index 8f1d18d12f..8f57bc787f 100644 --- a/airbyte_cdk/sources/declarative/extractors/__init__.py +++ b/airbyte_cdk/sources/declarative/extractors/__init__.py @@ -4,6 +4,10 @@ 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.nested_record_extractor import ( + NestedRecordExtractor, + ParentFieldPath, +) from airbyte_cdk.sources.declarative.extractors.record_filter import RecordFilter from airbyte_cdk.sources.declarative.extractors.record_selector import RecordSelector from airbyte_cdk.sources.declarative.extractors.response_to_file_extractor import ( @@ -15,6 +19,8 @@ "TypeTransformer", "HttpSelector", "DpathExtractor", + "NestedRecordExtractor", + "ParentFieldPath", "RecordFilter", "RecordSelector", "ResponseToFileExtractor", diff --git a/airbyte_cdk/sources/declarative/extractors/nested_record_extractor.py b/airbyte_cdk/sources/declarative/extractors/nested_record_extractor.py new file mode 100644 index 0000000000..5051bdb5eb --- /dev/null +++ b/airbyte_cdk/sources/declarative/extractors/nested_record_extractor.py @@ -0,0 +1,367 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import copy +import logging +from dataclasses import InitVar, dataclass +from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Union + +import requests + +from airbyte_cdk.models import FailureType +from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor +from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString +from airbyte_cdk.sources.types import Config +from airbyte_cdk.utils.traced_exception import AirbyteTracedException + +logger = logging.getLogger("airbyte") + +# The `*` wildcard is rejected rather than treated as a literal key, because a manifest that writes +# one almost certainly means it to glob. Every other character is a literal: unlike the `dpath`-based +# extractors, this component walks paths itself, so `?` and `[...]` address the keys that spell them. +_UNSUPPORTED_PATH_CHARACTER = "*" + +# Values of these types are immutable, so every child can share the one the parent holds. Anything +# else is copied per child — see `NestedRecordExtractor.extract_records`. +_IMMUTABLE_VALUE_TYPES = (str, int, float, bool, type(None)) + + +# Returned by `_descend` when a segment does not address anything on the node it is given, which is +# distinct from a segment that addresses a field holding `None`. +_MISSING = object() + + +def _is_list_index(segment: str) -> bool: + """Whether `segment` addresses a list position. `isdecimal` rather than `isdigit`: the latter is + true for superscripts, which `int` then refuses.""" + return (segment[1:] if segment.startswith("-") else segment).isdecimal() + + +def _descend(node: Any, segment: str) -> Any: + """ + Follow one path segment into `node`, returning `_MISSING` when the segment addresses nothing. + + Objects are addressed by key and lists by decimal index, including a negative one, which is what + the `dpath`-based extractors accept. + """ + if isinstance(node, Mapping): + return node[segment] if segment in node else _MISSING + if isinstance(node, list): + if not _is_list_index(segment): + return _MISSING + try: + return node[int(segment)] + except IndexError: + return _MISSING + return _MISSING + + +def _normalize_path(path: List[Any], field_name: str, component: str) -> List[str]: + """ + Validate an interpolated path and return it as field names. + + Segments arrive as whatever the interpolation produced: `InterpolatedString` runs its result + through `ast.literal_eval`, so a segment spelled `0` evaluates to the integer `0`. JSON object + keys are always strings and list indices are addressed by their decimal spelling, so the segment + is normalized to `str`. A segment that resolved to nothing cannot address a field and is + rejected, rather than silently reading or writing a field named "". + """ + normalized = [] + for raw_segment in path: + segment = "" if raw_segment is None else str(raw_segment) + if segment == "": + raise AirbyteTracedException( + message=f"The connector is configured with an invalid `{field_name}`. Check the connector's configuration.", + internal_message=( + f"{component} received an empty segment in `{field_name}` ({path!r}). An empty " + f"segment usually means an interpolated expression resolved to nothing." + ), + failure_type=FailureType.config_error, + ) + if _UNSUPPORTED_PATH_CHARACTER in segment: + raise AirbyteTracedException( + message=f"The connector is configured with an invalid `{field_name}`. Check the connector's configuration.", + internal_message=( + f"{component} does not support the '{_UNSUPPORTED_PATH_CHARACTER}' wildcard in " + f"`{field_name}` ({path!r}). Flatten across several collections in the " + f"`parent_extractor` instead." + ), + failure_type=FailureType.config_error, + ) + normalized.append(segment) + return normalized + + +@dataclass +class ParentFieldPath: + """ + A single copy instruction: read `parent_path` out of the parent record and write it into the + child record at `record_path`. + + Both paths are lists of field names, so a value nested inside the parent can be copied into a + nested position on the child. Intermediate objects on `record_path` are created as needed, and + an intermediate that is present but is not an object is replaced by one. + + Attributes: + parent_path (List[Union[InterpolatedString, str]]): Path of the field to read on the parent record + record_path (List[Union[InterpolatedString, str]]): Path to write the value to on the child record + """ + + parent_path: List[Union[InterpolatedString, str]] + record_path: List[Union[InterpolatedString, str]] + parameters: InitVar[Mapping[str, Any]] + + def __post_init__(self, parameters: Mapping[str, Any]) -> None: + if not self.parent_path: + raise ValueError("ParentFieldPath requires a non-empty `parent_path`") + if not self.record_path: + raise ValueError("ParentFieldPath requires a non-empty `record_path`") + self._parent_path = [ + InterpolatedString.create(path, parameters=parameters) for path in self.parent_path + ] + self._record_path = [ + InterpolatedString.create(path, parameters=parameters) for path in self.record_path + ] + + def eval_parent_path(self, config: Config) -> List[str]: + return _normalize_path( + [path.eval(config) for path in self._parent_path], "parent_path", "ParentFieldPath" + ) + + def eval_record_path(self, config: Config) -> List[str]: + return _normalize_path( + [path.eval(config) for path in self._record_path], "record_path", "ParentFieldPath" + ) + + +@dataclass +class NestedRecordExtractor(RecordExtractor): + """ + Record extractor that yields the records of a collection nested inside another record, while + carrying fields down from the record that contains it. + + A `DpathExtractor` pointed at a child collection has already discarded the node that held it, so + a record that needs a field from its parent cannot be produced by that component alone. This + component keeps the parent in scope: for every record its `parent_extractor` yields it reads + `child_field_path` out of that record, yields each element, and copies the `parent_fields` + entries from the parent onto every element first. + + ``` + extractor: + type: NestedRecordExtractor + parent_extractor: + type: DpathExtractor + field_path: ["data", "repository", "pullRequests", "nodes"] + child_field_path: ["reviews", "nodes"] + parent_fields: + - parent_path: ["url"] + record_path: ["pull_request_url"] + ``` + + For a single level of nesting whose child collection is always a list, `DpathExtractor` with a + `record_expander` plus an `AddFields` transformation can produce the same records. This component + exists for the three shapes that combination cannot express: a child field that is a single + object rather than a list, fields copied from more than one ancestor (`record_expander` is not + nestable), and copying named fields rather than deep-copying the whole parent onto every child. + + Documented behaviour, all of it deliberate: + + * Child records are **mutated in place** rather than copied, when the parent extractor yields + mutable ones. The dicts come out of the response body this extractor just decoded and nothing + else reads them, so a copy would only cost memory. The *values* copied from the parent are a + different matter and are copied per child — see `extract_records`. + * A `parent_path` that is absent on the parent, or whose traversal runs into a null, copies + `None`. Every record of the stream then carries the field, which keeps the record shape stable + instead of making the field's presence depend on the parent. + * A `parent_path` that runs into a value that is neither an object nor null raises, because the + path cannot address a field there and silently copying `None` would hide a `parent_path` that + does not match the response. + * A key already present on the child **is overwritten**, including when an intermediate segment + of `record_path` holds something that is not an object. The parent context is the authoritative + source for the fields the manifest names, and silently keeping the child's value would hide a + misconfigured `record_path`. + * A `child_field_path` that is missing, null, or resolves to an empty collection yields nothing + rather than raising, the same way `DpathExtractor` treats a path that does not resolve. + * Parent records and child elements that are not objects are skipped, because a scalar cannot + carry the parent fields and emitting it would produce a record that silently lacks them. + * Paths do not support the `*` wildcard, and it is rejected rather than read as a literal key. + Flattening across several collections belongs in the `parent_extractor`, whose `DpathExtractor` + supports it; the last hop has to stay a single collection for its parentage to be well defined. + + `parent_extractor` may itself be a `NestedRecordExtractor`, which is how a collection more than + one level down is reached, and how fields from two different ancestors can be copied onto the + same record. `type` is required on `parent_extractor`: unlike `RecordSelector.extractor` it has + no default, because the field accepts a `NestedRecordExtractor` as readily as a `DpathExtractor`. + + This component holds no per-read mutable state: one instance is shared by every partition of a + stream and the partitions are read concurrently. + + Attributes: + parent_extractor (RecordExtractor): Extractor producing the records that contain the collection + child_field_path (List[Union[InterpolatedString, str]]): Path to the nested collection on each parent record + config (Config): The user-provided configuration as specified by the source's spec + parent_fields (Optional[List[ParentFieldPath]]): Fields to copy from the parent onto every child record + """ + + parent_extractor: RecordExtractor + child_field_path: List[Union[InterpolatedString, str]] + config: Config + parameters: InitVar[Mapping[str, Any]] + parent_fields: Optional[List[ParentFieldPath]] = None + + def __post_init__(self, parameters: Mapping[str, Any]) -> None: + if not self.child_field_path: + raise ValueError("NestedRecordExtractor requires a non-empty `child_field_path`") + self._child_field_path = [ + InterpolatedString.create(path, parameters=parameters) for path in self.child_field_path + ] + self._parent_fields = self.parent_fields or [] + + def extract_records(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: + child_field_path = _normalize_path( + [path.eval(self.config) for path in self._child_field_path], + "child_field_path", + "NestedRecordExtractor", + ) + field_copies = [ + (parent_field.eval_parent_path(self.config), parent_field.eval_record_path(self.config)) + for parent_field in self._parent_fields + ] + for parent_path, _ in field_copies: + self._reject_self_referential_copy(parent_path, child_field_path) + + for parent in self.parent_extractor.extract_records(response): + if not isinstance(parent, Mapping): + logger.debug( + "NestedRecordExtractor skipped a parent record that is not an object (%s)", + type(parent).__name__, + ) + continue + + # Reading the parent is invariant across its children, so it happens once per parent. + # Doing it per child makes the extractor quadratic in the size of the child collection, + # because each read walks a parent that contains that collection. + values = [] + for parent_path, record_path in field_copies: + value = self._read_parent_value(parent, parent_path) + values.append((record_path, value, not isinstance(value, _IMMUTABLE_VALUE_TYPES))) + + for child in self._children(parent, child_field_path): + if not isinstance(child, MutableMapping): + # A parent extractor may legally yield read-only mappings: the `RecordExtractor` + # contract is `Mapping`. Copy rather than drop the record. + child = dict(child) + for record_path, value, needs_copy in values: + # A mutable value must be copied per child. Sharing one object across the + # collection means a later transformation that writes into it rewrites records + # that were already emitted — they are queued unserialized downstream. + self._write_child_value( + child, record_path, copy.deepcopy(value) if needs_copy else value + ) + yield child + + @staticmethod + def _read_parent_value(parent: Mapping[str, Any], parent_path: List[str]) -> Any: + """Read `parent_path` out of `parent`, returning `None` when the path does not resolve.""" + node: Any = parent + for index, segment in enumerate(parent_path): + if node is None: + return None + if not isinstance(node, Mapping) and not ( + isinstance(node, list) and _is_list_index(segment) + ): + raise AirbyteTracedException( + message="A record returned by the API does not have the shape the connector expects.", + internal_message=( + f"NestedRecordExtractor cannot read `parent_path` {parent_path!r}: the " + f"segment {segment!r} does not address a field on the " + f"{type(node).__name__} it was reached on. The path resolved up to " + f"{parent_path[:index]!r}." + ), + failure_type=FailureType.system_error, + ) + node = _descend(node, segment) + if node is _MISSING: + return None + return node + + @staticmethod + def _write_child_value( + child: MutableMapping[str, Any], record_path: List[str], value: Any + ) -> None: + """ + Write `value` into `child` at `record_path`, creating intermediate objects as needed. + + An intermediate that is present but is not an object is replaced by one, so the documented + "a key already present on the child is overwritten" holds for a nested `record_path` too. + Segments are always object keys — a numeric-looking segment does not create a list. + """ + node = child + for segment in record_path[:-1]: + existing = node.get(segment) + if isinstance(existing, MutableMapping): + node = existing + elif isinstance(existing, Mapping): + replacement = dict(existing) + node[segment] = replacement + node = replacement + else: + replacement = {} + node[segment] = replacement + node = replacement + node[record_path[-1]] = value + + @staticmethod + def _reject_self_referential_copy(parent_path: List[str], child_field_path: List[str]) -> None: + """ + Reject a `parent_path` that addresses the child collection, or something containing it. + + The child records are not copied out of the parent, so the value read by such a path holds + the very records being emitted. Copying it onto each of them gives every record a snapshot + of its whole sibling collection — quadratic in its size, and never what a manifest means. + """ + if child_field_path[: len(parent_path)] == parent_path: + raise AirbyteTracedException( + message="The connector is configured with an invalid `parent_path`. Check the connector's configuration.", + internal_message=( + f"NestedRecordExtractor cannot copy `parent_path` {parent_path!r} onto the " + f"records of `child_field_path` {child_field_path!r}: the path addresses the " + f"child collection itself, so every record would carry a copy of its siblings." + ), + failure_type=FailureType.config_error, + ) + + @staticmethod + def _children( + parent: Mapping[str, Any], child_field_path: List[str] + ) -> Iterable[Mapping[str, Any]]: + """Yield the elements of the nested collection without materialising them into a new list.""" + extracted: Any = parent + for segment in child_field_path: + extracted = _descend(extracted, segment) + if extracted is _MISSING: + return + + if isinstance(extracted, Mapping): + if extracted: + # A single object where a collection would also be allowed, the shape a GraphQL + # drill-down document returns. An empty object yields nothing, matching how + # `DpathExtractor` treats one. + yield extracted + elif isinstance(extracted, list): + for child in extracted: + if isinstance(child, Mapping): + yield child + else: + logger.debug( + "NestedRecordExtractor skipped a child element that is not an object (%s)", + type(child).__name__, + ) + elif extracted is not None: + logger.debug( + "NestedRecordExtractor resolved `child_field_path` %s to a %s, which holds no " + "records", + child_field_path, + type(extracted).__name__, + ) diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 21a92cc3be..1773053095 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -507,6 +507,23 @@ class Config: ) +class ParentFieldPath(BaseModel): + parent_path: List[str] = Field( + ..., + description='List of potentially nested fields describing the full path of the field to read on the parent record. The "*" wildcard is not supported and is rejected rather than read as a field name.', + examples=[["url"], ["id"], ["author", "login"]], + min_items=1, + title="Parent Path", + ) + record_path: List[str] = Field( + ..., + description="List of potentially nested fields describing where to write the value on the extracted record. Intermediate objects are created as needed, and an intermediate that is present but is not an object is replaced by one. Every segment is an object key, so a numeric segment does not create a list.", + examples=[["pull_request_url"], ["comment_id"], ["parent", "id"]], + min_items=1, + title="Record Path", + ) + + class ResponseToFileExtractor(BaseModel): type: Literal["ResponseToFileExtractor"] preserve_na_values: Optional[bool] = Field( @@ -2517,6 +2534,42 @@ class DpathExtractor(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class NestedRecordExtractor(BaseModel): + type: Literal["NestedRecordExtractor"] + parent_extractor: Union[DpathExtractor, NestedRecordExtractor, CustomRecordExtractor] = Field( + ..., + description="Extractor producing the records that contain the nested collection. Nest another NestedRecordExtractor here to reach a collection more than one level down, or to copy fields from more than one ancestor. The type is required here, because this field accepts a NestedRecordExtractor as readily as a DpathExtractor.", + title="Parent Extractor", + ) + child_field_path: List[str] = Field( + ..., + description='List of potentially nested fields describing the full path of the nested collection on each parent record. The collection may be a list or a single object. A path that does not resolve yields no records. The "*" wildcard is not supported and is rejected rather than read as a field name; flatten across several collections in the parent extractor instead. Every other character addresses the field that spells it.', + examples=[ + ["reviews", "nodes"], + ["comments"], + ["{{ parameters.child_connection }}", "nodes"], + ], + min_items=1, + title="Child Field Path", + ) + parent_fields: Optional[List[ParentFieldPath]] = Field( + None, + description="Fields to copy from the parent record onto every record extracted from the nested collection. A parent path that is absent on the parent, or that runs into a null, copies null; one that runs into a value holding no such field fails the sync rather than copying null silently. A key already present on the child record is overwritten.", + examples=[ + [{"parent_path": ["url"], "record_path": ["pull_request_url"]}], + [ + {"parent_path": ["id"], "record_path": ["pull_request_id"]}, + { + "parent_path": ["author", "login"], + "record_path": ["pull_request_author"], + }, + ], + ], + title="Parent Fields", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class ZipfileDecoder(BaseModel): class Config: extra = Extra.allow @@ -2531,7 +2584,7 @@ class Config: class RecordSelector(BaseModel): type: Literal["RecordSelector"] - extractor: Union[DpathExtractor, CustomRecordExtractor] + extractor: Union[DpathExtractor, NestedRecordExtractor, CustomRecordExtractor] record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( None, description="Responsible for filtering records to be emitted by the Source.", @@ -3476,6 +3529,7 @@ class DynamicDeclarativeStream(BaseModel): ComplexFieldType.update_forward_refs() +NestedRecordExtractor.update_forward_refs() GzipDecoder.update_forward_refs() CompositeErrorHandler.update_forward_refs() DeclarativeSource1.update_forward_refs() diff --git a/airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py b/airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py index 00b1a18ffb..d8807e78c4 100644 --- a/airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py +++ b/airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py @@ -79,6 +79,7 @@ "DefaultPaginator.pagination_strategy": "CustomPaginationStrategy", "HttpRequester.authenticator": "CustomAuthenticator", "HttpRequester.error_handler": "CustomErrorHandler", + "NestedRecordExtractor.parent_extractor": "CustomRecordExtractor", "RecordSelector.extractor": "CustomRecordExtractor", "SimpleRetriever.partition_router": "CustomPartitionRouter", } 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..d8741024d7 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 @@ ) from airbyte_cdk.sources.declarative.extractors import ( DpathExtractor, + NestedRecordExtractor, + ParentFieldPath, RecordFilter, RecordSelector, ResponseToFileExtractor, @@ -369,6 +371,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( MovingWindowCallRatePolicy as MovingWindowCallRatePolicyModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + NestedRecordExtractor as NestedRecordExtractorModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( NoAuth as NoAuthModel, ) @@ -390,6 +395,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ParametrizedComponentsResolver as ParametrizedComponentsResolverModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + ParentFieldPath as ParentFieldPathModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ParentStreamConfig as ParentStreamConfigModel, ) @@ -786,6 +794,8 @@ def _init_mappings(self) -> None: DefaultPaginatorModel: self.create_default_paginator, DpathExtractorModel: self.create_dpath_extractor, DpathValidatorModel: self.create_dpath_validator, + NestedRecordExtractorModel: self.create_nested_record_extractor, + ParentFieldPathModel: self.create_parent_field_path, ResponseToFileExtractorModel: self.create_response_to_file_extractor, ExponentialBackoffStrategyModel: self.create_exponential_backoff_strategy, SessionTokenAuthenticatorModel: self.create_session_token_authenticator, @@ -2422,13 +2432,34 @@ def create_default_error_handler( parameters=model.parameters or {}, ) + @staticmethod + def _extractor_model_for_page_size( + extractor_model: Optional[ + Union[CustomRecordExtractorModel, DpathExtractorModel, NestedRecordExtractorModel] + ], + ) -> Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]]: + """ + Resolve the extractor a pagination strategy should count with. + + `OffsetIncrement` and `PageIncrement` advance by the number of records the response held, so + the count has to be over the items the API's `offset`/`page` addresses. A + `NestedRecordExtractor` emits one record per element of a nested collection, so counting with + it over-advances the offset and skips records; the collection the API paginates is the one + its innermost `parent_extractor` reads. + """ + while isinstance(extractor_model, NestedRecordExtractorModel): + extractor_model = extractor_model.parent_extractor + return extractor_model + def create_default_paginator( self, model: DefaultPaginatorModel, config: Config, *, url_base: str, - extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None, + extractor_model: Optional[ + Union[CustomRecordExtractorModel, DpathExtractorModel, NestedRecordExtractorModel] + ] = None, decoder: Optional[Decoder] = None, cursor_used_for_stop_condition: Optional[Cursor] = None, ) -> Union[DefaultPaginator, PaginatorTestReadDecorator]: @@ -2453,7 +2484,7 @@ def create_default_paginator( model=model.pagination_strategy, config=config, decoder=decoder_to_use, - extractor_model=extractor_model, + extractor_model=self._extractor_model_for_page_size(extractor_model), ) if cursor_used_for_stop_condition: pagination_strategy = StopConditionPaginationStrategyDecorator( @@ -2500,6 +2531,70 @@ def create_dpath_extractor( record_expander=record_expander, ) + def create_nested_record_extractor( + self, + model: NestedRecordExtractorModel, + config: Config, + decoder: Optional[Decoder] = None, + **kwargs: Any, + ) -> NestedRecordExtractor: + # The decoder is forwarded so the parent extractor reads the response the same way it would + # if it were the stream's extractor. This component never reads the response itself. + parent_extractor = self._create_component_from_model( + model=model.parent_extractor, + config=config, + decoder=decoder, + ) + self._reject_wildcard_path(model.child_field_path, "child_field_path") + model_child_field_path: List[Union[InterpolatedString, str]] = [ + x for x in model.child_field_path + ] + parent_fields = [ + self.create_parent_field_path( + model=parent_field, config=config, parameters=model.parameters or {} + ) + for parent_field in model.parent_fields or [] + ] + + return NestedRecordExtractor( + parent_extractor=parent_extractor, + child_field_path=model_child_field_path, + parent_fields=parent_fields, + config=config, + parameters=model.parameters or {}, + ) + + @staticmethod + def _reject_wildcard_path(path: List[str], field_name: str) -> None: + """ + Reject a literal `*` in a `NestedRecordExtractor` path, which the schema documents as + unsupported. Rejecting at parse time is what `lazy_read_pointer` does, and it turns a + data-dependent read-time failure into an error raised before the first request. A `*` that + arrives through interpolation is caught by the component itself. + """ + if any("*" in segment for segment in path): + raise ValueError( + f"The '*' wildcard in a NestedRecordExtractor '{field_name}' is not supported — only direct paths are allowed. " + f"Flatten across several collections in the 'parent_extractor' instead." + ) + + @staticmethod + def create_parent_field_path( + model: ParentFieldPathModel, + config: Config, + parameters: Optional[Mapping[str, Any]] = None, + **kwargs: Any, + ) -> ParentFieldPath: + ModelToComponentFactory._reject_wildcard_path(model.parent_path, "parent_path") + ModelToComponentFactory._reject_wildcard_path(model.record_path, "record_path") + model_parent_path: List[Union[InterpolatedString, str]] = [x for x in model.parent_path] + model_record_path: List[Union[InterpolatedString, str]] = [x for x in model.record_path] + return ParentFieldPath( + parent_path=model_parent_path, + record_path=model_record_path, + parameters=parameters or {}, + ) + def create_record_expander( self, model: RecordExpanderModel, @@ -3072,7 +3167,9 @@ def create_offset_increment( model: OffsetIncrementModel, config: Config, decoder: Decoder, - extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None, + extractor_model: Optional[ + Union[CustomRecordExtractorModel, DpathExtractorModel, NestedRecordExtractorModel] + ] = None, **kwargs: Any, ) -> OffsetIncrement: if isinstance(decoder, PaginationDecoderDecorator): @@ -3121,7 +3218,9 @@ def create_page_increment( model: PageIncrementModel, config: Config, decoder: Optional[Decoder] = None, - extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None, + extractor_model: Optional[ + Union[CustomRecordExtractorModel, DpathExtractorModel, NestedRecordExtractorModel] + ] = None, **kwargs: Any, ) -> PageIncrement: # Like OffsetIncrement, we instantiate a separate extractor with identical behavior to the diff --git a/unit_tests/sources/declarative/extractors/test_nested_record_extractor.py b/unit_tests/sources/declarative/extractors/test_nested_record_extractor.py new file mode 100644 index 0000000000..5c12915860 --- /dev/null +++ b/unit_tests/sources/declarative/extractors/test_nested_record_extractor.py @@ -0,0 +1,899 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# +import io +import json +import logging +from types import MappingProxyType +from typing import Any, Dict, Iterable, List, Mapping, MutableMapping, Union + +import pytest +import requests + +from airbyte_cdk.models import FailureType +from airbyte_cdk.sources.declarative.extractors.dpath_extractor import DpathExtractor +from airbyte_cdk.sources.declarative.extractors.nested_record_extractor import ( + NestedRecordExtractor, + ParentFieldPath, +) +from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor +from airbyte_cdk.utils.traced_exception import AirbyteTracedException + +config = {"child_connection": "reviews"} + + +def create_response(body: Union[Dict, List, bytes]) -> requests.Response: + response = requests.Response() + response.raw = io.BytesIO(body if isinstance(body, bytes) else json.dumps(body).encode("utf-8")) + return response + + +def dpath(*field_path: str, parameters: Mapping[str, Any] = {}) -> DpathExtractor: + return DpathExtractor(field_path=list(field_path), config=config, parameters=parameters) + + +def nested( + parent_extractor: RecordExtractor, + child_field_path: List[str], + parent_fields: Union[List[Dict[str, List[str]]], None] = None, + parameters: Mapping[str, Any] = {}, +) -> NestedRecordExtractor: + return NestedRecordExtractor( + parent_extractor=parent_extractor, + child_field_path=child_field_path, + parent_fields=[ + ParentFieldPath( + parent_path=parent_field["parent_path"], + record_path=parent_field["record_path"], + parameters=parameters, + ) + for parent_field in parent_fields or [] + ], + config=config, + parameters=parameters, + ) + + +class CountingExtractor(RecordExtractor): + """Parent extractor that records how many parents have been pulled so far.""" + + def __init__(self, parents: List[Mapping[str, Any]]) -> None: + self._parents = parents + self.yielded = 0 + + def extract_records(self, response: requests.Response) -> Iterable[MutableMapping[Any, Any]]: + for parent in self._parents: + self.yielded += 1 + yield parent # type: ignore[misc] + + +def test_parent_as_list(): + response = create_response({"data": [{"id": "p1", "children": [{"id": "c1"}, {"id": "c2"}]}]}) + extractor = nested(dpath("data"), ["children"]) + + assert list(extractor.extract_records(response)) == [{"id": "c1"}, {"id": "c2"}] + + +def test_parent_as_single_object(): + # The shape the GraphQL drill-down document returns: one parent, not a connection. + response = create_response({"data": {"id": "p1", "children": [{"id": "c1"}]}}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["id"], "record_path": ["pid"]}] + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "pid": "p1"}] + + +def test_child_as_single_object(): + response = create_response({"data": {"id": "p1", "child": {"id": "c1"}}}) + extractor = nested(dpath("data"), ["child"], [{"parent_path": ["id"], "record_path": ["pid"]}]) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "pid": "p1"}] + + +@pytest.mark.parametrize( + "parent, child_field_path", + [ + pytest.param({"id": "p1"}, ["children"], id="child_field_path_absent"), + pytest.param({"id": "p1", "children": None}, ["children"], id="child_collection_null"), + pytest.param({"id": "p1", "children": []}, ["children"], id="child_collection_empty"), + pytest.param( + {"id": "p1", "children": {}}, ["children"], id="child_collection_empty_object" + ), + pytest.param( + {"id": "p1", "children": {"nodes": None}}, + ["children", "nodes"], + id="nested_collection_null", + ), + pytest.param( + {"id": "p1", "children": {"nodes": []}}, + ["children", "nodes"], + id="nested_collection_empty", + ), + pytest.param({"id": "p1", "children": "not-a-collection"}, ["children"], id="scalar"), + pytest.param({"id": "p1"}, ["children", "nodes"], id="intermediate_node_absent"), + ], +) +def test_unresolvable_or_empty_child_field_path_yields_nothing(parent, child_field_path): + response = create_response({"data": [parent]}) + extractor = nested(dpath("data"), child_field_path) + + assert list(extractor.extract_records(response)) == [] + + +def test_nested_child_field_path(): + response = create_response( + {"data": [{"id": "p1", "children": {"nodes": [{"id": "c1"}, {"id": "c2"}]}}]} + ) + extractor = nested(dpath("data"), ["children", "nodes"]) + + assert list(extractor.extract_records(response)) == [{"id": "c1"}, {"id": "c2"}] + + +def test_non_object_children_are_skipped(): + # A scalar cannot carry the parent fields, so emitting it would produce a record that silently + # lacks them. + response = create_response({"data": [{"id": "p1", "children": ["scalar", {"id": "c1"}, 7]}]}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["id"], "record_path": ["pid"]}] + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "pid": "p1"}] + + +def test_missing_parent_path_copies_none(): + # Documented choice: copy None, matching `parent.get(source)` in the source-github extractor, so + # every record of the stream carries the field and the record shape stays stable. + response = create_response({"data": [{"id": "p1", "children": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["url"], "record_path": ["parent_url"]}] + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "parent_url": None}] + + +def test_missing_nested_parent_path_copies_none(): + response = create_response({"data": [{"id": "p1", "author": {}, "children": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), + ["children"], + [{"parent_path": ["author", "login"], "record_path": ["parent_login"]}], + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "parent_login": None}] + + +def test_existing_key_on_child_is_overwritten(): + # Documented choice: the parent is authoritative for the fields the manifest names. Keeping the + # child's value would hide a misconfigured `record_path`. + response = create_response( + {"data": [{"id": "parent", "children": [{"id": "child", "url": "child-url"}]}]} + ) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["id"], "record_path": ["url"]}] + ) + + assert list(extractor.extract_records(response)) == [{"id": "child", "url": "parent"}] + + +def test_multiple_parent_fields(): + response = create_response( + { + "data": [ + { + "id": "p1", + "url": "https://example.com/p1", + "author": {"login": "octocat"}, + "children": [{"id": "c1"}], + } + ] + } + ) + extractor = nested( + dpath("data"), + ["children"], + [ + {"parent_path": ["id"], "record_path": ["parent_id"]}, + {"parent_path": ["url"], "record_path": ["parent_url"]}, + {"parent_path": ["author", "login"], "record_path": ["parent", "login"]}, + ], + ) + + assert list(extractor.extract_records(response)) == [ + { + "id": "c1", + "parent_id": "p1", + "parent_url": "https://example.com/p1", + "parent": {"login": "octocat"}, + } + ] + + +def test_record_path_creates_intermediate_objects(): + response = create_response({"data": [{"id": "p1", "children": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["id"], "record_path": ["a", "b", "c"]}] + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "a": {"b": {"c": "p1"}}}] + + +def test_paths_are_interpolated(): + response = create_response({"data": [{"id": "p1", "reviews": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), + ["{{ config['child_connection'] }}"], + [{"parent_path": ["{{ parameters['parent_field'] }}"], "record_path": ["parent_id"]}], + parameters={"parent_field": "id"}, + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "parent_id": "p1"}] + + +def test_children_are_mutated_in_place(): + # Deliberate, and documented on the component: the child dicts come out of the body this + # extractor just decoded, so copying them would only cost memory. Pinned so nobody "fixes" it + # into a copy without noticing they changed the contract. + child = {"id": "c1"} + parent_extractor = CountingExtractor([{"id": "p1", "children": [child]}]) + extractor = nested( + parent_extractor, ["children"], [{"parent_path": ["id"], "record_path": ["pid"]}] + ) + + records = list(extractor.extract_records(create_response({}))) + + assert records[0] is child + assert child == {"id": "c1", "pid": "p1"} + + +def test_is_lazy_and_does_not_materialise_children(): + parents = [{"id": f"p{index}", "children": [{"id": f"c{index}"}]} for index in range(5)] + parent_extractor = CountingExtractor(parents) + extractor = nested(parent_extractor, ["children"]) + + records = extractor.extract_records(create_response({})) + + assert parent_extractor.yielded == 0 + assert next(iter(records)) == {"id": "c0"} + # Exactly one parent was pulled to produce the first record. A parent whose children were + # materialised, or an implementation that built the full result before returning, would have + # drained all five. + assert parent_extractor.yielded == 1 + + +def test_no_mutable_state_on_the_component(): + # One instance is shared by every partition of a stream and the partitions are read + # concurrently, so two interleaved reads must not see each other. + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["id"], "record_path": ["pid"]}] + ) + first = extractor.extract_records( + create_response({"data": [{"id": "p1", "children": [{"id": "c1"}, {"id": "c2"}]}]}) + ) + second = extractor.extract_records( + create_response({"data": [{"id": "p2", "children": [{"id": "c3"}]}]}) + ) + + interleaved = [next(first), next(second), next(first)] + + assert interleaved == [ + {"id": "c1", "pid": "p1"}, + {"id": "c3", "pid": "p2"}, + {"id": "c2", "pid": "p1"}, + ] + + +def test_response_is_read_once(): + class SingleUseExtractor(RecordExtractor): + def __init__(self) -> None: + self.reads = 0 + + def extract_records( + self, response: requests.Response + ) -> Iterable[MutableMapping[Any, Any]]: + self.reads += 1 + yield from response.json()["data"] + + parent_extractor = SingleUseExtractor() + extractor = nested(parent_extractor, ["children"]) + + assert list( + extractor.extract_records(create_response({"data": [{"children": [{"id": "c1"}]}]})) + ) == [{"id": "c1"}] + assert parent_extractor.reads == 1 + + +def test_empty_child_field_path_is_rejected(): + with pytest.raises(ValueError, match="non-empty `child_field_path`"): + nested(dpath("data"), []) + + +@pytest.mark.parametrize( + "parent_path, record_path, expected_message", + [ + pytest.param([], ["x"], "non-empty `parent_path`", id="empty_parent_path"), + pytest.param(["x"], [], "non-empty `record_path`", id="empty_record_path"), + ], +) +def test_empty_parent_field_paths_are_rejected(parent_path, record_path, expected_message): + with pytest.raises(ValueError, match=expected_message): + ParentFieldPath(parent_path=parent_path, record_path=record_path, parameters={}) + + +def test_nested_parent_extractor(): + # Two levels without a third component type: the inner extractor's records are the outer + # extractor's parents. The outer level copies the grandparent field off the stamp the inner + # level left on the review, which is how fields from two ancestors reach the same record. + response = create_response( + { + "data": [ + { + "id": "p1", + "reviews": { + "nodes": [ + {"id": "r1", "comments": {"nodes": [{"id": "cm1"}, {"id": "cm2"}]}}, + {"id": "r2", "comments": {"nodes": []}}, + ] + }, + } + ] + } + ) + extractor = nested( + nested( + dpath("data"), + ["reviews", "nodes"], + [{"parent_path": ["id"], "record_path": ["pull_request_id"]}], + ), + ["comments", "nodes"], + [ + {"parent_path": ["id"], "record_path": ["review_id"]}, + {"parent_path": ["pull_request_id"], "record_path": ["pull_request_id"]}, + ], + ) + + assert list(extractor.extract_records(response)) == [ + {"id": "cm1", "review_id": "r1", "pull_request_id": "p1"}, + {"id": "cm2", "review_id": "r1", "pull_request_id": "p1"}, + ] + + +# --------------------------------------------------------------------------------------------- +# A realistic GraphQL document, end to end, and the deepest nesting the component supports. +# --------------------------------------------------------------------------------------------- + + +def _reaction(identifier: str) -> Dict[str, Any]: + return {"id": identifier, "content": "THUMBS_UP", "user": {"login": "octocat"}} + + +def _comment(identifier: str, *reaction_ids: str) -> Dict[str, Any]: + return { + "id": identifier, + "body": "a comment", + "reactions": { + "pageInfo": {"hasNextPage": False}, + "nodes": [_reaction(reaction_id) for reaction_id in reaction_ids], + }, + } + + +def _review(identifier: str, *comments: Dict[str, Any]) -> Dict[str, Any]: + return { + "id": identifier, + "comments": {"pageInfo": {"hasNextPage": False}, "nodes": list(comments)}, + } + + +def _pull_request(number: int, *reviews: Dict[str, Any]) -> Dict[str, Any]: + return { + "number": number, + "url": f"https://github.com/airbytehq/airbyte/pull/{number}", + "reviews": {"pageInfo": {"hasNextPage": False}, "nodes": list(reviews)}, + } + + +def test_graphql_connection_document(): + # A connection: `nodes` is a list, and the sibling `pageInfo` is not a record. + response = create_response( + { + "data": { + "repository": { + "pullRequests": { + "pageInfo": {"hasNextPage": False, "endCursor": "Y3Vy"}, + "nodes": [ + { + "number": 1, + "url": "https://github.com/airbytehq/airbyte/pull/1", + "reviews": { + "pageInfo": {"hasNextPage": False}, + "nodes": [ + {"id": "PRR_1", "state": "APPROVED"}, + {"id": "PRR_2", "state": "COMMENTED"}, + ], + }, + }, + { + "number": 2, + "url": "https://github.com/airbytehq/airbyte/pull/2", + "reviews": {"pageInfo": {"hasNextPage": False}, "nodes": []}, + }, + ], + } + } + } + } + ) + extractor = nested( + dpath("data", "repository", "pullRequests", "nodes"), + ["reviews", "nodes"], + [{"parent_path": ["url"], "record_path": ["pull_request_url"]}], + ) + + assert list(extractor.extract_records(response)) == [ + { + "id": "PRR_1", + "state": "APPROVED", + "pull_request_url": "https://github.com/airbytehq/airbyte/pull/1", + }, + { + "id": "PRR_2", + "state": "COMMENTED", + "pull_request_url": "https://github.com/airbytehq/airbyte/pull/1", + }, + ] + + +def test_three_levels_of_nesting(): + # `...pullRequests.nodes[*].reviews.nodes[*].comments.nodes[*].reactions.nodes[*]`. Three + # `NestedRecordExtractor`s reach the comment, and only the last hop keeps a parent in scope. + response = create_response( + { + "data": { + "repository": { + "pullRequests": { + "pageInfo": {"hasNextPage": False}, + "nodes": [ + _pull_request( + 1, + _review("PRR_1", _comment("PRRC_1", "RE_1", "RE_2")), + _review("PRR_2", _comment("PRRC_2")), + ), + _pull_request(2), + ], + } + } + } + } + ) + extractor = nested( + nested( + nested(dpath("data", "repository", "pullRequests", "nodes"), ["reviews", "nodes"]), + ["comments", "nodes"], + ), + ["reactions", "nodes"], + [{"parent_path": ["id"], "record_path": ["comment_id"]}], + ) + + assert [ + (record["id"], record["comment_id"]) for record in extractor.extract_records(response) + ] == [ + ("RE_1", "PRRC_1"), + ("RE_2", "PRRC_1"), + ] + + +# --------------------------------------------------------------------------------------------- +# Copied values are never shared between records. +# --------------------------------------------------------------------------------------------- + + +def test_copied_object_is_not_shared_between_records(): + # Sharing one object across the collection means a later transformation writing into it rewrites + # records that were already emitted — they are queued unserialized downstream. + response = create_response( + { + "data": [ + { + "author": {"login": "octocat"}, + "children": [{"id": "c1"}, {"id": "c2"}, {"id": "c3"}], + } + ] + } + ) + extractor = nested( + dpath("data"), + ["children"], + [{"parent_path": ["author"], "record_path": ["author"]}], + ) + + records = list(extractor.extract_records(response)) + assert records[0]["author"] is not records[1]["author"] + + # What a downstream `AddFields` writing `["author", "child_id"]` would do. + for record in records: + record["author"]["child_id"] = record["id"] + assert [record["author"]["child_id"] for record in records] == ["c1", "c2", "c3"] + + +def test_copied_list_is_not_shared_between_records(): + response = create_response( + {"data": [{"labels": ["bug"], "children": [{"id": "c1"}, {"id": "c2"}]}]} + ) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["labels"], "record_path": ["labels"]}] + ) + + records = list(extractor.extract_records(response)) + assert records[0]["labels"] is not records[1]["labels"] + + records[0]["labels"].append("added") + assert records[1]["labels"] == ["bug"] + + +def test_copied_scalar_is_not_copied(): + # Scalars are immutable, so every child shares the parent's object. Nothing can observe the + # difference, and copying per child would be the extractor's per-record cost for nothing. + response = create_response({"data": [{"url": "u", "children": [{"id": "c1"}, {"id": "c2"}]}]}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["url"], "record_path": ["url"]}] + ) + + records = list(extractor.extract_records(response)) + assert records[0]["url"] is records[1]["url"] + + +def test_copied_object_is_not_shared_with_the_parent(): + parent = {"author": {"login": "octocat"}, "children": [{"id": "c1"}]} + extractor = nested( + CountingExtractor([parent]), + ["children"], + [{"parent_path": ["author"], "record_path": ["author"]}], + ) + + record = next(iter(extractor.extract_records(create_response({})))) + record["author"]["login"] = "rewritten" + assert parent["author"] == {"login": "octocat"} + + +# --------------------------------------------------------------------------------------------- +# The parent is read once per parent, not once per child. +# --------------------------------------------------------------------------------------------- + + +class RecordingMapping(Dict[str, Any]): + """A record that counts how many times each of its keys has been read.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.reads: List[str] = [] + + def __getitem__(self, key: str) -> Any: + self.reads.append(key) + return super().__getitem__(key) + + +def test_parent_is_read_once_per_parent(): + # Reading the parent per child makes the extractor quadratic in the size of the collection, + # because each read walks a parent that contains that collection. + parent = RecordingMapping( + {"id": "p1", "children": [{"id": f"c{index}"} for index in range(25)]} + ) + extractor = nested( + CountingExtractor([parent]), + ["children"], + [ + {"parent_path": ["id"], "record_path": ["parent_id"]}, + {"parent_path": ["id"], "record_path": ["also_parent_id"]}, + ], + ) + + records = list(extractor.extract_records(create_response({}))) + + assert len(records) == 25 + # One read per `parent_fields` entry, plus the one that resolves `child_field_path`. + assert parent.reads == ["id", "id", "children"] + + +# --------------------------------------------------------------------------------------------- +# `record_path` writes, including over a segment that is already taken. +# --------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "child,expected", + [ + pytest.param({"i": 1}, {"i": 1, "parent": {"id": "p1"}}, id="absent"), + pytest.param( + {"parent": {"other": 1}}, + {"parent": {"other": 1, "id": "p1"}}, + id="existing_object_is_merged_into", + ), + pytest.param({"parent": 7}, {"parent": {"id": "p1"}}, id="scalar_is_replaced"), + pytest.param({"parent": [{"a": 1}]}, {"parent": {"id": "p1"}}, id="list_is_replaced"), + pytest.param({"parent": None}, {"parent": {"id": "p1"}}, id="null_is_replaced"), + ], +) +def test_nested_record_path_over_an_existing_key(child, expected): + # The documented "a key already present on the child is overwritten" holds for a nested + # `record_path` too: an intermediate that is not an object is replaced by one. + response = create_response({"data": [{"id": "p1", "children": [child]}]}) + extractor = nested( + dpath("data"), + ["children"], + [{"parent_path": ["id"], "record_path": ["parent", "id"]}], + ) + + assert list(extractor.extract_records(response)) == [expected] + + +def test_numeric_record_path_segment_creates_an_object(): + # Every segment is an object key. A numeric segment creating a list would emit an array where + # the declared schema says object. + response = create_response({"data": [{"id": "p1", "children": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["id"], "record_path": ["hist", "0"]}] + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "hist": {"0": "p1"}}] + + +# --------------------------------------------------------------------------------------------- +# Paths that cannot address a field are rejected rather than read as field names. +# --------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "child_field_path,parent_fields", + [ + pytest.param(["a", "*", "n"], None, id="child_field_path"), + pytest.param( + ["children"], [{"parent_path": ["a", "*"], "record_path": ["v"]}], id="parent_path" + ), + pytest.param( + ["children"], [{"parent_path": ["a"], "record_path": ["*"]}], id="record_path" + ), + ], +) +def test_wildcard_in_a_path_is_rejected(child_field_path, parent_fields): + # The wildcard fails on the data rather than the manifest: a body with one match under `a` + # would work, and the next page with two would not. It is rejected outright instead. + response = create_response( + {"data": [{"a": {"x": {"n": 1}}, "children": [{"id": "c1"}]}]}, + ) + extractor = nested(dpath("data"), child_field_path, parent_fields) + + with pytest.raises(AirbyteTracedException) as exception_info: + list(extractor.extract_records(response)) + + assert exception_info.value.failure_type == FailureType.config_error + assert "'*' wildcard" in str(exception_info.value.internal_message) + + +def test_other_glob_characters_are_literal_field_names(): + # Unlike the `dpath`-based extractors, this component walks paths itself, so `?` and `[...]` + # address the fields that spell them rather than matching several. + response = create_response({"data": [{"a?b": "v1", "c[1]": "v2", "children": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), + ["children"], + [ + {"parent_path": ["a?b"], "record_path": ["q"]}, + {"parent_path": ["c[1]"], "record_path": ["b"]}, + ], + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "q": "v1", "b": "v2"}] + + +@pytest.mark.parametrize( + "child_field_path,parent_fields,expected_field", + [ + pytest.param(["{{ parameters.missing }}"], None, "child_field_path", id="child_field_path"), + pytest.param( + ["children"], + [{"parent_path": ["{{ parameters.missing }}"], "record_path": ["v"]}], + "parent_path", + id="parent_path", + ), + pytest.param( + ["children"], + [{"parent_path": ["id"], "record_path": ["{{ parameters.missing }}"]}], + "record_path", + id="record_path", + ), + ], +) +def test_segment_that_resolves_to_nothing_is_rejected( + child_field_path, parent_fields, expected_field +): + # An unbound parameter evaluates to "", and writing a field literally named "" is rejected by + # most destinations with an error that names the destination rather than the manifest. + response = create_response({"data": [{"id": "p1", "children": [{"id": "c1"}]}]}) + extractor = nested(dpath("data"), child_field_path, parent_fields) + + with pytest.raises(AirbyteTracedException) as exception_info: + list(extractor.extract_records(response)) + + assert exception_info.value.failure_type == FailureType.config_error + assert expected_field in str(exception_info.value.internal_message) + + +def test_parent_path_addressing_the_child_collection_is_rejected(): + # The child records are not copied out of the parent, so such a path would give every record a + # snapshot of its whole sibling collection. + response = create_response({"data": [{"children": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["children"], "record_path": ["ctx"]}] + ) + + with pytest.raises(AirbyteTracedException) as exception_info: + list(extractor.extract_records(response)) + + assert exception_info.value.failure_type == FailureType.config_error + + +def test_parent_path_addressing_an_ancestor_of_the_child_collection_is_rejected(): + response = create_response({"data": [{"a": {"children": [{"id": "c1"}]}}]}) + extractor = nested( + dpath("data"), ["a", "children"], [{"parent_path": ["a"], "record_path": ["ctx"]}] + ) + + with pytest.raises(AirbyteTracedException) as exception_info: + list(extractor.extract_records(response)) + + assert exception_info.value.failure_type == FailureType.config_error + + +# --------------------------------------------------------------------------------------------- +# Reading a `parent_path` that does not match the response. +# --------------------------------------------------------------------------------------------- + + +def test_parent_path_through_a_value_holding_no_fields_raises(): + # Copying `None` here would be indistinguishable from the documented absent-key case, and a + # stamped field is a legitimate primary key or cursor. + response = create_response({"data": [{"tags": [{"name": "t1"}], "children": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["tags", "name"], "record_path": ["tag"]}] + ) + + with pytest.raises(AirbyteTracedException) as exception_info: + list(extractor.extract_records(response)) + + assert exception_info.value.failure_type == FailureType.system_error + assert "tags" in str(exception_info.value.internal_message) + + +def test_parent_path_through_a_null_copies_none(): + # A null is data, not a misconfigured path, so it is treated as the absent-key case. + response = create_response({"data": [{"a": None, "children": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["a", "b"], "record_path": ["v"]}] + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "v": None}] + + +@pytest.mark.parametrize( + "parent_path,expected", + [ + pytest.param(["tags", "0", "name"], "t1", id="index"), + pytest.param(["tags", "-1", "name"], "t2", id="negative_index"), + ], +) +def test_list_positions_can_be_addressed(parent_path, expected): + response = create_response( + {"data": [{"tags": [{"name": "t1"}, {"name": "t2"}], "children": [{"id": "c1"}]}]} + ) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": parent_path, "record_path": ["tag"]}] + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "tag": expected}] + + +def test_out_of_range_list_position_copies_none(): + response = create_response({"data": [{"tags": [], "children": [{"id": "c1"}]}]}) + extractor = nested( + dpath("data"), ["children"], [{"parent_path": ["tags", "0"], "record_path": ["tag"]}] + ) + + assert list(extractor.extract_records(response)) == [{"id": "c1", "tag": None}] + + +# --------------------------------------------------------------------------------------------- +# The `RecordExtractor` contract is `Mapping`, not `MutableMapping`. +# --------------------------------------------------------------------------------------------- + + +class ReadOnlyExtractor(RecordExtractor): + """A parent extractor honouring the declared `Iterable[Mapping[str, Any]]` return type.""" + + def __init__(self, parents: List[Mapping[str, Any]]) -> None: + self._parents = parents + + def extract_records(self, response: requests.Response) -> Iterable[Mapping[str, Any]]: + yield from self._parents + + +@pytest.mark.parametrize( + "parent", + [ + pytest.param( + MappingProxyType({"id": "p1", "children": [{"id": "c1"}]}), id="read_only_parent" + ), + pytest.param( + {"id": "p1", "children": [MappingProxyType({"id": "c1"})]}, id="read_only_child" + ), + pytest.param( + MappingProxyType({"id": "p1", "children": [MappingProxyType({"id": "c1"})]}), + id="read_only_throughout", + ), + ], +) +def test_read_only_mappings_are_not_dropped(parent): + # The ABC declares `Iterable[Mapping[str, Any]]`, and the schema admits a + # `CustomRecordExtractor` as a `parent_extractor`, so a custom extractor honouring the declared + # contract must not silently produce an empty stream. + extractor = nested( + ReadOnlyExtractor([parent]), + ["children"], + [{"parent_path": ["id"], "record_path": ["parent_id"]}], + ) + + records = list(extractor.extract_records(create_response({}))) + + assert records == [{"id": "c1", "parent_id": "p1"}] + # Downstream transformations mutate records in place, so what is yielded has to be mutable. + assert all(isinstance(record, MutableMapping) for record in records) + + +def test_read_only_child_is_copied_rather_than_mutated(): + child = MappingProxyType({"id": "c1"}) + extractor = nested( + ReadOnlyExtractor([{"id": "p1", "children": [child]}]), + ["children"], + [{"parent_path": ["id"], "record_path": ["parent_id"]}], + ) + + assert list(extractor.extract_records(create_response({}))) == [{"id": "c1", "parent_id": "p1"}] + assert dict(child) == {"id": "c1"} + + +# --------------------------------------------------------------------------------------------- +# Every skip path leaves a trace. +# --------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "body, child_field_path, expected_log", + [ + pytest.param( + {"data": ["scalar", {"id": "p1", "children": [{"id": "c1"}]}]}, + ["children"], + "skipped a parent record that is not an object (str)", + id="non_object_parent", + ), + pytest.param( + {"data": [{"id": "p1", "children": [7, {"id": "c1"}]}]}, + ["children"], + "skipped a child element that is not an object (int)", + id="non_object_child", + ), + pytest.param( + {"data": [{"id": "p1", "children": "not-a-collection"}]}, + ["children"], + "to a str, which holds no records", + id="child_field_path_resolves_to_a_scalar", + ), + ], +) +def test_skipped_records_are_logged(body, child_field_path, expected_log, caplog): + # A connector whose API shape shifts should not lose records with no signal at all. + extractor = nested(dpath("data"), child_field_path) + + with caplog.at_level(logging.DEBUG, logger="airbyte"): + list(extractor.extract_records(create_response(body))) + + assert expected_log in caplog.text diff --git a/unit_tests/sources/declarative/extractors/test_record_selector.py b/unit_tests/sources/declarative/extractors/test_record_selector.py index 2d32bc75cc..4c14f36f2f 100644 --- a/unit_tests/sources/declarative/extractors/test_record_selector.py +++ b/unit_tests/sources/declarative/extractors/test_record_selector.py @@ -10,9 +10,17 @@ from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder from airbyte_cdk.sources.declarative.extractors.dpath_extractor import DpathExtractor +from airbyte_cdk.sources.declarative.extractors.nested_record_extractor import ( + NestedRecordExtractor, + ParentFieldPath, +) from airbyte_cdk.sources.declarative.extractors.record_filter import RecordFilter from airbyte_cdk.sources.declarative.extractors.record_selector import RecordSelector from airbyte_cdk.sources.declarative.transformations import RecordTransformation +from airbyte_cdk.sources.declarative.transformations.add_fields import ( + AddedFieldDefinition, + AddFields, +) from airbyte_cdk.sources.types import Record, StreamSlice from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer @@ -312,3 +320,65 @@ def transformation_side_effect(record, config, stream_state, stream_slice): else: assert final_record_data[0]["id"] == 2 assert final_record_data[0]["myfield"] == 999 + + +def test_nested_record_extractor_survives_a_transformation_that_writes_into_a_copied_field(): + """A parent field copied onto several records must not be one object shared by all of them. + + `AddFields` mutates the record in place, and records are queued unserialized downstream, so a + shared object means writing into one record rewrites records that were already emitted. + """ + body = { + "data": [ + { + "author": {"login": "octocat"}, + "children": [{"id": "c1"}, {"id": "c2"}, {"id": "c3"}], + } + ] + } + response = requests.Response() + response._content = json.dumps(body).encode("utf-8") + + record_selector = RecordSelector( + extractor=NestedRecordExtractor( + parent_extractor=DpathExtractor(field_path=["data"], config={}, parameters={}), + child_field_path=["children"], + parent_fields=[ + ParentFieldPath(parent_path=["author"], record_path=["author"], parameters={}) + ], + config={}, + parameters={}, + ), + record_filter=None, + transformations=[ + AddFields( + fields=[ + AddedFieldDefinition( + path=["author", "child_id"], + value="{{ record['id'] }}", + value_type=str, + parameters={}, + ) + ], + parameters={}, + ) + ], + schema_normalization=TypeTransformer(TransformConfig.NoTransform), + config={}, + parameters={}, + ) + + records = list( + record_selector.select_records( + response=response, + stream_state={}, + records_schema={}, + stream_slice=StreamSlice(partition={}, cursor_slice={}), + ) + ) + + assert [(record["id"], record["author"]["child_id"]) for record in records] == [ + ("c1", "c1"), + ("c2", "c2"), + ("c3", "c3"), + ] 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..bb56b02a36 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 @@ -53,8 +53,17 @@ from airbyte_cdk.sources.declarative.checks import CheckStream 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 import ( + IterableDecoder, + JsonDecoder, + PaginationDecoderDecorator, +) +from airbyte_cdk.sources.declarative.extractors import ( + DpathExtractor, + NestedRecordExtractor, + RecordFilter, + RecordSelector, +) from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor from airbyte_cdk.sources.declarative.extractors.record_filter import ( ClientSideIncrementalRecordFilterDecorator, @@ -2032,6 +2041,302 @@ def test_create_record_selector(test_name, record_selector, expected_runtime_sel assert selector.record_filter.condition == "{{ record['id'] > stream_state['id'] }}" +def test_create_nested_record_extractor(): + content = """ + selector: + type: RecordSelector + $parameters: + child_connection: "reviews" + extractor: + type: NestedRecordExtractor + parent_extractor: + type: DpathExtractor + field_path: ["data", "repository", "pullRequests", "nodes"] + child_field_path: ["{{ parameters['child_connection'] }}", "nodes"] + parent_fields: + - parent_path: ["url"] + record_path: ["pull_request_url"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + 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, NestedRecordExtractor) + assert isinstance(extractor.parent_extractor, DpathExtractor) + assert [fp.eval(input_config) for fp in extractor.parent_extractor._field_path] == [ + "data", + "repository", + "pullRequests", + "nodes", + ] + # `$parameters` reaches the paths, so one YAML block can serve several streams. + assert [fp.eval(input_config) for fp in extractor._child_field_path] == ["reviews", "nodes"] + assert len(extractor.parent_fields) == 1 + assert extractor.parent_fields[0].eval_parent_path(input_config) == ["url"] + assert extractor.parent_fields[0].eval_record_path(input_config) == ["pull_request_url"] + + +def test_create_nested_record_extractor_nested_in_itself(): + content = """ + selector: + type: RecordSelector + extractor: + type: NestedRecordExtractor + parent_extractor: + type: NestedRecordExtractor + parent_extractor: + type: DpathExtractor + field_path: ["data", "repository", "pullRequests", "nodes"] + child_field_path: ["reviews", "nodes"] + child_field_path: ["comments", "nodes"] + parent_fields: + - parent_path: ["id"] + record_path: ["review_id"] + """ + 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, + ) + + extractor = selector.extractor + assert isinstance(extractor, NestedRecordExtractor) + assert isinstance(extractor.parent_extractor, NestedRecordExtractor) + assert isinstance(extractor.parent_extractor.parent_extractor, DpathExtractor) + + +def test_create_nested_record_extractor_rejects_empty_child_field_path(): + content = """ + selector: + type: RecordSelector + extractor: + type: NestedRecordExtractor + parent_extractor: + type: DpathExtractor + field_path: ["data"] + child_field_path: [] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + # Rejected at config time by the schema's `minItems: 1`, before any record is read. The + # component also guards its own constructor, see + # unit_tests/sources/declarative/extractors/test_nested_record_extractor.py. + with pytest.raises(ValidationError, match="child_field_path"): + factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + +def test_nested_record_extractor_forwards_the_decoder_to_its_parent_extractor(): + content = """ + selector: + type: RecordSelector + extractor: + type: NestedRecordExtractor + parent_extractor: + type: DpathExtractor + field_path: ["data"] + child_field_path: ["children"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + decoder = IterableDecoder(parameters={}) + + selector = factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=decoder, + transformations=[], + config=input_config, + ) + + assert selector.extractor.parent_extractor.decoder is decoder + + +@pytest.mark.parametrize( + "child_field_path, parent_fields, expected_field", + [ + pytest.param('["a", "*", "n"]', "", "child_field_path", id="child_field_path"), + pytest.param( + '["children"]', + """ + parent_fields: + - parent_path: ["a", "*"] + record_path: ["v"]""", + "parent_path", + id="parent_path", + ), + pytest.param( + '["children"]', + """ + parent_fields: + - parent_path: ["a"] + record_path: ["*"]""", + "record_path", + id="record_path", + ), + ], +) +def test_create_nested_record_extractor_rejects_the_wildcard( + child_field_path, parent_fields, expected_field +): + content = f""" + selector: + type: RecordSelector + extractor: + type: NestedRecordExtractor + parent_extractor: + type: DpathExtractor + field_path: ["data"] + child_field_path: {child_field_path}{parent_fields} + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + # A wildcard fails on the data rather than the manifest — one match resolves and two do not — + # so it is rejected before the first request instead. + with pytest.raises(ValueError, match=expected_field): + factory.create_component( + model_type=RecordSelectorModel, + name="test_stream", + component_definition=selector_manifest, + decoder=None, + transformations=[], + config=input_config, + ) + + +def test_nested_record_extractor_paginates_over_parents_not_children(): + content = """ + retriever: + type: SimpleRetriever + requester: + type: HttpRequester + url_base: "https://example.com" + path: "/" + record_selector: + type: RecordSelector + extractor: + type: NestedRecordExtractor + parent_extractor: + type: NestedRecordExtractor + parent_extractor: + type: DpathExtractor + field_path: ["data"] + child_field_path: ["reviews"] + child_field_path: ["comments"] + paginator: + type: DefaultPaginator + pagination_strategy: + type: OffsetIncrement + page_size: 2 + page_token_option: + type: RequestOption + field_name: "offset" + inject_into: "request_parameter" + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + retriever_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["retriever"], {} + ) + + retriever = factory.create_component( + model_type=SimpleRetrieverModel, + component_definition=retriever_manifest, + config=input_config, + name="test_stream", + primary_key=None, + stream_slicer=None, + request_options_provider=None, + transformations=[], + ) + + # `OffsetIncrement` advances by the number of records the response held, so it has to count the + # items the API's `offset` addresses. Counting with the `NestedRecordExtractor` counts the + # nested collection instead, which over-advances the offset and skips records. + strategy = retriever.paginator.pagination_strategy + assert isinstance(strategy.extractor, DpathExtractor) + + response = requests.Response() + response.raw = io.BytesIO( + json.dumps( + { + "data": [ + {"reviews": [{"comments": [{"id": 1}, {"id": 2}, {"id": 3}]}]}, + {"reviews": [{"comments": [{"id": 4}]}]}, + ] + } + ).encode("utf-8") + ) + + assert ( + strategy.next_page_token( + response=response, last_page_size=0, last_record=None, last_page_token_value=0 + ) + == 2 + ) + + +def test_nested_record_extractor_accepts_a_class_name_only_parent_extractor(): + content = """ + selector: + type: RecordSelector + extractor: + type: NestedRecordExtractor + parent_extractor: + class_name: unit_tests.sources.declarative.parsers.testing_components.TestingSomeComponent + child_field_path: ["children"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + selector_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["selector"], {} + ) + + # `RecordSelector.extractor` is in `CUSTOM_COMPONENTS_MAPPING`, so a custom extractor declared + # with only a `class_name` works as a stream extractor. It has to keep working when wrapped. + assert selector_manifest["extractor"]["parent_extractor"]["type"] == "CustomRecordExtractor" + + @pytest.mark.parametrize( "test_name, error_handler, expected_backoff_strategy_type", [