d7514f4e65
Add transactional file storage workflows, typed API contracts, recursive category handling, frontend component separation, and Files API coverage.
95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
"""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)
|