geMoldInsight/src/moldinsight/services/task_storage_service.py

# services/task_storage_service.py
"""任务与源文件生命周期存储——PostgreSQL（+ 源文件 RustFS 上传）。

批次 3 自 storage_integration_rustfs.py 按职责拆分（原 867 行混杂
写入/查询/历史三类职责）：
- 本模块：STPFile 生命周期 + ProcessingTask 创建/状态/参数
- 分析结果数据：analysis_storage_service.AnalysisStorageService
- 历史查询视图：file_history_service.FileHistoryService
"""
from pathlib import Path
from typing import Optional, Dict, Any
from datetime import datetime
import uuid

from sqlalchemy import update, select
from sqlalchemy.ext.asyncio import AsyncSession

from moldinsight.models import STPFile, ProcessingTask
from moldinsight.storage.rustfs_storage import rustfs_manager
from shared.utils.logger import get_logger

logger = get_logger(__name__)


class TaskStorageService:
    """STP 文件与处理任务的生命周期存储"""

    async def save_stp_file(self, session: AsyncSession,
                           file_path: Path,
                           original_filename: str,
                           user_id: Optional[int] = None,
                           upload_batch: Optional[str] = None) -> STPFile:
        """保存STP文件到PostgreSQL元数据 + RustFS对象存储

        支持同一文件多次上传，每次上传都会创建新记录
        """

        # 1. 上传到RustFS
        upload_result = await rustfs_manager.upload_file(
            file_type='stp_files',
            file_path=file_path,
            original_filename=original_filename,
            metadata={
                'original_filename': original_filename,
                'user_id': str(user_id) if user_id else 'anonymous',
                'upload_batch': upload_batch or str(uuid.uuid4())
            }
        )

        file_hash = upload_result['file_hash']
        batch_id = upload_batch or str(uuid.uuid4())

        # 2. 创建新PostgreSQL记录（每次上传都创建新记录）
        stp_file = STPFile(
            user_id=user_id,
            object_key=upload_result['object_key'],
            storage_bucket=upload_result['bucket'],
            original_filename=original_filename,
            file_size=upload_result['file_size'],
            file_hash=file_hash,
            upload_batch=batch_id,
            status="uploaded",
            file_path=str(file_path),
            upload_time=datetime.now()
        )

        session.add(stp_file)
        # D9：仅 flush，与 ProcessingTask 由路由层一并原子提交（避免孤儿文件记录）
        await session.flush()
        await session.refresh(stp_file)

        logger.info(f"STP文件保存成功 RustFS: {stp_file.id}, 批次: {batch_id}")
        return stp_file

    async def create_processing_task(
        self,
        session: AsyncSession,
        task_id: str,
        stp_file_id: int,
        task_type: str = "stp_parsing",
        parameters: Optional[Dict[str, Any]] = None,
        batch_id: Optional[str] = None,
    ) -> ProcessingTask:
        """创建处理任务记录（D9：仅 flush 不 commit，事务由调用方收口——
        与 STPFile 记录同批提交，避免留下无任务的孤儿文件记录；batch_id 用于批量任务聚合查询）
        """
        try:
            task = ProcessingTask(
                task_id=task_id,
                stp_file_id=stp_file_id,
                task_type=task_type,
                status="pending",
                started_time=datetime.now(),
                parameters=parameters or {},
                batch_id=batch_id,
            )

            session.add(task)
            await session.flush()

            logger.info(f"处理任务创建成功: {task_id}")
            return task

        except Exception as e:
            await session.rollback()
            logger.error(f"创建处理任务失败: {e}")
            raise

    async def update_task_status(
        self,
        session: AsyncSession,
        task_id: str,
        status: str,
        progress: Optional[int] = None,
        current_step: Optional[str] = None,
        error_message: Optional[str] = None
    ):
        """更新任务状态（保留即时 commit：进度/状态需跨事务对外可见，
        处理链路中的各阶段进度依赖它落库——D9 收口仅针对数据本体写方法）"""
        try:
            update_data = {
                "status": status,
                "completed_time": datetime.now() if status in ["completed", "failed"] else None,
                "error_message": error_message
            }

            if progress is not None:
                update_data["progress"] = progress
            if current_step is not None:
                update_data["current_step"] = current_step

            await session.execute(
                update(ProcessingTask)
                .where(ProcessingTask.task_id == task_id)
                .values(**update_data)
            )
            await session.commit()

            logger.info(f"任务状态更新: {task_id} -> {status}")

        except Exception as e:
            await session.rollback()
            logger.error(f"更新任务状态失败: {e}")
            raise

    async def update_task_parameters(
        self,
        session: AsyncSession,
        task_id: str,
        parameters: Dict[str, Any],
    ):
        """合并更新任务参数，便于保存阶段耗时等元数据。（D9：flush 不 commit，事务由调用方收口）"""
        try:
            task = await session.execute(
                select(ProcessingTask).where(ProcessingTask.task_id == task_id)
            )
            task = task.scalar_one_or_none()
            if task is None:
                return

            merged = dict(task.parameters or {})
            merged.update(parameters or {})
            task.parameters = merged
            await session.flush()
        except Exception as e:
            await session.rollback()
            logger.error(f"更新任务参数失败: {e}")
            raise

    async def update_stp_file_status(self, session: AsyncSession, stp_file_id: int, status: str):
        """更新STP文件状态（保留即时 commit，理由同 update_task_status）"""
        try:
            await session.execute(
                update(STPFile)
                .where(STPFile.id == stp_file_id)
                .values(
                    status=status,
                    processed_time=datetime.now() if status in ["completed", "failed"] else None
                )
            )
            await session.commit()

            logger.info(f"STP文件状态更新: ID {stp_file_id} -> {status}")

        except Exception as e:
            await session.rollback()
            logger.error(f"更新STP文件状态失败: {e}")
            raise

    async def update_stp_file_analysis_summary(
        self,
        session: AsyncSession,
        stp_file_id: int,
        volume: Optional[float] = None,
        surface_area: Optional[float] = None,
        product_weight: Optional[float] = None
    ):
        """更新STP文件的分析摘要字段（用于快速查询）"""
        try:
            update_data = {}
            if volume is not None:
                update_data['volume'] = volume
            if surface_area is not None:
                update_data['surface_area'] = surface_area
            if product_weight is not None:
                update_data['product_weight'] = product_weight

            if update_data:
                await session.execute(
                    update(STPFile)
                    .where(STPFile.id == stp_file_id)
                    .values(**update_data)
                )
                # D9：flush 不 commit，随结果包由编排层统一提交
                await session.flush()
                logger.info(f"STP文件分析摘要更新: ID {stp_file_id}")

        except Exception as e:
            await session.rollback()
            logger.error(f"更新STP文件分析摘要失败: {e}")
            raise