From 0e538acdef8c4c27e12c88dda4d43ed7052c644a Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Wed, 22 Jul 2026 15:36:50 -0700 Subject: [PATCH] Feature: add ZipExportSink with atomic finalize and manifest spooling --- src/documents/export/sinks.py | 119 ++++++++++++++++++++- src/documents/tests/export/test_sinks.py | 125 +++++++++++++++++++++++ 2 files changed, 243 insertions(+), 1 deletion(-) diff --git a/src/documents/export/sinks.py b/src/documents/export/sinks.py index b9fe55263..d33002be7 100644 --- a/src/documents/export/sinks.py +++ b/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 ``/.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 diff --git a/src/documents/tests/export/test_sinks.py b/src/documents/tests/export/test_sinks.py index 0a6577a57..a61b3bb84 100644 --- a/src/documents/tests/export/test_sinks.py +++ b/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