How to Build an AI SQL Agent That Understands Natural Language and Queries Databases

ai-sql-agent/
├── app/
│   ├── __init__.py
│   ├── ast_guard.py
│   ├── db_executor.py
│   └── sql_agent.py
├── tests/
│   └── test_sql_agent.py
└── requirements.txt

pip install sqlglot==25.24.0 sqlalchemy==2.0.35 pydantic==2.9.2 openai==1.51.0 pytest==8.3.3

from typing import Set
import sqlglot
from sqlglot import exp

FORBIDDEN_AST_NODES = (
    exp.Insert,
    exp.Update,
    exp.Delete,
    exp.Drop,
    exp.Create,
    exp.Alter,
    exp.Command,
    exp.Grant,
    exp.Revoke,
    exp.Transaction,
    exp.Commit,
    exp.Rollback,
    exp.Copy,
    exp.Pragma,
)

FORBIDDEN_FUNCTIONS = {
    "pg_read_file",
    "pg_ls_dir",
    "pg_execute_server_program",
    "pg_sleep",
    "dblink",
    "lo_import",
    "lo_export",
    "load_extension",
}

def validate_and_rewrite_readonly_sql(
    raw_sql: str,
    allowed_tables: Set[str],
    dialect: str = "postgres",
    max_limit: int = 100,
) -> str:
    """Validates SQL via SQLGlot AST inspection and enforces a hard LIMIT cap."""
    try:
        parsed_statements = sqlglot.parse(raw_sql, read=dialect)
    except sqlglot.errors.ParseError as exc:
        raise ValueError(f"SQL syntax parse error: {exc}") from exc

    # 1. Block stacked multi-statement queries (e.g. SELECT 1; DROP TABLE x;)
    non_empty_stmts = [s for s in parsed_statements if s is not None]
    if len(non_empty_stmts) != 1:
        raise PermissionError("Exactly one SQL statement is permitted per query.")

    ast = non_empty_stmts[0]

    # 2. Root expression must be a SELECT or UNION of SELECTs
    if not isinstance(ast, (exp.Select, exp.Union)):
        raise PermissionError(f"Only SELECT queries are permitted. Found root: {type(ast).__name__}")

    # 3. Walk the entire AST to block writable CTEs or DDL/DML nodes
    for node in ast.walk():
        if isinstance(node, FORBIDDEN_AST_NODES):
            raise PermissionError(f"Forbidden SQL operation detected in AST: {type(node).__name__}")
        if isinstance(node, exp.Anonymous) and node.name.lower() in FORBIDDEN_FUNCTIONS:
            raise PermissionError(f"Forbidden database function called: {node.name}")
        if isinstance(node, exp.Func) and node.sql_name().lower() in FORBIDDEN_FUNCTIONS:
            raise PermissionError(f"Forbidden database function called: {node.sql_name()}")

    # 4. Allowlist check on every referenced table (excluding CTE aliases)
    cte_aliases = {cte.alias_or_name.lower() for cte in ast.find_all(exp.CTE)}
    allowed_lower = {t.lower() for t in allowed_tables}

    for table_node in ast.find_all(exp.Table):
        tbl_name = table_node.name.lower()
        if tbl_name not in cte_aliases and tbl_name not in allowed_lower:
            raise PermissionError(f"Query references unauthorized or unknown table: '{table_node.name}'")

    # 5. Enforce or cap LIMIT clause on top-level SELECT
    if isinstance(ast, exp.Select):
        existing_limit = ast.args.get("limit")
        if existing_limit is None:
            ast = ast.limit(max_limit)
        else:
            try:
                limit_val = int(existing_limit.expression.name)
                if limit_val > max_limit:
                    ast = ast.limit(max_limit)
            except (ValueError, AttributeError):
                ast = ast.limit(max_limit)

    return ast.sql(dialect=dialect)

import sqlite3
from typing import List, Dict, Any, Set
from app.ast_guard import validate_and_rewrite_readonly_sql

def execute_safe_sqlite_query(
    db_path: str,
    raw_sql: str,
    allowed_tables: Set[str],
    max_limit: int = 100,
) -> Dict[str, Any]:
    """Validates SQL via SQLGlot AST and executes it against a read-only SQLite connection."""
    safe_sql = validate_and_rewrite_readonly_sql(
        raw_sql=raw_sql,
        allowed_tables=allowed_tables,
        dialect="sqlite",
        max_limit=max_limit,
    )
    conn = sqlite3.connect(db_path, timeout=5.0)
    conn.row_factory = sqlite3.Row
    try:
        conn.execute("PRAGMA query_only = ON;")
        cursor = conn.execute(safe_sql)
        rows = [dict(row) for row in cursor.fetchall()]
        columns = [desc[0] for desc in cursor.description] if cursor.description else []
        return {
            "executed_sql": safe_sql,
            "columns": columns,
            "row_count": len(rows),
            "rows": rows,
        }
    finally:
        conn.close()

import sqlite3
import pytest
from app.ast_guard import validate_and_rewrite_readonly_sql
from app.db_executor import execute_safe_sqlite_query

ALLOWED = {"customers", "orders", "order_items"}

def test_ast_guard_blocks_stacked_queries_and_writable_ctes():
    # Block stacked DROP TABLE
    with pytest.raises(PermissionError, match="Exactly one SQL statement"):
        validate_and_rewrite_readonly_sql("SELECT * FROM orders; DROP TABLE orders;", ALLOWED)

    # Block writable CTE
    with pytest.raises(PermissionError, match="Forbidden SQL operation"):
        validate_and_rewrite_readonly_sql(
            "WITH d AS (DELETE FROM orders RETURNING *) SELECT * FROM d", ALLOWED
        )

    # Block unauthorized table (e.g., payroll_salaries)
    with pytest.raises(PermissionError, match="unauthorized or unknown table"):
        validate_and_rewrite_readonly_sql("SELECT * FROM payroll_salaries", ALLOWED)

def test_ast_guard_allows_cte_and_injects_limit():
    sql = """
    WITH regional_totals AS (
        SELECT customer_id, SUM(amount) AS total_spent
        FROM orders
        GROUP BY customer_id
    )
    SELECT c.name, r.total_spent
    FROM customers c
    JOIN regional_totals r ON c.id = r.customer_id
    ORDER BY r.total_spent DESC
    """
    safe_sql = validate_and_rewrite_readonly_sql(sql, ALLOWED, dialect="sqlite", max_limit=50)
    assert "LIMIT 50" in safe_sql

def test_end_to_end_readonly_sqlite_execution(tmp_path):
    db_file = str(tmp_path / "analytics.db")
    conn = sqlite3.connect(db_file)
    conn.executescript(
        """
        CREATE TABLE customers (id INTEGER PRIMARY KEY, name TEXT, region TEXT);
        CREATE TABLE orders (id INTEGER PRIMARY KEY, customer_id INTEGER, amount REAL);
        INSERT INTO customers VALUES (1, 'Acme Corp', 'APAC'), (2, 'Globex', 'NA');
        INSERT INTO orders VALUES (101, 1, 450.0), (102, 1, 550.0), (103, 2, 300.0);
        """
    )
    conn.close()

    res = execute_safe_sqlite_query(
        db_file,
        "SELECT c.name, SUM(o.amount) AS revenue FROM customers c JOIN orders o ON c.id = o.customer_id GROUP BY c.name",
        allowed_tables={"customers", "orders"},
    )
    assert res["row_count"] == 2
    assert res["rows"][0]["revenue"] == 1000.0