"""上传服务:流式落盘 + 原子改名 + 元数据入库 + sha256 去重。 两条上传路径(HTTP 整文件 / 分片拼接)共享本类的「存储路径生成 + 落库提交 + sha256 去重」逻辑,避免重复实现。SFTP 服务器仅作为文件暂存通道(chroot 到 upload_root),不再有 HTTP 登记接口。 """ from __future__ import annotations import hashlib import logging import os import time 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") # 磁盘扫描频率限制:两次扫描至少间隔此秒数,否则跳过(只返回 DB 缓存)。 _SCAN_MIN_INTERVAL = 3.0 _last_scan_time: float = 0.0 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 scan_sftp_files(self, force: bool = False) -> int: """扫描 uploads/ 目录,把磁盘上有但 DB 未记录的文件补录入库。 SFTP 上传的文件直接落盘到 upload_root(chroot 根),不经过 HTTP 路径, 因此没有 DB 记录。本方法遍历磁盘文件,与 DB 已有 storage_path 比对, 为缺失项创建记录(source="sftp",计算 size + sha256)。 频率限制:两次扫描间隔 < _SCAN_MIN_INTERVAL(3s)则跳过,force=True 强制扫。 返回本次新补录的条数(跳过时返回 0)。 """ global _last_scan_time now = time.monotonic() if not force and (now - _last_scan_time) < _SCAN_MIN_INTERVAL: return 0 _last_scan_time = now known = self.dao.list_storage_paths() new_count = 0 # 排除分片会话暂存目录与 .part 残品 skip_dirs = {".work"} for abs_path in self.upload_root.rglob("*"): if not abs_path.is_file(): continue if abs_path.suffix == ".part": continue # 跳过 .work 目录下的任何文件 rel = abs_path.relative_to(self.upload_root) if rel.parts and rel.parts[0] in skip_dirs: continue storage_path = str(rel).replace("\\", "/") if storage_path in known: continue # 磁盘有、DB 无:补录 try: size = abs_path.stat().st_size sha256 = self._hash_file(abs_path) except Exception as exc: logger.warning("扫描文件失败 path=%s: %s", abs_path, exc) continue entity = UploadedFile( storage_path=storage_path, original_filename=abs_path.name, content_type="", size_bytes=size, sha256=sha256, source="sftp", uploaded_by="sftp", ) try: self.dao.create(entity) new_count += 1 known.add(storage_path) except Exception as exc: logger.warning("补录 SFTP 文件失败 path=%s: %s", storage_path, exc) if new_count: logger.info("SFTP 文件扫描补录 %d 个", new_count) return new_count def _hash_file(self, path: Path) -> str: """流式计算文件 sha256(避免大文件一次性读入内存)。""" h = hashlib.sha256() with path.open("rb") as f: while chunk := f.read(self.chunk_bytes): h.update(chunk) return h.hexdigest() 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 "")