Last active
April 23, 2025 14:48
-
-
Save travishathaway/b67c32f6fed3bcc9cb8ac72e611961bb to your computer and use it in GitHub Desktop.
Postgres Connections with Python Decorators
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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!