Merge pull request #718 from christianlouis/perf-upload-15584535240168516797
⚡ Optimize file upload endpoint to use asynchronous aiofiles
This commit is contained in:
+4
-3
@@ -11,6 +11,7 @@ import zipfile
|
|||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from typing import Annotated, List, Optional
|
from typing import Annotated, List, Optional
|
||||||
|
|
||||||
|
import aiofiles
|
||||||
from fastapi import APIRouter, Depends, File, HTTPException, Query, Request, UploadFile, status
|
from fastapi import APIRouter, Depends, File, HTTPException, Query, Request, UploadFile, status
|
||||||
from fastapi.responses import StreamingResponse
|
from fastapi.responses import StreamingResponse
|
||||||
from sqlalchemy import asc, desc
|
from sqlalchemy import asc, desc
|
||||||
@@ -1277,7 +1278,7 @@ async def ui_upload(request: Request, db: DbSession, file: UploadFile = File(...
|
|||||||
# enforcing the size limit during the read so memory usage stays bounded.
|
# enforcing the size limit during the read so memory usage stays bounded.
|
||||||
try:
|
try:
|
||||||
written_size = 0
|
written_size = 0
|
||||||
with open(target_path, "wb") as f:
|
async with aiofiles.open(target_path, "wb") as f:
|
||||||
chunk_size = 65536 # 64 KB chunks
|
chunk_size = 65536 # 64 KB chunks
|
||||||
while True:
|
while True:
|
||||||
chunk = await file.read(chunk_size)
|
chunk = await file.read(chunk_size)
|
||||||
@@ -1286,14 +1287,14 @@ async def ui_upload(request: Request, db: DbSession, file: UploadFile = File(...
|
|||||||
written_size += len(chunk)
|
written_size += len(chunk)
|
||||||
if written_size > max_size:
|
if written_size > max_size:
|
||||||
# Exceeded limit mid-stream; clean up and reject
|
# Exceeded limit mid-stream; clean up and reject
|
||||||
f.close()
|
await f.close()
|
||||||
os.remove(target_path)
|
os.remove(target_path)
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=413,
|
status_code=413,
|
||||||
detail=f"File too large: exceeded {max_size} bytes during upload. "
|
detail=f"File too large: exceeded {max_size} bytes during upload. "
|
||||||
f"See SECURITY_AUDIT.md for configuration details.",
|
f"See SECURITY_AUDIT.md for configuration details.",
|
||||||
)
|
)
|
||||||
f.write(chunk)
|
await f.write(chunk)
|
||||||
except HTTPException:
|
except HTTPException:
|
||||||
raise
|
raise
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -58,3 +58,4 @@ sentry-sdk[fastapi,celery,sqlalchemy]>=2.20.0,<3.0.0
|
|||||||
|
|
||||||
# GraphQL API
|
# GraphQL API
|
||||||
strawberry-graphql[fastapi]>=0.243.0,<1.0.0
|
strawberry-graphql[fastapi]>=0.243.0,<1.0.0
|
||||||
|
aiofiles>=24.1.0 # Asynchronous file I/O support
|
||||||
|
|||||||
@@ -345,7 +345,7 @@ class TestUploadErrorHandling:
|
|||||||
|
|
||||||
def test_upload_disk_write_failure(self, client: TestClient):
|
def test_upload_disk_write_failure(self, client: TestClient):
|
||||||
"""Test handling of disk write failures."""
|
"""Test handling of disk write failures."""
|
||||||
with patch("builtins.open", side_effect=IOError("Disk full")):
|
with patch("aiofiles.open", side_effect=IOError("Disk full")):
|
||||||
pdf_content = b"%PDF-1.4\n%EOF"
|
pdf_content = b"%PDF-1.4\n%EOF"
|
||||||
|
|
||||||
response = client.post(
|
response = client.post(
|
||||||
|
|||||||
Reference in New Issue
Block a user