updated after the review + maximes proposal

Author: Oleksandr Bazarnov

airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py

@@ -7,44 +7,31 @@
 import io
 import json
 import logging
-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, Optional
 
 import orjson
 import requests
 
 from airbyte_cdk.models import FailureType
-from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
+from airbyte_cdk.sources.declarative.decoders.decoder import DECODER_OUTPUT_TYPE, Decoder
+from airbyte_cdk.sources.declarative.decoders.decoder_parser import (
+    PARSER_OUTPUT_TYPE,
+    PARSERS_BY_HEADER_TYPE,
+    PARSERS_TYPE,
+    Parser,
+)
 from airbyte_cdk.utils import AirbyteTracedException
 
 logger = logging.getLogger("airbyte")
 
 
-@dataclass
-class Parser(ABC):
-    @abstractmethod
-    def parse(
-        self,
-        data: BufferedIOBase,
-        compressed: Optional[bool] = False,
-    ) -> Generator[MutableMapping[str, Any], None, None]:
-        """
-        Parse data and yield dictionaries.
-        """
-        pass
-
-
 @dataclass
 class GzipParser(Parser):
     inner_parser: Parser
 
-    def parse(
-        self,
-        data: BufferedIOBase,
-        compressed: Optional[bool] = False,
-    ) -> Generator[MutableMapping[str, Any], None, None]:
+    def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE:
         """
         Decompress gzipped bytes and pass decompressed data to the inner parser.
 
@@ -55,22 +42,15 @@ 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
 class JsonParser(Parser):
     encoding: str = "utf-8"
 
-    def parse(
-        self,
-        data: BufferedIOBase,
-        compressed: Optional[bool] = False,
-    ) -> Generator[MutableMapping[str, Any], None, None]:
+    def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE:
         """
         Attempts to deserialize data using orjson library. As an extra layer of safety we fallback on the json library to deserialize the data.
         """
@@ -110,11 +90,7 @@ def _parse_json(self, raw_data: bytes) -> Optional[Any]:
 class JsonLineParser(Parser):
     encoding: Optional[str] = "utf-8"
 
-    def parse(
-        self,
-        data: BufferedIOBase,
-        compressed: Optional[bool] = False,
-    ) -> Generator[MutableMapping[str, Any], None, None]:
+    def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE:
         for line in data:
             try:
                 yield json.loads(line.decode(encoding=self.encoding or "utf-8"))
@@ -138,11 +114,7 @@ def _get_delimiter(self) -> Optional[str]:
 
         return self.delimiter
 
-    def parse(
-        self,
-        data: BufferedIOBase,
-        compressed: Optional[bool] = False,
-    ) -> Generator[MutableMapping[str, Any], None, None]:
+    def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE:
         """
         Parse CSV data from decompressed bytes.
         """
@@ -152,10 +124,9 @@ def parse(
             yield row
 
 
-@dataclass
 class CompositeRawDecoder(Decoder):
     """
-    Decoder strategy to transform a requests.Response into a Generator[MutableMapping[str, Any], None, None]
+    Decoder strategy to transform a requests.Response into a PARSER_OUTPUT_TYPE
     passed response.raw to parser(s).
 
     Note: response.raw is not decoded/decompressed by default. Parsers should be instantiated recursively.
@@ -168,26 +139,80 @@ class CompositeRawDecoder(Decoder):
         )
     """
 
-    parser: Parser
-    stream_response: bool = True
+    def __init__(
+        self,
+        parser: Parser,
+        stream_response: bool = True,
+        parsers_by_header: PARSERS_BY_HEADER_TYPE = None,
+    ) -> None:
+        # since we moved from using `dataclass` to `__init__` method,
+        # we need to keep using the `parser` to be able to resolve the depenencies
+        # between the parsers correctly.
+        self.parser = parser
+
+        self._parsers_by_header = parsers_by_header if parsers_by_header else {}
+        self._stream_response = stream_response
+
+    @classmethod
+    def by_headers(
+        cls,
+        parsers: PARSERS_TYPE,
+        stream_response: bool,
+        fallback_parser: Parser,
+    ) -> "CompositeRawDecoder":
+        """
+        Create a CompositeRawDecoder instance based on header values.
+
+        Args:
+            parsers (PARSERS_TYPE): A list of tuples where each tuple contains headers, header values, and a parser.
+            stream_response (bool): A flag indicating whether the response should be streamed.
+            fallback_parser (Parser): A parser to use if no matching header is found.
+
+        Returns:
+            CompositeRawDecoder: An instance of CompositeRawDecoder configured with the provided parsers.
+        """
+        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)
 
     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]:
+    def decode(self, response: requests.Response) -> DECODER_OUTPUT_TYPE:
+        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:
+        """
+        Selects the appropriate parser based on the response headers.
+
+        This method iterates through the `_parsers_by_header` dictionary to find a matching parser
+        based on the headers in the response. If a matching header and header value are found,
+        the corresponding parser is returned. If no match is found, the default parser is returned.
+
+        Args:
+            response (requests.Response): The HTTP response object containing headers to check.
+
+        Returns:
+            Parser: The parser corresponding to the matched header value, or the default parser if no match is found.
+        """
+        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.parser

airbyte_cdk/sources/declarative/decoders/decoder.py

@@ -8,16 +8,7 @@
 
 import requests
 
-COMPRESSSED_RESPONSE_TYPES = [
-    "gzip",
-    "x-gzip",
-    "gzip, deflate",
-    "x-gzip, deflate",
-    "application/zip",
-    "application/gzip",
-    "application/x-gzip",
-    "application/x-zip-compressed",
-]
+DECODER_OUTPUT_TYPE = Generator[MutableMapping[str, Any], None, None]
 
 
 @dataclass
@@ -33,20 +24,9 @@ def is_stream_response(self) -> bool:
         """
 
     @abstractmethod
-    def decode(
-        self, response: requests.Response
-    ) -> Generator[MutableMapping[str, Any], None, None]:
+    def decode(self, response: requests.Response) -> DECODER_OUTPUT_TYPE:
         """
         Decodes a requests.Response into a Mapping[str, Any] or an array
         :param response: the response to decode
         :return: Generator of Mapping describing the response
         """
-
-    def is_compressed_response(self, response: requests.Response) -> bool:
-        """
-        Check if the response is compressed based on the `Content-Encoding` or `Content-Type` header.
-        """
-        return (
-            response.headers.get("Content-Encoding") in COMPRESSSED_RESPONSE_TYPES
-            or response.headers.get("Content-Type") in COMPRESSSED_RESPONSE_TYPES
-        )

airbyte_cdk/sources/declarative/decoders/decoder_parser.py

@@ -0,0 +1,30 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from io import BufferedIOBase
+from typing import Any, Dict, Generator, List, MutableMapping, Optional, Set, Tuple
+
+logger = logging.getLogger("airbyte")
+
+
+PARSER_OUTPUT_TYPE = Generator[MutableMapping[str, Any], None, None]
+
+
+@dataclass
+class Parser(ABC):
+    @abstractmethod
+    def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE:
+        """
+        Parse data and yield dictionaries.
+        """
+        pass
+
+
+# reusable parser types
+PARSERS_TYPE = List[Tuple[Set[str], Set[str], Parser]]
+PARSERS_BY_HEADER_TYPE = Optional[Dict[str, Dict[str, Parser]]]

airbyte_cdk/sources/declarative/decoders/zipfile_decoder.py

@@ -6,13 +6,13 @@
 import zipfile
 from dataclasses import dataclass
 from io import BytesIO
-from typing import Any, Generator, MutableMapping
 
 import requests
 
 from airbyte_cdk.models import FailureType
 from airbyte_cdk.sources.declarative.decoders import Decoder
 from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import Parser
+from airbyte_cdk.sources.declarative.decoders.decoder import DECODER_OUTPUT_TYPE
 from airbyte_cdk.utils import AirbyteTracedException
 
 logger = logging.getLogger("airbyte")
@@ -25,9 +25,7 @@ class ZipfileDecoder(Decoder):
     def is_stream_response(self) -> bool:
         return False
 
-    def decode(
-        self, response: requests.Response
-    ) -> Generator[MutableMapping[str, Any], None, None]:
+    def decode(self, response: requests.Response) -> DECODER_OUTPUT_TYPE:
         try:
             with zipfile.ZipFile(BytesIO(response.content)) as zip_file:
                 for file_name in zip_file.namelist():
@@ -36,7 +34,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

@@ -2204,9 +2204,29 @@ def create_jsonl_decoder(
     def create_gzip_decoder(
         self, model: GzipDecoderModel, config: Config, **kwargs: Any
     ) -> Decoder:
-        return CompositeRawDecoder(
-            parser=ModelToComponentFactory._get_parser(model, config),
-            stream_response=False if self._emit_connector_builder_messages else True,
+        _compressed_response_types = {
+            "gzip",
+            "x-gzip",
+            "gzip, deflate",
+            "x-gzip, deflate",
+            "application/zip",
+            "application/gzip",
+            "application/x-gzip",
+            "application/x-zip-compressed",
+        }
+
+        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)
+
+        return CompositeRawDecoder.by_headers(
+            [({"Content-Encoding", "Content-Type"}, _compressed_response_types, gzip_parser)],
+            stream_response=True,
+            fallback_parser=gzip_parser.inner_parser,
         )
 
     @staticmethod