import os import shutil import logging import tempfile from fastapi import APIRouter, HTTPException, Query, UploadFile, File from pydantic import BaseModel from services.auth_service import get_current_user from services.database_service import create_video, update_video_status, get_video from services.storage_service import get_local_upload_path logger = logging.getLogger(__name__) router = APIRouter(prefix="/api", tags=["upload"]) CHUNK_TMP_DIR = os.path.join(tempfile.gettempdir(), "uploads") class UploadRequest(BaseModel): filename: str mode: str = "balanced" file_size: int = 0 person_query: str = "" @router.post("/upload") async def api_upload(req: UploadRequest, token: str = Query(None)): if not token: raise HTTPException(status_code=401, detail="缺少认证 token") try: user = await get_current_user(token) except ValueError as e: raise HTTPException(status_code=401, detail=str(e)) ext = os.path.splitext(req.filename)[1].lower() if ext not in [".mp4", ".mov", ".avi", ".mkv", ".webm"]: raise HTTPException(status_code=400, detail="不支持的视频格式") if req.file_size > 1024 * 1024 * 1024: raise HTTPException(status_code=400, detail="文件大小超过1GB限制") video_id = await create_video( user_id=user["user_id"], original_filename=req.filename, mode=req.mode, file_size_bytes=req.file_size, person_query=req.person_query, ) return { "video_id": video_id, "storage_mode": "local", } @router.post("/upload/{video_id}/file") async def api_upload_file( video_id: str, token: str = Query(None), file: UploadFile = File(...), ): if not token: raise HTTPException(status_code=401, detail="缺少认证 token") try: user = await get_current_user(token) except ValueError as e: raise HTTPException(status_code=401, detail=str(e)) video = await get_video(video_id) if not video: raise HTTPException(status_code=404, detail="视频不存在") if video["user_id"] != user["user_id"]: raise HTTPException(status_code=403, detail="无权操作") local_path = get_local_upload_path(user["user_id"], video_id, video["original_filename"] or "video.mp4") chunk_size = 1024 * 1024 with open(local_path, "wb") as f: while True: chunk = await file.read(chunk_size) if not chunk: break f.write(chunk) await update_video_status(video_id, "uploaded") return {"video_id": video_id, "status": "uploaded", "message": "上传成功"} class ChunkCompleteRequest(BaseModel): total_chunks: int filename: str def _get_chunk_dir(video_id: str) -> str: return os.path.join(CHUNK_TMP_DIR, video_id) @router.post("/upload/{video_id}/chunk") async def api_upload_chunk( video_id: str, token: str = Query(None), chunk_index: int = Query(...), total_chunks: int = Query(...), chunk: UploadFile = File(...), ): if not token: raise HTTPException(status_code=401, detail="缺少认证 token") try: user = await get_current_user(token) except ValueError as e: raise HTTPException(status_code=401, detail=str(e)) video = await get_video(video_id) if not video: raise HTTPException(status_code=404, detail="视频不存在") if video["user_id"] != user["user_id"]: raise HTTPException(status_code=403, detail="无权操作") if chunk_index < 0 or chunk_index >= total_chunks: raise HTTPException(status_code=400, detail="分片索引越界") chunk_dir = _get_chunk_dir(video_id) os.makedirs(chunk_dir, exist_ok=True) chunk_path = os.path.join(chunk_dir, f"chunk_{chunk_index:04d}") chunk_size = 1024 * 1024 try: with open(chunk_path, "wb") as f: while True: data = await chunk.read(chunk_size) if not data: break f.write(data) except Exception as e: logger.exception("保存分片失败 video_id=%s chunk_index=%s", video_id, chunk_index) raise HTTPException(status_code=500, detail=f"保存分片失败: {e}") # 移除 listdir 调用以减少 IO 开销;前端通过 uploadedChunks 自行跟踪进度 return { "video_id": video_id, "chunk_index": chunk_index, "received": chunk_index + 1, "total_chunks": total_chunks, } @router.post("/upload/{video_id}/complete") async def api_upload_complete( video_id: str, token: str = Query(None), req: ChunkCompleteRequest = ..., ): if not token: raise HTTPException(status_code=401, detail="缺少认证 token") try: user = await get_current_user(token) except ValueError as e: raise HTTPException(status_code=401, detail=str(e)) video = await get_video(video_id) if not video: raise HTTPException(status_code=404, detail="视频不存在") if video["user_id"] != user["user_id"]: raise HTTPException(status_code=403, detail="无权操作") chunk_dir = _get_chunk_dir(video_id) if not os.path.isdir(chunk_dir): raise HTTPException(status_code=400, detail="未找到分片数据,请重新上传") expected = req.total_chunks chunk_files = [] for i in range(expected): p = os.path.join(chunk_dir, f"chunk_{i:04d}") if not os.path.exists(p): raise HTTPException(status_code=400, detail=f"分片 {i} 缺失,请重新上传该分片") chunk_files.append(p) original_filename = req.filename or video.get("original_filename") or "video.mp4" final_path = get_local_upload_path(user["user_id"], video_id, original_filename) try: with open(final_path, "wb") as out: for p in chunk_files: with open(p, "rb") as part: shutil.copyfileobj(part, out) except Exception as e: logger.exception("合并分片失败 video_id=%s", video_id) raise HTTPException(status_code=500, detail=f"合并分片失败: {e}") final_size = os.path.getsize(final_path) expected_size = video.get("file_size_bytes") or 0 if expected_size and final_size != expected_size: logger.warning( "合并后大小不一致 video_id=%s expected=%s actual=%s", video_id, expected_size, final_size, ) shutil.rmtree(chunk_dir, ignore_errors=True) await update_video_status(video_id, "uploaded") return { "video_id": video_id, "status": "uploaded", "message": "上传成功", "file_size": final_size, } @router.delete("/upload/{video_id}/chunk") async def api_upload_cancel( video_id: str, token: str = Query(None), ): if not token: raise HTTPException(status_code=401, detail="缺少认证 token") try: user = await get_current_user(token) except ValueError as e: raise HTTPException(status_code=401, detail=str(e)) video = await get_video(video_id) if not video: raise HTTPException(status_code=404, detail="视频不存在") if video["user_id"] != user["user_id"]: raise HTTPException(status_code=403, detail="无权操作") chunk_dir = _get_chunk_dir(video_id) if os.path.isdir(chunk_dir): shutil.rmtree(chunk_dir, ignore_errors=True) return {"video_id": video_id, "status": "cancelled", "message": "已清理分片"}