From 01b5d65a1da80f3cc6120e53bba46e4da2ff44ca Mon Sep 17 00:00:00 2001
From: Tomasz Buchert
Date: Mon, 26 Sep 2011 00:03:47 +0200
Subject: Issue 187: queries are captured and not executed when template panel
 is rendered.

---
 debug_toolbar/panels/template.py   | 15 ++++++++++++++-
 debug_toolbar/utils/tracking/db.py | 37 +++++++++++++++++++++++++++++++++++--
 tests/tests.py                     | 16 ++++++++++++++++
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/debug_toolbar/panels/template.py b/debug_toolbar/panels/template.py
index d620a38..ff8cd15 100644
--- a/debug_toolbar/panels/template.py
+++ b/debug_toolbar/panels/template.py
@@ -7,7 +7,9 @@ from django.template.context import get_standard_processors
 from django.template.loader import render_to_string
 from django.test.signals import template_rendered
 from django.utils.translation import ugettext_lazy as _
+from django.db.models.query import QuerySet
 from debug_toolbar.panels import DebugPanel
+from debug_toolbar.utils.tracking.db import recording, SQLQueryTriggered
 
 # Code taken and adapted from Simon Willison and Django Snippets:
 # http://www.djangosnippets.org/snippets/766/
@@ -66,8 +68,19 @@ class TemplateDebugPanel(DebugPanel):
                     # Replace LANGUAGES, which is available in i18n context processor
                     elif key == 'LANGUAGES' and isinstance(value, tuple):
                         temp_layer[key] = '<<languages>>'
+                    # QuerySet would trigger the database: user can run the query from SQL Panel
+                    elif isinstance(value, QuerySet):
+                        temp_layer[key] = '<<queryset>>'
                     else:
-                        temp_layer[key] = value
+                        try:
+                            recording(False)
+                            pformat(value)  # this MAY trigger a db query
+                        except SQLQueryTriggered:
+                            temp_layer[key] = '<<contains queryset>>'
+                        else:
+                            temp_layer[key] = value
+                        finally:
+                            recording(True)
             try:
                 context_list.append(pformat(temp_layer))
             except UnicodeEncodeError:
diff --git a/debug_toolbar/utils/tracking/db.py b/debug_toolbar/utils/tracking/db.py
index 7ffacef..1fa7b36 100644
--- a/debug_toolbar/utils/tracking/db.py
+++ b/debug_toolbar/utils/tracking/db.py
@@ -2,6 +2,7 @@ import inspect
 import sys
 
 from datetime import datetime
+from threading import local
 
 from django.conf import settings
 from django.template import Node
@@ -16,7 +17,39 @@ from debug_toolbar.utils.compat.db import connections
 SQL_WARNING_THRESHOLD = getattr(settings, 'DEBUG_TOOLBAR_CONFIG', {}) \
                             .get('SQL_WARNING_THRESHOLD', 500)
 
-class CursorWrapper(object):
+class SQLQueryTriggered(Exception):
+    """Thrown when template panel triggers a query"""
+    pass
+
+class ThreadLocalState(local):
+    def __init__(self):
+        self.enabled = True
+
+    @property
+    def Wrapper(self):
+        return NormalCursorWrapper if self.enabled else ExceptionCursorWrapper
+
+    def recording(self, v):
+        self.enabled = v
+
+state = ThreadLocalState()
+recording = state.recording # export function
+
+def CursorWrapper(*args, **kwds):  # behave like a class
+    return state.Wrapper(*args, **kwds)
+
+class ExceptionCursorWrapper(object):
+    """
+    Wraps a cursor and raises an exception on any operation.
+    Used in Templates panel.
+    """
+    def __init__(self, cursor, db, logger):
+        pass
+
+    def __getattr__(self, attr):
+        raise SQLQueryTriggered()
+
+class NormalCursorWrapper(object):
     """
     Wraps a cursor and logs queries.
     """
@@ -103,4 +136,4 @@ class CursorWrapper(object):
             return getattr(self.cursor, attr)
 
     def __iter__(self):
-        return iter(self.cursor)
\ No newline at end of file
+        return iter(self.cursor)
diff --git a/tests/tests.py b/tests/tests.py
index 60620bb..5a52670 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -1,6 +1,7 @@
 from debug_toolbar.middleware import DebugToolbarMiddleware
 from debug_toolbar.panels.sql import SQLDebugPanel
 from debug_toolbar.panels.request_vars import RequestVarsDebugPanel
+from debug_toolbar.panels.template import TemplateDebugPanel
 from debug_toolbar.toolbar.loader import DebugToolbar
 from debug_toolbar.utils import get_name_from_obj
 from debug_toolbar.utils.tracking import pre_dispatch, post_dispatch, callbacks
@@ -9,6 +10,7 @@ from django.conf import settings
 from django.contrib.auth.models import User
 from django.http import HttpResponse
 from django.test import TestCase
+from django.template import Template, Context
 
 from dingus import Dingus
 import thread
@@ -194,6 +196,20 @@ class SQLPanelTestCase(BaseTestCase):
         self.assertTrue('duration' in query[1])
         self.assertTrue('stacktrace' in query[1])
 
+class TemplatePanelTestCase(BaseTestCase):
+    def test_queryset_hook(self):
+        template_panel = self.toolbar.get_panel(TemplateDebugPanel)
+        sql_panel = self.toolbar.get_panel(SQLDebugPanel)
+        t = Template("No context variables here!")
+        c = Context({ 'queryset' : User.objects.all(), 'deep_queryset' : { 'queryset' : User.objects.all() } })
+        t.render(c)
+        # ensure the query was NOT logged
+        self.assertEquals(len(sql_panel._queries), 0)
+        ctx = template_panel.templates[0]['context'][0]
+        ctx = eval(ctx) # convert back to Python
+        self.assertEquals(ctx['queryset'], '<<queryset>>')
+        self.assertEquals(ctx['deep_queryset'], '<<contains queryset>>')
+
 def module_func(*args, **kwargs):
     """Used by dispatch tests"""
     return 'blah'
-- 
cgit v1.2.3