validator.py

"""SQL validator.

Guarantees (defense-in-depth on top of the read-only DB driver):
  - Exactly one statement.
  - Statement is a SELECT (no INSERT / UPDATE / DELETE / DDL / ATTACH / PRAGMA).
  - No multi-statement injection via ';'.
  - Adds a LIMIT clause if missing (DoS protection).

Returns the normalized SQL on success, raises `InvalidSQL` on failure.
"""
from __future__ import annotations

import sqlglot
from sqlglot import exp


class InvalidSQL(Exception):
    pass


_FORBIDDEN_NODES = (
    exp.Insert,
    exp.Update,
    exp.Delete,
    exp.Create,
    exp.Drop,
    exp.Alter,
    exp.AlterColumn,
    exp.TruncateTable,
    exp.Pragma,
    exp.Attach,
    exp.Command,  # generic catch-all for VACUUM, REINDEX, ...
)


def validate_and_normalize(sql: str, *, max_limit: int) -> str:
    sql = (sql or "").strip().rstrip(";").strip()
    if not sql:
        raise InvalidSQL("Empty SQL.")

    try:
        statements = sqlglot.parse(sql, read="sqlite")
    except sqlglot.errors.ParseError as e:
        raise InvalidSQL(f"Parse error: {e}") from e

    statements = [s for s in statements if s is not None]
    if len(statements) != 1:
        raise InvalidSQL(
            f"Expected exactly one statement, got {len(statements)}. "
            "Do not chain statements with ';'."
        )

    stmt = statements[0]

    if not isinstance(stmt, exp.Select) and stmt.find(exp.Select) is None:
        raise InvalidSQL("Only SELECT statements are allowed.")

    for node in stmt.walk():
        if isinstance(node, _FORBIDDEN_NODES):
            raise InvalidSQL(
                f"Disallowed statement type: {type(node).__name__}. Only SELECT is permitted."
            )

    # Inject LIMIT if missing on the outermost SELECT.
    outer = stmt if isinstance(stmt, exp.Select) else stmt.find(exp.Select)
    if outer is not None and outer.args.get("limit") is None:
        outer.set("limit", exp.Limit(expression=exp.Literal.number(max_limit)))

    return stmt.sql(dialect="sqlite")