Repository: jazzband/django-debug-toolbar

Issue:
Add Support for psycopg3 
With the release of django 4.2a we can use psycopg3. However [django-debug-toolbar](https://github.com/jazzband/django-debug-toolbar) is not compatible with it.


Discussion:

Resolving patch:
diff --git a/debug_toolbar/panels/sql/panel.py b/debug_toolbar/panels/sql/panel.py
--- a/debug_toolbar/panels/sql/panel.py
+++ b/debug_toolbar/panels/sql/panel.py
@@ -17,15 +17,31 @@
 
 def get_isolation_level_display(vendor, level):
     if vendor == "postgresql":
-        import psycopg2.extensions
-
-        choices = {
-            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT: _("Autocommit"),
-            psycopg2.extensions.ISOLATION_LEVEL_READ_UNCOMMITTED: _("Read uncommitted"),
-            psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED: _("Read committed"),
-            psycopg2.extensions.ISOLATION_LEVEL_REPEATABLE_READ: _("Repeatable read"),
-            psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE: _("Serializable"),
-        }
+        try:
+            import psycopg
+
+            choices = {
+                # AUTOCOMMIT level does not exists in psycopg3
+                psycopg.IsolationLevel.READ_UNCOMMITTED: _("Read uncommitted"),
+                psycopg.IsolationLevel.READ_COMMITTED: _("Read committed"),
+                psycopg.IsolationLevel.REPEATABLE_READ: _("Repeatable read"),
+                psycopg.IsolationLevel.SERIALIZABLE: _("Serializable"),
+            }
+        except ImportError:
+            import psycopg2.extensions
+
+            choices = {
+                psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT: _("Autocommit"),
+                psycopg2.extensions.ISOLATION_LEVEL_READ_UNCOMMITTED: _(
+                    "Read uncommitted"
+                ),
+                psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED: _("Read committed"),
+                psycopg2.extensions.ISOLATION_LEVEL_REPEATABLE_READ: _(
+                    "Repeatable read"
+                ),
+                psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE: _("Serializable"),
+            }
+
     else:
         raise ValueError(vendor)
     return choices.get(level)
@@ -33,15 +49,27 @@ def get_isolation_level_display(vendor, level):
 
 def get_transaction_status_display(vendor, level):
     if vendor == "postgresql":
-        import psycopg2.extensions
-
-        choices = {
-            psycopg2.extensions.TRANSACTION_STATUS_IDLE: _("Idle"),
-            psycopg2.extensions.TRANSACTION_STATUS_ACTIVE: _("Active"),
-            psycopg2.extensions.TRANSACTION_STATUS_INTRANS: _("In transaction"),
-            psycopg2.extensions.TRANSACTION_STATUS_INERROR: _("In error"),
-            psycopg2.extensions.TRANSACTION_STATUS_UNKNOWN: _("Unknown"),
-        }
+        try:
+            import psycopg
+
+            choices = {
+                psycopg.pq.TransactionStatus.IDLE: _("Idle"),
+                psycopg.pq.TransactionStatus.ACTIVE: _("Active"),
+                psycopg.pq.TransactionStatus.INTRANS: _("In transaction"),
+                psycopg.pq.TransactionStatus.INERROR: _("In error"),
+                psycopg.pq.TransactionStatus.UNKNOWN: _("Unknown"),
+            }
+        except ImportError:
+            import psycopg2.extensions
+
+            choices = {
+                psycopg2.extensions.TRANSACTION_STATUS_IDLE: _("Idle"),
+                psycopg2.extensions.TRANSACTION_STATUS_ACTIVE: _("Active"),
+                psycopg2.extensions.TRANSACTION_STATUS_INTRANS: _("In transaction"),
+                psycopg2.extensions.TRANSACTION_STATUS_INERROR: _("In error"),
+                psycopg2.extensions.TRANSACTION_STATUS_UNKNOWN: _("Unknown"),
+            }
+
     else:
         raise ValueError(vendor)
     return choices.get(level)
diff --git a/debug_toolbar/panels/sql/tracking.py b/debug_toolbar/panels/sql/tracking.py
--- a/debug_toolbar/panels/sql/tracking.py
+++ b/debug_toolbar/panels/sql/tracking.py
@@ -9,11 +9,17 @@
 from debug_toolbar.utils import get_stack_trace, get_template_info
 
 try:
-    from psycopg2._json import Json as PostgresJson
-    from psycopg2.extensions import STATUS_IN_TRANSACTION
+    import psycopg
+
+    PostgresJson = psycopg.types.json.Jsonb
+    STATUS_IN_TRANSACTION = psycopg.pq.TransactionStatus.INTRANS
 except ImportError:
-    PostgresJson = None
-    STATUS_IN_TRANSACTION = None
+    try:
+        from psycopg2._json import Json as PostgresJson
+        from psycopg2.extensions import STATUS_IN_TRANSACTION
+    except ImportError:
+        PostgresJson = None
+        STATUS_IN_TRANSACTION = None
 
 # Prevents SQL queries from being sent to the DB. It's used
 # by the TemplatePanel to prevent the toolbar from issuing
@@ -126,7 +132,13 @@ def _quote_params(self, params):
 
     def _decode(self, param):
         if PostgresJson and isinstance(param, PostgresJson):
-            return param.dumps(param.adapted)
+            # psycopg3
+            if hasattr(param, "obj"):
+                return param.dumps(param.obj)
+            # psycopg2
+            if hasattr(param, "adapted"):
+                return param.dumps(param.adapted)
+
         # If a sequence type, decode each element separately
         if isinstance(param, (tuple, list)):
             return [self._decode(element) for element in param]
@@ -149,7 +161,7 @@ def _record(self, method, sql, params):
         if vendor == "postgresql":
             # The underlying DB connection (as opposed to Django's wrapper)
             conn = self.db.connection
-            initial_conn_status = conn.status
+            initial_conn_status = conn.info.transaction_status
 
         start_time = time()
         try:
@@ -166,7 +178,10 @@ def _record(self, method, sql, params):
 
             # Sql might be an object (such as psycopg Composed).
             # For logging purposes, make sure it's str.
-            sql = str(sql)
+            if vendor == "postgresql" and not isinstance(sql, str):
+                sql = sql.as_string(conn)
+            else:
+                sql = str(sql)
 
             params = {
                 "vendor": vendor,
@@ -205,7 +220,7 @@ def _record(self, method, sql, params):
                 # case where Django can start a transaction before the first query
                 # executes, so in that case logger.current_transaction_id() will
                 # generate a new transaction ID since one does not already exist.
-                final_conn_status = conn.status
+                final_conn_status = conn.info.transaction_status
                 if final_conn_status == STATUS_IN_TRANSACTION:
                     if initial_conn_status == STATUS_IN_TRANSACTION:
                         trans_id = self.logger.current_transaction_id(alias)
@@ -217,7 +232,7 @@ def _record(self, method, sql, params):
                 params.update(
                     {
                         "trans_id": trans_id,
-                        "trans_status": conn.get_transaction_status(),
+                        "trans_status": conn.info.transaction_status,
                         "iso_level": iso_level,
                     }
                 )


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -79,6 +79,12 @@ jobs:
       matrix:
         python-version: ['3.8', '3.9', '3.10', '3.11']
         database: [postgresql, postgis]
+        # Add psycopg3 to our matrix for 3.10 and 3.11
+        include:
+            - python-version: '3.10'
+              database: psycopg3
+            - python-version: '3.11'
+              database: psycopg3
 
     services:
       postgres:
diff --git a/tests/panels/test_sql.py b/tests/panels/test_sql.py
--- a/tests/panels/test_sql.py
+++ b/tests/panels/test_sql.py
@@ -16,6 +16,11 @@
 import debug_toolbar.panels.sql.tracking as sql_tracking
 from debug_toolbar import settings as dt_settings
 
+try:
+    import psycopg
+except ImportError:
+    psycopg = None
+
 from ..base import BaseMultiDBTestCase, BaseTestCase
 from ..models import PostgresJSON
 
@@ -222,9 +227,13 @@ def test_json_param_conversion(self):
         )
 
     @unittest.skipUnless(
-        connection.vendor == "postgresql", "Test valid only on PostgreSQL"
+        connection.vendor == "postgresql" and psycopg is None,
+        "Test valid only on PostgreSQL with psycopg2",
     )
     def test_tuple_param_conversion(self):
+        """
+        Regression test for tuple parameter conversion.
+        """
         self.assertEqual(len(self.panel._queries), 0)
 
         list(
@@ -377,12 +386,15 @@ def test_erroneous_query(self):
     @unittest.skipUnless(
         connection.vendor == "postgresql", "Test valid only on PostgreSQL"
     )
-    def test_execute_with_psycopg2_composed_sql(self):
+    def test_execute_with_psycopg_composed_sql(self):
         """
-        Test command executed using a Composed psycopg2 object is logged.
-        Ref: http://initd.org/psycopg/docs/sql.html
+        Test command executed using a Composed psycopg object is logged.
+        Ref: https://www.psycopg.org/psycopg3/docs/api/sql.html
         """
-        from psycopg2 import sql
+        try:
+            from psycopg import sql
+        except ImportError:
+            from psycopg2 import sql
 
         self.assertEqual(len(self.panel._queries), 0)