From 41737736e6d68bbf4c71a03bf5b0a0dfa9ad5da6 Mon Sep 17 00:00:00 2001 From: "zane.hyatt" Date: Mon, 28 Sep 2026 06:19:18 +0000 Subject: [PATCH 1/6] feat(low-code): add JsonItemsDecoder `spool_to_disk` with retry-safe body download in HttpClient Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../concurrent_declarative_source.py | 35 +++++--- .../declarative_component_schema.yaml | 10 +++ .../decoders/composite_raw_decoder.py | 5 ++ .../sources/declarative/decoders/decoder.py | 6 ++ .../models/declarative_component_schema.py | 5 ++ .../parsers/model_to_component_factory.py | 28 +++++++ .../declarative/requesters/http_requester.py | 2 + .../sources/streams/http/http_client.py | 56 +++++++++++++ .../sources/streams/http/streamed_response.py | 84 +++++++++++++++++++ 9 files changed, 219 insertions(+), 12 deletions(-) create mode 100644 airbyte_cdk/sources/streams/http/streamed_response.py diff --git a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py index 52f83a644d..14bc041bb2 100644 --- a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py +++ b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py @@ -107,6 +107,17 @@ from airbyte_cdk.utils.traced_exception import AirbyteTracedException +def _decoder_spools_to_disk(decoder_definition: Optional[Dict[str, Any]]) -> bool: + """Whether a resolved decoder manifest dict enables spool_to_disk (GzipDecoder recurses).""" + if not isinstance(decoder_definition, dict): + return False + if decoder_definition.get("type") == "JsonItemsDecoder": + return bool(decoder_definition.get("spool_to_disk")) + if decoder_definition.get("type") == "GzipDecoder": + return _decoder_spools_to_disk(decoder_definition.get("decoder")) + return False + + @dataclass class TestLimits: __test__: ClassVar[bool] = False # Tell Pytest this is not a Pytest class, despite its name @@ -530,8 +541,12 @@ def _initialize_cache_for_parent_streams( """ parent_streams = set() - def _set_cache_if_not_disabled(requester: Dict[str, Any]) -> None: + def _set_cache_if_not_disabled(retriever: Dict[str, Any]) -> None: """Set use_cache to True only if not explicitly disabled.""" + if _decoder_spools_to_disk(retriever.get("decoder")): + # requests_cache consumes the whole body, leaving nothing to spool + return + requester = retriever["requester"] if requester.get("use_cache") is not False: requester["use_cache"] = True @@ -542,19 +557,19 @@ def update_with_cache_parent_configs( parent_streams.add(parent_config["stream"]["name"]) if parent_config["stream"]["type"] == "StateDelegatingStream": _set_cache_if_not_disabled( - parent_config["stream"]["full_refresh_stream"]["retriever"]["requester"] + parent_config["stream"]["full_refresh_stream"]["retriever"] ) _set_cache_if_not_disabled( - parent_config["stream"]["incremental_stream"]["retriever"]["requester"] + parent_config["stream"]["incremental_stream"]["retriever"] ) else: - _set_cache_if_not_disabled(parent_config["stream"]["retriever"]["requester"]) + _set_cache_if_not_disabled(parent_config["stream"]["retriever"]) for stream_config in stream_configs: if stream_config.get("incremental_sync", {}).get("parent_stream"): parent_streams.add(stream_config["incremental_sync"]["parent_stream"]["name"]) _set_cache_if_not_disabled( - stream_config["incremental_sync"]["parent_stream"]["retriever"]["requester"] + stream_config["incremental_sync"]["parent_stream"]["retriever"] ) elif stream_config.get("retriever", {}).get("partition_router", {}): @@ -581,14 +596,10 @@ def update_with_cache_parent_configs( for stream_config in stream_configs: if stream_config["name"] in parent_streams: if stream_config["type"] == "StateDelegatingStream": - _set_cache_if_not_disabled( - stream_config["full_refresh_stream"]["retriever"]["requester"] - ) - _set_cache_if_not_disabled( - stream_config["incremental_stream"]["retriever"]["requester"] - ) + _set_cache_if_not_disabled(stream_config["full_refresh_stream"]["retriever"]) + _set_cache_if_not_disabled(stream_config["incremental_stream"]["retriever"]) else: - _set_cache_if_not_disabled(stream_config["retriever"]["requester"]) + _set_cache_if_not_disabled(stream_config["retriever"]) return stream_configs def spec(self, logger: logging.Logger) -> ConnectorSpecification: diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 56d373afa2..e5c1a6c723 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2956,6 +2956,16 @@ definitions: description: Text encoding used to decode the streamed bytes before JSON parsing. type: string default: utf-8 + spool_to_disk: + title: Spool Response To Disk + description: >- + Download the whole response body to an anonymous temporary file before parsing it, instead of + parsing straight from the network socket. The download happens inside the request retry logic, so + a connection that breaks mid-body is retried exactly like a failed request and no record of the + page is emitted before the full page is on disk. Costs the decoded size of one page in temporary + disk per in-flight request; memory stays bounded. Recommended for very large pages consumed slowly. + type: boolean + default: false JsonlDecoder: title: JSON Lines description: Select 'JSON Lines' if the response consists of JSON objects separated by new lines ('\n') in JSONL format. diff --git a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py index e29f863c88..baab81df9b 100644 --- a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py +++ b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py @@ -242,6 +242,7 @@ def __init__( parser: Parser, stream_response: bool = True, parsers_by_header: PARSERS_BY_HEADER_TYPE = None, + spool_response: bool = False, ) -> None: # since we moved from using `dataclass` to `__init__` method, # we need to keep using the `parser` to be able to resolve the depenencies @@ -250,6 +251,7 @@ def __init__( self._parsers_by_header = parsers_by_header if parsers_by_header else {} self._stream_response = stream_response + self._spool_response = spool_response @classmethod def by_headers( @@ -278,6 +280,9 @@ def by_headers( def is_stream_response(self) -> bool: return self._stream_response + def spools_response(self) -> bool: + return self._spool_response + def decode(self, response: requests.Response) -> DECODER_OUTPUT_TYPE: parser = self._select_parser(response) if self.is_stream_response(): diff --git a/airbyte_cdk/sources/declarative/decoders/decoder.py b/airbyte_cdk/sources/declarative/decoders/decoder.py index 34d99db1f9..bbda5678d4 100644 --- a/airbyte_cdk/sources/declarative/decoders/decoder.py +++ b/airbyte_cdk/sources/declarative/decoders/decoder.py @@ -23,6 +23,12 @@ def is_stream_response(self) -> bool: Set to True if you'd like to use stream=True option in http requester """ + def spools_response(self) -> bool: + """ + Set to True if the full response body should be copied to disk before decoding. + """ + return False + @abstractmethod def decode(self, response: requests.Response) -> DECODER_OUTPUT_TYPE: """ diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index a62c3db2fe..20d5a02e25 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -846,6 +846,11 @@ class JsonItemsDecoder(BaseModel): description="The character encoding of the JSON data. Defaults to UTF-8.", title="Encoding", ) + spool_to_disk: Optional[bool] = Field( + False, + description="Download the whole response body to an anonymous temporary file before parsing it, instead of parsing straight from the network socket. The download happens inside the request retry logic, so a connection that breaks mid-body is retried exactly like a failed request and no record of the page is emitted before the full page is on disk. Costs the decoded size of one page in temporary disk per in-flight request; memory stays bounded. Recommended for very large pages consumed slowly.", + title="Spool Response To Disk", + ) class JsonlDecoder(BaseModel): diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index 8940f33fa3..a65acbb20f 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2648,6 +2648,30 @@ def create_http_requester( should_use_cache = (model.use_cache or bool(use_cache)) and not self._disable_cache + spools_response = decoder is not None and decoder.spools_response() + + if should_use_cache and spools_response: + raise ValueError( + f"Stream {name}: `use_cache` cannot be combined with a decoder configured with `spool_to_disk` " + f"({type(decoder).__name__}); requests_cache reads the whole body when storing a response, which " + "leaves nothing for the decoder to spool. " + "Set `use_cache: false` on the requester (including on parent streams, whose cache is enabled automatically)." + ) + + if ( + spools_response + and isinstance(model.error_handler, DefaultErrorHandlerModel) + and any( + f.predicate or f.error_message_contains + for f in model.error_handler.response_filters or [] + ) + ): + LOGGER.warning( + f"Stream {name}: `spool_to_disk` response bodies larger than 1 MiB are not " + "evaluated by body-based response filters (`predicate`, `error_message_contains`); " + "only `http_codes` apply to them." + ) + return HttpRequester( name=name, url=model.url, @@ -2665,6 +2689,7 @@ def create_http_requester( use_cache=should_use_cache, decoder=decoder, stream_response=decoder.is_stream_response() if decoder else False, + spool_response=spools_response, ) @staticmethod @@ -2808,6 +2833,7 @@ def create_csv_decoder(self, model: CsvDecoderModel, config: Config, **kwargs: A return CompositeRawDecoder( parser=ModelToComponentFactory._get_parser(model, config), stream_response=False if self._emit_connector_builder_messages else True, + spool_response=bool(model.spool_to_disk) and not self._emit_connector_builder_messages, ) def create_jsonl_decoder( @@ -2816,6 +2842,7 @@ def create_jsonl_decoder( return CompositeRawDecoder( parser=ModelToComponentFactory._get_parser(model, config), stream_response=False if self._emit_connector_builder_messages else True, + spool_response=bool(model.spool_to_disk) and not self._emit_connector_builder_messages, ) def create_json_items_decoder( @@ -2824,6 +2851,7 @@ def create_json_items_decoder( return CompositeRawDecoder( parser=ModelToComponentFactory._get_parser(model, config), stream_response=False if self._emit_connector_builder_messages else True, + spool_response=bool(model.spool_to_disk) and not self._emit_connector_builder_messages, ) def create_gzip_decoder( diff --git a/airbyte_cdk/sources/declarative/requesters/http_requester.py b/airbyte_cdk/sources/declarative/requesters/http_requester.py index 3ce4c85408..952afd5a01 100644 --- a/airbyte_cdk/sources/declarative/requesters/http_requester.py +++ b/airbyte_cdk/sources/declarative/requesters/http_requester.py @@ -69,6 +69,7 @@ class HttpRequester(Requester): _exit_on_rate_limit: bool = False stream_response: bool = False decoder: Decoder = field(default_factory=lambda: JsonDecoder(parameters={})) + spool_response: bool = False def __post_init__(self, parameters: Mapping[str, Any]) -> None: self._url = InterpolatedString.create( @@ -464,6 +465,7 @@ def send_request( next_page_token=next_page_token, ), request_kwargs={"stream": self.stream_response}, + spool_response=self.spool_response, headers=self._request_headers( stream_state, stream_slice, next_page_token, request_headers ), diff --git a/airbyte_cdk/sources/streams/http/http_client.py b/airbyte_cdk/sources/streams/http/http_client.py index f79276a7aa..c88f598587 100644 --- a/airbyte_cdk/sources/streams/http/http_client.py +++ b/airbyte_cdk/sources/streams/http/http_client.py @@ -4,6 +4,8 @@ import logging import os +import tempfile +import time import urllib from pathlib import Path from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union @@ -61,6 +63,11 @@ ResponseAwareAuthenticator, TokenRotatingAuthenticator, ) +from airbyte_cdk.sources.streams.http.streamed_response import ( + SpooledResponseBody, + mark_body_streamed, + set_spooled_body_size, +) from airbyte_cdk.sources.utils.types import JsonType from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH @@ -76,6 +83,11 @@ BODY_REQUEST_METHODS = ("GET", "POST", "PUT", "PATCH") +_SPOOL_CHUNK_SIZE = 1 << 20 +_SPOOL_IN_MEMORY_LIMIT = 8 << 20 +# Subclasses may override _send/_send_with_retry, so the spool option travels inside request_kwargs. +_SPOOL_RESPONSE_KWARG = "_airbyte_spool_response" + def monkey_patched_get_item(self, key): # type: ignore # this interface is a copy/paste from the requests_cache lib """ @@ -416,6 +428,37 @@ def _update_authenticator_from_response( "Authenticator failed to update quota state from response", exc_info=True ) + def _spool_response_body( + self, response: requests.Response, request: requests.PreparedRequest + ) -> None: + start = time.monotonic() + spool = tempfile.SpooledTemporaryFile( + max_size=_SPOOL_IN_MEMORY_LIMIT, prefix="airbyte-http-body-" + ) + size = 0 + try: + for chunk in response.iter_content(chunk_size=_SPOOL_CHUNK_SIZE): + spool.write(chunk) + size += len(chunk) + except BaseException: + spool.close() + response.close() # release the (broken) connection + raise + response.close() # body fully read: release the connection to the pool now + spool.seek(0) + set_spooled_body_size(response, size) + response.raw = SpooledResponseBody(spool) + self._logger.debug( + "Spooled response body to disk", + extra={ + "url": request.url, + "status": response.status_code, + "bytes": size, + "seconds": round(time.monotonic() - start, 3), + }, + ) + response._content_consumed = False # type: ignore[attr-defined] # iter_content set it True; the body is readable again from disk + def _send( self, request: requests.PreparedRequest, @@ -423,6 +466,8 @@ def _send( log_formatter: Optional[Callable[[requests.Response], Any]] = None, exit_on_rate_limit: Optional[bool] = False, ) -> requests.Response: + request_kwargs = dict(request_kwargs) + spool_response = bool(request_kwargs.pop(_SPOOL_RESPONSE_KWARG, False)) if request not in self._request_attempt_count: self._request_attempt_count[request] = 1 else: @@ -440,7 +485,12 @@ def _send( try: response = self._session.send(request, **request_kwargs) + if request_kwargs.get("stream"): + mark_body_streamed(response) + if spool_response: + self._spool_response_body(response, request) except requests.RequestException as e: + response = None exc = e if response is not None: @@ -449,6 +499,9 @@ def _send( error_resolution: ErrorResolution = self._error_handler.interpret_response( response if response is not None else exc ) + if response is not None and isinstance(getattr(response, "raw", None), SpooledResponseBody): + # body filters may have consumed the spooled body; rewind it for the decoder + response.raw.seek(0) # Evaluation of response.text can be heavy, for example, if streaming a large response # Do it only in debug mode @@ -817,6 +870,7 @@ def send_request( dedupe_query_params: bool = False, log_formatter: Optional[Callable[[requests.Response], Any]] = None, exit_on_rate_limit: Optional[bool] = False, + spool_response: bool = False, ) -> Tuple[requests.PreparedRequest, requests.Response]: """ Prepares and sends request and return request and response objects. @@ -840,6 +894,8 @@ def send_request( cert=request_kwargs.get("cert"), ) request_kwargs = {**request_kwargs, **env_settings} + if spool_response: + request_kwargs[_SPOOL_RESPONSE_KWARG] = True response: requests.Response = self._send_with_retry( request=request, diff --git a/airbyte_cdk/sources/streams/http/streamed_response.py b/airbyte_cdk/sources/streams/http/streamed_response.py new file mode 100644 index 0000000000..91d9f48067 --- /dev/null +++ b/airbyte_cdk/sources/streams/http/streamed_response.py @@ -0,0 +1,84 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import io +from typing import IO, Any, Optional + +import requests + +_BODY_STREAMED_ATTR = "_airbyte_body_streamed" +_DOCUMENT_REMAINDER_ATTR = "_airbyte_document_remainder" +_SPOOLED_BODY_SIZE_ATTR = "_airbyte_spooled_body_size" + + +def mark_body_streamed(response: requests.Response) -> None: + response.__dict__[_BODY_STREAMED_ATTR] = True + + +def is_body_streamed(response: requests.Response) -> bool: + return bool(response.__dict__.get(_BODY_STREAMED_ATTR, False)) + + +def set_document_remainder(response: requests.Response, remainder: Any) -> None: + response.__dict__[_DOCUMENT_REMAINDER_ATTR] = remainder + + +def get_document_remainder(response: requests.Response) -> Optional[Any]: + return response.__dict__.get(_DOCUMENT_REMAINDER_ATTR) + + +def set_spooled_body_size(response: requests.Response, size: int) -> None: + response.__dict__[_SPOOLED_BODY_SIZE_ATTR] = size + + +def get_spooled_body_size(response: requests.Response) -> Optional[int]: + return response.__dict__.get(_SPOOLED_BODY_SIZE_ATTR) + + +class _FileLikeRaw(io.RawIOBase): + """Raw adapter over a SpooledTemporaryFile so BufferedReader can wrap it without fileno().""" + + def __init__(self, fileobj: IO[bytes]) -> None: + self._fileobj = fileobj + + def readable(self) -> bool: + return True + + def seekable(self) -> bool: + return True + + def readinto(self, b: Any) -> int: + # SpooledTemporaryFile has no readinto on python 3.10 + data = self._fileobj.read(len(b)) + n = len(data) + b[:n] = data + return n + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + return self._fileobj.seek(offset, whence) + + def tell(self) -> int: + return self._fileobj.tell() + + def close(self) -> None: + try: + super().close() + finally: + self._fileobj.close() + + +class SpooledResponseBody(io.BufferedReader): + """File-backed replacement for `requests.Response.raw` once the body has been copied to disk.""" + + def __init__(self, spool: IO[bytes]) -> None: + # spool is a tempfile.SpooledTemporaryFile(); BytesIO below max_size, unlinked file above + self._spool = spool + super().__init__(_FileLikeRaw(spool)) + self.auto_close = False # CompositeRawDecoder sets this attribute; keep it settable + + def close(self) -> None: + try: + super().close() + finally: + self._spool.close() # last fd -> kernel frees the unlinked file From 789d68f39ba810c16457e741f8594dd92d21f0b1 Mon Sep 17 00:00:00 2001 From: "zane.hyatt" Date: Mon, 28 Sep 2026 06:19:29 +0000 Subject: [PATCH 2/6] feat(low-code): paginate JsonItemsDecoder responses from the captured document remainder Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../decoders/composite_raw_decoder.py | 62 ++++++++++++++++--- .../decoders/pagination_decoder_decorator.py | 11 +++- .../parsers/model_to_component_factory.py | 17 ++--- 3 files changed, 73 insertions(+), 17 deletions(-) diff --git a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py index baab81df9b..02b1aea3b8 100644 --- a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py +++ b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py @@ -8,9 +8,10 @@ import io import json import logging +from collections.abc import Mapping from dataclasses import dataclass from io import BufferedIOBase, TextIOWrapper -from typing import Any, List, Optional +from typing import Any, Callable, Dict, Generator, List, Optional, Tuple import ijson import orjson @@ -25,6 +26,7 @@ PARSERS_TYPE, Parser, ) +from airbyte_cdk.sources.streams.http.streamed_response import set_document_remainder from airbyte_cdk.utils import AirbyteTracedException logger = logging.getLogger("airbyte") @@ -63,7 +65,7 @@ def readinto(self, buffer: Buffer) -> int: class GzipParser(Parser): inner_parser: Parser - def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE: + def parse(self, data: BufferedIOBase, **kwargs: Any) -> PARSER_OUTPUT_TYPE: """Decompress gzipped data or pass uncompressed data through unchanged. Args: @@ -82,9 +84,9 @@ def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE: if prefix == b"\x1f\x8b": with gzip.GzipFile(fileobj=prefixed_data, mode="rb") as gzipobj: - yield from self.inner_parser.parse(gzipobj) + yield from self.inner_parser.parse(gzipobj, **kwargs) else: - yield from self.inner_parser.parse(prefixed_data) + yield from self.inner_parser.parse(prefixed_data, **kwargs) @dataclass @@ -176,7 +178,12 @@ class JsonItemsParser(Parser): items_path: str = "" encoding: Optional[str] = "utf-8" - def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE: + def parse( + self, + data: BufferedIOBase, + *, + on_document_remainder: Optional[Callable[[Any], None]] = None, + ) -> PARSER_OUTPUT_TYPE: if not self.items_path: raise ValueError("JsonItemsParser requires a non-empty items_path.") if self.encoding and codecs.lookup(self.encoding).name != "utf-8": @@ -190,7 +197,20 @@ def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE: # whole stream up front. # use_float=True yields floats for non-integer numbers instead of Decimal, matching # json.loads/orjson behavior so downstream JSON serialization doesn't choke on Decimal. - yield from ijson.items(data, f"{self.items_path}.item", use_float=True) + items_prefix = f"{self.items_path}.item" + if on_document_remainder is None: + yield from ijson.items(data, items_prefix, use_float=True) + return + builder = ijson.common.ObjectBuilder() + + def events() -> Generator[Tuple[str, str, Any], None, None]: + for prefix, event, value in ijson.parse(data, use_float=True): + if not (prefix == items_prefix or prefix.startswith(items_prefix + ".")): + builder.event(event, value) + yield prefix, event, value + + yield from ijson.items(events(), items_prefix) + on_document_remainder(builder.value) @dataclass @@ -222,6 +242,14 @@ def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE: yield row +def _parser_supports_document_remainder(parser: Parser) -> bool: + if isinstance(parser, JsonItemsParser): + return True + if isinstance(parser, GzipParser): + return _parser_supports_document_remainder(parser.inner_parser) + return False + + class CompositeRawDecoder(Decoder): """ Decoder strategy to transform a requests.Response into a PARSER_OUTPUT_TYPE @@ -252,6 +280,7 @@ def __init__( self._parsers_by_header = parsers_by_header if parsers_by_header else {} self._stream_response = stream_response self._spool_response = spool_response + self._capture_document_remainder = False @classmethod def by_headers( @@ -277,6 +306,10 @@ def by_headers( parsers_by_header[header] = {header_value: parser for header_value in header_values} return cls(fallback_parser, stream_response, parsers_by_header) + def enable_document_remainder_capture(self) -> None: + """Opt in to capturing the non-items document fields for pagination.""" + self._capture_document_remainder = True + def is_stream_response(self) -> bool: return self._stream_response @@ -285,6 +318,20 @@ def spools_response(self) -> bool: def decode(self, response: requests.Response) -> DECODER_OUTPUT_TYPE: parser = self._select_parser(response) + + parse_kwargs: Dict[str, Any] = {} + if self._capture_document_remainder and _parser_supports_document_remainder(parser): + + def on_document_remainder(doc: Any) -> None: + set_document_remainder(response, doc) + logger.debug( + "Captured pagination document remainder", + extra={ + "keys": list(doc.keys()) if isinstance(doc, Mapping) else type(doc).__name__ + }, + ) + + parse_kwargs["on_document_remainder"] = on_document_remainder if self.is_stream_response(): # urllib mentions that some interfaces don't play nice with auto_close # More info here: https://urllib3.readthedocs.io/en/stable/user-guide.html#using-io-wrappers-with-response-content @@ -293,10 +340,11 @@ def decode(self, response: requests.Response) -> DECODER_OUTPUT_TYPE: response.raw.auto_close = False yield from parser.parse( data=response.raw, # type: ignore[arg-type] + **parse_kwargs, ) response.raw.close() else: - yield from parser.parse(data=io.BytesIO(response.content)) + yield from parser.parse(data=io.BytesIO(response.content), **parse_kwargs) def _select_parser(self, response: requests.Response) -> Parser: """ diff --git a/airbyte_cdk/sources/declarative/decoders/pagination_decoder_decorator.py b/airbyte_cdk/sources/declarative/decoders/pagination_decoder_decorator.py index e5a1527110..a2741434e3 100644 --- a/airbyte_cdk/sources/declarative/decoders/pagination_decoder_decorator.py +++ b/airbyte_cdk/sources/declarative/decoders/pagination_decoder_decorator.py @@ -9,6 +9,7 @@ import requests from airbyte_cdk.sources.declarative.decoders import Decoder +from airbyte_cdk.sources.streams.http.streamed_response import get_document_remainder logger = logging.getLogger("airbyte") @@ -33,7 +34,13 @@ def decode( self, response: requests.Response ) -> Generator[MutableMapping[str, Any], None, None]: if self._decoder.is_stream_response(): - logger.warning("Response is streamed and therefore will not be decoded for pagination.") - yield {} + remainder = get_document_remainder(response) + if remainder is not None: + yield remainder + else: + logger.warning( + "Response is streamed and therefore will not be decoded for pagination." + ) + yield {} else: yield from self._decoder.decode(response) 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 a65acbb20f..d81c39439d 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2452,6 +2452,10 @@ def create_default_paginator( ) -> Union[DefaultPaginator, PaginatorTestReadDecorator]: if decoder: if self._is_supported_decoder_for_pagination(decoder): + if isinstance(decoder, CompositeRawDecoder): + # the retriever shares this decoder instance with the record selector, + # so opting in here makes decode() also capture the remainder + decoder.enable_document_remainder_capture() decoder_to_use = PaginationDecoderDecorator(decoder=decoder) else: raise ValueError(self._UNSUPPORTED_DECODER_ERROR.format(decoder_type=type(decoder))) @@ -2833,7 +2837,6 @@ def create_csv_decoder(self, model: CsvDecoderModel, config: Config, **kwargs: A return CompositeRawDecoder( parser=ModelToComponentFactory._get_parser(model, config), stream_response=False if self._emit_connector_builder_messages else True, - spool_response=bool(model.spool_to_disk) and not self._emit_connector_builder_messages, ) def create_jsonl_decoder( @@ -2842,7 +2845,6 @@ def create_jsonl_decoder( return CompositeRawDecoder( parser=ModelToComponentFactory._get_parser(model, config), stream_response=False if self._emit_connector_builder_messages else True, - spool_response=bool(model.spool_to_disk) and not self._emit_connector_builder_messages, ) def create_json_items_decoder( @@ -4921,8 +4923,8 @@ def create_parametrized_components_resolver( _UNSUPPORTED_DECODER_ERROR = ( "Specified decoder of {decoder_type} is not supported for pagination." - "Please set as `JsonDecoder`, `XmlDecoder`, or a `CompositeRawDecoder` with an inner_parser of `JsonParser` or `GzipParser` instead." - "If using `GzipParser`, please ensure that the lowest level inner_parser is a `JsonParser`." + "Please set as `JsonDecoder`, `XmlDecoder`, or a `CompositeRawDecoder` with an inner_parser of `JsonParser`, `JsonItemsParser`, or `GzipParser` instead." + "If using `GzipParser`, please ensure that the lowest level inner_parser is a `JsonParser` or `JsonItemsParser`." ) def _is_supported_decoder_for_pagination(self, decoder: Decoder) -> bool: @@ -4934,12 +4936,11 @@ def _is_supported_decoder_for_pagination(self, decoder: Decoder) -> bool: return False def _is_supported_parser_for_pagination(self, parser: Parser) -> bool: - if isinstance(parser, JsonParser): + if isinstance(parser, (JsonParser, JsonItemsParser)): return True elif isinstance(parser, GzipParser): - return isinstance(parser.inner_parser, JsonParser) - else: - return False + return isinstance(parser.inner_parser, (JsonParser, JsonItemsParser)) + return False def create_http_api_budget( self, model: HTTPAPIBudgetModel, config: Config, **kwargs: Any From b0652acbbf93a3c29564abac78ed999bc0a9bedd Mon Sep 17 00:00:00 2001 From: "zane.hyatt" Date: Mon, 28 Sep 2026 06:19:35 +0000 Subject: [PATCH 3/6] fix(low-code): skip body-based response filters only for spooled bodies over 1 MiB Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../error_handlers/http_response_filter.py | 30 +++++++++++++++---- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/airbyte_cdk/sources/declarative/requesters/error_handlers/http_response_filter.py b/airbyte_cdk/sources/declarative/requesters/error_handlers/http_response_filter.py index 866d6b7d25..9cd3fc45ef 100644 --- a/airbyte_cdk/sources/declarative/requesters/error_handlers/http_response_filter.py +++ b/airbyte_cdk/sources/declarative/requesters/error_handlers/http_response_filter.py @@ -18,8 +18,21 @@ ErrorResolution, ResponseAction, ) +from airbyte_cdk.sources.streams.http.streamed_response import ( + get_spooled_body_size, + is_body_streamed, +) from airbyte_cdk.sources.types import Config +_MAX_BODY_FILTER_SIZE = 1 << 20 + + +def _skips_body_filters(response: requests.Response) -> bool: + if not (response.ok and is_body_streamed(response)): + return False + size = get_spooled_body_size(response) + return size is not None and size > _MAX_BODY_FILTER_SIZE + @dataclass class HttpResponseFilter: @@ -130,10 +143,16 @@ def _matches_filter( :param response: The HTTP response to evaluate :return: The action to execute. None if the response does not match the filter """ - if isinstance(response_or_exception, requests.Response) and ( - response_or_exception.status_code in self.http_codes # type: ignore # http_codes set is always initialized to a value in __post_init__ - or self._response_matches_predicate(response_or_exception) - or self._response_contains_error_message(response_or_exception) + if not isinstance(response_or_exception, requests.Response): + return None + response = response_or_exception + if response.status_code in self.http_codes: # type: ignore # http_codes set is always initialized to a value in __post_init__ + return self.action # type: ignore # action is always cast to a ResponseAction not a str + if _skips_body_filters(response): + # body-based matching would consume a large spooled body the decoder has not read yet + return None + if self._response_matches_predicate(response) or self._response_contains_error_message( + response ): return self.action # type: ignore # action is always cast to a ResponseAction not a str return None @@ -151,8 +170,9 @@ def _create_error_message(self, response: requests.Response) -> Optional[str]: :param response: The HTTP response which can be used during interpolation :return: The evaluated error message string to be emitted """ + response_json = {} if _skips_body_filters(response) else self._safe_response_json(response) return self.error_message.eval( # type: ignore[no-any-return, union-attr] - self.config, response=self._safe_response_json(response), headers=response.headers + self.config, response=response_json, headers=response.headers ) def _response_matches_predicate(self, response: requests.Response) -> bool: From 539ba2e250590c8286b03648bec7382c21d1d8ce Mon Sep 17 00:00:00 2001 From: "zane.hyatt" Date: Mon, 28 Sep 2026 06:19:35 +0000 Subject: [PATCH 4/6] fix(low-code): decode transport encoding on the streamed decoder path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../sources/declarative/decoders/composite_raw_decoder.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py index 02b1aea3b8..7e88dbee48 100644 --- a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py +++ b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py @@ -338,6 +338,8 @@ def on_document_remainder(doc: Any) -> None: # We have indeed observed some issues with CSV parsing. # Hence, we will manage the closing of the file ourselves until we find a better solution. response.raw.auto_close = False + if hasattr(response.raw, "decode_content"): + response.raw.decode_content = True yield from parser.parse( data=response.raw, # type: ignore[arg-type] **parse_kwargs, From 03df3c38563aaa1f88a58833adba2b6651c7e596 Mon Sep 17 00:00:00 2001 From: "zane.hyatt" Date: Mon, 28 Sep 2026 06:19:35 +0000 Subject: [PATCH 5/6] test(low-code): cover spool_to_disk, JsonItems pagination, body-filter gating and compatibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../decoders/test_streamed_json_pagination.py | 222 +++++++++++ .../test_model_to_component_factory.py | 234 +++++++++++ .../test_http_response_filter.py | 99 +++++ .../requesters/test_http_requester.py | 35 ++ .../test_concurrent_declarative_source.py | 246 ++++++++++++ .../sources/streams/http/test_http_client.py | 362 ++++++++++++++++++ 6 files changed, 1198 insertions(+) create mode 100644 unit_tests/sources/declarative/decoders/test_streamed_json_pagination.py diff --git a/unit_tests/sources/declarative/decoders/test_streamed_json_pagination.py b/unit_tests/sources/declarative/decoders/test_streamed_json_pagination.py new file mode 100644 index 0000000000..9e20fa612b --- /dev/null +++ b/unit_tests/sources/declarative/decoders/test_streamed_json_pagination.py @@ -0,0 +1,222 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import io +import json + +import pytest +import requests +import urllib3 + +from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import ( + CompositeRawDecoder, + GzipParser, + JsonItemsParser, +) +from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder +from airbyte_cdk.sources.declarative.decoders.pagination_decoder_decorator import ( + PaginationDecoderDecorator, +) +from airbyte_cdk.sources.streams.http.streamed_response import ( + get_document_remainder, + set_document_remainder, +) + + +def _streamed_response(body: bytes, headers=None) -> requests.Response: + raw = urllib3.HTTPResponse( + body=io.BytesIO(body), + headers=headers or {"Content-Type": "application/json"}, + status=200, + preload_content=False, + ) + response = requests.Response() + response.status_code = 200 + response.raw = raw + response.url = "https://airbyte.io/" + return response + + +def test_json_items_parser_reports_document_remainder(): + body = b'{"a":1,"tickets":[{"id":1},{"id":2}],"after_url":"x","end_of_stream":false}' + remainders = [] + records = list( + JsonItemsParser(items_path="tickets").parse( + io.BytesIO(body), on_document_remainder=remainders.append + ) + ) + assert records == [{"id": 1}, {"id": 2}] + assert remainders == [{"a": 1, "tickets": [], "after_url": "x", "end_of_stream": False}] + + +def test_json_items_parser_remainder_nested_path(): + body = b'{"data":{"items":[{"id":1}],"meta":{"c":2}},"z":3}' + remainders = [] + records = list( + JsonItemsParser(items_path="data.items").parse( + io.BytesIO(body), on_document_remainder=remainders.append + ) + ) + assert records == [{"id": 1}] + assert remainders == [{"data": {"items": [], "meta": {"c": 2}}, "z": 3}] + + +def test_gzip_parser_forwards_remainder_kwarg(): + import gzip + + body = gzip.compress(b'{"data":[{"id":1}],"after_url":"x"}') + remainders = [] + records = list( + GzipParser(inner_parser=JsonItemsParser(items_path="data")).parse( + io.BytesIO(body), on_document_remainder=remainders.append + ) + ) + assert records == [{"id": 1}] + assert remainders == [{"data": [], "after_url": "x"}] + + +def test_json_items_parser_without_remainder_kwarg_unchanged(): + body = b'{"data":[{"id":1},{"id":2}]}' + assert list(JsonItemsParser(items_path="data").parse(io.BytesIO(body))) == [ + {"id": 1}, + {"id": 2}, + ] + + +def test_composite_raw_decoder_sets_remainder_on_streamed_response(): + body = json.dumps( + {"tickets": [{"id": 1}, {"id": 2}], "after_url": "x", "end_of_stream": False} + ).encode() + response = _streamed_response(body) + decoder = CompositeRawDecoder(parser=JsonItemsParser(items_path="tickets")) + decoder.enable_document_remainder_capture() + assert list(decoder.decode(response)) == [{"id": 1}, {"id": 2}] + assert get_document_remainder(response) == { + "tickets": [], + "after_url": "x", + "end_of_stream": False, + } + + +class _StreamingJsonDecoder(JsonDecoder): + def is_stream_response(self) -> bool: + return True + + +def test_pagination_decoder_decorator_yields_remainder_when_present(): + response = _streamed_response(b"{}") + set_document_remainder(response, {"after_url": "x", "end_of_stream": False}) + decorator = PaginationDecoderDecorator(decoder=_StreamingJsonDecoder(parameters={})) + assert next(decorator.decode(response)) == {"after_url": "x", "end_of_stream": False} + + +def test_pagination_decoder_decorator_yields_empty_and_warns_without_remainder(caplog): + response = _streamed_response(b"{}") + decorator = PaginationDecoderDecorator(decoder=_StreamingJsonDecoder(parameters={})) + with caplog.at_level("WARNING"): + assert next(decorator.decode(response)) == {} + assert "will not be decoded for pagination" in caplog.text + + +def test_pagination_decoder_decorator_delegates_when_not_streamed(): + body = b'{"data":[{"id":1}]}' + response = requests.Response() + response.status_code = 200 + response._content = body + decorator = PaginationDecoderDecorator(decoder=JsonDecoder(parameters={})) + assert next(decorator.decode(response)) == {"data": [{"id": 1}]} + + +import gzip +import tempfile + +from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import ( + CsvParser, + JsonLineParser, + JsonParser, +) +from airbyte_cdk.sources.streams.http.streamed_response import SpooledResponseBody + + +def _spooled_body(body: bytes, max_size: int) -> SpooledResponseBody: + spool = tempfile.SpooledTemporaryFile(max_size=max_size) + spool.write(body) + spool.seek(0) + return SpooledResponseBody(spool) + + +_JSON_BODY = json.dumps({"items": [{"id": 1}, {"id": 2}], "after_url": "x"}).encode() + + +@pytest.mark.parametrize( + "parser, body, expected", + [ + (JsonParser(), _JSON_BODY, [_JSON_BODY]), + (JsonItemsParser(items_path="items"), _JSON_BODY, [{"id": 1}, {"id": 2}]), + (JsonLineParser(), b'{"id":1}\n{"id":2}\n', [{"id": 1}, {"id": 2}]), + (CsvParser(), b"a,b\n1,2\n", [{"a": "1", "b": "2"}]), + ( + GzipParser(inner_parser=JsonItemsParser(items_path="items")), + gzip.compress(_JSON_BODY), + [{"id": 1}, {"id": 2}], + ), + ], + ids=["json", "json_items", "json_lines", "csv", "gzip_json_items"], +) +@pytest.mark.parametrize("max_size", [64 << 10, 1], ids=["in_memory", "rolled_to_disk"]) +def test_parsers_read_from_spooled_response_body(parser, body, expected, max_size): + raw = _spooled_body(body, max_size=max_size) + if parser.__class__ is JsonParser: + assert list(parser.parse(raw)) == [json.loads(body)] + else: + assert list(parser.parse(raw)) == expected + raw.close() + + +def test_spooled_response_body_small_body_stays_in_memory(): + spool = tempfile.SpooledTemporaryFile(max_size=8 << 20) + spool.write(b"hello") + spool.seek(0) + raw = SpooledResponseBody(spool) + assert raw.read() == b"hello" + assert spool._rolled is False + + +def test_streamed_decode_applies_decode_content_for_transport_gzip(): + body = gzip.compress(_JSON_BODY) + response = _streamed_response(body, headers={"Content-Encoding": "gzip"}) + decoder = CompositeRawDecoder(parser=JsonItemsParser(items_path="items")) + assert list(decoder.decode(response)) == [{"id": 1}, {"id": 2}] + + +def test_streamed_gzip_decoder_over_double_compressed_body(): + body = gzip.compress(gzip.compress(_JSON_BODY)) + response = _streamed_response(body, headers={"Content-Encoding": "gzip"}) + decoder = CompositeRawDecoder( + parser=GzipParser(inner_parser=JsonItemsParser(items_path="items")) + ) + assert list(decoder.decode(response)) == [{"id": 1}, {"id": 2}] + + +def test_streamed_gzip_decoder_over_plain_gzip_payload_unchanged(): + body = gzip.compress(_JSON_BODY) + response = _streamed_response(body) + decoder = CompositeRawDecoder( + parser=GzipParser(inner_parser=JsonItemsParser(items_path="items")) + ) + assert list(decoder.decode(response)) == [{"id": 1}, {"id": 2}] + + +def test_composite_raw_decoder_does_not_capture_remainder_without_opt_in(mocker): + body = json.dumps( + {"tickets": [{"id": 1}, {"id": 2}], "after_url": "x", "end_of_stream": False} + ).encode() + response = _streamed_response(body) + spy = mocker.patch( + "airbyte_cdk.sources.declarative.decoders.composite_raw_decoder.set_document_remainder" + ) + decoder = CompositeRawDecoder(parser=JsonItemsParser(items_path="tickets")) + assert list(decoder.decode(response)) == [{"id": 1}, {"id": 2}] + spy.assert_not_called() + assert get_document_remainder(response) is None diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index ea3b0cf151..9e90b4f838 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 @@ -7816,3 +7816,237 @@ def test_incremental_dependency_without_incremental_sync_warns( for record in caplog.records ) assert warning_emitted == expected_warning + + +def _json_items_decoder_component(items_path: str = "tickets"): + from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import ( + CompositeRawDecoder, + JsonItemsParser, + ) + + return CompositeRawDecoder(parser=JsonItemsParser(items_path=items_path)) + + +_PAGINATOR_MANIFEST = """ + paginator: + type: "DefaultPaginator" + page_token_option: + type: RequestPath + pagination_strategy: + type: "CursorPagination" + cursor_value: "{{ response.get('after_url') }}" + stop_condition: "{{ response.get('end_of_stream') }}" +""" + + +def _build_paginator(decoder): + parsed_manifest = YamlDeclarativeSource._parse(_PAGINATOR_MANIFEST) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + paginator_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["paginator"], {} + ) + return factory.create_component( + model_type=DefaultPaginatorModel, + component_definition=paginator_manifest, + config=input_config, + url_base="https://airbyte.io", + decoder=decoder, + ) + + +def test_create_paginator_with_json_items_decoder_is_supported(): + from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import ( + CompositeRawDecoder, + GzipParser, + JsonItemsParser, + ) + from airbyte_cdk.sources.streams.http.streamed_response import set_document_remainder + + paginator = _build_paginator(_json_items_decoder_component()) + assert isinstance(paginator, DefaultPaginator) + assert isinstance(paginator.page_token_option, RequestPath) + + response = requests.Response() + response.status_code = 200 + set_document_remainder(response, {"after_url": "https://next", "end_of_stream": False}) + assert ( + paginator.pagination_strategy.next_page_token( + response=response, last_page_size=2, last_record=None + ) + == "https://next" + ) + set_document_remainder(response, {"after_url": "https://next", "end_of_stream": True}) + assert ( + paginator.pagination_strategy.next_page_token( + response=response, last_page_size=2, last_record=None + ) + is None + ) + + gzip_decoder = CompositeRawDecoder( + parser=GzipParser(inner_parser=JsonItemsParser(items_path="tickets")) + ) + assert isinstance(_build_paginator(gzip_decoder), DefaultPaginator) + + +def test_create_paginator_decoder_support_matrix_unchanged(): + from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import ( + CompositeRawDecoder, + CsvParser, + ) + from airbyte_cdk.sources.declarative.decoders.xml_decoder import XmlDecoder + + assert isinstance(_build_paginator(XmlDecoder(parameters={})), DefaultPaginator) + assert isinstance(_build_paginator(JsonDecoder(parameters={})), DefaultPaginator) + with pytest.raises(ValueError, match="not supported for pagination"): + _build_paginator(CompositeRawDecoder(parser=CsvParser())) + + +_REQUESTER_MANIFEST = """ +requester: + type: HttpRequester + url_base: "https://airbyte.io" + path: "/v1/items" + http_method: "GET" + use_cache: {use_cache} +""" + + +def _build_requester(use_cache: bool, decoder, factory_=None): + content = _REQUESTER_MANIFEST.format(use_cache=str(use_cache).lower()) + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + requester_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["requester"], {} + ) + return (factory_ or factory).create_component( + model_type=HttpRequesterModel, + component_definition=requester_manifest, + config=input_config, + name="name", + decoder=decoder, + ) + + +def test_use_cache_with_streaming_decoder_builds_fine(): + from airbyte_cdk.sources.declarative.models import ( + JsonlDecoder as JsonlDecoderModel, + ) + + requester = _build_requester( + True, + factory.create_jsonl_decoder(JsonlDecoderModel(type="JsonlDecoder"), input_config), + ) + assert requester.use_cache is True + + +def test_use_cache_with_spool_to_disk_decoder_raises(): + from airbyte_cdk.sources.declarative.models import ( + JsonItemsDecoder as JsonItemsDecoderModel, + ) + + with pytest.raises(ValueError, match="use_cache.*spool_to_disk"): + _build_requester( + True, + factory.create_json_items_decoder( + JsonItemsDecoderModel( + type="JsonItemsDecoder", items_path="tickets", spool_to_disk=True + ), + input_config, + ), + ) + + +def test_spool_to_disk_propagates_to_requester(): + from airbyte_cdk.sources.declarative.models import ( + JsonItemsDecoder as JsonItemsDecoderModel, + ) + + requester = _build_requester( + False, + factory.create_json_items_decoder( + JsonItemsDecoderModel( + type="JsonItemsDecoder", items_path="tickets", spool_to_disk=True + ), + input_config, + ), + ) + assert requester.stream_response is True + assert requester.spool_response is True + + +def test_builder_mode_json_items_decoder_is_not_streamed_or_spooled(): + from airbyte_cdk.sources.declarative.models import ( + JsonItemsDecoder as JsonItemsDecoderModel, + ) + + builder_factory = ModelToComponentFactory(emit_connector_builder_messages=True) + decoder = builder_factory.create_json_items_decoder( + JsonItemsDecoderModel(type="JsonItemsDecoder", items_path="tickets", spool_to_disk=True), + input_config, + ) + assert decoder.is_stream_response() is False + assert decoder.spools_response() is False + + +def test_default_paginator_enables_remainder_capture_on_shared_decoder(): + decoder = _json_items_decoder_component() + _build_paginator(decoder) + assert decoder._capture_document_remainder is True + + decoder_without_paginator = _json_items_decoder_component() + assert decoder_without_paginator._capture_document_remainder is False + + +_SPOOL_BODY_FILTER_MANIFEST = """ +requester: + type: HttpRequester + url_base: "https://airbyte.io" + path: "/v1/items" + http_method: "GET" + error_handler: + type: DefaultErrorHandler + response_filters: + - action: IGNORE + error_message_contains: "You do not have access" +""" + + +def _build_requester_from_manifest(content: str, decoder): + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + requester_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["requester"], {} + ) + return factory.create_component( + model_type=HttpRequesterModel, + component_definition=requester_manifest, + config=input_config, + name="name", + decoder=decoder, + ) + + +def _json_items_spooling_decoder(): + from airbyte_cdk.sources.declarative.models import ( + JsonItemsDecoder as JsonItemsDecoderModel, + ) + + return factory.create_json_items_decoder( + JsonItemsDecoderModel(type="JsonItemsDecoder", items_path="tickets", spool_to_disk=True), + input_config, + ) + + +def test_spool_to_disk_with_body_filters_logs_warning(caplog): + with caplog.at_level(logging.WARNING): + _build_requester_from_manifest(_SPOOL_BODY_FILTER_MANIFEST, _json_items_spooling_decoder()) + assert any( + "spool_to_disk" in r.message and "response filters" in r.message for r in caplog.records + ) + + +def test_no_spool_to_disk_no_body_filter_warning(caplog): + with caplog.at_level(logging.WARNING): + _build_requester_from_manifest(_SPOOL_BODY_FILTER_MANIFEST, _json_items_decoder_component()) + assert not any("spool_to_disk" in r.message for r in caplog.records) diff --git a/unit_tests/sources/declarative/requesters/error_handlers/test_http_response_filter.py b/unit_tests/sources/declarative/requesters/error_handlers/test_http_response_filter.py index 87e522d4ab..308d420c9e 100644 --- a/unit_tests/sources/declarative/requesters/error_handlers/test_http_response_filter.py +++ b/unit_tests/sources/declarative/requesters/error_handlers/test_http_response_filter.py @@ -13,6 +13,7 @@ ErrorResolution, ResponseAction, ) +from airbyte_cdk.sources.streams.http.streamed_response import mark_body_streamed @pytest.mark.parametrize( @@ -237,3 +238,101 @@ def test_matches( assert actual_response_status.error_message == expected_error_resolution.error_message else: assert actual_response_status is None + + +def _access_denied_filter(**kwargs) -> HttpResponseFilter: + return HttpResponseFilter( + action=ResponseAction.IGNORE, + failure_type=None, + config={}, + parameters={}, + http_codes=set(), + predicate="", + error_message_contains="You do not have access", + error_message="", + **kwargs, + ) + + +def _streamed_2xx_with_access_denied_body(requests_mock): + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + text=json.dumps({"error": "You do not have access"}), + status_code=200, + ) + return requests.get("https://airbyte.io/", stream=True) + + +def test_error_message_contains_matches_403_despite_streamed_marker(requests_mock): + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + text=json.dumps({"error": "You do not have access"}), + status_code=403, + ) + response = requests.get("https://airbyte.io/", stream=True) + mark_body_streamed(response) + resolution = _access_denied_filter().matches(response) + assert resolution.response_action == ResponseAction.IGNORE + + +def test_unspooled_streamed_2xx_evaluates_body_filters(requests_mock): + response = _streamed_2xx_with_access_denied_body(requests_mock) + mark_body_streamed(response) + resolution = _access_denied_filter().matches(response) + assert resolution.response_action == ResponseAction.IGNORE + + +def test_unmarked_2xx_still_matches_error_message_contains(requests_mock): + response = _streamed_2xx_with_access_denied_body(requests_mock) + resolution = _access_denied_filter().matches(response) + assert resolution.response_action == ResponseAction.IGNORE + + +def test_unspooled_streamed_2xx_evaluates_predicate_filter(requests_mock): + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + text=json.dumps({"flag": "match me"}), + status_code=200, + ) + response = requests.get("https://airbyte.io/", stream=True) + mark_body_streamed(response) + response_filter = HttpResponseFilter( + action=ResponseAction.IGNORE, + failure_type=None, + config={}, + parameters={}, + http_codes=set(), + predicate='{{ response.flag == "match me" }}', + error_message_contains="", + error_message="", + ) + resolution = response_filter.matches(response) + assert resolution.response_action == ResponseAction.IGNORE + + +def test_http_codes_filter_matches_marked_2xx_without_consuming_body(requests_mock): + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + text=json.dumps({"error": "whatever"}), + status_code=200, + ) + response = requests.get("https://airbyte.io/", stream=True) + mark_body_streamed(response) + response_filter = HttpResponseFilter( + action=ResponseAction.IGNORE, + failure_type=None, + config={}, + parameters={}, + http_codes={200}, + predicate="", + error_message_contains="", + error_message="matched: {{ response.get('error') }}", + ) + resolution = response_filter.matches(response) + assert resolution.response_action == ResponseAction.IGNORE + # like main: the error_message template interpolates the actual body + assert resolution.error_message == "matched: whatever" diff --git a/unit_tests/sources/declarative/requesters/test_http_requester.py b/unit_tests/sources/declarative/requesters/test_http_requester.py index 211b348908..df6263042d 100644 --- a/unit_tests/sources/declarative/requesters/test_http_requester.py +++ b/unit_tests/sources/declarative/requesters/test_http_requester.py @@ -960,3 +960,38 @@ def test_http_requester_with_mock_api_budget(http_requester_factory, monkeypatch assert response.status_code == 200 assert mock_budget.acquire_call.call_count == 1 + + +def test_spool_response_is_keyword_only_in_practice(): + from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import ( + CompositeRawDecoder, + JsonItemsParser, + ) + from airbyte_cdk.sources.message import NoopMessageRepository + + decoder = CompositeRawDecoder(parser=JsonItemsParser(items_path="tickets")) + # Positional order matching the fields before spool_response, as on main: + # name, config, parameters, url, url_base, path, authenticator, http_method, + # request_options_provider, error_handler, api_budget, disable_retries, + # message_repository, use_cache, _exit_on_rate_limit, stream_response, decoder + requester = HttpRequester( + "tickets", # name + {}, # config + {}, # parameters + None, # url + "https://api.example.com/", # url_base + "tickets.json", # path + None, # authenticator + HttpMethod.GET, + None, # request_options_provider + None, # error_handler + None, # api_budget + False, # disable_retries + NoopMessageRepository(), + False, # use_cache + False, # _exit_on_rate_limit + False, # stream_response + decoder, + ) + assert requester.decoder is decoder + assert requester.spool_response is False diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index c8aaca355b..3e76456661 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -6299,3 +6299,249 @@ def capturing_resolve(resolver_self, *args, **kwargs): "HttpComponentsResolver's requester should have api_budget set during dynamic stream " "discovery, but it was None. This means discovery HTTP requests are not rate-limited." ) + + +def _parent_substream_manifest(decoder_definition, parent_requester_extra=None): + parent_retriever = { + "requester": { + "path": "parents", + "authenticator": { + "type": "BasicHttpAuthenticator", + "username": "{{ config['api_key'] }}", + }, + **(parent_requester_extra or {}), + }, + "record_selector": {"extractor": {"type": "DpathExtractor", "field_path": []}}, + } + if decoder_definition is not None: + parent_retriever["decoder"] = decoder_definition + parent_stream = { + "type": "DeclarativeStream", + "$parameters": { + "name": "parents", + "primary_key": "id", + "url_base": "https://example.com/v1/", + }, + "schema_loader": {"type": "InlineSchemaLoader", "schema": {}}, + "retriever": parent_retriever, + } + child_stream = { + "type": "DeclarativeStream", + "name": "children", + "$parameters": { + "name": "children", + "primary_key": "id", + "url_base": "https://example.com/v1/", + }, + "schema_loader": {"type": "InlineSchemaLoader", "schema": {}}, + "retriever": { + "requester": { + "path": "children", + "authenticator": { + "type": "BasicHttpAuthenticator", + "username": "{{ config['api_key'] }}", + }, + }, + "record_selector": {"extractor": {"type": "DpathExtractor", "field_path": []}}, + "partition_router": { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "parent_key": "id", + "partition_field": "parent_id", + "stream": deepcopy(parent_stream), + } + ], + }, + }, + } + return { + "version": "0.29.3", + "definitions": {}, + "streams": [deepcopy(parent_stream), child_stream], + "check": {"type": "CheckStream", "stream_names": ["parents"]}, + } + + +def _stream_uses_cache(stream) -> bool: + return stream._stream_partition_generator._partition_factory._retriever.requester.use_cache + + +@pytest.mark.parametrize( + "decoder_definition", + [ + {"type": "JsonItemsDecoder", "items_path": "parents"}, + {"type": "JsonlDecoder"}, + {"type": "CsvDecoder"}, + {"type": "GzipDecoder", "decoder": {"type": "JsonlDecoder"}}, + ], + ids=["json_items", "jsonl", "csv", "gzip_jsonl"], +) +def test_parent_stream_with_streaming_decoder_still_gets_cache(decoder_definition): + manifest = _parent_substream_manifest(decoder_definition) + source = ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=create_catalog("parents"), state=None + ) + streams = source.streams({}) + assert _stream_uses_cache(streams[0]) + parent = streams[1]._stream_partition_generator._stream_slicer.parent_stream_configs[0].stream + assert _stream_uses_cache(parent) + + +@pytest.mark.parametrize( + "decoder_definition", + [ + {"type": "JsonItemsDecoder", "items_path": "parents", "spool_to_disk": True}, + { + "type": "GzipDecoder", + "decoder": { + "type": "JsonItemsDecoder", + "items_path": "parents", + "spool_to_disk": True, + }, + }, + ], + ids=["json_items_spool", "gzip_json_items_spool"], +) +def test_parent_stream_with_spool_to_disk_decoder_does_not_get_cache(decoder_definition): + manifest = _parent_substream_manifest(decoder_definition) + source = ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=create_catalog("parents"), state=None + ) + streams = source.streams({}) + assert not _stream_uses_cache(streams[0]) + parent = streams[1]._stream_partition_generator._stream_slicer.parent_stream_configs[0].stream + assert not _stream_uses_cache(parent) + + +def test_parent_stream_with_json_decoder_still_gets_cache(): + manifest = _parent_substream_manifest({"type": "JsonDecoder"}) + source = ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=create_catalog("parents"), state=None + ) + streams = source.streams({}) + assert _stream_uses_cache(streams[0]) + parent = streams[1]._stream_partition_generator._stream_slicer.parent_stream_configs[0].stream + assert _stream_uses_cache(parent) + + +def test_parent_stream_with_explicit_use_cache_and_spooling_decoder_raises(): + manifest = _parent_substream_manifest( + {"type": "JsonItemsDecoder", "items_path": "parents", "spool_to_disk": True}, + parent_requester_extra={"use_cache": True}, + ) + source = ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=create_catalog("parents"), state=None + ) + with pytest.raises(ValueError, match="use_cache.*cannot be combined"): + source.streams({}) + + +def test_state_delegating_parent_cache_only_applies_to_non_streaming_branch(): + parent_stream = { + "type": "StateDelegatingStream", + "name": "delegating_parent", + "full_refresh_stream": { + "type": "DeclarativeStream", + "$parameters": { + "name": "delegating_parent", + "primary_key": "id", + "url_base": "https://example.com/v1/", + }, + "schema_loader": {"type": "InlineSchemaLoader", "schema": {}}, + "retriever": { + "requester": { + "path": "parents", + "authenticator": { + "type": "BasicHttpAuthenticator", + "username": "{{ config['api_key'] }}", + }, + }, + "record_selector": {"extractor": {"type": "DpathExtractor", "field_path": []}}, + "decoder": {"type": "JsonDecoder"}, + }, + }, + "incremental_stream": { + "type": "DeclarativeStream", + "$parameters": { + "name": "delegating_parent", + "primary_key": "id", + "url_base": "https://example.com/v1/", + }, + "schema_loader": {"type": "InlineSchemaLoader", "schema": {}}, + "retriever": { + "requester": { + "path": "parents/incremental", + "authenticator": { + "type": "BasicHttpAuthenticator", + "username": "{{ config['api_key'] }}", + }, + }, + "record_selector": {"extractor": {"type": "DpathExtractor", "field_path": []}}, + "decoder": { + "type": "JsonItemsDecoder", + "items_path": "parents", + "spool_to_disk": True, + }, + }, + "incremental_sync": { + "type": "DatetimeBasedCursor", + "cursor_field": "updated_at", + "cursor_datetime_formats": ["%s"], + "datetime_format": "%s", + "start_datetime": {"datetime": "{{ config.get('start_date', '1970-01-01') }}"}, + }, + }, + } + child_stream = { + "type": "DeclarativeStream", + "name": "children", + "$parameters": { + "name": "children", + "primary_key": "id", + "url_base": "https://example.com/v1/", + }, + "schema_loader": {"type": "InlineSchemaLoader", "schema": {}}, + "retriever": { + "requester": { + "path": "children", + "authenticator": { + "type": "BasicHttpAuthenticator", + "username": "{{ config['api_key'] }}", + }, + }, + "record_selector": {"extractor": {"type": "DpathExtractor", "field_path": []}}, + "partition_router": { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "parent_key": "id", + "partition_field": "parent_id", + "stream": deepcopy(parent_stream), + } + ], + }, + }, + } + manifest = { + "version": "0.29.3", + "definitions": {}, + "streams": [deepcopy(parent_stream), child_stream], + "check": {"type": "CheckStream", "stream_names": ["delegating_parent"]}, + } + configs = ConcurrentDeclarativeSource._initialize_cache_for_parent_streams( + deepcopy(manifest["streams"]) + ) + delegating = [c for c in configs if c.get("type") == "StateDelegatingStream"] + delegating += [ + cfg["stream"] + for c in configs + for r in [c.get("retriever", {}).get("partition_router") or []] + for r in (r if isinstance(r, list) else [r]) + for cfg in r.get("parent_stream_configs", []) + if cfg["stream"].get("type") == "StateDelegatingStream" + ] + assert delegating + for cfg in delegating: + assert cfg["full_refresh_stream"]["retriever"]["requester"].get("use_cache") is True + assert cfg["incremental_stream"]["retriever"]["requester"].get("use_cache") is not True diff --git a/unit_tests/sources/streams/http/test_http_client.py b/unit_tests/sources/streams/http/test_http_client.py index 8724c48e49..ba241ed4f9 100644 --- a/unit_tests/sources/streams/http/test_http_client.py +++ b/unit_tests/sources/streams/http/test_http_client.py @@ -1,19 +1,32 @@ # Copyright (c) 2024 Airbyte, Inc., all rights reserved. +import gzip +import inspect +import io import json import logging import os +import tempfile import time from datetime import timedelta +from typing import Any, Callable, Mapping, Optional from unittest.mock import MagicMock, patch import pytest import requests +import urllib3 from pympler import asizeof from requests_cache import CachedRequest from airbyte_cdk.models import FailureType, Level from airbyte_cdk.sources.declarative.auth.oauth import DeclarativeOauth2Authenticator +from airbyte_cdk.sources.declarative.requesters.error_handlers import ( + DefaultErrorHandler, + HttpResponseFilter, +) +from airbyte_cdk.sources.declarative.requesters.error_handlers.backoff_strategies import ( + WaitTimeFromHeaderBackoffStrategy, +) from airbyte_cdk.sources.http_logger import format_http_message from airbyte_cdk.sources.message import InMemoryMessageRepository from airbyte_cdk.sources.streams.call_rate import CachedLimiterSession, LimiterSession @@ -38,6 +51,12 @@ Oauth2Authenticator, TokenAuthenticator, ) +from airbyte_cdk.sources.streams.http.streamed_response import ( + SpooledResponseBody, + is_body_streamed, + mark_body_streamed, + set_spooled_body_size, +) from airbyte_cdk.utils.datetime_helpers import ab_datetime_now from airbyte_cdk.utils.traced_exception import AirbyteTracedException @@ -1838,3 +1857,346 @@ def test_given_no_reduce_page_size_action_then_log_the_response_as_a_page(): logged = [json.loads(message.log.message) for message in message_repository.consume_queue()] assert [entry["http"].get("is_auxiliary") for entry in logged] == [None] + + +class _BreakingBody(io.RawIOBase): + """File-like body that raises a urllib3 error after `limit` bytes.""" + + def __init__(self, payload: bytes, limit: int, error: Exception): + self._fh = io.BytesIO(payload) + self._limit = limit + self._error = error + self._served = 0 + + def readinto(self, b) -> int: + if self._served >= self._limit: + raise self._error + n = self._fh.readinto(memoryview(b)[: self._limit - self._served]) + self._served += n + return n + + def readable(self): + return True + + +def _raw_streaming_response(payload: bytes, headers=None, breaking=None): + raw = urllib3.HTTPResponse( + body=breaking if breaking is not None else io.BytesIO(payload), + headers=headers or {}, + status=200, + preload_content=False, + decode_content=False, + ) + response = requests.Response() + response.status_code = 200 + response.raw = raw + response.url = "https://test_base_url.com/v1/endpoint" + for k, v in (headers or {}).items(): + response.headers[k] = v + return response, raw + + +def _http_client_with_session(session, error_handler=None): + return HttpClient( + name="test", + logger=MagicMock(), + error_handler=error_handler, + session=session, + ) + + +_GZIP_PAYLOAD = gzip.compress(json.dumps({"tickets": [{"id": 1}], "after_url": "x"}).encode()) + + +def test_send_request_with_spool_replaces_raw_with_spooled_body(): + mocked_session = MagicMock(spec=requests.Session) + mocked_session.merge_environment_settings.return_value = {} + response, original_raw = _raw_streaming_response( + _GZIP_PAYLOAD, headers={"Content-Encoding": "gzip"} + ) + mocked_session.send.return_value = response + http_client = _http_client_with_session(mocked_session) + + before = set(os.listdir(tempfile.gettempdir())) + request, returned_response = http_client.send_request( + http_method="get", + url="https://test_base_url.com/v1/endpoint", + request_kwargs={"stream": True}, + spool_response=True, + ) + + assert isinstance(returned_response.raw, SpooledResponseBody) + assert returned_response.raw.read() == gzip.decompress(_GZIP_PAYLOAD) + returned_response.raw.seek(0) + assert returned_response.json() == {"tickets": [{"id": 1}], "after_url": "x"} + assert is_body_streamed(returned_response) is True + assert original_raw.closed + returned_response.raw.close() + assert set(os.listdir(tempfile.gettempdir())) == before + + +def _zendesk_like_error_handler(): + return DefaultErrorHandler( + backoff_strategies=[ + WaitTimeFromHeaderBackoffStrategy(header="Retry-After", config={}, parameters={}) + ], + response_filters=[ + HttpResponseFilter( + action=ResponseAction.IGNORE, + error_message_contains="You do not have access", + config={}, + parameters={}, + ), + HttpResponseFilter( + action=ResponseAction.FAIL, + failure_type=FailureType.config_error, + http_codes={403, 404}, + config={}, + parameters={}, + ), + ], + config={}, + parameters={}, + ) + + +@pytest.mark.parametrize( + "injected_error, expected_requests_error", + [ + ( + urllib3.exceptions.IncompleteRead(10, 10), + requests.exceptions.ChunkedEncodingError, + ), + ( + urllib3.exceptions.DecodeError("/x", Exception("corrupt")), + requests.exceptions.ContentDecodingError, + ), + ( + urllib3.exceptions.ReadTimeoutError(None, "/", "timed out"), + requests.exceptions.ConnectionError, + ), + ], + ids=["incomplete_read", "content_decoding", "read_timeout"], +) +def test_spool_mid_body_failure_retries_like_a_failed_request( + mocker, injected_error, expected_requests_error +): + payload = json.dumps({"tickets": [{"id": 1}]}).encode() + breaking = _BreakingBody(payload, len(payload) // 2, injected_error) + broken_response, _ = _raw_streaming_response(payload, breaking=breaking) + good_response, _ = _raw_streaming_response(payload) + + mocked_session = MagicMock(spec=requests.Session) + mocked_session.merge_environment_settings.return_value = {} + mocked_session.send.side_effect = [broken_response, good_response] + mocker.patch("time.sleep") + + http_client = _http_client_with_session( + mocked_session, error_handler=_zendesk_like_error_handler() + ) + + _, returned_response = http_client.send_request( + http_method="get", + url="https://test_base_url.com/v1/endpoint", + request_kwargs={"stream": True}, + spool_response=True, + ) + + assert returned_response is good_response + assert mocked_session.send.call_count == 2 + assert isinstance(returned_response.raw, SpooledResponseBody) + + +def test_without_stream_no_marker_and_spool_is_ignored(): + mocked_session = MagicMock(spec=requests.Session) + mocked_session.merge_environment_settings.return_value = {} + payload = b'{"ok": true}' + response, _ = _raw_streaming_response(payload) + mocked_session.send.return_value = response + http_client = _http_client_with_session(mocked_session) + + _, returned_response = http_client.send_request( + http_method="get", + url="https://test_base_url.com/v1/endpoint", + request_kwargs={}, + spool_response=True, + ) + + assert not isinstance(returned_response.raw, SpooledResponseBody) + assert is_body_streamed(returned_response) is False + + +def test_stream_without_spool_marks_body_streamed(): + mocked_session = MagicMock(spec=requests.Session) + mocked_session.merge_environment_settings.return_value = {} + response, _ = _raw_streaming_response(b"{}") + mocked_session.send.return_value = response + http_client = _http_client_with_session(mocked_session) + + _, returned_response = http_client.send_request( + http_method="get", + url="https://test_base_url.com/v1/endpoint", + request_kwargs={"stream": True}, + ) + + assert is_body_streamed(returned_response) is True + + +def _spooled_response(payload: bytes): + spool = tempfile.TemporaryFile() + spool.write(payload) + spool.seek(0) + response, _ = _raw_streaming_response(payload) + response.raw = SpooledResponseBody(spool) + mark_body_streamed(response) + set_spooled_body_size(response, len(payload)) + return response + + +def test_small_spooled_body_is_evaluated_by_error_message_contains(): + payload = json.dumps({"error": "You do not have access"}).encode() + response = _spooled_response(payload) + filter_ = HttpResponseFilter( + action=ResponseAction.IGNORE, + error_message_contains="You do not have access", + config={}, + parameters={}, + ) + resolution = filter_.matches(response) + assert resolution is not None + assert resolution.response_action == ResponseAction.IGNORE + + +def test_small_spooled_body_is_evaluated_by_predicate(): + payload = json.dumps({"error": "You do not have access"}).encode() + response = _spooled_response(payload) + filter_ = HttpResponseFilter( + action=ResponseAction.IGNORE, + predicate="{{ response.get('error') }}", + config={}, + parameters={}, + ) + resolution = filter_.matches(response) + assert resolution is not None + assert resolution.response_action == ResponseAction.IGNORE + + +def test_large_spooled_body_skips_body_filters(): + payload = json.dumps({"error": "You do not have access", "pad": "x" * (2 << 20)}).encode() + response = _spooled_response(payload) + filter_ = HttpResponseFilter( + action=ResponseAction.IGNORE, + error_message_contains="You do not have access", + config={}, + parameters={}, + ) + assert filter_.matches(response) is None + + +def test_unspooled_streamed_body_is_evaluated_by_body_filters(): + payload = json.dumps({"error": "You do not have access"}).encode() + response, _ = _raw_streaming_response(payload) + mark_body_streamed(response) + filter_ = HttpResponseFilter( + action=ResponseAction.IGNORE, + error_message_contains="You do not have access", + config={}, + parameters={}, + ) + resolution = filter_.matches(response) + assert resolution is not None + assert resolution.response_action == ResponseAction.IGNORE + + +def test_send_request_rewinds_small_spooled_body_after_ignore(): + payload = json.dumps({"error": "You do not have access"}).encode() + assert len(payload) < 1 << 20 + response, _ = _raw_streaming_response(payload) + mocked_session = MagicMock(spec=requests.Session) + mocked_session.merge_environment_settings.return_value = {} + mocked_session.send.return_value = response + http_client = _http_client_with_session( + mocked_session, error_handler=_zendesk_like_error_handler() + ) + + _, returned_response = http_client.send_request( + http_method="get", + url="https://test_base_url.com/v1/endpoint", + request_kwargs={"stream": True}, + spool_response=True, + ) + + assert isinstance(returned_response.raw, SpooledResponseBody) + assert returned_response.raw.read() == payload + + +def test_send_request_large_spooled_body_skips_filter_and_is_rewound(): + payload = json.dumps({"error": "You do not have access", "pad": "x" * (2 << 20)}).encode() + response, _ = _raw_streaming_response(payload) + mocked_session = MagicMock(spec=requests.Session) + mocked_session.merge_environment_settings.return_value = {} + mocked_session.send.return_value = response + http_client = _http_client_with_session( + mocked_session, error_handler=_zendesk_like_error_handler() + ) + + _, returned_response = http_client.send_request( + http_method="get", + url="https://test_base_url.com/v1/endpoint", + request_kwargs={"stream": True}, + spool_response=True, + ) + + assert isinstance(returned_response.raw, SpooledResponseBody) + assert returned_response._content_consumed is False + assert returned_response.raw.tell() == 0 + assert returned_response.raw.read() == payload + + +class LegacyOverrideHttpClient(HttpClient): + def _send_with_retry( + self, + request: requests.PreparedRequest, + request_kwargs: Mapping[str, Any], + log_formatter: Optional[Callable[[requests.Response], Any]] = None, + exit_on_rate_limit: Optional[bool] = False, + ) -> requests.Response: + return super()._send_with_retry(request, request_kwargs, log_formatter, exit_on_rate_limit) + + +def test_legacy_send_with_retry_override_without_spool_param_still_works(): + assert "spool_response" not in inspect.signature(HttpClient._send_with_retry).parameters + assert "spool_response" not in inspect.signature(HttpClient._send).parameters + + mocked_session = MagicMock(spec=requests.Session) + mocked_session.merge_environment_settings.return_value = {} + payload = gzip.compress(json.dumps({"tickets": [{"id": 1}]}).encode()) + response, _ = _raw_streaming_response(payload, headers={"Content-Encoding": "gzip"}) + mocked_session.send.return_value = response + http_client = LegacyOverrideHttpClient( + name="test", + logger=MagicMock(), + session=mocked_session, + ) + + _, resp = http_client.send_request( + http_method="get", + url="https://test_base_url.com/v1/endpoint", + request_kwargs={"stream": True}, + spool_response=False, + ) + assert resp.status_code == 200 + assert not isinstance(resp.raw, SpooledResponseBody) + + response2, _ = _raw_streaming_response(payload, headers={"Content-Encoding": "gzip"}) + mocked_session.send.return_value = response2 + _, resp2 = http_client.send_request( + http_method="get", + url="https://test_base_url.com/v1/endpoint", + request_kwargs={"stream": True}, + spool_response=True, + ) + assert resp2.status_code == 200 + assert isinstance(resp2.raw, SpooledResponseBody) + sent_kwargs = mocked_session.send.call_args.kwargs + assert "_airbyte_spool_response" not in sent_kwargs From 6a82cf6764a5c619868d1e2de3000c9fe9ba8014 Mon Sep 17 00:00:00 2001 From: "zane.hyatt" Date: Mon, 5 Oct 2026 23:16:07 +0000 Subject: [PATCH 6/6] fix(low-code): fall back to a full orjson parse when a spooled JSON body holds an integer outside int64 ijson's yajl backend raises 'integer overflow' for integers outside the signed 64-bit range, which the orjson-based JsonParser accepts. When the body has been spooled (seekable), JsonItemsParser now rewinds, re-parses the document with orjson, yields the items not yet emitted, and reports the pagination remainder from the full document. Non-seekable bodies and any other malformed JSON still raise unchanged. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../decoders/composite_raw_decoder.py | 66 ++++++++++--- .../decoders/test_streamed_json_pagination.py | 92 +++++++++++++++++++ 2 files changed, 145 insertions(+), 13 deletions(-) diff --git a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py index 7e88dbee48..357e8f6fdd 100644 --- a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py +++ b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py @@ -160,6 +160,21 @@ def read(self, size: int = -1) -> bytes: return self._decoder.decode(chunk, final=not chunk).encode("utf-8") +_INTEGER_OVERFLOW_MESSAGE = "integer overflow" + + +def _resolve_items_path(document: Any, items_path: str) -> Tuple[Any, str, List[Any]]: + """Walks an ijson dotted `items_path` and returns (parent, last_key, items_list).""" + *parents, key = items_path.split(".") + node = document + for part in parents: + node = node[part] + items = node[key] + if not isinstance(items, list): + raise ValueError(f"items_path {items_path!r} does not point at a JSON array.") + return node, key, items + + @dataclass class JsonItemsParser(Parser): """Streaming JSON parser that yields each element of a nested array. @@ -192,25 +207,50 @@ def parse( # so ijson keeps using its fast byte backend rather than a (deprecated) text # stream. The recoder decodes lazily in chunks, preserving bounded memory. data = _Utf8Recoder(data, self.encoding) # type: ignore[assignment] + source = data if not isinstance(data, _Utf8Recoder) else data._stream # ijson auto-selects the best available backend (yajl2_c when present) # and reads from `data` lazily — it does not call `.read()` on the # whole stream up front. # use_float=True yields floats for non-integer numbers instead of Decimal, matching # json.loads/orjson behavior so downstream JSON serialization doesn't choke on Decimal. items_prefix = f"{self.items_path}.item" - if on_document_remainder is None: - yield from ijson.items(data, items_prefix, use_float=True) - return - builder = ijson.common.ObjectBuilder() - - def events() -> Generator[Tuple[str, str, Any], None, None]: - for prefix, event, value in ijson.parse(data, use_float=True): - if not (prefix == items_prefix or prefix.startswith(items_prefix + ".")): - builder.event(event, value) - yield prefix, event, value - - yield from ijson.items(events(), items_prefix) - on_document_remainder(builder.value) + yielded = 0 + try: + if on_document_remainder is None: + for item in ijson.items(data, items_prefix, use_float=True): + yield item + yielded += 1 + return + builder = ijson.common.ObjectBuilder() + + def events() -> Generator[Tuple[str, str, Any], None, None]: + for prefix, event, value in ijson.parse(data, use_float=True): + if not (prefix == items_prefix or prefix.startswith(items_prefix + ".")): + builder.event(event, value) + yield prefix, event, value + + for item in ijson.items(events(), items_prefix): + yield item + yielded += 1 + on_document_remainder(builder.value) + except ijson.common.IncompleteJSONError as exc: + if _INTEGER_OVERFLOW_MESSAGE not in str(exc) or not source.seekable(): + raise + # yajl (ijson's C backend) rejects integers outside int64. The spooled body is + # rewindable, so fall back to a full (non-streaming) orjson parse of the same bytes, + # as JsonParser does, and skip the items this response already emitted. + logger.warning( + "JsonItemsParser: integer outside int64 in response; re-parsing spooled body " + "with orjson after %d items", + yielded, + ) + source.seek(0) + document = orjson.loads(source.read().decode(self.encoding or "utf-8")) + parent, key, items = _resolve_items_path(document, self.items_path) + yield from items[yielded:] + if on_document_remainder is not None: + parent[key] = [] + on_document_remainder(document) @dataclass diff --git a/unit_tests/sources/declarative/decoders/test_streamed_json_pagination.py b/unit_tests/sources/declarative/decoders/test_streamed_json_pagination.py index 9e20fa612b..6dace64ded 100644 --- a/unit_tests/sources/declarative/decoders/test_streamed_json_pagination.py +++ b/unit_tests/sources/declarative/decoders/test_streamed_json_pagination.py @@ -5,6 +5,8 @@ import io import json +import ijson +import orjson import pytest import requests import urllib3 @@ -220,3 +222,93 @@ def test_composite_raw_decoder_does_not_capture_remainder_without_opt_in(mocker) assert list(decoder.decode(response)) == [{"id": 1}, {"id": 2}] spy.assert_not_called() assert get_document_remainder(response) is None + + +@pytest.mark.parametrize("max_size", [64 << 10, 1], ids=["in_memory", "rolled_to_disk"]) +@pytest.mark.parametrize("with_remainder", [True, False], ids=["with_remainder", "no_remainder"]) +def test_json_items_parser_int64_overflow_falls_back_on_spooled_body( + max_size, with_remainder, caplog +): + body = b'{"items":[1,9223372036854775808,3],"after_url":"x","end_of_stream":false}' + remainders = [] + raw = _spooled_body(body, max_size=max_size) + kwargs = {"on_document_remainder": remainders.append} if with_remainder else {} + with caplog.at_level("WARNING"): + records = list(JsonItemsParser(items_path="items").parse(raw, **kwargs)) + assert records == [1, 2**63, 3] + assert type(records[1]) is int + if with_remainder: + assert remainders == [{"items": [], "after_url": "x", "end_of_stream": False}] + warnings = [r for r in caplog.records if r.levelname == "WARNING"] + assert len(warnings) == 1 + assert "re-parsing spooled body" in warnings[0].getMessage() + assert "9223372036854775808" not in warnings[0].getMessage() + + +def test_json_items_parser_int64_overflow_mid_list_no_dups_or_gaps(): + items = [{"id": i} for i in range(50)] + items[37]["big"] = 2**63 + body = json.dumps({"items": items}).encode() + raw = _spooled_body(body, max_size=1) + records = list(JsonItemsParser(items_path="items").parse(raw)) + assert [r["id"] for r in records] == list(range(50)) + assert records[37]["big"] == 2**63 + + +def test_json_items_parser_int64_overflow_in_remainder_field(): + body = b'{"items":[{"id":1},{"id":2}],"after_url":"x","total":9223372036854775808}' + remainders = [] + raw = _spooled_body(body, max_size=1) + records = list( + JsonItemsParser(items_path="items").parse(raw, on_document_remainder=remainders.append) + ) + assert records == [{"id": 1}, {"id": 2}] + assert remainders == [{"items": [], "after_url": "x", "total": 2**63}] + + +def test_int64_overflow_in_remainder_field_full_decoder_path(): + body = b'{"items":[{"id":1},{"id":2}],"after_url":"x","total":9223372036854775808}' + response = _streamed_response(body) + response.raw = _spooled_body(body, max_size=1) + decoder = CompositeRawDecoder(parser=JsonItemsParser(items_path="items")) + decoder.enable_document_remainder_capture() + assert list(decoder.decode(response)) == [{"id": 1}, {"id": 2}] + assert get_document_remainder(response) == { + "items": [], + "after_url": "x", + "total": 2**63, + } + + +def test_json_items_parser_int64_overflow_non_seekable_raises(): + response = _streamed_response(b'{"items":[1, 9223372036854775808]}') + decoder = CompositeRawDecoder(parser=JsonItemsParser(items_path="items")) + with pytest.raises(ijson.common.IncompleteJSONError, match="integer overflow"): + list(decoder.decode(response)) + + +@pytest.mark.parametrize( + "body", + [b'{"items":[1,2', b'{"items":[1, tru]}'], + ids=["truncated", "bad_token"], +) +def test_json_items_parser_malformed_json_still_raises(body, caplog): + raw = _spooled_body(body, max_size=1) + with caplog.at_level("WARNING"): + with pytest.raises(ijson.common.IncompleteJSONError) as exc_info: + list(JsonItemsParser(items_path="items").parse(raw)) + assert "integer overflow" not in str(exc_info.value) + assert not any("re-parsing spooled body" in r.getMessage() for r in caplog.records) + + +def test_json_items_parser_int64_overflow_fallback_matches_orjson_for_huge_ints(): + body = b'{"items":[1,18446744073709551616,3],"total":1e30}' + remainders = [] + raw = _spooled_body(body, max_size=1) + records = list( + JsonItemsParser(items_path="items").parse(raw, on_document_remainder=remainders.append) + ) + expected = orjson.loads(body) + assert records == expected["items"] + assert type(records[1]) is float + assert remainders == [{"items": [], "total": 1e30}]