Flask-SQLAlchemy has some nice built-ins (e.g. accessing query directly on classes). To continue leveraging these nicities while still inside of a Celery worker, we need to make sure we setup/teardown in a similar fashion to Flask-SQLAlchemy does on Flask.
Flask-SQLAlchemy uses create_scoped_session at startup which avoids any setup on a per-request basis.
https://github.com/mitsuhiko/flask-sqlalchemy/blob/2.0/flask_sqlalchemy/__init__.py#L668
This means Celery can piggyback off of this initialization.
Flask-SQLAlchemy tears down when we leave the request/application context.
https://github.com/mitsuhiko/flask-sqlalchemy/blob/2.0/flask_sqlalchemy/__init__.py#L747-L753
To create a similar experience, we can run teardown when a Celery task ends:
We put in a clause about
CELERY_ALWAYS_EAGERto prevent conflicts with Flask's normal execution.
from celery.signals import task_postrun
def handle_celery_postrun(retval=None, *args, **kwargs):
"""After each Celery task, teardown our db session"""
if app.config['SQLALCHEMY_COMMIT_ON_TEARDOWN']:
if not isinstance(retval, Exception):
db.session.commit()
# If we aren't in an eager request (i.e. Flask will perform teardown), then teardown
if not app.config['CELERY_ALWAYS_EAGER']:
db.session.remove()
task_postrun.connect(handle_celery_postrun)
Since celery hooks like this fail silently (in our version of celery at least), I highly recommend using
app.config.gethere to prevent unset configuration values from raising hiddenKeyErrors. (We hadSQLALCHEMY_COMMIT_ON_TEARDOWNget unset and silently started leaking sessions between celery tasks.)