airbyte_cdk.sources.declarative.expanders

45class OnNoRecords(Enum):
46    """
47    Behavior when record expansion produces no records.
48    """
49
50    skip = "skip"
51    emit_parent = "emit_parent"

 54@dataclass
 55class RecordExpander:
 56    """Expands records by extracting items from a nested array field.
 57
 58    When configured, this component extracts items from a specified nested array path
 59    within each record and emits each item as a separate record. Set `remain_original_record: true`
 60    to embed the full parent record under `original_record` in each expanded item when you need
 61    downstream transformations to access parent context.
 62
 63    The expand_records_from_field path supports wildcards (*) for matching multiple arrays.
 64    When wildcards are used, items from all matched arrays are extracted and emitted.
 65
 66    Examples of instantiating this component:
 67    ```
 68      record_expander:
 69        type: RecordExpander
 70        expand_records_from_field:
 71          - "lines"
 72          - "data"
 73        remain_original_record: true
 74    ```
 75
 76    ```
 77      record_expander:
 78        type: RecordExpander
 79        expand_records_from_field:
 80          - "sections"
 81          - "*"
 82          - "items"
 83        on_no_records: emit_parent
 84    ```
 85
 86    Attributes:
 87        expand_records_from_field: Path to a nested array field within each record.
 88            Items from this array will be extracted and emitted as separate records.
 89            Supports wildcards (*).
 90        remain_original_record: If True, each expanded record will include the original
 91            parent record in an "original_record" field. Defaults to False.
 92        on_no_records: Behavior when expansion produces no records. "skip" (default)
 93            emits nothing. "emit_parent" emits the original parent record unchanged.
 94        truncation_indicator_path: Path within each record to a field indicating that the
 95            embedded nested list is truncated (e.g. a `has_more` flag on the list object).
 96            When the indicator is truthy and no `truncated_list_retriever` is configured, the
 97            embedded items are expanded as normal and a WARNING is logged (once per stream
 98            instance) describing the expansion path and the embedded item count, so that the
 99            data loss is visible instead of silent. Glob metacharacters (`*`, `?`, `[`) are
100            rejected in this path, and in `expand_records_from_field` when a retriever is
101            configured; the check runs on the interpolated values.
102        truncated_list_retriever: Retriever used to fetch the complete list of items when
103            the field at `truncation_indicator_path` is truthy. The record being expanded is
104            exposed to the retriever's interpolation context as `stream_slice['parent_record']`.
105            One fetch is issued per truncated parent record; enable `use_cache` on its requester
106            when the same list can be fetched repeatedly. Without a `paginator` only the first
107            page is read. If the retriever returns no records, the embedded items are expanded
108            as a fallback; if it returns fewer records than the `total_count` field next to the
109            indicator, a WARNING is logged once per stream instance. Request failures surface
110            through the retriever's error handler and fail the stream like any other request.
111            `$parameters` of the enclosing stream propagate into this retriever's components.
112            In Connector Builder test reads the page limit applies to each fetch independently,
113            so the fetched list may be shorter than `total_count`; no incomplete-fetch warning
114            is emitted there.
115        message_repository: Optional repository through which the truncation warnings are emitted
116            as Airbyte LOG messages so they are visible in the Connector Builder. When it is not
117            set, the warnings go to the `airbyte` logger instead.
118        suppress_incomplete_fetch_warning: Skip the incomplete-fetch WARNING. Set by the factory
119            for Connector Builder test reads when the page limit caps the retriever's paginator,
120            so a shortfall against `total_count` is expected. The truncated-without-retriever
121            warning is not affected.
122        config: The user-provided configuration as specified by the source's spec.
123    """
124
125    expand_records_from_field: Sequence[str]
126    config: Config
127    parameters: InitVar[Mapping[str, Any]]
128    remain_original_record: bool = False
129    on_no_records: OnNoRecords = OnNoRecords.skip
130    truncation_indicator_path: Optional[Sequence[str]] = None
131    truncated_list_retriever: Optional["Retriever"] = None
132    message_repository: Optional[MessageRepository] = None
133    suppress_incomplete_fetch_warning: bool = False
134
135    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
136        self._expand_path: list[InterpolatedString] = [
137            InterpolatedString.create(path, parameters=parameters)
138            for path in self.expand_records_from_field
139        ]
140        if self.truncated_list_retriever and not self.truncation_indicator_path:
141            raise ValueError(
142                "`truncation_indicator_path` is required when `truncated_list_retriever` is configured."
143            )
144        self._truncation_indicator_path: list[InterpolatedString] = [
145            InterpolatedString.create(path, parameters=parameters)
146            for path in (self.truncation_indicator_path or [])
147        ]
148        # The paths only interpolate `config`, so their evaluated values are validated up front.
149        self._reject_globs(self._evaluated_indicator_path(), "truncation_indicator_path")
150        if self.truncated_list_retriever:
151            self._reject_globs(self._evaluated_expand_path(), "expand_records_from_field")
152        self._warning_lock = threading.Lock()
153        self._warned_truncation_without_retriever = False
154        self._warned_incomplete_fetch = False
155
156    @staticmethod
157    def _reject_globs(path: Sequence[Any], field_name: str) -> None:
158        if any(
159            isinstance(segment, str) and any(char in segment for char in _GLOB_METACHARACTERS)
160            for segment in path
161        ):
162            raise ValueError(
163                f"Glob characters {_GLOB_METACHARACTERS} are not supported in `{field_name}` when "
164                "truncation handling is configured: the path must identify a single field."
165            )
166
167    def _evaluated_indicator_path(self) -> list[Any]:
168        return [segment.eval(self.config) for segment in self._truncation_indicator_path]
169
170    def _evaluated_expand_path(self) -> list[Any]:
171        return [segment.eval(self.config) for segment in self._expand_path]
172
173    def expand_record(self, record: MutableMapping[Any, Any]) -> Iterable[MutableMapping[Any, Any]]:
174        """Expand a record by extracting items from a nested array field."""
175        if not isinstance(record, Mapping):
176            # If the input isn't a mapping, expansion can't proceed; yield as-is.
177            yield record
178            return
179
180        if not self._expand_path:
181            yield record
182            return
183
184        parent_record = record
185
186        expand_path = self._evaluated_expand_path()
187        truncated = bool(self._truncation_indicator_path) and self._is_truncated(parent_record)
188
189        if truncated and self.truncated_list_retriever:
190            # Streamed, so the shortfall is only known once the retriever is exhausted. If the
191            # consumer stops early the fetch was cut short by it, not by the API, and no warning
192            # would be accurate anyway.
193            fetched_count = 0
194            for fetched in self._fetch_complete_list(parent_record):
195                fetched_count += 1
196                yield fetched
197            self._warn_if_fetch_incomplete(parent_record, expand_path, fetched_count)
198            if fetched_count > 0:
199                return
200
201        try:
202            extracted_values = dpath.values(parent_record, expand_path)
203        except KeyError:
204            extracted_values = []
205
206        embedded_lists = [
207            extracted for extracted in extracted_values if isinstance(extracted, list)
208        ]
209        embedded_count = sum(len(items) for items in embedded_lists)
210
211        if truncated and not self.truncated_list_retriever:
212            self._warn_truncated_without_retriever(parent_record, expand_path, embedded_count)
213
214        for items in embedded_lists:
215            for item in items:
216                if isinstance(item, dict):
217                    expanded_record = dict(item)
218                    self._apply_parent_context(parent_record, expanded_record)
219                    yield expanded_record
220                elif self.remain_original_record:
221                    yield {
222                        "value": item,
223                        "original_record": copy.deepcopy(parent_record),
224                    }
225                else:
226                    yield item
227
228        if embedded_count == 0 and self.on_no_records == OnNoRecords.emit_parent:
229            yield parent_record
230
231    def _warn_truncated_without_retriever(
232        self, parent_record: Mapping[str, Any], expand_path: list[Any], embedded_count: int
233    ) -> None:
234        with self._warning_lock:
235            if self._warned_truncation_without_retriever:
236                return
237            self._warned_truncation_without_retriever = True
238
239        indicator_path = self._evaluated_indicator_path()
240        total_count = self._get_sibling_total_count(parent_record, indicator_path)
241        total_fragment = f" of {total_count} total" if total_count is not None else ""
242        self._emit_warning(
243            f"The nested list at {expand_path} is marked as truncated (the field at {indicator_path} "
244            f"is truthy) but no `truncated_list_retriever` is configured, so only the "
245            f"{embedded_count} embedded item(s){total_fragment} were expanded and the remaining "
246            "items are not emitted. Configure `truncated_list_retriever` to fetch the complete list "
247            "if the API provides an endpoint for it. This warning is emitted once per stream; other "
248            "records may be truncated as well."
249        )
250
251    def _warn_if_fetch_incomplete(
252        self, parent_record: Mapping[str, Any], expand_path: list[Any], fetched_count: int
253    ) -> None:
254        if self.suppress_incomplete_fetch_warning:
255            return
256        indicator_path = self._evaluated_indicator_path()
257        total_count = self._get_sibling_total_count(parent_record, indicator_path)
258        if total_count is None or fetched_count >= total_count:
259            return
260        with self._warning_lock:
261            if self._warned_incomplete_fetch:
262                return
263            self._warned_incomplete_fetch = True
264        fallback_fragment = (
265            " The embedded items were expanded as a fallback." if fetched_count == 0 else ""
266        )
267        self._emit_warning(
268            f"The `truncated_list_retriever` for the nested list at {expand_path} returned "
269            f"{fetched_count} record(s) but the `total_count` field next to {indicator_path} reports "
270            f"{total_count}, so the fetched list is still incomplete.{fallback_fragment} Check that "
271            "the retriever has a `paginator` configured and that its request matches the "
272            "complete-list endpoint. This warning is emitted once per stream; other records may be "
273            "affected as well."
274        )
275
276    def _emit_warning(self, message: str) -> None:
277        if self.message_repository:
278            self.message_repository.emit_message(
279                AirbyteMessage(
280                    type=MessageType.LOG,
281                    log=AirbyteLogMessage(level=Level.WARN, message=message),
282                )
283            )
284        else:
285            logger.warning(message)
286
287    def _get_sibling_total_count(
288        self, parent_record: Mapping[str, Any], indicator_path: list[Any]
289    ) -> Optional[int]:
290        """Best-effort lookup of a `total_count` field next to the truncation indicator."""
291        if not indicator_path:
292            return None
293        try:
294            total = dpath.get(dict(parent_record), [*indicator_path[:-1], "total_count"])
295        except (KeyError, ValueError):
296            return None
297        return total if isinstance(total, int) and not isinstance(total, bool) else None
298
299    def _is_truncated(self, parent_record: MutableMapping[Any, Any]) -> bool:
300        indicator_path = self._evaluated_indicator_path()
301        try:
302            return bool(dpath.get(parent_record, indicator_path))
303        except (KeyError, ValueError):
304            return False
305
306    def _fetch_complete_list(self, parent_record: Mapping[str, Any]) -> Iterable[Any]:
307        if not self.truncated_list_retriever:
308            return
309        stream_slice = StreamSlice(partition={"parent_record": parent_record}, cursor_slice={})
310        for item in self.truncated_list_retriever.read_records(
311            records_schema={}, stream_slice=stream_slice
312        ):
313            if isinstance(item, AirbyteMessage):
314                if item.type != MessageType.RECORD or item.record is None:
315                    continue
316                data: Any = item.record.data
317            elif isinstance(item, Record):
318                data = item.data
319            elif isinstance(item, _BARE_PROTOCOL_MESSAGES):
320                continue
321            else:
322                data = item
323            if isinstance(data, Mapping):
324                expanded_record = dict(data)
325                self._apply_parent_context(parent_record, expanded_record)
326                yield expanded_record
327            elif self.remain_original_record:
328                yield {"value": data, "original_record": copy.deepcopy(parent_record)}
329            else:
330                yield data
331
332    def _apply_parent_context(
333        self, parent_record: Mapping[str, Any], child_record: MutableMapping[str, Any]
334    ) -> None:
335        """Apply parent context to a child record."""
336        if self.remain_original_record:
337            child_record["original_record"] = copy.deepcopy(parent_record)