Spaces:
Running on Zero
Running on Zero
File size: 4,755 Bytes
a067ada | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """
DuckDB-based SQL executor for in-memory analytical queries.
Accepts pandas DataFrames or CSV/JSON paths and exposes them as
queryable tables in a single in-memory connection.
"""
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union
import duckdb
import pandas as pd
logger = logging.getLogger(__name__)
class SQLExecutor:
"""Execute SQL queries against an in-memory DuckDB connection."""
def __init__(self) -> None:
self.con = duckdb.connect(database=":memory:")
self._tables: Dict[str, int] = {}
def register_dataframe(self, name: str, df: pd.DataFrame) -> None:
"""Register a DataFrame as a queryable table."""
safe = self._sanitize_name(name)
self.con.register(f"_tmp_{safe}", df)
self.con.execute(f'CREATE OR REPLACE TABLE "{safe}" AS SELECT * FROM _tmp_{safe}')
self.con.unregister(f"_tmp_{safe}")
self._tables[safe] = len(df)
logger.info(f"Registered table '{safe}' ({len(df):,} rows, {len(df.columns)} cols)")
def register_file(self, path: Union[str, Path], name: Optional[str] = None) -> str:
"""Load a CSV/JSON/Parquet file into a table. Returns the table name used."""
path = Path(path)
if not path.exists():
raise FileNotFoundError(path)
safe = self._sanitize_name(name or path.stem)
ext = path.suffix.lower()
if ext == ".csv":
self.con.execute(
f"CREATE OR REPLACE TABLE \"{safe}\" AS SELECT * FROM read_csv_auto('{path}')"
)
elif ext == ".json":
self.con.execute(
f"CREATE OR REPLACE TABLE \"{safe}\" AS SELECT * FROM read_json_auto('{path}')"
)
elif ext in (".parquet", ".pq"):
self.con.execute(
f"CREATE OR REPLACE TABLE \"{safe}\" AS SELECT * FROM read_parquet('{path}')"
)
elif ext in (".xls", ".xlsx"):
df = pd.read_excel(path)
self.register_dataframe(safe, df)
return safe
else:
raise ValueError(f"Unsupported file extension: {ext}")
rows = self.con.execute(f'SELECT COUNT(*) FROM "{safe}"').fetchone()[0]
self._tables[safe] = rows
logger.info(f"Loaded '{path.name}' as table '{safe}' ({rows:,} rows)")
return safe
def execute(self, query: str) -> Tuple[List[Dict[str, Any]], List[Dict[str, str]]]:
"""Execute a query and return (rows, column_info)."""
if not query or not query.strip():
raise ValueError("Query cannot be empty")
query = query.strip().rstrip(";")
logger.info(f"Executing: {query[:120]}...")
try:
cur = self.con.execute(query)
rows = cur.fetchall()
descriptions = cur.description or []
columns = [
{"name": d[0], "type": str(d[1]) if d[1] else "VARCHAR"}
for d in descriptions
]
results = [dict(zip([c["name"] for c in columns], row)) for row in rows]
logger.info(f"Returned {len(results):,} rows × {len(columns)} cols")
return results, columns
except duckdb.Error as e:
logger.error(f"DuckDB error: {e}")
raise ValueError(f"SQL error: {e}")
def validate_query(self, query: str) -> bool:
"""Check that a query parses and references valid tables, without executing."""
if not query or not query.strip():
return False
try:
self.con.execute(f"EXPLAIN {query.strip().rstrip(';')}")
return True
except Exception as e:
logger.warning(f"Validation failed: {e}")
return False
def get_table_names(self) -> List[str]:
rows = self.con.execute("SHOW TABLES").fetchall()
return [r[0] for r in rows]
def get_table_schema(self, table: str) -> List[Dict[str, Any]]:
safe = self._sanitize_name(table)
rows = self.con.execute(f'DESCRIBE "{safe}"').fetchall()
return [
{"name": r[0], "type": r[1], "nullable": r[2] != "NO" if r[2] else True}
for r in rows
]
def get_sample(self, table: str, n: int = 5) -> pd.DataFrame:
safe = self._sanitize_name(table)
return self.con.execute(f'SELECT * FROM "{safe}" LIMIT {n}').df()
def close(self) -> None:
self.con.close()
@staticmethod
def _sanitize_name(name: str) -> str:
"""Make a string safe to use as an unquoted table identifier fallback."""
s = "".join(c if c.isalnum() or c == "_" else "_" for c in name)
if s and s[0].isdigit():
s = "t_" + s
return s or "table"
|