Managing Sessions with SQLAlchemy for Multiple Databases in Flask

Effectively Handling SQLAlchemy Sessions for Multiple Database Connections in Flask

In web applications, it’s not uncommon to need access to multiple databases to store and retrieve data. Flask, a lightweight Python web framework, can be easily integrated with SQLAlchemy, a popular Object Relational Mapper (ORM), to manage database connections and operations. In this blog post, we’ll discuss how to manage sessions in a Flask application with SQLAlchemy when connecting to multiple databases.

Setting Up The Flask Application

First, let’s set up a basic Flask application. To do this, you need to have Flask and SQLAlchemy installed in your Python environment. You can install them using pip:

pip install Flask SQLAlchemy

Once the libraries are installed, you can create a new Python file (e.g., app.py) and import the required libraries:

from flask import Flask, jsonify
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker

Creating Database Connections with SQLAlchemy

To connect to multiple databases, we’ll use SQLAlchemy’s create_engine function. The create_engine function takes a database URI as an argument and returns a new Engine instance:

# Define your database URIs
DATABASE_URI_1 = 'your_database_uri_1'
DATABASE_URI_2 = 'your_database_uri_2'

# Create engines for each database
engine_1 = create_engine(DATABASE_URI_1)
engine_2 = create_engine(DATABASE_URI_2)

Sessions in SQLAlchemy manage the persistence operations for ORM-mapped objects. To create a session for each database, we’ll use the sessionmaker function, which creates a factory for producing new Session instances:

# Create session factories for each engine
session_factory_1 = sessionmaker(bind=engine_1)
session_factory_2 = sessionmaker(bind=engine_2)

Scoped sessions are used to ensure that each thread maintains a separate session. We’ll create a scoped session for each session factory using the scoped_session function:

# Create scoped sessions for each session factory
Session1 = scoped_session(session_factory_1)
Session2 = scoped_session(session_factory_2)

Using Sessions with Multiple Databases

Now that we’ve set up the sessions for each database, let’s see how to use them in a Flask endpoint:

app = Flask(__name__)

@app.route('/my_endpoint')
def my_endpoint():
    # Query database 1
    session1 = Session1()
    data1 = session1.query(MyModel1).all()

    # Query database 2
    session2 = Session2()
    data2 = session2.query(MyModel2).all()

    # Process data and create response
    response = jsonify(data1=data1, data2=data2)

    # Close the sessions
    session1.close()
    session2.close()

    return response

if __name__ == '__main__':
    app.run()

In this example, we create a session for each database, query the data, and return a JSON response. It’s crucial to close the sessions after completing the operations to free up resources.

Handling Sessions in Case of Exceptions

If an exception occurs during a query, we should ensure that the sessions are properly rolled back and closed. We can use a try-except-finally block for this:

@app.route('/my_endpoint')
def my_endpoint():
    session1 = Session1()
    session2 = Session2()

    try:
        # Query database 1
        data1 = session1.query(MyModel1).all()

        # Query database 2
        data2 = session2.query(MyModel2).all()

    # Process data and create response
    response_data = {'data1': [item.serialize() for item in data1],
                     'data2': [item.serialize() for item in data2]}
    response = jsonify(response_data)
except Exception as e:
    # Rollback and log the error
    session1.rollback()
    session2.rollback()
    app.logger.error(f'Error occurred while querying databases: {e}')
    response = jsonify({'error': 'An error occurred while processing your request'})
    response.status_code = 500
finally:
    # Close the sessions
    session1.close()
    session2.close()

return response

In this example, if any exception occurs during the query or processing of the data, the sessions are rolled back, the error is logged, and an error response is sent to the client.

What’s important is to ensure that seperate sessions are created for each database connection, using the corresponding session object when querying the databases, and handling exceptions to ensure sessions are properly closed. By following the steps outlined in this blog post, you can effectively manage sessions for multiple database connections in your Flask application using SQLAlchemy.