perf(api): detect MIME types without copying the upload

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.1.12
+
+### Improvements
+
+- Detect unknown upload MIME types from the request's upload stream, passing the uploaded filename as the extension hint and rewinding the stream on success or failure. Once `unstructured` includes [unstructured#4419](https://github.com/Unstructured-IO/unstructured/pull/4419) (expected in 0.27.12), the detector no longer makes a full in-memory copy of the upload.
+
 ## 0.1.11
 
 ### Fixes

prepline_general/api/__version__.py

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

prepline_general/api/filetypes.py

@@ -1,6 +1,4 @@
-import os
 from typing import Optional
-from io import BytesIO
 
 from fastapi import HTTPException, UploadFile
 
@@ -36,15 +34,13 @@ def get_validated_mimetype(file: UploadFile, content_type_hint: str | None = Non
 
     filetype = FileType.from_mime_type(content_type)
 
-    # If content_type was not specified, use the library to identify the file
-    # We inspect the bytes to do this, so we need to buffer the file
+    # If content_type was not specified, use the library to identify the file. The upload's own
+    # `.name` is its temporary file, so the uploaded filename is passed as the extension hint.
     if not filetype or filetype == FileType.UNK:
-        file_buffer = BytesIO(file.file.read())
-        file.file.seek(0)
-
-        file_buffer.name = file.filename
-
-        filetype = detect_filetype(file=file_buffer)
+        try:
+            filetype = detect_filetype(file=file.file, metadata_file_path=file.filename)
+        finally:
+            file.file.seek(0)
 
     if not filetype.is_partitionable:
         raise HTTPException(

test_general/api/test_filetypes.py

@@ -0,0 +1,96 @@
+from io import BytesIO
+from pathlib import Path
+from tempfile import SpooledTemporaryFile
+
+import pytest
+from fastapi import UploadFile
+
+from prepline_general.api import filetypes
+from unstructured.file_utils.model import FileType
+
+
+def test_unknown_mimetype_is_detected_from_existing_upload_stream(monkeypatch):
+    upload_stream = SpooledTemporaryFile()
+    upload_stream.write(b"sample text")
+    upload_stream.seek(0)
+    upload = UploadFile(file=upload_stream, filename="sample.txt")
+
+    def fake_detect_filetype(*, file, metadata_file_path):
+        assert file is upload_stream
+        assert metadata_file_path == "sample.txt"
+        file.seek(4)
+        return FileType.TXT
+
+    monkeypatch.setattr(filetypes, "detect_filetype", fake_detect_filetype)
+
+    assert filetypes.get_validated_mimetype(upload) == "text/plain"
+    assert upload_stream.tell() == 0
+
+
+def test_unknown_mimetype_rewinds_upload_stream_when_detection_fails(monkeypatch):
+    upload_stream = SpooledTemporaryFile()
+    upload_stream.write(b"sample text")
+    upload_stream.seek(0)
+    upload = UploadFile(file=upload_stream, filename="sample.txt")
+
+    def fake_detect_filetype(*, file, metadata_file_path):
+        file.seek(4)
+        raise RuntimeError("detection failed")
+
+    monkeypatch.setattr(filetypes, "detect_filetype", fake_detect_filetype)
+
+    with pytest.raises(RuntimeError, match="detection failed"):
+        filetypes.get_validated_mimetype(upload)
+
+    assert upload_stream.tell() == 0
+
+
+def _assert_matches_copied_upload(filename: str | None, payload: bytes, max_size: int):
+    copied = BytesIO(payload)
+    copied.name = filename
+    expected = filetypes.detect_filetype(file=copied)
+    with SpooledTemporaryFile(max_size=max_size) as stream:
+        stream.write(payload)
+        stream.seek(0)
+        upload = UploadFile(file=stream, filename=filename)
+        assert filetypes.get_validated_mimetype(upload) == expected.mime_type
+        assert stream.tell() == 0
+        assert not stream.closed
+
+
+@pytest.mark.parametrize("max_size", [1, 1024 * 1024])
+@pytest.mark.parametrize(
+    "filename,payload",
+    [
+        ("sample.txt", b"A sample paragraph of ordinary text."),
+        ("sample.html", b"<!doctype html><html><body><p>Hello</p></body></html>"),
+        ("sample.csv", b"name,value\nAlice,1\nBob,2\n"),
+        ("sample.json", b'{"name": "Alice", "value": 1}'),
+        (None, b"A sample paragraph of ordinary text."),
+    ],
+)
+def test_real_detector_matches_copied_upload(filename, payload, max_size):
+    _assert_matches_copied_upload(filename, payload, max_size)
+
+
+@pytest.mark.parametrize("max_size", [1, 1024 * 1024])
+@pytest.mark.parametrize(
+    "filename,payload,expected",
+    [
+        ("sample.md", b"# Heading\n\nSome *markdown* text.\n", FileType.MD),
+        ("sample.tsv", b"name\tvalue\nAlice\t1\nBob\t2\n", FileType.TSV),
+    ],
+)
+def test_real_detector_uses_uploaded_filename_extension(filename, payload, expected, max_size):
+    with SpooledTemporaryFile(max_size=max_size) as stream:
+        stream.write(payload)
+        stream.seek(0)
+        upload = UploadFile(file=stream, filename=filename)
+        assert filetypes.get_validated_mimetype(upload) == expected.mime_type
+
+
+@pytest.mark.parametrize("filename", ["layout-parser-paper.pdf", "notes.pptx", "stanley-cups.xlsx"])
+@pytest.mark.parametrize("max_size", [1, 10 * 1024 * 1024])
+def test_real_detector_matches_copied_binary_upload(filename, max_size):
+    payload = (Path("sample-docs") / filename).read_bytes()
+    _assert_matches_copied_upload(filename, payload, max_size)