Added psycopg3 support (#1739)

* Added support and test cases for psycopg3

updated trans_id to rely on transaction status instead of connection status as it works with pyscorg3

Use correct json property based on psycopg version

Skip test_tuple_param_conversion if using psycopg3

Move psycopg3 skip logic into decorator.

Author: nofalx

.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:

debug_toolbar/panels/sql/panel.py

@@ -17,31 +17,59 @@
 
 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)
 
 
 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)

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,
                     }
                 )

docs/changes.rst

@@ -14,6 +14,7 @@ Pending
   memory leaks and sometimes very verbose and hard to silence output in some
   environments (but not others). The maintainers judged that time and effort is
   better invested elsewhere.
+* Added support for psycopg3.
 
 3.8.1 (2022-12-03)
 ------------------

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)
 

tox.ini

@@ -3,9 +3,11 @@ isolated_build = true
 envlist =
     docs
     packaging
-    py{38,39,310}-dj{32,41,42}-{sqlite,postgresql,postgis,mysql}
+    py{38,39,310}-dj{32}-{sqlite,postgresql,postgis,mysql}
     py{310}-dj{40}-{sqlite}
-    py{310,311}-dj{41,42,main}-{sqlite,postgresql,postgis,mysql}
+    py{310,311}-dj{41}-{sqlite,postgresql,postgis,mysql}
+    py{310,311}-dj{42,main}-{sqlite,postgresql,postgis,mysql}
+    py{310,311}-dj{42,main}-psycopg3
 
 [testenv]
 deps =
@@ -14,6 +16,7 @@ deps =
     dj41: django~=4.1.3
     dj42: django>=4.2a1,<5
     postgresql: psycopg2-binary
+    psycopg3: psycopg[binary]
     postgis: psycopg2-binary
     mysql: mysqlclient
     djmain: https://github.com/django/django/archive/main.tar.gz
@@ -47,12 +50,13 @@ allowlist_externals = make
 pip_pre = True
 commands = python -b -W always -m coverage run -m django test -v2 {posargs:tests}
 
-[testenv:py{38,39,310,311}-dj{32,40,41,42,main}-postgresql]
+[testenv:py{38,39,310,311}-dj{32,40,41,42,main}-{postgresql,psycopg3}]
 setenv =
     {[testenv]setenv}
     DB_BACKEND = postgresql
     DB_PORT = {env:DB_PORT:5432}
 
+
 [testenv:py{38,39,310,311}-dj{32,40,41,42,main}-postgis]
 setenv =
     {[testenv]setenv}
@@ -97,5 +101,6 @@ python =
 DB_BACKEND =
     mysql: mysql
     postgresql: postgresql
-    postgis: postgresql
+    psycopg3: psycopg3
+    postgis: postgis
     sqlite3: sqlite