Feature: add ZipExportSink with atomic finalize and manifest spooling

src/documents/export/sinks.py

@@ -3,10 +3,17 @@ from __future__ import annotations
 import abc
 import hashlib
 import json
+import os
+import shutil
+import tempfile
+import zipfile
 from contextlib import AbstractContextManager
 from contextlib import contextmanager
+from pathlib import Path
+from pathlib import PurePosixPath
 from typing import TYPE_CHECKING
 
+from django.conf import settings
 from django.core.serializers.json import DjangoJSONEncoder
 
 from documents.file_handling import delete_empty_directories
@@ -15,7 +22,6 @@ from documents.utils import copy_file_with_basic_stats
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
-    from pathlib import Path
     from typing import TextIO
 
 
@@ -216,3 +222,114 @@ class DirectoryExportSink(ExportSink):
         # Folder mode is in-place/incremental: streamed .tmp files are already
         # cleaned in stream(); leave everything else intact and skip the prune.
         return None
+
+
+class ZipExportSink(ExportSink):
+    """Writes a single zip archive, produced atomically only on success.
+
+    Builds into ``<target>/<zip_name>.zip.tmp`` and renames to ``.zip`` on clean
+    finalize. The manifest stream is spooled to a temp file in SCRATCH_DIR and
+    added as an entry at finalize (a zip entry cannot be interleaved with others).
+    """
+
+    def __init__(self, target: Path, zip_name: str, *, delete: bool = False) -> None:
+        self._target = target.resolve()
+        self._zip_path = (self._target / zip_name).with_suffix(".zip")
+        self._tmp_path = self._zip_path.with_name(self._zip_path.name + ".tmp")
+        self._delete = delete
+        self._zip: zipfile.ZipFile | None = None
+        self._dirs: set[str] = set()
+        self._pending_manifest: tuple[Path, str] | None = None
+        self._stream_open = False
+
+    def _open(self) -> None:
+        settings.SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
+        self._zip = zipfile.ZipFile(
+            self._tmp_path,
+            "w",
+            compression=zipfile.ZIP_DEFLATED,
+            allowZip64=True,
+        )
+
+    def _ensure_dirs(self, arcname: str) -> None:
+        parts = PurePosixPath(arcname).parts[:-1]
+        for i in range(len(parts)):
+            dir_arc = "/".join(parts[: i + 1]) + "/"
+            if dir_arc not in self._dirs:
+                self._dirs.add(dir_arc)
+                assert self._zip is not None
+                self._zip.mkdir(dir_arc)
+
+    def add_file(
+        self,
+        source: Path,
+        arcname: str,
+        *,
+        checksum: str | None = None,
+    ) -> None:
+        assert self._zip is not None
+        self._ensure_dirs(arcname)
+        self._zip.write(source, arcname=arcname)
+
+    def add_json(self, content: list | dict, arcname: str) -> None:
+        assert self._zip is not None
+        self._ensure_dirs(arcname)
+        self._zip.writestr(arcname, _dumps(content))
+
+    @contextmanager
+    def stream(self, arcname: str) -> Iterator[TextIO]:
+        if self._stream_open:
+            raise RuntimeError("A stream is already open on this sink")
+        settings.SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
+        fd, tmp_name = tempfile.mkstemp(
+            dir=settings.SCRATCH_DIR,
+            prefix="export-manifest-",
+            suffix=".json",
+        )
+        tmp = Path(tmp_name)
+        handle = os.fdopen(fd, "w", encoding="utf-8")
+        self._stream_open = True
+        try:
+            yield handle
+        except BaseException:
+            handle.close()
+            tmp.unlink(missing_ok=True)
+            raise
+        else:
+            handle.close()
+            self._pending_manifest = (tmp, arcname)
+        finally:
+            self._stream_open = False
+
+    def _finalize(self) -> None:
+        assert self._zip is not None
+        if self._pending_manifest is not None:
+            tmp, arcname = self._pending_manifest
+            self._ensure_dirs(arcname)
+            self._zip.write(tmp, arcname=arcname)
+            tmp.unlink(missing_ok=True)
+            self._pending_manifest = None
+        self._zip.close()
+        self._zip = None
+        if self._delete:
+            self._wipe_destination()
+        self._tmp_path.rename(self._zip_path)
+
+    def _wipe_destination(self) -> None:
+        skip = {self._zip_path.resolve(), self._tmp_path.resolve()}
+        for item in self._target.glob("*"):
+            if item.resolve() in skip:
+                continue
+            if item.is_dir():
+                shutil.rmtree(item)
+            else:
+                item.unlink()
+
+    def _abort(self) -> None:
+        if self._zip is not None:
+            self._zip.close()
+            self._zip = None
+        self._tmp_path.unlink(missing_ok=True)
+        if self._pending_manifest is not None:
+            self._pending_manifest[0].unlink(missing_ok=True)
+            self._pending_manifest = None

src/documents/tests/export/test_sinks.py

@@ -1,12 +1,15 @@
 import io
 import json
 import os
+import zipfile
 from pathlib import Path
 
 import pytest
 
 from documents.export.sinks import DirectoryExportSink
+from documents.export.sinks import ExportSink
 from documents.export.sinks import StreamingManifestWriter
+from documents.export.sinks import ZipExportSink
 from documents.export.sinks import _dumps
 
 
@@ -168,3 +171,125 @@ class TestDirectoryExportSink:
         ) as sink:
             sink.add_file(source_file, "originals/doc.pdf")
         assert stale.exists()
+
+
+class TestZipExportSink:
+    @pytest.fixture()
+    def source_file(self, tmp_path: Path) -> Path:
+        src: Path = tmp_path / "src" / "doc.pdf"
+        src.parent.mkdir(parents=True)
+        src.write_bytes(b"PDF-CONTENT")
+        return src
+
+    def test_round_trip_files_json_and_stream(
+        self,
+        tmp_path: Path,
+        source_file: Path,
+    ) -> None:
+        target: Path = tmp_path / "out"
+        target.mkdir()
+        with ZipExportSink(target, "export", delete=False) as sink:
+            sink.add_file(source_file, "originals/doc.pdf")
+            sink.add_json({"version": "x"}, "metadata.json")
+            with sink.stream("manifest.json") as handle:
+                writer = StreamingManifestWriter(handle)
+                writer.write_record({"pk": 1})
+                writer.close()
+        zip_path: Path = target / "export.zip"
+        assert zip_path.exists()
+        assert not (target / "export.zip.tmp").exists()
+        with zipfile.ZipFile(zip_path) as zf:
+            names = set(zf.namelist())
+            assert {"originals/doc.pdf", "metadata.json", "manifest.json"} <= names
+            assert zf.read("originals/doc.pdf") == b"PDF-CONTENT"
+            assert json.loads(zf.read("manifest.json")) == [{"pk": 1}]
+
+    def test_nested_arcname_emits_directory_marker(
+        self,
+        tmp_path: Path,
+        source_file: Path,
+    ) -> None:
+        target: Path = tmp_path / "out"
+        target.mkdir()
+        with ZipExportSink(target, "export", delete=False) as sink:
+            sink.add_file(source_file, "originals/doc.pdf")
+        with zipfile.ZipFile(target / "export.zip") as zf:
+            assert "originals/" in zf.namelist()
+
+    def test_flat_arcname_has_no_directory_markers(
+        self,
+        tmp_path: Path,
+        source_file: Path,
+    ) -> None:
+        target: Path = tmp_path / "out"
+        target.mkdir()
+        with ZipExportSink(target, "export", delete=False) as sink:
+            sink.add_file(source_file, "doc.pdf")
+        with zipfile.ZipFile(target / "export.zip") as zf:
+            assert all(not n.endswith("/") for n in zf.namelist())
+
+    def test_exception_leaves_no_zip_and_no_tmp(
+        self,
+        tmp_path: Path,
+        source_file: Path,
+    ) -> None:
+        target: Path = tmp_path / "out"
+        target.mkdir()
+        with pytest.raises(RuntimeError):
+            with ZipExportSink(target, "export", delete=False) as sink:
+                sink.add_file(source_file, "doc.pdf")
+                raise RuntimeError("boom")
+        assert not (target / "export.zip").exists()
+        assert not (target / "export.zip.tmp").exists()
+
+    def test_delete_wipes_destination_on_success(
+        self,
+        tmp_path: Path,
+        source_file: Path,
+    ) -> None:
+        target: Path = tmp_path / "out"
+        target.mkdir()
+        (target / "preexisting.txt").write_text("old")
+        (target / "olddir").mkdir()
+        with ZipExportSink(target, "export", delete=True) as sink:
+            sink.add_file(source_file, "doc.pdf")
+        assert (target / "export.zip").exists()
+        assert not (target / "preexisting.txt").exists()
+        assert not (target / "olddir").exists()
+
+    def test_abort_with_delete_does_not_wipe_destination(
+        self,
+        tmp_path: Path,
+        source_file: Path,
+    ) -> None:
+        target: Path = tmp_path / "out"
+        target.mkdir()
+        (target / "preexisting.txt").write_text("old")
+        with pytest.raises(RuntimeError):
+            with ZipExportSink(target, "export", delete=True) as sink:
+                sink.add_file(source_file, "doc.pdf")
+                raise RuntimeError("boom")
+        assert (target / "preexisting.txt").exists()
+        assert not (target / "export.zip").exists()
+
+
+class TestStreamContract:
+    @pytest.fixture(params=["dir", "zip"])
+    def sink(self, request: pytest.FixtureRequest, tmp_path: Path) -> ExportSink:
+        target: Path = tmp_path / "out"
+        target.mkdir()
+        if request.param == "dir":
+            return DirectoryExportSink(
+                target,
+                compare_checksums=False,
+                compare_json=False,
+                delete=False,
+            )
+        return ZipExportSink(target, "export", delete=False)
+
+    def test_second_concurrent_stream_is_rejected(self, sink: ExportSink) -> None:
+        with sink:
+            with sink.stream("manifest.json"):
+                with pytest.raises(RuntimeError, match="already open"):
+                    with sink.stream("other.json"):
+                        pass