-
-
Save hrabryi/6b208b5e92b658bceddb8a51f92679c0 to your computer and use it in GitHub Desktop.
Postgres Connections with Python Decorators
This file contains 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
#################### | |
# 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