diff --git a/src/documents/management/commands/document_exporter.py b/src/documents/management/commands/document_exporter.py index cc0aa4d1b..0d0f3f681 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 @@ -191,6 +196,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: -131072..22; ignored for " + "stored/lzma." + ), + ) + parser.add_argument( "--data-only", default=False, @@ -246,12 +273,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/tests/test_management_exporter.py b/src/documents/tests/test_management_exporter.py index 44456535a..7734de7f8 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,87 @@ class TestExportImport( skip_checks=True, ) + def test_compression_flags_require_zip(self) -> None: + 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: + 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: + with self.assertRaises(CommandError): + call_command( + "document_exporter", + self.target, + "--zip", + "--zip-compression", + "stored", + "--zip-compression-level", + "5", + skip_checks=True, + ) + + def test_zip_lzma_compression_round_trips(self) -> None: + shutil.rmtree(Path(self.dirs.media_dir) / "documents") + shutil.copytree( + Path(__file__).parent / "samples" / "documents", + Path(self.dirs.media_dir) / "documents", + ) + call_command( + "document_exporter", + self.target, + "--zip", + "--zip-compression", + "lzma", + skip_checks=True, + ) + expected = str( + self.target / f"export-{timezone.localdate().isoformat()}.zip", + ) + self.assertIsFile(expected) + with ZipFile(expected) as zip_file: + info = zip_file.getinfo("manifest.json") + # manifest.json carries the chosen method; deflated is the default + self.assertEqual(info.compress_type, ZIP_LZMA) + + def test_default_zip_uses_deflate(self) -> None: + shutil.rmtree(Path(self.dirs.media_dir) / "documents") + shutil.copytree( + Path(__file__).parent / "samples" / "documents", + Path(self.dirs.media_dir) / "documents", + ) + call_command( + "document_exporter", + self.target, + "--zip", + skip_checks=True, + ) + expected = str( + self.target / f"export-{timezone.localdate().isoformat()}.zip", + ) + with ZipFile(expected) as zip_file: + info = zip_file.getinfo("manifest.json") + self.assertEqual(info.compress_type, ZIP_DEFLATED) + @pytest.mark.management class TestCryptExportImport(