Python, decorators, and database idioms

Now that I’m a co-maintainer of the MySQL driver for Python I get to see a lot of confused people pop in on the support forums with some elementary problems. These seem to be mostly young developers who haven’t quite mastered the “start transaction, work with the database, commit” pattern that so many of us have become used to after a few years.

I still find the overhead of doing that sort of thing repeatedly to be tiresome. It doesn’t take long for your code to become a mess of try/except blocks, making code reading confusing and arduous. To avoid this, I’ve been using a Python decorator to wrap my SQL functions. It goes a little something like this (simplified form):

def dbwrap(func):
    """Wrap a function in an idomatic SQL transaction.  The wrapped function
    should take a cursor as its first argument; other arguments will be
    preserved.
    """
    def new_func(conn, *args, **kwargs):
        cursor = conn.cursor()
        try:
            cursor.execute("BEGIN")
            retval = func(cursor, *args, **kwargs)
            cursor.execute("COMMIT")
        except:
            cursor.execute("ROLLBACK")
            raise
        finally:
            cursor.close()

        return retval

    # Tidy up the help()-visible docstrings to be nice
    new_func.__name__ = func.__name__
    new_func.__doc__ = func.__doc__

    return new_func


@dbwrap
def do_something(cursor, val1=1, val2=2):
    """Do that database thing."""
    cursor.execute("SELECT %s, %s", (val1, val2))
    return cursor.fetchall()


class SomeClass(object):
    def __init__(self):
        conn = MySQLdb.connect(db='test')
        self.instance_var = SomeClass.get_stuff(conn)
        print self.instance_var
        conn.close()

    @staticmethod
    @dbwrap
    def get_stuff(cursor):
        """Load something meaningful from the database."""
        cursor.execute("SELECT 'blah'")
        return cursor.fetchall()[0][0]


if '__main__' == __name__:
    conn = MySQLdb.connect(db='test')
    print do_something(conn)
    print do_something(conn, 3, 4)
    print do_something(conn, 5)
    print do_something(conn, val2=6)
    #help(do_something)

    s = SomeClass()
    #help(s)

Now, there are obviously several additions you could make here. You could add retry/reconnect logic, be a little more careful when closing the cursor (which can throw exceptions), or even attempt some deadlock recovery. But I think this simple example shows how you can effectively use decorators to implement a useful DB interaction idiom on both functions and object static methods.

And, damn, if it doesn’t save a lot of typing.