RagFlow_Plugs/utils/minio/minio_util.py

from minio import Minio
from typing import BinaryIO
from datetime import timedelta
from conf.config import MinioConfig
from .file_utils import generate_unique_filename

class MinIOUtil:
    def __init__(self):
        self.client = Minio(
            endpoint=MinioConfig.get_minio_endpoint(),
            access_key=MinioConfig.get_minio_access_key(),
            secret_key=MinioConfig.get_minio_secret_key(),
            secure=False
        )
        self.bucket_name = MinioConfig.get_minio_bucket_name()
        self._ensure_bucket_exists()

    def _ensure_bucket_exists(self):
        """确保存储桶存在，若不存在则创建"""
        try:
            if not self.client.bucket_exists(self.bucket_name):
                self.client.make_bucket(self.bucket_name)
                print(f"Bucket '{self.bucket_name}' created successfully.")
            else:
                print(f"Bucket '{self.bucket_name}' already exists.")
        except Exception as e:
            raise RuntimeError(f"Failed to create bucket: {e}")

    def upload_file(self, file: BinaryIO, original_filename: str) -> str:
        """上传文件并返回URL"""
        try:
            # 生成唯一文件名，防止冲突
            unique_filename = generate_unique_filename(original_filename)
            content_type = self._get_content_type(original_filename)
            
            # 获取文件长度
            if hasattr(file, 'getbuffer'):
                # 对于BytesIO对象，获取其缓冲区大小
                length = file.getbuffer().nbytes
            elif hasattr(file, 'tell') and hasattr(file, 'seek'):
                # 对于支持seek/tell的文件对象，获取其大小
                current_pos = file.tell()
                file.seek(0, 2)  # 移动到文件末尾
                length = file.tell()
                file.seek(current_pos)  # 恢复到原始位置
            else:
                # 对于其他类型，使用-1让MinIO自动处理
                length = -1
            
            # 上传文件（支持大文件分块上传）
            self.client.put_object(
                bucket_name=self.bucket_name,
                object_name=unique_filename,
                data=file,
                length=length,
                content_type=content_type
            )
            
            # 生成公开可访问的URL（可选：设置过期时间或私有访问）
            url = self.client.get_presigned_url(
                method="GET",
                bucket_name=self.bucket_name,
                object_name=unique_filename,
                expires=timedelta(seconds=3600)  # 1小时有效期（可设为永久或更短）
            )
            
            return url
        except Exception as e:
            raise RuntimeError(f"File upload failed: {e}")

    def download_file(self, object_name: str) -> BinaryIO:
        """下载文件并返回文件流"""
        try:
            response = self.client.get_object(
                bucket_name=self.bucket_name,
                object_name=object_name
            )
            return response
        except Exception as e:
            raise RuntimeError(f"File download failed: {e}")

    def delete_file(self, object_name: str) -> bool:
        """删除文件"""
        try:
            self.client.remove_object(
                bucket_name=self.bucket_name,
                object_name=object_name
            )
            return True
        except Exception as e:
            print(f"Delete failed: {e}")
            return False

    def _get_content_type(self, filename: str) -> str:
        """根据文件后缀推断MIME类型"""
        ext = filename.split('.')[-1].lower()
        mime_map = {
            'jpg': 'image/jpeg', 'jpeg': 'image/jpeg',
            'png': 'image/png', 'gif': 'image/gif',
            'pdf': 'application/pdf',
            'txt': 'text/plain',
            'mp4': 'video/mp4',
            'mp3': 'audio/mpeg'
        }
        return mime_map.get(ext, 'application/octet-stream')