Spaces:
Sleeping
Sleeping
File size: 3,499 Bytes
dc59b01 | 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 | import sqlite3
import re
from pathlib import Path
class SQLValidator:
def __init__(self, db_root):
self.db_root = Path(db_root)
# ---------------------------
# Load schema
# ---------------------------
def load_schema(self, db_id):
db_path = self.db_root / db_id / f"{db_id}.sqlite"
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
tables = cursor.execute(
"SELECT name FROM sqlite_master WHERE type='table';"
).fetchall()
schema = {}
for (table,) in tables:
cols = cursor.execute(f"PRAGMA table_info({table});").fetchall()
schema[table.lower()] = [c[1].lower() for c in cols]
conn.close()
return schema
# ---------------------------
# Basic syntax check
# ---------------------------
def basic_structure_valid(self, sql):
s = sql.lower()
if "select" not in s or "from" not in s:
return False, "Missing SELECT or FROM"
if len(s.split()) < 4:
return False, "Too short to be SQL"
return True, None
# ---------------------------
# Extract identifiers
# ---------------------------
def extract_identifiers(self, sql):
tokens = re.findall(r"[A-Za-z_]+", sql.lower())
return set(tokens)
# ---------------------------
# Table validation
# ---------------------------
def validate_tables(self, sql, schema):
words = self.extract_identifiers(sql)
tables = set(schema.keys())
used_tables = [w for w in words if w in tables]
if not used_tables:
return False, "No valid table used"
return True, None
# ---------------------------
# Column validation
# ---------------------------
def validate_columns(self, sql, schema):
words = self.extract_identifiers(sql)
valid_columns = set()
for cols in schema.values():
valid_columns.update(cols)
# ignore SQL keywords
keywords = {
"select","from","where","join","on","group","by",
"order","limit","count","sum","avg","min","max",
"and","or","in","like","distinct","asc","desc"
}
invalid = []
for w in words:
if w not in valid_columns and w not in schema and w not in keywords:
if not w.isdigit():
invalid.append(w)
# allow small hallucinations but block many
if len(invalid) > 3:
return False, f"Too many unknown identifiers: {invalid[:5]}"
return True, None
# ---------------------------
# Dangerous query protection
# ---------------------------
def block_dangerous(self, sql):
bad = ["drop", "delete", "update", "insert", "alter"]
s = sql.lower()
for b in bad:
if b in s:
return False, f"Dangerous keyword detected: {b}"
return True, None
# ---------------------------
# Main validation
# ---------------------------
def validate(self, sql, db_id):
schema = self.load_schema(db_id)
checks = [
self.block_dangerous(sql),
self.basic_structure_valid(sql),
self.validate_tables(sql, schema),
self.validate_columns(sql, schema),
]
for ok, msg in checks:
if not ok:
return False, msg
return True, None
|