"""上传服务:流式落盘 + 原子改名 + 元数据入库 + sha256 去重。 两条上传路径(HTTP 整文件 / 分片拼接)共享本类的「存储路径生成 + 落库提交 + sha256 去重」逻辑,避免重复实现。SFTP 服务器仅作为文件暂存通道(chroot 到 upload_root),不再有 HTTP 登记接口。 """ from __future__ import annotations import hashlib import logging 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 logger = logging.getLogger("zikai.upload") 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 清理。返回是否命中。 供单删/批删复用,保证删除语义一致。磁盘删除失败仅记日志,仍清 DB 行 (保证列表不再显示),避免磁盘文件泄漏却无任何记录。 """ out, path = self.get_out_with_disk_path(file_id) if out is None: return False if path is not None: try: Path(path).unlink(missing_ok=True) except Exception as exc: logger.warning("删除磁盘文件失败 file_id=%s path=%s: %s", file_id, path, exc) 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 "")