Add Middleware and improve transparent patching

Author: Ignacio Avas

request_query_count/__init__.py

@@ -1 +1,2 @@
 __version__ = '0.1.0'
+default_app_config = 'request_query_count.apps.RequestQueryCountConfig'

request_query_count/apps.py

@@ -1,50 +1,142 @@
 # -*- coding: utf-8
+import json
+import os
+import os.path
+import threading
+
 from django.apps import AppConfig
 from django.conf import settings
-from django.core.checks import Error, register
-from django.utils import inspect
-from django.utils.module_loading import import_string
+from django.test import TransactionTestCase
+from django.test.utils import get_runner
+
+from request_query_count.query_count import (TestCaseQueryContainer,
+                                             TestResultQueryContainer)
 
-from request_query_count.query_count import Middleware
+local = threading.local()
 
 
 class RequestQueryCountConfig(AppConfig):
+    LOCAL_TESTCASE_CONTAINER_NAME = 'querycount_test_case_container'
+    LOCAL_RESULT_CONTAINER_NAME = 'querycount_result_container'
+
     name = 'request_query_count'
     verbose_name = 'Request Query Count'
 
-    def ready(self):
-        pass
-
-
-@register
-def check_middleware(_app_configs, **_kwargs):
-    errors = []
-
-    setting = getattr(settings, 'MIDDLEWARE', None)
-    setting_name = 'MIDDLEWARE'
-    if setting is None:
-        setting = settings.MIDDLEWARE_CLASSES
-        setting_name = 'MIDDLEWARE_CLASSES'
-
-    if not any(is_middleware_class(Middleware, middleware)
-               for middleware in enumerate(setting)):
-        errors.append(
-            Error(
-                "debug_toolbar.middleware.DebugToolbarMiddleware is missing "
-                "from %s." % setting_name,
-                hint="Add debug_toolbar.middleware.DebugToolbarMiddleware to "
-                "%s." % setting_name,
-            )
+    setting_name = 'REQUEST_QUERY_COUNT'
+
+    default_settings = {
+        'ENABLE': True,
+        'ENABLE_STACKTRACES': True,
+        'DETAIL_PATH': 'reports/query_count_detail.json',
+        'SUMMARY_PATH': 'reports/query_count.json'
+    }
+
+    @classmethod
+    def get_setting(cls, setting_name):
+        return (getattr(settings, cls.setting_name, {})
+                .get(setting_name, cls.default_settings[setting_name]))
+
+    @classmethod
+    def stacktraces_enabled(cls):
+        return cls.get_setting('ENABLE_STACKTRACES')
+
+    @classmethod
+    def enabled(cls):
+        return cls.get_setting('ENABLE')
+
+    @classmethod
+    def get_testcase_container(cls):
+        return getattr(local, cls.LOCAL_TESTCASE_CONTAINER_NAME)
+
+    @classmethod
+    def add_middleware(cls):
+        setting = getattr(settings, 'MIDDLEWARE', None)
+        setting_name = 'MIDDLEWARE'
+        if setting is None:
+            setting = settings.MIDDLEWARE_CLASSES
+            setting_name = 'MIDDLEWARE_CLASSES'
+
+        setattr(
+            settings,
+            setting_name,
+            setting + ('request_query_count.middleware.Middleware',)
         )
-    return errors
-
-
-def is_middleware_class(middleware_class, middleware_path):
-    try:
-        middleware_cls = import_string(middleware_path)
-    except ImportError:
-        return
-    return (
-        inspect.isclass(middleware_cls) and
-        issubclass(middleware_cls, middleware_class)
-    )
+
+    @classmethod
+    def wrap_set_up(cls, set_up):
+        def wrapped(self, *args, **kwargs):
+            result = set_up(self, *args, **kwargs)
+            setattr(local, cls.LOCAL_TESTCASE_CONTAINER_NAME,
+                    TestCaseQueryContainer())
+            return result
+
+        return wrapped
+
+    @classmethod
+    def wrap_tear_down(cls, tear_down):
+        def wrapped(self, *args, **kwargs):
+            if not hasattr(cls, 'test_result_container'):
+                return tear_down(self, *args, *kwargs)
+
+            container = cls.get_testcase_container()
+
+            test_method = getattr(self, self._testMethodName)
+
+            exclusions = (
+                getattr(self.__class__, "__querycount_exclude__", []) +
+                getattr(test_method, "__querycount_exclude__", [])
+            )
+
+            all_queries = cls.test_result_container
+            current_queries = container.filter_by(exclusions)
+            all_queries.add(self.id(), current_queries)
+
+            return tear_down(self, *args, *kwargs)
+
+        return wrapped
+
+    @classmethod
+    def patch_test_case(cls):
+        TransactionTestCase.setUp = cls.wrap_set_up(TransactionTestCase.setUp)
+        TransactionTestCase.tearDown = cls.wrap_tear_down(
+            TransactionTestCase.tearDown)
+
+    @classmethod
+    def save_json(cls, setting_name, container, detail):
+        summary_path = os.path.realpath(cls.get_setting(setting_name))
+        os.makedirs(os.path.dirname(summary_path), exist_ok=True)
+
+        with open(summary_path, 'w') as json_file:
+            json.dump(container.get_json(detail=detail), json_file,
+                      ensure_ascii=False, indent=4, sort_keys=True)
+
+    @classmethod
+    def wrap_testrunner_run(cls, func):
+        def wrapped(self, *args, **kwargs):
+            cls.test_result_container = TestResultQueryContainer()
+
+            result = func(self, *args, **kwargs)
+
+            cls.save_json('SUMMARY_PATH', cls.test_result_container, False)
+            cls.save_json('DETAIL_PATH', cls.test_result_container, True)
+
+            result.queries = cls.test_result_container
+            del cls.test_result_container
+
+            return result
+
+        return wrapped
+
+    @classmethod
+    def patch_runner(cls):
+        # FIXME: this is incompatible with --parallel and --test-runner
+        # command arguments
+        django_test_runner = get_runner(settings)
+        test_runner = django_test_runner.test_runner
+        test_runner.run = cls.wrap_testrunner_run(test_runner.run)
+
+    def ready(self):
+        if self.enabled():
+            self.add_middleware()
+            self.patch_test_case()
+            self.patch_runner()

request_query_count/middleware.py

@@ -0,0 +1,22 @@
+from django.db import DEFAULT_DB_ALIAS, connections
+from django.test.utils import CaptureQueriesContext
+
+from request_query_count.apps import RequestQueryCountConfig
+
+
+class Middleware(object):
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        if not RequestQueryCountConfig.enabled():
+            # It is not necessary to capture queries
+            return self.get_response(request)
+
+        with CaptureQueriesContext(connections[DEFAULT_DB_ALIAS]) as context:
+            response = self.get_response(request)
+
+        query_container = RequestQueryCountConfig.get_testcase_container()
+        query_container.add(request, context.captured_queries)
+
+        return response