Skip to content

Instantly share code, notes, and snippets.

@travishathaway
Last active July 9, 2023 16:38
Show Gist options
  • Star 10 You must be signed in to star a gist
  • Fork 2 You must be signed in to fork a gist
  • Save travishathaway/b67c32f6fed3bcc9cb8ac72e611961bb to your computer and use it in GitHub Desktop.
Save travishathaway/b67c32f6fed3bcc9cb8ac72e611961bb to your computer and use it in GitHub Desktop.
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

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