"""上传服务：流式落盘 + 原子改名 + 元数据入库。""" from __future__ import annotations import hashlib import os import uuid from datetime import datetime, timezone from pathlib import Path from fastapi import HTTPException, 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 # SFTP 客户端登记前必须先把文件写入此目录（chroot 内） SFTP_INCOMING_DIR = "incoming" 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: """流式落盘并写元数据。 先写 ``.part``，DB 行提交成功后再 ``os.replace`` 成正式名。 任何失败仅会留下可识别的 ``.part``，由 start.sh 启动时自动清理。 """ rel_dir = self._relative_dir() (self.upload_root / rel_dir).mkdir(parents=True, exist_ok=True) ext = self._safe_ext(file.filename or "") rel_path = rel_dir / f"{uuid.uuid4().hex}{ext}" abs_path = self.upload_root / rel_path 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, ) try: saved = self.dao.create(entity) except Exception: part_path.unlink(missing_ok=True) raise try: os.replace(part_path, abs_path) except Exception: try: self.dao.delete(saved.id) finally: part_path.unlink(missing_ok=True) raise return self._to_response(saved) # ---------------- 查询 ---------------- def list_files(self, limit: int = 100, offset: int = 0) -> tuple[int, list[UploadedFileOut]]: rows = self.dao.list(limit=limit, offset=offset) items = [UploadedFileOut.model_validate(r) for r in rows] return len(items), 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 resolve_disk_path(self, file_id: int) -> Path | None: row = self.dao.get_by_id(file_id) return (self.upload_root / row.storage_path).resolve() if row else None # ---------------- SFTP 登记 ---------------- def register_sftp( self, filename: str, original_filename: str, uploaded_by: str = "sftp", ) -> FileUploadResponse: """把一个已经通过 SFTP 落到 incoming/ 下的文件登记入库。 失败 / 拒绝场景： - filename 路径穿越或不在 incoming/ 下 → 400 - 文件不存在 / 不是普通文件 → 404 - sha256 已存在 → 返回已有行（去重），同时删除新上传的副本 - 否则把文件从 incoming/ 原子改名到 YYYY/MM/.，落 DB 行（source='sftp'） """ src_abs = self._validate_incoming_path(filename) size, digest = self._hash_disk_file(src_abs) existing = self.dao.get_by_sha256(digest) if existing is not None: # 内容已存在 -> 丢弃新副本，避免 uploads/ 越积越多。 src_abs.unlink(missing_ok=True) return self._to_response(existing) rel_dir = self._relative_dir() (self.upload_root / rel_dir).mkdir(parents=True, exist_ok=True) ext = self._safe_ext(original_filename or filename) rel_path = rel_dir / f"{uuid.uuid4().hex}{ext}" abs_path = self.upload_root / rel_path entity = UploadedFile( storage_path=str(rel_path), original_filename=os.path.basename(original_filename or src_abs.name), content_type="", size_bytes=size, sha256=digest, source="sftp", uploaded_by=uploaded_by, ) saved = self.dao.create(entity) try: os.replace(src_abs, abs_path) except Exception: # 改名失败 -> 回滚 DB 行，避免出现孤儿元数据 self.dao.delete(saved.id) raise return self._to_response(saved) # ---------------- 内部 ---------------- @staticmethod def _relative_dir() -> Path: now = datetime.now(timezone.utc) return Path(f"{now:%Y}/{now:%m}") @staticmethod def _safe_ext(filename: str) -> str: return os.path.splitext(os.path.basename(filename))[1] @staticmethod def _to_response(row: UploadedFile) -> 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, ) def _validate_incoming_path(self, filename: str) -> Path: """把客户端传来的相对路径解析成 upload_root 下的绝对路径。 要求落在 ``upload_root/incoming/`` 之下且为普通文件，否则抛 HTTPException。 """ if not filename or filename.startswith(("/", "\\")): raise HTTPException(400, "filename 必须是 incoming/ 下的相对路径") incoming_root = (self.upload_root / SFTP_INCOMING_DIR).resolve() try: abs_path = (self.upload_root / filename).resolve() abs_path.relative_to(incoming_root) except ValueError: raise HTTPException(400, f"filename 必须落在 {SFTP_INCOMING_DIR}/ 之下") if not abs_path.is_file(): raise HTTPException(404, f"文件不存在或不是普通文件：{filename}") return abs_path def _hash_disk_file(self, path: Path) -> tuple[int, str]: """以流式方式读盘上文件，返回 (size, sha256)。""" h = hashlib.sha256() size = 0 with path.open("rb") as fh: while chunk := fh.read(self.chunk_bytes): size += len(chunk) h.update(chunk) return size, h.hexdigest() 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 "")