airbyte-cdk 0.0.0.dev0__py3-none-any.whl

airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py

@@ -0,0 +1,176 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+import logging
+import os
+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
+import requests
+from numpy import nan
+
+from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor
+
+DEFAULT_ENCODING: str = "utf-8"
+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
+    a tradeoff.
+
+    Eventually, we want to support multiple file type by re-using the file based CDK parsers if possible. However, the lift is too high for
+    a first iteration so we will only support CSV parsing using pandas as salesforce and sendgrid were doing.
+    """
+
+    parameters: InitVar[Mapping[str, Any]]
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        self.logger = logging.getLogger("airbyte")
+
+    def _get_response_encoding(self, headers: Dict[str, Any]) -> str:
+        """
+        Get the encoding of the response based on the provided headers. This method is heavily inspired by the requests library
+        implementation.
+
+        Args:
+            headers (Dict[str, Any]): The headers of the response.
+        Returns:
+            str: The encoding of the response.
+        """
+
+        content_type = headers.get("content-type")
+
+        if not content_type:
+            return DEFAULT_ENCODING
+
+        content_type, params = requests.utils.parse_header_links(content_type)
+
+        if "charset" in params:
+            return params["charset"].strip("'\"")  # type: ignore  # we assume headers are returned as str
+
+        return DEFAULT_ENCODING
+
+    def _filter_null_bytes(self, b: bytes) -> bytes:
+        """
+        Filter out null bytes from a bytes object.
+
+        Args:
+            b (bytes): The input bytes object.
+        Returns:
+            bytes: The filtered bytes object with null bytes removed.
+
+        Referenced Issue:
+            https://github.com/airbytehq/airbyte/issues/8300
+        """
+
+        res = b.replace(b"\x00", b"")
+        if len(res) < len(b):
+            self.logger.warning(
+                "Filter 'null' bytes from string, size reduced %d -> %d chars", len(b), len(res)
+            )
+        return res
+
+    def _save_to_file(self, response: requests.Response) -> Tuple[str, str]:
+        """
+        Saves the binary data from the given response to a temporary file and returns the filepath and response encoding.
+
+        Args:
+            response (Optional[requests.Response]): The response object containing the binary data. Defaults to None.
+
+        Returns:
+            Tuple[str, str]: A tuple containing the filepath of the temporary file and the response encoding.
+
+        Raises:
+            ValueError: If the temporary file does not exist after saving the binary data.
+        """
+        # 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:
+                        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):
+            return tmp_file, response_encoding
+        else:
+            raise ValueError(
+                f"The IO/Error occured while verifying binary data. Tmp file {tmp_file} doesn't exist."
+            )
+
+    def _read_with_chunks(
+        self, path: str, file_encoding: str, chunk_size: int = 100
+    ) -> Iterable[Mapping[str, Any]]:
+        """
+        Reads data from a file in chunks and yields each row as a dictionary.
+
+        Args:
+            path (str): The path to the file to be read.
+            file_encoding (str): The encoding of the file.
+            chunk_size (int, optional): The size of each chunk to be read. Defaults to 100.
+
+        Yields:
+            Mapping[str, Any]: A dictionary representing each row of data.
+
+        Raises:
+            ValueError: If an IO/Error occurs while reading the temporary data.
+        """
+
+        try:
+            with open(path, "r", encoding=file_encoding) as data:
+                chunks = pd.read_csv(
+                    data, chunksize=chunk_size, iterator=True, dialect="unix", dtype=object
+                )
+                for chunk in chunks:
+                    chunk = chunk.replace({nan: None}).to_dict(orient="records")
+                    for row in chunk:
+                        yield row
+        except pd.errors.EmptyDataError as e:
+            self.logger.info(f"Empty data received. {e}")
+            yield from []
+        except IOError as ioe:
+            raise ValueError(f"The IO/Error occured while reading tmp data. Called: {path}", ioe)
+        finally:
+            # remove binary tmp file, after data is read
+            os.remove(path)
+
+    def extract_records(
+        self, response: Optional[requests.Response] = None
+    ) -> Iterable[Mapping[str, Any]]:
+        """
+        Extracts records from the given response by:
+            1) Saving the result to a tmp file
+            2) Reading from saved file by chunks to avoid OOM
+
+        Args:
+            response (Optional[requests.Response]): The response object containing the data. Defaults to None.
+
+        Yields:
+            Iterable[Mapping[str, Any]]: An iterable of mappings representing the extracted records.
+
+        Returns:
+            None
+        """
+        if response:
+            file_path, encoding = self._save_to_file(response)
+            yield from self._read_with_chunks(file_path, encoding)
+        else:
+            yield from []

airbyte_cdk/sources/declarative/extractors/type_transformer.py

@@ -0,0 +1,55 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Dict, Mapping
+
+
+@dataclass
+class TypeTransformer(ABC):
+    """
+    Abstract base class for implementing type transformation logic.
+
+    This class provides a blueprint for defining custom transformations
+    on data records based on a provided schema. Implementing classes
+    must override the `transform` method to specify the transformation
+    logic.
+
+    Attributes:
+        None explicitly defined, as this is a dataclass intended to be
+        subclassed.
+
+    Methods:
+        transform(record: Dict[str, Any], schema: Mapping[str, Any]) -> None:
+            Abstract method that must be implemented by subclasses.
+            It performs a transformation on a given data record based
+            on the provided schema.
+
+    Usage:
+        To use this class, create a subclass that implements the
+        `transform` method with the desired transformation logic.
+    """
+
+    @abstractmethod
+    def transform(
+        self,
+        record: Dict[str, Any],
+        schema: Mapping[str, Any],
+    ) -> None:
+        """
+        Perform a transformation on a data record based on a given schema.
+
+        Args:
+            record (Dict[str, Any]): The data record to be transformed.
+            schema (Mapping[str, Any]): The schema that dictates how
+                the record should be transformed.
+
+        Returns:
+            None
+
+        Raises:
+            NotImplementedError: If the method is not implemented
+                by a subclass.
+        """

airbyte_cdk/sources/declarative/incremental/__init__.py

@@ -0,0 +1,37 @@
+#
+# Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+#
+
+from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import (
+    ConcurrentCursorFactory,
+    ConcurrentPerPartitionCursor,
+)
+from airbyte_cdk.sources.declarative.incremental.datetime_based_cursor import DatetimeBasedCursor
+from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
+from airbyte_cdk.sources.declarative.incremental.global_substream_cursor import (
+    GlobalSubstreamCursor,
+)
+from airbyte_cdk.sources.declarative.incremental.per_partition_cursor import (
+    CursorFactory,
+    PerPartitionCursor,
+)
+from airbyte_cdk.sources.declarative.incremental.per_partition_with_global import (
+    PerPartitionWithGlobalCursor,
+)
+from airbyte_cdk.sources.declarative.incremental.resumable_full_refresh_cursor import (
+    ChildPartitionResumableFullRefreshCursor,
+    ResumableFullRefreshCursor,
+)
+
+__all__ = [
+    "CursorFactory",
+    "ConcurrentCursorFactory",
+    "ConcurrentPerPartitionCursor",
+    "DatetimeBasedCursor",
+    "DeclarativeCursor",
+    "GlobalSubstreamCursor",
+    "PerPartitionCursor",
+    "PerPartitionWithGlobalCursor",
+    "ResumableFullRefreshCursor",
+    "ChildPartitionResumableFullRefreshCursor",
+]