fix(gzip): bound decompression memory

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.1.11
+
+### Fixes
+
+- **Bound gzip decompression memory**: gzip uploads are decompressed incrementally into a spooled temporary file instead of being materialized as a complete in-memory `bytes` value. Expanded outputs larger than 1 MiB remain disk-backed through MIME detection and partitioning, avoiding downstream document-sized copies while preserving the partition API.
+
 ## 0.1.10
 
 ### Fixes

prepline_general/api/__version__.py

@@ -1 +1 @@
-__version__ = "0.1.10"  # pragma: no cover
+__version__ = "0.1.11"  # pragma: no cover

prepline_general/api/general.py

@@ -7,10 +7,12 @@
 import mimetypes
 import os
 import secrets
+import shutil
+import tempfile
 from base64 import b64encode
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
-from typing import IO, Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union, cast
+from typing import IO, Any, BinaryIO, Dict, List, Mapping, Optional, Sequence, Tuple, Union, cast
 
 import backoff
 import pandas as pd
@@ -31,6 +33,7 @@
 from starlette.datastructures import Headers
 from starlette.types import Send
 
+from prepline_general.api import __version__ as api_version
 from prepline_general.api.filetypes import get_validated_mimetype
 from prepline_general.api.models.form_params import GeneralFormParams
 from unstructured.documents.elements import Element
@@ -41,11 +44,37 @@
     elements_from_json,
 )
 from unstructured_inference.models.base import UnknownModelException
-from prepline_general.api import __version__ as api_version
 
 app = FastAPI()
 router = APIRouter()
 
+_GZIP_COPY_CHUNK_SIZE = 1024 * 1024
+_GZIP_SPOOL_MAX_MEMORY_BYTES = 1024 * 1024
+
+
+class _SpooledFileProxy:
+    """Expose a spooled file without its concrete type triggering downstream copies."""
+
+    def __init__(self, file: IO[bytes], name: str | None):
+        self._file = file
+        self.name = name
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._file, name)
+
+    def __enter__(self) -> _SpooledFileProxy:
+        self._file.__enter__()
+        return self
+
+    def __exit__(self, *args: Any) -> Any:
+        return self._file.__exit__(*args)
+
+    def __iter__(self):
+        return iter(self._file)
+
+    def __next__(self) -> bytes:
+        return next(self._file)
+
 
 def is_compatible_response_type(media_type: str, response_type: type) -> bool:
     """True when `response_type` can be converted to `media_type` for HTTP Response."""
@@ -613,13 +642,30 @@ def return_content_type(filename: str):
     if filename.endswith(".gz"):
         filename = filename[:-3]
 
-    gzip_file = gzip.open(file.file).read()
-    return UploadFile(
-        file=io.BytesIO(gzip_file),
-        size=len(gzip_file),
-        filename=filename,
-        headers=Headers({"content-type": return_content_type(filename)}),
+    output_file = tempfile.SpooledTemporaryFile(max_size=_GZIP_SPOOL_MAX_MEMORY_BYTES)
+    try:
+        with gzip.open(file.file) as gzip_file:
+            shutil.copyfileobj(gzip_file, output_file, length=_GZIP_COPY_CHUNK_SIZE)
+        uncompressed_size = output_file.tell()
+        output_file.seek(0)
+    except Exception:
+        output_file.close()
+        raise
+
+    # Reuse the request-owned UploadFile so FastAPI closes the decompressed spool when the
+    # request finishes. A newly-created UploadFile would not belong to the parsed FormData and
+    # would therefore remain open after the response.
+    file.file.close()
+    decompressed_file = (
+        _SpooledFileProxy(output_file, filename)
+        if uncompressed_size > _GZIP_SPOOL_MAX_MEMORY_BYTES
+        else output_file
     )
+    file.file = cast(BinaryIO, decompressed_file)
+    file.size = uncompressed_size
+    file.filename = filename
+    file.headers = Headers({"content-type": return_content_type(filename)})
+    return file
 
 
 @router.get("/general/v0/general", include_in_schema=False)
@@ -747,7 +793,7 @@ def join_responses(
         frames = [
             pd.read_csv(io.BytesIO(response.body))  # pyright: ignore[reportUnknownMemberType]
             for response in responses
-            if response.body.strip()
+            if cast(bytes, response.body).strip()
         ]
         if not frames:
             return PlainTextResponse(responses[0].body)

test_general/api/test_gzip.py

@@ -9,13 +9,58 @@
 import pandas as pd
 import pytest
 from deepdiff import DeepDiff
+from fastapi import UploadFile
 from fastapi.testclient import TestClient
 
 from prepline_general.api.app import app
+from prepline_general.api.general import _GZIP_SPOOL_MAX_MEMORY_BYTES, ungz_file
 
 MAIN_API_ROUTE = "general/v0/general"
 
 
+def _gzip_upload(content: bytes, filename: str = "sample.txt.gz") -> UploadFile:
+    compressed = io.BytesIO(gzip.compress(content))
+    return UploadFile(file=compressed, filename=filename)
+
+
+@pytest.mark.parametrize(
+    ("content", "spills_to_disk"),
+    [
+        pytest.param(b"small gzip payload", False, id="small-stays-in-memory"),
+        pytest.param(b"x" * (_GZIP_SPOOL_MAX_MEMORY_BYTES + 1), True, id="large-spills-to-disk"),
+    ],
+)
+def test_ungz_file_bounds_decompression_memory(content: bytes, spills_to_disk: bool):
+    upload = _gzip_upload(content)
+    compressed_file = upload.file
+
+    result = ungz_file(upload)
+
+    try:
+        assert result is upload
+        assert compressed_file.closed
+        assert result.filename == "sample.txt"
+        assert result.content_type == "text/plain"
+        assert result.size == len(content)
+        assert result.file.read() == content
+        assert result.file._rolled is spills_to_disk
+        assert isinstance(result.file, tempfile.SpooledTemporaryFile) is not spills_to_disk
+    finally:
+        result.file.close()
+        assert result.file.closed
+
+
+def test_ungz_file_proxy_preserves_special_file_methods():
+    result = ungz_file(_gzip_upload(b"first\nsecond\n" * _GZIP_SPOOL_MAX_MEMORY_BYTES))
+
+    try:
+        assert iter(result.file) is not result.file
+        assert next(result.file) == b"first\n"
+        assert result.file.__enter__() is result.file
+    finally:
+        result.file.close()
+
+
 @pytest.mark.xfail(reason="The outputs are different as of unstructured==0.13.5")
 @pytest.mark.parametrize("output_format", ["application/json", "text/csv"])
 @pytest.mark.parametrize(