SQLAlchemy Session Stats

sessstats.py

import sys
import inspect
import time
import traceback

from sqlalchemy import event
from sqlalchemy.orm.session import sessionmaker
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine

class QueryStats(object):
    def __init__(self):
        self.count = 0
        self.total_time = 0.0

        self._query_log = []
        self.query_info = None

    def start_query(self, statement, parameters):
        if self.query_info is not None:
            self._query_log.append( self.query_info + (float('nan'),) )

        stack = extract_user_stack(2)

        self.query_info = (time.time(), statement, parameters, stack)
        self.count += 1

    def end_query(self):
        assert self.query_info is not None

        query_end = time.time()
        self.total_time += query_end - self.query_info[0]
        self._query_log.append( self.query_info + (query_end - self.query_info[0],) )

        self.query_info = None

    @property
    def query_log(self):
        if self.query_info is not None:
            return self._query_log + [ self.query_info + (float('nan'),) ]
        return self._query_log

    def __repr__(self):
        return '<QueryStats count=%d time=%.2fs>' % (self.count, self.total_time)

@event.listens_for(Engine, 'before_cursor_execute')
def on_query_start(conn, cursor, statement, parameters, context, executemany):
    if hasattr(conn, '_query_stats'):
        conn._query_stats.start_query(statement, parameters)

@event.listens_for(Engine, 'after_cursor_execute')
def on_query_end(conn, cursor, statement, parameters, context, executemany):
    if hasattr(conn, '_query_stats'):
        conn._query_stats.end_query()

def setup_session_tracking(session):
    """ Enables the stats tracking for a Session or a Session factory """
    @event.listens_for(session, 'after_begin')
    def on_begin(session, transaction, connection):
        if not hasattr(session, '_query_stats'):
            session._query_stats = QueryStats()

        connection._query_stats = session._query_stats

    @event.listens_for(session, 'after_commit')
    @event.listens_for(session, 'after_rollback')
    def on_commit(session):
        print "Done %s" % session._query_stats

def extract_user_stack(skip=1):
    """Like traceback.extract_stack, but skips any sqlalchemy modules and only returns filenames """
    try:
        raise ZeroDivisionError
    except ZeroDivisionError:
        f = sys.exc_info()[2].tb_frame.f_back 
    list = []

    while f is not None and skip > 0:
        f = f.f_back
        skip -= 1

    skipping = True
    while f is not None:
        lineno = f.f_lineno
        co = f.f_code
        module = inspect.getmodule(co)
        if skipping and not module.__name__.startswith('sqlalchemy'):
            skipping = False
        if not skipping:
            name = co.co_name  
            list.append((module.__name__, name, lineno))
        f = f.f_back
    list.reverse()
    return list

from sqlalchemy import event
from sqlalchemy.schema import Table, Column, MetaData
from sqlalchemy.types import Integer
from sqlalchemy import select

def test():
    engine = create_engine('postgres:///kleptog')
    SessionMaker = sessionmaker(bind=engine)

    setup_session_tracking(SessionMaker)

    m = MetaData()
    some_table = Table('pg_stats', m, Column('avg_width', Integer))

    q = select([some_table.c.avg_width]).where(some_table.c.avg_width == 1)

    s1 = SessionMaker()
    s2 = SessionMaker()

    s1.execute("SELECT 1")
    try:
        s2.execute(q)
    except:
        pass

    s1.commit()
    s2.commit()

    print s1._query_stats.query_log
    print s2._query_stats.query_log

test()