feat: s3适配切片上传

This commit is contained in:
Lan
2025-12-01 15:13:33 +08:00
parent a20a0beb79
commit c1966b3942
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)}")