diff --git a/airbyte_cdk/sources/file_based/config/excel_format.py b/airbyte_cdk/sources/file_based/config/excel_format.py index 632a0bc387..038f8605a4 100644 --- a/airbyte_cdk/sources/file_based/config/excel_format.py +++ b/airbyte_cdk/sources/file_based/config/excel_format.py @@ -2,6 +2,8 @@ # Copyright (c) 2024 Airbyte, Inc., all rights reserved. # +from typing import Optional + from pydantic.v1 import BaseModel, Field from airbyte_cdk.utils.oneof_option_config import OneOfOptionConfig @@ -16,3 +18,13 @@ class Config(OneOfOptionConfig): "excel", const=True, ) + sheet_name: Optional[str] = Field( + None, + title="Sheet Name", + description=( + "The worksheet to read from each workbook. Leave empty to read only the first " + 'worksheet. Enter an exact worksheet name to read that worksheet, or "*" to read ' + "every worksheet in the workbook. Worksheet names are case-sensitive." + ), + examples=["Sheet1", "*"], + ) 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..750ddd718b 100644 --- a/airbyte_cdk/sources/file_based/file_types/excel_parser.py +++ b/airbyte_cdk/sources/file_based/file_types/excel_parser.py @@ -6,7 +6,7 @@ import warnings from io import IOBase from pathlib import Path -from typing import Any, Dict, Iterable, Mapping, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterable, Iterator, Mapping, Optional, Tuple, Union import orjson import pandas as pd @@ -33,6 +33,8 @@ class ExcelParser(FileTypeParser): ENCODING = None + ALL_SHEETS = "*" + ab_sheet_name_col = "_ab_source_sheet_name" def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[str]]: """ @@ -62,18 +64,23 @@ async def infer_schema( # Validate the format of the config self.validate_format(config.format, logger) + excel_format = config.format + if not isinstance(excel_format, ExcelFormat): + raise ConfigValidationError(FileBasedSourceError.CONFIG_VALIDATION_ERROR) 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._iter_worksheets(fp, excel_format, logger, file): + 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, + ) + + if self._reads_all_sheets(excel_format): + fields[self.ab_sheet_name_col] = "string" schema = { field: ( @@ -109,19 +116,29 @@ def parse_records( # Validate the format of the config self.validate_format(config.format, logger) + excel_format = config.format + if not isinstance(excel_format, ExcelFormat): + raise ConfigValidationError(FileBasedSourceError.CONFIG_VALIDATION_ERROR) + + include_sheet_name = self._reads_all_sheets(excel_format) 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 sheet, df in self._iter_worksheets(fp, excel_format, logger, file): + # DataFrame.to_dict() returns pandas.Timestamp values not serializable by orjson. + # DataFrame.to_json() serializes datetimes to iso8601 with microseconds. + records = orjson.loads( + df.to_json(orient="records", date_format="iso", date_unit="us") + ) + if include_sheet_name: + for record in records: + record[self.ab_sheet_name_col] = sheet + yield from records + + except ConfigValidationError: + # A missing worksheet is a configuration problem, not an unparseable record. + raise except Exception as exc: # Raise a RecordParseError if any exception occurs during parsing raise RecordParseError( @@ -187,7 +204,8 @@ def _open_and_parse_file_with_calamine( fp: Union[IOBase, str, Path], logger: logging.Logger, file: RemoteFile, - ) -> pd.DataFrame: + sheet_name: Union[int, str, None] = 0, + ) -> Union[pd.DataFrame, Dict[Union[int, str], pd.DataFrame]]: """Opens and parses Excel file using Calamine engine. Args: @@ -202,7 +220,7 @@ def _open_and_parse_file_with_calamine( ExcelCalamineParsingError: If Calamine fails to parse the file. """ try: - return pd.ExcelFile(fp, engine="calamine").parse() # type: ignore [arg-type, call-overload, no-any-return] + return pd.ExcelFile(fp, engine="calamine").parse(sheet_name=sheet_name) # type: ignore [arg-type, call-overload, no-any-return] except BaseException as exc: # Calamine engine raises PanicException(child of BaseException) if Calamine fails to parse the file. # Checking if ValueError in exception arg to know if it was actually an error during parsing due to invalid values in cells. @@ -222,7 +240,8 @@ def _open_and_parse_file_with_openpyxl( fp: Union[IOBase, str, Path], logger: logging.Logger, file: RemoteFile, - ) -> pd.DataFrame: + sheet_name: Union[int, str, None] = 0, + ) -> Union[pd.DataFrame, Dict[Union[int, str], pd.DataFrame]]: """Opens and parses Excel file using Openpyxl engine. Args: @@ -233,6 +252,16 @@ def _open_and_parse_file_with_openpyxl( Returns: pd.DataFrame: Parsed data from the Excel file. """ + self._rewind(fp, logger, file) + return self._with_openpyxl_warnings_logged( # type: ignore [no-any-return] + lambda: pd.ExcelFile(fp, engine="openpyxl").parse(sheet_name=sheet_name), # type: ignore [arg-type, call-overload] + logger, + file, + ) + + @staticmethod + def _rewind(fp: Union[IOBase, str, Path], logger: logging.Logger, file: RemoteFile) -> None: + """Rewinds the stream before handing it to a second engine, where possible.""" # Some file-like objects are not seekable. if hasattr(fp, "seek"): try: @@ -243,21 +272,27 @@ def _open_and_parse_file_with_openpyxl( f"proceeding with openpyxl from current position: {exc}" ) + @staticmethod + def _with_openpyxl_warnings_logged( + call: Callable[[], Any], logger: logging.Logger, file: RemoteFile + ) -> Any: + """Runs an openpyxl call, surfacing the warnings it raises as log lines.""" with warnings.catch_warnings(record=True) as warning_records: warnings.simplefilter("always") - df = pd.ExcelFile(fp, engine="openpyxl").parse() # type: ignore [arg-type, call-overload] + result = call() for warning in warning_records: logger.warning(f"Openpyxl warning for {file.file_uri_for_logging}: {warning.message}") - return df # type: ignore [no-any-return] + return result def open_and_parse_file( self, fp: Union[IOBase, str, Path], logger: logging.Logger, file: RemoteFile, - ) -> pd.DataFrame: + sheet_name: Union[int, str, None] = 0, + ) -> Union[pd.DataFrame, Dict[Union[int, str], pd.DataFrame]]: """Opens and parses the Excel file with Calamine-first and Openpyxl fallback. Args: @@ -269,6 +304,156 @@ def open_and_parse_file( pd.DataFrame: Parsed data from the Excel file. """ try: - return self._open_and_parse_file_with_calamine(fp, logger, file) + return self._open_and_parse_file_with_calamine(fp, logger, file, sheet_name) + except ExcelCalamineParsingError: + return self._open_and_parse_file_with_openpyxl(fp, logger, file, sheet_name) + + def _iter_worksheets( + self, + fp: Union[IOBase, str, Path], + excel_format: ExcelFormat, + logger: logging.Logger, + file: RemoteFile, + ) -> Iterator[Tuple[Union[int, str], pd.DataFrame]]: + """Yields (worksheet name, frame), one worksheet at a time. + + Callers always see the same shape whether one worksheet or all of them were + selected, and the frames arrive lazily so a caller that streams records can let + each one go before the next is parsed. + """ + sheet_name = self._resolve_sheet_name(excel_format) + if sheet_name is not None: + yield self._parse_one_worksheet(fp, excel_format, sheet_name, logger, file) + return + yield from self._iter_every_worksheet(fp, logger, file) + + def _parse_one_worksheet( + self, + fp: Union[IOBase, str, Path], + excel_format: ExcelFormat, + sheet_name: Union[int, str], + logger: logging.Logger, + file: RemoteFile, + ) -> Tuple[Union[int, str], pd.DataFrame]: + """Parses the single worksheet the config selected.""" + try: + parsed = self.open_and_parse_file(fp, logger, file, sheet_name) + except ValueError as exc: + # pandas raises ValueError("Worksheet named 'x' not found") for both engines. + if "not found" not in str(exc): + raise + raise ConfigValidationError( + f"Worksheet {excel_format.sheet_name!r} was not found in the workbook. " + f"{self._describe_available_sheets(fp, logger)}" + "Worksheet names are case-sensitive. " + f'Set the "Sheet Name" option to an exact worksheet name, or to "*" to read ' + "every worksheet.", + filename=file.uri, + ) from exc + return (sheet_name if isinstance(sheet_name, str) else 0), parsed # type: ignore [return-value] + + def _iter_every_worksheet( + self, + fp: Union[IOBase, str, Path], + logger: logging.Logger, + file: RemoteFile, + ) -> Iterator[Tuple[Union[int, str], pd.DataFrame]]: + """Parses each worksheet in turn so only one frame is held at a time. + + Asking pandas for every worksheet at once (`sheet_name=None`) materializes one + frame per worksheet up front, so peak memory tracks the whole workbook rather + than its largest sheet. + + The engine fallback only covers the first worksheet. Once an earlier worksheet's + records are downstream, re-reading the workbook with openpyxl would emit them a + second time, so a later failure is left to propagate. The stream logs it and + retries the file on the next sync, which is how every other streaming parser + already behaves on a mid-file error. + """ + try: + yield from self._iter_worksheets_with_calamine(fp, logger, file) + return except ExcelCalamineParsingError: - return self._open_and_parse_file_with_openpyxl(fp, logger, file) + pass + yield from self._iter_worksheets_with_openpyxl(fp, logger, file) + + def _iter_worksheets_with_calamine( + self, + fp: Union[IOBase, str, Path], + logger: logging.Logger, + file: RemoteFile, + ) -> Iterator[Tuple[Union[int, str], pd.DataFrame]]: + """Walks the workbook with Calamine, holding the handle open across worksheets.""" + workbook = pd.ExcelFile(fp, engine="calamine") # type: ignore [arg-type] + emitted = False + for name in workbook.sheet_names: + try: + df = workbook.parse(sheet_name=name) + except BaseException as exc: + # Calamine raises PanicException (a BaseException) on a cell it cannot + # read. Falling back is only safe while nothing has been emitted yet. + if not emitted and "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 + emitted = True + yield name, df + + def _iter_worksheets_with_openpyxl( + self, + fp: Union[IOBase, str, Path], + logger: logging.Logger, + file: RemoteFile, + ) -> Iterator[Tuple[Union[int, str], pd.DataFrame]]: + """Walks the workbook with Openpyxl, holding the handle open across worksheets.""" + self._rewind(fp, logger, file) + workbook = self._with_openpyxl_warnings_logged( + lambda: pd.ExcelFile(fp, engine="openpyxl"), # type: ignore [arg-type] + logger, + file, + ) + for name in workbook.sheet_names: + yield ( + name, + self._with_openpyxl_warnings_logged( + lambda sheet=name: workbook.parse(sheet_name=sheet), # type: ignore [misc] + logger, + file, + ), + ) + + def _resolve_sheet_name(self, excel_format: ExcelFormat) -> Union[int, str, None]: + """Converts the config value to a pandas-compatible `sheet_name` argument. + + Unset means the first worksheet, which is the behavior that predates this option. + `*` means every worksheet; pandas spells that `None`. Any other value is used as a + literal worksheet name -- notably, a numeric-looking name like "2026" stays a name + rather than becoming a positional index. + """ + value = excel_format.sheet_name + if value is None: + return 0 + if value == self.ALL_SHEETS: + return None + return value + + def _reads_all_sheets(self, excel_format: ExcelFormat) -> bool: + """Whether the config selects every worksheet in the workbook.""" + return excel_format.sheet_name == self.ALL_SHEETS + + @staticmethod + def _describe_available_sheets(fp: Union[IOBase, str, Path], logger: logging.Logger) -> str: + """Best-effort listing of the worksheets present, for the not-found error message.""" + try: + if hasattr(fp, "seek"): + fp.seek(0) # type: ignore [union-attr] + names = pd.ExcelFile(fp, engine="calamine").sheet_names # type: ignore [arg-type] + except Exception as exc: + logger.info(f"Could not list worksheets for the error message: {exc}") + return "" + return f"Available worksheets: {', '.join(repr(str(n)) for n in names)}. " 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..fa6239a380 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,9 +3,11 @@ # +import asyncio import datetime import warnings from io import BytesIO +from typing import Any, Dict, List from unittest.mock import MagicMock, Mock, mock_open, patch import pandas as pd @@ -152,7 +154,7 @@ def test_open_and_parse_file_falls_back_to_openpyxl(mock_logger): calamine_excel_file = MagicMock() - def calamine_parse_side_effect(): + def calamine_parse_side_effect(**kwargs): raise FakePanic( "failed to construct date: PyErr { type: , value: ValueError('year 20225 is out of range'), traceback: None }" ) @@ -161,7 +163,7 @@ def calamine_parse_side_effect(): openpyxl_excel_file = MagicMock() - def openpyxl_parse_side_effect(): + def openpyxl_parse_side_effect(**kwargs): warnings.warn("Cell A146 has invalid date", UserWarning) return fallback_df @@ -238,3 +240,240 @@ 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") + openpyxl_excel_file.parse.assert_called_once_with(sheet_name=0) + + +SHEET_NAME_COL = ExcelParser.ab_sheet_name_col + + +def _make_excel_bytes(sheets: Dict[str, pd.DataFrame]) -> bytes: + """Creates an in-memory Excel workbook from a mapping of worksheet name to frame.""" + buf = BytesIO() + with pd.ExcelWriter(buf, engine="xlsxwriter") as writer: + for name, frame in sheets.items(): + frame.to_excel(writer, index=False, sheet_name=name) + return buf.getvalue() + + +def _make_multisheet_excel_bytes() -> bytes: + """Creates an in-memory Excel workbook with two sheets for testing.""" + return _make_excel_bytes( + { + "First": pd.DataFrame({"col_a": ["first"], "shared": [1]}), + "Second": pd.DataFrame({"col_b": [2.5], "shared": [2]}), + } + ) + + +def _stream_reader_for(excel_bytes: bytes) -> MagicMock: + reader = MagicMock(spec=AbstractFileBasedStreamReader) + reader.open_file.return_value = BytesIO(excel_bytes) + return reader + + +def _parse(excel_bytes: bytes, remote_file, **format_kwargs) -> List[Dict[str, Any]]: + parser = ExcelParser() + config = FileBasedStreamConfig(name="test_stream", format=ExcelFormat(**format_kwargs)) + return list( + parser.parse_records(config, remote_file, _stream_reader_for(excel_bytes), MagicMock()) + ) + + +def _infer(excel_bytes: bytes, remote_file, **format_kwargs) -> Dict[str, Any]: + parser = ExcelParser() + config = FileBasedStreamConfig(name="test_stream", format=ExcelFormat(**format_kwargs)) + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete( + parser.infer_schema(config, remote_file, _stream_reader_for(excel_bytes), MagicMock()) + ) + finally: + loop.close() + + +@pytest.mark.parametrize( + "format_kwargs,expected_records", + [ + pytest.param( + {}, + [{"col_a": "first", "shared": 1}], + id="unset_reads_first_sheet", + ), + pytest.param( + {"sheet_name": "Second"}, + [{"col_b": 2.5, "shared": 2}], + id="sheet_by_name", + ), + pytest.param( + {"sheet_name": "*"}, + [ + {"col_a": "first", "shared": 1, SHEET_NAME_COL: "First"}, + {"col_b": 2.5, "shared": 2, SHEET_NAME_COL: "Second"}, + ], + id="all_sheets_tagged_with_origin", + ), + ], +) +def test_parse_records_selects_configured_sheet(format_kwargs, expected_records, remote_file): + assert _parse(_make_multisheet_excel_bytes(), remote_file, **format_kwargs) == expected_records + + +def test_parse_records_without_sheet_name_adds_no_provenance(remote_file): + """The pre-existing single-sheet contract must be byte-for-byte unchanged.""" + records = _parse(_make_multisheet_excel_bytes(), remote_file) + + assert all(SHEET_NAME_COL not in record for record in records) + + +@pytest.mark.parametrize( + "format_kwargs,expected_schema", + [ + pytest.param( + {}, + {"col_a": {"type": "string"}, "shared": {"type": "number"}}, + id="first_sheet_schema", + ), + pytest.param( + {"sheet_name": "*"}, + { + "col_a": {"type": "string"}, + "col_b": {"type": "number"}, + "shared": {"type": "number"}, + SHEET_NAME_COL: {"type": "string"}, + }, + id="all_sheets_merged_schema", + ), + ], +) +def test_infer_schema_with_sheet_selection(format_kwargs, expected_schema, remote_file): + assert _infer(_make_multisheet_excel_bytes(), remote_file, **format_kwargs) == expected_schema + + +@pytest.mark.parametrize( + "config_value,expected", + [ + pytest.param(None, 0, id="unset_means_first_sheet"), + pytest.param("MySheet", "MySheet", id="named_sheet"), + pytest.param("2026", "2026", id="numeric_name_stays_a_name"), + pytest.param("0", "0", id="zero_is_a_name_not_an_index"), + pytest.param("*", None, id="all_sheets"), + ], +) +def test_resolve_sheet_name(config_value, expected): + parser = ExcelParser() + fmt = ExcelFormat(sheet_name=config_value) + assert parser._resolve_sheet_name(fmt) == expected + + +def test_parse_records_reads_worksheet_with_numeric_name(remote_file): + """Worksheets named for fiscal years must resolve by name, not by position.""" + excel_bytes = _make_excel_bytes( + { + "2026": pd.DataFrame({"month": ["Jan"], "total": [10]}), + "2025": pd.DataFrame({"month": ["Feb"], "total": [8]}), + } + ) + + assert _parse(excel_bytes, remote_file, sheet_name="2026") == [{"month": "Jan", "total": 10}] + assert _parse(excel_bytes, remote_file, sheet_name="2025") == [{"month": "Feb", "total": 8}] + + +def test_parse_records_preserves_workbook_sheet_order(remote_file): + excel_bytes = _make_excel_bytes( + { + "Alpha": pd.DataFrame({"v": [1, 2]}), + "Beta": pd.DataFrame({"v": [3]}), + "Gamma": pd.DataFrame({"v": [4]}), + } + ) + + records = _parse(excel_bytes, remote_file, sheet_name="*") + + assert [record[SHEET_NAME_COL] for record in records] == [ + "Alpha", + "Alpha", + "Beta", + "Gamma", + ] + + +def test_all_sheets_on_single_sheet_workbook(remote_file): + excel_bytes = _make_excel_bytes({"Only": pd.DataFrame({"a": [1]})}) + + assert _parse(excel_bytes, remote_file, sheet_name="*") == [{"a": 1, SHEET_NAME_COL: "Only"}] + + +@pytest.mark.parametrize( + "sheet_name", + [ + pytest.param("Missing", id="absent_worksheet"), + pytest.param("first", id="wrong_case"), + ], +) +def test_missing_worksheet_raises_config_error(sheet_name, remote_file): + """A bad worksheet name is a config mistake and must not be reported as a parse failure.""" + with pytest.raises(ConfigValidationError) as exc_info: + _parse(_make_multisheet_excel_bytes(), remote_file, sheet_name=sheet_name) + + message = str(exc_info.value) + assert repr(sheet_name) in message + assert "'First'" in message and "'Second'" in message + assert "case-sensitive" in message + assert "mismatch between the config's file type" not in message + + +def test_unparseable_file_still_raises_record_parse_error(remote_file): + """Narrowing the error handling must not mask genuinely corrupt files.""" + with pytest.raises(RecordParseError): + _parse(b"this is not an excel file at all", remote_file) + + +def _spy_on_sheet_parsing(monkeypatch) -> List[Any]: + """Records the sheet_name argument of every pandas parse call, in order.""" + parsed: List[Any] = [] + original = pd.ExcelFile.parse + + def spy(self, *args, **kwargs): + parsed.append(kwargs.get("sheet_name", args[0] if args else None)) + return original(self, *args, **kwargs) + + monkeypatch.setattr(pd.ExcelFile, "parse", spy) + return parsed + + +def test_all_sheets_parses_worksheets_lazily(remote_file, monkeypatch): + """Worksheets are parsed on demand. + + Asking pandas for every worksheet at once builds one frame per worksheet before a + single record is emitted, so peak memory tracks the whole workbook. Parsing lazily + keeps it at one worksheet regardless of how many the workbook holds. + """ + excel_bytes = _make_excel_bytes( + { + "A": pd.DataFrame({"v": [1]}), + "B": pd.DataFrame({"v": [2]}), + "C": pd.DataFrame({"v": [3]}), + } + ) + parsed = _spy_on_sheet_parsing(monkeypatch) + parser = ExcelParser() + config = FileBasedStreamConfig(name="test_stream", format=ExcelFormat(sheet_name="*")) + stream = parser.parse_records(config, remote_file, _stream_reader_for(excel_bytes), MagicMock()) + + next(stream) + assert parsed == ["A"], "later worksheets must not be parsed before they are requested" + + assert list(stream) == [ + {"v": 2, SHEET_NAME_COL: "B"}, + {"v": 3, SHEET_NAME_COL: "C"}, + ] + assert parsed == ["A", "B", "C"] + + +def test_single_sheet_parses_only_the_selected_worksheet(monkeypatch, remote_file): + """Selecting one worksheet must not walk the rest of the workbook.""" + parsed = _spy_on_sheet_parsing(monkeypatch) + + _parse(_make_multisheet_excel_bytes(), remote_file, sheet_name="Second") + + assert parsed == ["Second"] diff --git a/unit_tests/sources/file_based/in_memory_files_source.py b/unit_tests/sources/file_based/in_memory_files_source.py index 4cd4658953..ae9635b9cf 100644 --- a/unit_tests/sources/file_based/in_memory_files_source.py +++ b/unit_tests/sources/file_based/in_memory_files_source.py @@ -279,11 +279,13 @@ def open_file( def _make_file_contents(self, file_name: str) -> bytes: contents = self.files[file_name]["contents"] - df = pd.DataFrame(contents) + # A mapping declares one worksheet per key; a bare list is a single "Sheet1". + sheets = contents if isinstance(contents, Mapping) else {"Sheet1": contents} with io.BytesIO() as fp: writer = pd.ExcelWriter(fp, engine="xlsxwriter") - df.to_excel(writer, index=False, sheet_name="Sheet1") + for sheet_name, rows in sheets.items(): + pd.DataFrame(rows).to_excel(writer, index=False, sheet_name=sheet_name) writer._save() fp.seek(0) return fp.read() diff --git a/unit_tests/sources/file_based/scenarios/csv_scenarios.py b/unit_tests/sources/file_based/scenarios/csv_scenarios.py index f16d83e200..97715c821b 100644 --- a/unit_tests/sources/file_based/scenarios/csv_scenarios.py +++ b/unit_tests/sources/file_based/scenarios/csv_scenarios.py @@ -472,7 +472,13 @@ "default": "excel", "const": "excel", "type": "string", - } + }, + "sheet_name": { + "title": "Sheet Name", + "description": 'The worksheet to read from each workbook. Leave empty to read only the first worksheet. Enter an exact worksheet name to read that worksheet, or "*" to read every worksheet in the workbook. Worksheet names are case-sensitive.', + "examples": ["Sheet1", "*"], + "type": "string", + }, }, "required": ["filetype"], }, diff --git a/unit_tests/sources/file_based/scenarios/excel_scenarios.py b/unit_tests/sources/file_based/scenarios/excel_scenarios.py index 94ccc67646..f39c0998a6 100644 --- a/unit_tests/sources/file_based/scenarios/excel_scenarios.py +++ b/unit_tests/sources/file_based/scenarios/excel_scenarios.py @@ -467,3 +467,123 @@ } ) ).build() + + +_multi_sheet_excel_file = { + "a.xlsx": { + "contents": { + "Sales": [ + {"region": "West", "revenue": 100}, + {"region": "East", "revenue": 200}, + ], + "Inventory": [ + {"sku": "A1", "qty": 5}, + ], + }, + "last_modified": "2023-06-05T03:54:07.000Z", + } +} + + +def _multi_sheet_scenario_builder(name, sheet_name, records, properties): + """Builds a multi-worksheet scenario differing only in the `sheet_name` option.""" + excel_format = {"filetype": "excel"} + if sheet_name is not None: + excel_format["sheet_name"] = sheet_name + return ( + TestScenarioBuilder() + .set_name(name) + .set_config( + { + "streams": [ + { + "name": "stream1", + "format": excel_format, + "globs": ["*"], + "validation_policy": "Emit Record", + } + ] + } + ) + .set_source_builder( + FileBasedSourceBuilder() + .set_stream_reader( + TemporaryExcelFilesStreamReader(files=_multi_sheet_excel_file, file_type="excel") + ) + .set_file_type("excel") + ) + .set_expected_check_status("SUCCEEDED") + .set_expected_records( + [ + { + "data": { + **record, + "_ab_source_file_last_modified": "2023-06-05T03:54:07.000000Z", + "_ab_source_file_url": "a.xlsx", + }, + "stream": "stream1", + } + for record in records + ] + ) + .set_expected_catalog( + { + "streams": [ + { + "default_cursor_field": ["_ab_source_file_last_modified"], + "json_schema": { + "type": "object", + "properties": { + **properties, + "_ab_source_file_last_modified": {"type": "string"}, + "_ab_source_file_url": {"type": "string"}, + }, + }, + "name": "stream1", + "source_defined_cursor": True, + "supported_sync_modes": ["full_refresh", "incremental"], + "is_resumable": True, + } + ] + } + ) + ).build() + + +# Unset `sheet_name` must behave exactly as it did before the option existed: first +# worksheet only, and no worksheet-name field on the records. +multi_sheet_excel_default_scenario = _multi_sheet_scenario_builder( + "multi_sheet_excel_default", + None, + [ + {"region": "West", "revenue": 100}, + {"region": "East", "revenue": 200}, + ], + {"region": {"type": ["null", "string"]}, "revenue": {"type": ["null", "number"]}}, +) + +# A named worksheet reaches past the first one. +multi_sheet_excel_named_scenario = _multi_sheet_scenario_builder( + "multi_sheet_excel_named", + "Inventory", + [{"sku": "A1", "qty": 5}], + {"sku": {"type": ["null", "string"]}, "qty": {"type": ["null", "number"]}}, +) + +# `*` reads every worksheet, unions their columns, and tags each record with its origin. +multi_sheet_excel_all_scenario = _multi_sheet_scenario_builder( + "multi_sheet_excel_all", + "*", + [ + {"region": "West", "revenue": 100, "_ab_source_sheet_name": "Sales"}, + {"region": "East", "revenue": 200, "_ab_source_sheet_name": "Sales"}, + {"sku": "A1", "qty": 5, "_ab_source_sheet_name": "Inventory"}, + ], + { + "region": {"type": ["null", "string"]}, + "revenue": {"type": ["null", "number"]}, + "sku": {"type": ["null", "string"]}, + "qty": {"type": ["null", "number"]}, + "_ab_source_sheet_name": {"type": ["null", "string"]}, + }, +) diff --git a/unit_tests/sources/file_based/test_file_based_scenarios.py b/unit_tests/sources/file_based/test_file_based_scenarios.py index 36c9f42fc2..a9596571c6 100644 --- a/unit_tests/sources/file_based/test_file_based_scenarios.py +++ b/unit_tests/sources/file_based/test_file_based_scenarios.py @@ -91,6 +91,9 @@ ) from unit_tests.sources.file_based.scenarios.excel_scenarios import ( excel_all_types_scenario, + multi_sheet_excel_all_scenario, + multi_sheet_excel_default_scenario, + multi_sheet_excel_named_scenario, multiple_excel_combine_schema_scenario, multiple_streams_excel_scenario, single_excel_scenario, @@ -249,6 +252,9 @@ multiple_streams_avro_scenario, avro_file_with_double_as_number_scenario, excel_all_types_scenario, + multi_sheet_excel_all_scenario, + multi_sheet_excel_default_scenario, + multi_sheet_excel_named_scenario, multiple_excel_combine_schema_scenario, multiple_streams_excel_scenario, single_excel_scenario,