feat: add Composite Raw Decoder (#179)

Signed-off-by: Artem Inzhyyants <[email protected]>

Author: artem1205

.github/workflows/pypi_publish.yml

@@ -10,7 +10,7 @@ name: Packaging and Publishing
 on:
   push:
     tags:
-      - 'v*'
+      - "v*"
   workflow_dispatch:
     inputs:
       version:

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -2766,9 +2766,64 @@ definitions:
           - "$ref": "#/definitions/IterableDecoder"
           - "$ref": "#/definitions/XmlDecoder"
           - "$ref": "#/definitions/GzipJsonDecoder"
+          - "$ref": "#/definitions/CompositeRawDecoder"
       $parameters:
         type: object
         additionalProperties: true
+  CompositeRawDecoder:
+    description: "(This is experimental, use at your own risk)"
+    type: object
+    required:
+      - type
+      - parser
+    properties:
+      type:
+        type: string
+        enum: [CompositeRawDecoder]
+      parser:
+        anyOf:
+          - "$ref": "#/definitions/GzipParser"
+          - "$ref": "#/definitions/JsonLineParser"
+          - "$ref": "#/definitions/CsvParser"
+  #    PARSERS
+  GzipParser:
+    type: object
+    required:
+      - type
+      - inner_parser
+    properties:
+      type:
+        type: string
+        enum: [GzipParser]
+      inner_parser:
+        anyOf:
+          - "$ref": "#/definitions/JsonLineParser"
+          - "$ref": "#/definitions/CsvParser"
+  JsonLineParser:
+    type: object
+    required:
+      - type
+    properties:
+      type:
+        type: string
+        enum: [JsonLineParser]
+      encoding:
+        type: string
+        default: utf-8
+  CsvParser:
+    type: object
+    required:
+      - type
+    properties:
+      type:
+        type: string
+        enum: [CsvParser]
+      encoding:
+        type: string
+        default: utf-8
+      delimiter:
+        type: string
+        default: ","
   AsyncJobStatusMap:
     description: Matches the api job status to Async Job Status.
     type: object

airbyte_cdk/sources/declarative/decoders/__init__.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import CompositeRawDecoder
 from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
 from airbyte_cdk.sources.declarative.decoders.json_decoder import (
     GzipJsonDecoder,
@@ -17,6 +18,7 @@
 
 __all__ = [
     "Decoder",
+    "CompositeRawDecoder",
     "JsonDecoder",
     "JsonlDecoder",
     "IterableDecoder",

airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py

@@ -0,0 +1,97 @@
+import csv
+import gzip
+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
+
+import requests
+
+from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
+
+logger = logging.getLogger("airbyte")
+
+
+@dataclass
+class Parser(ABC):
+    @abstractmethod
+    def parse(
+        self,
+        data: BufferedIOBase,
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        """
+        Parse data and yield dictionaries.
+        """
+        pass
+
+
+@dataclass
+class GzipParser(Parser):
+    inner_parser: Parser
+
+    def parse(
+        self,
+        data: BufferedIOBase,
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        """
+        Decompress gzipped bytes and pass decompressed data to the inner parser.
+        """
+        with gzip.GzipFile(fileobj=data, mode="rb") as gzipobj:
+            yield from self.inner_parser.parse(gzipobj)
+
+
+@dataclass
+class JsonLineParser(Parser):
+    encoding: Optional[str] = "utf-8"
+
+    def parse(
+        self,
+        data: BufferedIOBase,
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        for line in data:
+            try:
+                yield json.loads(line.decode(encoding=self.encoding or "utf-8"))
+            except json.JSONDecodeError as e:
+                logger.warning(f"Cannot decode/parse line {line!r} as JSON, error: {e}")
+
+
+@dataclass
+class CsvParser(Parser):
+    # TODO: migrate implementation to re-use file-base classes
+    encoding: Optional[str] = "utf-8"
+    delimiter: Optional[str] = ","
+
+    def parse(
+        self,
+        data: BufferedIOBase,
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        """
+        Parse CSV data from decompressed bytes.
+        """
+        text_data = TextIOWrapper(data, encoding=self.encoding)  # type: ignore
+        reader = csv.DictReader(text_data, delimiter=self.delimiter or ",")
+        yield from reader
+
+
+@dataclass
+class CompositeRawDecoder(Decoder):
+    """
+    Decoder strategy to transform a requests.Response into a Generator[MutableMapping[str, Any], None, None]
+    passed response.raw to parser(s).
+    Note: response.raw is not decoded/decompressed by default.
+    parsers should be instantiated recursively.
+    Example:
+    composite_raw_decoder = CompositeRawDecoder(parser=GzipParser(inner_parser=JsonLineParser(encoding="iso-8859-1")))
+    """
+
+    parser: Parser
+
+    def is_stream_response(self) -> bool:
+        return True
+
+    def decode(
+        self, response: requests.Response
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        yield from self.parser.parse(data=response.raw)  # type: ignore[arg-type]

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -1125,6 +1125,17 @@ class LegacySessionTokenAuthenticator(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class JsonLineParser(BaseModel):
+    type: Literal["JsonLineParser"]
+    encoding: Optional[str] = "utf-8"
+
+
+class CsvParser(BaseModel):
+    type: Literal["CsvParser"]
+    encoding: Optional[str] = "utf-8"
+    delimiter: Optional[str] = ","
+
+
 class AsyncJobStatusMap(BaseModel):
     type: Optional[Literal["AsyncJobStatusMap"]] = None
     running: List[str]
@@ -1208,6 +1219,8 @@ class ComponentMappingDefinition(BaseModel):
             "{{ components_values['updates'] }}",
             "{{ components_values['MetaData']['LastUpdatedTime'] }}",
             "{{ config['segment_id'] }}",
+            "{{ stream_slice['parent_id'] }}",
+            "{{ stream_slice['extra_fields']['name'] }}",
         ],
         title="Value",
     )
@@ -1504,6 +1517,11 @@ class RecordSelector(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class GzipParser(BaseModel):
+    type: Literal["GzipParser"]
+    inner_parser: Union[JsonLineParser, CsvParser]
+
+
 class Spec(BaseModel):
     type: Literal["Spec"]
     connection_specification: Dict[str, Any] = Field(
@@ -1534,6 +1552,11 @@ class CompositeErrorHandler(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class CompositeRawDecoder(BaseModel):
+    type: Literal["CompositeRawDecoder"]
+    parser: Union[GzipParser, JsonLineParser, CsvParser]
+
+
 class DeclarativeSource1(BaseModel):
     class Config:
         extra = Extra.forbid
@@ -1936,6 +1959,7 @@ class SimpleRetriever(BaseModel):
             IterableDecoder,
             XmlDecoder,
             GzipJsonDecoder,
+            CompositeRawDecoder,
         ]
     ] = Field(
         None,

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -67,6 +67,12 @@
     PaginationDecoderDecorator,
     XmlDecoder,
 )
+from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import (
+    CompositeRawDecoder,
+    CsvParser,
+    GzipParser,
+    JsonLineParser,
+)
 from airbyte_cdk.sources.declarative.extractors import (
     DpathExtractor,
     RecordFilter,
@@ -125,6 +131,9 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     CompositeErrorHandler as CompositeErrorHandlerModel,
 )
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    CompositeRawDecoder as CompositeRawDecoderModel,
+)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ConcurrencyLevel as ConcurrencyLevelModel,
 )
@@ -134,6 +143,9 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ConstantBackoffStrategy as ConstantBackoffStrategyModel,
 )
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    CsvParser as CsvParserModel,
+)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     CursorPagination as CursorPaginationModel,
 )
@@ -203,6 +215,9 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     GzipJsonDecoder as GzipJsonDecoderModel,
 )
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    GzipParser as GzipParserModel,
+)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     HttpComponentsResolver as HttpComponentsResolverModel,
 )
@@ -227,6 +242,9 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     JsonlDecoder as JsonlDecoderModel,
 )
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    JsonLineParser as JsonLineParserModel,
+)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     JwtAuthenticator as JwtAuthenticatorModel,
 )
@@ -455,6 +473,7 @@ def _init_mappings(self) -> None:
             BearerAuthenticatorModel: self.create_bearer_authenticator,
             CheckStreamModel: self.create_check_stream,
             CompositeErrorHandlerModel: self.create_composite_error_handler,
+            CompositeRawDecoderModel: self.create_composite_raw_decoder,
             ConcurrencyLevelModel: self.create_concurrency_level,
             ConstantBackoffStrategyModel: self.create_constant_backoff_strategy,
             CursorPaginationModel: self.create_cursor_pagination,
@@ -485,7 +504,9 @@ def _init_mappings(self) -> None:
             InlineSchemaLoaderModel: self.create_inline_schema_loader,
             JsonDecoderModel: self.create_json_decoder,
             JsonlDecoderModel: self.create_jsonl_decoder,
+            JsonLineParserModel: self.create_json_line_parser,
             GzipJsonDecoderModel: self.create_gzipjson_decoder,
+            GzipParserModel: self.create_gzip_parser,
             KeysToLowerModel: self.create_keys_to_lower_transformation,
             KeysToSnakeCaseModel: self.create_keys_to_snake_transformation,
             FlattenFieldsModel: self.create_flatten_fields,
@@ -1701,6 +1722,12 @@ def create_jsonl_decoder(
     ) -> JsonlDecoder:
         return JsonlDecoder(parameters={})
 
+    @staticmethod
+    def create_json_line_parser(
+        model: JsonLineParserModel, config: Config, **kwargs: Any
+    ) -> JsonLineParser:
+        return JsonLineParser(encoding=model.encoding)
+
     @staticmethod
     def create_iterable_decoder(
         model: IterableDecoderModel, config: Config, **kwargs: Any
@@ -1717,6 +1744,22 @@ def create_gzipjson_decoder(
     ) -> GzipJsonDecoder:
         return GzipJsonDecoder(parameters={}, encoding=model.encoding)
 
+    def create_gzip_parser(
+        self, model: GzipParserModel, config: Config, **kwargs: Any
+    ) -> GzipParser:
+        inner_parser = self._create_component_from_model(model=model.inner_parser, config=config)
+        return GzipParser(inner_parser=inner_parser)
+
+    @staticmethod
+    def create_csv_parser(model: CsvParserModel, config: Config, **kwargs: Any) -> CsvParser:
+        return CsvParser(encoding=model.encoding, delimiter=model.delimiter)
+
+    def create_composite_raw_decoder(
+        self, model: CompositeRawDecoderModel, config: Config, **kwargs: Any
+    ) -> CompositeRawDecoder:
+        parser = self._create_component_from_model(model=model.parser, config=config)
+        return CompositeRawDecoder(parser=parser)
+
     @staticmethod
     def create_json_file_schema_loader(
         model: JsonFileSchemaLoaderModel, config: Config, **kwargs: Any