""" 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, )