refactor: 抽出 _commit / _make_storage_path 复用两条上传路径

2026-06-23 16:29:39 +00:00
parent 6c7d3f1272
commit fda635ccbd
1 changed files with 61 additions and 85 deletions
--- a/app/services/upload_service.py
+++ b/app/services/upload_service.py
@@ -15,7 +15,7 @@ from ..dao.uploaded_file_dao import UploadedFileDAO
 from ..models.uploaded_file import UploadedFile
 from ..schemas.file import FileUploadResponse, UploadedFileOut

-# SFTP 客户端登记前必须先把文件写入此目录（chroot 内）
+# SFTP 客户端登记前必须把文件先落到此目录（chroot 内）
 SFTP_INCOMING_DIR = "incoming"


@@ -27,24 +27,17 @@ class UploadService:
        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 流式上传：先写 ``<final>.part``，DB 行提交后再 ``os.replace`` 成正式名。

-        先写 ``<final>.part``，DB 行提交成功后再 ``os.replace`` 成正式名。
-        任何失败仅会留下可识别的 ``.part``，由 start.sh 启动时自动清理。
+        失败仅会留下可识别的 ``.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
+        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(
@@ -56,22 +49,34 @@ class UploadService:
            source=source,
            uploaded_by=uploaded_by,
        )
-        try:
-            saved = self.dao.create(entity)
-        except Exception:
-            part_path.unlink(missing_ok=True)
-            raise
+        return self._commit(entity, part_path, abs_path)

-        try:
-            os.replace(part_path, abs_path)
-        except Exception:
-            try:
-                self.dao.delete(saved.id)
-            finally:
-                part_path.unlink(missing_ok=True)
-            raise
+    def register_sftp(
+        self, filename: str, original_filename: str, uploaded_by: str = "sftp",
+    ) -> FileUploadResponse:
+        """登记一个已通过 SFTP 落到 ``incoming/`` 下的文件。

-        return self._to_response(saved)
+        相同 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)

    # ---------------- 查询 ----------------

@@ -92,61 +97,8 @@ class UploadService:
        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/<uuid>.<ext>，落 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]
@@ -162,11 +114,16 @@ class UploadService:
            uploaded_at=row.uploaded_at,
        )

-    def _validate_incoming_path(self, filename: str) -> Path:
-        """把客户端传来的相对路径解析成 upload_root 下的绝对路径。
+    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

-        要求落在 ``upload_root/incoming/`` 之下且为普通文件，否则抛 HTTPException。
-        """
+    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()
@@ -179,8 +136,27 @@ class UploadService:
            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)。"""
+        """流式读取磁盘文件，返回 (size, sha256)。"""
        h = hashlib.sha256()
        size = 0
        with path.open("rb") as fh:
@@ -190,7 +166,7 @@ class UploadService:
        return size, h.hexdigest()

    def _write_part(self, file: UploadFile, part_path: Path) -> tuple[int, str]:
-        """流式写到 part_path 并 fsync；返回 (size, sha256)。失败时清理残品。"""
+        """把上传流写到 part_path 并 fsync；返回 (size, sha256)。失败时清理残品。"""
        hasher = hashlib.sha256() if self.hash_on_upload else None
        size = 0
        try: