Postgres Connections with Python Decorators

Summary #

In this post, I will walk through creating a decorator function in Python to handle the set up and tear down of PostgreSQL connections. For the impatient who need only the code sans banter, I have conveniently created this Gist for you.

The Problem #

Recently, I was pouring over a module written for a CLI ascript in Python. This particular CLI script required a connection to a PostgreSQL database and the module handled it in the following way:

import psycopg2
import os

connection = psycopg2.connect(
host=os.getenv('DB_HOST', ''),
port=os.getenv('DB_PORT', 5432),
cursor = connection.cursor()
def process_db_stuff(data):
# Do some DB stuff

There are couple problems with this approach. Here are a few that immediately come to mind:

There must be a better way #

Enter the world of Python decorators. A Python decorator is a programing pattern where you wrap the function that handles the main logic being performed. One of the better known use cases of the decorator pattern in the wild is the Flask's app routing decorator. This is how flask registers functions to application routes (Read this great post for more information on how this works). Another great example of decorators in action are Django authentication decorators

For my use case, I will need to handle both the set up and tear down of a database connection. I will also be using a decorator where I can pass in arguments. In this case, the argument will be the information necessary to make a database connection.

With all of those requirements in mind, here is the decorator function that fit my use case. Below is the definition of this decorator and a simple use case:

def psycopg2_cur(conn_info):
"""Wrap function to setup and tear down a Postgres connection while
providing a cursor object to make queries with.

def wrap(f):
def wrapper(*args, **kwargs):
# Setup postgres connection
connection = psycopg2.connect(**conn_info)
cursor = connection.cursor()
# Call function passing in cursor
return_val = f(cursor, *args, **kwargs)
# Close connection

return return_val
return wrapper
return wrap
# Define the psycopg2 kwargs here
'host': '',
'port': '5432',
'user': 'postgres',
'password': '',
'dbname': 'postgres'
def tester(cursor):
"""Test function that uses our psycopg2 decorator

cursor.execute('SELECT 1 + 1')
return cursor.fetchall()

With this design pattern, we have moved ourselves much closer to explicit, modular code that eliminates a lot of setup and teardown boilerplate and makes our code much more readible. The above code could be extended in many ways such adding better logic for transaction handling or using a callable object instead to manage single shared connection. For my use case, command line scripts that perform a single action, setting up and tearing down a single database connection per function works perfectly.

Hope you found this useful!