"""分片上传服务：会话管理 + 分片落盘 + 拼接 + 创建转写任务。 存储布局:: //0.part 分片暂存 //1.part ... ///. complete 后的正式视频 与 server 的区别：视频无需 sha256 去重（每个视频都转写），complete 直接创建 Task。 管线触发由 controller 调用 pipeline.enqueue_task，本服务不依赖 pipeline。 """ from __future__ import annotations import logging import os import shutil import uuid from datetime import datetime, timedelta, timezone from pathlib import Path from fastapi import HTTPException from sqlalchemy.orm import Session from ..config import get_settings from ..models.task import Task from ..models.upload_session import UploadSession from ..schemas.task import ( ChunkUploadResponse, CompleteResponse, CreateSessionRequest, CreateSessionResponse, SessionStatusResponse, ) logger = logging.getLogger("audio2text.upload") class UploadService: def __init__(self, db: Session) -> None: self.db = db s = get_settings() self.upload_root = s.upload_dir() self.work_root = s.work_dir() self.chunk_bytes = s.storage.chunk_bytes self.session_ttl = s.storage.chunk_session_ttl_seconds # ---------------- 会话生命周期 ---------------- def create_session(self, body: CreateSessionRequest) -> CreateSessionResponse: upload_id = uuid.uuid4().hex session = UploadSession( upload_id=upload_id, filename=body.filename, size_bytes=body.size_bytes, chunk_size=body.chunk_size, total_chunks=body.total_chunks, uploaded_chunks=[], status="pending", ) self.db.add(session) self.db.commit() self._session_dir(upload_id).mkdir(parents=True, exist_ok=True) return CreateSessionResponse( upload_id=upload_id, filename=body.filename, size_bytes=body.size_bytes, chunk_size=body.chunk_size, total_chunks=body.total_chunks, ) def get_status(self, upload_id: str) -> SessionStatusResponse: session = self._require_session(upload_id) return SessionStatusResponse( upload_id=session.upload_id, filename=session.filename, size_bytes=session.size_bytes, chunk_size=session.chunk_size, total_chunks=session.total_chunks, uploaded_chunks=list(session.uploaded_chunks or []), completed=(session.status == "completed"), task_id=session.task_id, ) # ---------------- 分片写入 ---------------- def write_chunk(self, upload_id: str, index: int, data: bytes) -> list[int]: session = self._require_session(upload_id) self._validate_index(session, index) session_dir = self._session_dir(upload_id) session_dir.mkdir(parents=True, exist_ok=True) chunk_path = session_dir / f"{index}.part" try: with chunk_path.open("wb") as out: out.write(data) out.flush() os.fsync(out.fileno()) except Exception: chunk_path.unlink(missing_ok=True) raise uploaded = list(session.uploaded_chunks or []) if index not in uploaded: uploaded.append(index) session.uploaded_chunks = uploaded session.updated_at = datetime.now(timezone.utc) self.db.commit() return sorted(uploaded) # ---------------- 拼接 + 创建任务 ---------------- def complete(self, upload_id: str) -> CompleteResponse: session = self._require_session(upload_id) # 幂等：已 complete 则返回已建任务 if session.status == "completed" and session.task_id is not None: task = self.db.get(Task, session.task_id) if task is not None: return CompleteResponse( task_id=task.id, filename=task.filename, size_bytes=session.size_bytes, status=task.status, ) uploaded = set(session.uploaded_chunks or []) missing = [i for i in range(session.total_chunks) if i not in uploaded] if missing: raise HTTPException( 409, f"分片不齐全：缺失 {len(missing)} 个，例如 {sorted(missing)[:10]}", ) final_path = self._assemble(session) rel = str(final_path.relative_to(self.upload_root)) task = Task( filename=session.filename, source_path=rel, status="queued", progress=0.0, ) self.db.add(task) session.status = "completed" session.final_path = rel session.updated_at = datetime.now(timezone.utc) self.db.commit() self.db.refresh(task) # 正向关联：session → task（替代旧的 source_path 反向查找） session.task_id = task.id self.db.commit() # 清理分片暂存 self._cleanup_session_dir(upload_id) logger.info("上传完成 task_id=%s file=%s size=%d", task.id, session.filename, session.size_bytes) return CompleteResponse( task_id=task.id, filename=task.filename, size_bytes=session.size_bytes, status=task.status, ) # ---------------- 过期会话清理 ---------------- def reap_stale_sessions(self) -> int: """清理被放弃的会话：pending 且 updated_at 超过 ttl。 时间比较统一用 aware UTC datetime，避免 naive datetime 的 .timestamp() 按本地时区算导致的偏移。 """ cutoff = datetime.now(timezone.utc) - timedelta(seconds=self.session_ttl) sessions = ( self.db.query(UploadSession) .filter(UploadSession.status == "pending") .all() ) n = 0 for session in sessions: updated = session.updated_at if updated is None: continue # SQLite 存 naive datetime，统一补 UTC 后比较 if updated.tzinfo is None: updated = updated.replace(tzinfo=timezone.utc) if updated < cutoff: self._cleanup_session_dir(session.upload_id) self.db.delete(session) n += 1 logger.info("清理过期分片会话 upload_id=%s file=%s", session.upload_id, session.filename) if n: self.db.commit() return n # ---------------- 内部 ---------------- def _require_session(self, upload_id: str) -> UploadSession: if not upload_id: raise HTTPException(400, "upload_id 不能为空") session = ( self.db.query(UploadSession) .filter(UploadSession.upload_id == upload_id) .first() ) if session is None: raise HTTPException(404, f"会话不存在或已过期：{upload_id}") return session @staticmethod def _validate_index(session: UploadSession, index: int) -> None: if index < 0 or index >= session.total_chunks: raise HTTPException(400, f"分片下标越界：{index} 不在 [0, {session.total_chunks})") def _session_dir(self, upload_id: str) -> Path: return self.work_root / upload_id def _assemble(self, session: UploadSession) -> Path: """按 index 顺序拼接全部分片为正式视频文件。""" ext = Path(session.filename).suffix or ".mp4" now = datetime.now(timezone.utc) sub = self.upload_root / f"{now:%Y}" / f"{now:%m}" sub.mkdir(parents=True, exist_ok=True) final = sub / f"{uuid.uuid4().hex}{ext}" part_path = final.with_suffix(final.suffix + ".part") session_dir = self._session_dir(session.upload_id) try: with part_path.open("wb") as out: for index in range(session.total_chunks): chunk_path = session_dir / f"{index}.part" if not chunk_path.is_file(): raise HTTPException(409, f"拼接时发现分片缺失：{index}.part") with chunk_path.open("rb") as src: while buf := src.read(self.chunk_bytes): out.write(buf) out.flush() os.fsync(out.fileno()) os.replace(part_path, final) except Exception: part_path.unlink(missing_ok=True) raise return final def _cleanup_session_dir(self, upload_id: str) -> None: session_dir = self._session_dir(upload_id) try: if session_dir.exists(): shutil.rmtree(session_dir, ignore_errors=True) except Exception as exc: # pragma: no cover logger.warning("清理会话目录失败 upload_id=%s: %s", upload_id, exc)