x

src/moldinsight/services/shape_loader.py

@@ -0,0 +1,88 @@
+# 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