Psycopg2 & Flask - tying connection to before_request & teardown_appcontext

Question

Cheers guys, refactoring my Flask app I got stuck at tying the db connection to @app.before_request and closing it at @app.teardown_appcontext. I am using plain Psycopg2 and the app factory pattern.

First I created a function to call wihtin the app factory so I could use @app as suggested by Miguel Grinberg here:

def create_app(test_config=None):
    app = Flask(__name__, instance_relative_config=True)

    --

    from shop.db import connect_and_close_db
    connect_and_close_db(app)

    --

    return app

Then I tried this pattern suggested on http://flask.pocoo.org/docs/1.0/appcontext/#storing-data:

def connect_and_close_db(app):

    @app.before_request
    def get_db_test():
        conn_string = "dbname=testdb user=testuser password=test host=localhost"
        if 'db' not in g:
            g.db = psycopg2.connect(conn_string)
        return g.db

    @app.teardown_appcontext
    def close_connection(exception):
        db = g.pop('db', None)

        if db is not None:
            db.close()

It resulted in:

TypeError: 'psycopg2.extensions.connection' object is not callable

Anyone has an idea what happend and how to make it work?

Furthermore I wonder how I would access the connection object for creating a cursor once its creation is tied to before_request?

Answer 1

This solution is probably far from perfect, and it's not really DRY. I'd welcome comments, or other answers that build on this.

To implement for raw psycopg2 support, you probably need to take a look at the connection pooler. There's also a good guide on how to implement this outwith Flask.

The basic idea is to create your connection pool first. You want this to be established when the flask application initializes (This could within the python interpreter or via gunicorn worker of which there may be several - in which case each worker has its own connection pool). I chose to store the returned pool in the config:

from flask import Flask, g, jsonify

import psycopg2
from psycopg2 import pool

app = Flask(__name__)

app.config['postgreSQL_pool'] = psycopg2.pool.SimpleConnectionPool(1, 20,
    user = "postgres",
    password = "very_secret",
    host = "127.0.0.1",
    port = "5432",
    database = "postgres")

Note the first two arguments to SimpleConnectionPool are the min & max connections. That's the number of connections going to your database server, bwtween 1 & 20 in this case.

Next define a get_db function:

def get_db():
    if 'db' not in g:
        g.db = app.config['postgreSQL_pool'].getconn()
    return g.db

The SimpleConnectionPool.getconn() method used here simply returns a connection from the pool, which we assign to g.db and return. This means when we call get_db() anywhere in the code it returns the same connection, or creates a connection if not present. There's no need for a before.context decorator.

Do define your teardown function:

@app.teardown_appcontext
def close_conn(e):
    db = g.pop('db', None)
    if db is not None:
        app.config['postgreSQL_pool'].putconn(db)

This runs when the application context is destroyed, and uses SimpleConnectionPool.putconn() to put away the connection.

Finally define a route:

@app.route('/')
def index():
    db = get_db()
    cursor = db.cursor()

    cursor.execute("select 1;")
    result = cursor.fetchall()
    print (result)

    cursor.close()
    return jsonify(result)

This code works for me tested against postgres runnning in a docker container. A few things which probably should be improved:

• This view isn't very DRY. Perhaps you could move some of this into the get_db function so it returns a cursor. (!!!)

• When the python interpreter exits, you should also find away to close the connection with app.config['postgreSQL_pool'].closeall

• Although tested some kind of way to monitor the pool would be good, so that you could watch pool/db connections under load and make sure the pooler behaves as expected.

(!!!)In another land, the sqlalchemy.scoped_session documentation explains more things relating to this, with some theory on how its 'sessions' work in relation to requests. They have implemented it in such a way that you can call Session.query('SELECT 1') and it will create the session if it doesn't already exist.

---

EDIT: Here's a gist with your app factory pattern, and sample usage in the comment.