simbarag/blueprints/ynab/sql_tool.py

"""Admin-only read-only SQL over the YNAB mirror tables.

Defense in depth: the query is parsed and must be a single SELECT touching
only ``ynab_*`` tables with no dangerous functions, then runs in a READ ONLY
transaction with a statement timeout and a row cap.
"""

import re
from decimal import Decimal

import sqlglot
from langchain_core.runnables import RunnableConfig
from langchain.tools import tool
from sqlglot import exp
from sqlglot.optimizer.scope import traverse_scope
from tortoise.transactions import in_transaction

from . import models
from .sync import ensure_fresh

MAX_ROWS = 200
MAX_OUTPUT_CHARS = 12000
STATEMENT_TIMEOUT = "10s"

ALLOWED_TABLES = frozenset(
    m._meta.db_table
    for m in vars(models).values()
    if isinstance(m, type) and issubclass(m, models.Model) and m is not models.Model
)

FORBIDDEN_NODES = (
    exp.Insert,
    exp.Update,
    exp.Delete,
    exp.Merge,
    exp.Create,
    exp.Drop,
    exp.Alter,
    exp.Command,
    exp.Into,
)

FORBIDDEN_FUNCTIONS = re.compile(
    r"^(pg_|lo_|dblink|set_config|current_setting|"
    r"(query|table|cursor|schema|database)_to_xml|txid_|nextval|setval)",
    re.IGNORECASE,
)


def validate_sql(query: str) -> str | None:
    """Return an error message if ``query`` isn't allowed, else None."""
    try:
        statements = [s for s in sqlglot.parse(query, read="postgres") if s]
    except sqlglot.errors.ParseError as e:
        return f"Could not parse SQL: {e}"
    if len(statements) != 1:
        return "Exactly one SQL statement is allowed."
    root = statements[0]
    if not isinstance(root, exp.Query):
        return "Only SELECT queries are allowed."
    if root.find(*FORBIDDEN_NODES) or any(
        s.args.get("locks") for s in root.find_all(exp.Select)
    ):
        return "Only read-only SELECT queries are allowed."

    # A name matching a CTE isn't necessarily a CTE reference (e.g. inside the
    # CTE's own body, or outside the scope that defines it), so resolve each
    # FROM source through sqlglot's scopes. Any table node that isn't a CTE
    # name at all is checked too, in case scope analysis misses a construct.
    try:
        scopes = traverse_scope(root)
    except Exception as e:
        return f"Could not analyze SQL: {e}"
    cte_names = {cte.alias_or_name for cte in root.find_all(exp.CTE)}
    tables = [
        source
        for scope in scopes
        for source in scope.sources.values()
        if isinstance(source, exp.Table)
    ]
    tables += [t for t in root.find_all(exp.Table) if t.name not in cte_names]
    for table in tables:
        if isinstance(table.this, exp.Func):  # e.g. generate_series(...)
            continue
        if table.catalog or table.db not in ("", "public"):
            return f"Schema-qualified table '{table.sql()}' is not allowed."
        if table.name not in ALLOWED_TABLES:
            return (
                f"Table '{table.name}' is not allowed. "
                f"Available: {', '.join(sorted(ALLOWED_TABLES))}"
            )

    for func in root.find_all(exp.Func):
        name = func.name if isinstance(func, exp.Anonymous) else func.sql_name()
        if FORBIDDEN_FUNCTIONS.match(name or ""):
            return f"Function '{name}' is not allowed."
    return None


def _format_value(value) -> str:
    if value is None:
        return ""
    if isinstance(value, Decimal):
        value = (
            format(value.normalize(), "f") if value == value.to_integral() else value
        )
    return str(value).replace("\t", " ").replace("\n", " ")


def format_rows(rows: list[dict]) -> str:
    if not rows:
        return "Query returned no rows."
    header = list(rows[0])
    lines = ["\t".join(header)]
    lines += ["\t".join(_format_value(r[c]) for c in header) for r in rows]
    text = "\n".join(lines)
    if len(text) > MAX_OUTPUT_CHARS:
        text = text[:MAX_OUTPUT_CHARS] + "\n...(output truncated)"
    if len(rows) == MAX_ROWS:
        text += f"\n(capped at {MAX_ROWS} rows; aggregate or add LIMIT)"
    return text


@tool
async def ynab_sql(query: str, config: RunnableConfig) -> str:
    """Run a read-only Postgres SELECT over the YNAB mirror for questions the other YNAB tools can't answer. Tables: ynab_transactions (date, amount, payee_name, category_name, account_name, memo, transfer_account_id, deleted), ynab_subtransactions (split lines), ynab_accounts, ynab_categories, ynab_category_groups, ynab_payees, ynab_months, ynab_month_categories (month, category_id, budgeted=assigned, activity, balance=available), ynab_changes (edit history: synced_at, entity_type, action, label, changes JSON). Amounts are dollars; outflows negative. Filter NOT deleted.

    Args:
        query: A single PostgreSQL SELECT statement
    """
    if not (config or {}).get("configurable", {}).get("is_admin"):
        return "YNAB SQL access is restricted to admin users."
    error = validate_sql(query)
    if error:
        return error
    if await ensure_fresh() is None:
        return "YNAB data is unavailable: the budget has not been synced yet."

    query = query.strip().rstrip(";")
    try:
        async with in_transaction() as conn:
            await conn.execute_query("SET TRANSACTION READ ONLY")
            await conn.execute_query(
                f"SET LOCAL statement_timeout = '{STATEMENT_TIMEOUT}'"
            )
            rows = await conn.execute_query_dict(
                f"SELECT * FROM ({query}) AS q LIMIT {MAX_ROWS}"
            )
    except Exception as e:
        return f"SQL error: {e}"
    return format_rows(rows)