refactor: add contextvar to enable/disable tracking of DDT models

Author: theShinigami

debug_toolbar/panels/sql/tracking.py

@@ -8,6 +8,7 @@
 from django.utils.encoding import force_str
 
 from debug_toolbar.utils import get_stack_trace, get_template_info
+from django.apps import apps
 
 try:
     import psycopg
@@ -27,6 +28,17 @@
 # additional queries.
 allow_sql = contextvars.ContextVar("debug-toolbar-allow-sql", default=True)
 
+# Prevents tracking of DDT models
+allow_ddt_models_tracking = contextvars.ContextVar(
+    "debug-toolbar-allow-ddt-models", default=False
+)
+
+DDT_MODELS = {
+    m._meta.db_table
+    for m in apps.get_app_config("debug_toolbar").get_models()
+    if m._meta.app_label == "debug_toolbar"
+}
+
 
 class SQLQueryTriggered(Exception):
     """Thrown when template panel triggers a query"""
@@ -224,6 +236,12 @@ def _record(self, method, sql, params):
                     }
                 )
 
+            # Skip recording if query includes DDT models.
+            if not allow_ddt_models_tracking.get() and any(
+                table in sql for table in DDT_MODELS
+            ):
+                return
+
             # We keep `sql` to maintain backwards compatibility
             self.logger.record(**kwargs)
 

tests/panels/test_sql.py

@@ -16,6 +16,7 @@
 
 import debug_toolbar.panels.sql.tracking as sql_tracking
 from debug_toolbar.panels.sql import SQLPanel
+from debug_toolbar.models import HistoryEntry
 
 try:
     import psycopg
@@ -33,6 +34,20 @@ def sql_call(*, use_iterator=False):
     return list(qs)
 
 
+def sql_call_ddt(*, use_iterator=False):
+    qs = HistoryEntry.objects.all()
+    if use_iterator:
+        qs = qs.iterator()
+    return list(qs)
+
+
+async def async_sql_call_ddt(*, use_iterator=False):
+    qs = HistoryEntry.objects.all()
+    if use_iterator:
+        qs = qs.iterator()
+    return await sync_to_async(list)(qs)
+
+
 async def async_sql_call(*, use_iterator=False):
     qs = User.objects.all()
     if use_iterator:
@@ -104,6 +119,64 @@ async def test_recording_concurrent_async(self):
         # ensure the stacktrace is populated
         self.assertTrue(len(query["stacktrace"]) > 0)
 
+    def test_ddt_models_tracking(self):
+        self.assertEqual(len(self.panel._queries), 0)
+
+        # enable ddt tracking
+        sql_tracking.allow_ddt_models_tracking.set(True)
+
+        sql_call_ddt()
+
+        # ensure query was logged
+        self.assertEqual(len(self.panel._queries), 1)
+        query = self.panel._queries[0]
+        self.assertEqual(query["alias"], "default")
+        self.assertTrue("sql" in query)
+        self.assertTrue("duration" in query)
+        self.assertTrue("stacktrace" in query)
+
+        # ensure the stacktrace is populated
+        self.assertTrue(len(self.panel._queries), 0)
+
+    async def test_ddt_models_tracking_async(self):
+        self.assertEqual(len(self.panel._queries), 0)
+
+        # enable ddt tracking
+        sql_tracking.allow_ddt_models_tracking.set(True)
+
+        await async_sql_call_ddt()
+
+        # ensure query was logged
+        self.assertEqual(len(self.panel._queries), 1)
+        query = self.panel._queries[0]
+        self.assertEqual(query["alias"], "default")
+        self.assertTrue("sql" in query)
+        self.assertTrue("duration" in query)
+        self.assertTrue("stacktrace" in query)
+
+        # ensure the stacktrace is populated
+        self.assertTrue(len(self.panel._queries), 0)
+
+    def test_ddt_models_untracking(self):
+        self.assertEqual(len(self.panel._queries), 0)
+
+        # disable ddt tracking
+        sql_tracking.allow_ddt_models_tracking.set(False)
+
+        sql_call_ddt()
+
+        self.assertEqual(len(self.panel._queries), 0)
+
+    async def test_ddt_models_untracking_async(self):
+        self.assertEqual(len(self.panel._queries), 0)
+
+        # enable ddt tracking
+        sql_tracking.allow_ddt_models_tracking.set(False)
+
+        await async_sql_call_ddt()
+
+        self.assertEqual(len(self.panel._queries), 0)
+
     @unittest.skipUnless(
         connection.vendor == "postgresql", "Test valid only on PostgreSQL"
     )