diff --git a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py index 7acb5c1e20..fbfa7153b6 100644 --- a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py +++ b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py @@ -26,24 +26,84 @@ logger = logging.getLogger("airbyte") +_GZIP_MAGIC_BYTES = b"\x1f\x8b" + + +class _PrefixedStream(io.RawIOBase): + """Prepend already-read bytes back onto a stream without buffering the entire payload.""" + + def __init__(self, prefix: bytes, stream: BufferedIOBase) -> None: + self._prefix = io.BytesIO(prefix) + self._stream = stream + self._prefix_done = False + + def readable(self) -> bool: + return True + + def read(self, n: int = -1) -> bytes: + if not self._prefix_done: + chunk = self._prefix.read(n) + if chunk: + if n != -1 and len(chunk) >= n: + return chunk + self._prefix_done = True + remaining = self._stream.read(n - len(chunk) if n != -1 else -1) + return chunk + (remaining or b"") + self._prefix_done = True + return self._stream.read(n) + + def readinto(self, b: bytearray | memoryview) -> int: # type: ignore[override] + data = self.read(len(b)) + n = len(data) + b[:n] = data + return n + + def close(self) -> None: # type: ignore[override] + try: + self._prefix.close() + finally: + try: + self._stream.close() + finally: + super().close() + + def __enter__(self) -> "_PrefixedStream": + return self + + def __exit__( + self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any + ) -> None: + self.close() + @dataclass class GzipParser(Parser): inner_parser: Parser def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE: - """ - Decompress gzipped bytes and pass decompressed data to the inner parser. + """Decompress gzipped bytes and pass decompressed data to the inner parser. - IMPORTANT: - - If the data is not gzipped, reset the pointer and pass the data to the inner parser as is. + Auto-detects gzip content by checking for magic bytes (0x1f 0x8b) at the start. + If the data is not gzip-compressed, it is passed directly to the inner parser. + This handles APIs that return gzip-compressed bodies without `Content-Encoding` header. - Note: - - The data is not decoded by default. + Note: In builder mode (`stream_response=False`), the `requests` library auto-decompresses + responses that carry `Content-Encoding: gzip`. The already-decompressed bytes will not + start with the gzip magic bytes, so they pass through to the inner parser unchanged. + A false-positive (decompressed content coincidentally starting with 0x1f 0x8b) is + theoretically possible but extremely unlikely for structured formats (CSV, JSON, JSONL). """ + header = data.read(2) + if not header: + return - with gzip.GzipFile(fileobj=data, mode="rb") as gzipobj: - yield from self.inner_parser.parse(gzipobj) + stream: BufferedIOBase = _PrefixedStream(header, data) # type: ignore[assignment] + + if header == _GZIP_MAGIC_BYTES: + with gzip.GzipFile(fileobj=stream, mode="rb") as gzipobj: + yield from self.inner_parser.parse(gzipobj) + else: + yield from self.inner_parser.parse(stream) @dataclass 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 560dd4056a..b58742f1c9 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2640,15 +2640,16 @@ def create_gzip_decoder( gzip_parser: GzipParser = ModelToComponentFactory._get_parser(model, config) # type: ignore # based on the model, we know this will be a GzipParser if self._emit_connector_builder_messages: - # This is very surprising but if the response is not streamed, - # CompositeRawDecoder calls response.content and the requests library actually uncompress the data as opposed to response.raw, - # which uses urllib3 directly and does not uncompress the data. - return CompositeRawDecoder(gzip_parser.inner_parser, False) + # When not streaming, CompositeRawDecoder uses response.content which the requests + # library auto-decompresses when Content-Encoding is set. However, some APIs return + # gzip data without Content-Encoding headers. Using gzip_parser (which auto-detects + # gzip magic bytes) ensures decompression works in both cases. + return CompositeRawDecoder(gzip_parser, False) return CompositeRawDecoder.by_headers( [({"Content-Encoding", "Content-Type"}, _compressed_response_types, gzip_parser)], stream_response=True, - fallback_parser=gzip_parser.inner_parser, + fallback_parser=gzip_parser, ) @staticmethod diff --git a/unit_tests/sources/declarative/decoders/test_composite_decoder.py b/unit_tests/sources/declarative/decoders/test_composite_decoder.py index d92d6c605d..2a8a228078 100644 --- a/unit_tests/sources/declarative/decoders/test_composite_decoder.py +++ b/unit_tests/sources/declarative/decoders/test_composite_decoder.py @@ -362,3 +362,124 @@ def test_given_response_is_not_streamed_when_decode_then_can_be_called_multiple_ content_second_time = list(composite_raw_decoder.decode(response)) assert content == content_second_time + + +class TestGzipParserAutoDetection: + """Tests for GzipParser auto-detection of gzip content via magic bytes.""" + + def test_gzip_csv_without_content_encoding_header(self, requests_mock): + """GzipParser should decompress gzip data even without Content-Encoding header.""" + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + content=generate_csv(should_compress=True), + # No Content-Encoding header set + ) + response = requests.get("https://airbyte.io/", stream=True) + parser = GzipParser(inner_parser=CsvParser()) + decoder = CompositeRawDecoder(parser=parser) + records = list(decoder.decode(response)) + assert len(records) == 3 + assert records[0] == {"id": "1", "name": "John", "age": "28"} + + def test_gzip_jsonl_without_content_encoding_header(self, requests_mock): + """GzipParser should decompress gzip JSONL data without Content-Encoding header.""" + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + content=generate_compressed_jsonlines(), + # No Content-Encoding header set + ) + response = requests.get("https://airbyte.io/", stream=True) + parser = GzipParser(inner_parser=JsonLineParser()) + decoder = CompositeRawDecoder(parser=parser) + records = list(decoder.decode(response)) + assert len(records) == 3 + assert records[0] == {"id": 1, "message": "Hello, World!"} + + def test_non_gzip_data_passthrough(self, requests_mock): + """GzipParser should pass non-gzip data through to the inner parser unchanged.""" + plain_csv = generate_csv(should_compress=False) + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + content=plain_csv, + ) + response = requests.get("https://airbyte.io/", stream=True) + parser = GzipParser(inner_parser=CsvParser()) + decoder = CompositeRawDecoder(parser=parser) + records = list(decoder.decode(response)) + assert len(records) == 3 + assert records[0] == {"id": "1", "name": "John", "age": "28"} + + def test_non_gzip_jsonl_passthrough(self, requests_mock): + """GzipParser should pass plain JSONL data through to the inner parser.""" + plain_jsonl = "".join(generate_jsonlines()).encode("utf-8") + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + content=plain_jsonl, + ) + response = requests.get("https://airbyte.io/", stream=True) + parser = GzipParser(inner_parser=JsonLineParser()) + decoder = CompositeRawDecoder(parser=parser) + records = list(decoder.decode(response)) + assert len(records) == 3 + assert records[0] == {"id": 1, "message": "Hello, World!"} + + def test_empty_data_returns_no_records(self): + """GzipParser should gracefully handle empty data.""" + parser = GzipParser(inner_parser=CsvParser()) + records = list(parser.parse(BytesIO(b""))) + assert records == [] + + def test_gzip_fallback_in_by_headers_mode(self, requests_mock): + """When used as fallback_parser in by_headers mode, GzipParser should auto-detect gzip.""" + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + content=generate_compressed_jsonlines(), + headers={"Content-Encoding": "identity"}, # Not gzip, so fallback is used + ) + response = requests.get("https://airbyte.io/", stream=True) + gzip_parser = GzipParser(inner_parser=JsonLineParser()) + decoder = CompositeRawDecoder.by_headers( + [({"Content-Encoding"}, {"gzip"}, gzip_parser)], + stream_response=True, + fallback_parser=gzip_parser, + ) + records = list(decoder.decode(response)) + assert len(records) == 3 + + def test_non_streamed_gzip_without_content_encoding(self, requests_mock): + """GzipParser should handle gzip data in non-streamed mode (response.content).""" + requests_mock.register_uri( + "GET", + "https://airbyte.io/", + content=generate_compressed_jsonlines(), + ) + response = requests.get("https://airbyte.io/") # Not streamed + parser = GzipParser(inner_parser=JsonLineParser()) + decoder = CompositeRawDecoder(parser=parser, stream_response=False) + records = list(decoder.decode(response)) + assert len(records) == 3 + + def test_prefixed_stream_closes_wrapped_stream(self): + """GzipParser should close the underlying stream once parsing completes.""" + + class CloseTrackingBytesIO(BytesIO): + def __init__(self, data: bytes) -> None: + super().__init__(data) + self.closed_flag = False + + def close(self) -> None: + self.closed_flag = True + super().close() + + stream = CloseTrackingBytesIO(compress_with_gzip(json.dumps({"hello": "world"}))) + parser = GzipParser(inner_parser=JsonParser()) + + records = list(parser.parse(stream)) + + assert records == [{"hello": "world"}] + assert stream.closed_flag is True