Spaces:
Sleeping
Sleeping
| 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() | |