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 utils.mcp_tools import admin_only_message

from . import models
from .sync import ensure_fresh
from .tools import NOT_SYNCED, ynab_tool_errors

MAX_ROWS = 100
# Tool results are cut at 4000 chars in blueprints/conversation/logic.py
MAX_OUTPUT_CHARS = 3800
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."""
    return _check(query)[0]


def _check(query: str) -> tuple[str | None, str | None]:
    """Validate ``query``; return ``(error, sql_to_run)``.

    The SQL that runs is regenerated from the validated parse tree (minus
    comments), so Postgres executes exactly what was checked.
    """
    try:
        statements = [
            s
            for s in sqlglot.parse(query, read="postgres")
            # a trailing "; -- comment" parses as a bare Semicolon
            if s and not isinstance(s, exp.Semicolon)
        ]
    except sqlglot.errors.SqlglotError as e:
        return f"Could not parse SQL: {e}", None
    if len(statements) != 1:
        return "Exactly one SQL statement is allowed.", None
    root = statements[0]
    if not isinstance(root, exp.Query):
        return "Only SELECT queries are allowed.", None
    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.", None

    # 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}", None
    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.", None
        if table.name not in ALLOWED_TABLES:
            return (
                f"Table '{table.name}' is not allowed. "
                f"Available: {', '.join(sorted(ALLOWED_TABLES))}",
                None,
            )

    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.", None
    return None, root.sql(dialect="postgres", comments=False).rstrip("; ")


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(columns: list[str], rows: list[tuple]) -> str:
    if not rows:
        return "Query returned no rows."
    lines = ["\t".join(columns)]
    lines += ["\t".join(_format_value(v) for v in row) for row in rows]
    text = "\n".join(lines)
    if len(text) > MAX_OUTPUT_CHARS:
        text = text[:MAX_OUTPUT_CHARS] + "\n...(output truncated; aggregate more)"
    elif len(rows) == MAX_ROWS:
        text += f"\n(capped at {MAX_ROWS} rows; aggregate or add LIMIT)"
    return text


@tool
@ynab_tool_errors
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 admin_only_message("YNAB SQL")
    error, sql = _check(query)
    if error:
        return error
    if await ensure_fresh() is None:
        return NOT_SYNCED

    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}'"
            )
            # Raw records rather than dicts, which would merge same-named
            # columns (e.g. c.name and g.name)
            _, records = await conn.execute_query(
                f"SELECT * FROM ({sql}) AS q LIMIT {MAX_ROWS}"
            )
    except Exception as e:
        return f"SQL error: {e}"
    columns = list(records[0].keys()) if records else []
    return format_rows(columns, [tuple(r.values()) for r in records])