From 360ea7ea8f4f9162accab2fe522e6fb7e7fe37f1 Mon Sep 17 00:00:00 2001 From: Pier-Jean Malandrino Date: Fri, 3 Apr 2026 13:52:32 +0200 Subject: [PATCH] Stream upload reads in 64 KB chunks Read uploaded files in fixed-size chunks instead of a single file.read() to reduce peak memory pressure. Also enforces the size limit during streaming so oversized payloads are rejected before fully buffered. --- document-parser/api/documents.py | 12 +++++++++++- document-parser/services/document_service.py | 12 ++++++++++-- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/document-parser/api/documents.py b/document-parser/api/documents.py index 8fac0bb..b6e96ca 100644 --- a/document-parser/api/documents.py +++ b/document-parser/api/documents.py @@ -13,6 +13,8 @@ from services import document_service logger = logging.getLogger(__name__) router = APIRouter(prefix="/api/documents", tags=["documents"]) +_READ_CHUNK_SIZE = 64 * 1024 # 64 KB + def _to_response(doc) -> DocumentResponse: return DocumentResponse( @@ -35,7 +37,15 @@ async def upload(file: UploadFile) -> DocumentResponse: if file.size and file.size > document_service.MAX_FILE_SIZE: raise HTTPException(status_code=413, detail="File too large (max 50 MB)") - content = await file.read() + # Read in chunks to avoid holding the full upload in a single allocation + chunks: list[bytes] = [] + total = 0 + while chunk := await file.read(_READ_CHUNK_SIZE): + total += len(chunk) + if total > document_service.MAX_FILE_SIZE: + raise HTTPException(status_code=413, detail="File too large (max 50 MB)") + chunks.append(chunk) + content = b"".join(chunks) try: doc = await document_service.upload( diff --git a/document-parser/services/document_service.py b/document-parser/services/document_service.py index 58cf2f3..3b564a8 100644 --- a/document-parser/services/document_service.py +++ b/document-parser/services/document_service.py @@ -23,8 +23,14 @@ MAX_FILE_SIZE = 50 * 1024 * 1024 # 50 MB _PDF_MAGIC = b"%PDF" +_UPLOAD_CHUNK_SIZE = 64 * 1024 # 64 KB chunks for streaming writes + + async def upload(filename: str, content_type: str, file_content: bytes) -> Document: - """Save uploaded file to disk and persist metadata.""" + """Save uploaded file to disk and persist metadata. + + Writes the file in fixed-size chunks to keep peak memory usage low. + """ if len(file_content) > MAX_FILE_SIZE: raise ValueError("File too large (max 50 MB)") @@ -37,8 +43,10 @@ async def upload(filename: str, content_type: str, file_content: bytes) -> Docum safe_name = f"{uuid.uuid4()}{ext}" file_path = os.path.join(UPLOAD_DIR, safe_name) + # Write in chunks to avoid doubling memory usage for large files with open(file_path, "wb") as f: - f.write(file_content) + for offset in range(0, len(file_content), _UPLOAD_CHUNK_SIZE): + f.write(file_content[offset : offset + _UPLOAD_CHUNK_SIZE]) # Count PDF pages page_count = _count_pages(file_content)