From 5405757bfd16e2a1284c7dc4ff27d0b36c34b88d Mon Sep 17 00:00:00 2001
From: tschilling <schillingt@better-simple.com>
Date: Thu, 19 Dec 2019 21:22:09 -0600
Subject: [PATCH] Handle postgres CursorDebugWrapper in debugsqlshell.

---
 debug_toolbar/management/commands/debugsqlshell.py | 11 ++++++++---
 tests/commands/test_debugsqlshell.py               | 11 ++++++++---
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/debug_toolbar/management/commands/debugsqlshell.py b/debug_toolbar/management/commands/debugsqlshell.py
index ea39f3e1c..a532ea5ba 100644
--- a/debug_toolbar/management/commands/debugsqlshell.py
+++ b/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
diff --git a/tests/commands/test_debugsqlshell.py b/tests/commands/test_debugsqlshell.py
index 54cd248e0..9939c5ca9 100644
--- a/tests/commands/test_debugsqlshell.py
+++ b/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()