Transaction Aware Celery Task - SQLAlchemy

tasks.py

# Based on Django version found here:
# https://gist.github.com/tapanpandita/a84938354d32b0f12423d69259e06c2c
from celery import Task
from sqlalchemy import event

from some_module import DBSession


class TransactionAwareTask(Task):
    """
    Use this task as a base task if you want to guarantee that a task runs only AFTER a transaction is
    successfully committed.
    This will help avoid race conditions (when the task runs before a transaction finishes).
    """
    abstract = True

    def apply_async(self, *args, **kwargs):
        """
        NB: This task doesn't return AsyncResult like its parent normally does.
        """
        event.listens_for(DBSession, 'after_commit')(
            lambda _session: super(TransactionAwareTask, self).apply_async(*args, **kwargs)
        )