222 lines
7.9 KiB
Python
222 lines
7.9 KiB
Python
|
# 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
|