Show context where SQL query originated from template

Signed-off-by: Chris Lamb <lamby@debian.org>
Signed-off-by: Rob Hudson <rob@cogit8.org>

1 files changed, 54 insertions, 0 deletions

diff --git a/debug_toolbar/panels/sql.py b/debug_toolbar/panels/sql.py
index e1a2b7d..ce43328 100644
--- a/debug_toolbar/panels/sql.py
+++ b/debug_toolbar/panels/sql.py
@@ -1,5 +1,6 @@
 from datetime import datetime
 import os
+import sys
 import SocketServer
 import traceback
 
@@ -7,6 +8,8 @@ import django
 from django.conf import settings
 from django.db import connection
 from django.db.backends import util
+from django.views.debug import linebreak_iter
+from django.template import Node
 from django.template.loader import render_to_string
 from django.utils import simplejson
 from django.utils.encoding import force_unicode
@@ -42,6 +45,40 @@ def tidy_stacktrace(strace):
         trace.append((s[0], s[1], s[2], s[3]))
     return trace
 
+def get_template_info(source, context_lines=3):
+    line = 0
+    upto = 0
+    source_lines = []
+    before = during = after = ""
+
+    origin, (start, end) = source
+    template_source = origin.reload()
+
+    for num, next in enumerate(linebreak_iter(template_source)):
+        if start >= upto and end <= next:
+            line = num
+            before = template_source[upto:start]
+            during = template_source[start:end]
+            after = template_source[end:next]
+        source_lines.append((num, template_source[upto:next]))
+        upto = next
+
+    top = max(1, line - context_lines)
+    bottom = min(len(source_lines), line + 1 + context_lines)
+
+    context = []
+    for num, content in source_lines[top:bottom]:
+        context.append({
+            'num': num,
+            'content': content,
+            'highlight': (num == line),
+        })
+
+    return {
+        'name': origin.name,
+        'context': context,
+    }
+
 class DatabaseStatTracker(util.CursorDebugWrapper):
     """
     Replacement for CursorDebugWrapper which stores additional information
@@ -60,6 +97,22 @@ class DatabaseStatTracker(util.CursorDebugWrapper):
                 _params = simplejson.dumps([force_unicode(x) for x in params])
             except TypeError:
                 pass # object not JSON serializable
+
+            template_info = None
+            cur_frame = sys._getframe().f_back
+            try:
+                while cur_frame is not None:
+                    if cur_frame.f_code.co_name == 'render':
+                        node = cur_frame.f_locals['self']
+                        if isinstance(node, Node):
+                            template_info = get_template_info(node.source)
+                            break
+                    cur_frame = cur_frame.f_back
+            except:
+                pass
+            finally:
+                del cur_frame
+
             # We keep `sql` to maintain backwards compatibility
             self.db.queries.append({
                 'sql': self.db.ops.last_executed_query(self.cursor, sql, params),
@@ -72,6 +125,7 @@ class DatabaseStatTracker(util.CursorDebugWrapper):
                 'stop_time': stop,
                 'is_slow': (duration > SQL_WARNING_THRESHOLD),
                 'is_select': sql.lower().strip().startswith('select'),
+                'template_info': template_info,
             })
 util.CursorDebugWrapper = DatabaseStatTracker