CapacityReport/app/api/routers/database.py

from contextlib import suppress
from datetime import datetime
from pathlib import Path
from typing import Optional

import pandas as pd
from fastapi import APIRouter, Body, HTTPException
from fastapi.responses import FileResponse
from starlette.background import BackgroundTask

from app import state
from app.config import CACHE_DIR
from app.database import DatabaseManager


router = APIRouter(tags=["database"])


def _remove_file(path: Path) -> None:
    with suppress(OSError):
        path.unlink()


@router.post("/api/database/test")
async def test_database():
    db = DatabaseManager(state.config)
    try:
        success, message = db.test_connection()
        return {"success": success, "message": message}
    finally:
        db.dispose()


@router.get("/api/database/info")
async def get_database_info():
    db = DatabaseManager(state.config)
    try:
        return {"success": True, **db.get_server_info()}
    except Exception as exc:
        return {"success": False, "error": str(exc)}
    finally:
        db.dispose()


@router.get("/api/database/tables")
@router.post("/api/database/tables")
async def get_tables():
    db = DatabaseManager(state.config)
    try:
        return {"tables": db.get_tables()}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        db.dispose()


@router.post("/api/database/table/info")
async def get_table_info(table_name: str = Body(..., embed=True)):
    db = DatabaseManager(state.config)
    try:
        return db.get_table_info(table_name)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        db.dispose()


@router.post("/api/database/table/data")
async def query_table_data(
    table_name: str = Body(..., embed=True),
    page: int = Body(1),
    page_size: int = Body(50),
    order_by: Optional[str] = Body(None),
    order_dir: str = Body("ASC"),
):
    db = DatabaseManager(state.config)
    try:
        return db.query_table(table_name, page, page_size, order_by=order_by, order_dir=order_dir)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        db.dispose()


@router.post("/api/database/table/query")
async def query_table_with_filter(
    table_name: str = Body(..., embed=True),
    page: int = Body(1),
    page_size: int = Body(50),
    filters: Optional[dict[str, str]] = Body(None),
    order_by: Optional[str] = Body(None),
    order_dir: str = Body("ASC"),
):
    db = DatabaseManager(state.config)
    try:
        return db.query_table(
            table_name,
            page,
            page_size,
            filters=filters or {},
            order_by=order_by,
            order_dir=order_dir,
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        db.dispose()


@router.post("/api/database/table/truncate")
async def truncate_table(table_name: str = Body(..., embed=True)):
    db = DatabaseManager(state.config)
    try:
        db.truncate_table(table_name)
        return {"success": True, "message": f"表 {table_name} 已清空"}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        db.dispose()


@router.post("/api/database/table/drop")
async def drop_table(table_name: str = Body(..., embed=True)):
    db = DatabaseManager(state.config)
    try:
        db.drop_table(table_name)
        return {"success": True, "message": f"表 {table_name} 已删除"}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        db.dispose()


@router.post("/api/database/table/drop-all")
async def drop_all_tables():
    db = DatabaseManager(state.config)
    try:
        result = db.drop_all_tables()
        return {
            "success": True,
            "message": f"已删除 {result['dropped_count']} 个表",
            "dropped_count": result["dropped_count"],
            "tables": result["tables"],
        }
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        db.dispose()


@router.post("/api/database/execute")
async def execute_sql(sql: str = Body(..., embed=True)):
    db = DatabaseManager(state.config)
    try:
        success, result = db.execute_sql(sql)
        if success:
            return {"success": True, "result": result}
        raise HTTPException(status_code=400, detail=result)
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        db.dispose()


@router.post("/api/download")
async def download_table(
    table_name: str = Body(..., embed=True),
    file_format: str = Body("csv", alias="format"),
):
    if file_format not in {"csv", "xlsx"}:
        raise HTTPException(status_code=400, detail="不支持的导出格式")

    db = DatabaseManager(state.config)
    try:
        result = db.query_table(table_name, page=1, page_size=1000000)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        db.dispose()

    df = pd.DataFrame(result["data"])
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{table_name}_{timestamp}.{file_format}"
    filepath = CACHE_DIR / filename
    CACHE_DIR.mkdir(parents=True, exist_ok=True)

    try:
        if file_format == "csv":
            df.to_csv(filepath, index=False, encoding="utf-8-sig")
            media_type = "text/csv"
        else:
            df.to_excel(filepath, index=False)
            media_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    except Exception:
        _remove_file(filepath)
        raise

    return FileResponse(
        path=str(filepath),
        filename=filename,
        media_type=media_type,
        background=BackgroundTask(_remove_file, filepath),
    )