diff --git a/cabinet/management/commands/__init__.py b/cabinet/management/commands/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cabinet/management/commands/archive_cabinet_folder.py b/cabinet/management/commands/archive_cabinet_folder.py new file mode 100644 index 0000000..40f4e35 --- /dev/null +++ b/cabinet/management/commands/archive_cabinet_folder.py @@ -0,0 +1,48 @@ +import random +import string +from pathlib import Path +from zipfile import ZIP_DEFLATED, ZipFile + +from cabinet.models import Folder +from django.core.management import BaseCommand + + +def _get_random_suffix(): + return random.choices(string.ascii_lowercase, k=4) + + +class Command(BaseCommand): + help = "Create archive with contents of a cabinet folder, using the data structure of the db instead of the disk." + + def add_arguments(self, parser): + parser.add_argument("--folder-id", type=int, required=True) + parser.add_argument("--output", type=Path, required=True) + + def handle(self, **options): + folder = Folder.objects.get(id=options["folder_id"]) + output = options["output"] + + arc_paths = set() + with ZipFile(output, "w", ZIP_DEFLATED) as zip_file: + for file, path in self._walk(folder, path=()): + arc_path = Path(*path) / file.file_name + + if arc_path in arc_paths: + filename = Path(file.file_name) + arc_path = Path(*path) / "".join([ + filename.stem, + "_", + *_get_random_suffix(), + *filename.suffixes, + ]) + + zip_file.write(file.file.path, arc_path) + arc_paths.add(arc_path) + + def _walk(self, folder, *, path): + path = (*tuple(path), folder.name) + for file_ in folder.files.all(): + yield (file_, path) + + for child in folder.children.all(): + yield from self._walk(child, path=path) diff --git a/tests/testapp/test_cabinet.py b/tests/testapp/test_cabinet.py index 63d622b..97b6ec7 100644 --- a/tests/testapp/test_cabinet.py +++ b/tests/testapp/test_cabinet.py @@ -1,3 +1,4 @@ +import tempfile import io import itertools import json @@ -9,13 +10,17 @@ from django.contrib.auth.models import User from django.core.exceptions import ImproperlyConfigured, ValidationError from django.core.files.base import ContentFile +from django.core.management import call_command from django.test import Client, TestCase from django.test.utils import override_settings from django.urls import reverse from cabinet.base import AbstractFile, DownloadMixin, determine_accept_file_functions from cabinet.models import File, Folder, get_file_model +from pathlib import Path from testapp.models import Stuff +from unittest.mock import patch +from zipfile import ZipFile class CabinetTestCase(TestCase): @@ -547,3 +552,28 @@ class CustomFile(AbstractFile, NonModelMixin, DownloadMixin): } ], ) + + @patch("cabinet.management.commands.archive_cabinet_folder._get_random_suffix", return_value="asdf") + def test_archive_management_command(self, patched__get_random_suffix): + output = "archive.zip" + folder = Folder.objects.create(name="Top") + subfolder = Folder.objects.create(parent=folder, name="Sub") + + for _ in range(2): + file = File(folder=subfolder) + file.download_file.save("hello.txt", ContentFile("Hello")) + + # enforce duplicate names + File.objects.all().update(file_name="hello.txt") + + with tempfile.TemporaryDirectory() as tmp_dir: + output = Path(tmp_dir) / 'output.zip' + call_command('archive_cabinet_folder', folder_id=folder.id, output=output) + with ZipFile(output, "r") as zip_file: + self.assertEqual( + zip_file.namelist(), + [ + "Top/Sub/hello.txt", + "Top/Sub/hello_asdf.txt", + ] + )