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

airbyte_cdk/sources/file_based/file_types/excel_parser.py

@@ -0,0 +1,196 @@
+#
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+from io import IOBase
+from pathlib import Path
+from typing import Any, Dict, Iterable, Mapping, Optional, Tuple, Union
+
+import orjson
+import pandas as pd
+from numpy import datetime64, issubdtype
+from numpy import dtype as dtype_
+from pydantic.v1 import BaseModel
+
+from airbyte_cdk.sources.file_based.config.file_based_stream_config import (
+    ExcelFormat,
+    FileBasedStreamConfig,
+)
+from airbyte_cdk.sources.file_based.exceptions import (
+    ConfigValidationError,
+    FileBasedSourceError,
+    RecordParseError,
+)
+from airbyte_cdk.sources.file_based.file_based_stream_reader import (
+    AbstractFileBasedStreamReader,
+    FileReadMode,
+)
+from airbyte_cdk.sources.file_based.file_types.file_type_parser import FileTypeParser
+from airbyte_cdk.sources.file_based.remote_file import RemoteFile
+from airbyte_cdk.sources.file_based.schema_helpers import SchemaType
+
+
+class ExcelParser(FileTypeParser):
+    ENCODING = None
+
+    def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[str]]:
+        """
+        ExcelParser does not require config checks, implicit pydantic validation is enough.
+        """
+        return True, None
+
+    async def infer_schema(
+        self,
+        config: FileBasedStreamConfig,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+    ) -> SchemaType:
+        """
+        Infers the schema of the Excel file by examining its contents.
+
+        Args:
+            config (FileBasedStreamConfig): Configuration for the file-based stream.
+            file (RemoteFile): The remote file to be read.
+            stream_reader (AbstractFileBasedStreamReader): Reader to read the file.
+            logger (logging.Logger): Logger for logging information and errors.
+
+        Returns:
+            SchemaType: Inferred schema of the Excel file.
+        """
+
+        # Validate the format of the config
+        self.validate_format(config.format, logger)
+
+        fields: Dict[str, str] = {}
+
+        with stream_reader.open_file(file, self.file_read_mode, self.ENCODING, logger) as fp:
+            df = self.open_and_parse_file(fp)
+            for column, df_type in df.dtypes.items():
+                # Choose the broadest data type if the column's data type differs in dataframes
+                prev_frame_column_type = fields.get(column)  # type: ignore [call-overload]
+                fields[column] = self.dtype_to_json_type(  # type: ignore [index]
+                    prev_frame_column_type,
+                    df_type,
+                )
+
+        schema = {
+            field: (
+                {"type": "string", "format": "date-time"}
+                if fields[field] == "date-time"
+                else {"type": fields[field]}
+            )
+            for field in fields
+        }
+        return schema
+
+    def parse_records(
+        self,
+        config: FileBasedStreamConfig,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+        discovered_schema: Optional[Mapping[str, SchemaType]] = None,
+    ) -> Iterable[Dict[str, Any]]:
+        """
+        Parses records from an Excel file based on the provided configuration.
+
+        Args:
+            config (FileBasedStreamConfig): Configuration for the file-based stream.
+            file (RemoteFile): The remote file to be read.
+            stream_reader (AbstractFileBasedStreamReader): Reader to read the file.
+            logger (logging.Logger): Logger for logging information and errors.
+            discovered_schema (Optional[Mapping[str, SchemaType]]): Discovered schema for validation.
+
+        Yields:
+            Iterable[Dict[str, Any]]: Parsed records from the Excel file.
+        """
+
+        # Validate the format of the config
+        self.validate_format(config.format, logger)
+
+        try:
+            # Open and parse the file using the stream reader
+            with stream_reader.open_file(file, self.file_read_mode, self.ENCODING, logger) as fp:
+                df = self.open_and_parse_file(fp)
+                # Yield records as dictionaries
+                # DataFrame.to_dict() method returns datetime values in pandas.Timestamp values, which are not serializable by orjson
+                # DataFrame.to_json() returns string with datetime values serialized to iso8601 with microseconds to align with pydantic behavior
+                # see PR description: https://github.com/airbytehq/airbyte/pull/44444/
+                yield from orjson.loads(
+                    df.to_json(orient="records", date_format="iso", date_unit="us")
+                )
+
+        except Exception as exc:
+            # Raise a RecordParseError if any exception occurs during parsing
+            raise RecordParseError(
+                FileBasedSourceError.ERROR_PARSING_RECORD, filename=file.uri
+            ) from exc
+
+    @property
+    def file_read_mode(self) -> FileReadMode:
+        """
+        Returns the file read mode for the Excel file.
+
+        Returns:
+            FileReadMode: The file read mode (binary).
+        """
+        return FileReadMode.READ_BINARY
+
+    @staticmethod
+    def dtype_to_json_type(
+        current_type: Optional[str],
+        dtype: dtype_,  # type: ignore [type-arg]
+    ) -> str:
+        """
+        Convert Pandas DataFrame types to Airbyte Types.
+
+        Args:
+            current_type (Optional[str]): One of the previous types based on earlier dataframes.
+            dtype: Pandas DataFrame type.
+
+        Returns:
+            str: Corresponding Airbyte Type.
+        """
+        number_types = ("int64", "float64")
+        if current_type == "string":
+            # Previous column values were of the string type, no need to look further.
+            return current_type
+        if dtype is object:
+            return "string"
+        if dtype in number_types and (not current_type or current_type == "number"):
+            return "number"
+        if dtype == "bool" and (not current_type or current_type == "boolean"):
+            return "boolean"
+        if issubdtype(dtype, datetime64):
+            return "date-time"
+        return "string"
+
+    @staticmethod
+    def validate_format(excel_format: BaseModel, logger: logging.Logger) -> None:
+        """
+        Validates if the given format is of type ExcelFormat.
+
+        Args:
+            excel_format (Any): The format to be validated.
+
+        Raises:
+            ConfigValidationError: If the format is not ExcelFormat.
+        """
+        if not isinstance(excel_format, ExcelFormat):
+            logger.info(f"Expected ExcelFormat, got {excel_format}")
+            raise ConfigValidationError(FileBasedSourceError.CONFIG_VALIDATION_ERROR)
+
+    @staticmethod
+    def open_and_parse_file(fp: Union[IOBase, str, Path]) -> pd.DataFrame:
+        """
+        Opens and parses the Excel file.
+
+        Args:
+            fp: File pointer to the Excel file.
+
+        Returns:
+            pd.DataFrame: Parsed data from the Excel file.
+        """
+        return pd.ExcelFile(fp, engine="calamine").parse()  # type: ignore [arg-type, call-overload, no-any-return]

airbyte_cdk/sources/file_based/file_types/file_transfer.py

@@ -0,0 +1,30 @@
+#
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+#
+import logging
+from typing import Iterable, Tuple
+
+from airbyte_cdk.models import AirbyteRecordMessageFileReference
+from airbyte_cdk.sources.file_based.file_based_stream_reader import AbstractFileBasedStreamReader
+from airbyte_cdk.sources.file_based.file_record_data import FileRecordData
+from airbyte_cdk.sources.file_based.remote_file import RemoteFile
+from airbyte_cdk.sources.utils.files_directory import get_files_directory
+
+
+class FileTransfer:
+    def __init__(self) -> None:
+        self._local_directory = get_files_directory()
+
+    def upload(
+        self,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+    ) -> Iterable[Tuple[FileRecordData, AirbyteRecordMessageFileReference]]:
+        try:
+            yield stream_reader.upload(
+                file=file, local_directory=self._local_directory, logger=logger
+            )
+        except Exception as ex:
+            logger.error("An error has occurred while getting file: %s", str(ex))
+            raise ex

airbyte_cdk/sources/file_based/file_types/file_type_parser.py

@@ -0,0 +1,86 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Iterable, Mapping, Optional, Tuple
+
+from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
+from airbyte_cdk.sources.file_based.file_based_stream_reader import (
+    AbstractFileBasedStreamReader,
+    FileReadMode,
+)
+from airbyte_cdk.sources.file_based.remote_file import RemoteFile
+from airbyte_cdk.sources.file_based.schema_helpers import SchemaType
+
+Record = Dict[str, Any]
+
+
+class FileTypeParser(ABC):
+    """
+    An abstract class containing methods that must be implemented for each
+    supported file type.
+    """
+
+    @property
+    def parser_max_n_files_for_schema_inference(self) -> Optional[int]:
+        """
+        The discovery policy decides how many files are loaded for schema inference. This method can provide a parser-specific override. If it's defined, the smaller of the two values will be used.
+        """
+        return None
+
+    @property
+    def parser_max_n_files_for_parsability(self) -> Optional[int]:
+        """
+        The availability policy decides how many files are loaded for checking whether parsing works correctly. This method can provide a parser-specific override. If it's defined, the smaller of the two values will be used.
+        """
+        return None
+
+    def get_parser_defined_primary_key(self, config: FileBasedStreamConfig) -> Optional[str]:
+        """
+        The parser can define a primary key. If no user-defined primary key is provided, this will be used.
+        """
+        return None
+
+    @abstractmethod
+    def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[str]]:
+        """
+        Check whether the config is valid for this file type. If it is, return True and None. If it's not, return False and an error message explaining why it's invalid.
+        """
+        return True, None
+
+    @abstractmethod
+    async def infer_schema(
+        self,
+        config: FileBasedStreamConfig,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+    ) -> SchemaType:
+        """
+        Infer the JSON Schema for this file.
+        """
+        ...
+
+    @abstractmethod
+    def parse_records(
+        self,
+        config: FileBasedStreamConfig,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+        discovered_schema: Optional[Mapping[str, SchemaType]],
+    ) -> Iterable[Record]:
+        """
+        Parse and emit each record.
+        """
+        ...
+
+    @property
+    @abstractmethod
+    def file_read_mode(self) -> FileReadMode:
+        """
+        The mode in which the file should be opened for reading.
+        """
+        ...

airbyte_cdk/sources/file_based/file_types/jsonl_parser.py

@@ -0,0 +1,145 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import json
+import logging
+from typing import Any, Dict, Iterable, Mapping, Optional, Tuple, Union
+
+import orjson
+
+from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
+from airbyte_cdk.sources.file_based.exceptions import FileBasedSourceError, RecordParseError
+from airbyte_cdk.sources.file_based.file_based_stream_reader import (
+    AbstractFileBasedStreamReader,
+    FileReadMode,
+)
+from airbyte_cdk.sources.file_based.file_types.file_type_parser import FileTypeParser
+from airbyte_cdk.sources.file_based.remote_file import RemoteFile
+from airbyte_cdk.sources.file_based.schema_helpers import (
+    PYTHON_TYPE_MAPPING,
+    SchemaType,
+    merge_schemas,
+)
+
+
+class JsonlParser(FileTypeParser):
+    MAX_BYTES_PER_FILE_FOR_SCHEMA_INFERENCE = 1_000_000
+    ENCODING = "utf8"
+
+    def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[str]]:
+        """
+        JsonlParser does not require config checks, implicit pydantic validation is enough.
+        """
+        return True, None
+
+    async def infer_schema(
+        self,
+        config: FileBasedStreamConfig,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+    ) -> SchemaType:
+        """
+        Infers the schema for the file by inferring the schema for each line, and merging
+        it with the previously-inferred schema.
+        """
+        inferred_schema: Mapping[str, Any] = {}
+
+        for entry in self._parse_jsonl_entries(file, stream_reader, logger, read_limit=True):
+            line_schema = self._infer_schema_for_record(entry)
+            inferred_schema = merge_schemas(inferred_schema, line_schema)
+
+        return inferred_schema
+
+    def parse_records(
+        self,
+        config: FileBasedStreamConfig,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+        discovered_schema: Optional[Mapping[str, SchemaType]],
+    ) -> Iterable[Dict[str, Any]]:
+        """
+        This code supports parsing json objects over multiple lines even though this does not align with the JSONL format. This is for
+        backward compatibility reasons i.e. the previous source-s3 parser did support this. The drawback is:
+        * performance as the way we support json over multiple lines is very brute forced
+        * given that we don't have `newlines_in_values` config to scope the possible inputs, we might parse the whole file before knowing if
+          the input is improperly formatted or if the json is over multiple lines
+
+        The goal is to run the V4 of source-s3 in production, track the warning log emitted when there are multiline json objects and
+        deprecate this feature if it's not a valid use case.
+        """
+        yield from self._parse_jsonl_entries(file, stream_reader, logger)
+
+    @classmethod
+    def _infer_schema_for_record(cls, record: Dict[str, Any]) -> Dict[str, Any]:
+        record_schema = {}
+        for key, value in record.items():
+            if value is None:
+                record_schema[key] = {"type": "null"}
+            else:
+                record_schema[key] = {"type": PYTHON_TYPE_MAPPING[type(value)]}
+
+        return record_schema
+
+    @property
+    def file_read_mode(self) -> FileReadMode:
+        return FileReadMode.READ
+
+    def _parse_jsonl_entries(
+        self,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+        read_limit: bool = False,
+    ) -> Iterable[Dict[str, Any]]:
+        with stream_reader.open_file(file, self.file_read_mode, self.ENCODING, logger) as fp:
+            read_bytes = 0
+
+            had_json_parsing_error = False
+            has_warned_for_multiline_json_object = False
+            yielded_at_least_once = False
+
+            accumulator = None
+            for line in fp:
+                if not accumulator:
+                    accumulator = self._instantiate_accumulator(line)
+                read_bytes += len(line)
+                accumulator += line  # type: ignore [operator]  # In reality, it's either bytes or string and we add the same type
+                try:
+                    record = orjson.loads(accumulator)
+                    if had_json_parsing_error and not has_warned_for_multiline_json_object:
+                        logger.warning(
+                            f"File at {file.uri} is using multiline JSON. Performance could be greatly reduced"
+                        )
+                        has_warned_for_multiline_json_object = True
+
+                    yield record
+                    yielded_at_least_once = True
+                    accumulator = self._instantiate_accumulator(line)
+                except orjson.JSONDecodeError:
+                    had_json_parsing_error = True
+
+                if (
+                    read_limit
+                    and yielded_at_least_once
+                    and read_bytes >= self.MAX_BYTES_PER_FILE_FOR_SCHEMA_INFERENCE
+                ):
+                    logger.warning(
+                        f"Exceeded the maximum number of bytes per file for schema inference ({self.MAX_BYTES_PER_FILE_FOR_SCHEMA_INFERENCE}). "
+                        f"Inferring schema from an incomplete set of records."
+                    )
+                    break
+
+            if had_json_parsing_error and not yielded_at_least_once:
+                raise RecordParseError(
+                    FileBasedSourceError.ERROR_PARSING_RECORD, filename=file.uri, lineno=line
+                )
+
+    @staticmethod
+    def _instantiate_accumulator(line: Union[bytes, str]) -> Union[bytes, str]:
+        if isinstance(line, bytes):
+            return bytes("", json.detect_encoding(line))
+        elif isinstance(line, str):
+            return ""