How to count sqlalchemy queries in unit tests

Question

In Django I often assert the number of queries that should be made so that unit tests catch new N+1 query problems

from django import db
from django.conf import settings
settings.DEBUG=True

class SendData(TestCase):
    def test_send(self):
        db.connection.queries = []
        event = Events.objects.all()[1:]
        s = str(event) # QuerySet is lazy, force retrieval
        self.assertEquals(len(db.connection.queries), 2)

In in SQLAlchemy tracing to STDOUT is enabled by setting the echo flag on engine

engine.echo=True

What is the best way to write tests that count the number of queries made by SQLAlchemy?

class SendData(TestCase):
    def test_send(self):
        event = session.query(Events).first()
        s = str(event)
        self.assertEquals( ... , 2)

Answer 1

I've created a context manager class for this purpose:

class DBStatementCounter(object):
    """
    Use as a context manager to count the number of execute()'s performed
    against the given sqlalchemy connection.

    Usage:
        with DBStatementCounter(conn) as ctr:
            conn.execute("SELECT 1")
            conn.execute("SELECT 1")
        assert ctr.get_count() == 2
    """
    def __init__(self, conn):
        self.conn = conn
        self.count = 0
        # Will have to rely on this since sqlalchemy 0.8 does not support
        # removing event listeners
        self.do_count = False
        sqlalchemy.event.listen(conn, 'after_execute', self.callback)

    def __enter__(self):
        self.do_count = True
        return self

    def __exit__(self, *_):
        self.do_count = False

    def get_count(self):
        return self.count

    def callback(self, *_):
        if self.do_count:
            self.count += 1

Answer 2

what about the approach of using flask_sqlalchemy.get_debug_queries() btw. this is the methodology used by internal of Flask Debug Toolbar check its source

from flask_sqlalchemy import get_debug_queries


def test_list_with_assuring_queries_count(app, client):
    with app.app_context():
        # here generating some test data
        for _ in range(10):
            notebook = create_test_scheduled_notebook_based_on_notebook_file(
                db.session, owner='testing_user',
                schedule={"kind": SCHEDULE_FREQUENCY_DAILY}
            )
            for _ in range(100):
                create_test_scheduled_notebook_run(db.session, notebook_id=notebook.id)
    with app.app_context():
        # after resetting the context call actual view we want asserNumOfQueries
        client.get(url_for('notebooks.personal_notebooks'))
        assert len(get_debug_queries()) == 3

keep in mind that for having reset context and count you have to call with app.app_context() before the exact stuff you want to measure.

Answer 3

Slightly modified version of @omar-tarabai's solution that removes the event listener when exiting the context:

from sqlalchemy import event

class QueryCounter(object):
    """Context manager to count SQLALchemy queries."""

    def __init__(self, connection):
        self.connection = connection.engine
        self.count = 0

    def __enter__(self):
        event.listen(self.connection, "before_cursor_execute", self.callback)
        return self

    def __exit__(self, *args, **kwargs):
        event.remove(self.connection, "before_cursor_execute", self.callback)

    def callback(self, *args, **kwargs):
        self.count += 1

Usage:

with QueryCounter(session.connection()) as counter:
     session.query(XXX).all()
     session.query(YYY).all()

print(counter.count)  # 2