Fix multi rom zipping + streaming

backend/endpoints/rom.py

@@ -1,11 +1,11 @@
-import io
-import tempfile
-import zipfile
+from datetime import datetime
 from fastapi import APIRouter, Request, status, HTTPException
 from fastapi_pagination.ext.sqlalchemy import paginate
 from fastapi_pagination.cursor import CursorPage, CursorParams
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, StreamingResponse
 from pydantic import BaseModel, BaseConfig
+from stat import S_IFREG
+from stream_zip import ZIP_64, stream_zip
 
 from logger.logger import log
 from handler import dbh
@@ -76,31 +76,24 @@ def download_rom(id: int, files: str):
     rom_path = f"{LIBRARY_BASE_PATH}/{rom.full_path}"
 
     if not rom.multi:
-        return FileResponse(
-            path=rom_path,
-            filename=rom.file_name,
-            media_type="application/octet-stream",
-        )
+        return FileResponse(path=rom_path, filename=rom.file_name)
 
-    mf = io.BytesIO()
-    with zipfile.ZipFile(mf, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
-        try:
-            for file_name in files.split(","):
-                zf.write(f"{rom_path}/{file_name}", file_name)
-        except FileNotFoundError as e:
-            log.error(str(e))
-        finally:
-            zf.close()
+    # Builds a generator of tuples for each member file
+    def local_files():
+        def contents(file_name):
+            with open(f"{rom_path}/{file_name}", "rb") as f:
+                while chunk := f.read(65536):
+                    yield chunk
 
-    tmp = tempfile.NamedTemporaryFile(delete=False)
-    with open(tmp.name, "wb") as f:
-        f.write(mf.getvalue())
+        return [
+            (file_name, datetime.now(), S_IFREG | 0o600, ZIP_64, contents(file_name))
+            for file_name in files.split(",")
+        ]
 
-        return FileResponse(
-            path=tmp.name,
-            filename=f"{rom.r_name}.zip",
-            media_type="application/octet-stream",
-        )
+    zipped_chunks = stream_zip(local_files())
+
+    # Streams the zip file to the client
+    return StreamingResponse(zipped_chunks, media_type="application/zip")
 
 
 @router.get("/platforms/{p_slug}/roms", status_code=200)

poetry.lock

@@ -1490,6 +1490,21 @@ anyio = ">=3.4.0,<5"
 [package.extras]
 full = ["httpx (>=0.22.0)", "itsdangerous", "jinja2", "python-multipart", "pyyaml"]
 
+[[package]]
+name = "stream-zip"
+version = "0.0.67"
+description = "Python function to construct a ZIP archive with stream processing - without having to store the entire ZIP in memory or disk"
+optional = false
+python-versions = ">=3.6.7"
+files = [
+    {file = "stream_zip-0.0.67-py3-none-any.whl", hash = "sha256:e4200f9f18f48f759742dd9bc980f4712568dd0815ff618b3facf8573aba3800"},
+    {file = "stream_zip-0.0.67.tar.gz", hash = "sha256:492ddfb338701fbf4d744126027515c9ae8daca692c0acbc1719dc37ba15f634"},
+]
+
+[package.extras]
+ci = ["coverage (==6.2)", "pytest (==6.2.5)", "pytest-cov (==3.0.0)", "stream-unzip (==0.0.86)"]
+dev = ["coverage (>=6.2)", "pytest (>=6.2.5)", "pytest-cov (>=3.0.0)", "stream-unzip (>=0.0.86)"]
+
 [[package]]
 name = "tomli"
 version = "2.0.1"
@@ -1918,4 +1933,4 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "8000e74d2aeed13c72744d23d6075103398d8d255ac3f7317683920f7441015f"
+content-hash = "c973b508a9bf80ff3875173c2c9b505e5c598fba6ef82f432936ce2a270cb292"

pyproject.toml

@@ -33,6 +33,7 @@ types-pyyaml = "^6.0.12.11"
 types-requests = "^2.31.0.2"
 mypy = "^1.4.1"
 types-redis = "^4.6.0.3"
+stream-zip = "^0.0.67"
 
 [build-system]
 requires = ["poetry-core"]