ThothII/harness/tht/execute/__init__.py

import json
import time
from dataclasses import dataclass

import sqlglot
from sqlalchemy import Engine, text
from sqlalchemy.exc import DBAPIError


class ExecutionError(Exception):
    """Errore di esecuzione con messaggio leggibile per il reviewer."""


@dataclass
class ExecResult:
    columns: list[str]
    rows: list[tuple]
    execution_ms: int
    truncated: bool


@dataclass
class PlanSummary:
    total_cost: float
    plan_rows: int
    node_types: list[str]


def _inject_limit(sql: str, limit: int) -> tuple[str, bool]:
    """Aggiunge LIMIT limit+1 se assente (il +1 serve a rilevare il troncamento).
    Se la query ha gia' un suo LIMIT, lo si rispetta."""
    ast = sqlglot.parse_one(sql, read="postgres")
    # solo le query (SELECT/UNION & co.) accettano un LIMIT: per qualunque altro
    # statement non si inietta nulla e la transazione READ ONLY lo rifiutera' (rete 2).
    if not isinstance(ast, sqlglot.exp.Query) or ast.args.get("limit") is not None:
        return sql, False
    return ast.limit(limit + 1).sql(dialect="postgres"), True


def assert_read_only(sql: str) -> None:
    """Guard read-only strutturale condiviso dai due codepath di esecuzione (D7).

    Difesa in profondita': oltre alla rete server (utente RO / READ ONLY tx / RPC
    SELECT-only), ogni esecuzione passa da qui, cosi' anche un riuso diretto delle API
    Python (es. dal backend) non puo' bypassare il single-statement + SELECT-only.
    Usa sqlcheck.validate_sql senza schema fisico: parse + un solo statement +
    read-only per struttura + funzioni vietate.
    """
    from tht.sqlcheck import validate_sql

    check = validate_sql(sql)
    if not check.ok:
        raise ExecutionError(
            "SQL rifiutato (read-only enforcement): " + "; ".join(check.errors)
        )


def _translate_error(e: DBAPIError) -> ExecutionError:
    msg = str(e.orig)
    if "statement timeout" in msg or "canceling statement" in msg:
        return ExecutionError("timeout: la query ha superato statement_timeout_ms")
    if "read-only" in msg:
        return ExecutionError("rifiutato: statement non eseguibile in transazione read-only")
    return ExecutionError(f"errore SQL: {msg}")


def run_controlled(engine: Engine, sql: str, *, limit: int, timeout_ms: int) -> ExecResult:
    """Esecuzione nelle 4 reti: utente RO (a monte), transazione READ ONLY,
    statement_timeout, LIMIT iniettato via AST. Guard read-only client-side a monte."""
    assert_read_only(sql)
    final_sql, injected = _inject_limit(sql, limit)
    start = time.monotonic()
    with engine.connect() as conn:
        trans = conn.begin()
        try:
            conn.execute(text("SET TRANSACTION READ ONLY"))
            conn.execute(text(f"SET LOCAL statement_timeout = {int(timeout_ms)}"))
            result = conn.execute(text(final_sql))
            columns = list(result.keys())
            rows = [tuple(r) for r in result.fetchall()]
        except DBAPIError as e:
            raise _translate_error(e) from e
        finally:
            trans.rollback()
    elapsed_ms = int((time.monotonic() - start) * 1000)
    truncated = injected and len(rows) > limit
    return ExecResult(
        columns=columns, rows=rows[:limit], execution_ms=elapsed_ms, truncated=truncated
    )


def explain(engine: Engine, sql: str, *, timeout_ms: int) -> PlanSummary:
    """EXPLAIN (FORMAT JSON), mai ANALYZE: il piano si stima, non si esegue.

    assert_read_only a monte: l'EXPLAIN interpola lo SQL, quindi il single-statement +
    SELECT-only va garantito anche qui (non solo nel chiamante CLI)."""
    assert_read_only(sql)
    with engine.connect() as conn:
        trans = conn.begin()
        try:
            conn.execute(text("SET TRANSACTION READ ONLY"))
            conn.execute(text(f"SET LOCAL statement_timeout = {int(timeout_ms)}"))
            row = conn.execute(text(f"EXPLAIN (FORMAT JSON) {sql}")).fetchone()
        except DBAPIError as e:
            raise _translate_error(e) from e
        finally:
            trans.rollback()
    payload = row[0]
    if isinstance(payload, str):
        payload = json.loads(payload)
    plan = payload[0]["Plan"]

    node_types: list[str] = []

    def collect(node: dict) -> None:
        node_types.append(node.get("Node Type", "?"))
        for child in node.get("Plans", []):
            collect(child)

    collect(plan)
    return PlanSummary(
        total_cost=float(plan.get("Total Cost", 0.0)),
        plan_rows=int(plan.get("Plan Rows", 0)),
        node_types=node_types,
    )