Spaces:
Sleeping
Sleeping
File size: 4,759 Bytes
49ba367 794e918 49ba367 | 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 | """
SQLite in-memory execution engine.
Provides deterministic query execution with timeout protection.
"""
from __future__ import annotations
import sqlite3
import threading
import time
from typing import List, Optional, Tuple
class SQLiteEngine:
"""In-memory SQLite executor with schema setup and safe query running."""
def __init__(self):
self._conn: Optional[sqlite3.Connection] = None
self._create_connection()
def _create_connection(self):
self._conn = sqlite3.connect(":memory:", check_same_thread=False)
self._conn.execute("PRAGMA foreign_keys = ON")
def setup(self, schema_sql: str, seed_data_sql: str) -> None:
"""Execute DDL and seed data to prepare the database."""
cursor = self._conn.cursor()
cursor.executescript(schema_sql)
cursor.executescript(seed_data_sql)
self._conn.commit()
def execute(self, query: str, timeout_s: int = 5) -> Tuple[Optional[List[tuple]], Optional[str]]:
"""
Execute a SQL query and return (rows, None) on success or (None, error) on failure.
Uses SQLite's progress_handler for thread-safe timeout.
"""
deadline = time.monotonic() + timeout_s
def _check_timeout():
if time.monotonic() > deadline:
return 1 # Non-zero cancels the operation
return 0
try:
# Check every 1000 SQLite VM instructions
self._conn.set_progress_handler(_check_timeout, 1000)
cursor = self._conn.cursor()
cursor.execute(query)
rows = cursor.fetchall()
return rows, None
except sqlite3.OperationalError as e:
if "interrupted" in str(e).lower():
return None, "Query execution timed out (exceeded 5 seconds)"
return None, str(e)
except sqlite3.Error as e:
return None, str(e)
except Exception as e:
return None, f"Unexpected error: {str(e)}"
finally:
self._conn.set_progress_handler(None, 0)
def explain_query(self, query: str) -> Optional[str]:
"""Return the SQLite query plan for the given query as a string.
Executes ``EXPLAIN QUERY PLAN {query}`` and formats each row of the
plan into a human-readable string. Returns None if an error occurs.
"""
try:
cursor = self._conn.cursor()
cursor.execute(f"EXPLAIN QUERY PLAN {query}")
rows = cursor.fetchall()
return "\n".join(str(row) for row in rows)
except sqlite3.Error:
return None
def get_table_names(self) -> List[str]:
"""Return a list of all table names in the current database."""
try:
cursor = self._conn.cursor()
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
return [row[0] for row in cursor.fetchall()]
except sqlite3.Error:
return []
def get_column_names(self, table_name: str) -> List[str]:
"""Return a list of column names for the given table.
Uses ``PRAGMA table_info`` to retrieve column metadata. Returns an
empty list if the table does not exist or an error occurs.
"""
try:
cursor = self._conn.cursor()
cursor.execute(f"PRAGMA table_info({table_name})")
return [row[1] for row in cursor.fetchall()]
except sqlite3.Error:
return []
@staticmethod
def categorize_error(error_msg: str) -> str:
"""Categorize a SQLite error message into a broad error type.
Args:
error_msg: The raw error string returned by SQLite.
Returns:
One of: "syntax", "missing_table", "missing_column",
"ambiguous_column", "type_error", "constraint", "timeout",
"unknown".
"""
msg = error_msg.lower()
if "no such table" in msg:
return "missing_table"
if "no such column" in msg:
return "missing_column"
if "ambiguous" in msg:
return "ambiguous_column"
if "near" in msg or "syntax error" in msg:
return "syntax"
if "datatype mismatch" in msg or "type" in msg:
return "type_error"
if "constraint" in msg or "unique" in msg or "not null" in msg or "foreign key" in msg:
return "constraint"
if "timed out" in msg or "timeout" in msg or "interrupted" in msg:
return "timeout"
return "unknown"
def reset(self) -> None:
"""Drop the current connection and create a fresh one."""
if self._conn:
self._conn.close()
self._create_connection()
|