Debug queries in Django

debug_queries.py

import csv
import functools
import logging
import time

from django.db import reset_queries, connection
from django.conf import settings

logger = logging.getLogger(__name__)


def debug_queries(func=None, *, delimiter=','):
    """
    Allow storing the query that are run on a function.
    Examples:
        >>>
        >>> @debug_queries
        >>> def func():
        >>>     pass
        >>>
        >>> @debug_queries(arg='foo')
        >>> def func():
        >>>     pass
        >>>
    """

    def store_queries_csv(func_name, queries, delimiter=','):
        store_queries_csv.counter += 1
        csv_file = f"{func_name}-{store_queries_csv.counter}.csv"
        csv_columns = ['sql', 'time']

        with open(csv_file, 'w') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=csv_columns, delimiter=delimiter)
            writer.writeheader()
            for data in queries:
                writer.writerow(data)

    store_queries_csv.counter = 0
    # 1. Decorator arguments are applied to itself as partial arguments
    if func is None:
        return functools.partial(debug_queries, delimiter=delimiter)

    # 2. logic with the arguments
    #    ...

    # 3. Handles the actual decorating
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # Write decorator function logic here
        # Before function call
        # ...
        # reset all django.db.connection.queries
        reset_queries()

        start_queries = len(connection.queries)
        start_time = time.perf_counter()

        old_debug_value = settings.DEBUG
        settings.DEBUG = True
        result = func(*args, **kwargs)
        settings.DEBUG = old_debug_value
        total_queries_execution_time = sum([float(q['time']) for q in connection.queries])

        end_time = time.perf_counter()
        end_queries = len(connection.queries)

        queries = end_queries - start_queries
        time_taken = end_time - start_time
        store_queries_csv(func.__name__, connection.queries, delimiter)
        logger.info(f">>> DEBUG QUERY - Function : {func.__name__}")
        logger.info(f">>> DEBUG QUERY - Number of Queries : {queries}")
        logger.info(f">>> DEBUG QUERY - Finished in : {(time_taken):.2f}s")
        logger.info(f">>> DEBUG QUERY - Total time : {(total_queries_execution_time):.2f}s")
        # After function call
        # ...
        return result

    return wrapper