89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
|
# services/shape_loader.py
|
||
|
|
"""按 task_id 从持久化存储重建 OCC 几何形状。
|
||
|
|
|
||
|
|
任务完成后 TopoDS_Shape 不驻留内存/Redis（原生内存与体积原因），
|
||
|
|
需要几何的端点（倒扣检测、按需重导出等）通过 STP 原件重建：
|
||
|
|
PG(object_key) -> RustFS 下载 -> 临时文件 -> OCC 单线程 executor 解析。
|
||
|
|
"""
|
||
|
|
import tempfile
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
from sqlalchemy import select
|
||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
|
|
||
|
|
from shared.models.database import ProcessingTask, STPFile
|
||
|
|
from shared.utils.logger import get_logger
|
||
|
|
|
||
|
|
logger = get_logger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class ShapeLoader:
|
||
|
|
"""任务几何重建器"""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
from moldinsight.core.stp_parser import STPParser
|
||
|
|
from moldinsight.services.processing_service import processing_service
|
||
|
|
|
||
|
|
self._parser = STPParser()
|
||
|
|
self._processing = processing_service
|
||
|
|
|
||
|
|
async def load_shape_for_task(
|
||
|
|
self, db_session: AsyncSession, task_id: str
|
||
|
|
) -> Optional["object"]:
|
||
|
|
"""重建任务的产品几何。任务不存在或 STP 原件不可用时返回 None。"""
|
||
|
|
result = await db_session.execute(
|
||
|
|
select(ProcessingTask, STPFile)
|
||
|
|
.join(STPFile, ProcessingTask.stp_file_id == STPFile.id)
|
||
|
|
.where(ProcessingTask.task_id == task_id)
|
||
|
|
)
|
||
|
|
row = result.first()
|
||
|
|
if not row:
|
||
|
|
logger.warning(f"几何重建失败：任务不存在 {task_id}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
_, stp_file = row
|
||
|
|
if not stp_file.object_key:
|
||
|
|
logger.warning(f"几何重建失败：任务缺少 object_key {task_id}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
from moldinsight.storage.rustfs_storage import rustfs_manager
|
||
|
|
|
||
|
|
try:
|
||
|
|
data = await rustfs_manager.download_file(
|
||
|
|
file_type="stp_files", object_key=stp_file.object_key
|
||
|
|
)
|
||
|
|
except Exception as exc:
|
||
|
|
logger.error(f"几何重建失败：STP 原件下载失败 {task_id}: {exc}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
with tempfile.NamedTemporaryFile(suffix=".stp", delete=False) as tmp:
|
||
|
|
tmp.write(data)
|
||
|
|
tmp_path = Path(tmp.name)
|
||
|
|
|
||
|
|
try:
|
||
|
|
shape = await self._processing.run_occ(
|
||
|
|
self._parser.load_step_file, tmp_path
|
||
|
|
)
|
||
|
|
return shape
|
||
|
|
except Exception as exc:
|
||
|
|
logger.error(f"几何重建失败：STP 解析失败 {task_id}: {exc}")
|
||
|
|
return None
|
||
|
|
finally:
|
||
|
|
try:
|
||
|
|
tmp_path.unlink(missing_ok=True)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
# 惰性单例：__init__ 会实例化 STPParser 并校验 OCC 可用性，
|
||
|
|
# 延迟到首次真实使用，避免模块导入期失败拖垮路由加载
|
||
|
|
_shape_loader: Optional[ShapeLoader] = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_shape_loader() -> ShapeLoader:
|
||
|
|
global _shape_loader
|
||
|
|
if _shape_loader is None:
|
||
|
|
_shape_loader = ShapeLoader()
|
||
|
|
return _shape_loader
|