geMoldInsight/src/api/v1/upload_router.py

# api/v1/upload_router.py
from fastapi import APIRouter, UploadFile, File, HTTPException, BackgroundTasks, Depends
from typing import Optional
import uuid
from datetime import datetime
from pathlib import Path

from models.schemas import ProcessingStatus, create_task_info
from utils.file_handler import FileHandler
from services.storage_integration_rustfs import StorageIntegrationService
from services.redis_task_manager import redis_task_manager
from database.database import get_db_session
from utils.logger import get_logger
from sqlalchemy.ext.asyncio import AsyncSession
from services.auth_service import get_current_active_user
from services.processing_service import processing_service
from models.database import User

logger = get_logger(__name__)

router = APIRouter()

file_handler = FileHandler()


@router.post("/upload")
async def upload_stp(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    material: Optional[str] = "ABS",
    db_session: AsyncSession = Depends(get_db_session),
    current_user: User = Depends(get_current_active_user)
):
    """上传STP文件并存储到数据库"""
    logger.info(
        f"[UPLOAD] 用户={current_user.username}(id={current_user.id}) "
        f"文件={file.filename} 材料={material} "
        f"大小={file.size if hasattr(file, 'size') else 'unknown'}"
    )

    if not file.filename.lower().endswith(('.stp', '.step')):
        logger.warning(f"[UPLOAD] 拒绝: 不支持的文件类型 - {file.filename}")
        raise HTTPException(400, "只支持STP/STEP文件")

    task_id = str(uuid.uuid4())

    file_path, file_size = await file_handler.save_uploaded_file(file)
    logger.info(f"[UPLOAD] 文件已保存: {file_path} ({file_size} bytes), task_id={task_id}")

    storage_service = StorageIntegrationService()

    stp_file = await storage_service.save_stp_file(
        session=db_session,
        file_path=file_path,
        original_filename=file.filename,
        user_id=current_user.id
    )
    logger.info(f"[UPLOAD] STP文件已存入RustFS+PG: stp_file.id={stp_file.id}")

    await storage_service.create_processing_task(db_session, task_id, stp_file.id)

    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())
    )
    await redis_task_manager.set_task(task_id, task_info)

    background_tasks.add_task(
        processing_service.process_file_with_storage,
        task_id, file_path, stp_file.id, material
    )
    logger.info(f"[UPLOAD] 后台处理已调度: task_id={task_id}")

    return {
        "task_id": task_id,
        "status": "processing",
        "message": "文件上传成功，开始处理并存储到数据库",
        "file_info": {
            "filename": file.filename,
            "size": file_size,
            "pythonocc_available": True,
            "database_file_id": stp_file.id
        }
    }