296 lines
10 KiB
Python
296 lines
10 KiB
Python
# 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 |