#214 Use Django 2.0 execute_wrapper()

Author: SebCorbin

project/tests/test_db.py

@@ -2,6 +2,7 @@
 Test profiling of DB queries without mocking, to catch possible
 incompatibility
 """
+
 from django.shortcuts import reverse
 from django.test import Client, TestCase
 
@@ -20,24 +21,28 @@ def setUpClass(cls):
         BlindFactory.create_batch(size=5)
         SilkyConfig().SILKY_META = False
 
+    def setUp(self):
+        DataCollector().clear()
+
     def test_profile_request_to_db(self):
         DataCollector().configure(Request(reverse('example_app:index')))
 
         with silk_profile(name='test_profile'):
             resp = self.client.get(reverse('example_app:index'))
 
-        DataCollector().profiles.values()
-        assert len(resp.context['blinds']) == 5
+        self.assertEqual(len(DataCollector().queries), 1, [q['query'] for q in DataCollector().queries.values()])
+        self.assertEqual(len(resp.context['blinds']), 5)
 
     def test_profile_request_to_db_with_constraints(self):
         DataCollector().configure(Request(reverse('example_app:create')))
 
         resp = self.client.post(reverse('example_app:create'), {'name': 'Foo'})
+        self.assertEqual(len(DataCollector().queries), 2)
+        self.assertTrue(list(DataCollector().queries.values())[-1]['query'].startswith('INSERT'))
         self.assertEqual(resp.status_code, 302)
 
 
 class TestAnalyzeQueries(TestCase):
-
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
@@ -48,7 +53,7 @@ def setUpClass(cls):
     @classmethod
     def tearDownClass(cls):
         super().tearDownClass()
-        SilkyConfig().SILKLY_ANALYZE_QUERIES = False
+        SilkyConfig().SILKY_ANALYZE_QUERIES = False
 
     def test_analyze_queries(self):
         DataCollector().configure(Request(reverse('example_app:index')))
@@ -59,16 +64,3 @@ def test_analyze_queries(self):
 
         DataCollector().profiles.values()
         assert len(resp.context['blinds']) == 5
-
-
-class TestAnalyzeQueriesExplainParams(TestAnalyzeQueries):
-
-    @classmethod
-    def setUpClass(cls):
-        super().setUpClass()
-        SilkyConfig().SILKY_EXPLAIN_FLAGS = {'verbose': True}
-
-    @classmethod
-    def tearDownClass(cls):
-        super().tearDownClass()
-        SilkyConfig().SILKY_EXPLAIN_FLAGS = None

project/tests/test_execute_sql.py

@@ -1,6 +1,14 @@
 from django.apps import AppConfig
+from django.db import connection
+
+from silk.sql import SilkQueryWrapper
 
 
 class SilkAppConfig(AppConfig):
     default_auto_field = "django.db.models.AutoField"
     name = "silk"
+
+    def ready(self):
+        # Add wrapper to db connection
+        if not any(isinstance(wrapper, SilkQueryWrapper) for wrapper in connection.execute_wrappers):
+            connection.execute_wrappers.append(SilkQueryWrapper())

silk/apps.py

@@ -2,7 +2,6 @@
 import random
 
 from django.db import DatabaseError, transaction
-from django.db.models.sql.compiler import SQLCompiler
 from django.urls import NoReverseMatch, reverse
 from django.utils import timezone
 
@@ -11,7 +10,6 @@
 from silk.model_factory import RequestModelFactory, ResponseModelFactory
 from silk.profiling import dynamic
 from silk.profiling.profiler import silk_meta_profiler
-from silk.sql import execute_sql
 
 Logger = logging.getLogger('silk.middleware')
 
@@ -85,15 +83,11 @@ def _apply_dynamic_mappings(self):
             name = conf.get('name')
             if module and function:
                 if start_line and end_line:  # Dynamic context manager
-                    dynamic.inject_context_manager_func(module=module,
-                                                        func=function,
-                                                        start_line=start_line,
-                                                        end_line=end_line,
-                                                        name=name)
+                    dynamic.inject_context_manager_func(
+                        module=module, func=function, start_line=start_line, end_line=end_line, name=name
+                    )
                 else:  # Dynamic decorator
-                    dynamic.profile_function_or_method(module=module,
-                                                       func=function,
-                                                       name=name)
+                    dynamic.profile_function_or_method(module=module, func=function, name=name)
             else:
                 raise KeyError('Invalid dynamic mapping %s' % conf)
 
@@ -107,9 +101,6 @@ def process_request(self, request):
         Logger.debug('process_request')
         request.silk_is_intercepted = True
         self._apply_dynamic_mappings()
-        if not hasattr(SQLCompiler, '_execute_sql'):
-            SQLCompiler._execute_sql = SQLCompiler.execute_sql
-            SQLCompiler.execute_sql = execute_sql
 
         silky_config = SilkyConfig()
 

silk/middleware.py

@@ -1,26 +1,16 @@
 import logging
 import traceback
 
-from django.core.exceptions import EmptyResultSet
+from django.apps import apps
+from django.db import connection
 from django.utils import timezone
 from django.utils.encoding import force_str
 
-from silk.collector import DataCollector
 from silk.config import SilkyConfig
 
 Logger = logging.getLogger('silk.sql')
 
 
-def _should_wrap(sql_query):
-    if not DataCollector().request:
-        return False
-
-    for ignore_str in SilkyConfig().SILKY_IGNORE_QUERIES:
-        if ignore_str in sql_query:
-            return False
-    return True
-
-
 def _unpack_explanation(result):
     for row in result:
         if not isinstance(row, str):
@@ -34,16 +24,14 @@ def _explain_query(connection, q, params):
         if SilkyConfig().SILKY_ANALYZE_QUERIES:
             # Work around some DB engines not supporting analyze option
             try:
-                prefix = connection.ops.explain_query_prefix(
-                    analyze=True, **(SilkyConfig().SILKY_EXPLAIN_FLAGS or {})
-                )
+                prefix = connection.ops.explain_query_prefix(analyze=True, **(SilkyConfig().SILKY_EXPLAIN_FLAGS or {}))
             except ValueError as error:
                 error_str = str(error)
                 if error_str.startswith("Unknown options:"):
                     Logger.warning(
-                        "Database does not support analyzing queries with provided params. %s."
+                        "Database does not support analyzing queries with provided params. %s. "
                         "SILKY_ANALYZE_QUERIES option will be ignored",
-                        error_str
+                        error_str,
                     )
                     prefix = connection.ops.explain_query_prefix()
                 else:
@@ -61,40 +49,53 @@ def _explain_query(connection, q, params):
     return None
 
 
-def execute_sql(self, *args, **kwargs):
-    """wrapper around real execute_sql in order to extract information"""
+class SilkQueryWrapper:
+    def __init__(self):
+        # Local import to prevent messing app.ready()
+        from silk.collector import DataCollector
 
-    try:
-        q, params = self.as_sql()
-        if not q:
-            raise EmptyResultSet
-    except EmptyResultSet:
-        try:
-            result_type = args[0]
-        except IndexError:
-            result_type = kwargs.get('result_type', 'multi')
-        if result_type == 'multi':
-            return iter([])
-        else:
-            return
-    tb = ''.join(reversed(traceback.format_stack()))
-    sql_query = q % tuple(force_str(param) for param in params)
-    if _should_wrap(sql_query):
-        query_dict = {
-            'query': sql_query,
-            'start_time': timezone.now(),
-            'traceback': tb
-        }
+        self.data_collector = DataCollector()
+        self.silk_model_table_names = [model._meta.db_table for model in apps.get_app_config('silk').get_models()]
+
+    def __call__(self, execute, sql, params, many, context):
+        sql_query = sql % tuple(force_str(param) for param in params) if params else sql
+        query_dict = None
+        if self._should_wrap(sql_query):
+            tb = ''.join(reversed(traceback.format_stack()))
+            query_dict = {'query': sql_query, 'start_time': timezone.now(), 'traceback': tb}
         try:
-            return self._execute_sql(*args, **kwargs)
+            return execute(sql, params, many, context)
         finally:
-            query_dict['end_time'] = timezone.now()
-            request = DataCollector().request
-            if request:
-                query_dict['request'] = request
-            if getattr(self.query.model, '__module__', '') != 'silk.models':
-                query_dict['analysis'] = _explain_query(self.connection, q, params)
-                DataCollector().register_query(query_dict)
-            else:
-                DataCollector().register_silk_query(query_dict)
-    return self._execute_sql(*args, **kwargs)
+            if query_dict:
+                query_dict['end_time'] = timezone.now()
+                request = self.data_collector.request
+                if request:
+                    query_dict['request'] = request
+                if not any(table_name in sql_query for table_name in self.silk_model_table_names):
+                    query_dict['analysis'] = _explain_query(connection, sql, params)
+                    self.data_collector.register_query(query_dict)
+                else:
+                    self.data_collector.register_silk_query(query_dict)
+
+    def _should_wrap(self, sql_query):
+        # Must have a request ongoing
+        if not self.data_collector.request:
+            return False
+
+        # Must not try to explain 'EXPLAIN' queries or transaction stuff
+        if any(
+            sql_query.startswith(keyword)
+            for keyword in [
+                'SAVEPOINT',
+                'RELEASE SAVEPOINT',
+                'ROLLBACK TO SAVEPOINT',
+                'PRAGMA',
+                connection.ops.explain_query_prefix(),
+            ]
+        ):
+            return False
+
+        for ignore_str in SilkyConfig().SILKY_IGNORE_QUERIES:
+            if ignore_str in sql_query:
+                return False
+        return True