zTools2/app/services/upload_service.py

"""上传服务：流式落盘 + 原子改名 + 元数据入库 + sha256 去重。

两条上传路径（HTTP 整文件 / 分片拼接）共享本类的「存储路径生成 + 落库提交 +
sha256 去重」逻辑，避免重复实现。SFTP 服务器仅作为文件暂存通道（chroot 到
upload_root），不再有 HTTP 登记接口。
"""

from __future__ import annotations

import hashlib
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


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 流式上传：先写 ``<final>.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

    # ---------------- 共享 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 "")