"""上传服务：流式落盘 + 原子改名 + 元数据入库。""" 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: """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, part_path, abs_path) def register_sftp( self, filename: str, original_filename: str, uploaded_by: str = "sftp", ) -> FileUploadResponse: """登记一个已通过 SFTP 落到 ``incoming/`` 下的文件。 相同 sha256 已存在时直接返回旧行并删掉新副本（去重）。 """ 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: src_abs.unlink(missing_ok=True) return self._to_response(existing) rel_path, abs_path = self._make_storage_path(original_filename or filename) 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, ) return self._commit(entity, src_abs, abs_path) # ---------------- 查询 ---------------- 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 # ---------------- 内部 ---------------- @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 _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 _validate_incoming_path(self, filename: str) -> Path: """校验客户端给出的相对路径必须落在 upload_root/incoming/ 之下且是文件。""" 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 _commit( 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 _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 "")