Spaces:
Sleeping
Sleeping
| """SQL safety guards for DB-ER. | |
| Three classes of irreversible-damage checks (spec section 2.3): | |
| 1. Schema Destruction: DROP TABLE on a core table - episode ends, reward -1.0 | |
| 2. Catastrophic Loss: core table loses >10% of rows - episode ends, reward -1.0 | |
| 3. Truncation: DELETE without a WHERE clause - episode ends, reward -1.0 | |
| ``check_pre_execution(query, core_tables)`` runs *before* the SQL hits the DB. | |
| ``check_post_execution(conn, initial_counts)`` runs *after* a mutation succeeds. | |
| """ | |
| from __future__ import annotations | |
| import re | |
| from dataclasses import dataclass | |
| from typing import Dict, List, Optional | |
| class SafetyViolation: | |
| """Carries the reason for a safety violation.""" | |
| kind: str # "SCHEMA_DESTRUCTION" | "CATASTROPHIC_LOSS" | "TRUNCATION" | |
| message: str | |
| # Query normalisation | |
| def _strip_comments(sql: str) -> str: | |
| """Remove SQL line comments and block comments.""" | |
| sql = re.sub(r"--[^\n]*", " ", sql) | |
| sql = re.sub(r"/\*.*?\*/", " ", sql, flags=re.DOTALL) | |
| return sql | |
| def _normalise(sql: str) -> str: | |
| """Collapse whitespace and upper-case for safe pattern matching.""" | |
| return re.sub(r"\s+", " ", _strip_comments(sql)).strip().upper() | |
| # Pre-execution checks | |
| def check_pre_execution( | |
| query: str, core_tables: List[str] | |
| ) -> Optional[SafetyViolation]: | |
| """Analyse the raw SQL *before* running it. | |
| Returns a :class:`SafetyViolation` if the query is forbidden, else ``None``. | |
| Uses regex as the primary detection layer (more reliable than sqlparse | |
| for edge cases like ``DROP TABLE IF EXISTS``) with sqlparse as a secondary | |
| confirmation for the DELETE+WHERE check. | |
| """ | |
| v = _check_drop_regex(query, core_tables) | |
| if v: | |
| return v | |
| v = _check_bare_delete_regex(query) | |
| return v | |
| # DROP TABLE check | |
| # Matches: DROP TABLE [IF EXISTS] <name> | |
| # Handles optional backtick / quote / bracket quoting around the table name. | |
| _DROP_RE = re.compile( | |
| r"\bDROP\s+TABLE\s+(?:IF\s+EXISTS\s+)?([`\"\[]?\w+[`\"\]]?)", | |
| re.IGNORECASE, | |
| ) | |
| def _check_drop_regex(query: str, core_tables: List[str]) -> Optional[SafetyViolation]: | |
| q_clean = _strip_comments(query) | |
| core_lower = {t.lower() for t in core_tables} | |
| for match in _DROP_RE.finditer(q_clean): | |
| raw_name = match.group(1) | |
| table_name = raw_name.strip("`\"[]").lower() | |
| if table_name in core_lower: | |
| return SafetyViolation( | |
| kind="SCHEMA_DESTRUCTION", | |
| message=( | |
| f"INSTANT FAIL: DROP TABLE on core table '{table_name}' is forbidden. " | |
| "Core tables are protected from schema destruction. " | |
| "You may create temporary helper tables and drop those instead." | |
| ), | |
| ) | |
| return None | |
| # DELETE without WHERE check | |
| def _check_bare_delete_regex(query: str) -> Optional[SafetyViolation]: | |
| """Detect DELETE FROM <table> without a WHERE clause. | |
| Strategy: | |
| 1. If the statement has no DELETE keyword at all, skip. | |
| 2. If it has a WHERE keyword (anywhere in the statement), assume it's safe. | |
| 3. Otherwise, flag as a bare delete. | |
| This conservative approach may flag rare edge-cases (e.g. DELETE inside a CTE) | |
| but that is acceptable: the agent should always use WHERE in DELETE statements. | |
| """ | |
| q_norm = _normalise(query) | |
| if not re.search(r"\bDELETE\b", q_norm): | |
| return None | |
| if re.search(r"\bWHERE\b", q_norm): | |
| return None # WHERE clause present - safe | |
| return SafetyViolation( | |
| kind="TRUNCATION", | |
| message=( | |
| "INSTANT FAIL: DELETE without a WHERE clause detected. " | |
| "Wholesale table truncation is forbidden. " | |
| "Always specify a WHERE condition to target only the rows you intend to remove." | |
| ), | |
| ) | |
| # Post-execution checks | |
| def check_post_execution( | |
| conn, | |
| initial_counts: Dict[str, int], | |
| ) -> Optional[SafetyViolation]: | |
| """Check row counts *after* a mutation to detect catastrophic data loss. | |
| Returns a :class:`SafetyViolation` if any core table has lost >10% of its | |
| original rows, else ``None``. | |
| """ | |
| for table, initial in initial_counts.items(): | |
| if initial == 0: | |
| continue # nothing to protect | |
| try: | |
| row = conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone() | |
| current = row[0] | |
| except Exception: | |
| continue # table may not exist yet - handled elsewhere | |
| if current < 0.90 * initial: | |
| pct_lost = round((1 - current / initial) * 100, 1) | |
| return SafetyViolation( | |
| kind="CATASTROPHIC_LOSS", | |
| message=( | |
| f"INSTANT FAIL: Table '{table}' lost {pct_lost}% of its rows " | |
| f"(was {initial}, now {current}). " | |
| "Core tables must retain at least 90% of initial row count." | |
| ), | |
| ) | |
| return None | |
| # SQL type classification | |
| def classify_sql(query: str) -> str: | |
| """Return SELECT | MUTATION | META for budget/reward purposes.""" | |
| q = _normalise(query) | |
| if re.match(r"\b(SELECT|WITH|EXPLAIN|VALUES)\b", q): | |
| return "SELECT" | |
| if re.match(r"\bPRAGMA\b", q): | |
| return "META" | |
| if re.match(r"\b(INSERT|UPDATE|DELETE|ALTER|CREATE|DROP|REPLACE|UPSERT)\b", q): | |
| return "MUTATION" | |
| return "SELECT" # conservative default | |
| def budget_cost(sql_type: str) -> int: | |
| """Return the budget cost for a given SQL type (spec section 2.1).""" | |
| return 1 if sql_type in ("SELECT", "META") else 5 | |