Spaces:
Running on Zero
Running on Zero
File size: 3,147 Bytes
3fbd1fd | 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 | """Safely execute SQL against Spider SQLite databases and compare results."""
from __future__ import annotations
import sqlite3
import time
from collections import Counter
from contextlib import closing
from dataclasses import dataclass, field
from src.data import db_path
ALLOWED_START = ("select", "with")
@dataclass
class ExecResult:
ok: bool
rows: list = field(default_factory=list)
columns: list = field(default_factory=list)
error: str | None = None
def execute_sql(db_id: str, sql: str, timeout_s: float = 2.0,
max_rows: int = 10_000) -> ExecResult:
"""Run one read-only query with a timeout. Never raises."""
sql = sql.strip().rstrip(";").strip()
if not sql.lower().startswith(ALLOWED_START):
return ExecResult(ok=False, error="Only SELECT/WITH queries are allowed")
uri = f"file:{db_path(db_id)}?mode=ro"
try:
with closing(sqlite3.connect(uri, uri=True)) as conn:
# Some Spider DBs contain invalid UTF-8; don't crash on it
conn.text_factory = lambda b: b.decode(errors="replace")
# Abort the query once the deadline passes
deadline = time.monotonic() + timeout_s
conn.set_progress_handler(lambda: int(time.monotonic() > deadline), 10_000)
cur = conn.execute(sql)
rows = cur.fetchmany(max_rows + 1)
columns = [d[0] for d in cur.description or []]
except (sqlite3.Error, sqlite3.Warning) as e:
return ExecResult(ok=False, error=f"{type(e).__name__}: {e}")
if len(rows) > max_rows:
return ExecResult(ok=False, error=f"Result exceeded {max_rows} rows")
return ExecResult(ok=True, rows=rows, columns=columns)
def _normalize(value):
return round(value, 6) if isinstance(value, float) else value
def results_match(pred_rows: list, gold_rows: list, order_matters: bool = False) -> bool:
"""Compare two result sets. Order only matters if the gold SQL uses ORDER BY."""
pred = [tuple(_normalize(v) for v in row) for row in pred_rows]
gold = [tuple(_normalize(v) for v in row) for row in gold_rows]
if order_matters:
return pred == gold
return Counter(pred) == Counter(gold) # same rows, any order, duplicates counted
if __name__ == "__main__":
from src.data import load_split
ex = load_split("dev")[0]
db = ex["db_id"]
gold = execute_sql(db, ex["gold_sql"])
checks = {
"gold query": ex["gold_sql"],
"broken SQL": "SELEC count(*) FROM singer",
"write attempt": "DROP TABLE singer",
"sneaky 2nd statement": "SELECT 1; DROP TABLE singer",
"infinite query": "WITH RECURSIVE c(x) AS (SELECT 1 UNION ALL SELECT x+1 FROM c) "
"SELECT count(*) FROM c",
}
for name, sql in checks.items():
start = time.monotonic()
res = execute_sql(db, sql)
ms = (time.monotonic() - start) * 1000
status = f"rows={res.rows[:3]}" if res.ok else f"error={res.error}"
print(f"{name:22s} ok={res.ok!s:5s} {ms:6.0f}ms {status}")
print("\ngold matches itself:", results_match(gold.rows, gold.rows)) |