56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
"""Execute read-only SQL against the financial warehouse.
|
|
|
|
Defence in depth:
|
|
- Engine-level `search_path` and `statement_timeout` via libpq options in
|
|
`api.tools.db.get_engine` — no SQL-side SETs.
|
|
- `execution_options(postgresql_readonly=True, postgresql_deferrable=True)`
|
|
— SQLAlchemy emits `SET TRANSACTION READ ONLY DEFERRABLE` internally; we
|
|
don't compose a SQL string for it.
|
|
- Hard refusal of any statement that doesn't begin with SELECT or WITH.
|
|
- Row cap on the returned result set.
|
|
- No commit — read-only transactions roll back on context exit.
|
|
|
|
The SQL passed in IS a string (it's the LLM's output); nothing on the Python
|
|
side concatenates SQL fragments around it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from sqlalchemy import text
|
|
|
|
from api import langfuse_client as lf
|
|
from api.tools.db import get_engine
|
|
from api.tools.types import QueryResult
|
|
|
|
logger = logging.getLogger("nvi.tools.execute_sql")
|
|
|
|
MAX_ROWS = 1000
|
|
|
|
|
|
def execute_sql(sql: str) -> QueryResult:
|
|
sql_clean = sql.strip().rstrip(";").strip()
|
|
head = sql_clean.split(None, 1)[0].upper() if sql_clean else ""
|
|
if head not in {"SELECT", "WITH"}:
|
|
raise ValueError(f"refusing non-SELECT statement (starts with {head!r})")
|
|
|
|
engine = get_engine()
|
|
with lf.span("execute_sql", input={"sql": sql_clean}) as span:
|
|
ro_engine = engine.execution_options(
|
|
postgresql_readonly=True,
|
|
postgresql_deferrable=True,
|
|
)
|
|
with ro_engine.connect() as conn:
|
|
res = conn.execute(text(sql_clean))
|
|
columns = list(res.keys())
|
|
fetched = res.fetchmany(MAX_ROWS + 1)
|
|
|
|
truncated = len(fetched) > MAX_ROWS
|
|
rows = [list(r) for r in fetched[:MAX_ROWS]]
|
|
result = QueryResult(
|
|
columns=columns, rows=rows, row_count=len(rows), truncated=truncated
|
|
)
|
|
span.update(output={"row_count": result.row_count, "truncated": truncated})
|
|
return result
|