src/nodes/sql_safety.py

import re
from src.state import AgentState

BLOCKED = r'\b(DROP|DELETE|INSERT|UPDATE|ALTER|TRUNCATE|CREATE|EXEC|GRANT|REVOKE|UNION)\b'
SYSTEM_TABLES = ["sqlite_master", "sqlite_sequence", "information_schema"]

def check_sql_safety(state: AgentState) -> dict:
    """Validate SQL query for safety before execution"""
    query = state.get("sql_query", "")

    if not query:
        return {"sql_result": {"error": "no_query"}}

    # Check for dangerous keywords
    if re.search(BLOCKED, query, re.IGNORECASE):
        return {
            "sql_query": None,
            "sql_result": {"error": "unsafe_query"},
            "error_log": state.get("error_log", []) + ["[sql_safety] Blocked dangerous SQL keyword"]
        }

    # Check for system table access
    for sys_table in SYSTEM_TABLES:
        if sys_table in query.lower():
            return {
                "sql_query": None,
                "sql_result": {"error": "system_table_access"},
                "error_log": state.get("error_log", []) + [f"[sql_safety] Blocked system table access: {sys_table}"]
            }

    # Ensure SELECT only
    stripped = query.strip().upper()
    if not stripped.startswith("SELECT"):
        return {
            "sql_query": None,
            "sql_result": {"error": "non_select"},
            "error_log": state.get("error_log", []) + ["[sql_safety] Non-SELECT query blocked"]
        }

    return {}