Writing transactions for non-DB code in Python

Question

A piece of python's code is writing/creating multiple files. I would like to put it in a transaction-ish way, so everything is rolled-back if one of the file's creation fails. This would help to save hard-drive space and avoid unuseful files from failed transactions.

Concretely, by sending a picture to my python server:

1. It first saves the original file, full res
2. Creates a smaller version of the pic
3. Saves this smaller verion
4. Transaction succeded

However step 2 or 3 could fail for some reason. How to rollback step 1 ?

The solution should be generic as there might be more steps.

Is there a common transaction trick in Python ?

Answer 1

My solution

It creates a temporary directory in the system's default temporary directory (/tmp on linux) and returns a pathlib.Path object where you can write files to, if all is successful it will eventually move the contents to the destination, if not nothing will happen.

Either way it cleans up after it and removes the data it has written to the temporary directory

import tempfile
import shutil
from typing import ContextManager, Union
from pathlib import Path
from contextlib import contextmanager

@contextmanager
def transactional_mkdir(destination: Union[str, Path], overwrite: bool = False) -> ContextManager[Path]:
    """
    Create a directory and get reference to it, while inside the context if any unhandled failure is met all the writing
    operations that were made to the directory will be undone

    :param destination: Destination directory to write create
    :param overwrite: If True, will overwrite destination if it exists 
    :return: Path object to destination directory
    """

    # Validate input
    destination = Path(destination)
    if not overwrite and destination.exists():
        raise FileExistsError(f'destination already exists. To overwrite use "overwrite=True"')

    # Create temporary directory
    temp_dir = Path(tempfile.mkdtemp())

    try:
        yield temp_dir  # Give control

        # Remove existing
        if destination.exists():
            if destination.is_dir():
                shutil.rmtree(destination)  # Remove directory
            else:
                destination.unlink()  # Remove file

        # Save to destination
        shutil.move(temp_dir, destination)

    finally:
        shutil.rmtree(temp_dir)  # Remove directory