implement rollbacks on the data directory

This helps if the upload/editing process fails, so that we don't end up
with orphan data directories (probably the most important change!), or
superfluous image files.

4 files changed, 141 insertions, 28 deletions

diff --git a/fietsboek/actions.py b/fietsboek/actions.py
index f1a32fc..bba4185 100644
--- a/fietsboek/actions.py
+++ b/fietsboek/actions.py
@@ -8,14 +8,14 @@ the test functions.
 import datetime
 import logging
 import re
-from typing import List
+from typing import List, Optional
 
 from pyramid.request import Request
 from sqlalchemy import select
 from sqlalchemy.orm.session import Session
 
 from fietsboek import models, util
-from fietsboek.data import DataManager
+from fietsboek.data import DataManager, TrackDataDir
 from fietsboek.models.track import TrackType, Visibility
 
 LOGGER = logging.getLogger(__name__)
@@ -81,22 +81,22 @@ def add_track(
 
     # Save the GPX data
     LOGGER.debug("Creating a new data folder for %d", track.id)
-    manager = data_manager.initialize(track.id)
-    LOGGER.debug("Saving GPX to %s", manager.gpx_path())
-    manager.compress_gpx(gpx_data)
-    manager.backup()
-
-    manager.engrave_metadata(
-        title=track.title,
-        description=track.description,
-        author_name=track.owner.name,
-        time=track.date,
-    )
+    with data_manager.initialize(track.id) as manager:
+        LOGGER.debug("Saving GPX to %s", manager.gpx_path())
+        manager.compress_gpx(gpx_data)
+        manager.backup()
+
+        manager.engrave_metadata(
+            title=track.title,
+            description=track.description,
+            author_name=track.owner.name,
+            time=track.date,
+        )
 
     return track
 
 
-def edit_images(request: Request, track: models.Track):
+def edit_images(request: Request, track: models.Track, *, manager: Optional[TrackDataDir] = None):
     """Edit the images according to the given request.
 
     This deletes and adds images and image descriptions as needed, based on the
@@ -104,9 +104,11 @@ def edit_images(request: Request, track: models.Track):
 
     :param request: The request.
     :param track: The track to edit.
+    :param manager: The track's data manager. If not given, it will
+        automatically be opened from the request.
     """
     LOGGER.debug("Editing images for %d", track.id)
-    manager = request.data_manager.open(track.id)
+    manager = manager if manager else request.data_manager.open(track.id)
 
     # Delete requested images
     for image in request.params.getall("delete-image[]"):
diff --git a/fietsboek/data.py b/fietsboek/data.py
index 1a1b66b..cbe20e9 100644
--- a/fietsboek/data.py
+++ b/fietsboek/data.py
@@ -10,7 +10,7 @@ import shutil
 import string
 import uuid
 from pathlib import Path
-from typing import BinaryIO, List, Optional
+from typing import BinaryIO, List, Literal, Optional
 
 import brotli
 import gpxpy
@@ -77,7 +77,7 @@ class DataManager:
         """
         path = self._track_data_dir(track_id)
         path.mkdir(parents=True)
-        return TrackDataDir(track_id, path)
+        return TrackDataDir(track_id, path, journal=True, is_fresh=True)
 
     def purge(self, track_id: int):
         """Forcefully purges all data from the given track.
@@ -101,11 +101,94 @@ class DataManager:
 
 
 class TrackDataDir:
-    """Manager for a single track's data."""
+    """Manager for a single track's data.
+
+    If initialized with ``journal = True``, then you can use :meth:`rollback`
+    to roll back the changes in case of an error. In case of no error, use
+    :meth:`commit` to commit the changes. If you don't want the "journalling"
+    semantics, use ``journal = False``.
+    """
 
-    def __init__(self, track_id: int, path: Path):
+    def __init__(self, track_id: int, path: Path, *, journal: bool = False, is_fresh: bool = False):
         self.track_id: int = track_id
         self.path: Path = path
+        self.journal: Optional[list] = [] if journal else None
+        self.is_fresh = is_fresh
+
+    def __enter__(self) -> "TrackDataDir":
+        if self.journal is None:
+            self.journal = []
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> Literal[False]:
+        if exc_type is None and exc_val is None and exc_tb is None:
+            self.commit()
+        else:
+            self.rollback()
+        return False
+
+    def rollback(self):
+        """Rolls back the journal, e.g. in case of error.
+
+        :raises ValueError: If the data directory was opened without the
+            journal, this raises :exc:`ValueError`.
+        """
+        LOGGER.debug("Rolling back state of %s", self.path)
+
+        if self.journal is None:
+            raise ValueError("Rollback on a non-journalling data directory")
+
+        if self.is_fresh:
+            # Shortcut if the directory is fresh, simply remove everything
+            self.journal = None
+            self.purge()
+            return
+
+        for action, *rest in reversed(self.journal):
+            if action == "purge":
+                (new_name,) = rest
+                shutil.move(new_name, self.path)
+            elif action == "compress_gpx":
+                (old_data,) = rest
+                if old_data is None:
+                    self.gpx_path().unlink()
+                else:
+                    self.gpx_path().write_bytes(old_data)
+            elif action == "add_image":
+                (image_path,) = rest
+                image_path.unlink()
+            elif action == "delete_image":
+                path, data = rest
+                path.write_bytes(data)
+
+        self.journal = None
+
+    def commit(self):
+        """Commits all changed and deletes the journal.
+
+        Note that this function will do nothing if the journal is disabled,
+        meaning it can always be called.
+        """
+        LOGGER.debug("Committing journal for %s", self.path)
+
+        if self.journal is None:
+            return
+
+        for action, *rest in reversed(self.journal):
+            if action == "purge":
+                (new_name,) = rest
+                shutil.rmtree(new_name, ignore_errors=False, onerror=self._log_deletion_error)
+            elif action == "compress_gpx":
+                # Nothing to do here, the new data is already on the disk
+                pass
+            elif action == "add_image":
+                # Nothing to do here, the image is already saved
+                pass
+            elif action == "delete_image":
+                # Again, nothing to do here, we simply discard the in-memory image data
+                pass
+
+        self.journal = None
 
     def lock(self) -> FileLock:
         """Returns a FileLock that can be used to lock access to the track's
@@ -115,18 +198,23 @@ class TrackDataDir:
         """
         return FileLock(self.path / "lock")
 
+    @staticmethod
+    def _log_deletion_error(_, path, exc_info):
+        LOGGER.warning("Failed to remove %s", path, exc_info=exc_info)
+
     def purge(self):
         """Purge all data pertaining to the track.
 
         This function logs errors but raises no exception, as such it can
         always be used to clean up after a track.
         """
-
-        def log_error(_, path, exc_info):
-            LOGGER.warning("Failed to remove %s", path, exc_info=exc_info)
-
-        if self.path.is_dir():
-            shutil.rmtree(self.path, ignore_errors=False, onerror=log_error)
+        if self.journal is None:
+            if self.path.is_dir():
+                shutil.rmtree(self.path, ignore_errors=False, onerror=self._log_deletion_error)
+        else:
+            new_name = self.path.with_name("trash-" + self.path.name)
+            shutil.move(self.path, new_name)
+            self.journal.append(("purge", new_name))
 
     def gpx_path(self) -> Path:
         """Returns the path of the GPX file.
@@ -147,6 +235,16 @@ class TrackDataDir:
         :param quality: Compression quality, from 0 to 11 - 11 is highest
             quality but slowest compression speed.
         """
+        if self.journal is not None:
+            # First, we check if we already saved an old state of the GPX data
+            for action, *_ in self.journal:
+                if action == "compress_gpx":
+                    break
+            else:
+                # We did not save a state yet
+                old_data = None if not self.gpx_path().is_file() else self.gpx_path().read_bytes()
+                self.journal.append(("compress_gpx", old_data))
+
         compressed = brotli.compress(data, quality=quality)
         self.gpx_path().write_bytes(compressed)
 
@@ -233,6 +331,9 @@ class TrackDataDir:
         with open(path, "wb") as fobj:
             shutil.copyfileobj(image, fobj)
 
+        if self.journal is not None:
+            self.journal.append(("add_image", path))
+
         return filename
 
     def delete_image(self, image_id: str):
@@ -246,4 +347,8 @@ class TrackDataDir:
         if "/" in image_id or "\\" in image_id:
             return
         path = self.image_path(image_id)
+
+        if self.journal is not None:
+            self.journal.append(("delete_image", path, path.read_bytes()))
+
         path.unlink()
diff --git a/fietsboek/views/edit.py b/fietsboek/views/edit.py
index e3fe99d..61741e5 100644
--- a/fietsboek/views/edit.py
+++ b/fietsboek/views/edit.py
@@ -82,7 +82,7 @@ def do_edit(request):
     data: TrackDataDir = request.data_manager.open(track.id)
     tz_offset = datetime.timedelta(minutes=int(request.params["date-tz"]))
     date = datetime.datetime.fromisoformat(request.params["date"])
-    with data.lock():
+    with data, data.lock():
         track.date = date.replace(tzinfo=datetime.timezone(tz_offset))
 
         track.tagged_people = tagged_people
@@ -94,7 +94,7 @@ def do_edit(request):
         tags = request.params.getall("tag[]")
         track.sync_tags(tags)
 
-        actions.edit_images(request, request.context)
+        actions.edit_images(request, request.context, manager=data)
         data.engrave_metadata(
             title=track.title,
             description=track.description,
diff --git a/fietsboek/views/upload.py b/fietsboek/views/upload.py
index e172339..ed8f285 100644
--- a/fietsboek/views/upload.py
+++ b/fietsboek/views/upload.py
@@ -176,7 +176,13 @@ def do_finish_upload(request):
 
     # Don't forget to add the images
     LOGGER.debug("Starting to edit images for the upload")
-    actions.edit_images(request, track)
+    try:
+        actions.edit_images(request, track)
+    except Exception:
+        # We just created the folder, so we'll be fine deleting it
+        LOGGER.info("Deleting partially created folder for track %d", track.id)
+        request.data_manager.open(track.id).purge()
+        raise
 
     request.session.flash(request.localizer.translate(_("flash.upload_success")))