71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
"""Document and binary storage helpers."""
|
|
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import pymupdf
|
|
from fastapi import HTTPException, UploadFile
|
|
|
|
from app.config import settings
|
|
|
|
_COPY_CHUNK_SIZE = 1024 * 1024
|
|
|
|
|
|
def validate_pdf(filepath: Path) -> int:
|
|
"""Validate a PDF's magic bytes and return its page count."""
|
|
with filepath.open("rb") as file:
|
|
if file.read(5) != b"%PDF-":
|
|
raise ValueError("Not a valid PDF file (missing magic bytes)")
|
|
|
|
try:
|
|
with pymupdf.open(filepath) as document:
|
|
page_count = len(document)
|
|
except Exception as err:
|
|
raise ValueError(f"Failed to open PDF with PyMuPDF: {err}") from err
|
|
|
|
if page_count == 0:
|
|
raise ValueError("PDF does not contain any pages")
|
|
return page_count
|
|
|
|
|
|
def safe_filename(filename: str | None) -> str:
|
|
"""Return a display-safe basename without allowing path components."""
|
|
name = Path(filename or "Untitled.pdf").name.strip()
|
|
return name or "Untitled.pdf"
|
|
|
|
|
|
async def save_upload_file(upload_file: UploadFile) -> tuple[str, int, int]:
|
|
"""Stream an uploaded PDF to storage, enforcing the configured size limit."""
|
|
file_id = str(uuid.uuid4())
|
|
settings.PDF_STORAGE_PATH.mkdir(parents=True, exist_ok=True)
|
|
filepath = settings.PDF_STORAGE_PATH / f"{file_id}.pdf"
|
|
max_bytes = settings.MAX_UPLOAD_MB * 1024 * 1024
|
|
total_bytes = 0
|
|
|
|
try:
|
|
with filepath.open("wb") as destination:
|
|
while chunk := upload_file.file.read(_COPY_CHUNK_SIZE):
|
|
total_bytes += len(chunk)
|
|
if total_bytes > max_bytes:
|
|
raise HTTPException(status_code=413, detail="File too large")
|
|
destination.write(chunk)
|
|
except HTTPException:
|
|
filepath.unlink(missing_ok=True)
|
|
raise
|
|
except Exception as err:
|
|
filepath.unlink(missing_ok=True)
|
|
raise HTTPException(status_code=500, detail="Failed to save file") from err
|
|
|
|
try:
|
|
page_count = validate_pdf(filepath)
|
|
except ValueError as err:
|
|
filepath.unlink(missing_ok=True)
|
|
raise HTTPException(status_code=400, detail=str(err)) from err
|
|
|
|
return file_id, total_bytes, page_count
|
|
|
|
|
|
def delete_pdf_file(file_id: str) -> None:
|
|
"""Delete a PDF file from storage."""
|
|
filepath = settings.PDF_STORAGE_PATH / f"{file_id}.pdf"
|
|
filepath.unlink(missing_ok=True)
|