Django database query logger

django_query_logger.py

# Credits to https://stackoverflow.com/a/68194120

import time
from django.db import connection


class QueryLogger:
    def __init__(self):
        self.queries = []

    def __call__(self, execute, sql, params, many, context):
        current_query = {"sql": sql, "params": params, "many": many}
        start = time.time()
        try:
            result = execute(sql, params, many, context)
        except Exception as exc:
            current_query["status"] = "error"
            current_query["exception"] = exc
            raise
        else:
            current_query["status"] = "ok"
            return result
        finally:
            current_query["duration"] = time.time() - start
            self.queries.append(current_query)


class db_query_logger:  # noqa: N801
    def __init__(self):
        self._logger = QueryLogger()

    def __enter__(self):
        connection.execute_wrappers.append(self._logger)

    def __exit__(self, exc_type, exc_value, exc_traceback):
        connection.execute_wrappers.pop()
        print(
            f"Total queries: {len(self._logger.queries)}",
            f"Total duration: {sum(query['duration'] for query in self._logger.queries):.6f}s",
            sep="\n",
        )
        print()
        for query in self._logger.queries:
            print(
                f"Query: {query['sql']}",
                f"Params: {query['params']}",
                f"Many: {query['many']}",
                f"Duration: {query['duration']:.6f}s",
                sep="\n",
            )
            print()


"""
Usage example:
    
with db_query_logger():
    # put any Django queryset or function calls including Django queries here
    ...

"""