Handle postgres CursorDebugWrapper in debugsqlshell.

Author: tim-schilling

debug_toolbar/management/commands/debugsqlshell.py

@@ -2,12 +2,17 @@
 
 import sqlparse
 from django.core.management.commands.shell import Command  # noqa
-from django.db.backends import utils as db_backends_utils
+from django.db import connection
+
+if connection.vendor == "postgresql":
+    from django.db.backends.postgresql import base as base_module
+else:
+    from django.db.backends import utils as base_module
 
 # 'debugsqlshell' is the same as the 'shell'.
 
 
-class PrintQueryWrapper(db_backends_utils.CursorDebugWrapper):
+class PrintQueryWrapper(base_module.CursorDebugWrapper):
     def execute(self, sql, params=()):
         start_time = time()
         try:
@@ -20,4 +25,4 @@ def execute(self, sql, params=()):
             print("{} [{:.2f}ms]".format(formatted_sql, duration))
 
 
-db_backends_utils.CursorDebugWrapper = PrintQueryWrapper
+base_module.CursorDebugWrapper = PrintQueryWrapper

tests/commands/test_debugsqlshell.py

@@ -3,15 +3,20 @@
 
 from django.contrib.auth.models import User
 from django.core import management
-from django.db.backends import utils as db_backends_utils
+from django.db import connection
 from django.test import TestCase
 from django.test.utils import override_settings
 
+if connection.vendor == "postgresql":
+    from django.db.backends.postgresql import base as base_module
+else:
+    from django.db.backends import utils as base_module
+
 
 @override_settings(DEBUG=True)
 class DebugSQLShellTestCase(TestCase):
     def setUp(self):
-        self.original_cursor_wrapper = db_backends_utils.CursorDebugWrapper
+        self.original_wrapper = base_module.CursorDebugWrapper
         # Since debugsqlshell monkey-patches django.db.backends.utils, we can
         # test it simply by loading it, without executing it. But we have to
         # undo the monkey-patch on exit.
@@ -20,7 +25,7 @@ def setUp(self):
         management.load_command_class(app_name, command_name)
 
     def tearDown(self):
-        db_backends_utils.CursorDebugWrapper = self.original_cursor_wrapper
+        base_module.CursorDebugWrapper = self.original_wrapper
 
     def test_command(self):
         original_stdout, sys.stdout = sys.stdout, io.StringIO()