Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Postgres Connections with Python Decorators
from functools import wraps
import psycopg2
####################
# Define Decorator #
####################
def psycopg2_cursor(conn_info):
"""Wrap function to setup and tear down a Postgres connection while
providing a cursor object to make queries with.
"""
def wrap(f):
@wraps(f)
def wrapper(*args, **kwargs):
try:
# Setup postgres connection
connection = psycopg2.connect(**conn_info)
cursor = connection.cursor()
# Call function passing in cursor
return_val = f(cursor, *args, **kwargs)
finally:
# Close connection
connection.close()
return return_val
return wrapper
return wrap
#################
# Example Usage #
#################
# Define the psycopg2 kwargs here
PSQL_CONN = {
'host': '127.0.0.1',
'port': '5432',
'user': 'postgres',
'password': '',
'dbname': 'postgres'
}
@psycopg2_cursor(PSQL_CONN)
def tester(cursor):
"""Test function that uses our psycopg2 decorator
"""
cursor.execute('SELECT 1 + 1')
return cursor.fetchall()
@marcospaterson
Copy link

marcospaterson commented May 26, 2020

i thought of pulling the code, amending and pushing it back but it is such a simple change that i dont think it deserves the work. you forgot to include
from functools import wraps

thanks for the code though. really helpful!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment