""" File upload loader -- supports CSV, SQLite/DB files, and Excel workbooks. CSV: single table "data_table" -> WikiSQL-format schema (no FK/PK) SQLite/DB: copy to temp path, read multi-table schema via PRAGMA -> Spider-format schema Excel: each sheet -> one table in temp SQLite -> Spider-format schema The schema string format matches what models were trained on: CSV -> WikiSQL format (process_wikisql.build_schema_string) SQLite/Excel -> Spider format (table ( col*:type , ... ) | FK: ...) """ import shutil import sqlite3 import sys import tempfile from pathlib import Path import pandas as pd _SRC_DIR = Path(__file__).resolve().parent.parent # src/ if str(_SRC_DIR) not in sys.path: sys.path.insert(0, str(_SRC_DIR)) from process_wikisql import build_schema_string as _wikisql_schema # noqa: E402 TABLE_NAME = "data_table" _UPLOAD_DB_PATH = Path(tempfile.gettempdir()) / "_uploaded_data.sqlite" MAX_PREVIEW_ROWS = 500 # --------------------------------------------------------------------------- # Shared helpers # --------------------------------------------------------------------------- def _infer_wikisql_types(df: pd.DataFrame) -> list: return ["real" if pd.api.types.is_numeric_dtype(dt) else "text" for dt in df.dtypes] def _lowercase_text_columns(df: pd.DataFrame, types: list) -> pd.DataFrame: """WikiSQL-trained models expect lowercase string literals in WHERE clauses.""" df = df.copy() for col, t in zip(df.columns, types): if t == "text": df[col] = df[col].apply(lambda v: v.lower() if isinstance(v, str) else v) return df def _sqlite_col_type(declared_type: str) -> str: """Map SQLite declared type -> 'number' or 'text' for schema string.""" upper = declared_type.upper() if any(k in upper for k in ("INT", "REAL", "FLOAT", "NUMERIC", "DECIMAL", "DOUBLE")): return "number" return "text" def _preview_from_db(db_path: str) -> dict: """Return preview rows from first table and total row count across all tables.""" conn = sqlite3.connect(db_path) cur = conn.cursor() tables = [r[0] for r in cur.execute( "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name" ).fetchall()] total_rows = 0 first_table_df = None for tbl in tables: cnt = cur.execute(f'SELECT COUNT(*) FROM "{tbl}"').fetchone()[0] total_rows += cnt if first_table_df is None: first_table_df = pd.read_sql_query( f'SELECT * FROM "{tbl}" LIMIT {MAX_PREVIEW_ROWS}', conn ) conn.close() if first_table_df is None or first_table_df.empty: return {"columns": [], "rows": [], "row_count": 0, "truncated": False} preview_df = first_table_df.astype(object).where(first_table_df.notna(), None) return { "columns": list(first_table_df.columns), "rows": preview_df.values.tolist(), "row_count": total_rows, "truncated": total_rows > MAX_PREVIEW_ROWS, } def get_table_names(db_path: str) -> list: """Return all user table names in a SQLite database (alphabetical order).""" conn = sqlite3.connect(db_path) tables = [r[0] for r in conn.execute( "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name" ).fetchall()] conn.close() return tables def get_preview_for_table(db_path: str, table: str) -> dict: """Return preview rows (up to MAX_PREVIEW_ROWS) from a specific table.""" conn = sqlite3.connect(db_path) total = conn.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0] df = pd.read_sql_query(f'SELECT * FROM "{table}" LIMIT {MAX_PREVIEW_ROWS}', conn) conn.close() df = df.astype(object).where(df.notna(), None) return { "columns": list(df.columns), "rows": df.values.tolist(), "row_count": total, "truncated": total > MAX_PREVIEW_ROWS, } def _build_spider_schema_from_sqlite(db_path: str) -> str: """Read schema from a SQLite file and return a Spider-format schema string. Format: table ( col*:type , col2:type ) | table2 ( ... ) | foreign_keys: t.c = t2.c , ... """ conn = sqlite3.connect(db_path) cur = conn.cursor() tables = [r[0] for r in cur.execute( "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name" ).fetchall()] table_parts = [] fk_parts = [] for tbl in tables: tbl_lower = tbl.lower() cur.execute(f'PRAGMA table_info("{tbl}")') cols = cur.fetchall() # (cid, name, type, notnull, default, pk) col_strs = [] for col in cols: col_name = col[1].lower() col_type = _sqlite_col_type(col[2]) pk_marker = "*" if col[5] > 0 else "" col_strs.append(f"{col_name}{pk_marker}:{col_type}") table_parts.append(f"{tbl_lower} ( {' , '.join(col_strs)} )") cur.execute(f'PRAGMA foreign_key_list("{tbl}")') for fk in cur.fetchall(): # (id, seq, ref_table, from_col, to_col, ...) ref_table = fk[2].lower() from_col = fk[3].lower() to_col = fk[4].lower() fk_parts.append(f"{tbl_lower}.{from_col} = {ref_table}.{to_col}") conn.close() schema = " | ".join(table_parts) if fk_parts: schema += " | foreign_keys: " + " , ".join(fk_parts) return schema # --------------------------------------------------------------------------- # CSV loader (unchanged from original, single-table WikiSQL format) # --------------------------------------------------------------------------- def load_csv_to_sqlite(file_obj_or_path) -> dict: """Load a CSV file into a single-table SQLite DB and return metadata.""" df = pd.read_csv(file_obj_or_path) header = list(df.columns) types = _infer_wikisql_types(df) df = _lowercase_text_columns(df, types) conn = sqlite3.connect(str(_UPLOAD_DB_PATH)) try: df.to_sql(TABLE_NAME, conn, index=False, if_exists="replace") finally: conn.close() schema_string = _wikisql_schema(header, types=types) preview_df = df.head(MAX_PREVIEW_ROWS).astype(object).where(df.head(MAX_PREVIEW_ROWS).notna(), None) return { "db_path": str(_UPLOAD_DB_PATH), "schema_string": schema_string, "columns": header, "row_count": len(df), "rows": preview_df.values.tolist(), "truncated": len(df) > MAX_PREVIEW_ROWS, } # --------------------------------------------------------------------------- # SQLite / .db file loader (multi-table, Spider format) # --------------------------------------------------------------------------- def load_sqlite_file(file_obj) -> dict: """Copy an uploaded .sqlite/.db file to temp location and build schema.""" # Write uploaded bytes to temp file with open(str(_UPLOAD_DB_PATH), "wb") as out: shutil.copyfileobj(file_obj, out) schema_string = _build_spider_schema_from_sqlite(str(_UPLOAD_DB_PATH)) preview = _preview_from_db(str(_UPLOAD_DB_PATH)) return { "db_path": str(_UPLOAD_DB_PATH), "schema_string": schema_string, **preview, } # --------------------------------------------------------------------------- # Excel loader (each sheet = one table, Spider format) # --------------------------------------------------------------------------- def load_excel_to_sqlite(file_obj) -> dict: """Load an Excel workbook: each sheet becomes a table in the temp SQLite DB.""" xl = pd.ExcelFile(file_obj) sheet_names = xl.sheet_names conn = sqlite3.connect(str(_UPLOAD_DB_PATH)) try: for sheet in sheet_names: df = xl.parse(sheet) # Sanitize table name: lowercase, spaces to underscores tbl_name = sheet.lower().replace(" ", "_") df.to_sql(tbl_name, conn, index=False, if_exists="replace") conn.commit() finally: conn.close() schema_string = _build_spider_schema_from_sqlite(str(_UPLOAD_DB_PATH)) preview = _preview_from_db(str(_UPLOAD_DB_PATH)) return { "db_path": str(_UPLOAD_DB_PATH), "schema_string": schema_string, **preview, } # --------------------------------------------------------------------------- # Model input builder # --------------------------------------------------------------------------- def build_input(question: str, schema_string: str) -> str: """Build model input string from a pre-built schema string (any format).""" return f"question: {question.strip()} | schema: {schema_string}"