Fixing celery w/ the gevent worker -- Flask, celery, gevent, app_context

celery_gevent_flask_fix.py

"""
 If you're using Celery with the gevent worker and a Flask app,
 you may have noticed that you're unsuccessful in having the 
 tasks execute within your Flask app's app_context.  Normally,
 you can do this
 
 with app.app_context():
   celery.start()

 ...but...doesn't appear to work with the gevent worker pool.
 
  I didn't find an elegant solution, instead, I wrap the tasks
  in a decorator that tests whether or not an app is configured
  on my Flask-SQLAlchemy instance.  If it is not, I know I don't
  have the app context.
  
  WARNING: pasted here just as an example in case you hit this
  issue...not complete.  See also
  https://github.com/kljensen/async-flask-sqlalchemy-example
  
  woot
"""

# Assumed to be a Flask-SQLAlchemy SQAlchemy instance,
# e.g. in app.py you did:
# from flask.ext.sqlalchemy import SQLAlchemy
# db = SQLAlchemy(app)
#
from app import db
from functools import wraps
import logging

# Assumed to be your configured celery object
from app import celery

# This is our decorator
#
def with_app_context(fn):
    """ Ensure that tasks are executed with our main application's context.
        Celery has some issues with context when it is running using the
        gevent worker pool class.  In particular, calling `celery.start()`
        under a Flask app's context does not
    """

    # Store our app, so we only have to create it once
    #
    memo = {}
    memo["app"] = None


    # Wrap the function
    #
    @wraps(fn)
    def wrapped_fn(*args, **kwargs):

        # See if we have a valid Flask app context
        # for database operations
        #
        try:
            db.get_app()
        except RuntimeError:

            # If we don't create the app, configure Flask-SQLAlchemy
            # and create the app context
            #
            from app import my_app_creating_function

            #  See if we've already created the app
            #
            if not memo["app"]:
                logging.warning("Had to create app!")
                app = my_app_creating_function()
                memo["app"] = app
            else:
                app = memo["app"]

            # Return a copy of `fn` that is wrapped in the
            # Flask app context.
            #
            with app.app_context():
                return fn(*args, **kwargs)
        return fn(*args, **kwargs)
    return wrapped_fn
    
# Now you can decorate your tasks like
#
@celery.task(name="do_my_stuff_woot")
@with_app_context
def add(x, y):
    result = db.session.query(MyModel).all()
    logging.info(result)