Files
geMoldInsight/src/services/storage_integration_rustfs.py
T
2026-02-16 19:06:41 +08:00

570 lines
21 KiB
Python

# services/storage_integration_rustfs.py
"""存储集成服务 - 协调 PostgreSQL 和 RustFS"""
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, update
from pathlib import Path
from typing import Optional, Dict, Any
import json
from datetime import datetime
from models.database import (
STPFile, GeometryData, MeshData, MoldCavityData,
HTMLFile, ProcessingTask, User,
FeatureDetection, DesignRecommendation,
UserActivity, SystemLog
)
from storage.rustfs_storage import rustfs_manager
from utils.logger import get_logger
logger = get_logger(__name__)
class StorageIntegrationService:
"""存储集成服务 - PostgreSQL + RustFS"""
async def save_stp_file(self, session: AsyncSession,
file_path: Path,
original_filename: str,
user_id: Optional[int] = 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'
}
)
file_hash = upload_result['file_hash']
# 2. 检查是否已存在相同文件
existing_file = await session.execute(
select(STPFile).where(STPFile.file_hash == file_hash)
)
existing_file = existing_file.scalar_one_or_none()
if existing_file:
logger.info(f"文件已存在,返回现有记录: {existing_file.id}")
return existing_file
# 3. 创建新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,
status="uploaded",
file_path=str(file_path) # 保留本地路径以兼容
)
session.add(stp_file)
await session.commit()
await session.refresh(stp_file)
logger.info(f"STP文件保存成功 RustFS: {stp_file.id}")
return stp_file
async def create_processing_task(self, session: AsyncSession,
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()
)
session.add(task)
await session.commit()
await session.refresh(task)
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
):
"""更新任务状态"""
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_stp_file_status(self, session: AsyncSession, stp_file_id: int, status: str):
"""更新STP文件状态"""
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 save_geometry_data(self, session: AsyncSession,
stp_file_id: int,
geometry_json: Dict[str, Any],
analysis_method: str = "pythonocc") -> GeometryData:
"""保存几何数据到PostgreSQL元数据 + RustFS对象存储"""
# 1. 获取文件哈希
stp_file = await session.get(STPFile, stp_file_id)
file_hash = stp_file.file_hash
# 2. 上传到RustFS
upload_result = await rustfs_manager.upload_json_data(
file_type='geometry_data',
json_data=geometry_json,
file_hash=file_hash
)
# 3. 提取几何数据
if 'geometry_data' in geometry_json:
geo_data = geometry_json['geometry_data']
else:
geo_data = geometry_json
# 4. 创建PostgreSQL记录
geometry_data = GeometryData(
stp_file_id=stp_file_id,
object_key=upload_result['object_key'],
storage_bucket=upload_result['bucket'],
analysis_method=analysis_method,
# 提取摘要字段
volume=geo_data.get('volume'),
surface_area=geo_data.get('surface_area'),
bounding_box_min=geo_data.get('bounding_box', {}).get('min'),
bounding_box_max=geo_data.get('bounding_box', {}).get('max'),
center_of_mass=geo_data.get('center_of_mass'),
topology_faces=geo_data.get('topology', {}).get('faces'),
topology_edges=geo_data.get('topology', {}).get('edges'),
topology_vertices=geo_data.get('topology', {}).get('vertices')
)
session.add(geometry_data)
await session.commit()
await session.refresh(geometry_data)
logger.info(f"几何数据保存成功 RustFS: {geometry_data.id}")
return geometry_data
async def save_mesh_data(
self,
session: AsyncSession,
stp_file_id: int,
mesh_json: Dict[str, Any],
quality: str = "medium"
) -> MeshData:
"""保存网格数据到 PostgreSQL 元数据 + RustFS 对象存储
mesh_json 为完整网格 JSON(顶点、面、点云等),
PostgreSQL 只存 object_key 和一些摘要字段,详细数据放在 RustFS。
"""
# 1. 获取文件哈希
stp_file = await session.get(STPFile, stp_file_id)
file_hash = stp_file.file_hash
# 2. 上传网格 JSON 到 RustFS
upload_result = await rustfs_manager.upload_json_data(
file_type='mesh_data',
json_data=mesh_json,
file_hash=file_hash
)
# 3. 提取摘要信息
mesh_section = mesh_json.get('mesh', {})
pointcloud_section = mesh_json.get('pointcloud', {})
bbox = mesh_json.get('bounding_box', {})
vertices = mesh_section.get('vertices') or []
faces = mesh_section.get('faces') or []
vertex_count = len(vertices)
face_count = len(faces)
point_count = pointcloud_section.get('count')
# 4. 创建 PostgreSQL 记录
mesh_data = MeshData(
stp_file_id=stp_file_id,
object_key=upload_result['object_key'],
storage_bucket=upload_result['bucket'],
quality=quality,
vertex_count=vertex_count,
face_count=face_count,
point_count=point_count,
bounding_box_min=bbox.get('min'),
bounding_box_max=bbox.get('max')
)
session.add(mesh_data)
await session.commit()
await session.refresh(mesh_data)
logger.info(f"网格数据保存成功 RustFS: {mesh_data.id}")
return mesh_data
async def save_mold_cavity_data(self, session: AsyncSession,
stp_file_id: int,
cavity_json: Dict[str, Any]) -> MoldCavityData:
"""保存模具型腔数据到PostgreSQL元数据 + RustFS对象存储"""
# 1. 获取文件哈希
stp_file = await session.get(STPFile, stp_file_id)
file_hash = stp_file.file_hash
# 2. 上传到RustFS
upload_result = await rustfs_manager.upload_json_data(
file_type='mold_cavities',
json_data=cavity_json,
file_hash=file_hash
)
# 3. 提取关键信息
metadata = cavity_json.get('metadata', {})
product_analysis = cavity_json.get('product_analysis', {})
manufacturing_info = cavity_json.get('manufacturing_info', {})
mold_size = manufacturing_info.get('estimated_mold_size', {})
key_info = cavity_json.get('mold_cavities', {}).get('cavity_key_info', {})
# 4. 创建PostgreSQL记录
mold_cavity = MoldCavityData(
stp_file_id=stp_file_id,
detailed_object_key=upload_result['object_key'],
storage_bucket=upload_result['bucket'],
# 模具参数
mold_material=manufacturing_info.get('recommended_material', 'Aluminum Alloy 7075'),
shrinkage_rate=metadata.get('shrinkage_rate', 0.005),
draft_angle=metadata.get('draft_angle', 2.0),
# 提取的摘要字段
cavity_key_info=key_info,
mold_size_length=mold_size.get('length'),
mold_size_width=mold_size.get('width'),
mold_size_height=mold_size.get('height'),
estimated_clamping_force=manufacturing_info.get('estimated_clamping_force'),
product_volume=product_analysis.get('volume'),
# 从key_info中提取(如果存在)
product_weight=key_info.get('geometric_characteristics', {}).get('product_weight'),
wall_thickness_range=key_info.get('geometric_characteristics', {}).get('wall_thickness_range'),
complexity_score=key_info.get('geometric_characteristics', {}).get('complexity_score'),
# 质量评估
weld_line_risk=key_info.get('quality_considerations', {}).get('potential_weld_lines'),
sink_mark_risk=key_info.get('quality_considerations', {}).get('sink_mark_areas'),
warpage_risk=key_info.get('quality_considerations', {}).get('warpage_risk')
)
session.add(mold_cavity)
await session.commit()
await session.refresh(mold_cavity)
logger.info(f"模具型腔数据保存成功 RustFS: {mold_cavity.id}")
return mold_cavity
async def save_html_file(self, session: AsyncSession,
stp_file_id: int,
filename: str,
file_path: str,
html_content: Optional[str] = None,
visualization_type: str = "3d_viewer") -> HTMLFile:
"""保存HTML文件到PostgreSQL元数据 + RustFS对象存储"""
# 1. 获取文件哈希
stp_file = await session.get(STPFile, stp_file_id)
file_hash = stp_file.file_hash
# 2. 读取HTML内容(如果未提供)
if html_content is None:
try:
with open(file_path, 'r', encoding='utf-8') as f:
html_content = f.read()
except Exception as e:
logger.error(f"读取HTML文件失败: {e}")
html_content = ""
# 3. 上传到RustFS
html_json = {'content': html_content, 'filename': filename}
upload_result = await rustfs_manager.upload_json_data(
file_type='html_files',
json_data=html_json,
file_hash=file_hash
)
# 4. 创建PostgreSQL记录
html_file = HTMLFile(
stp_file_id=stp_file_id,
object_key=upload_result['object_key'],
storage_bucket=upload_result['bucket'],
filename=filename,
file_path=file_path, # 保留本地路径
html_content=html_content, # 保留内容以兼容
visualization_type=visualization_type
)
session.add(html_file)
await session.commit()
await session.refresh(html_file)
logger.info(f"HTML文件保存成功 RustFS: {html_file.id}")
return html_file
async def save_features_and_recommendations(
self, session: AsyncSession,
stp_file_id: int,
features: list,
recommendations: list
):
"""保存特征检测结果和设计建议"""
# 1. 保存特征
for feature in features:
feature_record = FeatureDetection(
stp_file_id=stp_file_id,
feature_type=feature.get('feature_type'),
confidence=feature.get('confidence'),
location=feature.get('location'),
dimensions=feature.get('dimensions'),
parameters=feature.get('parameters')
)
session.add(feature_record)
# 2. 保存建议
for rec in recommendations:
rec_record = DesignRecommendation(
stp_file_id=stp_file_id,
rec_type=rec.get('rec_type'),
priority=rec.get('priority'),
description=rec.get('description'),
reason=rec.get('reason'),
parameters=rec.get('parameters')
)
session.add(rec_record)
await session.commit()
logger.info(f"保存了 {len(features)} 个特征和 {len(recommendations)} 个建议")
async def log_user_activity(self, session: AsyncSession,
user_id: int,
activity_type: str,
resource_type: Optional[str] = None,
resource_id: Optional[int] = None,
description: Optional[str] = None,
metadata: Optional[Dict] = None,
ip_address: Optional[str] = None,
user_agent: Optional[str] = None):
"""记录用户活动"""
activity = UserActivity(
user_id=user_id,
activity_type=activity_type,
resource_type=resource_type,
resource_id=resource_id,
description=description,
metadata=metadata,
ip_address=ip_address,
user_agent=user_agent
)
session.add(activity)
await session.commit()
logger.debug(f"用户活动记录: {activity_type} by user {user_id}")
async def get_stp_file_with_data(self, session: AsyncSession,
stp_file_id: int) -> Dict[str, Any]:
"""获取STP文件及其所有关联数据"""
# 1. 获取STP文件记录
stp_file = await session.get(STPFile, stp_file_id)
if not stp_file:
raise ValueError(f"STP文件不存在: {stp_file_id}")
result = {
'metadata': {
'id': stp_file.id,
'original_filename': stp_file.original_filename,
'file_size': stp_file.file_size,
'file_hash': stp_file.file_hash,
'upload_time': stp_file.upload_time.isoformat() if stp_file.upload_time else None,
'status': stp_file.status,
'user_id': stp_file.user_id
},
'geometry_data': None,
'mesh_data': None,
'mold_cavity_data': None,
'html_content': None,
'features': [],
'recommendations': []
}
# 2. 从RustFS获取数据
try:
# 几何数据
if stp_file.geometry_data:
geo_data_bytes = await rustfs_manager.download_file(
file_type='geometry_data',
object_key=stp_file.geometry_data.object_key
)
result['geometry_data'] = json.loads(geo_data_bytes.decode('utf-8'))
# 模具型腔数据
if stp_file.mold_cavity_data:
cavity_data_bytes = await rustfs_manager.download_file(
file_type='mold_cavities',
object_key=stp_file.mold_cavity_data.detailed_object_key
)
result['mold_cavity_data'] = json.loads(cavity_data_bytes.decode('utf-8'))
# 网格数据
if stp_file.mesh_data:
mesh_bytes = await rustfs_manager.download_file(
file_type='mesh_data',
object_key=stp_file.mesh_data.object_key
)
result['mesh_data'] = json.loads(mesh_bytes.decode('utf-8'))
# HTML文件
if stp_file.html_file:
html_bytes = await rustfs_manager.download_file(
file_type='html_files',
object_key=stp_file.html_file.object_key
)
html_json = json.loads(html_bytes.decode('utf-8'))
result['html_content'] = html_json.get('content', '')
except Exception as e:
logger.error(f"从RustFS获取数据失败: {e}")
# 3. 从PostgreSQL获取特征和建议
features = await session.execute(
select(FeatureDetection).where(FeatureDetection.stp_file_id == stp_file_id)
)
result['features'] = [
{
'feature_type': f.feature_type,
'confidence': f.confidence,
'location': f.location,
'dimensions': f.dimensions,
'parameters': f.parameters
}
for f in features.scalars().all()
]
recommendations = await session.execute(
select(DesignRecommendation).where(DesignRecommendation.stp_file_id == stp_file_id)
)
result['recommendations'] = [
{
'rec_type': r.rec_type,
'priority': r.priority,
'description': r.description,
'reason': r.reason,
'parameters': r.parameters
}
for r in recommendations.scalars().all()
]
return result
async def delete_stp_file_cascade(self, session: AsyncSession,
stp_file_id: int):
"""级联删除STP文件及其所有关联数据"""
stp_file = await session.get(STPFile, stp_file_id)
if not stp_file:
raise ValueError(f"STP文件不存在: {stp_file_id}")
# 1. 删除RustFS中的文件
try:
if stp_file.object_key:
await rustfs_manager.delete_file('stp_files', stp_file.object_key)
except Exception as e:
logger.error(f"删除RustFS文件失败: {e}")
try:
if stp_file.geometry_data:
await rustfs_manager.delete_file('geometry_data', stp_file.geometry_data.object_key)
except Exception as e:
logger.error(f"删除几何数据失败: {e}")
try:
if stp_file.mold_cavity_data:
await rustfs_manager.delete_file('mold_cavities', stp_file.mold_cavity_data.detailed_object_key)
except Exception as e:
logger.error(f"删除型腔数据失败: {e}")
try:
if stp_file.mesh_data:
await rustfs_manager.delete_file('mesh_data', stp_file.mesh_data.object_key)
except Exception as e:
logger.error(f"删除网格数据失败: {e}")
try:
if stp_file.html_file:
await rustfs_manager.delete_file('html_files', stp_file.html_file.object_key)
except Exception as e:
logger.error(f"删除HTML文件失败: {e}")
# 2. 级联删除PostgreSQL记录(通过外键自动处理)
await session.delete(stp_file)
await session.commit()
logger.info(f"STP文件及其关联数据已删除: {stp_file_id}")
# 全局存储集成服务实例
storage_integration = StorageIntegrationService()