Postgres Connections with Python Decorators

pyscopg2_decorator.py

####################
# 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()