370 lines
11 KiB
Python
370 lines
11 KiB
Python
|
import asyncio
|
|||
|
|
from typing import Optional, Dict, Any, BinaryIO, Union
|
|||
|
|
from pathlib import Path
|
|||
|
|
from loguru import logger
|
|||
|
|
|
|||
|
|
from config import config_manager
|
|||
|
|
|
|||
|
|
# 可选导入aioboto3
|
|||
|
|
try:
|
|||
|
|
import aioboto3
|
|||
|
|
from botocore.exceptions import ClientError, NoCredentialsError
|
|||
|
|
HAS_BOTO3 = True
|
|||
|
|
except ImportError:
|
|||
|
|
aioboto3 = None
|
|||
|
|
ClientError = Exception
|
|||
|
|
NoCredentialsError = Exception
|
|||
|
|
HAS_BOTO3 = False
|
|||
|
|
|
|||
|
|
|
|||
|
|
class S3Client:
|
|||
|
|
"""异步S3客户端"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self._session = None
|
|||
|
|
self._client = None
|
|||
|
|
self._config = None
|
|||
|
|
|
|||
|
|
if not HAS_BOTO3:
|
|||
|
|
logger.warning("aioboto3未安装，S3客户端功能不可用")
|
|||
|
|
|
|||
|
|
def _get_s3_config(self):
|
|||
|
|
"""获取S3配置"""
|
|||
|
|
if self._config is None:
|
|||
|
|
self._config = config_manager.get_settings().s3
|
|||
|
|
return self._config
|
|||
|
|
|
|||
|
|
async def _get_client(self):
|
|||
|
|
"""获取S3客户端"""
|
|||
|
|
if not HAS_BOTO3:
|
|||
|
|
raise RuntimeError("aioboto3未安装，无法使用S3客户端功能。请运行: pip install aioboto3")
|
|||
|
|
|
|||
|
|
if self._client is None:
|
|||
|
|
config = self._get_s3_config()
|
|||
|
|
|
|||
|
|
# 验证必要的配置
|
|||
|
|
if not config.access_key_id or not config.secret_access_key:
|
|||
|
|
raise ValueError("S3 access_key_id 和 secret_access_key 不能为空")
|
|||
|
|
|
|||
|
|
if not config.bucket_name:
|
|||
|
|
raise ValueError("S3 bucket_name 不能为空")
|
|||
|
|
|
|||
|
|
if self._session is None:
|
|||
|
|
self._session = aioboto3.Session()
|
|||
|
|
|
|||
|
|
self._client = self._session.client(
|
|||
|
|
's3',
|
|||
|
|
aws_access_key_id=config.access_key_id,
|
|||
|
|
aws_secret_access_key=config.secret_access_key,
|
|||
|
|
endpoint_url=config.endpoint_url,
|
|||
|
|
region_name=config.region_name,
|
|||
|
|
use_ssl=config.use_ssl,
|
|||
|
|
config=aioboto3.Config(signature_version=config.signature_version)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return self._client
|
|||
|
|
|
|||
|
|
async def upload_file(
|
|||
|
|
self,
|
|||
|
|
file_path: Union[str, Path],
|
|||
|
|
key: str,
|
|||
|
|
bucket: Optional[str] = None,
|
|||
|
|
extra_args: Optional[Dict[str, Any]] = None
|
|||
|
|
) -> bool:
|
|||
|
|
"""
|
|||
|
|
上传文件到S3
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
file_path: 本地文件路径
|
|||
|
|
key: S3对象键名
|
|||
|
|
bucket: 存储桶名称，如果为None则使用配置中的默认bucket
|
|||
|
|
extra_args: 额外参数，如metadata, ACL等
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
bool: 是否上传成功
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
config = self._get_s3_config()
|
|||
|
|
bucket = bucket or config.bucket_name
|
|||
|
|
|
|||
|
|
if not bucket:
|
|||
|
|
raise ValueError("bucket名称不能为空")
|
|||
|
|
|
|||
|
|
async with await self._get_client() as s3:
|
|||
|
|
await s3.upload_file(
|
|||
|
|
Filename=str(file_path),
|
|||
|
|
Bucket=bucket,
|
|||
|
|
Key=key,
|
|||
|
|
ExtraArgs=extra_args or {}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
logger.info(f"文件上传成功: {file_path} -> s3://{bucket}/{key}")
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
except FileNotFoundError:
|
|||
|
|
logger.error(f"文件不存在: {file_path}")
|
|||
|
|
return False
|
|||
|
|
except NoCredentialsError:
|
|||
|
|
logger.error("S3凭据未配置或无效")
|
|||
|
|
return False
|
|||
|
|
except ClientError as e:
|
|||
|
|
logger.error(f"S3客户端错误: {e}")
|
|||
|
|
return False
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"上传文件失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def download_file(
|
|||
|
|
self,
|
|||
|
|
key: str,
|
|||
|
|
file_path: Union[str, Path],
|
|||
|
|
bucket: Optional[str] = None
|
|||
|
|
) -> bool:
|
|||
|
|
"""
|
|||
|
|
从S3下载文件
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
key: S3对象键名
|
|||
|
|
file_path: 本地保存路径
|
|||
|
|
bucket: 存储桶名称，如果为None则使用配置中的默认bucket
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
bool: 是否下载成功
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
config = self._get_s3_config()
|
|||
|
|
bucket = bucket or config.bucket_name
|
|||
|
|
|
|||
|
|
if not bucket:
|
|||
|
|
raise ValueError("bucket名称不能为空")
|
|||
|
|
|
|||
|
|
# 确保目标目录存在
|
|||
|
|
Path(file_path).parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
async with await self._get_client() as s3:
|
|||
|
|
await s3.download_file(
|
|||
|
|
Bucket=bucket,
|
|||
|
|
Key=key,
|
|||
|
|
Filename=str(file_path)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
logger.info(f"文件下载成功: s3://{bucket}/{key} -> {file_path}")
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
except ClientError as e:
|
|||
|
|
if e.response['Error']['Code'] == 'NoSuchKey':
|
|||
|
|
logger.error(f"S3对象不存在: s3://{bucket}/{key}")
|
|||
|
|
else:
|
|||
|
|
logger.error(f"S3客户端错误: {e}")
|
|||
|
|
return False
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"下载文件失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def upload_bytes(
|
|||
|
|
self,
|
|||
|
|
data: bytes,
|
|||
|
|
key: str,
|
|||
|
|
bucket: Optional[str] = None,
|
|||
|
|
content_type: Optional[str] = None,
|
|||
|
|
extra_args: Optional[Dict[str, Any]] = None
|
|||
|
|
) -> bool:
|
|||
|
|
"""
|
|||
|
|
上传字节数据到S3
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
data: 要上传的字节数据
|
|||
|
|
key: S3对象键名
|
|||
|
|
bucket: 存储桶名称，如果为None则使用配置中的默认bucket
|
|||
|
|
content_type: 内容类型
|
|||
|
|
extra_args: 额外参数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
bool: 是否上传成功
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
config = self._get_s3_config()
|
|||
|
|
bucket = bucket or config.bucket_name
|
|||
|
|
|
|||
|
|
if not bucket:
|
|||
|
|
raise ValueError("bucket名称不能为空")
|
|||
|
|
|
|||
|
|
extra_args = extra_args or {}
|
|||
|
|
if content_type:
|
|||
|
|
extra_args['ContentType'] = content_type
|
|||
|
|
|
|||
|
|
async with await self._get_client() as s3:
|
|||
|
|
await s3.put_object(
|
|||
|
|
Bucket=bucket,
|
|||
|
|
Key=key,
|
|||
|
|
Body=data,
|
|||
|
|
**extra_args
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
logger.info(f"数据上传成功: s3://{bucket}/{key}")
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"上传数据失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def download_bytes(
|
|||
|
|
self,
|
|||
|
|
key: str,
|
|||
|
|
bucket: Optional[str] = None
|
|||
|
|
) -> Optional[bytes]:
|
|||
|
|
"""
|
|||
|
|
从S3下载字节数据
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
key: S3对象键名
|
|||
|
|
bucket: 存储桶名称，如果为None则使用配置中的默认bucket
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Optional[bytes]: 下载的字节数据，失败时返回None
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
config = self._get_s3_config()
|
|||
|
|
bucket = bucket or config.bucket_name
|
|||
|
|
|
|||
|
|
if not bucket:
|
|||
|
|
raise ValueError("bucket名称不能为空")
|
|||
|
|
|
|||
|
|
async with await self._get_client() as s3:
|
|||
|
|
response = await s3.get_object(Bucket=bucket, Key=key)
|
|||
|
|
async with response['Body'] as stream:
|
|||
|
|
data = await stream.read()
|
|||
|
|
|
|||
|
|
logger.info(f"数据下载成功: s3://{bucket}/{key}")
|
|||
|
|
return data
|
|||
|
|
|
|||
|
|
except ClientError as e:
|
|||
|
|
if e.response['Error']['Code'] == 'NoSuchKey':
|
|||
|
|
logger.error(f"S3对象不存在: s3://{bucket}/{key}")
|
|||
|
|
else:
|
|||
|
|
logger.error(f"S3客户端错误: {e}")
|
|||
|
|
return None
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"下载数据失败: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
async def delete_object(
|
|||
|
|
self,
|
|||
|
|
key: str,
|
|||
|
|
bucket: Optional[str] = None
|
|||
|
|
) -> bool:
|
|||
|
|
"""
|
|||
|
|
删除S3对象
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
key: S3对象键名
|
|||
|
|
bucket: 存储桶名称，如果为None则使用配置中的默认bucket
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
bool: 是否删除成功
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
config = self._get_s3_config()
|
|||
|
|
bucket = bucket or config.bucket_name
|
|||
|
|
|
|||
|
|
if not bucket:
|
|||
|
|
raise ValueError("bucket名称不能为空")
|
|||
|
|
|
|||
|
|
async with await self._get_client() as s3:
|
|||
|
|
await s3.delete_object(Bucket=bucket, Key=key)
|
|||
|
|
|
|||
|
|
logger.info(f"对象删除成功: s3://{bucket}/{key}")
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"删除对象失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def list_objects(
|
|||
|
|
self,
|
|||
|
|
prefix: str = "",
|
|||
|
|
bucket: Optional[str] = None,
|
|||
|
|
max_keys: int = 1000
|
|||
|
|
) -> list:
|
|||
|
|
"""
|
|||
|
|
列出S3对象
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
prefix: 对象键前缀
|
|||
|
|
bucket: 存储桶名称，如果为None则使用配置中的默认bucket
|
|||
|
|
max_keys: 最大返回对象数量
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
list: 对象列表
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
config = self._get_s3_config()
|
|||
|
|
bucket = bucket or config.bucket_name
|
|||
|
|
|
|||
|
|
if not bucket:
|
|||
|
|
raise ValueError("bucket名称不能为空")
|
|||
|
|
|
|||
|
|
async with await self._get_client() as s3:
|
|||
|
|
response = await s3.list_objects_v2(
|
|||
|
|
Bucket=bucket,
|
|||
|
|
Prefix=prefix,
|
|||
|
|
MaxKeys=max_keys
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
objects = response.get('Contents', [])
|
|||
|
|
logger.info(f"列出对象成功: s3://{bucket}/{prefix}* ({len(objects)}个对象)")
|
|||
|
|
return objects
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"列出对象失败: {e}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
async def object_exists(
|
|||
|
|
self,
|
|||
|
|
key: str,
|
|||
|
|
bucket: Optional[str] = None
|
|||
|
|
) -> bool:
|
|||
|
|
"""
|
|||
|
|
检查S3对象是否存在
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
key: S3对象键名
|
|||
|
|
bucket: 存储桶名称，如果为None则使用配置中的默认bucket
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
bool: 对象是否存在
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
config = self._get_s3_config()
|
|||
|
|
bucket = bucket or config.bucket_name
|
|||
|
|
|
|||
|
|
if not bucket:
|
|||
|
|
raise ValueError("bucket名称不能为空")
|
|||
|
|
|
|||
|
|
async with await self._get_client() as s3:
|
|||
|
|
await s3.head_object(Bucket=bucket, Key=key)
|
|||
|
|
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
except ClientError as e:
|
|||
|
|
if e.response['Error']['Code'] == '404':
|
|||
|
|
return False
|
|||
|
|
else:
|
|||
|
|
logger.error(f"检查对象存在性失败: {e}")
|
|||
|
|
return False
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"检查对象存在性失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def close(self):
|
|||
|
|
"""关闭S3客户端"""
|
|||
|
|
if self._client:
|
|||
|
|
await self._client.close()
|
|||
|
|
self._client = None
|
|||
|
|
if self._session:
|
|||
|
|
await self._session.close()
|
|||
|
|
self._session = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局S3客户端实例
|
|||
|
|
s3_client = S3Client()
|