storage/db_routers.py

import os
import io

import sqlite3
import zipfile

from pathlib import Path

from fastapi import APIRouter, UploadFile, File, Query, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse

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

    # Create in-memory ZIP
    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 StreamingResponse(
        zip_buffer,
        media_type="application/zip",
        headers={
            "Content-Disposition": 'attachment; 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})