wizardchen
11811f9bf8
DuckDB's st_read (spatial) only reads the first layer/sheet of a .xlsx workbook, so every sheet beyond Sheet1 was silently dropped from the DuckDB table the Data Analysis tool builds. Users trying to analyse multi-sheet workbooks could only see the first sheet. Switch to DuckDB's dedicated 'excel' extension (read_xlsx) for the actual data load, enumerate sheets via the spatial extension's st_read_meta, and UNION ALL BY NAME the rows of every sheet into one table. A synthetic __sheet_name column records the source so the LLM can still filter/aggregate per sheet; schema drift between sheets is tolerated via UNION BY NAME. If enumeration fails (older DuckDB, local filesystem errors, …) we fall back to reading the first sheet so the tool stays usable. - Install & LOAD the 'excel' extension alongside 'spatial' both at startup (internal/container) and in the offline prefetch binary (cmd/download/duckdb). - Harden sheet/path handling against single quotes. - Update the tool description so the agent knows about __sheet_name. - Add unit tests for the CREATE TABLE SQL builder covering 0 / 1 / N sheets and quote escaping. Refs: https://github.com/Tencent/WeKnora/issues/1007
91 lines
3.2 KiB
Go
91 lines
3.2 KiB
Go
package tools
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestBuildExcelCreateTableSQL_NoSheets(t *testing.T) {
|
|
got := buildExcelCreateTableSQL("tbl", "/tmp/data.xlsx", nil)
|
|
want := `CREATE TABLE "tbl" AS SELECT * FROM read_xlsx('/tmp/data.xlsx')`
|
|
if got != want {
|
|
t.Fatalf("mismatch.\n got: %s\nwant: %s", got, want)
|
|
}
|
|
}
|
|
|
|
func TestBuildExcelCreateTableSQL_SingleSheetTagsSource(t *testing.T) {
|
|
got := buildExcelCreateTableSQL("tbl", "/tmp/data.xlsx", []string{"Sheet1"})
|
|
|
|
// Must use read_xlsx (excel extension) with explicit sheet param.
|
|
if !strings.Contains(got, "FROM read_xlsx('/tmp/data.xlsx', sheet = 'Sheet1')") {
|
|
t.Fatalf("expected read_xlsx with sheet param, got: %s", got)
|
|
}
|
|
// Must tag the source sheet name via the synthetic column so downstream
|
|
// SQL behaves consistently between single- and multi-sheet workbooks.
|
|
if !strings.Contains(got, "'Sheet1' AS "+excelSheetNameColumn) {
|
|
t.Fatalf("expected sheet-name column, got: %s", got)
|
|
}
|
|
}
|
|
|
|
func TestBuildExcelCreateTableSQL_MultiSheetUsesUnionAllByName(t *testing.T) {
|
|
got := buildExcelCreateTableSQL("tbl", "/tmp/data.xlsx", []string{"Sheet1", "Sheet2", "报表"})
|
|
|
|
// Each sheet must appear as a SELECT reading that specific sheet, and
|
|
// the __sheet_name column must carry its name for per-sheet filtering.
|
|
for _, sheet := range []string{"Sheet1", "Sheet2", "报表"} {
|
|
needleRead := "FROM read_xlsx('/tmp/data.xlsx', sheet = '" + sheet + "')"
|
|
needleTag := "'" + sheet + "' AS " + excelSheetNameColumn
|
|
if !strings.Contains(got, needleRead) {
|
|
t.Fatalf("missing read_xlsx for sheet %q in:\n%s", sheet, got)
|
|
}
|
|
if !strings.Contains(got, needleTag) {
|
|
t.Fatalf("missing __sheet_name tag for sheet %q in:\n%s", sheet, got)
|
|
}
|
|
}
|
|
|
|
// Must combine with UNION ALL BY NAME so schema drift between sheets is
|
|
// tolerated.
|
|
if !strings.Contains(got, "UNION ALL BY NAME") {
|
|
t.Fatalf("expected UNION ALL BY NAME in multi-sheet SQL, got:\n%s", got)
|
|
}
|
|
|
|
// Exactly N-1 UNIONs for N sheets.
|
|
if strings.Count(got, "UNION ALL BY NAME") != 2 {
|
|
t.Fatalf("expected 2 UNION ALL BY NAME separators, got %d in:\n%s",
|
|
strings.Count(got, "UNION ALL BY NAME"), got)
|
|
}
|
|
}
|
|
|
|
func TestBuildExcelCreateTableSQL_EscapesSingleQuotes(t *testing.T) {
|
|
// Sheet name and file path both contain single quotes, which must be
|
|
// doubled to produce a valid SQL literal.
|
|
sheets := []string{"Jo's data"}
|
|
got := buildExcelCreateTableSQL("tbl", "/tmp/O'Brien/data.xlsx", sheets)
|
|
|
|
if !strings.Contains(got, "sheet = 'Jo''s data'") {
|
|
t.Fatalf("sheet name was not escaped, got:\n%s", got)
|
|
}
|
|
if !strings.Contains(got, "read_xlsx('/tmp/O''Brien/data.xlsx'") {
|
|
t.Fatalf("file path was not escaped, got:\n%s", got)
|
|
}
|
|
if !strings.Contains(got, "'Jo''s data' AS "+excelSheetNameColumn) {
|
|
t.Fatalf("sheet-name literal was not escaped, got:\n%s", got)
|
|
}
|
|
}
|
|
|
|
func TestSqlSingleQuoteEscape(t *testing.T) {
|
|
cases := map[string]string{
|
|
"": "",
|
|
"no_quote": "no_quote",
|
|
"a'b": "a''b",
|
|
"''": "''''",
|
|
"mix'ed'quote": "mix''ed''quote",
|
|
"中文 with 'quote": "中文 with ''quote",
|
|
}
|
|
for in, want := range cases {
|
|
if got := sqlSingleQuoteEscape(in); got != want {
|
|
t.Errorf("sqlSingleQuoteEscape(%q) = %q, want %q", in, got, want)
|
|
}
|
|
}
|
|
}
|