Simple Python Transactional Object

main.py

def delete_instance(instance, args):
    """Given an instance, deletes it!"""
    if instance:
        print "Rolling back %s" % instance.id
        instance.terminate()


@Rollback(delete_instance)
def create_instance(conn, kwargs):
    """Performs creation of instance by wrapping over the AWS API and injecting
    connection & arguments. We should use this instead of the RAW API because
    we can Rollback operations that fail.
    """
    reserv = conn.run_instances(**kwargs)
    assert len(reserv.instances) == 1
    instance, = reserv.instances
    return instance
    
if __name__ == "__main__":
    import boto.ec2
    conn = boto.ec2.connect_to_region('us-east-1')
    
    with RollbackManager():
        create_instance(conn, {'image_id': 'ami-fe147796'})
        create_instance(conn, {'image_id': 'ami-fe147796'})
        create_instance(conn, {'image_id': 'ami-fe147796'})
        create_instance(conn, {'image_id': 'ami-fe147796'})

        # This exception will cause all 4 operations to rollback in reverse order
        throw Exception("testing rollback")
        
        

state.py

import sys
import traceback

# Global (yes globals) for managing the current transactions. The execution of
# transactions are potentially nested, so this keeps a list of transactions.
transactions = []


class Rollback(object):
    """A stateful operation is a way of wrapping functions that are currently
    executed with a dual rollback operation to undo the changes if necessary.

    The rollback function takes two arguments: the first is the original
    arguments to invoke the function, the second is the return value of the
    function. This only works in a context of a Transaction.

    For an undo operation to work, we should intercept calls to these methods,
    which receive the same args used to construct this object.
    """
    def __init__(self, undo_fn):
        self.undo_fn = undo_fn

    def __call__(self, f):
        def wrapped_f(*args):
            if len(transactions) == 0:
                raise Exception("Cannot perform @Rollbackable function without"
                                "a RollbackManager() context")

            current_transaction = transactions[-1]
            result = f(*args)
            current_transaction.append([self.undo_fn, result, args])
            return result

        return wrapped_f


class RollbackManager():
    """Context Manager or Decorator used to define the transaction
    boundary. `@Rollback` operations are applied in reverse if an exception
    bubbles through ``RollbackManager``

    Usage as a context manager:
        conn = boto.ec2.connect_to_region('us-west-1')

        with RollbackManager():
            create_instance(conn, {"foo": "bar"})

    Usage as a decorator:
    @RollbackManager
    def stateful(self, ..):
      # ...
    """
    def __call__(self, f):

        def wrapped_f(*args):
            with RollbackManager():
                f(*args)

        return wrapped_f

    def __enter__(self):
        # Create a new empty transaction object
        transactions.append([])

    def __exit__(self, *exc_info):
        undo_operations = transactions.pop()
        if exc_info != (None, None, None):
            exc_type, exc_value, tb = exc_info
            stacktrace = traceback.format_exc(tb)
            # Handle all operations that were marked stateful in reverse order
            # If they fail -- then it gets messy
            undo_operations.reverse()
            rollback_exceptions = []
            sys.stderr.write("%s\nRolling back with RollbackManager\n")
            for undo_fn, retval, args in undo_operations:
                sys.stderr.write("Rollback: %s %s\n" % (undo_fn.__name__, args))
                try:
                    undo_fn(retval, args)
                except Exception as e:
                    rollback_exceptions.add(e)

            raise RollbackException(exc_info, stacktrace, rollback_exceptions)

class RollbackException(Exception):
    def __init__(self, cause, stacktrace, rollback_exceptions):
        super(RollbackException, self).__init__("Exception occured during rollback.")
        self.cause = cause
        self.stacktrace = stacktrace
        self.rollback_exceptions = rollback_exceptions

    def __str__(self):
        return super(RollbackException, self).__str__() + "\n" + self.stacktrace