Commit 5640d5f4 authored by Mickaël Desfrênes's avatar Mickaël Desfrênes
Browse files

when making a snapshot with files, pre-check available disk space. Also handle OSError.

parent 215145a0

src/resources/snapshots.py

+31 −1
Changes for src/resources/snapshots.py: 31 added lines, 1 removed line.
Original line number Diff line number Diff line
@@ -147,7 +147,7 @@ def _project_collections_ordered(project: Project) -> list[Collection]: 
    return ordered





def export_project_snapshot(

def _export_project_snapshot(

    project_id: int,

    output_dir: str | Path,

    *,
@@ -163,6 +163,14 @@ def export_project_snapshot( 
    if output_dir.exists() and any(output_dir.iterdir()) and not overwrite:

        raise SnapshotError(f"Snapshot directory {output_dir} is not empty")

    output_dir.mkdir(parents=True, exist_ok=True)

    if include_files:

        required_space = project.files_size()

        available_space = shutil.disk_usage(output_dir).free

        if required_space > available_space:

            raise SnapshotError(

                f"Not enough free space in {output_dir}: "

                f"{required_space} bytes required, {available_space} bytes available"

            )



    file_ids = set(File.objects.filter(project=project).values_list("pk", flat=True))

    counts = {}
@@ -439,6 +447,28 @@ def export_project_snapshot( 
    return manifest





def export_project_snapshot(

    project_id: int,

    output_dir: str | Path,

    *,

    include_files: bool = False,

    overwrite: bool = False,

    chunk_size: int = 2000,

) -> dict:

    try:

        return _export_project_snapshot(

            project_id,

            output_dir,

            include_files=include_files,

            overwrite=overwrite,

            chunk_size=chunk_size,

        )

    except OSError as error:

        raise SnapshotError(

            f"Could not write snapshot to {output_dir}: {error}"

        ) from error





@contextmanager

def _restore_without_model_side_effects():

    """Disable save signals that would enqueue ARK work mid-restore."""

src/resources/tests.py

+30 −0
Changes for src/resources/tests.py: 30 added lines, 0 removed lines.
Original line number Diff line number Diff line
@@ -1190,6 +1190,36 @@ class ProjectSnapshotTestCase(TestCase): 
                        )

                self.assertEqual(source_path.read_bytes(), content)



    def test_snapshot_with_files_checks_available_disk_space_before_writing(self):

        with TemporaryDirectory() as temp_dir:

            snapshot_dir = Path(temp_dir) / "snapshot"

            with patch.object(

                snapshots.shutil,

                "disk_usage",

                return_value=SimpleNamespace(total=10, used=7, free=3),

            ):

                with self.assertRaisesRegex(

                    snapshots.SnapshotError,

                    "4 bytes required, 3 bytes available",

                ):

                    snapshots.export_project_snapshot(

                        self.project.pk, snapshot_dir, include_files=True

                    )



            self.assertFalse(snapshot_dir.joinpath("project.ndjson").exists())



    def test_snapshot_converts_os_errors_to_snapshot_errors(self):

        with TemporaryDirectory() as snapshot_dir:

            error = OSError(errno.ENOSPC, "No space left on device")

            with patch.object(snapshots, "_write_jsonl", side_effect=error):

                with self.assertRaisesRegex(

                    snapshots.SnapshotError,

                    "Could not write snapshot.*No space left on device",

                ) as context:

                    snapshots.export_project_snapshot(self.project.pk, snapshot_dir)



        self.assertIs(context.exception.__cause__, error)



    def test_replace_in_place_preserves_project_id_and_writes_safety_snapshot(self):

        with TemporaryDirectory() as temp_dir:

            temp_dir = Path(temp_dir)