2 files changed, 18 insertions, 10 deletions
diff --git a/debug_toolbar/panels/sql/panel.py b/debug_toolbar/panels/sql/panel.py
index cb80901..84ac992 100644
--- a/debug_toolbar/panels/sql/panel.py
+++ b/debug_toolbar/panels/sql/panel.py
@@ -11,7 +11,7 @@ from debug_toolbar.panels import DebugPanel
 from debug_toolbar.panels.sql.forms import SQLSelectForm
 from debug_toolbar.utils import render_stacktrace
 from debug_toolbar.panels.sql.utils import reformat_sql
-from debug_toolbar.panels.sql.tracking import CursorWrapper
+from debug_toolbar.panels.sql.tracking import wrap_cursor, unwrap_cursor
 
 
 def get_isolation_level_display(engine, level):
@@ -128,16 +128,11 @@ class SQLDebugPanel(DebugPanel):
     def enable_instrumentation(self):
         # This is thread-safe because database connections are thread-local.
         for connection in connections.all():
-            if not hasattr(connection, '_djdt_cursor'):
-                connection._djdt_cursor = connection.cursor
-                connection.cursor = lambda: CursorWrapper(
-                    connection._djdt_cursor(), connection, self)
+            wrap_cursor(connection, self)
 
     def disable_instrumentation(self):
         for connection in connections.all():
-            if hasattr(connection, '_djdt_cursor'):
-                del connection._djdt_cursor
-                del connection.cursor
+            unwrap_cursor(connection)
 
     def process_response(self, request, response):
         if self._queries:
diff --git a/debug_toolbar/panels/sql/tracking.py b/debug_toolbar/panels/sql/tracking.py
index 283e410..989ce3d 100644
--- a/debug_toolbar/panels/sql/tracking.py
+++ b/debug_toolbar/panels/sql/tracking.py
@@ -37,8 +37,21 @@ state = ThreadLocalState()
 recording = state.recording  # export function
 
 
-def CursorWrapper(*args, **kwds):  # behave like a class
-    return state.Wrapper(*args, **kwds)
+def wrap_cursor(connection, panel):
+    if not hasattr(connection, '_djdt_cursor'):
+        connection._djdt_cursor = connection.cursor
+
+        def cursor():
+            return state.Wrapper(connection._djdt_cursor(), connection, panel)
+
+        connection.cursor = cursor
+        return cursor
+
+
+def unwrap_cursor(connection):
+    if hasattr(connection, '_djdt_cursor'):
+        del connection._djdt_cursor
+        del connection.cursor
 
 
 class ExceptionCursorWrapper(object):