Move archive extraction from memory to stream (#965)

shelfmark/download/archive.py

@@ -1,18 +1,21 @@
 """Archive extraction utilities for downloaded book archives."""
 
+import shutil
+import tempfile
 import zipfile
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from shelfmark.core.logger import setup_logger
 from shelfmark.core.utils import is_audiobook as check_audiobook
-from shelfmark.download.fs import atomic_write
+from shelfmark.download.fs import atomic_move
 from shelfmark.download.postprocess.policy import (
     get_supported_audiobook_formats,
     get_supported_formats,
 )
 
 logger = setup_logger(__name__)
+_ARCHIVE_COPY_CHUNK_SIZE = 1024 * 1024
 
 if TYPE_CHECKING:
     import rarfile
@@ -208,9 +211,25 @@ def _extract_files_from_archive(archive: ArchiveType, output_dir: Path) -> list[
             logger.warning("Path traversal attempt blocked: %r", info.filename)
             continue
 
-        with archive.open(info) as src:
-            data = src.read()
-        final_path = atomic_write(target_path, data)
+        temp_path: Path | None = None
+        try:
+            with (
+                archive.open(info) as src,
+                tempfile.NamedTemporaryFile(
+                    dir=output_dir,
+                    prefix=".shelfmark-extract-",
+                    suffix=".tmp",
+                    delete=False,
+                ) as temp_file,
+            ):
+                temp_path = Path(temp_file.name)
+                shutil.copyfileobj(src, temp_file, length=_ARCHIVE_COPY_CHUNK_SIZE)
+
+            final_path = atomic_move(cast("Path", temp_path), target_path)
+        except Exception:
+            if temp_path is not None:
+                temp_path.unlink(missing_ok=True)
+            raise
         extracted_files.append(final_path)
         logger.debug("Extracted: %s", filename)
 

tests/download/test_archive_extraction.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from shelfmark.download import archive as archive_mod
+
+
+class _FakeZipInfo:
+    filename = "book.epub"
+    flag_bits = 0
+
+    def is_dir(self) -> bool:
+        return False
+
+
+class _ChunkOnlyStream:
+    def __init__(self, content: bytes) -> None:
+        self._content = content
+        self._offset = 0
+        self.whole_read_called = False
+
+    def __enter__(self) -> _ChunkOnlyStream:
+        return self
+
+    def __exit__(self, *_args: object) -> None:
+        return None
+
+    def read(self, size: int = -1) -> bytes:
+        if size < 0:
+            self.whole_read_called = True
+            msg = "archive member was read into memory"
+            raise AssertionError(msg)
+
+        chunk = self._content[self._offset : self._offset + size]
+        self._offset += len(chunk)
+        return chunk
+
+
+class _FakeZipFile:
+    stream: _ChunkOnlyStream
+
+    def __init__(self, _path: Path, _mode: str) -> None:
+        self.stream = _ChunkOnlyStream(b"streamed archive content")
+
+    def __enter__(self) -> _FakeZipFile:
+        return self
+
+    def __exit__(self, *_args: object) -> None:
+        return None
+
+    def infolist(self) -> list[_FakeZipInfo]:
+        return [_FakeZipInfo()]
+
+    def testzip(self) -> None:
+        return None
+
+    def open(self, _info: _FakeZipInfo) -> _ChunkOnlyStream:
+        return self.stream
+
+
+def test_extract_archive_raw_streams_members_without_whole_read(
+    tmp_path: Path, monkeypatch
+) -> None:
+    fake_archives: list[_FakeZipFile] = []
+
+    def fake_zip_file(path: Path, mode: str) -> _FakeZipFile:
+        archive = _FakeZipFile(path, mode)
+        fake_archives.append(archive)
+        return archive
+
+    monkeypatch.setattr(archive_mod.zipfile, "ZipFile", fake_zip_file)
+
+    extracted_files, warnings = archive_mod.extract_archive_raw(tmp_path / "book.zip", tmp_path)
+
+    assert warnings == []
+    assert [path.name for path in extracted_files] == ["book.epub"]
+    assert extracted_files[0].read_bytes() == b"streamed archive content"
+    assert fake_archives[0].stream.whole_read_called is False