545 lines
18 KiB
Python
545 lines
18 KiB
Python
"""
|
||
腾讯云COS存储提供者实现
|
||
|
||
本模块实现了腾讯云COS的具体存储操作，继承自抽象存储接口。
|
||
提供完整的COS文件上传、下载、删除等功能。
|
||
|
||
特性:
|
||
- 支持文件、字节数据、PyTorch张量的上传
|
||
- 完整的错误处理和重试机制
|
||
- 统一的日志记录
|
||
- 与S3接口保持一致的API设计
|
||
"""
|
||
|
||
import os
|
||
from typing import Optional, Dict, Any, List
|
||
import torch
|
||
import loguru
|
||
from qcloud_cos import CosConfig, CosS3Client, CosClientError, CosServiceError
|
||
from ..storage_interface import (
|
||
StorageProvider,
|
||
StorageFactory,
|
||
UploadResult,
|
||
DownloadResult,
|
||
)
|
||
from ...config_utils import config
|
||
|
||
|
||
class COSStorageProvider(StorageProvider):
|
||
"""
|
||
腾讯云COS存储提供者实现
|
||
|
||
实现了StorageProvider接口的所有方法，提供完整的COS存储功能。
|
||
采用懒加载模式初始化COS客户端，提高性能。
|
||
"""
|
||
|
||
def __init__(self, config: Dict[str, Any]):
|
||
"""
|
||
初始化COS存储提供者
|
||
|
||
Args:
|
||
config: COS配置字典，必须包含secret_id、secret_key和region
|
||
"""
|
||
super().__init__(config)
|
||
self.bucket_name = config.get("bucket_name", "bwkj-cos-1324682537")
|
||
self.region = config["region"]
|
||
self.secret_id = config["secret_id"]
|
||
self.secret_key = config["secret_key"]
|
||
self._client = None
|
||
|
||
def _validate_config(self) -> None:
|
||
"""
|
||
验证COS配置的完整性
|
||
|
||
Raises:
|
||
ValueError: 配置信息缺失时抛出异常
|
||
"""
|
||
required_keys = ["secret_id", "secret_key", "region"]
|
||
missing_keys = [key for key in required_keys if not self.config.get(key)]
|
||
|
||
if missing_keys:
|
||
raise ValueError(
|
||
f"COS配置缺失必要参数: {missing_keys}. " f"请检查配置文件或环境变量"
|
||
)
|
||
|
||
@property
|
||
def client(self):
|
||
"""
|
||
获取COS客户端实例（懒加载模式）
|
||
|
||
Returns:
|
||
CosS3Client: COS客户端实例
|
||
"""
|
||
if self._client is None:
|
||
try:
|
||
cos_config = CosConfig(
|
||
Region=self.region,
|
||
SecretId=self.secret_id,
|
||
SecretKey=self.secret_key,
|
||
)
|
||
self._client = CosS3Client(cos_config)
|
||
loguru.logger.info(f"COS客户端初始化成功，区域: {self.region}")
|
||
except Exception as e:
|
||
loguru.logger.error(f"COS客户端初始化失败: {e}")
|
||
raise
|
||
|
||
return self._client
|
||
|
||
def upload_file(
|
||
self,
|
||
local_path: str,
|
||
remote_key: str,
|
||
content_type: Optional[str] = None,
|
||
**kwargs,
|
||
) -> UploadResult:
|
||
"""
|
||
上传本地文件到COS
|
||
|
||
Args:
|
||
local_path: 本地文件路径
|
||
remote_key: COS中的键名
|
||
content_type: 文件内容类型
|
||
**kwargs: 额外的上传参数
|
||
|
||
Returns:
|
||
UploadResult: 上传操作结果
|
||
"""
|
||
try:
|
||
if not os.path.exists(local_path):
|
||
error_msg = f"本地文件不存在: {local_path}"
|
||
loguru.logger.error(error_msg)
|
||
return UploadResult(
|
||
success=False,
|
||
key=remote_key,
|
||
message=error_msg,
|
||
error=FileNotFoundError(error_msg),
|
||
)
|
||
|
||
file_size = os.path.getsize(local_path)
|
||
loguru.logger.info(
|
||
f"开始上传文件到COS: {local_path} -> cos://{self.bucket_name}/{remote_key} "
|
||
f"({file_size} bytes)"
|
||
)
|
||
|
||
# 执行上传操作，包含重试机制
|
||
for attempt in range(3): # 最多重试3次
|
||
try:
|
||
response = self.client.upload_file(
|
||
Bucket=self.bucket_name,
|
||
Key=remote_key,
|
||
LocalFilePath=local_path,
|
||
**kwargs,
|
||
)
|
||
break
|
||
except (CosClientError, CosServiceError) as e:
|
||
if attempt == 2: # 最后一次尝试失败
|
||
raise e
|
||
loguru.logger.warning(
|
||
f"COS上传尝试 {attempt + 1} 失败，重试中: {e}"
|
||
)
|
||
|
||
success_msg = f"文件上传成功: cos://{self.bucket_name}/{remote_key}"
|
||
loguru.logger.info(success_msg)
|
||
|
||
return UploadResult(
|
||
success=True, key=remote_key, size=file_size, message=success_msg
|
||
)
|
||
|
||
except Exception as e:
|
||
error_msg = f"COS文件上传失败: {str(e)}"
|
||
loguru.logger.error(f"{error_msg} (文件: {local_path}, 键: {remote_key})")
|
||
return UploadResult(
|
||
success=False, key=remote_key, message=error_msg, error=e
|
||
)
|
||
|
||
def upload_bytes(
|
||
self, data: bytes, remote_key: str, content_type: Optional[str] = None, **kwargs
|
||
) -> UploadResult:
|
||
"""
|
||
上传字节数据到COS
|
||
|
||
Args:
|
||
data: 字节数据
|
||
remote_key: COS中的键名
|
||
content_type: 文件内容类型
|
||
**kwargs: 额外的上传参数
|
||
|
||
Returns:
|
||
UploadResult: 上传操作结果
|
||
"""
|
||
try:
|
||
loguru.logger.info(
|
||
f"开始上传字节数据到COS: cos://{self.bucket_name}/{remote_key} "
|
||
f"({len(data)} bytes)"
|
||
)
|
||
|
||
# 执行上传操作，包含重试机制
|
||
for attempt in range(3):
|
||
try:
|
||
response = self.client.put_object(
|
||
Bucket=self.bucket_name,
|
||
Key=remote_key,
|
||
Body=data,
|
||
ContentType=content_type,
|
||
**kwargs,
|
||
)
|
||
break
|
||
except (CosClientError, CosServiceError) as e:
|
||
if attempt == 2:
|
||
raise e
|
||
loguru.logger.warning(
|
||
f"COS字节数据上传尝试 {attempt + 1} 失败，重试中: {e}"
|
||
)
|
||
|
||
success_msg = f"字节数据上传成功: cos://{self.bucket_name}/{remote_key}"
|
||
loguru.logger.info(success_msg)
|
||
|
||
return UploadResult(
|
||
success=True, key=remote_key, size=len(data), message=success_msg
|
||
)
|
||
|
||
except Exception as e:
|
||
error_msg = f"COS字节数据上传失败: {str(e)}"
|
||
loguru.logger.error(
|
||
f"{error_msg} (键: {remote_key}, 大小: {len(data)} bytes)"
|
||
)
|
||
return UploadResult(
|
||
success=False, key=remote_key, message=error_msg, error=e
|
||
)
|
||
|
||
def upload_tensor(
|
||
self, tensor: torch.Tensor, remote_key: str, format: str = "PNG", **kwargs
|
||
) -> UploadResult:
|
||
"""
|
||
上传PyTorch张量作为图像到COS
|
||
|
||
Args:
|
||
tensor: PyTorch张量
|
||
remote_key: COS中的键名
|
||
format: 图像格式（PNG, JPEG等）
|
||
**kwargs: 额外的上传参数
|
||
|
||
Returns:
|
||
UploadResult: 上传操作结果
|
||
"""
|
||
try:
|
||
from ...image_utils import tensor_to_tempfile
|
||
|
||
loguru.logger.info(
|
||
f"开始上传张量到COS: cos://{self.bucket_name}/{remote_key} "
|
||
f"(形状: {tensor.shape}, 格式: {format})"
|
||
)
|
||
|
||
# 将张量转换为临时文件
|
||
temp_file = tensor_to_tempfile(tensor, format=format)
|
||
temp_path = temp_file.name
|
||
|
||
try:
|
||
# 设置内容类型
|
||
content_type = f"image/{format.lower()}"
|
||
if format.upper() == "JPEG":
|
||
content_type = "image/jpeg"
|
||
|
||
# 上传临时文件
|
||
result = self.upload_file(
|
||
temp_path, remote_key, content_type=content_type, **kwargs
|
||
)
|
||
|
||
if result.success:
|
||
loguru.logger.info(
|
||
f"张量上传成功: cos://{self.bucket_name}/{remote_key}"
|
||
)
|
||
|
||
return result
|
||
|
||
finally:
|
||
# 清理临时文件
|
||
if os.path.exists(temp_path):
|
||
os.unlink(temp_path)
|
||
loguru.logger.debug(f"临时文件已清理: {temp_path}")
|
||
|
||
except Exception as e:
|
||
error_msg = f"COS张量上传失败: {str(e)}"
|
||
loguru.logger.error(
|
||
f"{error_msg} (键: {remote_key}, 张量形状: {tensor.shape})"
|
||
)
|
||
return UploadResult(
|
||
success=False, key=remote_key, message=error_msg, error=e
|
||
)
|
||
|
||
def download_file(
|
||
self, remote_key: str, local_path: str, **kwargs
|
||
) -> DownloadResult:
|
||
"""
|
||
从COS下载文件到本地
|
||
|
||
Args:
|
||
remote_key: COS中的键名
|
||
local_path: 本地保存路径
|
||
**kwargs: 额外的下载参数
|
||
|
||
Returns:
|
||
DownloadResult: 下载操作结果
|
||
"""
|
||
try:
|
||
loguru.logger.info(
|
||
f"开始从COS下载文件: cos://{self.bucket_name}/{remote_key} -> {local_path}"
|
||
)
|
||
|
||
# 确保本地目录存在
|
||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||
|
||
# 检查文件是否存在
|
||
if not self.file_exists(remote_key):
|
||
error_msg = f"COS文件不存在: cos://{self.bucket_name}/{remote_key}"
|
||
loguru.logger.error(error_msg)
|
||
return DownloadResult(
|
||
success=False, message=error_msg, error=FileNotFoundError(error_msg)
|
||
)
|
||
|
||
# 执行下载操作，包含重试机制
|
||
for attempt in range(3):
|
||
try:
|
||
response = self.client.download_file(
|
||
Bucket=self.bucket_name,
|
||
Key=remote_key,
|
||
DestFilePath=local_path,
|
||
**kwargs,
|
||
)
|
||
break
|
||
except (CosClientError, CosServiceError) as e:
|
||
if attempt == 2:
|
||
raise e
|
||
loguru.logger.warning(
|
||
f"COS下载尝试 {attempt + 1} 失败，重试中: {e}"
|
||
)
|
||
|
||
success_msg = f"文件下载成功: {local_path}"
|
||
loguru.logger.info(success_msg)
|
||
|
||
return DownloadResult(
|
||
success=True, local_path=local_path, message=success_msg
|
||
)
|
||
|
||
except Exception as e:
|
||
error_msg = f"COS文件下载失败: {str(e)}"
|
||
loguru.logger.error(
|
||
f"{error_msg} (键: {remote_key}, 本地路径: {local_path})"
|
||
)
|
||
return DownloadResult(success=False, message=error_msg, error=e)
|
||
|
||
def download_bytes(self, remote_key: str, **kwargs) -> DownloadResult:
|
||
"""
|
||
从COS下载文件为字节数据
|
||
|
||
Args:
|
||
remote_key: COS中的键名
|
||
**kwargs: 额外的下载参数
|
||
|
||
Returns:
|
||
DownloadResult: 下载操作结果，数据包含在data字段中
|
||
"""
|
||
try:
|
||
loguru.logger.info(
|
||
f"开始从COS下载字节数据: cos://{self.bucket_name}/{remote_key}"
|
||
)
|
||
|
||
# 检查文件是否存在
|
||
if not self.file_exists(remote_key):
|
||
error_msg = f"COS文件不存在: cos://{self.bucket_name}/{remote_key}"
|
||
loguru.logger.error(error_msg)
|
||
return DownloadResult(
|
||
success=False, message=error_msg, error=FileNotFoundError(error_msg)
|
||
)
|
||
|
||
# 执行下载操作，包含重试机制
|
||
for attempt in range(3):
|
||
try:
|
||
response = self.client.get_object(
|
||
Bucket=self.bucket_name, Key=remote_key, **kwargs
|
||
)
|
||
data = response["Body"].read()
|
||
break
|
||
except (CosClientError, CosServiceError) as e:
|
||
if attempt == 2:
|
||
raise e
|
||
loguru.logger.warning(
|
||
f"COS字节数据下载尝试 {attempt + 1} 失败，重试中: {e}"
|
||
)
|
||
|
||
success_msg = f"字节数据下载成功: {len(data)} bytes"
|
||
loguru.logger.info(success_msg)
|
||
|
||
return DownloadResult(success=True, data=data, message=success_msg)
|
||
|
||
except Exception as e:
|
||
error_msg = f"COS字节数据下载失败: {str(e)}"
|
||
loguru.logger.error(f"{error_msg} (键: {remote_key})")
|
||
return DownloadResult(success=False, message=error_msg, error=e)
|
||
|
||
def delete_file(self, remote_key: str, **kwargs) -> bool:
|
||
"""
|
||
删除COS中的文件
|
||
|
||
Args:
|
||
remote_key: COS中的键名
|
||
**kwargs: 额外的删除参数
|
||
|
||
Returns:
|
||
bool: 删除是否成功
|
||
"""
|
||
try:
|
||
loguru.logger.info(f"删除COS文件: cos://{self.bucket_name}/{remote_key}")
|
||
|
||
# 执行删除操作，包含重试机制
|
||
for attempt in range(3):
|
||
try:
|
||
self.client.delete_object(
|
||
Bucket=self.bucket_name, Key=remote_key, **kwargs
|
||
)
|
||
break
|
||
except (CosClientError, CosServiceError) as e:
|
||
if attempt == 2:
|
||
raise e
|
||
loguru.logger.warning(
|
||
f"COS删除尝试 {attempt + 1} 失败，重试中: {e}"
|
||
)
|
||
|
||
loguru.logger.info(f"文件删除成功: cos://{self.bucket_name}/{remote_key}")
|
||
return True
|
||
|
||
except Exception as e:
|
||
loguru.logger.error(f"COS文件删除失败: {e} (键: {remote_key})")
|
||
return False
|
||
|
||
def file_exists(self, remote_key: str, **kwargs) -> bool:
|
||
"""
|
||
检查COS中文件是否存在
|
||
|
||
Args:
|
||
remote_key: COS中的键名
|
||
**kwargs: 额外的检查参数
|
||
|
||
Returns:
|
||
bool: 文件是否存在
|
||
"""
|
||
try:
|
||
self.client.head_object(Bucket=self.bucket_name, Key=remote_key, **kwargs)
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
def list_files(
|
||
self, prefix: str = "", max_keys: int = 1000, **kwargs
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
列出COS中的文件
|
||
|
||
Args:
|
||
prefix: 文件前缀过滤
|
||
max_keys: 最大返回数量
|
||
**kwargs: 额外的列表参数
|
||
|
||
Returns:
|
||
List[Dict[str, Any]]: 文件信息列表
|
||
"""
|
||
try:
|
||
loguru.logger.info(
|
||
f"列出COS文件: cos://{self.bucket_name}/{prefix} (最大: {max_keys})"
|
||
)
|
||
|
||
response = self.client.list_objects(
|
||
Bucket=self.bucket_name, Prefix=prefix, MaxKeys=max_keys, **kwargs
|
||
)
|
||
|
||
files = []
|
||
if "Contents" in response:
|
||
for obj in response["Contents"]:
|
||
files.append(
|
||
{
|
||
"key": obj["Key"],
|
||
"size": obj["Size"],
|
||
"last_modified": obj["LastModified"],
|
||
"etag": obj.get("ETag", "").strip('"'),
|
||
"storage_class": obj.get("StorageClass", "STANDARD"),
|
||
}
|
||
)
|
||
|
||
loguru.logger.info(f"找到 {len(files)} 个文件")
|
||
return files
|
||
|
||
except Exception as e:
|
||
loguru.logger.error(f"COS文件列表获取失败: {e} (前缀: {prefix})")
|
||
return []
|
||
|
||
def get_file_url(self, remote_key: str, expires_in: int = 3600, **kwargs) -> str:
|
||
"""
|
||
获取COS文件的访问URL
|
||
|
||
Args:
|
||
remote_key: COS中的键名
|
||
expires_in: URL过期时间（秒）
|
||
**kwargs: 额外的URL生成参数
|
||
|
||
Returns:
|
||
str: 文件访问URL
|
||
"""
|
||
try:
|
||
# 生成预签名URL
|
||
url = self.client.get_presigned_url(
|
||
Method="GET",
|
||
Bucket=self.bucket_name,
|
||
Key=remote_key,
|
||
Expired=expires_in,
|
||
**kwargs,
|
||
)
|
||
|
||
return url
|
||
|
||
except Exception as e:
|
||
loguru.logger.error(f"COS URL生成失败: {e} (键: {remote_key})")
|
||
raise
|
||
|
||
|
||
class COSStorageFactory(StorageFactory):
|
||
"""
|
||
COS存储工厂实现
|
||
|
||
负责创建COS存储提供者实例，处理配置验证和初始化。
|
||
"""
|
||
|
||
def create_provider(self, config_dict: Dict[str, Any]) -> StorageProvider:
|
||
"""
|
||
创建COS存储提供者实例
|
||
|
||
Args:
|
||
config_dict: COS配置字典
|
||
|
||
Returns:
|
||
StorageProvider: COS存储提供者实例
|
||
"""
|
||
# 如果没有提供配置，尝试从全局配置获取
|
||
if not config_dict or not config_dict.get("secret_id"):
|
||
if config.has_cos_config():
|
||
cos_config = config.get_cos_config()
|
||
config_dict = {
|
||
"secret_id": cos_config["secret_id"],
|
||
"secret_key": cos_config["secret_key"],
|
||
"region": cos_config["region"],
|
||
"bucket_name": cos_config.get("bucket_name"),
|
||
**config_dict,
|
||
}
|
||
else:
|
||
raise ValueError("未提供有效的COS配置，且全局配置中也没有COS配置")
|
||
|
||
return COSStorageProvider(config_dict)
|
||
|
||
def get_supported_types(self) -> List[str]:
|
||
"""
|
||
获取支持的存储类型列表
|
||
|
||
Returns:
|
||
List[str]: 支持的存储类型
|
||
"""
|
||
return ["cos", "qcloud", "tencent"]
|