From cdda191f91058238b7e7487e5b552319e2ff5463 Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Thu, 23 Jul 2026 13:27:06 -0700 Subject: [PATCH] Feature: Allow configuring the compression type and compression levels during export Building on the zip export improvements, this now allows users to further configure the zip to fit their needs. A simple stored zip for speed, or a high compression zstd for the smallest archive. Full validation of the method and levels at the command line --- docs/administration.md | 15 ++ src/documents/export/compression.py | 108 ++++++++++ src/documents/export/sinks.py | 15 +- .../management/commands/document_exporter.py | 54 +++++ .../management/commands/document_importer.py | 16 ++ .../tests/export/test_compression.py | 190 ++++++++++++++++++ src/documents/tests/export/test_sinks.py | 43 ++++ .../tests/test_management_exporter.py | 182 +++++++++++++++++ .../tests/test_management_importer.py | 29 +++ 9 files changed, 650 insertions(+), 2 deletions(-) create mode 100644 src/documents/export/compression.py create mode 100644 src/documents/tests/export/test_compression.py diff --git a/docs/administration.md b/docs/administration.md index fbc636935..6d668241f 100644 --- a/docs/administration.md +++ b/docs/administration.md @@ -299,6 +299,8 @@ optional arguments: -sm, --split-manifest -z, --zip -zn, --zip-name +--zip-compression +--zip-compression-level --data-only --no-progress-bar --passphrase @@ -361,6 +363,19 @@ If `-z` or `--zip` is provided, the export will be a zip file in the target directory, named according to the current local date or the value set in `-zn` or `--zip-name`. +The compression method for the zip can be set with `--zip-compression` +(`stored`, `deflated` (default), `bzip2`, `lzma`, or `zstd`) and tuned with +`--zip-compression-level` (deflated: 0–9, bzip2: 1–9, zstd: -22–22; ignored +for `stored` and `lzma`). Both options require `--zip`. + +!!! warning + + `zstd` compression requires Python 3.14 or newer on **both** the machine + creating the export and any machine importing it. An archive compressed with + `zstd` (or `lzma`/`bzip2` where those modules are unavailable) cannot be + imported on a runtime that lacks the codec; the importer will refuse it with + a clear error. The default `deflated` is universally readable. + If `--data-only` is provided, only the database will be exported. This option is intended to facilitate database upgrades without needing to clean documents and thumbnails from the media directory. diff --git a/src/documents/export/compression.py b/src/documents/export/compression.py new file mode 100644 index 000000000..e61eebac7 --- /dev/null +++ b/src/documents/export/compression.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import importlib +import zipfile + +# ZIP_ZSTANDARD exists only on Python 3.14+ (PEP 784). None elsewhere. +ZSTD: int | None = getattr(zipfile, "ZIP_ZSTANDARD", None) + +# CLI choices are fixed across runtimes so argparse never hides zstd; runtime +# availability is enforced separately in compression_available(). +COMPRESSION_CHOICES: tuple[str, ...] = ( + "stored", + "deflated", + "bzip2", + "lzma", + "zstd", +) + +# Method name -> zipfile compression constant (zstd only when supported). +COMPRESSION_METHODS: dict[str, int] = { + "stored": zipfile.ZIP_STORED, + "deflated": zipfile.ZIP_DEFLATED, + "bzip2": zipfile.ZIP_BZIP2, + "lzma": zipfile.ZIP_LZMA, +} +if ZSTD is not None: + COMPRESSION_METHODS["zstd"] = ZSTD + +# Inclusive (min, max) level bounds per method; None => level not applicable. +# Verified on CPython 3.14.3. +# +# zstd's raw library bounds are (-131072, 22) +# (compression.zstd.CompressionParameter.compression_level.bounds()) — the +# minimum is an internal implementation constant (-ZSTD_TARGETLENGTH_MAX), +# not a meaningful distinct "level"; deeper negative values than -22 buy +# nothing over -22 in practice. We expose the conventional zstd CLI range +# instead of the raw library bounds. +LEVEL_BOUNDS: dict[str, tuple[int, int] | None] = { + "stored": None, + "deflated": (0, 9), + "bzip2": (1, 9), + "lzma": None, + "zstd": (-22, 22), +} + +# zipfile compress_type id -> method name. 93 = current zstd id, 20 = legacy +# zstd id that zipfile can still read. +_COMPRESS_TYPE_TO_METHOD: dict[int, str] = { + zipfile.ZIP_STORED: "stored", + zipfile.ZIP_DEFLATED: "deflated", + zipfile.ZIP_BZIP2: "bzip2", + zipfile.ZIP_LZMA: "lzma", + 93: "zstd", + 20: "zstd", +} + + +def compression_available(method: str) -> bool: + """Whether the running interpreter can actually use the given method.""" + if method in ("stored", "deflated"): + # zlib is a hard CPython dependency; stored needs nothing. + return True + if method == "bzip2": + return _module_importable("bz2") + if method == "lzma": + return _module_importable("lzma") + if method == "zstd": + return ZSTD is not None and _module_importable("compression.zstd") + return False + + +def _module_importable(name: str) -> bool: + try: + importlib.import_module(name) + except ImportError: + return False + return True + + +def level_error(method: str, level: int | None) -> str | None: + """Return a human message if (method, level) is invalid, else None.""" + if level is None: + return None + bounds = LEVEL_BOUNDS[method] + if bounds is None: + return f"--zip-compression-level has no effect for '{method}'" + low, high = bounds + if not (low <= level <= high): + return ( + f"--zip-compression-level for '{method}' must be between {low} and {high}" + ) + return None + + +def compress_type_readable(compress_type: int) -> bool: + """Whether this interpreter can decompress an entry of the given type.""" + method = _COMPRESS_TYPE_TO_METHOD.get(compress_type) + if method is None: + return False + return compression_available(method) + + +def unreadable_method_names(compress_types: set[int]) -> set[str]: + """Map a set of compress_type ids to human method names for error messages.""" + names: set[str] = set() + for ct in compress_types: + names.add(_COMPRESS_TYPE_TO_METHOD.get(ct, f"method {ct}")) + return names diff --git a/src/documents/export/sinks.py b/src/documents/export/sinks.py index 43b666cbe..a2e85d116 100644 --- a/src/documents/export/sinks.py +++ b/src/documents/export/sinks.py @@ -243,11 +243,21 @@ class ZipExportSink(ExportSink): 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: + def __init__( + self, + target: Path, + zip_name: str, + *, + delete: bool = False, + compression: int = zipfile.ZIP_DEFLATED, + compresslevel: int | None = None, + ) -> 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._compression = compression + self._compresslevel = compresslevel self._zip: zipfile.ZipFile | None = None self._dirs: set[str] = set() self._pending_manifest: tuple[Path, str] | None = None @@ -258,7 +268,8 @@ class ZipExportSink(ExportSink): self._zip = zipfile.ZipFile( self._tmp_path, "w", - compression=zipfile.ZIP_DEFLATED, + compression=self._compression, + compresslevel=self._compresslevel, allowZip64=True, ) diff --git a/src/documents/management/commands/document_exporter.py b/src/documents/management/commands/document_exporter.py index da98e88cd..320605235 100644 --- a/src/documents/management/commands/document_exporter.py +++ b/src/documents/management/commands/document_exporter.py @@ -29,6 +29,11 @@ if TYPE_CHECKING: if settings.AUDIT_LOG_ENABLED: from auditlog.models import LogEntry +from documents.export.compression import COMPRESSION_CHOICES +from documents.export.compression import COMPRESSION_METHODS +from documents.export.compression import ZSTD +from documents.export.compression import compression_available +from documents.export.compression import level_error from documents.export.sinks import DirectoryExportSink from documents.export.sinks import ExportSink from documents.export.sinks import StreamingManifestWriter @@ -192,6 +197,28 @@ class Command(CryptMixin, PaperlessCommand): help="Sets the export zip file name", ) + parser.add_argument( + "--zip-compression", + choices=COMPRESSION_CHOICES, + default=None, + help=( + "Compression method for the export zip (requires --zip). " + "Default: deflated. 'zstd' requires Python 3.14+ on both the " + "exporting and importing machine." + ), + ) + + parser.add_argument( + "--zip-compression-level", + type=int, + default=None, + help=( + "Compression level for the export zip (requires --zip). " + "deflated: 0-9, bzip2: 1-9, zstd: -22..22; ignored for " + "stored/lzma." + ), + ) + parser.add_argument( "--data-only", default=False, @@ -247,12 +274,39 @@ class Command(CryptMixin, PaperlessCommand): if not os.access(self.target, os.W_OK): raise CommandError("That path doesn't appear to be writable") + zip_compression: str | None = options["zip_compression"] + zip_compression_level: int | None = options["zip_compression_level"] + + if not self.zip_export and ( + zip_compression is not None or zip_compression_level is not None + ): + raise CommandError( + "--zip-compression and --zip-compression-level require --zip", + ) + + compression_method = zip_compression or "deflated" + if self.zip_export: + if not compression_available(compression_method): + if compression_method == "zstd" and ZSTD is None: + raise CommandError( + "zstd compression requires Python 3.14 or newer", + ) + raise CommandError( + f"Compression method '{compression_method}' is not " + f"available on this Python runtime", + ) + level_msg = level_error(compression_method, zip_compression_level) + if level_msg is not None: + raise CommandError(level_msg) + sink: ExportSink if self.zip_export: sink = ZipExportSink( self.target, options["zip_name"], delete=self.delete, + compression=COMPRESSION_METHODS[compression_method], + compresslevel=zip_compression_level, ) else: sink = DirectoryExportSink( diff --git a/src/documents/management/commands/document_importer.py b/src/documents/management/commands/document_importer.py index 0e54b0ced..7dcc03e5f 100644 --- a/src/documents/management/commands/document_importer.py +++ b/src/documents/management/commands/document_importer.py @@ -32,6 +32,8 @@ from django.db.models.signals import post_save from filelock import FileLock from guardian.shortcuts import clear_ct_cache +from documents.export.compression import compress_type_readable +from documents.export.compression import unreadable_method_names from documents.file_handling import create_source_path_directory from documents.management.commands.base import PaperlessCommand from documents.management.commands.mixins import CryptMixin @@ -460,6 +462,20 @@ class Command(CryptMixin, PaperlessCommand): with tempfile.TemporaryDirectory() as tmp_dir: if is_zipfile(self.source): with ZipFile(self.source) as zf: + unsupported = { + info.compress_type + for info in zf.infolist() + if not compress_type_readable(info.compress_type) + } + if unsupported: + names = sorted(unreadable_method_names(unsupported)) + message = ( + f"This archive uses compression this Python cannot " + f"read ({', '.join(names)})." + ) + if "zstd" in names: + message += " zstd archives require Python 3.14+." + raise CommandError(message) zf.extractall(tmp_dir) self.source = Path(tmp_dir) self._run_import() diff --git a/src/documents/tests/export/test_compression.py b/src/documents/tests/export/test_compression.py new file mode 100644 index 000000000..83dd4020a --- /dev/null +++ b/src/documents/tests/export/test_compression.py @@ -0,0 +1,190 @@ +import sys +import zipfile + +import pytest + +from documents.export import compression + + +class TestCompressionMethods: + def test_choices_always_include_zstd(self) -> None: + """ + GIVEN: + - The compression policy module's CLI choices list + WHEN: + - Read on any runtime + THEN: + - zstd is always present; availability is checked separately so + argparse never hides it based on the current Python version + """ + assert compression.COMPRESSION_CHOICES == ( + "stored", + "deflated", + "bzip2", + "lzma", + "zstd", + ) + + @pytest.mark.parametrize( + ("name", "constant"), + [ + ("stored", zipfile.ZIP_STORED), + ("deflated", zipfile.ZIP_DEFLATED), + ("bzip2", zipfile.ZIP_BZIP2), + ("lzma", zipfile.ZIP_LZMA), + ], + ) + def test_method_maps_to_zipfile_constant(self, name: str, constant: int) -> None: + """ + GIVEN: + - A compression method name + WHEN: + - Looked up in COMPRESSION_METHODS + THEN: + - It maps to the matching zipfile compression constant + """ + assert compression.COMPRESSION_METHODS[name] == constant + + def test_stored_and_deflated_always_available(self) -> None: + """ + GIVEN: + - The stored and deflated compression methods + WHEN: + - Checked with compression_available() + THEN: + - Both are always available (zlib is a hard CPython dependency) + """ + assert compression.compression_available("stored") + assert compression.compression_available("deflated") + + def test_zstd_availability_tracks_runtime(self) -> None: + """ + GIVEN: + - The zstd compression method + WHEN: + - Checked with compression_available() on this runtime + THEN: + - Availability matches whether Python is 3.14+ + """ + expected: bool = sys.version_info >= (3, 14) + assert compression.compression_available("zstd") == expected + + +class TestLevelError: + @pytest.mark.parametrize( + ("method", "level"), + [ + ("deflated", 0), + ("deflated", 9), + ("bzip2", 1), + ("bzip2", 9), + ("zstd", -22), + ("zstd", 22), + ("deflated", None), + ("stored", None), + ], + ) + def test_valid_levels_return_none(self, method: str, level: int | None) -> None: + """ + GIVEN: + - A method and a level within its valid bounds (or no level) + WHEN: + - Checked with level_error() + THEN: + - No error message is returned + """ + assert compression.level_error(method, level) is None + + @pytest.mark.parametrize( + ("method", "level"), + [ + ("deflated", 10), + ("deflated", -1), + ("bzip2", 0), + ("bzip2", 10), + ("zstd", -23), + ("zstd", 23), + ], + ) + def test_out_of_range_levels_return_message( + self, + method: str, + level: int, + ) -> None: + """ + GIVEN: + - A method and a level outside its valid bounds + WHEN: + - Checked with level_error() + THEN: + - An error message naming the valid range is returned + """ + msg: str | None = compression.level_error(method, level) + assert msg is not None + assert "between" in msg + + @pytest.mark.parametrize("method", ["stored", "lzma"]) + def test_level_on_levelless_method_is_rejected(self, method: str) -> None: + """ + GIVEN: + - A method that ignores compression level (stored, lzma) + WHEN: + - A level is passed to level_error() anyway + THEN: + - An error message noting the level has no effect is returned + """ + msg: str | None = compression.level_error(method, 5) + assert msg is not None + assert "no effect" in msg + + +class TestCompressTypeReadable: + @pytest.mark.parametrize("ct", [zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED]) + def test_stored_and_deflated_always_readable(self, ct: int) -> None: + """ + GIVEN: + - A stored or deflated compress_type id + WHEN: + - Checked with compress_type_readable() + THEN: + - It is always readable + """ + assert compression.compress_type_readable(ct) + + def test_zstd_compress_type_readability_tracks_runtime(self) -> None: + """ + GIVEN: + - The current (93) and legacy (20) zstd compress_type ids + WHEN: + - Checked with compress_type_readable() on this runtime + THEN: + - Readability matches whether Python is 3.14+ + """ + # 93 = ZIP_ZSTANDARD; 20 = legacy zstd method id (read-only) + expected: bool = sys.version_info >= (3, 14) + assert compression.compress_type_readable(93) == expected + assert compression.compress_type_readable(20) == expected + + def test_unknown_compress_type_is_unreadable(self) -> None: + """ + GIVEN: + - An unrecognized compress_type id + WHEN: + - Checked with compress_type_readable() + THEN: + - It is reported as unreadable + """ + assert not compression.compress_type_readable(9999) + + def test_unreadable_method_names_lists_methods(self) -> None: + """ + GIVEN: + - A set containing an unknown compress_type id + WHEN: + - Passed to unreadable_method_names() + THEN: + - It is reported generically as "method " + """ + # An unknown method id maps to no name and is reported generically. + names: set[str] = compression.unreadable_method_names({9999}) + assert names == {"method 9999"} diff --git a/src/documents/tests/export/test_sinks.py b/src/documents/tests/export/test_sinks.py index fe421b234..78c2936d0 100644 --- a/src/documents/tests/export/test_sinks.py +++ b/src/documents/tests/export/test_sinks.py @@ -5,6 +5,7 @@ import zipfile from pathlib import Path import pytest +import pytest_mock from pytest_django.fixtures import SettingsWrapper from documents.export.sinks import DirectoryExportSink @@ -305,6 +306,48 @@ class TestZipExportSink: assert not (target / "export.zip").exists() +class TestZipExportSinkCompression: + @pytest.mark.parametrize( + ("method", "constant"), + [ + ("stored", zipfile.ZIP_STORED), + ("deflated", zipfile.ZIP_DEFLATED), + ("bzip2", zipfile.ZIP_BZIP2), + ("lzma", zipfile.ZIP_LZMA), + ], + ) + def test_compression_and_level_forwarded_to_zipfile( + self, + mocker: pytest_mock.MockerFixture, + tmp_path: Path, + method: str, + constant: int, + ) -> None: + """ + GIVEN: + - A ZipExportSink constructed with a compression method and level + WHEN: + - The sink is opened + THEN: + - zipfile.ZipFile is constructed with those values forwarded + unchanged (whether ZipFile actually compresses is Python's own + contract, not ours, so this checks the call args, not a real + archive) + """ + target: Path = tmp_path / "out" + target.mkdir() + zip_cls = mocker.patch("documents.export.sinks.zipfile.ZipFile") + sink = ZipExportSink(target, "export", compression=constant, compresslevel=5) + sink._open() + zip_cls.assert_called_once_with( + mocker.ANY, + "w", + compression=constant, + compresslevel=5, + allowZip64=True, + ) + + class TestStreamContract: @pytest.fixture(params=["dir", "zip"]) def sink(self, request: pytest.FixtureRequest, tmp_path: Path) -> ExportSink: diff --git a/src/documents/tests/test_management_exporter.py b/src/documents/tests/test_management_exporter.py index 44456535a..9666da5d6 100644 --- a/src/documents/tests/test_management_exporter.py +++ b/src/documents/tests/test_management_exporter.py @@ -6,6 +6,8 @@ from datetime import timedelta from io import StringIO from pathlib import Path from unittest import mock +from zipfile import ZIP_DEFLATED +from zipfile import ZIP_LZMA from zipfile import ZipFile import pytest @@ -1078,6 +1080,186 @@ class TestExportImport( skip_checks=True, ) + def test_compression_flags_require_zip(self) -> None: + """ + GIVEN: + - A request to export without --zip + WHEN: + - --zip-compression or --zip-compression-level is passed anyway + THEN: + - A CommandError is raised (the flags are meaningless without --zip) + """ + for args in ( + ["--zip-compression", "lzma"], + ["--zip-compression-level", "5"], + ): + with self.assertRaises(CommandError): + call_command( + "document_exporter", + self.target, + *args, + skip_checks=True, + ) + + def test_zip_compression_level_out_of_range_raises(self) -> None: + """ + GIVEN: + - A request to export to a zip file + WHEN: + - --zip-compression-level is outside the chosen method's valid range + THEN: + - A CommandError is raised + """ + with self.assertRaises(CommandError): + call_command( + "document_exporter", + self.target, + "--zip", + "--zip-compression", + "deflated", + "--zip-compression-level", + "99", + skip_checks=True, + ) + + def test_zip_compression_level_rejected_for_stored(self) -> None: + """ + GIVEN: + - A request to export to a zip file with --zip-compression stored + WHEN: + - --zip-compression-level is also passed + THEN: + - A CommandError is raised (stored ignores level entirely) + """ + with self.assertRaises(CommandError): + call_command( + "document_exporter", + self.target, + "--zip", + "--zip-compression", + "stored", + "--zip-compression-level", + "5", + skip_checks=True, + ) + + def test_zip_compression_level_rejected_for_lzma(self) -> None: + """ + GIVEN: + - A request to export to a zip file with --zip-compression lzma + WHEN: + - --zip-compression-level is also passed + THEN: + - A CommandError is raised (lzma ignores level entirely) + """ + with self.assertRaises(CommandError): + call_command( + "document_exporter", + self.target, + "--zip", + "--zip-compression", + "lzma", + "--zip-compression-level", + "5", + skip_checks=True, + ) + + def test_zstd_unavailable_raises_friendly_error(self) -> None: + """ + GIVEN: + - A Python runtime without zstd support (< 3.14) + WHEN: + - --zip-compression zstd is requested + THEN: + - A CommandError naming the Python version requirement is raised + + zstd availability is mocked rather than relying on the actual + runtime: on a Python 3.14+ CI leg, ZSTD is not None, so without the + mock this check is skipped and the command falls through into the + real export, which fails on missing document files instead of + raising the expected CommandError. + """ + with ( + mock.patch( + "documents.management.commands.document_exporter.ZSTD", + None, + ), + mock.patch( + "documents.management.commands.document_exporter.compression_available", + return_value=False, + ), + self.assertRaises(CommandError) as e, + ): + call_command( + "document_exporter", + self.target, + "--zip", + "--zip-compression", + "zstd", + skip_checks=True, + ) + self.assertIn("3.14", str(e.exception)) + + def test_zip_compression_flag_resolves_to_sink_constant(self) -> None: + """ + GIVEN: + - A request to export to a zip file with --zip-compression lzma + WHEN: + - The export runs + THEN: + - ZipExportSink is constructed with the resolved ZIP_LZMA constant + (whether zipfile actually compresses with the chosen method is + Python's own contract, and ZipExportSink's own tests already + cover the forwarding; what this command owns is resolving the + CLI string to the right constant, so assert that resolution + directly) + """ + with mock.patch( + "documents.management.commands.document_exporter.ZipExportSink", + ) as sink_cls: + call_command( + "document_exporter", + self.target, + "--zip", + "--zip-compression", + "lzma", + skip_checks=True, + ) + sink_cls.assert_called_once_with( + mock.ANY, + mock.ANY, + delete=False, + compression=ZIP_LZMA, + compresslevel=None, + ) + + def test_default_zip_compression_resolves_to_deflate(self) -> None: + """ + GIVEN: + - A request to export to a zip file with no --zip-compression flag + WHEN: + - The export runs + THEN: + - ZipExportSink is constructed with the default ZIP_DEFLATED + constant and compresslevel=None, matching pre-existing behavior + """ + with mock.patch( + "documents.management.commands.document_exporter.ZipExportSink", + ) as sink_cls: + call_command( + "document_exporter", + self.target, + "--zip", + skip_checks=True, + ) + sink_cls.assert_called_once_with( + mock.ANY, + mock.ANY, + delete=False, + compression=ZIP_DEFLATED, + compresslevel=None, + ) + @pytest.mark.management class TestCryptExportImport( diff --git a/src/documents/tests/test_management_importer.py b/src/documents/tests/test_management_importer.py index 33eec7267..243de9d64 100644 --- a/src/documents/tests/test_management_importer.py +++ b/src/documents/tests/test_management_importer.py @@ -525,6 +525,35 @@ class TestCommandImport( self.assertEqual(doc.tags.count(), 1) self.assertEqual(doc.tags.first().name, "batch-flush-tag") + def test_import_rejects_unreadable_compression(self) -> None: + """ + GIVEN: + - A zip archive with an entry whose compression this Python can't read + WHEN: + - Import is attempted + THEN: + - A CommandError naming the issue is raised, before extraction + """ + import zipfile + from unittest import mock + + archive = Path(self.dirs.scratch_dir) / "export.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("manifest.json", "[]") + + with mock.patch( + "documents.management.commands.document_importer.compress_type_readable", + return_value=False, + ): + with self.assertRaises(CommandError) as e: + call_command( + "document_importer", + str(archive), + "--no-progress-bar", + skip_checks=True, + ) + self.assertIn("compression", str(e.exception)) + @pytest.mark.management @pytest.mark.django_db