feat: new Columnar upload form and API (#28192)

tests/unit_tests/commands/databases/excel_reader_test.py

@@ -20,6 +20,9 @@ from typing import Any
 
 import numpy as np
 import pytest
+import xlsxwriter
+from werkzeug.datastructures import FileStorage
+from xlsxwriter.workbook import Worksheet
 
 from superset.commands.database.exceptions import DatabaseUploadFailed
 from superset.commands.database.uploaders.excel_reader import (
@@ -50,6 +53,18 @@ EXCEL_DATA_DECIMAL_CHAR = {
 }
 
 
+def write_data_to_worksheet(
+    worksheet: Worksheet, header: list[str], data: list[list[Any]]
+):
+    all_data = [header] + data
+    row = 0
+    col = 0
+    for name, age in all_data:
+        worksheet.write(row, col, name)
+        worksheet.write(row, col + 1, age)
+        row += 1
+
+
 @pytest.mark.parametrize(
     "file, options, expected_cols, expected_values",
     [
@@ -175,6 +190,23 @@ def test_excel_reader_file_to_dataframe(file, options, expected_cols, expected_v
     file.close()
 
 
+def test_excel_reader_index_column():
+    excel_reader = ExcelReader(
+        options=ExcelReaderOptions(index_column="Name"),
+    )
+    df = excel_reader.file_to_dataframe(create_excel_file(EXCEL_DATA))
+    assert df.index.name == "Name"
+
+
+def test_excel_reader_wrong_index_column():
+    excel_reader = ExcelReader(
+        options=ExcelReaderOptions(index_column="wrong"),
+    )
+    with pytest.raises(DatabaseUploadFailed) as ex:
+        excel_reader.file_to_dataframe(create_excel_file(EXCEL_DATA))
+    assert str(ex.value) == ("Parsing error: Index wrong invalid (sheet: 0)")
+
+
 def test_excel_reader_wrong_columns_to_read():
     excel_reader = ExcelReader(
         options=ExcelReaderOptions(columns_read=["xpto"]),
@@ -203,7 +235,60 @@ def test_excel_reader_invalid_file():
         options=ExcelReaderOptions(),
     )
     with pytest.raises(DatabaseUploadFailed) as ex:
-        excel_reader.file_to_dataframe(io.StringIO("c1"))
+        excel_reader.file_to_dataframe(FileStorage(io.BytesIO(b"c1")))
     assert str(ex.value) == (
         "Parsing error: Excel file format cannot be determined, you must specify an engine manually."
     )
+
+
+def test_excel_reader_metadata():
+    excel_reader = ExcelReader(
+        options=ExcelReaderOptions(),
+    )
+    file = create_excel_file(EXCEL_DATA)
+    metadata = excel_reader.file_metadata(file)
+    assert metadata == {
+        "items": [
+            {"column_names": ["Name", "Age", "City", "Birth"], "sheet_name": "Sheet1"}
+        ]
+    }
+    file.close()
+
+
+def test_excel_reader_metadata_mul_sheets():
+    buffer = io.BytesIO()
+    workbook = xlsxwriter.Workbook(buffer)
+
+    worksheet1 = workbook.add_worksheet("Sheet1")
+    header1 = ["col11", "col12"]
+    data1 = [["v11", "v12"]]
+    write_data_to_worksheet(worksheet1, header1, data1)
+
+    worksheet2 = workbook.add_worksheet("Sheet2")
+    header2 = ["col21", "col22"]
+    data2 = [["v21", "v22"]]
+    write_data_to_worksheet(worksheet2, header2, data2)
+    workbook.close()
+
+    file = FileStorage(stream=buffer, filename="test.xls")
+
+    excel_reader = ExcelReader(
+        options=ExcelReaderOptions(),
+    )
+    metadata = excel_reader.file_metadata(file)
+    assert metadata == {
+        "items": [
+            {"column_names": ["col11", "col12"], "sheet_name": "Sheet1"},
+            {"column_names": ["col21", "col22"], "sheet_name": "Sheet2"},
+        ]
+    }
+    file.close()
+
+
+def test_excel_reader_file_metadata_invalid_file():
+    excel_reader = ExcelReader(
+        options=ExcelReaderOptions(),
+    )
+    with pytest.raises(DatabaseUploadFailed) as ex:
+        excel_reader.file_metadata(FileStorage(io.BytesIO(b"1")))
+    assert str(ex.value) == ("Excel file format cannot be determined")