6c7d3f1272
- DAO 增加 get_by_sha256
- UploadService 抽出 _to_response 复用；新增 find_by_sha256、register_sftp
- 控制器新增两条路由（注意排在 /{file_id} 之前以避免被吞）
- 引入 SftpRegisterRequest schema
- SFTP 服务启动时自动创建 uploads/incoming/
- 行为：路径穿越 400、文件不存在 404、sha256 已存在则去重返回已有行
209 lines
7.6 KiB
Python
209 lines
7.6 KiB
Python
"""上传服务：流式落盘 + 原子改名 + 元数据入库。"""
|
||
|
||
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:
|
||
"""流式落盘并写元数据。
|
||
|
||
先写 ``<final>.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/<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]
|
||
|
||
@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 "")
|