Update storage/db_routers.py

storage/db_routers.py

@@ -1,95 +1,134 @@
-import os
-
-import sqlite3
-
-from pathlib import Path
-
-from fastapi import APIRouter, UploadFile, File, Query, HTTPException
-from fastapi.responses import JSONResponse
-
-from storage.common import validate_token
-
-router = APIRouter(prefix="/db", tags=["Database Manager"])
-HF_TOKEN = os.getenv("HF_TOKEN")
-DB_PATH = Path("/data/db")
-
-@router.post("/upload_db_file", tags=["Database Manager"])
-async def upload_db_file(
-    file: UploadFile = File(...),
-    token: str = Query(..., description="Token required for authorization")
-):
-    """
-    Upload a file to the /data/db folder.
-
-    Steps:
-    - Validate the token.
-    - Ensure /data/db exists (create it if missing).
-    - Save the uploaded file inside /data/db using its original filename.
-    - Return a JSON response confirming the upload.
-    """
-    validate_token(token)
-
-    # Crear carpeta /data/db si no existe
-    DB_PATH.mkdir(parents=True, exist_ok=True)
-
-    final_path = DB_PATH / file.filename
-
-    try:
-        # Leer contenido en memoria y guardar
-        file_bytes = await file.read()
-        with open(final_path, "wb") as f:
-            f.write(file_bytes)
-    except Exception as exc:
-        raise HTTPException(status_code=500, detail=f"Failed to save file: {exc}")
-
-    return JSONResponse(
-        status_code=200,
-        content={
-            "status": "ok",
-            "saved_to": str(final_path)
-        }
-    )
-
-@router.post("/execute_query", tags=["Database Manager"])
-async def execute_query(
-    db_filename: str = Query(..., description="SQLite .db file name"),
-    query: str = Query(..., description="SQL query to execute"),
-    token: str = Query(..., description="Token required for authorization")
-):
-    """
-    Execute a SQL query against a specified SQLite database file in /data/db.
-
-    Steps:
-    - Validate the token.
-    - Ensure the requested .db file exists inside /data/db.
-    - Connect to the database using sqlite3.
-    - Execute the provided query.
-    - Return the results as a list of dictionaries (column: value).
-    - Capture and return errors if query execution fails.
-    """
-    validate_token(token)
-
-    db_file_path = DB_PATH / db_filename
-
-    if not db_file_path.exists() or not db_file_path.is_file():
-        raise HTTPException(status_code=404, detail=f"Database file {db_filename} not found")
-
-    try:
-        conn = sqlite3.connect(db_file_path)
-        conn.row_factory = sqlite3.Row  # para devolver columnas por nombre
-        cur = conn.cursor()
-
-        cur.execute(query)
-        rows = cur.fetchall()
-
-        # Convert rows to list of dicts
-        results = [dict(row) for row in rows]
-
-        conn.commit()
-        conn.close()
-    except sqlite3.Error as e:
-        raise HTTPException(status_code=400, detail=f"SQL error: {e}")
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Unexpected error: {e}")
-
+import os
+
+import sqlite3
+import zipfile
+
+from pathlib import Path
+
+from fastapi import APIRouter, UploadFile, File, Query, HTTPException
+from fastapi.responses import JSONResponse
+
+from storage.common import validate_token
+
+router = APIRouter(prefix="/db", tags=["Database Manager"])
+HF_TOKEN = os.getenv("HF_TOKEN")
+DB_PATH = Path("/data/db")
+
+@router.get("/download_all_db_files", tags=["Database Manager"])
+async def download_all_db_files(
+    token: str = Query(..., description="Token required for authorization")
+):
+    """
+    Download all .db files from /data/db as a ZIP archive.
+
+    Steps:
+    - Validate the token.
+    - Verify that /data/db exists and contains .db files.
+    - Create an in-memory ZIP containing all .db files.
+    - Return the ZIP as a downloadable file.
+    """
+    validate_token(token)
+
+    if not DB_PATH.exists():
+        raise HTTPException(status_code=404, detail="DB folder does not exist")
+
+    db_files = list(DB_PATH.glob("*.db"))
+    if not db_files:
+        raise HTTPException(status_code=404, detail="No .db files found")
+
+    # Crear ZIP en memoria (sin escribir en disco)
+    zip_buffer = io.BytesIO()
+
+    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for file_path in db_files:
+            zipf.write(file_path, arcname=file_path.name)
+
+    zip_buffer.seek(0)
+
+    return FileResponse(
+        zip_buffer,
+        media_type="application/zip",
+        filename="db_files.zip"
+    )
+
+
+@router.post("/upload_db_file", tags=["Database Manager"])
+async def upload_db_file(
+    file: UploadFile = File(...),
+    token: str = Query(..., description="Token required for authorization")
+):
+    """
+    Upload a file to the /data/db folder.
+
+    Steps:
+    - Validate the token.
+    - Ensure /data/db exists (create it if missing).
+    - Save the uploaded file inside /data/db using its original filename.
+    - Return a JSON response confirming the upload.
+    """
+    validate_token(token)
+
+    # Crear carpeta /data/db si no existe
+    DB_PATH.mkdir(parents=True, exist_ok=True)
+
+    final_path = DB_PATH / file.filename
+
+    try:
+        # Leer contenido en memoria y guardar
+        file_bytes = await file.read()
+        with open(final_path, "wb") as f:
+            f.write(file_bytes)
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"Failed to save file: {exc}")
+
+    return JSONResponse(
+        status_code=200,
+        content={
+            "status": "ok",
+            "saved_to": str(final_path)
+        }
+    )
+
+@router.post("/execute_query", tags=["Database Manager"])
+async def execute_query(
+    db_filename: str = Query(..., description="SQLite .db file name"),
+    query: str = Query(..., description="SQL query to execute"),
+    token: str = Query(..., description="Token required for authorization")
+):
+    """
+    Execute a SQL query against a specified SQLite database file in /data/db.
+
+    Steps:
+    - Validate the token.
+    - Ensure the requested .db file exists inside /data/db.
+    - Connect to the database using sqlite3.
+    - Execute the provided query.
+    - Return the results as a list of dictionaries (column: value).
+    - Capture and return errors if query execution fails.
+    """
+    validate_token(token)
+
+    db_file_path = DB_PATH / db_filename
+
+    if not db_file_path.exists() or not db_file_path.is_file():
+        raise HTTPException(status_code=404, detail=f"Database file {db_filename} not found")
+
+    try:
+        conn = sqlite3.connect(db_file_path)
+        conn.row_factory = sqlite3.Row  # para devolver columnas por nombre
+        cur = conn.cursor()
+
+        cur.execute(query)
+        rows = cur.fetchall()
+
+        # Convert rows to list of dicts
+        results = [dict(row) for row in rows]
+
+        conn.commit()
+        conn.close()
+    except sqlite3.Error as e:
+        raise HTTPException(status_code=400, detail=f"SQL error: {e}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error: {e}")
+
     return JSONResponse(content={"results": results})