Spaces:
Running
Running
| """ | |
| 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 | |
| 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, | |
| ) | |