feat: add download_decoder + download_extractor

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -1331,6 +1331,23 @@ definitions:
       $parameters:
         type: object
         additionalProperties: true
+  ResponseToFileExtractor:
+    title: CSV To File Extractor
+    description: Record extractor that downloads a headless
+    type: object
+    required:
+      - type
+      - needs_decompression
+    properties:
+      type:
+        type: string
+        enum: [ResponseToFileExtractor]
+      needs_decompression:
+        type: boolean
+        default: true
+      $parameters:
+        type: object
+        additionalProperties: true
   ExponentialBackoffStrategy:
     title: Exponential Backoff
     description: Backoff strategy with an exponential backoff interval. The interval is defined as factor * 2^attempt_count.
@@ -2513,6 +2530,12 @@ definitions:
         anyOf:
           - "$ref": "#/definitions/CustomRecordExtractor"
           - "$ref": "#/definitions/DpathExtractor"
+      download_extractor:
+        description: Responsible for fetching the records from provided urls.
+        anyOf:
+          - "$ref": "#/definitions/CustomRecordExtractor"
+          - "$ref": "#/definitions/DpathExtractor"
+          - "$ref": "#/definitions/ResponseToFileExtractor"
       creation_requester:
         description: Requester component that describes how to prepare HTTP requests to send to the source API to create the async server-side job.
         anyOf:
@@ -2567,6 +2590,16 @@ definitions:
           - "$ref": "#/definitions/IterableDecoder"
           - "$ref": "#/definitions/XmlDecoder"
           - "$ref": "#/definitions/GzipJsonDecoder"
+      download_decoder:
+        title: Download Decoder
+        description: Component decoding the download response so records can be extracted.
+        anyOf:
+          - "$ref": "#/definitions/CustomDecoder"
+          - "$ref": "#/definitions/JsonDecoder"
+          - "$ref": "#/definitions/JsonlDecoder"
+          - "$ref": "#/definitions/IterableDecoder"
+          - "$ref": "#/definitions/XmlDecoder"
+          - "$ref": "#/definitions/GzipJsonDecoder"
       $parameters:
         type: object
         additionalProperties: true

airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py

@@ -6,6 +6,7 @@
 import uuid
 import zlib
 from contextlib import closing
+from dataclasses import InitVar, dataclass
 from typing import Any, Dict, Iterable, Mapping, Optional, Tuple
 
 import pandas as pd
@@ -19,6 +20,7 @@
 DOWNLOAD_CHUNK_SIZE: int = 1024 * 10
 
 
+@dataclass
 class ResponseToFileExtractor(RecordExtractor):
     """
     This class is used when having very big HTTP responses (usually streamed) which would require too much memory so we use disk space as
@@ -28,7 +30,10 @@ class ResponseToFileExtractor(RecordExtractor):
     a first iteration so we will only support CSV parsing using pandas as salesforce and sendgrid were doing.
     """
 
-    def __init__(self) -> None:
+    parameters: InitVar[Mapping[str, Any]]
+    needs_decompression: bool = True
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self.logger = logging.getLogger("airbyte")
 
     def _get_response_encoding(self, headers: Dict[str, Any]) -> str:
@@ -89,21 +94,18 @@ def _save_to_file(self, response: requests.Response) -> Tuple[str, str]:
         """
         # set filepath for binary data from response
         decompressor = zlib.decompressobj(zlib.MAX_WBITS | 32)
-        needs_decompression = True  # we will assume at first that the response is compressed and change the flag if not
 
         tmp_file = str(uuid.uuid4())
         with closing(response) as response, open(tmp_file, "wb") as data_file:
             response_encoding = self._get_response_encoding(dict(response.headers or {}))
             for chunk in response.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
                 try:
-                    if needs_decompression:
+                    if self.needs_decompression:
                         data_file.write(decompressor.decompress(chunk))
-                        needs_decompression = True
                     else:
                         data_file.write(self._filter_null_bytes(chunk))
                 except zlib.error:
                     data_file.write(self._filter_null_bytes(chunk))
-                    needs_decompression = False
 
         # check the file exists
         if os.path.isfile(tmp_file):

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -567,6 +567,12 @@ class DpathExtractor(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class ResponseToFileExtractor(BaseModel):
+    type: Literal["ResponseToFileExtractor"]
+    needs_decompression: bool
+    parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
+
+
 class ExponentialBackoffStrategy(BaseModel):
     type: Literal["ExponentialBackoffStrategy"]
     factor: Optional[Union[float, str]] = Field(
@@ -1676,6 +1682,9 @@ class AsyncRetriever(BaseModel):
         ...,
         description="Responsible for fetching the final result `urls` provided by the completed / finished / ready async job.",
     )
+    download_extractor: Optional[
+        Union[CustomRecordExtractor, DpathExtractor, ResponseToFileExtractor]
+    ] = Field(None, description="Responsible for fetching the records from provided urls.")
     creation_requester: Union[CustomRequester, HttpRequester] = Field(
         ...,
         description="Requester component that describes how to prepare HTTP requests to send to the source API to create the async server-side job.",
@@ -1726,6 +1735,20 @@ class AsyncRetriever(BaseModel):
         description="Component decoding the response so records can be extracted.",
         title="Decoder",
     )
+    download_decoder: Optional[
+        Union[
+            CustomDecoder,
+            JsonDecoder,
+            JsonlDecoder,
+            IterableDecoder,
+            XmlDecoder,
+            GzipJsonDecoder,
+        ]
+    ] = Field(
+        None,
+        description="Component decoding the download response so records can be extracted.",
+        title="Download Decoder",
+    )
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -270,6 +270,9 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     RequestPath as RequestPathModel,
 )
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    ResponseToFileExtractor as ResponseToFileExtractorModel,
+)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     SelectiveAuthenticator as SelectiveAuthenticatorModel,
 )
@@ -427,6 +430,7 @@ def _init_mappings(self) -> None:
             DefaultErrorHandlerModel: self.create_default_error_handler,
             DefaultPaginatorModel: self.create_default_paginator,
             DpathExtractorModel: self.create_dpath_extractor,
+            ResponseToFileExtractorModel: self.create_response_to_file_extractor,
             ExponentialBackoffStrategyModel: self.create_exponential_backoff_strategy,
             SessionTokenAuthenticatorModel: self.create_session_token_authenticator,
             HttpRequesterModel: self.create_http_requester,
@@ -1447,6 +1451,16 @@ def create_dpath_extractor(
             parameters=model.parameters or {},
         )
 
+    def create_response_to_file_extractor(
+        self,
+        model: ResponseToFileExtractorModel,
+        needs_decompression: Optional[bool] = True,
+        **kwargs: Any,
+    ):
+        return ResponseToFileExtractor(
+            needs_decompression=needs_decompression, parameters=model.parameters or {}
+        )
+
     @staticmethod
     def create_exponential_backoff_strategy(
         model: ExponentialBackoffStrategyModel, config: Config
@@ -2025,16 +2039,33 @@ def create_async_retriever(
             name=f"job polling - {name}",
         )
         job_download_components_name = f"job download - {name}"
+        download_decoder = (
+            self._create_component_from_model(model=model.download_decoder, config=config)
+            if model.download_decoder
+            else JsonDecoder(parameters={})
+        )
+        download_extractor = (
+            self._create_component_from_model(
+                model=model.download_extractor,
+                config=config,
+                decoder=download_decoder,
+                parameters=model.parameters,
+            )
+            if model.download_extractor
+            else DpathExtractor(
+                [], config=config, decoder=download_decoder, parameters=model.parameters
+            )
+        )
         download_requester = self._create_component_from_model(
             model=model.download_requester,
-            decoder=decoder,
+            decoder=download_decoder,
             config=config,
             name=job_download_components_name,
         )
         download_retriever = SimpleRetriever(
             requester=download_requester,
             record_selector=RecordSelector(
-                extractor=ResponseToFileExtractor(),
+                extractor=download_extractor,
                 record_filter=None,
                 transformations=[],
                 schema_normalization=TypeTransformer(TransformConfig.NoTransform),

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -4,11 +4,12 @@
 
 # mypy: ignore-errors
 import datetime
-from typing import Any, Mapping
+from typing import Any, Iterable, Mapping
 
 import freezegun
 import pendulum
 import pytest
+import requests
 
 from airbyte_cdk import AirbyteTracedException
 from airbyte_cdk.models import FailureType, Level
@@ -27,6 +28,7 @@
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.decoders import JsonDecoder, PaginationDecoderDecorator
 from airbyte_cdk.sources.declarative.extractors import DpathExtractor, RecordFilter, RecordSelector
+from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor
 from airbyte_cdk.sources.declarative.extractors.record_filter import (
     ClientSideIncrementalRecordFilterDecorator,
 )
@@ -47,6 +49,9 @@
 from airbyte_cdk.sources.declarative.models import (
     CustomPartitionRouter as CustomPartitionRouterModel,
 )
+from airbyte_cdk.sources.declarative.models import (
+    CustomRecordExtractor as CustomRecordExtractorModel,
+)
 from airbyte_cdk.sources.declarative.models import CustomSchemaLoader as CustomSchemaLoaderModel
 from airbyte_cdk.sources.declarative.models import DatetimeBasedCursor as DatetimeBasedCursorModel
 from airbyte_cdk.sources.declarative.models import DeclarativeStream as DeclarativeStreamModel
@@ -3270,3 +3275,20 @@ def test_create_concurrent_cursor_uses_min_max_datetime_format_if_defined():
         "state_type": "date-range",
         "legacy": {},
     }
+
+
+class CustomRecordExtractor(RecordExtractor):
+    def extract_records(
+        self,
+        response: requests.Response,
+    ) -> Iterable[Mapping[str, Any]]:
+        yield from response.json()
+
+
+def test_create_custom_record_extractor():
+    definition = {
+        "type": "CustomRecordExtractor",
+        "class_name": "unit_tests.sources.declarative.parsers.test_model_to_component_factory.CustomRecordExtractor",
+    }
+    component = factory.create_component(CustomRecordExtractorModel, definition, {})
+    assert isinstance(component, CustomRecordExtractor)