"""Physical storage operations for uploaded files.""" import os import uuid from dataclasses import dataclass from pathlib import Path from fastapi import UploadFile UPLOAD_CHUNK_SIZE = 1024 * 1024 class FileTooLargeError(Exception): """Raised when a streamed upload exceeds its configured limit.""" class InvalidStorageNameError(Exception): """Raised when a storage name could escape the configured root.""" @dataclass(frozen=True) class PendingUpload: temp_path: Path storage_name: str file_size: int @dataclass(frozen=True) class StagedDeletion: original_path: Path staged_path: Path class FileStorage: """Store files below one root and provide rollback-friendly operations.""" def __init__(self, root: Path): self.root = root self.root.mkdir(parents=True, exist_ok=True) def path_for(self, storage_name: str) -> Path: if not storage_name or Path(storage_name).name != storage_name: raise InvalidStorageNameError(storage_name) return self.root / storage_name async def write_upload(self, upload: UploadFile, storage_name: str, max_bytes: int) -> PendingUpload: self.path_for(storage_name) temp_path = self.root / f".upload-{uuid.uuid4()}.tmp" file_size = 0 try: with temp_path.open("xb") as output: while chunk := await upload.read(UPLOAD_CHUNK_SIZE): file_size += len(chunk) if file_size > max_bytes: raise FileTooLargeError output.write(chunk) except Exception: temp_path.unlink(missing_ok=True) raise finally: await upload.close() return PendingUpload(temp_path=temp_path, storage_name=storage_name, file_size=file_size) def promote(self, pending: PendingUpload) -> Path: final_path = self.path_for(pending.storage_name) os.replace(pending.temp_path, final_path) return final_path def discard(self, path: Path) -> None: path.unlink(missing_ok=True) def stage_deletions(self, storage_names: list[str]) -> list[StagedDeletion]: staged: list[StagedDeletion] = [] try: for storage_name in storage_names: original_path = self.path_for(storage_name) if not original_path.exists(): continue staged_path = self.root / f".delete-{uuid.uuid4()}.tmp" os.replace(original_path, staged_path) staged.append(StagedDeletion(original_path=original_path, staged_path=staged_path)) except Exception: self.restore_deletions(staged) raise return staged def restore_deletions(self, staged: list[StagedDeletion]) -> None: for item in reversed(staged): if item.staged_path.exists(): os.replace(item.staged_path, item.original_path) def purge_deletions(self, staged: list[StagedDeletion]) -> None: for item in staged: item.staged_path.unlink(missing_ok=True)