geMoldInsight/src/moldinsight/services/storage_service.py

# services/storage_service.py
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, update
from datetime import datetime
import hashlib
import json
from pathlib import Path
from typing import Optional, Dict, Any

from shared.models.database import STPFile, GeometryData, HTMLFile, ProcessingTask
from shared.utils.logger import get_logger

from shared.models.database import MoldCavityData

logger = get_logger(__name__)

class StorageService:
    """数据存储服务"""
    
    def __init__(self, db_session: AsyncSession):
        self.db_session = db_session
    
    async def save_stp_file(
        self, 
        filename: str, 
        file_path: str, 
        file_size: int,
        file_content: Optional[bytes] = None
    ) -> STPFile:
        """保存STP文件信息到数据库"""
        try:
            # 计算文件哈希
            file_hash = self._calculate_file_hash(file_path, file_content)
            
            # 检查是否已存在相同文件
            existing_file = await self.db_session.execute(
                select(STPFile).where(STPFile.file_hash == file_hash)
            )
            existing_file = existing_file.scalar_one_or_none()
            
            if existing_file:
                logger.info(f"文件已存在，跳过保存: {filename}")
                return existing_file
            
            # 创建新的STP文件记录
            stp_file = STPFile(
                filename=filename,
                original_filename=filename,
                file_path=file_path,
                file_size=file_size,
                file_hash=file_hash,
                file_content=file_content,
                upload_time=datetime.now(),
                status="pending",
                # 必填字段提供默认值
                object_key=f"stp_files/{file_hash}",
                storage_bucket="default",
                object_url=None
            )
            
            self.db_session.add(stp_file)
            await self.db_session.commit()
            await self.db_session.refresh(stp_file)
            
            logger.info(f"STP文件保存成功: {filename} (ID: {stp_file.id})")
            return stp_file
            
        except Exception as e:
            await self.db_session.rollback()
            logger.error(f"保存STP文件失败: {e}")
            raise
    
    async def save_geometry_data(
        self, 
        stp_file_id: int, 
        geometry_json: Dict[str, Any], 
        analysis_method: str
    ) -> GeometryData:
        """保存几何数据JSON到数据库"""
        try:
            # 提取关键几何属性用于快速查询
            volume = geometry_json.get("volume")
            surface_area = geometry_json.get("surface_area")
            bounding_box = geometry_json.get("bounding_box", {})
            
            geometry_data = GeometryData(
                stp_file_id=stp_file_id,
                analysis_method=analysis_method,
                volume=volume,
                surface_area=surface_area,
                bounding_box_min=bounding_box.get("min"),
                bounding_box_max=bounding_box.get("max"),
                created_time=datetime.now(),
                # 必填字段提供默认值
                object_key=f"geometry_data/{stp_file_id}",
                storage_bucket="default",
                object_url=None
            )
            
            self.db_session.add(geometry_data)
            await self.db_session.commit()
            await self.db_session.refresh(geometry_data)
            
            logger.info(f"几何数据保存成功: STP文件ID {stp_file_id}")
            return geometry_data
            
        except Exception as e:
            await self.db_session.rollback()
            logger.error(f"保存几何数据失败: {e}")
            raise
    
    async def save_html_file(
        self, 
        stp_file_id: int, 
        filename: str, 
        file_path: str, 
        html_content: Optional[str] = None,
        visualization_type: str = "3d_viewer"
    ) -> HTMLFile:
        """保存HTML文件信息到数据库"""
        try:
            html_file = HTMLFile(
                stp_file_id=stp_file_id,
                filename=filename,
                file_path=file_path,
                html_content=html_content,
                visualization_type=visualization_type,
                has_interactive_elements=True,
                generated_time=datetime.now(),
                # 必填字段提供默认值
                object_key=f"html_files/{stp_file_id}",
                storage_bucket="default",
                object_url=None
            )
            
            self.db_session.add(html_file)
            await self.db_session.commit()
            await self.db_session.refresh(html_file)
            
            logger.info(f"HTML文件保存成功: {filename} (STP文件ID: {stp_file_id})")
            return html_file
            
        except Exception as e:
            await self.db_session.rollback()
            logger.error(f"保存HTML文件失败: {e}")
            raise
    
    async def create_processing_task(
        self, 
        task_id: str, 
        stp_file_id: int, 
        task_type: str = "stp_parsing"
    ) -> ProcessingTask:
        """创建处理任务记录"""
        try:
            task = ProcessingTask(
                task_id=task_id,
                stp_file_id=stp_file_id,
                task_type=task_type,
                status="pending",
                started_time=datetime.now()
            )
            
            self.db_session.add(task)
            await self.db_session.commit()
            await self.db_session.refresh(task)
            
            logger.info(f"处理任务创建成功: {task_id}")
            return task
            
        except Exception as e:
            await self.db_session.rollback()
            logger.error(f"创建处理任务失败: {e}")
            raise
    
    async def update_task_status(
        self, 
        task_id: str, 
        status: str, 
        progress: Optional[int] = None,
        current_step: Optional[str] = None,
        error_message: Optional[str] = None
    ):
        """更新任务状态"""
        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 self.db_session.execute(
                update(ProcessingTask)
                .where(ProcessingTask.task_id == task_id)
                .values(**update_data)
            )
            await self.db_session.commit()
            
            logger.info(f"任务状态更新: {task_id} -> {status}")
            
        except Exception as e:
            await self.db_session.rollback()
            logger.error(f"更新任务状态失败: {e}")
            raise
    
    async def update_stp_file_status(self, stp_file_id: int, status: str):
        """更新STP文件状态"""
        try:
            await self.db_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 self.db_session.commit()
            
            logger.info(f"STP文件状态更新: ID {stp_file_id} -> {status}")
            
        except Exception as e:
            await self.db_session.rollback()
            logger.error(f"更新STP文件状态失败: {e}")
            raise
    
    async def get_stp_file_by_id(self, stp_file_id: int) -> Optional[STPFile]:
        """根据ID获取STP文件"""
        try:
            result = await self.db_session.execute(
                select(STPFile).where(STPFile.id == stp_file_id)
            )
            return result.scalar_one_or_none()
        except Exception as e:
            logger.error(f"获取STP文件失败: {e}")
            return None
    
    async def get_geometry_data_by_stp_file_id(self, stp_file_id: int) -> Optional[GeometryData]:
        """根据STP文件ID获取几何数据"""
        try:
            result = await self.db_session.execute(
                select(GeometryData).where(GeometryData.stp_file_id == stp_file_id)
            )
            return result.scalar_one_or_none()
        except Exception as e:
            logger.error(f"获取几何数据失败: {e}")
            return None
    
    def _calculate_file_hash(self, file_path: str, file_content: Optional[bytes] = None) -> str:
        """计算文件哈希值"""
        sha256_hash = hashlib.sha256()
        
        if file_content:
            sha256_hash.update(file_content)
        else:
            # 从文件路径读取内容计算哈希
            with open(file_path, "rb") as f:
                for chunk in iter(lambda: f.read(4096), b""):
                    sha256_hash.update(chunk)
        
        return sha256_hash.hexdigest()

    async def save_mold_cavity_data(
            self,
            stp_file_id: int,
            cavity_json: Dict[str, Any],
            key_info: Dict[str, Any]
    ) -> MoldCavityData:
        """保存模具型腔数据"""
        try:
            mold_data = MoldCavityData(
                stp_file_id=stp_file_id,
                cavity_key_info=key_info,
                shrinkage_rate=cavity_json["metadata"]["shrinkage_rate"],
                draft_angle=cavity_json["metadata"]["draft_angle"],
                generated_time=datetime.now(),
                # 必填字段提供默认值
                detailed_object_key=f"mold_cavity/{stp_file_id}",
                storage_bucket="default"
            )

            self.db_session.add(mold_data)
            await self.db_session.commit()
            await self.db_session.refresh(mold_data)

            logger.info(f"模具型腔数据保存成功: STP文件ID {stp_file_id}")
            return mold_data

        except Exception as e:
            await self.db_session.rollback()
            logger.error(f"保存模具型腔数据失败: {e}")
            raise