Refactor tracker, add stack traces to queries

debug_toolbar_mongo/operation_tracker.py

 import functools
+import traceback
 import time
+import inspect
+import os
+import SocketServer
 
+import django
+from django.conf import settings
+
+import pymongo
 import pymongo.collection
 import pymongo.cursor
 
+__all__ = ['queries', 'inserts', 'updates', 'removes', 'install_tracker',
+           'uninstall_tracker', 'reset']
 
-class MongoOperationTracker(object):
-    """Track operations sent to MongoDB via PyMongo.
-    """
-
-    def __init__(self):
-        self.queries = []
-        self.inserts = []
-        self.updates = []
-        self.removes = []
-        self.active = False
 
-    def install(self):
-        self.original_methods = {
+_original_methods = {
     'insert': pymongo.collection.Collection.insert,
     'update': pymongo.collection.Collection.update,
     'remove': pymongo.collection.Collection.remove,
     'refresh': pymongo.cursor.Cursor._refresh,
-        }
+}
 
-        # Wrap Cursor._refresh for getting queries
-        @functools.wraps(self.original_methods['insert'])
-        def insert(collection_self, doc_or_docs, manipulate=True,
+queries = []
+inserts = []
+updates = []
+removes = []
+
+
+# Wrap Cursor._refresh for getting queries
+@functools.wraps(_original_methods['insert'])
+def _insert(collection_self, doc_or_docs, manipulate=True,
            safe=False, check_keys=True, **kwargs):
     start_time = time.time()
-            result = self.original_methods['insert'](
+    result = _original_methods['insert'](
         collection_self,
         doc_or_docs,
         safe=safe,
@@ -37,22 +42,19 @@ class MongoOperationTracker(object):
     )
     total_time = (time.time() - start_time) * 1000
 
-            if self.active:
-                self.inserts.append({
+    inserts.append({
         'document': doc_or_docs,
         'safe': safe,
         'time': total_time,
     })
     return result
 
-        pymongo.collection.Collection.insert = insert
-
-        # Wrap Cursor._refresh for getting queries
-        @functools.wraps(self.original_methods['update'])
-        def update(collection_self, spec, document, upsert=False,
+# Wrap Cursor._refresh for getting queries
+@functools.wraps(_original_methods['update'])
+def _update(collection_self, spec, document, upsert=False,
            maniuplate=False, safe=False, multi=False, **kwargs):
     start_time = time.time()
-            result = self.original_methods['update'](
+    result = _original_methods['update'](
         collection_self,
         spec,
         document,
@@ -63,8 +65,7 @@ class MongoOperationTracker(object):
     )
     total_time = (time.time() - start_time) * 1000
 
-            if self.active:
-                self.updates.append({
+    updates.append({
         'document': document,
         'upsert': upsert,
         'multi': multi,
@@ -74,13 +75,11 @@ class MongoOperationTracker(object):
     })
     return result
 
-        pymongo.collection.Collection.update = update
-
-        # Wrap Cursor._refresh for getting queries
-        @functools.wraps(self.original_methods['remove'])
-        def remove(collection_self, spec_or_id, safe=False, **kwargs):
+# Wrap Cursor._refresh for getting queries
+@functools.wraps(_original_methods['remove'])
+def _remove(collection_self, spec_or_id, safe=False, **kwargs):
     start_time = time.time()
-            result = self.original_methods['insert'](
+    result = _original_methods['insert'](
         collection_self,
         spec_or_id,
         safe=safe,
@@ -88,33 +87,30 @@ class MongoOperationTracker(object):
     )
     total_time = (time.time() - start_time) * 1000
 
-            if self.active:
-                self.removes.append({
+    removes.append({
         'spec_or_id': spec_or_id,
         'safe': safe,
         'time': total_time,
     })
     return result
 
-        pymongo.collection.Collection.remove = remove
-
-        # Wrap Cursor._refresh for getting queries
-        @functools.wraps(self.original_methods['refresh'])
-        def cursor_refresh(cursor_self):
+# Wrap Cursor._refresh for getting queries
+@functools.wraps(_original_methods['refresh'])
+def _cursor_refresh(cursor_self):
     # Look up __ private instance variables
     def privar(name):
         return getattr(cursor_self, '_Cursor__{0}'.format(name))
 
-            if not self.active or privar('id') is not None:
-                # Inactive or getMore - move on
-                return self.original_methods['refresh'](cursor_self)
+    if privar('id') is not None:
+        # getMore not query - move on
+        return _original_methods['refresh'](cursor_self)
 
     # NOTE: See pymongo/cursor.py+557 [_refresh()] and
     # pymongo/message.py for where information is stored
 
     # Time the actual query
     start_time = time.time()
-            result = self.original_methods['refresh'](cursor_self)
+    result = _original_methods['refresh'](cursor_self)
     total_time = (time.time() - start_time) * 1000
 
     query_son = privar('query_spec')()
@@ -122,6 +118,7 @@ class MongoOperationTracker(object):
     query_data = {
         'time': total_time,
         'operation': 'query',
+        'stack_trace': _tidy_stacktrace(reversed(inspect.stack())),
     }
 
     # Collection in format <db_name>.<collection_name>
@@ -143,31 +140,81 @@ class MongoOperationTracker(object):
         query_data['skip'] = privar('skip')
         query_data['limit'] = privar('limit')
         query_data['query'] = query_son['$query']
-                query_data['ordering'] = self._get_ordering(query_son)
+        query_data['ordering'] = _get_ordering(query_son)
 
-            self.queries.append(query_data)
+    queries.append(query_data)
 
     return result
 
-        pymongo.cursor.Cursor._refresh = cursor_refresh
-
-    def _get_ordering(self, son):
+def install_tracker():
+    if pymongo.collection.Collection.insert != _insert:
+        pymongo.collection.Collection.insert = _insert
+    if pymongo.collection.Collection.update != _update:
+        pymongo.collection.Collection.update = _update
+    if pymongo.collection.Collection.remove != _remove:
+        pymongo.collection.Collection.remove = _remove
+    if pymongo.cursor.Cursor._refresh != _cursor_refresh:
+        pymongo.cursor.Cursor._refresh = _cursor_refresh
+
+def uninstall_tracker():
+    if pymongo.collection.Collection.insert == _insert:
+        pymongo.collection.Collection.insert = _original_methods['insert']
+    if pymongo.collection.Collection.update == _update:
+        pymongo.collection.Collection.update = _original_methods['update']
+    if pymongo.collection.Collection.remove == _remove:
+        pymongo.collection.Collection.remove = _original_methods['remove']
+    if pymongo.cursor.Cursor._refresh == _cursor_refresh:
+        pymongo.cursor.Cursor._refresh = _original_methods['cursor_refresh']
+
+def reset():
+    global queries, inserts, updates, removes
+    queries = []
+    inserts = []
+    updates = []
+    removes = []
+
+def _get_ordering(son):
+    """Helper function to extract formatted ordering from dict.
+    """
     def fmt(field, direction):
         return '{0}{1}'.format({-1: '-', 1: '+'}[direction], field)
 
     if '$orderby' in son:
         return ', '.join(fmt(f, d) for f, d in son['$orderby'].items())
 
-    def uninstall(self):
-        pymongo.cursor.Cursor._refresh = self.original_methods['refresh']
-        pymongo.cursor.Collection.insert = self.original_methods['insert']
-        pymongo.cursor.Collection.update = self.original_methods['update']
-
-    def reset(self):
-        self.queries = []
+# Taken from Django Debug Toolbar 0.8.6
+def _tidy_stacktrace(stack):
+    """
+    Clean up stacktrace and remove all entries that:
+    1. Are part of Django (except contrib apps)
+    2. Are part of SocketServer (used by Django's dev server)
+    3. Are the last entry (which is part of our stacktracing code)
 
-    def start(self):
-        self.active = True
+    ``stack`` should be a list of frame tuples from ``inspect.stack()``
+    """
+    django_path = os.path.realpath(os.path.dirname(django.__file__))
+    django_path = os.path.normpath(os.path.join(django_path, '..'))
+    socketserver_path = os.path.realpath(os.path.dirname(SocketServer.__file__))
+    pymongo_path = os.path.realpath(os.path.dirname(pymongo.__file__))
+
+    trace = []
+    for frame, path, line_no, func_name, text in (f[:5] for f in stack):
+        s_path = os.path.realpath(path)
+        # Support hiding of frames -- used in various utilities that provide
+        # inspection.
+        if '__traceback_hide__' in frame.f_locals:
+            continue
+        if getattr(settings, 'DEBUG_TOOLBAR_CONFIG', {}).get('HIDE_DJANGO_SQL', True) \
+            and django_path in s_path and not 'django/contrib' in s_path:
+            continue
+        if socketserver_path in s_path:
+            continue
+        if pymongo_path in s_path:
+            continue
+        if not text:
+            text = ''
+        else:
+            text = (''.join(text)).strip()
+        trace.append((path, line_no, func_name, text))
+    return trace
 
-    def stop(self):
-        self.active = False

debug_toolbar_mongo/panel.py

@@ -2,7 +2,7 @@ from django.template.loader import render_to_string
 
 from debug_toolbar.panels import DebugPanel
 
-from operation_tracker import MongoOperationTracker
+import operation_tracker
 
 
 class MongoDebugPanel(DebugPanel):
@@ -13,23 +13,18 @@ class MongoDebugPanel(DebugPanel):
 
     def __init__(self, *args, **kwargs):
         super(self.__class__, self).__init__(*args, **kwargs)
-        self.op_tracker = MongoOperationTracker()
-        self.op_tracker.install()
+        operation_tracker.install_tracker()
 
     def process_request(self, request):
-        self.op_tracker.reset()
-        self.op_tracker.start()
-
-    def process_response(self, request, response):
-        self.op_tracker.stop()
+        operation_tracker.reset()
 
     def nav_title(self):
         return 'MongoDB'
 
     def nav_subtitle(self):
-        num_queries = len(self.op_tracker.queries)
+        num_queries = len(operation_tracker.queries)
         attrs = ['queries', 'inserts', 'updates', 'removes']
-        total_time = sum(sum(o['time'] for o in getattr(self.op_tracker, a))
+        total_time = sum(sum(o['time'] for o in getattr(operation_tracker, a))
                          for a in attrs)
         return '{0} operations in {1:.2f}ms'.format(num_queries, total_time)
 
@@ -41,10 +36,10 @@ class MongoDebugPanel(DebugPanel):
 
     def content(self):
         context = self.context.copy()
-        context['queries'] = self.op_tracker.queries
-        context['inserts'] = self.op_tracker.inserts
-        context['updates'] = self.op_tracker.updates
-        context['removes'] = self.op_tracker.removes
+        context['queries'] = operation_tracker.queries
+        context['inserts'] = operation_tracker.inserts
+        context['updates'] = operation_tracker.updates
+        context['removes'] = operation_tracker.removes
         return render_to_string('mongo-panel.html', context)
 
 

debug_toolbar_mongo/templates/mongo-panel.html

@@ -5,11 +5,23 @@ pre.mongo-highlight, pre.mongo-highlight span {
     font-family: 'Consolas', 'Deja Vu Sans Mono', 'Bitstream Vera Sans Mono', 'Monaco', 'Courier New', monospace !important;
     font-size: 1.05em !important;
 }
+pre.mongo-stack-trace {
+    line-height: 2.3em !important;
+}
+table.mongo-stack-trace th {
+    text-align: left !important;
+}
+table.mongo-stack-trace td.code, table.mongo-stack-trace td.code span, table.mongo-stack-trace td.function {
+    font-family: 'Consolas', 'Deja Vu Sans Mono', 'Bitstream Vera Sans Mono', 'Monaco', 'Courier New', monospace !important;
+}
+table.mongo-op-table tbody {
+    width: 100% !important;
+}
 </style>
 
 <h4>Queries</h4>
 {% if queries %}
-<table>
+<table class="mongo-op-table">
     <thead>
         <tr>
             <th>Time (ms)</th>
@@ -19,6 +31,7 @@ pre.mongo-highlight, pre.mongo-highlight span {
             <th>Ordering</th>
             <th>Skip</th>
             <th>Limit</th>
+            <th>Stack Trace</th>
         </tr>
     </thead>
     <tbody>
@@ -29,12 +42,38 @@ pre.mongo-highlight, pre.mongo-highlight span {
                 <td>{{ query.collection }}</td>
                 <td>
                     {% if query.query %}
-                    <pre class="mongo-highlight">{{ query.query|format_dict|highlight_json|safe }}</pre>
+                    <pre class="mongo-highlight">{{ query.query|format_dict|highlight:"javascript"|safe }}</pre>
                   {% endif %}
                 </td>
                 <td><pre class="mongo-highlight">{% if query.ordering %}{{ query.ordering }}{% endif %}</pre></td>
                 <td>{% if query.skip %}{{ query.skip }}{% endif %}</td>
                 <td>{% if query.limit %}{{ query.limit }}{% endif %}</td>
+                <td><a href="javascript:void(0);" class="mongo-toggle-trace" data-row="{{ forloop.counter }}">Toggle</a></td>
+            </tr>
+            <tr class="{% cycle 'djDebugOdd' 'djDebugEven' %}">
+                <td colspan="8">
+                    <table class="mongo-stack-trace "id="mongo-stack-trace-{{ forloop.counter }}">
+                        <thead>
+                            <tr>
+                                <th>Line</th>
+                                <th>File</th>
+                                <th>Function</th>
+                                <th>Code</th>
+                            </tr>
+                        </thead>
+                        <tbody>
+                        {% for line in query.stack_trace %}
+                            <tr>
+                                <td class="lineno">{{ line.1 }}</td>
+                                <td class="file">{{ line.0 }}</td>
+                                <td class="function">{{ line.2 }}</td>
+                                <td class="code">{{ line.3|highlight:"python"|safe }}</td>
+                            </tr>
+                        {% endfor %}
+                        </tbody>
+                    </table>
+                    <!--<pre class="mongo-stack-trace">{{ query.stack_trace|format_stack_trace }}</pre>-->
+                </td>
             </tr>
         {% endfor %}
     </tbody>
@@ -45,7 +84,7 @@ pre.mongo-highlight, pre.mongo-highlight span {
 
 <h4>Inserts</h4>
 {% if inserts %}
-<table>
+<table class="mongo-op-table">
     <thead>
         <tr>
             <th>Time (ms)</th>
@@ -58,7 +97,7 @@ pre.mongo-highlight, pre.mongo-highlight span {
             <tr class="{% cycle 'djDebugOdd' 'djDebugEven' %}">
                 <td>{{ insert.time|floatformat:"4" }}</td>
                 <td>
-                    <pre class="mongo-highlight">{{ insert.document|format_dict:120|highlight_json|safe }}</pre>
+                    <pre class="mongo-highlight">{{ insert.document|format_dict:120|highlight:"javascript"|safe }}</pre>
                 </td>
                 <td>{{ insert.safe }}</td>
             </tr>
@@ -71,7 +110,7 @@ pre.mongo-highlight, pre.mongo-highlight span {
 
 <h4>Removes</h4>
 {% if removes %}
-<table>
+<table class="mongo-op-table">
     <thead>
         <tr>
             <th>Time (ms)</th>
@@ -84,7 +123,7 @@ pre.mongo-highlight, pre.mongo-highlight span {
             <tr class="{% cycle 'djDebugOdd' 'djDebugEven' %}">
                 <td>{{ remove.time|floatformat:"4" }}</td>
                 <td>
-                    <pre class="mongo-highlight">{{ remove.spec_or_id|format_dict:120|highlight_json|safe }}</pre>
+                    <pre class="mongo-highlight">{{ remove.spec_or_id|format_dict:120|highlight:"javascript"|safe }}</pre>
                 </td>
                 <td>{{ remove.safe }}</td>
             </tr>
@@ -97,7 +136,7 @@ pre.mongo-highlight, pre.mongo-highlight span {
 
 <h4>Updates</h4>
 {% if updates %}
-<table>
+<table class="mongo-op-table">
     <thead>
         <tr>
             <th>Time (ms)</th>
@@ -113,10 +152,10 @@ pre.mongo-highlight, pre.mongo-highlight span {
             <tr class="{% cycle 'djDebugOdd' 'djDebugEven' %}">
                 <td>{{ update.time|floatformat:"4" }}</td>
                 <td>
-                    <pre class="mongo-highlight">{{ update.spec|format_dict:120|highlight_json|safe }}</pre>
+                    <pre class="mongo-highlight">{{ update.spec|format_dict:120|highlight:"javascript"|safe }}</pre>
                 </td>
                 <td>
-                    <pre class="mongo-highlight">{{ update.document|format_dict:120|highlight_json|safe }}</pre>
+                    <pre class="mongo-highlight">{{ update.document|format_dict:120|highlight:"javascript"|safe }}</pre>
                 </td>
                 <td>{{ update.safe }}</td>
                 <td>{{ update.multi }}</td>
@@ -129,3 +168,24 @@ pre.mongo-highlight, pre.mongo-highlight span {
 <p>No updates recorded</p>
 {% endif %}
 
+<script>
+(function() {
+    var toggleButtons = document.querySelectorAll('a.mongo-toggle-trace');
+    for (var i = 0; i < toggleButtons.length; i++) {
+        (function() {
+            var button = toggleButtons[i];
+            var index = button.dataset.row;
+            var trace = document.getElementById('mongo-stack-trace-' + index);
+            trace.style.display = "none";
+
+            button.addEventListener("click", function() {
+                if (trace.style.display == "none") {
+                    trace.style.display = "";
+                } else {
+                    trace.style.display = "none";
+                }
+            }, true);
+        })();
+    }
+})();
+</script>

debug_toolbar_mongo/templatetags/mongo_debug_tags.py

 from django import template
+from django.utils.html import escape
+from django.utils.safestring import mark_safe
 
 import pprint
 
 register = template.Library()
 
 @register.filter
+def format_stack_trace(value):
+    stack_trace = []
+    fmt = (
+        '<span class="path">{0}/</span>'
+        '<span class="file">{1}</span> in <span class="func">{3}</span>'
+        '(<span class="lineno">{2}</span>) <span class="code">{4}</span>'
+    )
+    for frame in value:
+        params = map(escape, frame[0].rsplit('/', 1) + list(frame[1:]))
+        stack_trace.append(fmt.format(*params))
+    return mark_safe('\n'.join(stack_trace))
+
+@register.filter
 def format_dict(value, width=60):
     return pprint.pformat(value, width=int(width))
 
 @register.filter
-def highlight_json(value):
+def highlight(value, language):
     try:
         from pygments import highlight
-        from pygments.lexers import JavascriptLexer
+        from pygments.lexers import get_lexer_by_name
         from pygments.formatters import HtmlFormatter
     except ImportError:
         return value
     # Can't use class-based colouring because the debug toolbar's css rules
     # are more specific so take precedence
     formatter = HtmlFormatter(style='friendly', nowrap=True, noclasses=True)
-    return highlight(value, JavascriptLexer(), formatter)
+    return highlight(value, get_lexer_by_name(language), formatter)

example/settings.py

@@ -37,10 +37,10 @@ DEBUG_TOOLBAR_PANELS = (
     'debug_toolbar.panels.headers.HeaderDebugPanel',
     'debug_toolbar.panels.request_vars.RequestVarsDebugPanel',
     'debug_toolbar.panels.template.TemplateDebugPanel',
+    'debug_toolbar_mongo.panel.MongoDebugPanel',
     'debug_toolbar.panels.sql.SQLDebugPanel',
     'debug_toolbar.panels.signals.SignalDebugPanel',
     'debug_toolbar.panels.logger.LoggingPanel',
-    'debug_toolbar_mongo.panel.MongoDebugPanel',
 )
 
 ROOT_URLCONF = 'example.urls'

example/views.py

@@ -8,7 +8,9 @@ conn = pymongo.Connection()
 db = conn.debug_test
 
 def index(request):
-    list(db.test.find({'name': 'test'}))
+    #list(db.test.find({'name': 'test'}))
+    db.test.find({'name': 'test'}).count()
+    db.test.find({'name': 'test'}).count()
     list(db.test.find({'name': 'test', 'age': {'$lt': 134234}}).skip(1))
     db.test.find({'name': 'test'}).count()
     db.test.find({'name': 'test'}).skip(1).count(with_limit_and_skip=True)