"""上传服务:流式落盘 + 原子改名 + 元数据入库 + sha256 去重。 两条上传路径(HTTP 整文件 / 分片拼接)共享本类的「存储路径生成 + 落库提交 + sha256 去重」逻辑,避免重复实现。SFTP 服务器仅作为文件暂存通道(chroot 到 upload_root),不再有 HTTP 登记接口。 """ from __future__ import annotations import hashlib import os import uuid from collections.abc import Iterator from datetime import datetime, timezone from pathlib import Path from fastapi import UploadFile from ..config import get_settings from ..dao.uploaded_file_dao import UploadedFileDAO from ..models.uploaded_file import UploadedFile from ..schemas.file import FileUploadResponse, UploadedFileOut def hash_stream(chunks: Iterator[bytes]) -> tuple[int, str]: """对一段字节块序列流式计算 (size, sha256)。供多条上传路径复用。""" h = hashlib.sha256() size = 0 for chunk in chunks: size += len(chunk) h.update(chunk) return size, h.hexdigest() class UploadService: def __init__(self, dao: UploadedFileDAO) -> None: s = get_settings() self.dao = dao self.upload_root = s.resolved_upload_dir() self.chunk_bytes = s.storage.chunk_bytes self.hash_on_upload = s.storage.sha256_on_upload # ---------------- 写入 ---------------- def stream_to_disk( self, file: UploadFile, source: str, uploaded_by: str, ) -> FileUploadResponse: """HTTP 流式上传:先写 ``.part``,DB 行提交后再 ``os.replace`` 成正式名。 失败仅会留下可识别的 ``.part`` 文件,由 start.sh 启动时统一清理。 """ rel_path, abs_path = self.make_storage_path(file.filename or "") part_path = abs_path.with_name(abs_path.name + ".part") size, digest = self._write_part(file, part_path) entity = UploadedFile( storage_path=str(rel_path), original_filename=os.path.basename(file.filename or rel_path.name), content_type=file.content_type or "", size_bytes=size, sha256=digest, source=source, uploaded_by=uploaded_by, ) return self.commit_entity(entity, part_path, abs_path) # ---------------- 查询 ---------------- def list_files(self, limit: int = 100, offset: int = 0) -> tuple[int, list[UploadedFileOut]]: total = self.dao.count() rows = self.dao.list(limit=limit, offset=offset) items = [UploadedFileOut.model_validate(r) for r in rows] return total, items def get_out(self, file_id: int) -> UploadedFileOut | None: row = self.dao.get_by_id(file_id) return UploadedFileOut.model_validate(row) if row else None def find_by_sha256(self, sha256: str) -> UploadedFileOut | None: row = self.dao.get_by_sha256(sha256) return UploadedFileOut.model_validate(row) if row else None def get_out_with_disk_path(self, file_id: int) -> tuple[UploadedFileOut | None, Path | None]: """合并查询:一次 DB 读取同时返回 (元数据, 磁盘绝对路径)。 供下载/删除路径复用,避免原先 get_out + resolve_disk_path 各查一次的重复读。 """ row = self.dao.get_by_id(file_id) if row is None: return None, None out = UploadedFileOut.model_validate(row) path = (self.upload_root / row.storage_path).resolve() return out, path def delete_file(self, file_id: int) -> bool: """硬删除:删磁盘文件 + 删 DB 行。磁盘缺失不阻断 DB 清理。返回是否命中。 供单删/批删复用,保证删除语义一致。 """ _, path = self.get_out_with_disk_path(file_id) row = self.dao.get_by_id(file_id) if row is None: return False if path is not None: try: Path(path).unlink(missing_ok=True) except Exception: # 即使磁盘删除失败也继续清 DB 行,保证列表不再显示 pass self.dao.delete(file_id) return True # ---------------- 共享 helper(供本类与 ChunkUploadService 复用) ---------------- @staticmethod def _safe_ext(filename: str) -> str: return os.path.splitext(os.path.basename(filename))[1] @staticmethod def to_response(row: UploadedFile, *, deduplicated: bool = False) -> FileUploadResponse: return FileUploadResponse( id=row.id, filename=row.original_filename, size_bytes=row.size_bytes, sha256=row.sha256, storage_path=row.storage_path, uploaded_at=row.uploaded_at, deduplicated=deduplicated, ) def make_storage_path(self, original_filename: str) -> tuple[Path, Path]: """生成 ``(rel_path, abs_path)``;abs_path 的父目录已创建。""" now = datetime.now(timezone.utc) rel_dir = Path(f"{now:%Y}/{now:%m}") (self.upload_root / rel_dir).mkdir(parents=True, exist_ok=True) rel_path = rel_dir / f"{uuid.uuid4().hex}{self._safe_ext(original_filename)}" return rel_path, self.upload_root / rel_path def commit_entity( self, entity: UploadedFile, src_path: Path, dest_path: Path, ) -> FileUploadResponse: """落 DB 行后把 src_path 原子改名到 dest_path;任一失败回滚已生成的副作用。""" try: saved = self.dao.create(entity) except Exception: src_path.unlink(missing_ok=True) raise try: os.replace(src_path, dest_path) except Exception: try: self.dao.delete(saved.id) finally: src_path.unlink(missing_ok=True) raise return self.to_response(saved) def dedup_or_commit( self, entity: UploadedFile, src_path: Path, ) -> FileUploadResponse: """公共尾部:按 entity.sha256 去重,命中则删 src 返回旧行;否则 commit。 供分片拼接等「先算出 sha256 再决定落盘」的路径复用,与 stream_to_disk (边写边算、无独立 src)的区别在于这里 sha256 已在 entity 上。 """ existing = self.dao.get_by_sha256(entity.sha256) if existing is not None: src_path.unlink(missing_ok=True) return self.to_response(existing, deduplicated=True) rel_path, abs_path = self.make_storage_path(entity.original_filename) entity.storage_path = str(rel_path) return self.commit_entity(entity, src_path, abs_path) # ---------------- 内部 ---------------- def _write_part(self, file: UploadFile, part_path: Path) -> tuple[int, str]: """把上传流写到 part_path 并 fsync;返回 (size, sha256)。失败时清理残品。""" hasher = hashlib.sha256() if self.hash_on_upload else None size = 0 try: with part_path.open("wb") as out: while chunk := file.file.read(self.chunk_bytes): out.write(chunk) size += len(chunk) if hasher is not None: hasher.update(chunk) out.flush() os.fsync(out.fileno()) except Exception: part_path.unlink(missing_ok=True) raise return size, (hasher.hexdigest() if hasher is not None else "")