Transaction.atomic with continuations

Question

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:

1. 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
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__?

Answer 1

1. 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.

2. 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?

3. The name could be improved: callbacks are not the same as continuations.

4. 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)
   

5. 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)
   

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

   1. If atomic.__exit__ is called with exc_val=None, it commits the transaction.

   2. If atomic.__exit__ is called with any other value of exc_val, it rolls back the transaction.

   3. 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.