count_queries.py

from decimal import Decimal

from django.db import connection
from django.utils.decorators import ContextDecorator


class CountQueries(ContextDecorator):
    '''
    Either logs(if logger instance has passed) or prints number of queries.

    Example:
        1)
            with CountQueries(log, 'with block'):
                obj = Model.objects.get(id=1)
                val = obj.name
            >>> DEBUG ... [with block] 1 query executed
        2)
            @CountQueries(log, 'decorator')
            def func():
                obj = Model.objects.get(id=1)
                val = obj.name
            func()
            >>> DEBUG ... [decorator] 1 query executed
    Params:
        - log - logger, not required
        - msg_prefix - str, not required
        - log_queries - bool, not required; to log or not SQL queries
    '''
    def __init__(self, log=None, msg_prefix=None, log_queries=False):
        self.log = log
        self.msg_prefix = msg_prefix
        self.log_queries = log_queries

    def __enter__(self):
        self.force_debug_cursor = connection.force_debug_cursor
        connection.force_debug_cursor = True
        connection.ensure_connection()
        self.start_len = len(connection.queries)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        connection.force_debug_cursor = self.force_debug_cursor
        if exc_type is not None:
            return
        count = len(connection.queries) - self.start_len

        # count time and output SQL queries if requested
        time_sum = 0
        for query in connection.queries[self.start_len:]:
            time_sum += Decimal(query['time'])
            if self.log_queries:
                msg = 'SQL query: {}'.format(query['sql'])
                if self.msg_prefix:
                    msg = '[{}] {}'.format(self.msg_prefix, msg)
                if self.log:
                    self.log.debug(msg)
                else:
                    print(msg)

        # final message
        msg = '{} quer{} executed in {} second{}'.format(
            count,
            'y' if count == 1 else 'ies',
            time_sum,
            '' if time_sum == 1 else 's',
        )
        if self.msg_prefix:
            msg = '[{}] {}'.format(self.msg_prefix, msg)
        if self.log:
            self.log.debug(msg)
        else:
            print(msg)