wip: 切片上传、文件秒传、断点续传陆续适配中当前进度（本地存储已适配）
This commit is contained in:
Lan
@@ -0,0 +1,28 @@
|
||||
from tortoise import connections
|
||||
|
||||
|
||||
async def create_upload_chunk_and_update_file_codes_table():
|
||||
conn = connections.get("default")
|
||||
await conn.execute_script(
|
||||
"""
|
||||
ALTER TABLE "filecodes" ADD "file_hash" VARCHAR(128);
|
||||
ALTER TABLE "filecodes" ADD "is_chunked" BOOL NOT NULL DEFAULT False;
|
||||
ALTER TABLE "filecodes" ADD "upload_id" VARCHAR(128);
|
||||
CREATE TABLE "uploadchunk" (
|
||||
id INTEGER not null primary key autoincrement,
|
||||
"upload_id" VARCHAR(36) NOT NULL,
|
||||
"chunk_index" INT NOT NULL,
|
||||
"chunk_hash" VARCHAR(128) NOT NULL,
|
||||
"total_chunks" INT NOT NULL,
|
||||
"file_size" BIGINT NOT NULL,
|
||||
"chunk_size" INT NOT NULL,
|
||||
"created_at" TIMESTAMPTZ NOT NULL,
|
||||
"file_name" VARCHAR(255) NOT NULL,
|
||||
"completed" BOOL NOT NULL
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
async def migrate():
|
||||
await create_upload_chunk_and_update_file_codes_table()
|
||||
+32
-94
@@ -2,119 +2,55 @@
|
||||
# @Author : Lan
|
||||
# @File : models.py
|
||||
# @Software: PyCharm
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from tortoise import fields
|
||||
from tortoise.models import Model
|
||||
from tortoise.contrib.pydantic import pydantic_model_creator
|
||||
|
||||
from tortoise import fields, models
|
||||
from datetime import datetime
|
||||
from core.utils import get_now
|
||||
|
||||
|
||||
class FileCodes(Model):
|
||||
id: Optional[int] = fields.IntField(pk=True)
|
||||
code: Optional[int] = fields.CharField(
|
||||
description="分享码", max_length=255, index=True, unique=True
|
||||
)
|
||||
prefix: Optional[str] = fields.CharField(
|
||||
max_length=255, description="前缀", default=""
|
||||
)
|
||||
suffix: Optional[str] = fields.CharField(
|
||||
max_length=255, description="后缀", default=""
|
||||
)
|
||||
uuid_file_name: Optional[str] = fields.CharField(
|
||||
max_length=255, description="uuid文件名", null=True
|
||||
)
|
||||
file_path: Optional[str] = fields.CharField(
|
||||
max_length=255, description="文件路径", null=True
|
||||
)
|
||||
size: Optional[int] = fields.IntField(description="文件大小", default=0)
|
||||
text: Optional[str] = fields.TextField(description="文本内容", null=True)
|
||||
expired_at: Optional[datetime] = fields.DatetimeField(
|
||||
null=True, description="过期时间"
|
||||
)
|
||||
expired_count: Optional[int] = fields.IntField(description="可用次数", default=0)
|
||||
used_count: Optional[int] = fields.IntField(description="已用次数", default=0)
|
||||
created_at: Optional[datetime] = fields.DatetimeField(
|
||||
auto_now_add=True, description="创建时间"
|
||||
)
|
||||
#
|
||||
# file_hash = fields.CharField(
|
||||
# max_length=128, # SHA-256需要64字符，这里预留扩展空间
|
||||
# description="文件哈希值",
|
||||
# unique=True,
|
||||
# null=True # 允许旧数据为空
|
||||
# )
|
||||
# hash_algorithm = fields.CharField(
|
||||
# max_length=20,
|
||||
# description="哈希算法类型",
|
||||
# null=True,
|
||||
# default="sha256"
|
||||
# )
|
||||
|
||||
# # 新增分片字段
|
||||
# chunk_size = fields.IntField(
|
||||
# description="分片大小(字节)",
|
||||
# default=0
|
||||
# )
|
||||
# total_chunks = fields.IntField(
|
||||
# description="总分片数",
|
||||
# default=0
|
||||
# )
|
||||
# uploaded_chunks = fields.IntField(
|
||||
# description="已上传分片数",
|
||||
# default=0
|
||||
# )
|
||||
# upload_status = fields.CharField(
|
||||
# max_length=20,
|
||||
# description="上传状态",
|
||||
# default="pending", # pending/in_progress/completed
|
||||
# choices=["pending", "in_progress", "completed"]
|
||||
# )
|
||||
# is_chunked = fields.BooleanField(
|
||||
# description="是否分片上传",
|
||||
# default=False
|
||||
# )
|
||||
class FileCodes(models.Model):
|
||||
id = fields.IntField(pk=True)
|
||||
code = fields.CharField(max_length=255, unique=True, index=True)
|
||||
prefix = fields.CharField(max_length=255, default="")
|
||||
suffix = fields.CharField(max_length=255, default="")
|
||||
uuid_file_name = fields.CharField(max_length=255, null=True)
|
||||
file_path = fields.CharField(max_length=255, null=True)
|
||||
size = fields.IntField(default=0)
|
||||
text = fields.TextField(null=True)
|
||||
expired_at = fields.DatetimeField(null=True)
|
||||
expired_count = fields.IntField(default=0)
|
||||
used_count = fields.IntField(default=0)
|
||||
created_at = fields.DatetimeField(auto_now_add=True)
|
||||
file_hash = fields.CharField(max_length=64, null=True)
|
||||
is_chunked = fields.BooleanField(default=False)
|
||||
upload_id = fields.CharField(max_length=36, null=True)
|
||||
|
||||
async def is_expired(self):
|
||||
# 按时间
|
||||
if self.expired_at is None:
|
||||
return False
|
||||
if self.expired_at and self.expired_count < 0:
|
||||
return self.expired_at < await get_now()
|
||||
# 按次数
|
||||
else:
|
||||
return self.expired_count <= 0
|
||||
return self.expired_count <= 0
|
||||
|
||||
async def get_file_path(self):
|
||||
return f"{self.file_path}/{self.uuid_file_name}"
|
||||
|
||||
|
||||
#
|
||||
# class FileChunks(Model):
|
||||
# id = fields.IntField(pk=True)
|
||||
# file_code = fields.ForeignKeyField(
|
||||
# "models.FileCodes",
|
||||
# related_name="chunks",
|
||||
# on_delete=fields.CASCADE
|
||||
# )
|
||||
# chunk_number = fields.IntField(description="分片序号")
|
||||
# chunk_hash = fields.CharField(
|
||||
# max_length=128,
|
||||
# description="分片哈希校验值"
|
||||
# )
|
||||
# chunk_path = fields.CharField(
|
||||
# max_length=255,
|
||||
# description="分片存储路径"
|
||||
# )
|
||||
# created_at = fields.DatetimeField(
|
||||
# auto_now_add=True,
|
||||
# description="上传时间"
|
||||
# )
|
||||
#
|
||||
# class Meta:
|
||||
# unique_together = [("file_code", "chunk_number")]
|
||||
class UploadChunk(models.Model):
|
||||
id = fields.IntField(pk=True)
|
||||
upload_id = fields.CharField(max_length=36, index=True)
|
||||
chunk_index = fields.IntField()
|
||||
chunk_hash = fields.CharField(max_length=64)
|
||||
total_chunks = fields.IntField()
|
||||
file_size = fields.BigIntField()
|
||||
chunk_size = fields.IntField()
|
||||
file_name = fields.CharField(max_length=255)
|
||||
created_at = fields.DatetimeField(auto_now_add=True)
|
||||
completed = fields.BooleanField(default=False)
|
||||
|
||||
|
||||
class KeyValue(Model):
|
||||
@@ -129,3 +65,5 @@ class KeyValue(Model):
|
||||
|
||||
|
||||
file_codes_pydantic = pydantic_model_creator(FileCodes, name="FileCodes")
|
||||
upload_chunk_pydantic = pydantic_model_creator(UploadChunk, name="UploadChunk")
|
||||
key_value_pydantic = pydantic_model_creator(KeyValue, name="KeyValue")
|
||||
|
||||
@@ -3,3 +3,15 @@ from pydantic import BaseModel
|
||||
|
||||
class SelectFileModel(BaseModel):
|
||||
code: str
|
||||
|
||||
|
||||
class InitChunkUploadModel(BaseModel):
|
||||
file_name: str
|
||||
chunk_size: int = 5 * 1024 * 1024
|
||||
file_size: int
|
||||
file_hash: str
|
||||
|
||||
|
||||
class CompleteUploadModel(BaseModel):
|
||||
expire_value: int
|
||||
expire_style: str
|
||||
|
||||
+25
-1
@@ -1,4 +1,5 @@
|
||||
import datetime
|
||||
import hashlib
|
||||
import os
|
||||
import uuid
|
||||
|
||||
@@ -29,8 +30,19 @@ async def get_file_path_name(file: UploadFile) -> Tuple[str, str, str, str, str]
|
||||
return path, suffix, prefix, file.filename, save_path
|
||||
|
||||
|
||||
async def get_chunk_file_path_name(file_name: str, upload_id: str) -> Tuple[str, str, str, str, str]:
|
||||
"""获取切片文件的路径和文件名"""
|
||||
today = datetime.datetime.now()
|
||||
storage_path = settings.storage_path.strip("/") # 移除开头和结尾的斜杠
|
||||
base_path = f"share/data/{today.strftime('%Y/%m/%d')}/{upload_id}"
|
||||
path = f"{storage_path}/{base_path}" if storage_path else base_path
|
||||
prefix, suffix = os.path.splitext(file_name)
|
||||
save_path = f"{path}/{prefix}{suffix}"
|
||||
return path, suffix, prefix, file_name, save_path
|
||||
|
||||
|
||||
async def get_expire_info(
|
||||
expire_value: int, expire_style: str
|
||||
expire_value: int, expire_style: str
|
||||
) -> Tuple[Optional[datetime.datetime], int, int, str]:
|
||||
"""获取过期信息"""
|
||||
expired_count, used_count = -1, 0
|
||||
@@ -86,6 +98,18 @@ async def get_random_code(style="num") -> str:
|
||||
return code
|
||||
|
||||
|
||||
async def calculate_file_hash(file: UploadFile, chunk_size=1024 * 1024) -> str:
|
||||
sha = hashlib.sha256()
|
||||
await file.seek(0)
|
||||
while True:
|
||||
chunk = await file.read(chunk_size)
|
||||
if not chunk:
|
||||
break
|
||||
sha.update(chunk)
|
||||
await file.seek(0)
|
||||
return sha.hexdigest()
|
||||
|
||||
|
||||
ip_limit = {
|
||||
"error": IPRateLimit(count=settings.uploadCount, minutes=settings.errorMinute),
|
||||
"upload": IPRateLimit(count=settings.errorCount, minutes=settings.errorMinute),
|
||||
|
||||
+142
-13
@@ -1,8 +1,13 @@
|
||||
import hashlib
|
||||
import uuid
|
||||
|
||||
from fastapi import APIRouter, Form, UploadFile, File, Depends, HTTPException
|
||||
from starlette import status
|
||||
|
||||
from apps.admin.dependencies import admin_required
|
||||
from apps.base.models import FileCodes
|
||||
from apps.base.schemas import SelectFileModel
|
||||
from apps.base.utils import get_expire_info, get_file_path_name, ip_limit
|
||||
from apps.base.models import FileCodes, UploadChunk
|
||||
from apps.base.schemas import SelectFileModel, InitChunkUploadModel, CompleteUploadModel
|
||||
from apps.base.utils import get_expire_info, get_file_path_name, ip_limit, get_chunk_file_path_name
|
||||
from core.response import APIResponse
|
||||
from core.settings import settings
|
||||
from core.storage import storages, FileStorageInterface
|
||||
@@ -25,10 +30,10 @@ async def create_file_code(code, **kwargs):
|
||||
|
||||
@share_api.post("/text/", dependencies=[Depends(admin_required)])
|
||||
async def share_text(
|
||||
text: str = Form(...),
|
||||
expire_value: int = Form(default=1, gt=0),
|
||||
expire_style: str = Form(default="day"),
|
||||
ip: str = Depends(ip_limit["upload"]),
|
||||
text: str = Form(...),
|
||||
expire_value: int = Form(default=1, gt=0),
|
||||
expire_style: str = Form(default="day"),
|
||||
ip: str = Depends(ip_limit["upload"]),
|
||||
):
|
||||
text_size = len(text.encode("utf-8"))
|
||||
max_txt_size = 222 * 1024
|
||||
@@ -53,10 +58,10 @@ async def share_text(
|
||||
|
||||
@share_api.post("/file/", dependencies=[Depends(admin_required)])
|
||||
async def share_file(
|
||||
expire_value: int = Form(default=1, gt=0),
|
||||
expire_style: str = Form(default="day"),
|
||||
file: UploadFile = File(...),
|
||||
ip: str = Depends(ip_limit["upload"]),
|
||||
expire_value: int = Form(default=1, gt=0),
|
||||
expire_style: str = Form(default="day"),
|
||||
file: UploadFile = File(...),
|
||||
ip: str = Depends(ip_limit["upload"]),
|
||||
):
|
||||
await validate_file_size(file, settings.uploadSize)
|
||||
|
||||
@@ -142,13 +147,137 @@ async def download_file(key: str, code: str, ip: str = Depends(ip_limit["error"]
|
||||
file_storage: FileStorageInterface = storages[settings.file_storage]()
|
||||
if await get_select_token(code) != key:
|
||||
ip_limit["error"].add_ip(ip)
|
||||
|
||||
has, file_code = await get_code_file_by_code(code, False)
|
||||
if not has:
|
||||
return APIResponse(code=404, detail="文件不存在")
|
||||
|
||||
return (
|
||||
APIResponse(detail=file_code.text)
|
||||
if file_code.text
|
||||
else await file_storage.get_file_response(file_code)
|
||||
)
|
||||
|
||||
|
||||
chunk_api = APIRouter(prefix="/chunk", tags=["切片"])
|
||||
|
||||
|
||||
@chunk_api.post("/upload/init/")
|
||||
async def init_chunk_upload(data: InitChunkUploadModel):
|
||||
# 秒传检查
|
||||
existing = await FileCodes.filter(file_hash=data.file_hash).first()
|
||||
if existing:
|
||||
if await existing.is_expired():
|
||||
file_storage: FileStorageInterface = storages[settings.file_storage](
|
||||
)
|
||||
await file_storage.delete_file(existing)
|
||||
await existing.delete()
|
||||
else:
|
||||
return APIResponse(detail={
|
||||
"code": existing.code,
|
||||
"existed": True,
|
||||
"name": f'{existing.prefix}{existing.suffix}'
|
||||
})
|
||||
|
||||
# 创建上传会话
|
||||
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,
|
||||
total_chunks=total_chunks,
|
||||
file_size=data.file_size,
|
||||
chunk_size=data.chunk_size,
|
||||
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
|
||||
})
|
||||
|
||||
|
||||
@chunk_api.post("/upload/chunk/{upload_id}/{chunk_index}")
|
||||
async def upload_chunk(
|
||||
upload_id: str,
|
||||
chunk_index: int,
|
||||
chunk: UploadFile = File(...),
|
||||
):
|
||||
# 获取上传会话信息
|
||||
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="上传会话不存在")
|
||||
|
||||
# 检查分片索引有效性
|
||||
if chunk_index < 0 or chunk_index >= chunk_info.total_chunks:
|
||||
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="无效的分片索引")
|
||||
|
||||
# 读取分片数据并计算哈希
|
||||
chunk_data = await chunk.read()
|
||||
chunk_hash = hashlib.sha256(chunk_data).hexdigest()
|
||||
|
||||
# 更新或创建分片记录
|
||||
await UploadChunk.update_or_create(
|
||||
upload_id=upload_id,
|
||||
chunk_index=chunk_index,
|
||||
defaults={
|
||||
'chunk_hash': chunk_hash,
|
||||
'completed': True,
|
||||
'file_size': chunk_info.file_size,
|
||||
'total_chunks': chunk_info.total_chunks,
|
||||
'chunk_size': chunk_info.chunk_size,
|
||||
'file_name': chunk_info.file_name
|
||||
}
|
||||
)
|
||||
# 获取文件路径
|
||||
_, _, _, _, 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})
|
||||
|
||||
|
||||
@chunk_api.post("/upload/complete/{upload_id}")
|
||||
async def complete_upload(upload_id: str, data: CompleteUploadModel, ip: str = Depends(ip_limit["upload"])):
|
||||
# 获取上传基本信息
|
||||
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="上传会话不存在")
|
||||
|
||||
storage = storages[settings.file_storage]()
|
||||
# 验证所有分片
|
||||
completed_chunks = await UploadChunk.filter(
|
||||
upload_id=upload_id,
|
||||
completed=True
|
||||
).count()
|
||||
if completed_chunks != chunk_info.total_chunks:
|
||||
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="分片不完整")
|
||||
# 获取文件路径
|
||||
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})
|
||||
|
||||
Reference in New Issue
Block a user