""" moldinsight/api/batch_router.py — 批量分析端点 - POST /api/batch-upload 批量上传多文件，返回 batch_id + 各 task_id - GET /api/batch/{batch_id} 聚合查询批量任务进度 批次 2（D7）：批量元数据以 PG 为单一事实源——ProcessingTask.batch_id 列聚合查询，替代此前的 Redis key + 进程内存降级存储。 """ import uuid from datetime import datetime from typing import List, Dict, Any from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Depends from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload from shared.database.database import get_db_session from shared.services.auth_service import get_current_active_user from shared.models.identity import User from moldinsight.models import ProcessingTask, STPFile from shared.models.schemas import ProcessingStatus, create_task_info from shared.services.redis_task_manager import redis_task_manager from shared.utils.file_handler import FileHandler from shared.utils.logger import get_logger from shared.config.settings import settings from moldinsight.services.task_storage_service import TaskStorageService from moldinsight.services.task_dispatcher import dispatch_processing logger = get_logger(__name__) router = APIRouter() # D14：上传限制接 settings（MAX_FILE_SIZE 此前为死配置，文件处理器硬编码 50MB） file_handler = FileHandler(upload_dir=settings.UPLOAD_DIR, max_file_size=settings.MAX_FILE_SIZE) @router.post("/batch-upload") async def batch_upload( files: List[UploadFile] = File(...), material: str = Form("ABS"), draft_angle: float = Form(2.0), shrinkage_rate: float = Form(0.5), parting_precision: float = Form(0.1), cavity_match: int = Form(95), db_session: AsyncSession = Depends(get_db_session), current_user: User = Depends(get_current_active_user), ): """批量上传多个 STP 文件，每个文件创建独立分析任务，用 batch_id 聚合。""" if not files: raise HTTPException(400, "请至少上传一个文件") if len(files) > 20: raise HTTPException(400, "单次批量上传最多 20 个文件") process_params = { "material": material, "draft_angle": float(draft_angle), "shrinkage_rate": float(shrinkage_rate), "parting_precision": float(parting_precision), "cavity_match": int(cavity_match), } batch_id = str(uuid.uuid4()) tasks: List[Dict[str, Any]] = [] storage_service = TaskStorageService() for file in files: # 文件类型检查 if not file.filename.lower().endswith(('.stp', '.step')): tasks.append({ "filename": file.filename, "task_id": None, "status": "rejected", "error": "不支持的文件类型", }) continue task_id = str(uuid.uuid4()) try: file_path, file_size, file_meta = await file_handler.save_uploaded_file(file) stp_file = await storage_service.save_stp_file( session=db_session, file_path=file_path, original_filename=file_meta["safe_original_name"], user_id=current_user.id, ) await storage_service.create_processing_task( db_session, task_id, stp_file.id, parameters=process_params, batch_id=batch_id, ) # D9：STPFile + ProcessingTask 原子提交，分派前置事务收口 await db_session.commit() task_info = create_task_info( task_id=task_id, status=ProcessingStatus.PROCESSING, filename=file.filename, file_path=str(file_path), file_size=file_size, upload_time=str(datetime.now()), ) task_info["material"] = material task_info["parameters"] = process_params task_info["batch_id"] = batch_id await redis_task_manager.set_task(task_id, task_info) # 调度处理 dispatch_processing(task_id, stp_file.id, process_params) tasks.append({ "filename": file.filename, "task_id": task_id, "status": "processing", "stp_file_id": stp_file.id, }) logger.info( f"[BATCH] batch_id={batch_id} task_id={task_id} " f"file={file.filename} user={current_user.username}" ) except Exception as exc: logger.warning(f"[BATCH] 文件 {file.filename} 上传失败: {exc}") tasks.append({ "filename": file.filename, "task_id": task_id, "status": "error", "error": str(exc), }) return { "batch_id": batch_id, "total": len(tasks), "accepted": sum(1 for t in tasks if t.get("status") != "rejected"), "tasks": tasks, } @router.get("/batch/{batch_id}") async def get_batch_status( batch_id: str, db_session: AsyncSession = Depends(get_db_session), current_user: User = Depends(get_current_active_user), ): """聚合查询批量任务进度（D7：以 PG 为单一事实源，按 batch_id 聚合；Redis 仅热缓存）""" rows = (await db_session.execute( select(ProcessingTask, STPFile) .join(STPFile, ProcessingTask.stp_file_id == STPFile.id) .where(ProcessingTask.batch_id == batch_id) .options(joinedload(STPFile.html_file)) .order_by(ProcessingTask.id) )).unique().all() if not rows: raise HTTPException(404, "批量任务不存在或已过期") # 归属校验：同批任务属于同一上传用户，任一不匹配即拒绝（无主不等于公共） if any(getattr(stp, "user_id", None) != current_user.id for _, stp in rows): raise HTTPException(403, "无权访问该批量任务") task_statuses = [] completed = 0 failed = 0 processing = 0 earliest_created = None for task, stp in rows: status = task.status or "unknown" if earliest_created is None or ( task.created_time and task.created_time < earliest_created ): earliest_created = task.created_time if status == ProcessingStatus.COMPLETED: completed += 1 elif status == ProcessingStatus.FAILED: failed += 1 else: processing += 1 html_file = "" if stp.html_file and stp.html_file.filename: html_file = f"/html/{stp.html_file.filename}" task_statuses.append({ "task_id": task.task_id, "status": status, "progress": task.progress or 0, "current_step": task.current_step, "filename": stp.original_filename or "", "error": task.error_message or "", "html_file": html_file, }) total = len(rows) return { "batch_id": batch_id, "created_at": earliest_created.isoformat() if earliest_created else None, "total": total, "completed": completed, "failed": failed, "processing": processing, "progress_percent": round((completed + failed) / max(total, 1) * 100, 1), "tasks": task_statuses, }