What's the proper way to use SQLAlchemy Sessions with Celery?

Viewed 933

I've been trying to figure out this for a while but a lot of answers I can find now are out-dated (posts from > 6 years ago), or less related.

The question is really how to deal with database sessions properly in celery. My current setting is: I have a global DbEngine object that contains dsn, engine, and Session. so every time if I want to use a session, I just call sess = db.Session() and start to use them inside my Flask app. It looks like the following:

#db.py
class DbEngine:
  def __init__(self, path, ...):
    self.dsn = self.create_dsn_from_file(path)
    self.engine = create_engine(self.dsn)
    self.Session = scoped_session(sessionmaker(bind=self.engine))

Until I introduced celery to my app, often I got various errors (Protocol error, sslSocket errors, etc). And I could not reproduce them locally, and it often gets fixed if I simply add retries to those tasks (it usually goes through with 3 retries). So I suspect that this might be caused by the Shared Sessions.

Then I decided to change my session for the celery app: every time if I need a session, I will actually create a new engine, create a new Session, and return the newly created Session(). However, the similar but not exactly the same problem happened again (various protocol errors with different codes).

I saw that celery has its own SessionManager, but I couldn't find sample code of usage. I was thinking about the following structure:

# celery_app.py

celery = ... 

session_manager = SessionManager()
engine, Session = session_manager.create_session(dsn)

and in the tasks I do:

# task_1.py
from celery_app import celery, Session

@celery.task
def tsk():
  sess = Session()
  sess.query(...)
  ... 
  sess.close()

But really not sure if it's the intended approach, since I am only calling the .create_session() once, and I don't know when and how I should call the other functions in the SessionManager.

For the background, I am using DB2.

If anyone knows how to properly use SessionManager or has experience in using session with SQLAlchemy that'd be much appreciated if you can provide any insights. Thank you in advance!

2 Answers

Maybe not the answer you're looking for, but flask_sqlalchemy provides a Flask-friendly layer over SQLAlchemy that almost entirely removes the need to manage sessions manually. And it plays well with celery once you set up a ContextTask per the example in the Flask docs.

Posting this hoping it helps someone out. In my case, we're using Fastapi, Postgres, Sqlalchemy and Celery.

I had the same issue and it was very challenging. Lots of cryptic errors from our postgres db looking like the connections were being stopped mid stream. PG_TUPLES_OK and no message from the libpq, AttributeError("'_NoResultMetaData' object has no attribute '_indexes_for_keys'") ... etc etc

TL;DR

Sqlalchemy pools connections by default in a non-threadsafe manner, Celery forks processes by default: one or the other needs to be changed.

Solution 1) Turn off Sqlalchemy pooling *we ended up going with this to maintain better concurrency

Sql Alchemy Docs

from sqlalchemy.pool import NullPool
engine = create_engine(
    SQLALCHEMY_DATABASE_URL, poolclass=NullPool
)

Solution 2) Make Celery run as a single process with no forking

celery -A celery_worker.celery worker -E --loglevel=info --pool=solo

Article on Celery pooling This solution worked, but meant that all of our celery tasks would be running serially, which eliminated the errors but we needed more bandwidth. This may be fine for certain applications.

Related