import csv import re from datetime import datetime from pathlib import Path from typing import Optional import pandas as pd from fastapi import APIRouter, Body, File, Form, HTTPException, UploadFile 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, detect_csv_encoding from app.services.platform import make_client from app.utils.files import remove_file_safely from app.warehouse import make_cell_data_warehouse, make_warehouse router = APIRouter(tags=["database"]) INVALID_SHEET_NAME_CHARS = re.compile(r"[:\\/?*\[\]]") def _resolve_requested_tables( table_name: Optional[str], table_names: Optional[list[str]], ) -> list[str]: names = table_names if table_names is not None else ([table_name] if table_name else []) return [name.strip() for name in names if isinstance(name, str) and name.strip()] def _make_sheet_name(table_name: str, used_names: set[str]) -> str: base = INVALID_SHEET_NAME_CHARS.sub("_", table_name).strip("'").strip() or "Sheet" base = base[:31] sheet_name = base index = 2 while sheet_name in used_names: suffix = f"_{index}" sheet_name = f"{base[:31 - len(suffix)]}{suffix}" or f"Sheet_{index}" index += 1 used_names.add(sheet_name) return sheet_name DatabaseSource = str def _dataframe_from_table(db: DatabaseManager, table_name: str) -> pd.DataFrame: result = db.query_table(table_name, page=1, page_size=1000000) table_info = db.get_table_info(table_name) columns = [str(column["Field"]) for column in table_info["columns"]] return pd.DataFrame(result["data"], columns=columns) def _db(database_source: DatabaseSource = "main"): """Direct MySQL DatabaseManager, or a Metrix-backed warehouse with the same interface.""" config = state.current_config() if database_source == "cell_data": return make_cell_data_warehouse(config) if database_source != "main": raise HTTPException(status_code=400, detail="不支持的数据库来源") return make_warehouse(config) @router.post("/api/database/test") def test_database(database_source: DatabaseSource = Body("main", embed=True)): db = _db(database_source) success, message = db.test_connection() return {"success": success, "message": message} @router.get("/api/database/info") def get_database_info(database_source: DatabaseSource = "main"): db = _db(database_source) try: return {"success": True, **db.get_server_info()} except Exception as exc: return {"success": False, "error": str(exc)} @router.get("/api/database/tables") @router.post("/api/database/tables") def get_tables(database_source: DatabaseSource = Body("main", embed=True)): db = _db(database_source) try: return {"tables": db.get_tables()} except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc @router.post("/api/database/table/info") def get_table_info( table_name: str = Body(..., embed=True), database_source: DatabaseSource = Body("main"), ): db = _db(database_source) try: return db.get_table_info(table_name) except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc @router.post("/api/database/table/data") def query_table_data( table_name: str = Body(..., embed=True), database_source: DatabaseSource = Body("main"), page: int = Body(1), page_size: int = Body(50), order_by: Optional[str] = Body(None), order_dir: str = Body("ASC"), ): db = _db(database_source) 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 @router.post("/api/database/table/query") def query_table_with_filter( table_name: str = Body(..., embed=True), database_source: DatabaseSource = Body("main"), 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 = _db(database_source) 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 @router.post("/api/database/table/truncate") def truncate_table( table_name: str = Body(..., embed=True), database_source: DatabaseSource = Body("main"), ): db = _db(database_source) 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 @router.post("/api/database/table/drop") def drop_table( table_name: str = Body(..., embed=True), database_source: DatabaseSource = Body("main"), ): db = _db(database_source) 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 @router.post("/api/database/table/drop-all") def drop_all_tables(database_source: DatabaseSource = Body("main", embed=True)): db = _db(database_source) 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 @router.post("/api/database/table/row/update") def update_table_row( table_name: str = Body(..., embed=True), database_source: DatabaseSource = Body("main"), identifier: dict = Body(...), values: dict = Body(...), ): if not identifier: raise HTTPException(status_code=400, detail="缺少行定位信息") db = _db(database_source) try: affected = db.update_row(table_name, identifier, values) except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc if affected == 0: raise HTTPException(status_code=404, detail="未找到匹配的数据行,可能已被修改或删除") return {"success": True, "message": "已更新该行", "affected_rows": affected} @router.post("/api/database/table/row/delete") def delete_table_row( table_name: str = Body(..., embed=True), database_source: DatabaseSource = Body("main"), identifier: dict = Body(...), ): if not identifier: raise HTTPException(status_code=400, detail="缺少行定位信息") db = _db(database_source) try: affected = db.delete_row(table_name, identifier) except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc if affected == 0: raise HTTPException(status_code=404, detail="未找到匹配的数据行,可能已被删除") return {"success": True, "message": "已删除该行", "affected_rows": affected} def _table_columns(db, table_name: str) -> list[str]: info = db.get_table_info(table_name) return [str(column["Field"]) for column in info.get("columns", [])] @router.post("/api/database/table/template") def download_table_template( table_name: str = Body(..., embed=True), database_source: DatabaseSource = Body("main"), ): db = _db(database_source) try: columns = _table_columns(db, table_name) except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc if not columns: raise HTTPException(status_code=400, detail="无法获取表字段,无法生成模板") CACHE_DIR.mkdir(parents=True, exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filepath = CACHE_DIR / f"{table_name}_template_{timestamp}.csv" try: with filepath.open("w", encoding="utf-8-sig", newline="") as handle: csv.writer(handle).writerow(columns) except Exception: remove_file_safely(filepath) raise return FileResponse( path=str(filepath), filename=f"{table_name}_模板.csv", media_type="text/csv", background=BackgroundTask(remove_file_safely, filepath), ) def _read_csv_header(path: Path, encoding: str) -> list[str]: with path.open("r", encoding=encoding, newline="") as handle: for row in csv.reader(handle): return [str(cell).strip() for cell in row] return [] # Sync def so FastAPI runs it in a threadpool: file read + DB insert are blocking. @router.post("/api/database/table/import") def import_table_csv( file: UploadFile = File(...), table_name: str = Form(...), database_source: str = Form("main"), ): if not file.filename or not file.filename.lower().endswith(".csv"): raise HTTPException(status_code=400, detail="仅支持 CSV 格式文件") db = _db(database_source) try: columns = _table_columns(db, table_name) except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc if not columns: raise HTTPException(status_code=400, detail="无法获取表字段") CACHE_DIR.mkdir(parents=True, exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") tmp_path = CACHE_DIR / f"import_{timestamp}.csv" tmp_path.write_bytes(file.file.read()) try: encoding = detect_csv_encoding(str(tmp_path)) header = _read_csv_header(tmp_path, encoding) if not header: raise HTTPException(status_code=400, detail="CSV 文件为空或缺少表头") missing = [name for name in columns if name not in header] extra = [name for name in header if name not in columns] if missing or extra: parts = [] if missing: parts.append("缺少字段: " + ", ".join(missing)) if extra: parts.append("多余字段: " + ", ".join(extra)) raise HTTPException(status_code=400, detail="CSV 字段与模板不一致,导入失败。" + ";".join(parts)) if database_source == "cell_data" and table_name.lower() == "sector": stats = db.upsert_csv(str(tmp_path), table_name, "CGI", encoding=encoding) return { "success": True, "message": ( f"导入成功,共 {stats['imported_rows']} 行:" f"新增 {stats['inserted_rows']} 行,更新 {stats['updated_rows']} 行," f"清理历史重复 {stats['removed_duplicate_rows']} 行" ), **stats, } imported = db.import_csv(str(tmp_path), table_name, encoding=encoding) except HTTPException: raise except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc finally: remove_file_safely(tmp_path) return {"success": True, "message": f"导入成功,共 {imported} 行", "imported_rows": imported} @router.post("/api/database/execute") def execute_sql( sql: str = Body(..., embed=True), database_source: DatabaseSource = Body("main"), ): db = _db(database_source) 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 # Sync def so FastAPI runs it in a threadpool: exporting large tables (and the Metrix # export-job polling) is blocking and would otherwise freeze the single-worker event loop. @router.post("/api/download") def download_table( table_name: Optional[str] = Body(None, embed=True), table_names: Optional[list[str]] = Body(None, embed=True), database_source: DatabaseSource = Body("main"), file_format: str = Body("csv", alias="format"), ): if file_format not in {"csv", "xlsx"}: raise HTTPException(status_code=400, detail="不支持的导出格式") requested_tables = _resolve_requested_tables(table_name, table_names) if not requested_tables: raise HTTPException(status_code=400, detail="请选择要导出的数据表") if file_format == "csv" and len(requested_tables) != 1: raise HTTPException(status_code=400, detail="CSV 每次只能导出一张表") config = state.current_config() if database_source == "main" and config.warehouse_type == "metrix": return _download_via_metrix(config, requested_tables, file_format) db = _db(database_source) try: available_tables = set(db.get_tables()) missing_tables = [name for name in requested_tables if name not in available_tables] if missing_tables: raise HTTPException(status_code=400, detail=f"数据表不存在: {', '.join(missing_tables)}") table_frames = { name: _dataframe_from_table(db, name) for name in requested_tables } except HTTPException: raise except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename_prefix = requested_tables[0] if len(requested_tables) == 1 else "tables" filename = f"{filename_prefix}_{timestamp}.{file_format}" filepath = CACHE_DIR / filename CACHE_DIR.mkdir(parents=True, exist_ok=True) try: if file_format == "csv": table_frames[requested_tables[0]].to_csv(filepath, index=False, encoding="utf-8-sig") media_type = "text/csv" else: used_sheet_names: set[str] = set() with pd.ExcelWriter(filepath) as writer: for name, df in table_frames.items(): df.to_excel(writer, sheet_name=_make_sheet_name(name, used_sheet_names), index=False) media_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" except Exception: remove_file_safely(filepath) raise return FileResponse( path=str(filepath), filename=filename, media_type=media_type, background=BackgroundTask(remove_file_safely, filepath), ) def _download_via_metrix(config, requested_tables: list[str], file_format: str) -> FileResponse: """Metrix 仓库模式:用平台导出任务生成文件后流式返回(避免分页上限丢行)。""" metrix = config.metrix.normalized() client = make_client(metrix) try: job_id = client.submit_export(metrix.database_conn_id, requested_tables, file_format, metrix.target_database) job = client.wait_job(job_id) if job.get("status") != "success": raise HTTPException(status_code=500, detail=f"导出失败: {job.get('error_code') or job.get('status')}") timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename_prefix = requested_tables[0] if len(requested_tables) == 1 else "tables" filename = f"{filename_prefix}_{timestamp}.{file_format}" filepath = CACHE_DIR / filename CACHE_DIR.mkdir(parents=True, exist_ok=True) client.download_job_file(job_id, filepath) except HTTPException: raise except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc media_type = "text/csv" if file_format == "csv" else "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" return FileResponse( path=str(filepath), filename=filename, media_type=media_type, background=BackgroundTask(remove_file_safely, filepath), )