Python Decorators and Exception Handling

On a recent project, I was using Python Flask (behind Apache/WSGI) to answer REST calls and accessed data in a relational database.  Both MySQL and PostgreSQL datastores were supported.  Following the coding DRY (Don’t Repeat Yourself) principle, I tried to re-use database access code as much as possible, but often I had to create yet another method with a tweak in the query logic.  This lead to the traditional try-catch-finally blocks appearing in many access methods to ensure things sailed smoothly on any exceptions from malformed REST calls any other unforeseen inconsistencies.

This seems tolerable at first, but when improving how things are logged or how the locking mechanism works, I quickly felt the pain.  I finally found the time to learn how to use decorators can help and decided to share it!

Here is an example of the old way:

def simple_select_all_query(conn, table_str):
    retval = None
    cursor = conn.cursor()
    try:
        cursor.execute("BEGIN")
        cursor.execute("select * from " + table_str)
        retval = cursor.fetchall()
        cursor.execute("COMMIT")
    except Exception, e:
        cursor.execute("ROLLBACK")
        print "Exception in simple_select_all_query():", str(e)
    finally:
         cursor.close()
    return retval

and is called via

result = simple_select_all_query(conn, table_str)

Making additional database access functions is a pain following this approach.  Try a decorator that will be defined further below.   Remove all the try-catch-finally mess,  and change the function’s first argument to be a cursor instead of a connection.

@database_query_wrapper_as_func
def simple_select_all_query(cursor, table_str):
    cursor.execute("select * from " + table_str)
    return cursor.fetchall()

The database query wrapper decorator can be a function or a class.  I’ll show you both. As a function it can look like:

def database_query_wrapper_as_func(func):
    def new_func(conn, *args, **kwargs):
        retval = None
        cursor = conn.cursor()
        try:
            cursor.execute("BEGIN")
            retval = func(cursor, *args, **kwargs)
            cursor.execute("COMMIT")
        except Exception, e:
            cursor.execute("ROLLBACK")
            print "Exception in simple_select_all_query():", str(e)
        finally:
            cursor.close()
        return retval

    return new_func

The new function created here contains all the exception handling instead.  It accepts a database connection, creates a cursor, and calls the wrapped function at line 7.  For anyone new to the *args or **kwargs, it is Python’s way of passing the remaining argument(s) to the new function.

It is important to emphasize that the function is called the same way as above but the function definition has the first argument as a cursor object instead of a connection object.

As promised, here is the wrapper as a class that has the same functionality:

class database_query_wrapper_as_class(object):
    def __init__(self, func):
        self.func = func
    def __get__(self, obj, type=None):
        return self.__class__(self.func.__get__(obj, type))
    def __call__(self, conn, *args, **kw):
        cursor = conn.cursor()
        try:
            cursor.execute("BEGIN")
            retval = self.func(cursor, *args, **kw)
            cursor.execute("COMMIT")
        except Exception, e :
            cursor.execute("ROLLBACK")
            retval = None
        finally:
            cursor.close()
        return retval

This has the advantage of being a little more organized, but it’s a bit of style choice in my opinion. The function being decorated with this class should have the @database_query_wrapper_as_class above the definition.

If your program is object-oriented the decorators above can wrap class methods without any modification. First, the call to a wrapped method:

    self.create_table(self.conn, table_str, schema_str)

and the full method is a single line!

    @database_query_wrapper_as_class
    def create_table(self, cursor, table_str, schema_str):
        cursor.execute("create table " + table_str + schema_str)

It’s a little strange to pass a class object to another class method, but it’s required with this decorator approach. I suppose you could have the wrapper as part of the class, but that topic may be for the next blog as this is more than I had planned to write.

Like I mentioned earlier, I wanted to improve the logging from regular print statements to using Python’s logging package. Here is snippet of code that could be placed into above except block. This would be a pain to copy and paste in to all those functions without decorators!

    ...
    except Exception, e:
        logging.warning('Exception in %s' % self.func)
        template = "An exception of type {0} occured. Arguments:\n{1!r}"
        message = template.format(type(e).__name__, e.args)
        logging.exception(message)
    ...

Lastly, I expanded the application of database query method wrapping to more generally handle exceptions in other functions with this “general purpose” wrapper to handle any exceptions. This just passes any number of arguments through to your wrapped function.

class general_function_handler(object):
    def __init__(self, func):
        self.func = func
    def __get__(self, obj, type=None):
        return self.__class__(self.func.__get__(obj, type))
    def __call__(self, *args, **kwargs):
        try:
            retval = self.func(*args, **kwargs)
        except Exception, e :
            ... # logging
            # TODO filter which exceptions are fatal and non-fatal for your purpose:
            # Fatal: AttributeError, KeyError, NameError, ProgrammingError, etc.
            sys.exit(1) # Exit on all exceptions for now
        return retval

There are many variants and opinions on how you can do this, but here is a simple place to start.

References (apologies if I left any out or misinterpreted your post!):

http://www.kylev.com/2009/05/22/python-decorators-and-database-idioms/

http://stackoverflow.com/questions/9823936/python-how-do-i-know-what-type-of-exception-occured

http://blog.dscpl.com.au/2014/01/how-you-implemented-your-python.html

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s