Blame for app/tools/sql_tool.py

"""
Outil SQL en lecture seule, exposé au SQL Agent.

Principe de sécurité : le LLM ne touche JAMAIS une vraie base de production.
Ici, les "tables" sont en réalité des vues DuckDB pointant vers des fichiers
Parquet stockés dans le dataset privé Hugging Face. Toute requête générée par
le LLM passe par `validate_sql` avant exécution :

  - une seule instruction ;
  - doit commencer par SELECT (ou WITH ... SELECT) ;
  - aucun mot-clé de modification (INSERT, UPDATE, DELETE, DROP, ...) ;
  - ne référence que les tables autorisées ;
  - le nombre de lignes renvoyées est plafonné.
"""
from __future__ import annotations

import re
import time
from dataclasses import dataclass, field

import duckdb
import sqlglot
from sqlglot import expressions as exp

from app.config import settings

FORBIDDEN_KEYWORDS = {
    "INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "TRUNCATE",
    "CREATE", "EXEC", "EXECUTE", "MERGE", "GRANT", "REVOKE",
    "ATTACH", "DETACH", "COPY", "PRAGMA", "CALL", "INSTALL", "LOAD",
}

ALLOWED_TABLES = {"sales", "products", "refunds", "discounts"}

MAX_ROWS = 500

SCHEMA_DESCRIPTION = """
Tables disponibles (lecture seule, moteur DuckDB) :

- products(product_id VARCHAR, product_name VARCHAR, category VARCHAR,
           list_price DOUBLE, unit_cost DOUBLE)
- sales(order_id BIGINT, order_date DATE, product_id VARCHAR, region VARCHAR,
        channel VARCHAR, units_sold INTEGER, unit_price DOUBLE,
        unit_cost DOUBLE, discount_rate DOUBLE, campaign_name VARCHAR)
  -- unit_price/unit_cost = prix et coût RÉELS de la commande (déjà remisés)
- refunds(refund_id BIGINT, order_id BIGINT, refund_date DATE,
          refund_amount DOUBLE, reason VARCHAR)
- discounts(campaign_name VARCHAR, category VARCHAR, region VARCHAR,
            channel VARCHAR, start_date DATE, end_date DATE,
            discount_rate DOUBLE, approved_by VARCHAR)

Notes :
- Le chiffre d'affaires net d'une ligne de vente = units_sold * unit_price.
- Le coût des ventes d'une ligne = units_sold * unit_cost.
- La marge brute = chiffre d'affaires net - coût des ventes - remboursements
  imputés sur la période (jointure sales -> refunds via order_id).
- Utilise EXTRACT(month FROM order_date) / EXTRACT(month FROM refund_date)
  pour comparer des mois entre eux.
""".strip()


class SqlValidationError(ValueError):
    pass


@dataclass
class SqlQueryResult:
    sql: str
    columns: list[str] = field(default_factory=list)
    rows: list[dict] = field(default_factory=list)
    row_count: int = 0
    truncated: bool = False
    duration_ms: float = 0.0


def validate_sql(query: str) -> str:
    """Valide une requête SQL et renvoie la requête nettoyée, ou lève SqlValidationError.

    La validation s'appuie sur une analyse syntaxique réelle (sqlglot) plutôt
    que sur de simples regex, afin d'éviter les faux positifs (par exemple
    `EXTRACT(month FROM order_date)` ne doit pas être interprété comme une
    référence à une table `order_date`) tout en bloquant fiablement toute
    tentative de DDL/DML ou de référence à une table non autorisée.
    """
    if not query or not query.strip():
        raise SqlValidationError("Requête SQL vide.")

    cleaned = query.strip().rstrip(";").strip()
    upper = cleaned.upper()

    if not (upper.startswith("SELECT") or upper.startswith("WITH")):
        raise SqlValidationError("Seules les requêtes SELECT (éventuellement précédées de WITH) sont autorisées.")

    # Garde-fou brut, en complément de l'analyse syntaxique ci-dessous.
    tokens = set(re.findall(r"[A-Za-z_]+", upper))
    forbidden_found = tokens & FORBIDDEN_KEYWORDS
    if forbidden_found:
        raise SqlValidationError(f"Mot-clé SQL interdit détecté : {', '.join(sorted(forbidden_found))}.")

    try:
        statements = sqlglot.parse(cleaned, read="duckdb")
    except sqlglot.errors.ParseError as exc:
        raise SqlValidationError(f"Requête SQL invalide : {exc}") from exc

    statements = [s for s in statements if s is not None]
    if len(statements) != 1:
        raise SqlValidationError("Une seule instruction SQL est autorisée par appel.")

    tree = statements[0]
    if not isinstance(tree, (exp.Select, exp.Union, exp.Subquery)):
        raise SqlValidationError("Seules les requêtes de lecture (SELECT/UNION) sont autorisées.")

    cte_names = {c.alias_or_name.lower() for c in tree.find_all(exp.CTE)}
    referenced_tables = {t.name.lower() for t in tree.find_all(exp.Table) if t.name}
    unknown_tables = referenced_tables - ALLOWED_TABLES - cte_names
    if unknown_tables:
        raise SqlValidationError(
            f"Table(s) non autorisée(s) : {', '.join(sorted(unknown_tables))}. "
            f"Tables autorisées : {', '.join(sorted(ALLOWED_TABLES))}."
        )

    return cleaned


def _connect() -> duckdb.DuckDBPyConnection:
    con = duckdb.connect(database=":memory:")
    table_paths = {
        "sales": settings.sales_parquet_path.parent / "sales.parquet",
        "products": settings.sales_parquet_path.parent / "products.parquet",
        "refunds": settings.sales_parquet_path.parent / "refunds.parquet",
        "discounts": settings.sales_parquet_path.parent / "discounts.parquet",
    }
    for name, path in table_paths.items():
        if path.exists():
            con.execute(f"CREATE VIEW {name} AS SELECT * FROM read_parquet('{path.as_posix()}')")
    return con


def run_sql(query: str) -> SqlQueryResult:
    """Valide puis exécute une requête SQL en lecture seule. Lève SqlValidationError
    si la requête est rejetée, ou duckdb.Error si l'exécution échoue."""
    cleaned = validate_sql(query)

    start = time.perf_counter()
    con = _connect()
    try:
        limited = f"SELECT * FROM ({cleaned}) AS _subquery LIMIT {MAX_ROWS + 1}"
        cursor = con.execute(limited)
        columns = [d[0] for d in cursor.description]
        raw_rows = cursor.fetchall()
    finally:
        con.close()
    duration_ms = (time.perf_counter() - start) * 1000

    truncated = len(raw_rows) > MAX_ROWS
    raw_rows = raw_rows[:MAX_ROWS]
    rows = [dict(zip(columns, row)) for row in raw_rows]
    # Rend les valeurs sérialisables en JSON (dates, Decimal, etc.).
    for row in rows:
        for key, value in row.items():
            if hasattr(value, "isoformat"):
                row[key] = value.isoformat()

    return SqlQueryResult(
        sql=cleaned,
        columns=columns,
        rows=rows,
        row_count=len(rows),
        truncated=truncated,
        duration_ms=duration_ms,
    )