oss_utils.py

# oss_utils.py
# OSS相关工具函数
import os
import oss2
from typing import List
import shutil

# OSS配置
OSS_CONFIG = {
    "access_key_id": os.getenv("OSS_ACCESS_KEY_ID"),
    "access_key_secret": os.getenv("OSS_ACCESS_KEY_SECRET"),
    "endpoint": os.getenv("OSS_ENDPOINT"),
    "bucket_name": os.getenv("OSS_BUCKET_NAME")
}

# 初始化OSS客户端
auth = oss2.Auth(OSS_CONFIG["access_key_id"], OSS_CONFIG["access_key_secret"])
bucket = oss2.Bucket(auth, OSS_CONFIG["endpoint"], OSS_CONFIG["bucket_name"])

# 临时文件根目录
TMP_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
os.makedirs(TMP_ROOT, exist_ok=True)

def list_oss_files(folder_path: str) -> List[str]:
    """列出OSS文件夹中的所有文件"""
    files = []
    try:
        for obj in oss2.ObjectIterator(bucket, prefix=folder_path):
            if not obj.key.endswith('/'):  # 排除目录本身
                files.append(obj.key)
        return sorted(files, key=lambda x: os.path.splitext(x)[0]) 
    except Exception as e:
        print(f"Error listing OSS files: {str(e)}")
        return []

def download_oss_file(oss_path: str, local_path: str):
    """从OSS下载文件到本地"""
    try:
        # 确保本地目录存在
        os.makedirs(os.path.dirname(local_path), exist_ok=True)
        bucket.get_object_to_file(oss_path, local_path)
    except Exception as e:
        print(f"Error downloading file {oss_path}: {str(e)}")
        raise

def oss_file_exists(oss_path: str) -> bool:
    """检查OSS文件是否存在"""
    try:
        return bucket.object_exists(oss_path)
    except Exception as e:
        print(f"Error checking if file exists in OSS: {str(e)}")
        return False

def get_user_tmp_dir(session_hash: str) -> str:
    """获取用户临时目录"""
    user_dir = os.path.join(TMP_ROOT, str(session_hash))
    os.makedirs(user_dir, exist_ok=True)
    return user_dir

def cleanup_user_tmp_dir(session_hash: str):
    """清理用户临时目录"""
    user_dir = os.path.join(TMP_ROOT, str(session_hash))
    if os.path.exists(user_dir):
        shutil.rmtree(user_dir)