In Python, how to make sure database connection will always close before leaving a code block?

Question

I want to prevent database connection being open as much as possible, because this code will run on an intensive used server and people here already told me database connections should always be closed as soon as possible.

def do_something_that_needs_database ():
    dbConnection = MySQLdb.connect(host=args['database_host'], user=args['database_user'], passwd=args['database_pass'], db=args['database_tabl'], cursorclass=MySQLdb.cursors.DictCursor)
    dbCursor = dbConnection.cursor()
    dbCursor.execute('SELECT COUNT(*) total FROM table')
    row = dbCursor.fetchone()
    if row['total'] == 0:
        print 'error: table have no records'
        dbCursor.execute('UPDATE table SET field="%s"', whatever_value)
        return None
    print 'table is ok'
    dbCursor.execute('UPDATE table SET field="%s"', another_value)

    # a lot more of workflow done here

    dbConnection.close()

    # even more stuff would come below

I believe that leaves a database connection open when there is no row on the table, tho I'm still really not sure how it works.

Anyway, maybe that is bad design in the sense that I could open and close a DB connection after each small block of execute. And sure, I could just add a close right before the return in that case...

But how could I always properly close the DB without having to worry if I have that return, or a raise, or continue, or whatever in the middle? I'm thinking in something like a code block, similar to using try, like in the following suggestion, which obviously doesn't work:

def do_something_that_needs_database ():
    dbConnection = MySQLdb.connect(host=args['database_host'], user=args['database_user'], passwd=args['database_pass'], db=args['database_tabl'], cursorclass=MySQLdb.cursors.DictCursor)
    try:
        dbCursor = dbConnection.cursor()
        dbCursor.execute('SELECT COUNT(*) total FROM table')
        row = dbCursor.fetchone()
        if row['total'] == 0:
            print 'error: table have no records'
            dbCursor.execute('UPDATE table SET field="%s"', whatever_value)
            return None
        print 'table is ok'
        dbCursor.execute('UPDATE table SET field="%s"', another_value)
        # again, that same lot of line codes done here
    except ExitingCodeBlock:
        closeDb(dbConnection)
    # still, that "even more stuff" from before would come below

I don't think there is anything similar to ExitingCodeBlock for an exception, tho I know there is the try else, but I hope Python already have a similar feature...

Or maybe someone can suggest me a paradigm move and tell me this is awful and highly advise me to never do that. Maybe this is just something to not worry about and let MySQLdb handle it, or is it?

Answer 1

The traditional approach is the try/finally statement:

def do_something_that_needs_database ():
    dbConnection = MySQLdb.connect(host=args['database_host'], user=args['database_user'], passwd=args['database_pass'], db=args['database_tabl'], cursorclass=MySQLdb.cursors.DictCursor)
    try:
       # as much work as you want, including return, raising exceptions, _whatever_
    finally:
       closeDb(dbConnection)

Since Python 2.6 (and 2.5 with a from __future__ import with_statement), there is an alternative (although try/finally still works perfectly well!): the with statement.

with somecontext as whatever:
   # the work goes here

A context has an __enter__ method, executed on entry (to return the whatever above, if you want) and an __exit__ method, executed on exit. Despite the elegance, since there is no existing context that works the way you want, the work needed to build one (although reduced in 2.6 with contextlib) should probably suggest that good old try/finally is best.

If you have 2.6 and want to try contextlib, this is one way you could do it to "hide" the try/finally...:

import contextlib

@contextlib.contextmanager
def dbconnect(**kwds):
  dbConnection = MySQLdb.connect(**kwds)
  try:
    yield dbConnection
  finally:
    closeDb(dbConnection)

to be used as:

def do_something_that_needs_database ():
    with dbconnect(host=args['database_host'], user=args['database_user'], 
                   passwd=args['database_pass'], db=args['database_tabl'], 
                   cursorclass=MySQLdb.cursors.DictCursor) as dbConnection:
       # as much work as you want, including return, raising exceptions, _whatever_

It may be worth it if you are going to use this many, many times, just to avoid repeating the try/finally over and over for each of those many uses.