geMoldInsight/src/moldinsight/api/batch_router.py

"""
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.database import User, 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 moldinsight.services.storage_integration_rustfs import StorageIntegrationService
from moldinsight.services.task_dispatcher import dispatch_processing

logger = get_logger(__name__)

router = APIRouter()

file_handler = FileHandler()


@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 = StorageIntegrationService()

    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,
    }