rag_server/applications/utils/oss/oss_client.py

import oss2
import os


class OSSClient:
    # 配置默认的 endpoint 地址
    DEFAULT_ENDPOINT = "oss-cn-hangzhou.aliyuncs.com"  # 默认华东1区（杭州）

    def __init__(self, access_key_id=None, access_key_secret=None, bucket_name=None):
        """
        初始化 OSS 客户端
        :param bucket_name: Bucket 名称
        """
        # 从环境变量中获取 Access Key 和 Secret
        if access_key_id is None or access_key_secret is None:
            access_key_id = "LTAIP6x1l3DXfSxm"
            access_key_secret = "KbTaM9ars4OX3PMS6Xm7rtxGr1FLon"
        if bucket_name is None:
            bucket_name = "art-pubbucket"

        # 检查是否有凭证
        if not access_key_id or not access_key_secret:
            raise ValueError(
                "ACCESS_KEY_ID and ACCESS_KEY_SECRET must be set in the environment variables."
            )

        # 使用默认的 endpoint 地址
        self.auth = oss2.Auth(access_key_id, access_key_secret)
        self.bucket = oss2.Bucket(self.auth, self.DEFAULT_ENDPOINT, bucket_name)

    def upload_file(self, local_file_path, oss_file_path):
        """
        上传文件到 OSS
        :param local_file_path: 本地文件路径
        :param oss_file_path: OSS 存储的文件路径（例如：pdfs/myfile.pdf）
        :return: 上传结果，成功返回文件信息，失败抛出异常
        """
        if not os.path.exists(local_file_path):
            raise FileNotFoundError(f"Local file {local_file_path} does not exist.")

        try:
            self.bucket.put_object_from_file(oss_file_path, local_file_path)
            return {"status": "success", "message": f"File uploaded to {oss_file_path}"}
        except Exception as e:
            raise Exception(f"Error uploading file to OSS: {str(e)}")

    def download_file(self, oss_file_path, local_file_path):
        """
        从 OSS 下载文件
        :param oss_file_path: OSS 文件路径（例如：pdfs/myfile.pdf）
        :param local_file_path: 本地保存路径
        :return: 下载结果，成功返回下载文件的路径，失败抛出异常
        """
        try:
            self.bucket.get_object_to_file(oss_file_path, local_file_path)
            return {
                "status": "success",
                "message": f"File downloaded to {local_file_path}",
            }
        except Exception as e:
            raise Exception(f"Error downloading file from OSS: {str(e)}")

    def delete_file(self, oss_file_path):
        """
        从 OSS 删除文件
        :param oss_file_path: OSS 文件路径（例如：pdfs/myfile.pdf）
        :return: 删除结果，成功返回消息，失败抛出异常
        """
        try:
            self.bucket.delete_object(oss_file_path)
            return {
                "status": "success",
                "message": f"File {oss_file_path} deleted from OSS",
            }
        except Exception as e:
            raise Exception(f"Error deleting file from OSS: {str(e)}")

    def file_exists(self, oss_file_path):
        """
        检查文件是否存在于 OSS 中
        :param oss_file_path: OSS 文件路径（例如：pdfs/myfile.pdf）
        :return: 布尔值，文件存在返回 True，文件不存在返回 False
        """
        try:
            self.bucket.get_object(oss_file_path)
            return True
        except oss2.exceptions.NoSuchKey:
            return False
        except Exception as e:
            raise Exception(f"Error checking file existence on OSS: {str(e)}")