geMoldInsight/src/storage/object_storage.py

# storage/object_storage.py
"""MinIO/S3 对象存储服务"""
from minio import Minio
from minio.error import S3Error
from pathlib import Path
from typing import Optional, BinaryIO
from io import BytesIO
from utils.logger import get_logger
import hashlib
import uuid

logger = get_logger(__name__)


class ObjectStorageManager:
    """对象存储管理器 - MinIO/S3兼容"""
    
    def __init__(self):
        self.client: Optional[Minio] = None
        self.is_connected = False
        
        # 桶名称
        self.buckets = {
            'stp_files': 'moldinsight-stp-files',      # STP/STEP文件
            'geometry_data': 'moldinsight-geometry',    # 几何数据JSON
            'mold_cavities': 'moldinsight-mold-cavities', # 模具型腔数据
            'html_files': 'moldinsight-html',          # HTML报告文件
            'user_files': 'moldinsight-user-files'      # 用户上传的其他文件
        }
    
    async def connect(self, endpoint: str, access_key: str, secret_key: str, 
                   secure: bool = False):
        """连接到MinIO/S3服务"""
        try:
            self.client = Minio(
                endpoint,
                access_key=access_key,
                secret_key=secret_key,
                secure=secure
            )
            
            # 测试连接
            self.client.list_buckets()
            
            self.is_connected = True
            logger.info(f"对象存储连接成功: {endpoint}")
            
            # 确保所有桶都存在
            await self._ensure_buckets()
            
        except S3Error as e:
            logger.error(f"对象存储连接失败: {e}")
            self.is_connected = False
            raise
    
    async def _ensure_buckets(self):
        """确保所有必要的桶都存在"""
        for bucket_name in self.buckets.values():
            try:
                if not self.client.bucket_exists(bucket_name):
                    self.client.make_bucket(bucket_name)
                    logger.info(f"创建存储桶: {bucket_name}")
                else:
                    logger.debug(f"存储桶已存在: {bucket_name}")
            except S3Error as e:
                logger.error(f"创建存储桶失败 {bucket_name}: {e}")
    
    def _generate_object_key(self, original_filename: str, prefix: str = '') -> str:
        """生成对象存储的唯一键名"""
        # 提取文件扩展名
        ext = Path(original_filename).suffix
        
        # 生成唯一ID
        unique_id = str(uuid.uuid4())
        
        # 生成键名: prefix/unique_id + original_ext
        if prefix:
            return f"{prefix}/{unique_id}{ext}"
        return f"{unique_id}{ext}"
    
    async def upload_stp_file(self, file_path: Path, 
                            original_filename: str) -> dict:
        """上传STP文件到对象存储"""
        if not self.is_connected:
            raise RuntimeError("对象存储未连接")
        
        bucket_name = self.buckets['stp_files']
        
        # 计算文件哈希
        file_hash = self._calculate_file_hash(file_path)
        
        # 检查是否已存在
        existing_key = await self._find_file_by_hash(bucket_name, file_hash)
        if existing_key:
            logger.info(f"文件已存在，跳过上传: {existing_key}")
            return {
                'object_key': existing_key,
                'file_hash': file_hash,
                'already_exists': True
            }
        
        # 生成唯一键名
        object_key = self._generate_object_key(
            original_filename, 
            prefix='stp'
        )
        
        # 上传文件
        try:
            result = self.client.fput_object(
                bucket_name,
                object_key,
                str(file_path),
                content_type='application/octet-stream'
            )
            
            logger.info(f"STP文件上传成功: {object_key}")
            
            return {
                'object_key': object_key,
                'file_hash': file_hash,
                'file_size': result.size,
                'etag': result.etag,
                'already_exists': False
            }
        except S3Error as e:
            logger.error(f"STP文件上传失败: {e}")
            raise
    
    async def upload_geometry_data(self, geometry_json: dict, 
                                 file_hash: str) -> dict:
        """上传几何数据JSON到对象存储"""
        if not self.is_connected:
            raise RuntimeError("对象存储未连接")
        
        bucket_name = self.buckets['geometry_data']
        
        # 使用文件哈希作为键名的一部分
        object_key = f"geometry/{file_hash}.json"
        
        # 转换为字节
        import json
        json_bytes = json.dumps(geometry_json, ensure_ascii=False).encode('utf-8')
        
        # 上传
        try:
            result = self.client.put_object(
                bucket_name,
                object_key,
                BytesIO(json_bytes),
                length=len(json_bytes),
                content_type='application/json'
            )
            
            logger.info(f"几何数据上传成功: {object_key}")
            
            return {
                'object_key': object_key,
                'file_size': result.size,
                'etag': result.etag
            }
        except S3Error as e:
            logger.error(f"几何数据上传失败: {e}")
            raise
    
    async def upload_mold_cavity_data(self, cavity_json: dict,
                                     file_hash: str) -> dict:
        """上传模具型腔数据到对象存储"""
        if not self.is_connected:
            raise RuntimeError("对象存储未连接")
        
        bucket_name = self.buckets['mold_cavities']
        object_key = f"mold-cavity/{file_hash}.json"
        
        import json
        json_bytes = json.dumps(cavity_json, ensure_ascii=False).encode('utf-8')
        
        try:
            result = self.client.put_object(
                bucket_name,
                object_key,
                BytesIO(json_bytes),
                length=len(json_bytes),
                content_type='application/json'
            )
            
            logger.info(f"模具型腔数据上传成功: {object_key}")
            
            return {
                'object_key': object_key,
                'file_size': result.size,
                'etag': result.etag
            }
        except S3Error as e:
            logger.error(f"模具型腔数据上传失败: {e}")
            raise
    
    async def upload_html_file(self, html_content: str,
                             original_filename: str,
                             file_hash: str) -> dict:
        """上传HTML文件到对象存储"""
        if not self.is_connected:
            raise RuntimeError("对象存储未连接")
        
        bucket_name = self.buckets['html_files']
        object_key = f"html/{file_hash}.html"
        
        html_bytes = html_content.encode('utf-8')
        
        try:
            result = self.client.put_object(
                bucket_name,
                object_key,
                BytesIO(html_bytes),
                length=len(html_bytes),
                content_type='text/html; charset=utf-8'
            )
            
            logger.info(f"HTML文件上传成功: {object_key}")
            
            return {
                'object_key': object_key,
                'file_size': result.size,
                'etag': result.etag
            }
        except S3Error as e:
            logger.error(f"HTML文件上传失败: {e}")
            raise
    
    async def download_file(self, bucket_type: str, 
                         object_key: str) -> bytes:
        """从对象存储下载文件"""
        if not self.is_connected:
            raise RuntimeError("对象存储未连接")
        
        bucket_name = self.buckets.get(bucket_type)
        if not bucket_name:
            raise ValueError(f"未知的桶类型: {bucket_type}")
        
        try:
            response = self.client.get_object(bucket_name, object_key)
            data = response.read()
            response.close()
            response.release_conn()
            
            logger.debug(f"文件下载成功: {object_key}")
            return data
        except S3Error as e:
            logger.error(f"文件下载失败 {object_key}: {e}")
            raise
    
    async def get_presigned_url(self, bucket_type: str, 
                              object_key: str,
                              expires: int = 3600) -> str:
        """生成预签名URL（临时访问链接）"""
        if not self.is_connected:
            raise RuntimeError("对象存储未连接")
        
        bucket_name = self.buckets.get(bucket_type)
        if not bucket_name:
            raise ValueError(f"未知的桶类型: {bucket_type}")
        
        try:
            url = self.client.presigned_get_object(
                bucket_name,
                object_key,
                expires=expires
            )
            return url
        except S3Error as e:
            logger.error(f"生成预签名URL失败: {e}")
            raise
    
    async def delete_file(self, bucket_type: str, object_key: str):
        """删除对象存储中的文件"""
        if not self.is_connected:
            raise RuntimeError("对象存储未连接")
        
        bucket_name = self.buckets.get(bucket_type)
        if not bucket_name:
            raise ValueError(f"未知的桶类型: {bucket_type}")
        
        try:
            self.client.remove_object(bucket_name, object_key)
            logger.info(f"文件删除成功: {object_key}")
        except S3Error as e:
            logger.error(f"文件删除失败 {object_key}: {e}")
            raise
    
    def _calculate_file_hash(self, file_path: Path) -> str:
        """计算文件的SHA256哈希"""
        sha256_hash = hashlib.sha256()
        with open(file_path, 'rb') as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256_hash.update(byte_block)
        return sha256_hash.hexdigest()
    
    async def _find_file_by_hash(self, bucket_name: str, 
                                 file_hash: str) -> Optional[str]:
        """根据哈希查找已存在的文件"""
        try:
            objects = self.client.list_objects(bucket_name, recursive=True)
            for obj in objects:
                # 从对象键中提取哈希（如果有）
                if file_hash in obj.object_name:
                    return obj.object_name
            return None
        except S3Error as e:
            logger.warning(f"查找文件哈希失败: {e}")
            return None
    
    async def get_file_info(self, bucket_type: str, 
                         object_key: str) -> dict:
        """获取文件信息"""
        if not self.is_connected:
            raise RuntimeError("对象存储未连接")
        
        bucket_name = self.buckets.get(bucket_type)
        if not bucket_name:
            raise ValueError(f"未知的桶类型: {bucket_type}")
        
        try:
            stat = self.client.stat_object(bucket_name, object_key)
            return {
                'size': stat.size,
                'etag': stat.etag,
                'content_type': stat.content_type,
                'last_modified': stat.last_modified
            }
        except S3Error as e:
            logger.error(f"获取文件信息失败: {e}")
            raise
    
    async def list_files(self, bucket_type: str, 
                      prefix: str = '') -> list:
        """列出存储桶中的文件"""
        if not self.is_connected:
            raise RuntimeError("对象存储未连接")
        
        bucket_name = self.buckets.get(bucket_type)
        if not bucket_name:
            raise ValueError(f"未知的桶类型: {bucket_type}")
        
        try:
            objects = self.client.list_objects(bucket_name, prefix=prefix)
            return [
                {
                    'object_key': obj.object_name,
                    'size': obj.size,
                    'etag': obj.etag,
                    'last_modified': obj.last_modified
                }
                for obj in objects
            ]
        except S3Error as e:
            logger.error(f"列出文件失败: {e}")
            raise


# 全局对象存储管理器实例
storage_manager = ObjectStorageManager()