Update app.py

app.py

@@ -4,6 +4,8 @@ import zipfile
 import re
 import difflib
 import tempfile
+import sqlite3
+import uuid
 from typing import List, Optional, Dict, Any
 
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form
@@ -15,26 +17,19 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from langdetect import detect
 from transformers import MarianMTModel, MarianTokenizer
 from openai import OpenAI
-import sys
-
-# --- forzar que el directorio actual (donde está app.py y sqlmanager.py) esté en sys.path ---
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-if BASE_DIR not in sys.path:
-    sys.path.append(BASE_DIR)
-
-from sqlmanager import SQLManager
 
 # ======================================================
-# 0) Configuración general
+# 0) Configuración general de paths
 # ======================================================
 
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+UPLOAD_DIR = os.path.join(BASE_DIR, "uploaded_dbs")
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+
 # Modelo NL→SQL entrenado por ti en Hugging Face
 MODEL_DIR = os.getenv("MODEL_DIR", "stvnnnnnn/t5-large-nl2sql-spider")
 DEVICE = torch.device("cpu")  # inferencia en CPU
 
-# Gestor de conexiones reales (MySQL/PostgreSQL)
-sql_manager = SQLManager()
-
 # Cliente OpenAI para transcripción de audio (Whisper / gpt-4o-transcribe)
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
@@ -42,29 +37,188 @@ if not OPENAI_API_KEY:
 openai_client = OpenAI(api_key=OPENAI_API_KEY) if OPENAI_API_KEY else None
 
 # ======================================================
-# 1) Inicialización de FastAPI
+# 1) SQLManager (versión actual: SQLite local)
+# ======================================================
+
+class SQLManager:
+    """
+    Gestor de "conexiones" a bases dinámicas.
+    Versión actual: cada conexión es un archivo SQLite en UPLOAD_DIR.
+    API pensada para poder cambiar después a Postgres/MySQL (Railway).
+    """
+
+    def __init__(self):
+        # connections[connection_id] = {
+        #   "label": str,
+        #   "engine": "sqlite",
+        #   "db_name": str,
+        #   "db_path": str
+        # }
+        self.connections: Dict[str, Dict[str, Any]] = {}
+
+    # ---------- utilidades internas ----------
+
+    def _new_connection_id(self) -> str:
+        return f"db_{uuid.uuid4().hex[:8]}"
+
+    def _get_info(self, connection_id: str) -> Dict[str, Any]:
+        if connection_id not in self.connections:
+            raise KeyError(f"connection_id '{connection_id}' no registrado")
+        return self.connections[connection_id]
+
+    # ---------- creación de BD desde dump ----------
+
+    def create_database_from_dump(self, label: str, sql_text: str) -> str:
+        """
+        Crea un archivo SQLite, ejecuta el dump SQL y
+        registra la conexión. Por ahora el dump debe ser
+        razonablemente compatible con SQLite.
+        """
+        connection_id = self._new_connection_id()
+        db_name = connection_id  # nombre lógico
+        db_path = os.path.join(UPLOAD_DIR, f"{db_name}.sqlite")
+
+        # Ejecutar todo el script. Si falla, borramos el archivo.
+        conn = sqlite3.connect(db_path)
+        try:
+            conn.executescript(sql_text)
+            conn.commit()
+        except Exception as e:
+            conn.close()
+            if os.path.exists(db_path):
+                os.remove(db_path)
+            raise RuntimeError(f"Error ejecutando dump SQL en SQLite: {e}")
+        finally:
+            conn.close()
+
+        self.connections[connection_id] = {
+            "label": label,
+            "engine": "sqlite",
+            "db_name": db_name,
+            "db_path": db_path,
+        }
+        return connection_id
+
+    # ---------- ejecución segura de SQL ----------
+
+    def execute_sql(self, connection_id: str, sql: str) -> Dict[str, Any]:
+        """
+        Ejecuta un SELECT sobre la BD asociada al connection_id.
+        Bloquea operaciones destructivas por seguridad.
+        """
+        info = self._get_info(connection_id)
+        db_path = info["db_path"]
+
+        forbidden = ["drop ", "delete ", "update ", "insert ", "alter ", "replace "]
+        sql_low = sql.lower()
+        if any(tok in sql_low for tok in forbidden):
+            return {
+                "ok": False,
+                "error": "Query bloqueada por seguridad (operación destructiva).",
+                "rows": None,
+                "columns": [],
+            }
+
+        try:
+            conn = sqlite3.connect(db_path)
+            cur = conn.cursor()
+            cur.execute(sql)
+            rows = cur.fetchall()
+            cols = [d[0] for d in cur.description] if cur.description else []
+            conn.close()
+            return {"ok": True, "error": None, "rows": [list(r) for r in rows], "columns": cols}
+        except Exception as e:
+            return {"ok": False, "error": str(e), "rows": None, "columns": []}
+
+    # ---------- introspección de esquema ----------
+
+    def get_schema(self, connection_id: str) -> Dict[str, Any]:
+        info = self._get_info(connection_id)
+        db_path = info["db_path"]
+
+        if not os.path.exists(db_path):
+            raise RuntimeError(f"SQLite no encontrado: {db_path}")
+
+        conn = sqlite3.connect(db_path)
+        cur = conn.cursor()
+
+        cur.execute("SELECT name FROM sqlite_master WHERE type='table';")
+        tables = [row[0] for row in cur.fetchall()]
+
+        tables_info: Dict[str, Dict[str, Any]] = {}
+        foreign_keys: List[Dict[str, Any]] = []
+
+        for t in tables:
+            cur.execute(f"PRAGMA table_info('{t}');")
+            rows = cur.fetchall()
+            cols = [r[1] for r in rows]
+            tables_info[t] = {"columns": cols}
+
+            cur.execute(f"PRAGMA foreign_key_list('{t}');")
+            fks = cur.fetchall()
+            for (id_, seq, ref_table, from_col, to_col, on_update, on_delete, match) in fks:
+                foreign_keys.append({
+                    "from_table": t,
+                    "from_column": from_col,
+                    "to_table": ref_table,
+                    "to_column": to_col,
+                })
+
+        conn.close()
+
+        return {
+            "tables": tables_info,
+            "foreign_keys": foreign_keys,
+        }
+
+    # ---------- preview de tabla ----------
+
+    def get_preview(self, connection_id: str, table: str, limit: int = 20) -> Dict[str, Any]:
+        info = self._get_info(connection_id)
+        db_path = info["db_path"]
+
+        conn = sqlite3.connect(db_path)
+        cur = conn.cursor()
+        try:
+            cur.execute(f'SELECT * FROM "{table}" LIMIT {int(limit)};')
+            rows = cur.fetchall()
+            cols = [d[0] for d in cur.description] if cur.description else []
+        finally:
+            conn.close()
+
+        return {
+            "columns": cols,
+            "rows": [list(r) for r in rows],
+        }
+
+
+# Instancia global de SQLManager
+sql_manager = SQLManager()
+
+# ======================================================
+# 2) Inicialización de FastAPI
 # ======================================================
 
 app = FastAPI(
-    title="NL2SQL T5-large Backend (MySQL/PostgreSQL)",
+    title="NL2SQL T5-large Backend (SQLite engine por ahora)",
     description=(
         "Intérprete NL→SQL (T5-large Spider) para usuarios no expertos. "
-        "El usuario sube sus dumps .sql (o ZIP con .sql) y se levantan "
-        "bases reales en MySQL/PostgreSQL; las consultas se ejecutan ahí."
+        "El usuario sube sus dumps .sql (o ZIP con .sql) y se crean "
+        "bases dinámicas (actualmente SQLite, futuro Postgres/MySQL)."
     ),
     version="2.0.0",
 )
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],   # en producción puedes acotar a tu dominio
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
 # ======================================================
-# 2) Modelo NL→SQL y traductor ES→EN
+# 3) Modelo NL→SQL y traductor ES→EN
 # ======================================================
 
 t5_tokenizer = None
@@ -124,7 +278,7 @@ def translate_es_to_en(text: str) -> str:
 
 
 # ======================================================
-# 3) Capa de reparación de SQL (usa el schema real)
+# 4) Capa de reparación de SQL (usa el schema real)
 # ======================================================
 
 def _normalize_name_for_match(name: str) -> str:
@@ -259,7 +413,7 @@ def try_repair_sql(sql: str, error: str, schema_meta: Dict[str, Any]) -> Optiona
 
 
 # ======================================================
-# 4) Construcción de prompt y NL→SQL + re-ranking
+# 5) Construcción de prompt y NL→SQL + re-ranking
 # ======================================================
 
 def build_prompt(question_en: str, db_id: str, schema_str: str) -> str:
@@ -274,7 +428,7 @@ def nl2sql_with_rerank(question: str, conn_id: str) -> Dict[str, Any]:
     if conn_id not in sql_manager.connections:
         raise HTTPException(status_code=404, detail=f"connection_id '{conn_id}' no registrado")
 
-    # Obtener esquema real desde MySQL/Postgres
+    # Obtener esquema real desde SQLite (futuro: Postgres/MySQL)
     meta = sql_manager.get_schema(conn_id)
     tables_info = meta["tables"]
 
@@ -331,10 +485,11 @@ def nl2sql_with_rerank(question: str, conn_id: str) -> Dict[str, Any]:
         exec_info = sql_manager.execute_sql(conn_id, raw_sql)
 
         # Intentar reparación solo si es error por tabla/columna
+        err_lower = (exec_info["error"] or "").lower()
         if (not exec_info["ok"]) and (
-            "no such table" in (exec_info["error"] or "").lower()
-            or "no such column" in (exec_info["error"] or "").lower()
-            or "does not exist" in (exec_info["error"] or "").lower()
+            "no such table" in err_lower
+            or "no such column" in err_lower
+            or "does not exist" in err_lower
         ):
             current_sql = raw_sql
             last_error = exec_info["error"] or ""
@@ -389,13 +544,13 @@ def nl2sql_with_rerank(question: str, conn_id: str) -> Dict[str, Any]:
 
 
 # ======================================================
-# 5) Schemas Pydantic
+# 6) Schemas Pydantic
 # ======================================================
 
 class UploadResponse(BaseModel):
     connection_id: str
     label: str
-    db_path: str  # ahora será un pseudo-path (engine://db_name)
+    db_path: str  # pseudo-path (engine://db_name o similar)
     note: Optional[str] = None
 
 
@@ -444,7 +599,7 @@ class SpeechInferResponse(BaseModel):
 
 
 # ======================================================
-# 6) Helpers para /upload (.sql y .zip)
+# 7) Helpers para /upload (.sql y .zip)
 # ======================================================
 
 def _combine_sql_files_from_zip(zip_bytes: bytes) -> str:
@@ -482,13 +637,13 @@ def _combine_sql_files_from_zip(zip_bytes: bytes) -> str:
 
 
 # ======================================================
-# 7) Endpoints FastAPI
+# 8) Endpoints FastAPI
 # ======================================================
 
 @app.on_event("startup")
 async def startup_event():
     load_nl2sql_model()
-    print("✅ Backend NL2SQL inicializado (MySQL/PostgreSQL).")
+    print("✅ Backend NL2SQL inicializado (engine SQLite por ahora).")
     print(f"MODEL_DIR={MODEL_DIR}, DEVICE={DEVICE}")
     print(f"Conexiones activas al inicio: {len(sql_manager.connections)}")
 
@@ -497,7 +652,7 @@ async def startup_event():
 async def upload_database(db_file: UploadFile = File(...)):
     """
     Subida de BD basada en dumps:
-      - .sql → dump MySQL/PostgreSQL (schema + data) → BD real
+      - .sql → dump (schema + data) → BD dinámica (SQLite por ahora)
       - .zip → debe contener uno o varios .sql (se concatenan)
     """
     filename = db_file.filename
@@ -553,7 +708,7 @@ async def upload_database(db_file: UploadFile = File(...)):
     engine = meta["engine"]
     db_name = meta["db_name"]
 
-    # db_path ahora es un pseudo-path para mantener compatibilidad
+    # db_path pseudo para mantener compatibilidad
     db_path = f"{engine}://{db_name}"
 
     return UploadResponse(
@@ -568,18 +723,12 @@ async def upload_database(db_file: UploadFile = File(...)):
 async def list_connections():
     return [
         ConnectionInfo(
-            connection_id=c["connection_id"],
-            label=c.get("label", ""),
-            engine=c.get("engine"),
-            db_name=c.get("db_name"),
+            connection_id=cid,
+            label=meta.get("label", ""),
+            engine=meta.get("engine"),
+            db_name=meta.get("db_name"),
         )
-        for c in [
-            {
-                "connection_id": cid,
-                **meta,
-            }
-            for cid, meta in sql_manager.connections.items()
-        ]
+        for cid, meta in sql_manager.connections.items()
     ]
 
 
@@ -681,13 +830,13 @@ async def health():
 @app.get("/")
 async def root():
     return {
-        "message": "NL2SQL T5-large backend is running with real MySQL/PostgreSQL engines.",
+        "message": "NL2SQL T5-large backend is running (engine SQLite, ready to upgrade to Postgres/MySQL).",
         "endpoints": [
             "POST /upload            (subir .sql o .zip con .sql → crear BD dinámica)",
             "GET  /connections       (listar BDs subidas)",
             "GET  /schema/{id}       (esquema resumido)",
             "GET  /preview/{id}/{t}  (preview de tabla)",
-            "POST /infer             (NL→SQL + ejecución en BD real)",
+            "POST /infer             (NL→SQL + ejecución en BD)",
             "POST /speech-infer      (NL por voz → SQL + ejecución)",
             "GET  /health            (estado del backend)",
             "GET  /docs              (OpenAPI UI)",