SQLAlchemy listen to connection/disconnection events

Question

TL;DR: Is it possible to listen to events of DB Connect and Disconnect with SQLAlchemy ORM

Background:

I have a Flask app that connects to a Postgres DB using SQLAlchemy ORM. (Flask-SQLAlchemy). Everything works well, there are no bugs with the below code

from flask import Flask
from flask_sqlalchemy import SQLAlchemy

app = Flask(__name__)
app.config['SQLALCHEMY_DATABASE_URI'] = 'postgresql://<username>:<password>@<host>:<port>/<db>'
db = SQLAlchemy(app)

Due to network issues, the connection to the DB fails or drop. Which generates errors when DB is accessed. By design SQLAlchemy handles all connections under the hood, and it is seamless when the connection drops and re-establishment are handled.

Requirement

To monitor and handle these connection drops, I need events triggered for connect and disconnect

Reading through docs I found connect, close and engine_connect events that can be tied to the engine

def connect(dbapi_connection, connection_record):
    print('DB Connected')


def disconnect(dbapi_connection, connection_record):
    print('DB Disconnected')


event.listen(db.engine, 'connect', connect)
event.listen(db.engine, 'engine_connect', connect)
event.listen(db.engine, 'close', connect)

Similarly, there are other events like close_detached first_connect etc.

Note that SQLAlchemy establishes a somewhat "lazy connection" i.e. only when a call to the DB is made for the first time will the actual connection be established hence db.engine.connect() needs to be called for the connect events to fire.

Expected Behaviour

1. Whenever a connection to the DB is established whether at the start of the application or during the lifecycle. The callback function connect() should be fired
2. Whenever a connection is lost due to DB going down or a network issue, the callback function disconnect() should be fired

Current Behaviour

1. On establishing a connection the first time connect() is called. No matter if the connection goes down and connects back connect() is not called again
2. On loss of connection disconnect() is never called, I might be listening to the wrong event possibly.

For testing, I have PG Server installed locally, I am starting and stopping the service. And verifying the connection using psql cli

---

Question

I have looked into SQLAlchemy Core to find the appropriate events and polling mechanism but have not found anything yet that works.

SQLAlchemy does not look at the connection when not in use, hence polling seems to be required. Does that mean I need to run a query every n seconds as db.connection() does not throw an error in my case when the connection is lost midway.

I understand the limitations of Flask running in a single process, not supporting async calls intrinsically. I would like to know how to achieve it even if it were not Flask, what kind of polling would be required

I am not ready to accept that a tool as mature as SQLAlchemy does not support this, but I spent a few days and have not found a working solution yet. Any help or direction is appreciated.

Bonus Question:

What would be the appropriate way to check if the connection to DB is alive on Flask startup. Currently, I call db.engine.connect() after initializing which throws an error if DB is not connected.

Related SO Question: How to verify SqlAlchemy engine object

More information on how SQLAlchemy Handles Disconnect - Dealing with Disconnects

Answer 1

Hmm, I am using the connect event of Sqlalchemy to set the time zone. I can confirm that the connect event works reliably when I manually restart the database. I'm using sqlalchemy 1.3.24 with flask_sqlalchemy 2.5.1 and MariaDB 10.3.17.

My code basically says:

app.config.update(SQLALCHEMY_DATABASE_URI="mysql+pymysql://..."))
db = SQLAlchemy(app)
event.listen(db.engine, 'connect', set_time_zone)

If I now start the database, start the application, and perform a request, my connect callback is called. If I restart the database now, the next request will fail (as expected) with:

sqlalchemy.exc.OperationalError: (pymysql.err.OperationalError) (1927, 'Connection was killed')

After that, the next request will call the connect callback again, and the applications works as normal.

The only thing I have to be extremely careful about is that I rollback the session when the OperationalError is raised. If I don't do that, flask needs to be restarted.

So every api endpoint looks like this:

def whatever():
    from .. import db
    session = db.session

    try:
        # do your job
    except Exception:
        logger.error("Error: Doing whatever failed")
        print_exc()
        session.rollback()