diff --git a/airbyte_cdk/sources/file_based/config/excel_format.py b/airbyte_cdk/sources/file_based/config/excel_format.py index 632a0bc387..2f2688030a 100644 --- a/airbyte_cdk/sources/file_based/config/excel_format.py +++ b/airbyte_cdk/sources/file_based/config/excel_format.py @@ -2,11 +2,19 @@ # Copyright (c) 2024 Airbyte, Inc., all rights reserved. # -from pydantic.v1 import BaseModel, Field +from enum import Enum +from typing import Any, Dict, List, Optional + +from pydantic.v1 import BaseModel, Field, root_validator from airbyte_cdk.utils.oneof_option_config import OneOfOptionConfig +class SheetsToSync(str, Enum): + FIRST_SHEET_ONLY = "first_sheet_only" + ALL_SHEETS = "all_sheets" + + class ExcelFormat(BaseModel): class Config(OneOfOptionConfig): title = "Excel Format" @@ -16,3 +24,31 @@ class Config(OneOfOptionConfig): "excel", const=True, ) + sheets_to_sync: SheetsToSync = Field( + default=SheetsToSync.FIRST_SHEET_ONLY, + title="Sheets to Sync", + description=( + "Controls which sheets (tabs) of the workbook are synced. " + "First Sheet Only reads only the first sheet. " + "All Sheets reads every sheet and adds _ab_sheet_name to each record." + ), + order=1, + ) + sheet_names: Optional[List[str]] = Field( + default=None, + title="Sheet Names", + description=( + "Optional list of sheet names to sync. When set, only these sheets are read and " + "_ab_sheet_name is added to each record." + ), + order=2, + ) + + @root_validator + def validate_sheet_selection( + cls, + values: Dict[str, Any], # noqa: N805 # Pydantic validators use cls, not self + ) -> Dict[str, Any]: + if values.get("sheet_names") == []: + raise ValueError("Sheet Names cannot be empty.") + return values diff --git a/airbyte_cdk/sources/file_based/file_types/excel_parser.py b/airbyte_cdk/sources/file_based/file_types/excel_parser.py index 93896f14f3..5e7be7906e 100644 --- a/airbyte_cdk/sources/file_based/file_types/excel_parser.py +++ b/airbyte_cdk/sources/file_based/file_types/excel_parser.py @@ -6,12 +6,13 @@ import warnings from io import IOBase from pathlib import Path -from typing import Any, Dict, Iterable, Mapping, Optional, Tuple, Union +from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Union import orjson import pandas as pd from pydantic.v1 import BaseModel +from airbyte_cdk.sources.file_based.config.excel_format import SheetsToSync from airbyte_cdk.sources.file_based.config.file_based_stream_config import ( ExcelFormat, FileBasedStreamConfig, @@ -33,6 +34,7 @@ class ExcelParser(FileTypeParser): ENCODING = None + SHEET_NAME_COLUMN = "_ab_sheet_name" def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[str]]: """ @@ -48,32 +50,21 @@ async def infer_schema( 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. + Infer the schema of the Excel file by examining its contents. """ - - # Validate the format of the config self.validate_format(config.format, logger) + excel_format = self._as_excel_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, logger, file) - 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, - ) + for df in self.parse_dataframes(fp, logger, file, excel_format): + for column, df_type in df.dtypes.items(): + 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: ( @@ -94,33 +85,21 @@ def parse_records( discovered_schema: Optional[Mapping[str, SchemaType]] = None, ) -> Iterable[Dict[str, Any]]: """ - Parses records from an Excel file with fallback error handling. - - 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. + Parse records from an Excel file with fallback error handling. """ - - # Validate the format of the config self.validate_format(config.format, logger) + excel_format = self._as_excel_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, logger, file) - # 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") - ) + for df in self.parse_dataframes(fp, logger, file, excel_format): + # 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 @@ -131,10 +110,7 @@ def parse_records( @property def file_read_mode(self) -> FileReadMode: """ - Returns the file read mode for the Excel file. - - Returns: - FileReadMode: The file read mode (binary). + Return the file read mode for Excel files. """ return FileReadMode.READ_BINARY @@ -170,18 +146,132 @@ def dtype_to_json_type( @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. + Validate that the format is `ExcelFormat`. """ if not isinstance(excel_format, ExcelFormat): logger.info(f"Expected ExcelFormat, got {excel_format}") raise ConfigValidationError(FileBasedSourceError.CONFIG_VALIDATION_ERROR) + @staticmethod + def _as_excel_format(excel_format: BaseModel, logger: logging.Logger) -> ExcelFormat: + ExcelParser.validate_format(excel_format, logger) + if isinstance(excel_format, ExcelFormat): + return excel_format + raise ConfigValidationError(FileBasedSourceError.CONFIG_VALIDATION_ERROR) + + def parse_dataframes( + self, + fp: Union[IOBase, str, Path], + logger: logging.Logger, + file: RemoteFile, + excel_format: ExcelFormat, + ) -> Iterable[pd.DataFrame]: + if self._uses_legacy_first_sheet_only(excel_format): + yield self.open_and_parse_file(fp, logger, file) + return + + yield from self._parse_sheets(fp, logger, file, excel_format) + + @staticmethod + def _uses_legacy_first_sheet_only(excel_format: ExcelFormat) -> bool: + return ( + not excel_format.sheet_names + and excel_format.sheets_to_sync == SheetsToSync.FIRST_SHEET_ONLY + ) + + def _parse_sheets( + self, + fp: Union[IOBase, str, Path], + logger: logging.Logger, + file: RemoteFile, + excel_format: ExcelFormat, + ) -> Iterable[pd.DataFrame]: + try: + dataframes = list(self._parse_sheets_with_calamine(fp, logger, file, excel_format)) + except ExcelCalamineParsingError: + dataframes = list(self._parse_sheets_with_openpyxl(fp, logger, file, excel_format)) + yield from dataframes + + def _parse_sheets_with_calamine( + self, + fp: Union[IOBase, str, Path], + logger: logging.Logger, + file: RemoteFile, + excel_format: ExcelFormat, + ) -> Iterable[pd.DataFrame]: + try: + with pd.ExcelFile(fp, engine="calamine") as workbook: # type: ignore [arg-type] + sheet_names = self._resolve_sheet_names(workbook, excel_format) + yield from self._parse_sheets_from_workbook(workbook, sheet_names, file) + except BaseException as exc: + if "ValueError" in str(exc): + logger.warning( + f"Calamine parsing failed for {file.file_uri_for_logging}, falling back to openpyxl: {exc}" + ) + raise ExcelCalamineParsingError( + f"Calamine engine failed to parse {file.file_uri_for_logging}", + filename=file.uri, + ) from exc + raise exc + + def _parse_sheets_with_openpyxl( + self, + fp: Union[IOBase, str, Path], + logger: logging.Logger, + file: RemoteFile, + excel_format: ExcelFormat, + ) -> Iterable[pd.DataFrame]: + # Some file-like objects are not seekable. + if hasattr(fp, "seek"): + try: + fp.seek(0) # type: ignore [union-attr] + except OSError as exc: + logger.info( + f"Could not rewind stream for {file.file_uri_for_logging}; " + f"proceeding with openpyxl from current position: {exc}" + ) + + with warnings.catch_warnings(record=True) as warning_records: + warnings.simplefilter("always") + with pd.ExcelFile(fp, engine="openpyxl") as workbook: # type: ignore [arg-type] + sheet_names = self._resolve_sheet_names(workbook, excel_format) + dataframes = list(self._parse_sheets_from_workbook(workbook, sheet_names, file)) + + for warning in warning_records: + logger.warning(f"Openpyxl warning for {file.file_uri_for_logging}: {warning.message}") + + yield from dataframes + + def _parse_sheets_from_workbook( + self, + workbook: pd.ExcelFile, + sheet_names: List[str], + file: RemoteFile, + ) -> Iterable[pd.DataFrame]: + workbook_sheet_names = [str(sheet_name) for sheet_name in workbook.sheet_names] + missing_sheet_names = sorted(set(sheet_names) - set(workbook_sheet_names)) + if missing_sheet_names: + raise ValueError( + f"Sheet names {missing_sheet_names} were not found in {file.file_uri_for_logging}" + ) + + for sheet_name in sheet_names: + df = workbook.parse(sheet_name=sheet_name) # type: ignore [call-overload] + if self.SHEET_NAME_COLUMN in df.columns: + raise ValueError( + f"Column {self.SHEET_NAME_COLUMN} is reserved for sheet metadata in {file.file_uri_for_logging}" + ) + df[self.SHEET_NAME_COLUMN] = sheet_name + yield df + + @staticmethod + def _resolve_sheet_names(workbook: pd.ExcelFile, excel_format: ExcelFormat) -> List[str]: + if excel_format.sheet_names: + return excel_format.sheet_names + if excel_format.sheets_to_sync == SheetsToSync.ALL_SHEETS: + return [str(sheet_name) for sheet_name in workbook.sheet_names] + return [str(sheet_name) for sheet_name in workbook.sheet_names[:1]] + def _open_and_parse_file_with_calamine( self, fp: Union[IOBase, str, Path], diff --git a/unit_tests/sources/file_based/file_types/test_excel_parser.py b/unit_tests/sources/file_based/file_types/test_excel_parser.py index 18850e9b0c..dac4af5c63 100644 --- a/unit_tests/sources/file_based/file_types/test_excel_parser.py +++ b/unit_tests/sources/file_based/file_types/test_excel_parser.py @@ -3,6 +3,7 @@ # +import asyncio import datetime import warnings from io import BytesIO @@ -11,12 +12,17 @@ import pandas as pd import pytest +from airbyte_cdk.sources.file_based.config.excel_format import SheetsToSync from airbyte_cdk.sources.file_based.config.file_based_stream_config import ( ExcelFormat, FileBasedStreamConfig, ValidationPolicy, ) -from airbyte_cdk.sources.file_based.exceptions import ConfigValidationError, RecordParseError +from airbyte_cdk.sources.file_based.exceptions import ( + ConfigValidationError, + ExcelCalamineParsingError, + RecordParseError, +) from airbyte_cdk.sources.file_based.file_based_stream_reader import AbstractFileBasedStreamReader from airbyte_cdk.sources.file_based.file_types.excel_parser import ExcelParser from airbyte_cdk.sources.file_based.remote_file import RemoteFile @@ -38,7 +44,7 @@ def file_config(): return FileBasedStreamConfig( name="test.xlsx", file_type="excel", - format=ExcelFormat(sheet_name="Sheet1"), + format=ExcelFormat(), validation_policy=ValidationPolicy.emit_record, ) @@ -52,7 +58,6 @@ def remote_file(): def setup_parser(remote_file): parser = ExcelParser() - # Sample data for the mock Excel file data = pd.DataFrame( { "column1": [1, 2, 3], @@ -62,13 +67,11 @@ def setup_parser(remote_file): } ) - # Convert the DataFrame to an Excel byte stream excel_bytes = BytesIO() with pd.ExcelWriter(excel_bytes, engine="xlsxwriter") as writer: data.to_excel(writer, index=False) excel_bytes.seek(0) - # Mock the stream_reader's open_file method to return the Excel byte stream stream_reader = MagicMock(spec=AbstractFileBasedStreamReader) stream_reader.open_file.return_value = BytesIO(excel_bytes.read()) @@ -83,17 +86,14 @@ def setup_parser(remote_file): @patch("pandas.ExcelFile") -@pytest.mark.asyncio -async def test_infer_schema(mock_excel_file, setup_parser): +def test_infer_schema(mock_excel_file, setup_parser): parser, config, file, stream_reader, logger, data = setup_parser - # Mock the parse method of the pandas ExcelFile object mock_excel_file.return_value.parse.return_value = data - # Call infer_schema - schema = await parser.infer_schema(config, file, stream_reader, logger) + loop = asyncio.get_event_loop() + schema = loop.run_until_complete(parser.infer_schema(config, file, stream_reader, logger)) - # Define the expected schema expected_schema: SchemaType = { "column1": {"type": "number"}, "column2": {"type": "string"}, @@ -101,15 +101,12 @@ async def test_infer_schema(mock_excel_file, setup_parser): "column4": {"type": "string", "format": "date-time"}, } - # Validate the schema assert schema == expected_schema - # Assert that the stream_reader's open_file was called correctly stream_reader.open_file.assert_called_once_with( file, parser.file_read_mode, parser.ENCODING, logger ) - # Assert that the logger was not used for warnings/errors logger.info.assert_not_called() logger.error.assert_not_called() @@ -238,3 +235,144 @@ def seek(self, *args, **kwargs): assert "Could not rewind stream" in msg assert remote_file.file_uri_for_logging in msg mock_excel.assert_called_once_with(fp, engine="openpyxl") + + +def _make_multi_sheet_workbook(sheet_rows): + excel_bytes = BytesIO() + with pd.ExcelWriter(excel_bytes, engine="xlsxwriter") as writer: + for sheet_name, rows in sheet_rows.items(): + pd.DataFrame(rows).to_excel(writer, index=False, sheet_name=sheet_name) + excel_bytes.seek(0) + return excel_bytes.read() + + +def _multi_sheet_setup(remote_file, excel_format, sheet_rows=None): + parser = ExcelParser() + workbook = _make_multi_sheet_workbook( + sheet_rows + or { + "People": [ + {"id": 1, "name": "alice"}, + {"id": 2, "name": "bob"}, + ], + "Orders": [ + {"id": 3, "amount": 10.5}, + {"id": 4, "amount": 20.0}, + {"id": 5, "amount": 30.25}, + ], + } + ) + stream_reader = MagicMock(spec=AbstractFileBasedStreamReader) + stream_reader.open_file.side_effect = lambda *args, **kwargs: BytesIO(workbook) + config = FileBasedStreamConfig(name="test_stream", format=excel_format) + return parser, config, remote_file, stream_reader, MagicMock() + + +def test_default_reads_only_first_sheet(remote_file): + parser, config, file, stream_reader, logger = _multi_sheet_setup(remote_file, ExcelFormat()) + + records = list(parser.parse_records(config, file, stream_reader, logger)) + + assert len(records) == 2 + assert {record["name"] for record in records} == {"alice", "bob"} + assert all(parser.SHEET_NAME_COLUMN not in record for record in records) + + +@pytest.mark.parametrize( + "excel_format, expected_sheets, expected_count", + [ + pytest.param( + ExcelFormat(sheets_to_sync=SheetsToSync.ALL_SHEETS), + {"People", "Orders"}, + 5, + id="all-sheets", + ), + pytest.param( + ExcelFormat(sheet_names=["Orders"]), + {"Orders"}, + 3, + id="explicit-sheet-names", + ), + pytest.param( + ExcelFormat(sheets_to_sync=SheetsToSync.FIRST_SHEET_ONLY, sheet_names=["Orders"]), + {"Orders"}, + 3, + id="explicit-sheet-names-take-precedence", + ), + ], +) +def test_multi_sheet_records_include_sheet_name( + remote_file, excel_format, expected_sheets, expected_count +): + parser, config, file, stream_reader, logger = _multi_sheet_setup(remote_file, excel_format) + + records = list(parser.parse_records(config, file, stream_reader, logger)) + + assert len(records) == expected_count + assert {record[parser.SHEET_NAME_COLUMN] for record in records} == expected_sheets + + +def test_multi_sheet_calamine_fallback_does_not_duplicate_partial_results(remote_file): + parser = ExcelParser() + fp = BytesIO(b"test") + logger = MagicMock() + excel_format = ExcelFormat(sheets_to_sync=SheetsToSync.ALL_SHEETS) + calamine_df = pd.DataFrame({"id": [1], ExcelParser.SHEET_NAME_COLUMN: ["Calamine"]}) + fallback_df = pd.DataFrame({"id": [2], ExcelParser.SHEET_NAME_COLUMN: ["Openpyxl"]}) + + def parse_with_calamine(*args, **kwargs): + yield calamine_df + raise ExcelCalamineParsingError("calamine failed", filename=remote_file.uri) + + with ( + patch.object(parser, "_parse_sheets_with_calamine", side_effect=parse_with_calamine), + patch.object(parser, "_parse_sheets_with_openpyxl", return_value=[fallback_df]), + ): + dataframes = list(parser._parse_sheets(fp, logger, remote_file, excel_format)) + + assert dataframes == [fallback_df] + + +def test_all_sheets_schema_merges_columns(remote_file): + parser, config, file, stream_reader, logger = _multi_sheet_setup( + remote_file, ExcelFormat(sheets_to_sync=SheetsToSync.ALL_SHEETS) + ) + + loop = asyncio.get_event_loop() + schema = loop.run_until_complete(parser.infer_schema(config, file, stream_reader, logger)) + + assert schema == { + "id": {"type": "number"}, + "name": {"type": "string"}, + "amount": {"type": "number"}, + parser.SHEET_NAME_COLUMN: {"type": "string"}, + } + + +def test_default_schema_only_first_sheet(remote_file): + parser, config, file, stream_reader, logger = _multi_sheet_setup(remote_file, ExcelFormat()) + + loop = asyncio.get_event_loop() + schema = loop.run_until_complete(parser.infer_schema(config, file, stream_reader, logger)) + + assert schema == {"id": {"type": "number"}, "name": {"type": "string"}} + + +def test_missing_explicit_sheet_name_raises_record_parse_error(remote_file): + parser, config, file, stream_reader, logger = _multi_sheet_setup( + remote_file, ExcelFormat(sheet_names=["Missing"]) + ) + + with pytest.raises(RecordParseError): + list(parser.parse_records(config, file, stream_reader, logger)) + + +def test_reserved_sheet_name_column_raises_record_parse_error(remote_file): + parser, config, file, stream_reader, logger = _multi_sheet_setup( + remote_file, + ExcelFormat(sheets_to_sync=SheetsToSync.ALL_SHEETS), + sheet_rows={"People": [{"id": 1, ExcelParser.SHEET_NAME_COLUMN: "source value"}]}, + ) + + with pytest.raises(RecordParseError): + list(parser.parse_records(config, file, stream_reader, logger)) diff --git a/unit_tests/sources/file_based/scenarios/csv_scenarios.py b/unit_tests/sources/file_based/scenarios/csv_scenarios.py index f16d83e200..31bcfae64f 100644 --- a/unit_tests/sources/file_based/scenarios/csv_scenarios.py +++ b/unit_tests/sources/file_based/scenarios/csv_scenarios.py @@ -472,7 +472,21 @@ "default": "excel", "const": "excel", "type": "string", - } + }, + "sheets_to_sync": { + "title": "Sheets to Sync", + "description": "Controls which sheets (tabs) of the workbook are synced. First Sheet Only reads only the first sheet. All Sheets reads every sheet and adds _ab_sheet_name to each record.", + "default": "first_sheet_only", + "enum": ["first_sheet_only", "all_sheets"], + "order": 1, + }, + "sheet_names": { + "title": "Sheet Names", + "description": "Optional list of sheet names to sync. When set, only these sheets are read and _ab_sheet_name is added to each record.", + "order": 2, + "type": "array", + "items": {"type": "string"}, + }, }, "required": ["filetype"], },