235 lines
8.0 KiB
Python
235 lines
8.0 KiB
Python
from datetime import datetime
|
||
from pathlib import Path
|
||
from threading import Thread
|
||
from typing import Any
|
||
|
||
from fastapi import APIRouter, Body, HTTPException
|
||
|
||
from app import state
|
||
from app.api.routers.task_runtime import (
|
||
apply_history_retention_safely,
|
||
log_license_check,
|
||
set_task_stage,
|
||
)
|
||
from app.config import AppConfig
|
||
from app.processor import DataProcessor, ProcessLogger
|
||
from app.config import CELLDATA_SCRIPT
|
||
from app.services.cell_data import copy_celldata_tables_to_capacity, execute_celldata_script, refresh_cell_data
|
||
from app.services.license import LicenseError, check_processing_allowed
|
||
|
||
|
||
router = APIRouter(tags=["tasks"])
|
||
|
||
|
||
@router.get("/api/task/status")
|
||
async def get_global_task_status():
|
||
if state.global_task_lock["locked"]:
|
||
task_id = state.global_task_lock["task_id"]
|
||
if task_id and _task_finished(task_id):
|
||
state.reset_task_lock()
|
||
return {"has_active": False}
|
||
|
||
return {
|
||
"has_active": True,
|
||
"task_id": task_id,
|
||
"stage": state.global_task_lock["stage"],
|
||
"started_at": state.global_task_lock["started_at"],
|
||
"logs": [],
|
||
}
|
||
|
||
active_tasks = {
|
||
task_id: task
|
||
for task_id, task in state.processing_tasks.items()
|
||
if task.get("status") == "processing"
|
||
}
|
||
if active_tasks:
|
||
task_id = next(iter(active_tasks))
|
||
return {
|
||
"has_active": True,
|
||
"task_id": task_id,
|
||
"stage": active_tasks[task_id].get("stage", "processing"),
|
||
"logs": active_tasks[task_id].get("logs", []),
|
||
}
|
||
|
||
return {"has_active": False}
|
||
|
||
|
||
@router.post("/api/task/lock")
|
||
async def lock_task(task_id: str = Body(..., embed=True)):
|
||
if state.global_task_lock["locked"]:
|
||
raise HTTPException(status_code=409, detail="已有任务在运行")
|
||
|
||
state.global_task_lock.update(
|
||
{
|
||
"locked": True,
|
||
"task_id": task_id,
|
||
"stage": "uploading",
|
||
"started_at": datetime.now().isoformat(),
|
||
}
|
||
)
|
||
return {"success": True, "message": "任务已锁定"}
|
||
|
||
|
||
@router.post("/api/task/unlock")
|
||
async def unlock_task(task_id: str | None = Body(None, embed=True)):
|
||
if task_id and state.global_task_lock["task_id"] != task_id:
|
||
raise HTTPException(status_code=403, detail="无权解锁此任务")
|
||
|
||
state.reset_task_lock()
|
||
return {"success": True, "message": "任务已解锁"}
|
||
|
||
|
||
@router.get("/api/process/active")
|
||
async def get_active_task():
|
||
return await get_global_task_status()
|
||
|
||
|
||
@router.post("/api/process/start")
|
||
async def start_processing(task_id: str = Body(..., embed=True)):
|
||
record = state.history_manager.get(task_id)
|
||
if not record:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
if record.status == "processing":
|
||
raise HTTPException(status_code=400, detail="任务正在处理中")
|
||
|
||
work_dir = Path(record.work_dir)
|
||
if not work_dir.exists():
|
||
raise HTTPException(status_code=400, detail="工作目录不存在")
|
||
|
||
logs: list[str] = []
|
||
current_stage = "processing"
|
||
|
||
def log_callback(message: str) -> None:
|
||
logs.append(message)
|
||
set_task_stage(task_id, current_stage, logs)
|
||
|
||
def stage_callback(stage: str) -> None:
|
||
nonlocal current_stage
|
||
current_stage = stage
|
||
set_task_stage(task_id, current_stage, logs)
|
||
|
||
logger = ProcessLogger(
|
||
log_file=work_dir / "log.txt",
|
||
callback=log_callback,
|
||
stage_callback=stage_callback,
|
||
)
|
||
app_config = state.current_config()
|
||
state.history_manager.update(task_id, status="processing")
|
||
state.processing_tasks[task_id] = {"logs": [], "status": "processing", "stage": current_stage}
|
||
state.global_task_lock.update(
|
||
{
|
||
"locked": True,
|
||
"task_id": task_id,
|
||
"stage": "processing",
|
||
"started_at": datetime.now().isoformat(),
|
||
}
|
||
)
|
||
|
||
thread = Thread(target=_run_processing, args=(task_id, work_dir, logger, app_config), daemon=True)
|
||
thread.start()
|
||
|
||
return {"success": True, "message": "处理任务已启动", "task_id": task_id}
|
||
|
||
|
||
@router.post("/api/process/status")
|
||
async def get_processing_status(task_id: str = Body(..., embed=True)):
|
||
if task_id in state.processing_tasks:
|
||
task_info = state.processing_tasks[task_id]
|
||
logs = task_info.get("logs") or state.history_manager.get_logs(task_id)
|
||
return {
|
||
"task_id": task_id,
|
||
"status": task_info["status"],
|
||
"stage": task_info.get("stage"),
|
||
"logs": logs,
|
||
"error": task_info.get("error"),
|
||
"error_detail": task_info.get("error_detail"),
|
||
}
|
||
|
||
record = state.history_manager.get(task_id)
|
||
if not record:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
|
||
return {
|
||
"task_id": task_id,
|
||
"status": record.status,
|
||
"stage": record.status,
|
||
"logs": state.history_manager.get_logs(task_id),
|
||
"elapsed_time": record.elapsed_time,
|
||
"error": record.error,
|
||
}
|
||
|
||
|
||
def _task_finished(task_id: str) -> bool:
|
||
record = state.history_manager.get(task_id)
|
||
if record and record.status in {"completed", "failed"}:
|
||
return True
|
||
|
||
task_info: dict[str, Any] | None = state.processing_tasks.get(task_id)
|
||
return bool(task_info and task_info.get("status") in {"completed", "failed"})
|
||
|
||
|
||
def _run_processing(task_id: str, work_dir: Path, logger: ProcessLogger, app_config: AppConfig) -> None:
|
||
try:
|
||
_try_refresh_cell_data(app_config, work_dir, logger)
|
||
if app_config.warehouse_type == "metrix":
|
||
import time
|
||
|
||
from app.services.pipeline import RESULT_TABLES, run_import_and_report
|
||
|
||
started = time.time()
|
||
run_import_and_report(work_dir, app_config, logger)
|
||
status, error, elapsed, result_tables = "completed", None, round(time.time() - started, 2), RESULT_TABLES
|
||
else:
|
||
logger.set_stage("license")
|
||
log_license_check(logger, check_processing_allowed(work_dir))
|
||
processor = DataProcessor(app_config, work_dir, logger)
|
||
result = processor.process()
|
||
status = "completed" if result.get("success") else "failed"
|
||
error = result.get("error")
|
||
elapsed = result.get("elapsed_time", 0)
|
||
result_tables = ["4G_结果表", "5G_结果表"]
|
||
state.history_manager.update(
|
||
task_id,
|
||
status=status,
|
||
elapsed_time=elapsed,
|
||
error=error,
|
||
result_tables=result_tables,
|
||
)
|
||
state.processing_tasks[task_id] = {
|
||
"logs": state.history_manager.get_logs(task_id),
|
||
"status": status,
|
||
"stage": status,
|
||
"error": error,
|
||
}
|
||
except Exception as exc:
|
||
error_detail = exc.to_detail() if isinstance(exc, LicenseError) else None
|
||
logger.error(str(exc))
|
||
state.history_manager.update(task_id, status="failed", error=str(exc))
|
||
state.processing_tasks[task_id] = {
|
||
"logs": state.history_manager.get_logs(task_id),
|
||
"status": "failed",
|
||
"stage": "failed",
|
||
"error": str(exc),
|
||
"error_detail": error_detail,
|
||
}
|
||
finally:
|
||
apply_history_retention_safely()
|
||
state.reset_task_lock()
|
||
|
||
|
||
def _try_refresh_cell_data(app_config: AppConfig, work_dir: Path, logger: ProcessLogger) -> None:
|
||
if not app_config.cell_data.remote_data.enabled:
|
||
return
|
||
logger.set_stage("cell_data")
|
||
logger.info("── CellData 更新 ──")
|
||
try:
|
||
result = refresh_cell_data(app_config, work_dir, logger)
|
||
logger.success(
|
||
f"CellData 更新完成：{result.imported_rows} 行"
|
||
f"（解析 {result.parsed_rows}，跳过 {result.skipped_rows}）"
|
||
)
|
||
execute_celldata_script(CELLDATA_SCRIPT, app_config, logger)
|
||
copy_celldata_tables_to_capacity(app_config, logger)
|
||
except Exception as exc:
|
||
logger.warning(f"CellData 更新失败，继续容量处理: {exc}")
|