# 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 import uuid from shared.models.database import ( STPFile, GeometryData, MeshData, MoldCavityData, HTMLFile, ProcessingTask, User, FeatureDetection, DesignRecommendation, UserActivity, SystemLog ) from moldinsight.storage.rustfs_storage import rustfs_manager from shared.utils.logger import get_logger logger = get_logger(__name__) class StorageIntegrationService: """存储集成服务 - PostgreSQL + RustFS""" @staticmethod def _resolve_best_scheme_payload(cavity_json: Dict[str, Any]) -> Dict[str, Any]: """从新多方案/旧单方案结构中解析推荐方案和型腔详情。""" if not isinstance(cavity_json, dict): return { "best_scheme_id": None, "best_scheme": {}, "best_cavity_data": {}, "key_info": {}, } candidate_schemes = cavity_json.get("candidate_schemes") or [] if not candidate_schemes: key_info = cavity_json.get("mold_cavities", {}).get("cavity_key_info", {}) return { "best_scheme_id": cavity_json.get("best_scheme_id"), "best_scheme": {}, "best_cavity_data": cavity_json, "key_info": key_info, } best_scheme_id = cavity_json.get("best_scheme_id") best_scheme = candidate_schemes[0] if best_scheme_id: for scheme in candidate_schemes: if scheme.get("scheme_id") == best_scheme_id: best_scheme = scheme break best_cavity_data = best_scheme.get("cavity_data", {}) if isinstance(best_scheme, dict) else {} key_info = best_scheme.get("key_info", {}) if isinstance(best_scheme, dict) else {} if not key_info: key_info = best_cavity_data.get("mold_cavities", {}).get("cavity_key_info", {}) return { "best_scheme_id": best_scheme.get("scheme_id") or best_scheme_id, "best_scheme": best_scheme, "best_cavity_data": best_cavity_data, "key_info": key_info, } @staticmethod def _parse_first_number(value: Any) -> Optional[float]: if value is None: return None try: return float(value) except (TypeError, ValueError): pass import re matches = re.findall(r"\d+(?:\.\d+)?", str(value)) if not matches: return None try: return float(matches[0]) except (TypeError, ValueError): return None async def save_stp_file(self, session: AsyncSession, file_path: Path, original_filename: str, user_id: Optional[int] = None, upload_batch: Optional[str] = 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', 'upload_batch': upload_batch or str(uuid.uuid4()) } ) file_hash = upload_result['file_hash'] batch_id = upload_batch or str(uuid.uuid4()) # 2. 创建新PostgreSQL记录（每次上传都创建新记录） from datetime import datetime 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, upload_batch=batch_id, status="uploaded", file_path=str(file_path), upload_time=datetime.now() ) session.add(stp_file) await session.commit() await session.refresh(stp_file) logger.info(f"STP文件保存成功 RustFS: {stp_file.id}, 批次: {batch_id}") return stp_file async def create_processing_task( self, session: AsyncSession, task_id: str, stp_file_id: int, task_type: str = "stp_parsing", parameters: Optional[Dict[str, Any]] = None, ) -> ProcessingTask: """创建处理任务记录""" try: task = ProcessingTask( task_id=task_id, stp_file_id=stp_file_id, task_type=task_type, status="pending", started_time=datetime.now(), parameters=parameters or {}, ) 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_task_parameters( self, session: AsyncSession, task_id: str, parameters: Dict[str, Any], ): """合并更新任务参数，便于保存阶段耗时等元数据。""" try: task = await session.execute( select(ProcessingTask).where(ProcessingTask.task_id == task_id) ) task = task.scalar_one_or_none() if task is None: return merged = dict(task.parameters or {}) merged.update(parameters or {}) task.parameters = merged await session.commit() 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. 提取关键信息（兼容多方案与单方案结构） payload = self._resolve_best_scheme_payload(cavity_json) best_scheme_id = payload.get("best_scheme_id") best_scheme = payload.get("best_scheme") or {} best_cavity_data = payload.get("best_cavity_data") or {} metadata = best_cavity_data.get('metadata', {}) product_analysis = best_cavity_data.get('product_analysis', {}) manufacturing_info = best_cavity_data.get('manufacturing_info', {}) mold_size = manufacturing_info.get('estimated_mold_size', {}) key_info = payload.get("key_info") or {} if not key_info: key_info = best_cavity_data.get('mold_cavities', {}).get('cavity_key_info', {}) mold_material = ( metadata.get("selected_material") or manufacturing_info.get("recommended_material") or 'Aluminum Alloy 7075' ) parting_line_length = self._parse_first_number( manufacturing_info.get("parting_line_length") ) # 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=mold_material, shrinkage_rate=metadata.get('shrinkage_rate', 0.005), draft_angle=metadata.get('draft_angle', 2.0), parting_line_length=parting_line_length, # 提取的摘要字段 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'), # 多方案可信化摘要 best_scheme_id=best_scheme_id, confidence_score=best_scheme.get("confidence_score"), is_fallback=best_scheme.get("is_fallback"), fallback_reason=best_scheme.get("fallback_reason"), ) 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 进 PG：读取路径走 RustFS（html_json.content）， # PG 仅存对象键与文件名，避免大文本撑爆表 html_content=None, 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('type') or 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, meta_data=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文件及其所有关联数据""" from sqlalchemy.orm import joinedload try: # 1. 获取STP文件记录（使用 joinedload 预加载关联数据） result = await session.execute( select(STPFile).options( joinedload(STPFile.geometry_data), joinedload(STPFile.mesh_data), joinedload(STPFile.mold_cavity_data), joinedload(STPFile.html_file), joinedload(STPFile.analysis_metrics) ).where(STPFile.id == stp_file_id) ) stp_file = result.scalar_one_or_none() if not stp_file: raise ValueError(f"STP文件不存在: {stp_file_id}") except Exception as e: logger.error(f"获取STP文件记录失败: {e}") raise 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': [], 'analysis_metrics': None # 新增分析指标字段 } # 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'] = [ { 'type': r.rec_type, # 改为 type 以匹配前端期望的字段名 'priority': r.priority, 'description': r.description, 'reason': r.reason, 'parameters': r.parameters } for r in recommendations.scalars().all() ] # 4. 获取分析指标 if stp_file.analysis_metrics: result['analysis_metrics'] = { 'volume_utilization': stp_file.analysis_metrics.volume_utilization, 'topology_complexity': stp_file.analysis_metrics.topology_complexity, 'wall_uniformity': stp_file.analysis_metrics.wall_uniformity, 'analysis_summary': stp_file.analysis_metrics.analysis_summary, 'verification_status': stp_file.analysis_metrics.verification_status, 'verification_volume_diff': stp_file.analysis_metrics.verification_volume_diff, 'verification_area_diff': stp_file.analysis_metrics.verification_area_diff, 'verification_details': stp_file.analysis_metrics.verification_details, } return result async def get_file_history_by_filename( self, session: AsyncSession, filename: str, user_id: Optional[int] = None, limit: int = 50 ) -> list: """获取同一文件名的所有上传历史记录""" from shared.models.database import ProcessingTask from sqlalchemy.orm import joinedload query = select(STPFile).options( joinedload(STPFile.processing_tasks) ).where( STPFile.original_filename == filename ).order_by(STPFile.upload_time.desc()) if user_id: query = query.where(STPFile.user_id == user_id) query = query.limit(limit) result = await session.execute(query) files = result.unique().scalars().all() return [ { 'id': f.id, 'task_id': f.processing_tasks[0].task_id if f.processing_tasks else None, 'upload_batch': f.upload_batch, 'upload_time': f.upload_time.strftime('%Y-%m-%d %H:%M:%S') if f.upload_time else None, 'file_size': f.file_size, 'status': f.status, 'volume': f.volume, 'surface_area': f.surface_area, 'product_weight': f.product_weight, 'has_analysis': f.status == 'completed' } for f in files ] async def get_all_file_groups( self, session: AsyncSession, user_id: Optional[int] = None, limit: int = 100 ) -> list: """获取所有文件分组（按文件名分组），包含每个文件的最新分析结果""" from sqlalchemy import func, desc from sqlalchemy.orm import joinedload from shared.models.database import ProcessingTask # 子查询：获取每个文件名的最新上传 subquery = ( select( STPFile.original_filename, func.max(STPFile.upload_time).label('latest_upload') ) .group_by(STPFile.original_filename) .order_by(desc('latest_upload')) .limit(limit) ) if user_id: subquery = subquery.where(STPFile.user_id == user_id) subquery = subquery.subquery() # 主查询：获取最新记录和统计信息 query = ( select(STPFile).options( joinedload(STPFile.processing_tasks) ) .join( subquery, (STPFile.original_filename == subquery.c.original_filename) & (STPFile.upload_time == subquery.c.latest_upload) ) .order_by(STPFile.upload_time.desc()) ) result = await session.execute(query) latest_files = result.unique().scalars().all() # 一次性聚合每个文件名的上传次数（替代逐文件 count 的 N+1 查询） count_subquery = ( select(STPFile.original_filename, func.count().label("upload_count")) .group_by(STPFile.original_filename) ) if user_id: count_subquery = count_subquery.where(STPFile.user_id == user_id) count_result = await session.execute(count_subquery) upload_counts = { row.original_filename: row.upload_count for row in count_result } # 获取每个文件名的上传次数 file_groups = [] for f in latest_files: task_id = f.processing_tasks[0].task_id if f.processing_tasks else None upload_count = upload_counts.get(f.original_filename, 1) file_groups.append({ 'filename': f.original_filename, 'latest_id': f.id, 'latest_task_id': task_id, 'latest_upload_time': f.upload_time.strftime('%Y-%m-%d %H:%M:%S') if f.upload_time else None, 'latest_status': f.status, 'upload_count': upload_count, 'file_size': f.file_size, 'volume': f.volume, 'surface_area': f.surface_area, 'product_weight': f.product_weight }) return file_groups async def update_stp_file_analysis_summary( self, session: AsyncSession, stp_file_id: int, volume: Optional[float] = None, surface_area: Optional[float] = None, product_weight: Optional[float] = None ): """更新STP文件的分析摘要字段（用于快速查询）""" try: update_data = {} if volume is not None: update_data['volume'] = volume if surface_area is not None: update_data['surface_area'] = surface_area if product_weight is not None: update_data['product_weight'] = product_weight if update_data: await session.execute( update(STPFile) .where(STPFile.id == stp_file_id) .values(**update_data) ) await session.commit() logger.info(f"STP文件分析摘要更新: ID {stp_file_id}") except Exception as e: await session.rollback() logger.error(f"更新STP文件分析摘要失败: {e}") raise 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()