import os
import oss2
from oss2.credentials import StaticCredentialsProvider
from pathlib import Path
from typing import Optional, Dict, Any, List, Tuple
import concurrent.futures
from threading import Lock
import time
endpoint = "https://oss-cn-beijing.aliyuncs.com"
region = "cn-beijing"
bucket_name = "image-browser"
env_access_key_id = "oss_image_browser_access_key_id"
env_access_key_secret = "oss_image_browser_access_key_secret"
class AsyncOSSUploader:
def __init__(self, workers: int = 10):
"""
初始化异步OSS上传器
Args:
workers: 最大并发工作线程数
"""
provider = StaticCredentialsProvider(
access_key_id=os.getenv(env_access_key_id),
access_key_secret=os.getenv(env_access_key_secret)
)
auth = oss2.ProviderAuthV4(provider)
self.bucket = oss2.Bucket(auth, endpoint, bucket_name, region=region)
self.workers = workers
self.lock = Lock()
self.uploaded_count = 0
self.failed_count = 0
def get_file_headers(self, file_path: str) -> Dict[str, str]:
"""根据文件类型获取合适的HTTP头"""
headers = {
'x-oss-storage-class': 'Standard'
}
# 根据文件扩展名设置Content-Type
ext = Path(file_path).suffix.lower()
content_types = {
'.html': 'text/html',
'.css': 'text/css',
'.js': 'application/javascript',
'.json': 'application/json',
'.png': 'image/png',
'.jpg': 'image/jpeg',
'.jpeg': 'image/jpeg',
'.gif': 'image/gif',
'.svg': 'image/svg+xml',
'.ico': 'image/x-icon',
'.webp': 'image/webp',
'.txt': 'text/plain',
'.md': 'text/markdown',
'.xml': 'application/xml',
'.pdf': 'application/pdf',
'.zip': 'application/zip',
'.woff': 'font/woff',
'.woff2': 'font/woff2',
'.ttf': 'font/ttf',
'.eot': 'application/vnd.ms-fontobject'
}
if ext in content_types:
headers['Content-Type'] = content_types[ext]
if Path(file_path).name == 'index.html':
headers['Cache-Control'] = 'no-cache'
return headers
def _upload_single_file(self, file_path: Path, oss_key: str) -> bool:
"""上传单个文件(线程安全)"""
try:
headers = self.get_file_headers(str(file_path))
with open(file_path, 'rb') as f:
self.bucket.put_object(oss_key, f, headers=headers)
with self.lock:
self.uploaded_count += 1
print(f"✅ 上传成功 ({self.uploaded_count}): {file_path.name} -> {oss_key}")
return True
except Exception as e:
with self.lock:
self.failed_count += 1
print(f"❌ 上传失败 ({self.failed_count}): {file_path.name} -> {e}")
return False
def upload(self, src_path: str, dst_path: str, include_root: bool = False) -> bool:
"""
上传文件或目录到OSS
Args:
src_path: 本地源路径(文件或目录)
dst_path: OSS目标路径
include_root:
- True: 整个目录上传到dst_path目录中 (类似 cp src dst)
- False: 目录内容上传到dst_path目录中 (类似 cp -rf src/* dst)
Returns:
bool: 上传是否成功
"""
src_path = Path(src_path)
dst_path = dst_path.rstrip('/')
if not src_path.exists():
print(f"❌ 源路径不存在: {src_path}")
return False
start_time = time.time()
try:
if src_path.is_file():
# 上传单个文件
return self._upload_single_file_sync(src_path, dst_path)
elif src_path.is_dir():
# 上传目录
return self._upload_directory(src_path, dst_path, include_root)
else:
print(f"❌ 不支持的路径类型: {src_path}")
return False
except Exception as e:
print(f"❌ 上传过程中发生错误: {e}")
return False
finally:
elapsed_time = time.time() - start_time
print(f"⏱️ 总耗时: {elapsed_time:.2f}秒")
def _upload_single_file_sync(self, file_path: Path, dst_path: str) -> bool:
"""同步上传单个文件"""
try:
if dst_path is None:
dst_path = ""
if dst_path.endswith('/') or not dst_path:
if dst_path:
oss_key = f"{dst_path}/{file_path.name}"
else:
oss_key = file_path.name
else:
oss_key = dst_path
oss_key = oss_key.lstrip('/')
print(f"📤 上传文件: {file_path} -> {oss_key}")
return self._upload_single_file(file_path, oss_key)
except Exception as e:
print(f"❌ 文件上传失败 {file_path}: {e}")
return False
def _upload_directory(self, dir_path: Path, dst_path: str, include_root: bool) -> bool:
"""上传目录"""
print(f"📁 开始上传目录: {dir_path} -> {dst_path}")
print(f"📋 模式: {'包含根目录' if include_root else '仅内容'}")
print(f"🚀 并发上传 (最大 {self.workers} 线程)")
# 收集所有需要上传的文件
upload_tasks = []
for file_path in dir_path.rglob('*'):
if file_path.is_file():
if include_root:
relative_path = str(file_path)
else:
relative_path = file_path.relative_to(dir_path)
relative_path = str(relative_path)
oss_key = f"{dst_path}/{relative_path}".replace("\\", "/")
oss_key = oss_key.lstrip('/')
upload_tasks.append((file_path, oss_key))
total_files = len(upload_tasks)
print(f"📊 发现 {total_files} 个文件需要上传")
if total_files == 0:
print("ℹ️ 没有文件需要上传")
return True
if total_files == 1:
# 单个文件直接上传
file_path, oss_key = upload_tasks[0]
return self._upload_single_file(file_path, oss_key)
else:
# 多个文件使用并发上传
return self._upload_concurrent(upload_tasks)
def _upload_concurrent(self, upload_tasks: List[Tuple[Path, str]]) -> bool:
"""并发上传文件"""
print("⚡ 使用并发上传模式")
with concurrent.futures.ThreadPoolExecutor(max_workers=self.workers) as executor:
# 提交所有上传任务
future_to_task = {
executor.submit(self._upload_single_file, file_path, oss_key): (file_path, oss_key)
for file_path, oss_key in upload_tasks
}
# 等待所有任务完成
for future in concurrent.futures.as_completed(future_to_task):
file_path, oss_key = future_to_task[future]
try:
future.result() # 获取结果,如果有异常会在这里抛出
except Exception as e:
print(f"❌ 任务异常 {file_path}: {e}")
print(f"📊 并发上传完成: 成功 {self.uploaded_count} 个文件,失败 {self.failed_count} 个文件")
return self.failed_count == 0
def upload(src_path: str, dst_path: str, include_root: bool = False, workers: int = 10) -> bool:
"""
便捷的上传函数
Args:
src_path: 本地源路径
dst_path: OSS目标路径
include_root: 是否包含根目录
workers: 最大并发工作线程数
Returns:
bool: 上传是否成功
"""
uploader = AsyncOSSUploader(workers=workers)
return uploader.upload(src_path, dst_path, include_root)
if __name__ == "__main__":
upload('node_modules', '/da1', include_root=True, workers=10)