File size: 4,278 Bytes
46cf14b 82a94f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
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}) |