traceback.extract_stack(): my new favourite python function

[iguana]

Ah, SQLAlchemy, my old nemesis. I'm still stuck using SQLAlchemy 0.4.x/0.5.x at work, so I don't know if this has changed in later releases, but I hope so. Whereas things in Python like file objects will automatically close themselves when they go out of scope, something like this will leave hanging connections open:

    def __init__(self, dburi, **sa_kw):
        self.engine = create_engine(dburi, **sa_kw)

    ...

    def my_thing(self, id_):
        row = self.engine.execute(
            my_table.select().where(my_table.c.id == id_)
        ).fetchone()
        # do stuff with row
        return

Instead, my_thing() needs to look like this:

    def my_thing(self, id_):
        conn = self.engine.connect()
        row = conn.execute(
            my_table.select().where(my_table.c.id == id_)
        ).fetchone()
        # do stuff with row
        conn.close()
        return

I'm not sure why the engine should even have a execute() function if it's going to end users up in a position where they're leaking connections, but there we go.

Now, this is annoying to debug because rather than getting an error immediately, you find out about it sometime in the future when the QueuePool queue exhausts all of its available connections and raises an error on some future unfortunate self.engine.execute().

Fortunately I can edit SQLAlchemy's code pretty easily, but throwing an error (my usual way of figuring out where code is misbehaving) when connections are established doesn't work due to the queue semantics above. After that, traceback.print_stack() is pretty useful to hack in, but the stacks end up huge through SQLAlchemy's internals, and I had a lot of database reconnections going on with all those engine.execute()s around.

So I decided to write something that would tell me where in my code the connections were established and closed to see where my leak was. The following code shows a self-contained file that demonstrates what's going on:

# tb.py:
import traceback

def debug():
    stack = traceback.extract_stack()
    items = []
    for item in stack:
        if 'tb.py' in item[0]:
            items.append(item)
    print ''.join(traceback.format_list(items))

def dosomething():
    debug()

if __name__ == '__main__':
    dosomething()

Essentially, I'm creating a stack trace that only shows the lines of the stack that I'm interested in.

The relevant part of SQLAlchemy to add these into are in pool.py's QueuePool methods do_get() and do_return_conn().

With this extra output it was relatively easy to spot which lines of my code were checking out more connections than they were putting back, and I found the bug quickly after that.

Anyway, I thought it was a neat hack and I'll probably end up using it again sometime, so in the blog it goes.

Edit: It seems that engine.execute(...).fetchall() does actually close the connection, whereas fetchone() doesn't (SQLA 0.5.8). Consistent!