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.
This commit is contained in:
parent
c2e91262c6
commit
360ea7ea8f
2 changed files with 21 additions and 3 deletions
|
|
@ -13,6 +13,8 @@ from services import document_service
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
router = APIRouter(prefix="/api/documents", tags=["documents"])
|
router = APIRouter(prefix="/api/documents", tags=["documents"])
|
||||||
|
|
||||||
|
_READ_CHUNK_SIZE = 64 * 1024 # 64 KB
|
||||||
|
|
||||||
|
|
||||||
def _to_response(doc) -> DocumentResponse:
|
def _to_response(doc) -> DocumentResponse:
|
||||||
return DocumentResponse(
|
return DocumentResponse(
|
||||||
|
|
@ -35,7 +37,15 @@ async def upload(file: UploadFile) -> DocumentResponse:
|
||||||
if file.size and file.size > document_service.MAX_FILE_SIZE:
|
if file.size and file.size > document_service.MAX_FILE_SIZE:
|
||||||
raise HTTPException(status_code=413, detail="File too large (max 50 MB)")
|
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:
|
try:
|
||||||
doc = await document_service.upload(
|
doc = await document_service.upload(
|
||||||
|
|
|
||||||
|
|
@ -23,8 +23,14 @@ MAX_FILE_SIZE = 50 * 1024 * 1024 # 50 MB
|
||||||
_PDF_MAGIC = b"%PDF"
|
_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:
|
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:
|
if len(file_content) > MAX_FILE_SIZE:
|
||||||
raise ValueError("File too large (max 50 MB)")
|
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}"
|
safe_name = f"{uuid.uuid4()}{ext}"
|
||||||
file_path = os.path.join(UPLOAD_DIR, safe_name)
|
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:
|
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
|
# Count PDF pages
|
||||||
page_count = _count_pages(file_content)
|
page_count = _count_pages(file_content)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue