Spaces:
Sleeping
Sleeping
File size: 6,963 Bytes
d3d0e0e | 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | from __future__ import annotations
import re
import sqlite3
from pathlib import Path
import pandas as pd
DEFAULT_SAMPLE_ROWS: int = 5
def _connect_read_only(path: Path) -> sqlite3.Connection:
uri = f"file:{path.resolve().as_posix()}?mode=ro"
return sqlite3.connect(uri, uri=True)
def inspect_sqlite_schema(path: Path, *, sample_rows: int = DEFAULT_SAMPLE_ROWS) -> dict[str, object]:
with _connect_read_only(path) as conn:
rows = conn.execute(
"""
SELECT name, sql
FROM sqlite_master
WHERE type = 'table' AND name NOT LIKE 'sqlite_%'
ORDER BY name
"""
).fetchall()
tables: list[dict[str, object]] = []
for name, create_sql in rows:
# Fetch sample rows to show actual data formats
try:
sample_cur = conn.execute(f'SELECT * FROM "{name}" LIMIT {sample_rows}')
col_names = [desc[0] for desc in sample_cur.description] if sample_cur.description else []
samples = sample_cur.fetchall()
except Exception:
col_names = []
samples = []
tables.append(
{
"name": name,
"create_sql": create_sql,
"columns": col_names,
"sample_rows": [list(r) for r in samples],
}
)
return {
"path": str(path),
"tables": tables,
}
def execute_read_only_sql(path: Path, sql: str, *, limit: int = 200) -> dict[str, object]:
normalized_sql = sql.lstrip().lower()
if not normalized_sql.startswith(("select", "with", "pragma")):
raise ValueError("Only read-only SQL statements are allowed.")
with _connect_read_only(path) as conn:
try:
cursor = conn.execute(sql)
except sqlite3.OperationalError as exc:
error_msg = str(exc)
# Provide helpful hints for common errors
hints: list[str] = []
if "no such column" in error_msg:
# Get available tables/columns for context
tables_cur = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
)
for (tbl,) in tables_cur.fetchall():
col_cur = conn.execute(f'PRAGMA table_info("{tbl}")')
cols = [r[1] for r in col_cur.fetchall()]
hints.append(f" Table '{tbl}' has columns: {cols}")
if "syntax error" in error_msg and "'" in sql:
hints.append(
" Hint: If a string value contains an apostrophe, escape it by doubling: "
"e.g. 'Women''s Soccer' instead of 'Women's Soccer'"
)
hint_text = "\n".join(hints)
raise sqlite3.OperationalError(
f"{error_msg}\n\nAvailable schema:\n{hint_text}" if hints else error_msg
) from exc
column_names = [item[0] for item in cursor.description or []]
rows = cursor.fetchmany(limit + 1)
truncated = len(rows) > limit
limited_rows = rows[:limit]
return {
"path": str(path),
"columns": column_names,
"rows": [list(row) for row in limited_rows],
"row_count": len(limited_rows),
"truncated": truncated,
}
# ---------------------------------------------------------------------------
# Option B: query_db helper injected into executor exec() namespace
# ---------------------------------------------------------------------------
def _fix_unescaped_quotes(sql: str) -> str:
"""Fix unescaped apostrophes inside SQL string literals."""
return re.sub(r"(\w)'(\w)", r"\1''\2", sql)
def _fix_column_case(sql: str, conn: sqlite3.Connection) -> str:
"""Attempt case-insensitive column name resolution in SQL."""
col_map: dict[str, str] = {}
try:
cur = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
)
for (tbl,) in cur.fetchall():
pragma_cur = conn.execute(f'PRAGMA table_info("{tbl}")')
for row in pragma_cur.fetchall():
actual_col = row[1]
col_map[actual_col.lower()] = actual_col
except Exception:
return sql
fixed = sql
for lower_name, actual_name in col_map.items():
pattern = re.compile(r'\b' + re.escape(actual_name) + r'\b', re.IGNORECASE)
fixed = pattern.sub(actual_name, fixed)
return fixed
def _get_schema_hint(conn: sqlite3.Connection) -> str:
"""Build a concise schema hint string."""
hints: list[str] = []
try:
cur = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
)
for (tbl,) in cur.fetchall():
col_cur = conn.execute(f'PRAGMA table_info("{tbl}")')
cols = [r[1] for r in col_cur.fetchall()]
hints.append(f" Table '{tbl}': {cols}")
except Exception:
pass
return "\n".join(hints)
def query_db(db_path: str, sql: str, *, limit: int = 200) -> pd.DataFrame:
"""High-level helper injected into executor namespace.
Runs read-only SQL against _auto.db and returns a DataFrame.
Auto-fixes common errors:
- Unescaped apostrophes in string literals (syntax error)
- Case-mismatched column names (no such column)
On failure, raises with full schema hint.
"""
path = Path(db_path)
conn = sqlite3.connect(f"file:{path.resolve().as_posix()}?mode=ro", uri=True)
try:
try:
df = pd.read_sql_query(sql, conn)
except Exception as exc:
error_msg = str(exc).lower()
# Auto-fix: syntax error from unescaped apostrophes
if "syntax error" in error_msg:
fixed_sql = _fix_unescaped_quotes(sql)
if fixed_sql != sql:
try:
df = pd.read_sql_query(fixed_sql, conn)
return df.head(limit) if len(df) > limit else df
except Exception:
pass
# Auto-fix: case-insensitive column resolution
if "no such column" in error_msg:
fixed_sql = _fix_column_case(sql, conn)
if fixed_sql != sql:
try:
df = pd.read_sql_query(fixed_sql, conn)
return df.head(limit) if len(df) > limit else df
except Exception:
pass
schema_hint = _get_schema_hint(conn)
raise type(exc)(
f"{exc}\n\nAvailable schema:\n{schema_hint}"
) from exc
raise
return df.head(limit) if len(df) > limit else df
finally:
conn.close()
|