Proposal for generic composite raw decoder

Author: maxi297

airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py

@@ -10,7 +10,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from io import BufferedIOBase, TextIOWrapper
-from typing import Any, Generator, MutableMapping, Optional
+from typing import Any, Dict, Generator, List, MutableMapping, Optional, Set, Tuple
 
 import orjson
 import requests
@@ -28,7 +28,6 @@ class Parser(ABC):
     def parse(
         self,
         data: BufferedIOBase,
-        compressed: Optional[bool] = False,
     ) -> Generator[MutableMapping[str, Any], None, None]:
         """
         Parse data and yield dictionaries.
@@ -43,7 +42,6 @@ class GzipParser(Parser):
     def parse(
         self,
         data: BufferedIOBase,
-        compressed: Optional[bool] = False,
     ) -> Generator[MutableMapping[str, Any], None, None]:
         """
         Decompress gzipped bytes and pass decompressed data to the inner parser.
@@ -55,11 +53,8 @@ def parse(
             - The data is not decoded by default.
         """
 
-        if compressed:
-            with gzip.GzipFile(fileobj=data, mode="rb") as gzipobj:
-                yield from self.inner_parser.parse(gzipobj)
-        else:
-            yield from self.inner_parser.parse(data)
+        with gzip.GzipFile(fileobj=data, mode="rb") as gzipobj:
+            yield from self.inner_parser.parse(gzipobj)
 
 
 @dataclass
@@ -69,7 +64,6 @@ class JsonParser(Parser):
     def parse(
         self,
         data: BufferedIOBase,
-        compressed: Optional[bool] = False,
     ) -> Generator[MutableMapping[str, Any], None, None]:
         """
         Attempts to deserialize data using orjson library. As an extra layer of safety we fallback on the json library to deserialize the data.
@@ -113,7 +107,6 @@ class JsonLineParser(Parser):
     def parse(
         self,
         data: BufferedIOBase,
-        compressed: Optional[bool] = False,
     ) -> Generator[MutableMapping[str, Any], None, None]:
         for line in data:
             try:
@@ -141,7 +134,6 @@ def _get_delimiter(self) -> Optional[str]:
     def parse(
         self,
         data: BufferedIOBase,
-        compressed: Optional[bool] = False,
     ) -> Generator[MutableMapping[str, Any], None, None]:
         """
         Parse CSV data from decompressed bytes.
@@ -152,7 +144,9 @@ def parse(
             yield row
 
 
-@dataclass
+_HEADER = str
+_HEADER_VALUE = str
+
 class CompositeRawDecoder(Decoder):
     """
     Decoder strategy to transform a requests.Response into a Generator[MutableMapping[str, Any], None, None]
@@ -168,26 +162,46 @@ class CompositeRawDecoder(Decoder):
         )
     """
 
-    parser: Parser
-    stream_response: bool = True
+    @classmethod
+    def by_headers(cls, parsers: List[Tuple[Set[_HEADER], Set[_HEADER_VALUE], Parser]], stream_response: bool, fallback_parser: Parser) -> "CompositeRawDecoder":
+        parsers_by_header = {}
+        for headers, header_values, parser in parsers:
+            for header in headers:
+                parsers_by_header[header] = {header_value: parser for header_value in header_values}
+        return cls(fallback_parser, stream_response, parsers_by_header)
+
+    @classmethod
+    def from_parser(cls, parser: Parser, stream_response: bool) -> "CompositeRawDecoder":
+        return cls(parser, stream_response, {})
+
+    def __init__(self, parser: Parser, stream_response: bool = True, parsers_by_header: Optional[Dict[_HEADER, Dict[_HEADER_VALUE, Parser]]] = None) -> None:
+        self._parsers_by_header = parsers_by_header if parsers_by_header else {}
+        self._fallback_parser = parser
+        self._stream_response = stream_response
 
     def is_stream_response(self) -> bool:
-        return self.stream_response
+        return self._stream_response
 
     def decode(
         self,
         response: requests.Response,
     ) -> Generator[MutableMapping[str, Any], None, None]:
+        parser = self._select_parser(response)
         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
             # 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
-            yield from self.parser.parse(
+            yield from parser.parse(
                 data=response.raw,  # type: ignore[arg-type]
-                compressed=self.is_compressed_response(response),
             )
             response.raw.close()
         else:
-            yield from self.parser.parse(data=io.BytesIO(response.content))
+            yield from parser.parse(data=io.BytesIO(response.content))
+
+    def _select_parser(self, response: requests.Response) -> Parser:
+        for header, parser_by_header_value in self._parsers_by_header.items():
+            if header in response.headers and response.headers[header] in parser_by_header_value.keys():
+                return parser_by_header_value[response.headers[header]]
+        return self._fallback_parser

airbyte_cdk/sources/declarative/decoders/zipfile_decoder.py

@@ -36,7 +36,6 @@ def decode(
                     try:
                         yield from self.parser.parse(
                             buffered_content,
-                            compressed=self.is_compressed_response(response),
                         )
                     except Exception as e:
                         logger.error(

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -513,6 +513,17 @@
     SchemaNormalizationModel.Default: TransformConfig.DefaultSchemaNormalization,
 }
 
+_COMPRESSED_RESPONSE_TYPES = {
+    "gzip",
+    "x-gzip",
+    "gzip, deflate",
+    "x-gzip, deflate",
+    "application/zip",
+    "application/gzip",
+    "application/x-gzip",
+    "application/x-zip-compressed",
+}
+
 
 class ModelToComponentFactory:
     EPOCH_DATETIME_FORMAT = "%s"
@@ -2204,9 +2215,13 @@ def create_jsonl_decoder(
     def create_gzip_decoder(
         self, model: GzipDecoderModel, config: Config, **kwargs: Any
     ) -> Decoder:
-        return CompositeRawDecoder(
-            parser=ModelToComponentFactory._get_parser(model, config),
+        gzip_parser: GzipParser = ModelToComponentFactory._get_parser(model, config)  # type: ignore  # based on the model, we know this will be a GzipParser
+        return CompositeRawDecoder.by_headers(
+            [
+                ({"Content-Encoding", "Content-Type"}, _COMPRESSED_RESPONSE_TYPES, gzip_parser)
+            ],
             stream_response=False if self._emit_connector_builder_messages else True,
+            fallback_parser=gzip_parser.inner_parser,
         )
 
     @staticmethod

unit_tests/sources/declarative/decoders/test_composite_decoder.py

@@ -8,7 +8,8 @@
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from io import BytesIO, StringIO
 from threading import Thread
-from unittest.mock import patch
+from typing import Iterable
+from unittest.mock import Mock, patch
 
 import pytest
 import requests
@@ -82,7 +83,7 @@ def test_composite_raw_decoder_gzip_csv_parser(requests_mock, encoding: str):
     assert counter == 3
 
 
-def generate_jsonlines():
+def generate_jsonlines() -> Iterable[str]:
     """
     Generator function to yield data in JSON Lines format.
     This is useful for streaming large datasets.
@@ -111,41 +112,46 @@ def test_composite_raw_decoder_gzip_jsonline_parser(requests_mock, encoding: str
         "GET",
         "https://airbyte.io/",
         content=generate_compressed_jsonlines(encoding=encoding),
-        headers={"Content-Encoding": "gzip"},
     )
     response = requests.get("https://airbyte.io/", stream=True)
 
     parser = GzipParser(inner_parser=JsonLineParser(encoding=encoding))
-    composite_raw_decoder = CompositeRawDecoder(parser=parser)
+    composite_raw_decoder = CompositeRawDecoder(parser)
     counter = 0
     for _ in composite_raw_decoder.decode(response):
         counter += 1
     assert counter == 3
 
 
-@pytest.mark.parametrize("encoding", ["utf-8", "utf", "iso-8859-1"])
-def test_composite_raw_decoder_gzip_jsonline_parser_decodes_non_gzipped_raw_response(
-    requests_mock, encoding: str
-) -> None:
-    """
-    Test the GzipParser with a non-compressed response.
-    """
+def test_given_header_match_when_decode_then_select_parser(requests_mock):
+    requests_mock.register_uri(
+        "GET",
+        "https://airbyte.io/",
+        content=generate_compressed_jsonlines(),
+        headers={"Content-Encoding": "gzip"},
+    )
+    response = requests.get("https://airbyte.io/", stream=True)
 
+    parser = GzipParser(inner_parser=JsonLineParser())
+    unused_parser = Mock()
+    composite_raw_decoder = CompositeRawDecoder.by_headers([({"Content-Encoding"}, {"gzip"}, parser)], stream_response=True, fallback_parser=unused_parser)
+    counter = 0
+    for _ in composite_raw_decoder.decode(response):
+        counter += 1
+    assert counter == 3
+
+
+def test_given_header_does_not_match_when_decode_then_select_fallback_parser(requests_mock):
     requests_mock.register_uri(
         "GET",
         "https://airbyte.io/",
-        # we encode the jsonl content as bytes here
-        content="".join(generate_jsonlines()).encode(encoding),
-        # we don't specify the `Content-Encoding` header here
-        # to simulate a non-compressed response
-        # but we still use the GzipParser to decode it
-        # to test the GzipParser's behavior with non-compressed data
-        # and to ensure it doesn't raise an error.
+        content="".join(generate_jsonlines()).encode("utf-8"),
+        headers={"Content-Encoding": "not gzip in order to expect fallback"},
     )
     response = requests.get("https://airbyte.io/", stream=True)
 
-    parser = GzipParser(inner_parser=JsonLineParser(encoding=encoding))
-    composite_raw_decoder = CompositeRawDecoder(parser=parser)
+    unused_parser = GzipParser(inner_parser=Mock())
+    composite_raw_decoder = CompositeRawDecoder.by_headers([({"Content-Encoding"}, {"gzip"}, unused_parser)], stream_response=True, fallback_parser=JsonLineParser())
     counter = 0
     for _ in composite_raw_decoder.decode(response):
         counter += 1