Transaction.atomic with continuations I have a view that looks like this:

@transaction.commit_on_success
def view(request):
    do_stuff(...)  
    if condition(...):
        result = do_other_stuff(...)
        transaction.commit()
        send_email(result)
    return HttpResponse(...)


Assume everything with (...) can touch the database.

There are currently two main cases:


condition is true, do_stuff() and do_other_stuff() are done in transaction 1, the transaction is saved, then the e-mail is sent and HttpResponse is done in transaction 2   
condition is false, everything happens in transaction 2


I want to rewrite this with transaction.atomic. The important requirement is: the e-mail should only be sent if do_other_stuff() really gets saved.

Obviously I can't use @atomic around the whole view, because I wouldn't be able to commit, so my first shot looks something like:

def view(request):
    should_send_email = False
    with transaction.atomic:
        do_stuff(...)
        if condition(...):
            result = do_other_stuff(...)
            should_send_email = True
        response = HttpResponse(...)
    if should_send_email:
        should_send_email(result)
    return response


It's kind of acceptable, but it clobbers the control flow. How to rewrite it better?

I was thinking about something like:

def view(request):
    with AtomicWithContinuation as a:  # new context manager
        do_stuff(...)
        if condition(...):
            result = do_other_stuff(...)
            a.after_commit(lambda: send_email(result))
        return HttpResponse(...)


which looks easier on the eyes, so I've rolled up a context manager like:

class AtomicWithContinuation(object):

    def __init__(self, *args, **kwargs):
        self.atomic = transaction.atomic(*args, **kwargs)
        self._after_commit = []
        self._after_rollback = []

    def after_commit(self, cb):
        self._after_commit.append(cb)

    def after_rollback(self, cb):
        self._after_rollback.append(cb)

    def __enter__(self):
        self.atomic.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            self.atomic.__exit__(exc_type, exc_val, exc_tb)
        except Exception:
            # atomic.__exit__ itself threw an exception
            # - assume a rollback has been done
            for cb in self._after_rollback:
                cb()
            raise
        if exc_val is not None:
            # Wrapped block raised an exception,
            # so atomic had to do a rollback
            for cb in self._after_rollback:
                cb()
        else:
            for cb in self._after_commit:
                cb()


Does this implementation make sense? Is this a good approach to take? Also, can I trust my assumptions in __exit__?
I don't think there's anything wrong with the should_send_email version of the code. But I can see that in more complex cases there would be a readability benefit in keeping the after-commit logic next to the condition.
There's no documentation. How do I use the AtomicWithContinuation class? What methods does it have? What guarantees does it provide about which callbacks will be called, when, and in what order?
The name could be improved: callbacks are not the same as continuations.
It would make sense for the after_commit and after_rollback methods to take *args and **kwargs to be passed to the callback:

def after_commit(self, cb, *args, **kwargs):
    """Call cb, passing *args and **kwargs, when this transction commits."""
    self._after_commit.append((cb, args, kwargs))


That would allow you to omit the lambda:

a.after_commit(send_email, result)

There's repeated code for calling all the callbacks in a list. This could be made into a function:

def run_callbacks(callbacks):
    """Run the callbacks in the iterable."""
    for cb, args, kwargs in callbacks:
        cb(*args, **kwargs)

You ask, "can I trust my assumptions in __exit__?" Well, there seem to be three assumptions:


If atomic.__exit__ is called with exc_val=None, it commits the transaction.
If atomic.__exit__ is called with any other value of exc_val, it rolls back the transaction.
If atomic.__exit__ raises an exception, then the transaction was rolled back.


Assumptions 1 and 2 are guaranteed by the docstring for the Atomic class, which says:


  __exit__ commits the transaction or releases the savepoint on normal exit, and rolls back the transaction or to the savepoint on exceptions.


It would be worth adding a comment to your __exit__ method quoting or linking to this guarantee.

Assumption 3 is not guaranteed anywhere I can find, and a look at the source for the Atomic.__exit__ method shows that it tries to roll back the transaction, but this might fail, possibly ending up with the connection being closed, or the needs_rollback flag being set.