Add utilities to inject and monitor functions. Change DB tracking to use new injection method on BaseDatabaseWrapper.cursor

7 files changed, 401 insertions, 181 deletions

diff --git a/debug_toolbar/panels/sql.py b/debug_toolbar/panels/sql.py
index 4c7b8a1..ba83455 100644
--- a/debug_toolbar/panels/sql.py
+++ b/debug_toolbar/panels/sql.py
@@ -1,188 +1,30 @@
-from datetime import datetime
-import os
 import re
-import sys
-import SocketServer
-import traceback
 
-import django
 from django.conf import settings
-try:
-    from django.db import connections
-except ImportError:
-    # Compat with < Django 1.2
-    from django.db import connection
-    connections = {'default': connection}
 from django.db.backends import BaseDatabaseWrapper
-from django.views.debug import linebreak_iter
-from django.template import Node
-from django.template.defaultfilters import escape
 from django.template.loader import render_to_string
-from django.utils import simplejson
-from django.utils.encoding import force_unicode
-from django.utils.hashcompat import sha_constructor
+from django.utils.html import escape
 from django.utils.safestring import mark_safe
 from django.utils.translation import ugettext_lazy as _, ungettext_lazy as __
 
+from debug_toolbar.utils.compat.db import connections
 from debug_toolbar.middleware import DebugToolbarMiddleware
 from debug_toolbar.panels import DebugPanel
 from debug_toolbar.utils import sqlparse
-
-# Figure out some paths
-django_path = os.path.realpath(os.path.dirname(django.__file__))
-socketserver_path = os.path.realpath(os.path.dirname(SocketServer.__file__))
-
-# TODO:This should be set in the toolbar loader as a default and panels should
-# get a copy of the toolbar object with access to its config dictionary
-SQL_WARNING_THRESHOLD = getattr(settings, 'DEBUG_TOOLBAR_CONFIG', {}) \
-                            .get('SQL_WARNING_THRESHOLD', 500)
-
-def tidy_stacktrace(strace):
-    """
-    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)
-    """
-    trace = []
-    for s in strace[:-1]:
-        s_path = os.path.realpath(s[0])
-        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
-        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 CursorWrapper(object):
-    """
-    Wraps a cursor and logs queries.
-    """
-    
-    def __init__(self, cursor, db):
-        self.cursor = cursor
-        self.db = db # Instance of a BaseDatabaseWrapper subclass
-
-    def execute(self, sql, params=()):
-        djdt = DebugToolbarMiddleware.get_current()
-        if not djdt:
-            return self.cursor.execute(sql, params)
-
-        panel = djdt.get_panel(SQLDebugPanel)
-        start = datetime.now()
-        try:
-            return self.cursor.execute(sql, params)
-        finally:
-            stop = datetime.now()
-            duration = ms_from_timedelta(stop - start)
-            stacktrace = tidy_stacktrace(traceback.extract_stack())
-            _params = ''
-            try:
-                _params = simplejson.dumps([force_unicode(x, strings_only=True) 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
-            del cur_frame
-
-            # We keep `sql` to maintain backwards compatibility
-            panel.record(**{
-                'alias': getattr(self, 'alias', 'default'),
-                'sql': self.db.ops.last_executed_query(self.cursor, sql, params),
-                'duration': duration,
-                'raw_sql': sql,
-                'params': _params,
-                'hash': sha_constructor(settings.SECRET_KEY + sql + _params).hexdigest(),
-                'stacktrace': stacktrace,
-                'start_time': start,
-                'stop_time': stop,
-                'is_slow': (duration > SQL_WARNING_THRESHOLD),
-                'is_select': sql.lower().strip().startswith('select'),
-                'template_info': template_info,
-            })
-
-    def executemany(self, sql, param_list):
-        return self.cursor.executemany(sql, param_list)
-
-    def __getattr__(self, attr):
-        if attr in self.__dict__:
-            return self.__dict__[attr]
-        else:
-            return getattr(self.cursor, attr)
-
-    def __iter__(self):
-        return iter(self.cursor)
-
-def inject_sql_tracker(cls):
-    """
-    Injects a replacement execute method which records queries within the SQLPanel.
-    """
-    import warnings
+from debug_toolbar.utils.tracking.db import CursorWrapper
+from debug_toolbar.utils.tracking import replace_call
+
+# Inject our tracking cursor
+@replace_call(BaseDatabaseWrapper.cursor)
+def cursor(func, self):
+    result = func(self)
+
+    djdt = DebugToolbarMiddleware.get_current()
+    if not djdt:
+        return result
+    logger = djdt.get_panel(SQLDebugPanel)
     
-    if not hasattr(cls, 'cursor'):
-        warnings.warn('Unable to patch %r: missing cursor method' % cls)
-
-    if getattr(cls.cursor, 'djdt_tracked', False):
-        return
-
-    def cursor(self):
-        result = cls.cursor.__wrapped__(self)
-        return CursorWrapper(result, self)
-
-    cursor.djdt_tracked = True
-    cursor.__wrapped__ = cls.cursor
-
-    cls.cursor = cursor
-
-# Inject our tracking code into the existing CursorWrapper's
-inject_sql_tracker(BaseDatabaseWrapper)
+    return CursorWrapper(result, self, logger=logger)
 
 class SQLDebugPanel(DebugPanel):
     """
@@ -286,12 +128,6 @@ class SQLDebugPanel(DebugPanel):
 
         return render_to_string('debug_toolbar/panels/sql.html', context)
 
-def ms_from_timedelta(td):
-    """
-    Given a timedelta object, returns a float representing milliseconds
-    """
-    return (td.seconds * 1000) + (td.microseconds / 1000.0)
-
 class BoldKeywordFilter(sqlparse.filters.Filter):
     """sqlparse filter to bold SQL keywords"""
     def process(self, stack, stream):
@@ -300,7 +136,7 @@ class BoldKeywordFilter(sqlparse.filters.Filter):
             is_keyword = token_type in sqlparse.tokens.Keyword
             if is_keyword:
                 yield sqlparse.tokens.Text, '<strong>'
-            yield token_type, django.utils.html.escape(value)
+            yield token_type, escape(value)
             if is_keyword:
                 yield sqlparse.tokens.Text, '</strong>'
 
diff --git a/debug_toolbar/tests/tests.py b/debug_toolbar/tests/tests.py
index cdc1947..476802c 100644
--- a/debug_toolbar/tests/tests.py
+++ b/debug_toolbar/tests/tests.py
@@ -1,6 +1,7 @@
 from debug_toolbar.middleware import DebugToolbarMiddleware
 from debug_toolbar.panels.sql import SQLDebugPanel
 from debug_toolbar.toolbar.loader import DebugToolbar
+from debug_toolbar.utils.tracking import pre_dispatch, post_dispatch, callbacks
 
 from django.contrib.auth.models import User
 from django.test import TestCase
@@ -36,4 +37,138 @@ class SQLPanelTestCase(BaseTestCase):
         self.assertTrue('sql' in query[1])
         self.assertTrue('duration' in query[1])
         self.assertTrue('stacktrace' in query[1])
-        
\ No newline at end of file
+
+def module_func(*args, **kwargs):
+    """Used by dispatch tests"""
+    return 'blah'
+
+class TrackingTestCase(BaseTestCase):
+    @classmethod
+    def class_method(cls, *args, **kwargs):
+        return 'blah'
+
+    def class_func(self, *args, **kwargs):
+        """Used by dispatch tests"""
+        return 'blah'
+    
+    def test_pre_hook(self):
+        foo = {}
+        
+        @pre_dispatch(module_func)
+        def test(**kwargs):
+            foo.update(kwargs)
+            
+        self.assertTrue(hasattr(module_func, '__wrapped__'))
+        self.assertEquals(len(callbacks['before']), 1)
+        
+        module_func('hi', foo='bar')
+        
+        self.assertTrue('sender' in foo, foo)
+        # best we can do
+        self.assertEquals(foo['sender'].__name__, 'module_func')
+        self.assertTrue('start' in foo, foo)
+        self.assertGreater(foo['start'], 0)
+        self.assertTrue('stop' not in foo, foo)
+        self.assertTrue('args' in foo, foo)
+        self.assertTrue(len(foo['args']), 1)
+        self.assertEquals(foo['args'][0], 'hi')
+        self.assertTrue('kwargs' in foo, foo)
+        self.assertTrue(len(foo['kwargs']), 1)
+        self.assertTrue('foo' in foo['kwargs'])
+        self.assertEquals(foo['kwargs']['foo'], 'bar')
+    
+        callbacks['before'] = {}
+    
+        @pre_dispatch(TrackingTestCase.class_func)
+        def test(**kwargs):
+            foo.update(kwargs)
+    
+        self.assertTrue(hasattr(TrackingTestCase.class_func, '__wrapped__'))
+        self.assertEquals(len(callbacks['before']), 1)
+
+        self.class_func('hello', foo='bar')
+
+        self.assertTrue('sender' in foo, foo)
+        # best we can do
+        self.assertEquals(foo['sender'].__name__, 'class_func')
+        self.assertTrue('start' in foo, foo)
+        self.assertGreater(foo['start'], 0)
+        self.assertTrue('stop' not in foo, foo)
+        self.assertTrue('args' in foo, foo)
+        self.assertTrue(len(foo['args']), 2)
+        self.assertEquals(foo['args'][1], 'hello')
+        self.assertTrue('kwargs' in foo, foo)
+        self.assertTrue(len(foo['kwargs']), 1)
+        self.assertTrue('foo' in foo['kwargs'])
+        self.assertEquals(foo['kwargs']['foo'], 'bar')
+
+        # callbacks['before'] = {}
+        #     
+        #         @pre_dispatch(TrackingTestCase.class_method)
+        #         def test(**kwargs):
+        #             foo.update(kwargs)
+        #     
+        #         self.assertTrue(hasattr(TrackingTestCase.class_method, '__wrapped__'))
+        #         self.assertEquals(len(callbacks['before']), 1)
+        # 
+        #         TrackingTestCase.class_method()
+        # 
+        #         self.assertTrue('sender' in foo, foo)
+        #         # best we can do
+        #         self.assertEquals(foo['sender'].__name__, 'class_method')
+        #         self.assertTrue('start' in foo, foo)
+        #         self.assertTrue('stop' not in foo, foo)
+        #         self.assertTrue('args' in foo, foo)
+
+    def test_post_hook(self):
+        foo = {}
+        
+        @post_dispatch(module_func)
+        def test(**kwargs):
+            foo.update(kwargs)
+            
+        self.assertTrue(hasattr(module_func, '__wrapped__'))
+        self.assertEquals(len(callbacks['after']), 1)
+        
+        module_func('hi', foo='bar')
+        
+        self.assertTrue('sender' in foo, foo)
+        # best we can do
+        self.assertEquals(foo['sender'].__name__, 'module_func')
+        self.assertTrue('start' in foo, foo)
+        self.assertGreater(foo['start'], 0)
+        self.assertTrue('stop' in foo, foo)
+        self.assertGreater(foo['stop'], foo['start'])
+        self.assertTrue('args' in foo, foo)
+        self.assertTrue(len(foo['args']), 1)
+        self.assertEquals(foo['args'][0], 'hi')
+        self.assertTrue('kwargs' in foo, foo)
+        self.assertTrue(len(foo['kwargs']), 1)
+        self.assertTrue('foo' in foo['kwargs'])
+        self.assertEquals(foo['kwargs']['foo'], 'bar')
+    
+        callbacks['after'] = {}
+    
+        @post_dispatch(TrackingTestCase.class_func)
+        def test(**kwargs):
+            foo.update(kwargs)
+    
+        self.assertTrue(hasattr(TrackingTestCase.class_func, '__wrapped__'))
+        self.assertEquals(len(callbacks['after']), 1)
+
+        self.class_func('hello', foo='bar')
+
+        self.assertTrue('sender' in foo, foo)
+        # best we can do
+        self.assertEquals(foo['sender'].__name__, 'class_func')
+        self.assertTrue('start' in foo, foo)
+        self.assertGreater(foo['start'], 0)
+        self.assertTrue('stop' in foo, foo)
+        self.assertGreater(foo['stop'], foo['start'])
+        self.assertTrue('args' in foo, foo)
+        self.assertTrue(len(foo['args']), 2)
+        self.assertEquals(foo['args'][1], 'hello')
+        self.assertTrue('kwargs' in foo, foo)
+        self.assertTrue(len(foo['kwargs']), 1)
+        self.assertTrue('foo' in foo['kwargs'])
+        self.assertEquals(foo['kwargs']['foo'], 'bar')
\ No newline at end of file
diff --git a/debug_toolbar/utils/__init__.py b/debug_toolbar/utils/__init__.py
index e69de29..61bb717 100644
--- a/debug_toolbar/utils/__init__.py
+++ b/debug_toolbar/utils/__init__.py
@@ -0,0 +1,68 @@
+import os.path
+import django
+import SocketServer
+
+from django.conf import settings
+from django.views.debug import linebreak_iter
+
+# Figure out some paths
+django_path = os.path.realpath(os.path.dirname(django.__file__))
+socketserver_path = os.path.realpath(os.path.dirname(SocketServer.__file__))
+
+def ms_from_timedelta(td):
+    """
+    Given a timedelta object, returns a float representing milliseconds
+    """
+    return (td.seconds * 1000) + (td.microseconds / 1000.0)
+
+def tidy_stacktrace(strace):
+    """
+    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)
+    """
+    trace = []
+    for s in strace[:-1]:
+        s_path = os.path.realpath(s[0])
+        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
+        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,
+    }
\ No newline at end of file
diff --git a/debug_toolbar/utils/compat/__init__.py b/debug_toolbar/utils/compat/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/debug_toolbar/utils/compat/__init__.py
diff --git a/debug_toolbar/utils/compat/db.py b/debug_toolbar/utils/compat/db.py
new file mode 100644
index 0000000..f3b37e6
--- /dev/null
+++ b/debug_toolbar/utils/compat/db.py
@@ -0,0 +1,6 @@
+try:
+    from django.db import connections
+except ImportError:
+    # Compat with < Django 1.2
+    from django.db import connection
+    connections = {'default': connection}
\ No newline at end of file
diff --git a/debug_toolbar/utils/tracking/__init__.py b/debug_toolbar/utils/tracking/__init__.py
new file mode 100644
index 0000000..db8ff18
--- /dev/null
+++ b/debug_toolbar/utils/tracking/__init__.py
@@ -0,0 +1,90 @@
+import logging
+import time
+import types
+
+def post_dispatch(func):
+    def wrapped(callback):
+        register_hook(func, 'after', callback)
+        return callback
+    return wrapped
+
+def pre_dispatch(func):
+    def wrapped(callback):
+        register_hook(func, 'before', callback)
+        return callback
+    return wrapped
+
+def replace_call(func):
+    def inner(callback):
+        def wrapped(*args, **kwargs):
+            return callback(func, *args, **kwargs)
+
+        actual = getattr(func, '__wrapped__', func)
+        wrapped.__wrapped__ = actual
+        wrapped.__doc__ = getattr(actual, '__doc__', None)
+        wrapped.__name__ = actual.__name__
+
+        _replace_function(func, wrapped)
+        return wrapped
+    return inner
+
+def fire_hook(hook, sender, **kwargs):
+    try:
+        for callback in callbacks[hook].get(id(sender), []):
+            callback(sender=sender, **kwargs)
+    except Exception, e:
+        # Log the exception, dont mess w/ the underlying function
+        logging.exception(e)
+
+def _replace_function(func, wrapped):
+    if isinstance(func, types.FunctionType):
+        if func.__module__ == '__builtin__':
+            # oh shit
+            __builtins__[func] = wrapped
+        else:
+            module = __import__(func.__module__, {}, {}, [func.__module__], 0)
+            setattr(module, func.__name__, wrapped)
+    elif getattr(func, 'im_self', None):
+        # TODO: classmethods
+        raise NotImplementedError
+    elif hasattr(func, 'im_class'):
+        # for unbound methods
+        setattr(func.im_class, func.__name__, wrapped)
+    else:
+        raise NotImplementedError
+
+callbacks = {
+    'before': {},
+    'after': {},
+}
+
+def register_hook(func, hook, callback):
+    """
+    def myhook(sender, args, kwargs):
+        print func, "executed
+        print "args:", args
+        print "kwargs:", kwargs
+    register_hook(BaseDatabaseWrapper.cursor, 'before', myhook)
+    """
+
+    assert hook in ('before', 'after')
+
+    def wrapped(*args, **kwargs):
+        start = time.time()
+        fire_hook('before', sender=wrapped.__wrapped__, args=args, kwargs=kwargs,
+                  start=start)
+        result = wrapped.__wrapped__(*args, **kwargs)
+        stop = time.time()
+        fire_hook('after', sender=wrapped.__wrapped__, args=args, kwargs=kwargs,
+                  result=result, start=start, stop=stop)
+    actual = getattr(func, '__wrapped__', func)
+    wrapped.__wrapped__ = actual
+    wrapped.__doc__ = getattr(actual, '__doc__', None)
+    wrapped.__name__ = actual.__name__
+    
+    id_ = id(actual)
+    if id_ not in callbacks[hook]:
+        callbacks[hook][id_] = []
+    callbacks[hook][id_].append(callback)
+
+    _replace_function(func, wrapped)
\ No newline at end of file
diff --git a/debug_toolbar/utils/tracking/db.py b/debug_toolbar/utils/tracking/db.py
new file mode 100644
index 0000000..4c9ee53
--- /dev/null
+++ b/debug_toolbar/utils/tracking/db.py
@@ -0,0 +1,85 @@
+import sys
+import traceback
+
+from datetime import datetime
+
+from django.conf import settings
+from django.template import Node
+from django.utils import simplejson
+from django.utils.encoding import force_unicode
+from django.utils.hashcompat import sha_constructor
+
+from debug_toolbar.utils import ms_from_timedelta, tidy_stacktrace, get_template_info
+
+# TODO:This should be set in the toolbar loader as a default and panels should
+# get a copy of the toolbar object with access to its config dictionary
+SQL_WARNING_THRESHOLD = getattr(settings, 'DEBUG_TOOLBAR_CONFIG', {}) \
+                            .get('SQL_WARNING_THRESHOLD', 500)
+
+class CursorWrapper(object):
+    """
+    Wraps a cursor and logs queries.
+    """
+    
+    def __init__(self, cursor, db, logger):
+        self.cursor = cursor
+        # Instance of a BaseDatabaseWrapper subclass
+        self.db = db
+        # logger must implement a ``record`` method
+        self.logger = logger
+
+    def execute(self, sql, params=()):
+        start = datetime.now()
+        try:
+            return self.cursor.execute(sql, params)
+        finally:
+            stop = datetime.now()
+            duration = ms_from_timedelta(stop - start)
+            stacktrace = tidy_stacktrace(traceback.extract_stack())
+            _params = ''
+            try:
+                _params = simplejson.dumps([force_unicode(x, strings_only=True) 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
+            del cur_frame
+
+            # We keep `sql` to maintain backwards compatibility
+            self.logger.record(**{
+                'alias': getattr(self, 'alias', 'default'),
+                'sql': self.db.ops.last_executed_query(self.cursor, sql, params),
+                'duration': duration,
+                'raw_sql': sql,
+                'params': _params,
+                'hash': sha_constructor(settings.SECRET_KEY + sql + _params).hexdigest(),
+                'stacktrace': stacktrace,
+                'start_time': start,
+                'stop_time': stop,
+                'is_slow': (duration > SQL_WARNING_THRESHOLD),
+                'is_select': sql.lower().strip().startswith('select'),
+                'template_info': template_info,
+            })
+
+    def executemany(self, sql, param_list):
+        return self.cursor.executemany(sql, param_list)
+
+    def __getattr__(self, attr):
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        else:
+            return getattr(self.cursor, attr)
+
+    def __iter__(self):
+        return iter(self.cursor)
\ No newline at end of file