feat: s3适配切片上传

This commit is contained in:
Lan
2025-12-01 15:13:33 +08:00
parent a3f1fd6a4c
commit 3ff4e9b6cf
6 changed files with 549 additions and 165 deletions
+1
View File
@@ -51,6 +51,7 @@ class UploadChunk(models.Model):
file_name = fields.CharField(max_length=255)
created_at = fields.DatetimeField(auto_now_add=True)
completed = fields.BooleanField(default=False)
actual_size = fields.IntField(null=True) # 实际接收的分片大小
class KeyValue(Model):
+179 -37
View File
@@ -164,6 +164,16 @@ chunk_api = APIRouter(prefix="/chunk", tags=["切片"])
@chunk_api.post("/upload/init/", dependencies=[Depends(share_required_login)])
async def init_chunk_upload(data: InitChunkUploadModel):
# 服务端校验:根据 total_chunks * chunk_size 计算理论最大上传量
total_chunks = (data.file_size + data.chunk_size - 1) // data.chunk_size
max_possible_size = total_chunks * data.chunk_size
if max_possible_size > settings.uploadSize:
max_size_mb = settings.uploadSize / (1024 * 1024)
raise HTTPException(
status_code=403,
detail=f"文件大小超过限制,最大为 {max_size_mb:.2f} MB"
)
# # 秒传检查
# existing = await FileCodes.filter(file_hash=data.file_hash).first()
# if existing:
@@ -179,9 +189,30 @@ async def init_chunk_upload(data: InitChunkUploadModel):
# "name": f'{existing.prefix}{existing.suffix}'
# })
# 创建上传会话
# 断点续传:检查是否存在相同文件的未完成上传会话
existing_session = await UploadChunk.filter(
chunk_hash=data.file_hash,
chunk_index=-1,
file_size=data.file_size,
file_name=data.file_name,
).first()
if existing_session:
# 复用已有会话,获取已上传的分片列表
uploaded_chunks = await UploadChunk.filter(
upload_id=existing_session.upload_id,
completed=True
).values_list('chunk_index', flat=True)
return APIResponse(detail={
"existed": False,
"upload_id": existing_session.upload_id,
"chunk_size": existing_session.chunk_size,
"total_chunks": existing_session.total_chunks,
"uploaded_chunks": list(uploaded_chunks)
})
# 创建新的上传会话
upload_id = uuid.uuid4().hex
total_chunks = (data.file_size + data.chunk_size - 1) // data.chunk_size
await UploadChunk.create(
upload_id=upload_id,
chunk_index=-1,
@@ -191,17 +222,12 @@ async def init_chunk_upload(data: InitChunkUploadModel):
chunk_hash=data.file_hash,
file_name=data.file_name,
)
# 获取已上传的分片列表
uploaded_chunks = await UploadChunk.filter(
upload_id=upload_id,
completed=True
).values_list('chunk_index', flat=True)
return APIResponse(detail={
"existed": False,
"upload_id": upload_id,
"chunk_size": data.chunk_size,
"total_chunks": total_chunks,
"uploaded_chunks": uploaded_chunks
"uploaded_chunks": []
})
@@ -220,11 +246,52 @@ async def upload_chunk(
if chunk_index < 0 or chunk_index >= chunk_info.total_chunks:
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="无效的分片索引")
# 检查是否已上传(支持断点续传)
existing_chunk = await UploadChunk.filter(
upload_id=upload_id,
chunk_index=chunk_index,
completed=True
).first()
if existing_chunk:
return APIResponse(detail={"chunk_hash": existing_chunk.chunk_hash, "skipped": True})
# 读取分片数据并计算哈希
chunk_data = await chunk.read()
chunk_size = len(chunk_data)
# 校验分片大小不超过声明的 chunk_size
if chunk_size > chunk_info.chunk_size:
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
detail=f"分片大小超过声明值: 最大 {chunk_info.chunk_size}, 实际 {chunk_size}"
)
# 计算已上传的实际总大小(防止绕过前端限制)
uploaded_chunks_list = await UploadChunk.filter(
upload_id=upload_id,
completed=True
).all()
uploaded_size = sum(c.actual_size for c in uploaded_chunks_list if c.actual_size)
if uploaded_size + chunk_size > settings.uploadSize:
max_size_mb = settings.uploadSize / (1024 * 1024)
raise HTTPException(
status_code=403,
detail=f"累计上传大小超过限制,最大为 {max_size_mb:.2f} MB"
)
chunk_hash = hashlib.sha256(chunk_data).hexdigest()
# 更新或创建分片记录
# 获取文件路径
_, _, _, _, save_path = await get_chunk_file_path_name(chunk_info.file_name, upload_id)
# 保存分片到存储
storage = storages[settings.file_storage]()
try:
await storage.save_chunk(upload_id, chunk_index, chunk_data, chunk_hash, save_path)
except Exception as e:
raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"分片保存失败: {str(e)}")
# 更新或创建分片记录(保存成功后再记录,包含实际大小)
await UploadChunk.update_or_create(
upload_id=upload_id,
chunk_index=chunk_index,
@@ -234,15 +301,59 @@ async def upload_chunk(
'file_size': chunk_info.file_size,
'total_chunks': chunk_info.total_chunks,
'chunk_size': chunk_info.chunk_size,
'file_name': chunk_info.file_name
'file_name': chunk_info.file_name,
'actual_size': chunk_size
}
)
return APIResponse(detail={"chunk_hash": chunk_hash})
@chunk_api.delete("/upload/{upload_id}", dependencies=[Depends(share_required_login)])
async def cancel_upload(upload_id: str):
"""取消上传并清理临时文件"""
chunk_info = await UploadChunk.filter(upload_id=upload_id, chunk_index=-1).first()
if not chunk_info:
raise HTTPException(status.HTTP_404_NOT_FOUND, detail="上传会话不存在")
# 获取文件路径
_, _, _, _, save_path = await get_chunk_file_path_name(chunk_info.file_name, upload_id)
# 保存分片到存储
# 清理存储中的临时文件
storage = storages[settings.file_storage]()
await storage.save_chunk(upload_id, chunk_index, chunk_data, chunk_hash, save_path)
return APIResponse(detail={"chunk_hash": chunk_hash})
try:
await storage.clean_chunks(upload_id, save_path)
except Exception as e:
# 记录错误但不阻止删除数据库记录
pass
# 清理数据库记录
await UploadChunk.filter(upload_id=upload_id).delete()
return APIResponse(detail={"message": "上传已取消"})
@chunk_api.get("/upload/status/{upload_id}", dependencies=[Depends(share_required_login)])
async def get_upload_status(upload_id: str):
"""获取上传状态"""
chunk_info = await UploadChunk.filter(upload_id=upload_id, chunk_index=-1).first()
if not chunk_info:
raise HTTPException(status.HTTP_404_NOT_FOUND, detail="上传会话不存在")
# 获取已上传的分片列表
uploaded_chunks = await UploadChunk.filter(
upload_id=upload_id,
completed=True
).values_list('chunk_index', flat=True)
return APIResponse(detail={
"upload_id": upload_id,
"file_name": chunk_info.file_name,
"file_size": chunk_info.file_size,
"chunk_size": chunk_info.chunk_size,
"total_chunks": chunk_info.total_chunks,
"uploaded_chunks": list(uploaded_chunks),
"progress": len(uploaded_chunks) / chunk_info.total_chunks * 100
})
@chunk_api.post("/upload/complete/{upload_id}", dependencies=[Depends(share_required_login)])
@@ -254,32 +365,63 @@ async def complete_upload(upload_id: str, data: CompleteUploadModel, ip: str = D
storage = storages[settings.file_storage]()
# 验证所有分片
completed_chunks = await UploadChunk.filter(
completed_chunks_list = await UploadChunk.filter(
upload_id=upload_id,
completed=True
).count()
if completed_chunks != chunk_info.total_chunks:
).all()
if len(completed_chunks_list) != chunk_info.total_chunks:
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="分片不完整")
# 计算实际上传的总大小并校验
actual_total_size = sum(c.actual_size for c in completed_chunks_list if c.actual_size)
if actual_total_size > settings.uploadSize:
# 清理已上传的分片
_, _, _, _, save_path = await get_chunk_file_path_name(chunk_info.file_name, upload_id)
try:
await storage.clean_chunks(upload_id, save_path)
except Exception:
pass
await UploadChunk.filter(upload_id=upload_id).delete()
max_size_mb = settings.uploadSize / (1024 * 1024)
raise HTTPException(
status_code=403,
detail=f"实际上传大小超过限制,最大为 {max_size_mb:.2f} MB"
)
# 获取文件路径
path, suffix, prefix, _, save_path = await get_chunk_file_path_name(chunk_info.file_name, upload_id)
# 合并文件并计算哈希
await storage.merge_chunks(upload_id, chunk_info, save_path)
# 创建文件记录
expired_at, expired_count, used_count, code = await get_expire_info(data.expire_value, data.expire_style)
await FileCodes.create(
code=code,
file_hash=chunk_info.chunk_hash,
is_chunked=True,
upload_id=upload_id,
size=chunk_info.file_size,
expired_at=expired_at,
expired_count=expired_count,
used_count=used_count,
file_path=path,
uuid_file_name=f"{prefix}{suffix}",
prefix=prefix,
suffix=suffix
)
# 清理临时文件
await storage.clean_chunks(upload_id, save_path)
return APIResponse(detail={"code": code, "name": chunk_info.file_name})
try:
# 合并文件并计算哈希
_, file_hash = await storage.merge_chunks(upload_id, chunk_info, save_path)
# 创建文件记录
expired_at, expired_count, used_count, code = await get_expire_info(data.expire_value, data.expire_style)
await FileCodes.create(
code=code,
file_hash=file_hash, # 使用合并后计算的哈希
is_chunked=True,
upload_id=upload_id,
size=actual_total_size, # 使用实际上传大小而非前端声明值
expired_at=expired_at,
expired_count=expired_count,
used_count=used_count,
file_path=path,
uuid_file_name=f"{prefix}{suffix}",
prefix=prefix,
suffix=suffix
)
# 清理临时文件
await storage.clean_chunks(upload_id, save_path)
# 清理数据库中的分片记录
await UploadChunk.filter(upload_id=upload_id).delete()
ip_limit["upload"].add_ip(ip)
return APIResponse(detail={"code": code, "name": chunk_info.file_name})
except ValueError as e:
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=str(e))
except Exception as e:
# 合并失败时清理临时文件
try:
await storage.clean_chunks(upload_id, save_path)
except Exception:
pass
raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"文件合并失败: {str(e)}")
+322 -122
View File
@@ -144,11 +144,20 @@ class SystemFileStorage(FileStorageInterface):
:param save_path: 文件保存路径
"""
chunk_dir = self.root_path / save_path
chunk_path = chunk_dir.parent / 'chunks' / f"{chunk_index}.part"
chunk_path = chunk_dir.parent / 'chunks' / upload_id / f"{chunk_index}.part"
if not chunk_path.parent.exists():
chunk_path.parent.mkdir(parents=True)
async with aiofiles.open(chunk_path, "wb") as f:
await f.write(chunk_data)
chunk_path.parent.mkdir(parents=True, exist_ok=True)
# 使用临时文件写入,确保原子性
temp_path = chunk_path.with_suffix('.tmp')
try:
async with aiofiles.open(temp_path, "wb") as f:
await f.write(chunk_data)
# 原子重命名
temp_path.rename(chunk_path)
except Exception as e:
if temp_path.exists():
temp_path.unlink()
raise e
async def merge_chunks(self, upload_id: str, chunk_info: UploadChunk, save_path: str) -> tuple[str, str]:
"""
@@ -158,22 +167,37 @@ class SystemFileStorage(FileStorageInterface):
:param save_path: 文件保存路径
:return: (文件路径, 文件哈希值)
"""
output_dir = self.root_path / save_path
output_dir.parent.mkdir(parents=True, exist_ok=True)
output_path = self.root_path / save_path
output_path.parent.mkdir(parents=True, exist_ok=True)
chunk_base_dir = output_path.parent / 'chunks' / upload_id
file_sha256 = hashlib.sha256()
async with aiofiles.open(output_dir, "wb") as out_file:
for i in range(chunk_info.total_chunks):
# 获取分片记录
chunk_record = await UploadChunk.get(upload_id=upload_id, chunk_index=i)
chunk_path = output_dir.parent / 'chunks' / f"{i}.part"
async with aiofiles.open(chunk_path, "rb") as in_file:
chunk_data = await in_file.read()
current_hash = hashlib.sha256(chunk_data).hexdigest()
if current_hash != chunk_record.chunk_hash:
raise ValueError(f"分片{i}哈希不匹配")
file_sha256.update(chunk_data)
await out_file.write(chunk_data)
return output_dir, file_sha256.hexdigest()
# 使用临时文件写入,确保原子性
temp_output = output_path.with_suffix('.merging')
try:
async with aiofiles.open(temp_output, "wb") as out_file:
for i in range(chunk_info.total_chunks):
# 获取分片记录
chunk_record = await UploadChunk.filter(upload_id=upload_id, chunk_index=i).first()
if not chunk_record:
raise ValueError(f"分片{i}记录不存在")
chunk_path = chunk_base_dir / f"{i}.part"
if not chunk_path.exists():
raise ValueError(f"分片{i}文件不存在")
async with aiofiles.open(chunk_path, "rb") as in_file:
chunk_data = await in_file.read()
current_hash = hashlib.sha256(chunk_data).hexdigest()
if current_hash != chunk_record.chunk_hash:
raise ValueError(f"分片{i}哈希不匹配: 期望 {chunk_record.chunk_hash}, 实际 {current_hash}")
file_sha256.update(chunk_data)
await out_file.write(chunk_data)
# 原子重命名
temp_output.rename(output_path)
except Exception as e:
if temp_output.exists():
temp_output.unlink()
raise e
return str(output_path), file_sha256.hexdigest()
async def clean_chunks(self, upload_id: str, save_path: str):
"""
@@ -181,12 +205,19 @@ class SystemFileStorage(FileStorageInterface):
:param upload_id: 上传会话ID
:param save_path: 文件保存路径
"""
chunk_dir = (self.root_path / save_path).parent / 'chunks'
chunk_dir = (self.root_path / save_path).parent / 'chunks' / upload_id
if chunk_dir.exists():
try:
shutil.rmtree(chunk_dir)
except Exception as e:
logger.info(f"清理本地分片目录失败: {e}")
# 清理父级 chunks 目录(如果为空)
chunks_parent = chunk_dir.parent
if chunks_parent.exists() and not any(chunks_parent.iterdir()):
try:
chunks_parent.rmdir()
except Exception as e:
logger.info(f"清理 chunks 父目录失败: {e}")
class S3FileStorage(FileStorageInterface):
@@ -293,61 +324,77 @@ class S3FileStorage(FileStorageInterface):
return result
async def save_chunk(self, upload_id: str, chunk_index: int, chunk_data: bytes, chunk_hash: str, save_path: str):
chunk_key = str(Path(save_path).parent / "chunks" /
upload_id / f"{chunk_index}.part")
async with self.session.client('s3') as s3:
response = await s3.upload_part(
"""
保存分片到 S3(使用独立对象存储每个分片)
注意:这里不使用 S3 原生的 multipart upload,而是将每个分片作为独立对象存储
"""
chunk_key = str(Path(save_path).parent / "chunks" / upload_id / f"{chunk_index}.part")
async with self.session.client(
's3',
endpoint_url=self.endpoint_url,
aws_session_token=self.aws_session_token,
region_name=self.region_name,
config=Config(signature_version=self.signature_version),
) as s3:
# 将分片作为独立对象上传
await s3.put_object(
Bucket=self.bucket_name,
Key=chunk_key,
PartNumber=chunk_index + 1,
UploadId=upload_id,
Body=chunk_data
)
await s3.put_object_tagging(
Bucket=self.bucket_name,
Key=chunk_key,
Tagging={
'TagSet': [
{'Key': 'ChunkHash', 'Value': chunk_hash},
{'Key': 'ETag', 'Value': response['ETag']}
]
Body=chunk_data,
Metadata={
'chunk-hash': chunk_hash,
'chunk-index': str(chunk_index)
}
)
async def merge_chunks(self, upload_id: str, chunk_info: UploadChunk, save_path: str) -> tuple[str, str]:
"""
合并 S3 上的分片文件
由于分片是独立对象存储的,需要下载后合并再上传
"""
file_sha256 = hashlib.sha256()
chunk_dir = str(Path(save_path).parent / "chunks" / upload_id)
async with self.session.client('s3') as s3:
# 获取所有分片
parts = await s3.list_parts(
Bucket=self.bucket_name,
Key=chunk_dir,
UploadId=upload_id
)
part_list = []
for part in parts['Parts']:
part_number = part['PartNumber']
chunk_index = part_number - 1
chunk_record = await UploadChunk.get(upload_id=upload_id, chunk_index=chunk_index)
response = await s3.get_object(
Bucket=self.bucket_name,
Key=f"{chunk_dir}/{chunk_index}.part",
PartNumber=part_number
)
chunk_data = await response['Body'].read()
merged_data = io.BytesIO()
async with self.session.client(
's3',
endpoint_url=self.endpoint_url,
aws_session_token=self.aws_session_token,
region_name=self.region_name,
config=Config(signature_version=self.signature_version),
) as s3:
# 按顺序读取并验证每个分片
for i in range(chunk_info.total_chunks):
chunk_key = f"{chunk_dir}/{i}.part"
chunk_record = await UploadChunk.filter(upload_id=upload_id, chunk_index=i).first()
if not chunk_record:
raise ValueError(f"分片{i}记录不存在")
try:
response = await s3.get_object(
Bucket=self.bucket_name,
Key=chunk_key
)
chunk_data = await response['Body'].read()
except Exception as e:
raise ValueError(f"分片{i}文件不存在: {e}")
current_hash = hashlib.sha256(chunk_data).hexdigest()
if current_hash != chunk_record.chunk_hash:
raise Exception(f"分片{chunk_index}哈希不匹配")
raise ValueError(f"分片{i}哈希不匹配: 期望 {chunk_record.chunk_hash}, 实际 {current_hash}")
file_sha256.update(chunk_data)
part_list.append(
{'PartNumber': part_number, 'ETag': part['ETag']})
# 完成合并
await s3.complete_multipart_upload(
merged_data.write(chunk_data)
# 上传合并后的文件
merged_data.seek(0)
await s3.put_object(
Bucket=self.bucket_name,
Key=save_path,
UploadId=upload_id,
MultipartUpload={'Parts': part_list}
Body=merged_data.getvalue(),
ContentType='application/octet-stream'
)
return save_path, file_sha256.hexdigest()
async def clean_chunks(self, upload_id: str, save_path: str):
@@ -357,32 +404,25 @@ class S3FileStorage(FileStorageInterface):
:param save_path: 文件保存路径
"""
chunk_dir = str(Path(save_path).parent / "chunks" / upload_id)
async with self.session.client('s3') as s3:
async with self.session.client(
's3',
endpoint_url=self.endpoint_url,
aws_session_token=self.aws_session_token,
region_name=self.region_name,
config=Config(signature_version=self.signature_version),
) as s3:
try:
# 终止未完成的分片上传会话
await s3.abort_multipart_upload(
Bucket=self.bucket_name,
Key=chunk_dir,
UploadId=upload_id
)
# 列出并删除所有分片对象
paginator = s3.get_paginator('list_objects_v2')
async for page in paginator.paginate(Bucket=self.bucket_name, Prefix=chunk_dir):
objects = page.get('Contents', [])
if objects:
delete_objects = [{'Key': obj['Key']} for obj in objects]
await s3.delete_objects(
Bucket=self.bucket_name,
Delete={'Objects': delete_objects}
)
except Exception as e:
# 如果上传会话不存在或其他错误,忽略
logger.info(f"清理 S3 分片时出错: {e}")
try:
# 清理已上传的分片数据
parts = await s3.list_parts(
Bucket=self.bucket_name,
Key=chunk_dir,
UploadId=upload_id
)
for part in parts.get('Parts', []):
await s3.delete_object(
Bucket=self.bucket_name,
Key=f"{chunk_dir}/{part['PartNumber'] - 1}.part"
)
except Exception as e:
# 如果分片数据不存在或其他错误,忽略
logger.info(f"清理 S3 分片数据时出错: {e}")
@@ -521,6 +561,105 @@ class OneDriveFileStorage(FileStorageInterface):
f"{file_code.prefix}{file_code.suffix}",
)
def _save_chunk(self, chunk_path: str, chunk_data: bytes):
"""同步保存分片到 OneDrive"""
path_parts = chunk_path.replace("\\", "/").split("/")
filename = path_parts[-1]
dir_path = "/".join(path_parts[:-1])
# 确保目录存在
current_folder = self.root_path
for part in dir_path.split("/"):
if part:
try:
current_folder = current_folder.get_by_path(part).get().execute_query()
except self._ClientRequestException as e:
if e.code == "itemNotFound":
current_folder = current_folder.create_folder(part).execute_query()
else:
raise e
# 上传分片
current_folder.upload(filename, chunk_data).execute_query()
async def save_chunk(self, upload_id: str, chunk_index: int, chunk_data: bytes, chunk_hash: str, save_path: str):
"""保存分片到 OneDrive"""
chunk_path = str(Path(save_path).parent / "chunks" / upload_id / f"{chunk_index}.part")
await asyncio.to_thread(self._save_chunk, chunk_path, chunk_data)
def _read_chunk(self, chunk_path: str) -> bytes:
"""同步读取分片"""
path = self._get_path_str(chunk_path)
file_obj = self.root_path.get_by_path(path).get().execute_query()
return file_obj.get_content().execute_query().value
def _upload_merged(self, save_path: str, data: bytes):
"""同步上传合并后的文件"""
path_parts = save_path.replace("\\", "/").split("/")
filename = path_parts[-1]
dir_path = "/".join(path_parts[:-1])
# 确保目录存在
current_folder = self.root_path
for part in dir_path.split("/"):
if part:
try:
current_folder = current_folder.get_by_path(part).get().execute_query()
except self._ClientRequestException as e:
if e.code == "itemNotFound":
current_folder = current_folder.create_folder(part).execute_query()
else:
raise e
current_folder.upload(filename, data).execute_query()
async def merge_chunks(self, upload_id: str, chunk_info: UploadChunk, save_path: str) -> tuple[str, str]:
"""合并 OneDrive 上的分片文件"""
file_sha256 = hashlib.sha256()
chunk_dir = str(Path(save_path).parent / "chunks" / upload_id)
merged_data = io.BytesIO()
for i in range(chunk_info.total_chunks):
chunk_path = f"{chunk_dir}/{i}.part"
chunk_record = await UploadChunk.filter(upload_id=upload_id, chunk_index=i).first()
if not chunk_record:
raise ValueError(f"分片{i}记录不存在")
try:
chunk_data = await asyncio.to_thread(self._read_chunk, chunk_path)
except Exception as e:
raise ValueError(f"分片{i}文件不存在: {e}")
current_hash = hashlib.sha256(chunk_data).hexdigest()
if current_hash != chunk_record.chunk_hash:
raise ValueError(f"分片{i}哈希不匹配: 期望 {chunk_record.chunk_hash}, 实际 {current_hash}")
file_sha256.update(chunk_data)
merged_data.write(chunk_data)
# 上传合并后的文件
merged_data.seek(0)
await asyncio.to_thread(self._upload_merged, save_path, merged_data.getvalue())
return save_path, file_sha256.hexdigest()
def _delete_chunk_dir(self, chunk_dir: str):
"""同步删除分片目录"""
try:
path = self._get_path_str(chunk_dir)
self.root_path.get_by_path(path).delete_object().execute_query()
except self._ClientRequestException as e:
if e.code != "itemNotFound":
raise e
async def clean_chunks(self, upload_id: str, save_path: str):
"""清理 OneDrive 上的临时分片文件"""
chunk_dir = str(Path(save_path).parent / "chunks" / upload_id)
try:
await asyncio.to_thread(self._delete_chunk_dir, chunk_dir)
except Exception as e:
logger.info(f"清理 OneDrive 分片时出错: {e}")
class OpenDALFileStorage(FileStorageInterface):
def __init__(self):
@@ -560,6 +699,50 @@ class OpenDALFileStorage(FileStorageInterface):
logger.info(e)
raise HTTPException(status_code=404, detail="文件已过期删除")
async def save_chunk(self, upload_id: str, chunk_index: int, chunk_data: bytes, chunk_hash: str, save_path: str):
"""保存分片到 OpenDAL 存储"""
chunk_path = str(Path(save_path).parent / "chunks" / upload_id / f"{chunk_index}.part")
await self.operator.write(chunk_path, chunk_data)
async def merge_chunks(self, upload_id: str, chunk_info: UploadChunk, save_path: str) -> tuple[str, str]:
"""合并 OpenDAL 存储上的分片文件"""
file_sha256 = hashlib.sha256()
chunk_dir = str(Path(save_path).parent / "chunks" / upload_id)
merged_data = io.BytesIO()
for i in range(chunk_info.total_chunks):
chunk_path = f"{chunk_dir}/{i}.part"
chunk_record = await UploadChunk.filter(upload_id=upload_id, chunk_index=i).first()
if not chunk_record:
raise ValueError(f"分片{i}记录不存在")
try:
chunk_data = await self.operator.read(chunk_path)
except Exception as e:
raise ValueError(f"分片{i}文件不存在: {e}")
current_hash = hashlib.sha256(chunk_data).hexdigest()
if current_hash != chunk_record.chunk_hash:
raise ValueError(f"分片{i}哈希不匹配: 期望 {chunk_record.chunk_hash}, 实际 {current_hash}")
file_sha256.update(chunk_data)
merged_data.write(chunk_data)
# 写入合并后的文件
merged_data.seek(0)
await self.operator.write(save_path, merged_data.getvalue())
return save_path, file_sha256.hexdigest()
async def clean_chunks(self, upload_id: str, save_path: str):
"""清理 OpenDAL 存储上的临时分片文件"""
chunk_dir = str(Path(save_path).parent / "chunks" / upload_id)
try:
# OpenDAL 支持递归删除
await self.operator.remove_all(chunk_dir)
except Exception as e:
logger.info(f"清理 OpenDAL 分片时出错: {e}")
class WebDAVFileStorage(FileStorageInterface):
_instance: Optional["WebDAVFileStorage"] = None
@@ -712,55 +895,72 @@ class WebDAVFileStorage(FileStorageInterface):
status_code=503, detail=f"WebDAV连接异常: {str(e)}")
async def save_chunk(self, upload_id: str, chunk_index: int, chunk_data: bytes, chunk_hash: str, save_path: str):
chunk_path = str(Path(save_path).parent / "chunks" / upload_id / f"{chunk_index}.part")
"""保存分片到 WebDAV"""
chunk_dir = str(Path(save_path).parent / "chunks" / upload_id)
chunk_path = f"{chunk_dir}/{chunk_index}.part"
# 先创建目录结构
await self._mkdir_p(chunk_dir)
chunk_url = self._build_url(chunk_path)
async with aiohttp.ClientSession(auth=self.auth) as session:
await session.put(chunk_url, data=chunk_data)
propfind_url = self._build_url(chunk_path)
headers = {
'Content-Type': 'application/xml; charset=utf-8', 'Depth': '0'}
body = f"""
<D:propertyupdate xmlns:D="DAV:">
<D:set>
<D:prop>
<ChunkHash xmlns="urn:filecodebox">{chunk_hash}</ChunkHash>
</D:prop>
</D:set>
</D:propertyupdate>
"""
await session.request('PROPPATCH', propfind_url, headers=headers, data=body)
async with session.put(chunk_url, data=chunk_data) as resp:
if resp.status not in (200, 201, 204):
content = await resp.text()
raise HTTPException(
status_code=resp.status,
detail=f"分片上传失败: {content[:200]}"
)
async def merge_chunks(self, upload_id: str, chunk_info: UploadChunk, save_path: str) -> tuple[str, str]:
"""
合并 WebDAV 上的分片文件
由于大多数 WebDAV 服务器不支持 PATCH 追加,这里下载所有分片后合并上传
"""
file_sha256 = hashlib.sha256()
output_url = self._build_url(save_path)
chunk_dir = str(Path(save_path).parent / "chunks" / upload_id)
merged_data = io.BytesIO()
async with aiohttp.ClientSession(auth=self.auth) as session:
await session.put(output_url, headers={'Content-Length': '0'})
# 按顺序读取并验证每个分片
for i in range(chunk_info.total_chunks):
chunk_path = f"{chunk_dir}/{i}.part"
chunk_url = self._build_url(chunk_path)
propfind_url = self._build_url(chunk_path)
headers = {
'Content-Type': 'application/xml; charset=utf-8', 'Depth': '0'}
body = """
<D:propfind xmlns:D="DAV:">
<D:prop>
<ChunkHash xmlns="urn:filecodebox"/>
</D:prop>
</D:propfind>
"""
async with session.request('PROPFIND', propfind_url, headers=headers, data=body) as resp:
xml_data = await resp.text()
chunk_hash = re.search(
r'<ChunkHash[^>]*>([^<]+)</ChunkHash>', xml_data).group(1)
file_sha256.update(bytes.fromhex(chunk_hash))
# 获取分片记录
chunk_record = await UploadChunk.filter(upload_id=upload_id, chunk_index=i).first()
if not chunk_record:
raise ValueError(f"分片{i}记录不存在")
# 下载分片数据
async with session.get(chunk_url) as resp:
if resp.status != 200:
raise ValueError(f"分片{i}文件不存在或无法访问")
chunk_data = await resp.read()
await session.request('PATCH', output_url, headers={
'Content-Type': 'application/octet-stream',
'Content-Length': str(len(chunk_data)),
'Content-Range': f'bytes */{chunk_info.file_size}'
}, data=chunk_data)
# 验证哈希
current_hash = hashlib.sha256(chunk_data).hexdigest()
if current_hash != chunk_record.chunk_hash:
raise ValueError(f"分片{i}哈希不匹配: 期望 {chunk_record.chunk_hash}, 实际 {current_hash}")
file_sha256.update(chunk_data)
merged_data.write(chunk_data)
# 确保目标目录存在
output_dir = str(Path(save_path).parent)
await self._mkdir_p(output_dir)
# 上传合并后的文件
output_url = self._build_url(save_path)
merged_data.seek(0)
async with session.put(output_url, data=merged_data.getvalue()) as resp:
if resp.status not in (200, 201, 204):
content = await resp.text()
raise HTTPException(
status_code=resp.status,
detail=f"合并文件上传失败: {content[:200]}"
)
return save_path, file_sha256.hexdigest()
async def clean_chunks(self, upload_id: str, save_path: str):
+42 -2
View File
@@ -3,13 +3,14 @@
# @File : tasks.py
# @Software: PyCharm
import asyncio
import datetime
import logging
import os
from tortoise.expressions import Q
from apps.base.models import FileCodes
from apps.base.utils import ip_limit
from apps.base.models import FileCodes, UploadChunk
from apps.base.utils import ip_limit, get_chunk_file_path_name
from core.settings import settings, data_root
from core.storage import FileStorageInterface, storages
from core.utils import get_now
@@ -42,3 +43,42 @@ async def delete_expire_files():
logging.error(e)
finally:
await asyncio.sleep(600)
async def clean_incomplete_uploads():
"""清理超时未完成的分片上传"""
file_storage: FileStorageInterface = storages[settings.file_storage]()
# 默认 24 小时未完成的上传视为过期
expire_hours = getattr(settings, 'chunk_expire_hours', 24)
while True:
try:
expire_time = datetime.datetime.now() - datetime.timedelta(hours=expire_hours)
# 查找所有过期的上传会话(chunk_index=-1 的记录)
expired_sessions = await UploadChunk.filter(
chunk_index=-1,
created_at__lt=expire_time
).all()
for session in expired_sessions:
try:
# 获取分片存储路径
_, _, _, _, save_path = await get_chunk_file_path_name(
session.file_name, session.upload_id
)
# 清理存储中的临时文件
await file_storage.clean_chunks(session.upload_id, save_path)
except Exception as e:
logging.error(f"清理分片文件失败 upload_id={session.upload_id}: {e}")
try:
# 删除该会话的所有数据库记录
await UploadChunk.filter(upload_id=session.upload_id).delete()
logging.info(f"已清理过期上传会话 upload_id={session.upload_id}")
except Exception as e:
logging.error(f"删除分片记录失败 upload_id={session.upload_id}: {e}")
except Exception as e:
logging.error(f"清理未完成上传任务异常: {e}")
finally:
await asyncio.sleep(3600) # 每小时执行一次
+5 -3
View File
@@ -19,7 +19,7 @@ from apps.admin.views import admin_api
from core.database import init_db
from core.response import APIResponse
from core.settings import data_root, settings, BASE_DIR, DEFAULT_CONFIG
from core.tasks import delete_expire_files
from core.tasks import delete_expire_files, clean_incomplete_uploads
from core.logger import logger
from contextlib import asynccontextmanager
@@ -42,6 +42,7 @@ async def lifespan(app: FastAPI):
# 启动后台任务
task = asyncio.create_task(delete_expire_files())
chunk_cleanup_task = asyncio.create_task(clean_incomplete_uploads())
logger.info("应用初始化完成")
try:
@@ -50,7 +51,8 @@ async def lifespan(app: FastAPI):
# 清理操作
logger.info("正在关闭应用...")
task.cancel()
await asyncio.gather(task, return_exceptions=True)
chunk_cleanup_task.cancel()
await asyncio.gather(task, chunk_cleanup_task, return_exceptions=True)
await Tortoise.close_connections()
logger.info("应用已关闭")
@@ -134,7 +136,7 @@ async def get_config():
"explain": settings.page_explain,
"uploadSize": settings.uploadSize,
"expireStyle": settings.expireStyle,
"enableChunk": settings.enableChunk if settings.file_storage == "local" and settings.enableChunk else 0,
"enableChunk": settings.enableChunk,
"openUpload": settings.openUpload,
"notify_title": settings.notify_title,
"notify_content": settings.notify_content,
-1
View File
@@ -22,7 +22,6 @@
## 🚀 更新计划
- [ ] 2025年皮肤
- [ ] 切片上传,同文件秒传,断点续传
- [ ] 文件收集功能
## 📝 项目简介