import time
from django.db import connection


class TimingMiddleware:
    """Middleware that measures request total time and DB query time.

    Note: `connection.queries` is populated only when Django debug is enabled
    (DEBUG=True) or when using a debug cursor. When not available, DB time
    will be left as None.
    """

    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        start = time.time()

        # record initial number of collected queries (if available)
        try:
            initial_q_count = len(connection.queries)
        except Exception:
            initial_q_count = 0

        response = self.get_response(request)

        total = time.time() - start

        db_time = None
        try:
            qs = getattr(connection, 'queries', None)
            if qs:
                # sum only queries executed during this request
                new_queries = qs[initial_q_count:]
                db_time = sum(float(q.get('time', 0)) for q in new_queries)
        except Exception:
            db_time = None

        # attach to request so context processor can read them
        request.db_time = db_time
        request.total_time = total

        return response