Passing connections to functions using a decorator

In many script I write, there are functions requiring database operations. Every time I need them, I try to write such code in functions like:

@with_connection
def do_some_job(cnn, arg1, arg2=None):
    cur = cnn.cursor()
    cur.execute(SQL, (arg1, arg2)) # or something else

do_some_job(42, arg2='hi')

There are many benefits for this pattern:

the client code doesn't depend on the details of connection creation and disposal
the client code doesn't either need to deal with the session management: no more pending transactions...
connection management is consistent throughout all the functions.

Implementation of such a decorator is straightforward, but because of the many possible variation it never made its way into a proper library. It goes something like:

def with_connection(f):
    def with_connection_(*args, **kwargs):
        # or use a pool, or a factory function...
        cnn = psycopg.connect(DSN)
        try:
            rv = f(cnn, *args, **kwargs)
        except Exception, e:
            cnn.rollback()
            raise
        else:
            cnn.commit() # or maybe not
        finally:
            cnn.close()

        return rv

    return with_connection_

There may be many variations: you may use use a connection pooler, you may want a read only or an autocommit connection... Best strategy I've found is just to keep one or a few decorator like this in the project I'm working at and use it for a consistent management of the database communications in the desired scope (the project, a package, a module...)

Comments (2 posted).

Another approach that I like even more is context managers. Something like this:

note: reformatted after introduced reST comments.

def do_some_job():
    do_something_non_db_related()

    with get_connection() as conn:
        cur = conn.cursor()
        cur.execute(stmt)

    do_something_more_without_db()

and the context manager itself is something like:

@contextmanager
def get_connection():
    conn = get_from_pool()
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise
    else:
        conn.commit()
    finally:
        return_to_pool(conn)

Thank you, Dan.

I've introduced reST syntax in comments (this is pretty much the Django comments system out of the box) and I've reformatted your comment: I hope you don't mind.

Feel free to add a new comment. You can use reST syntax.

psycopg

Passing connections to functions using a decorator

Articles archive