Fix psycopg3 tests

Several tests (such as SQLPanelTestCase.test_cursor_wrapper_singleton)
are written to ensure that only a single cursor wrapper is
instantiated during a test.  However, this fails when using psycopg3,
since the .last_executed_query() call in NormalCursorWrapper._record()
ends up creating an additional cursor (via [1]).  To avoid this, use a
._djdt_in_record attribute on the database wrapper.  Make the
NormalCursorWrapper._record() method set ._djdt_in_record to True on
entry and reset it to False on exit.  Then in the overridden database
wrapper .cursor() and .chunked_cursor() methods, check the
._djdt_in_record attribute and return the original cursor without
wrapping if the attribute is True.

[1] https://github.com/django/django/blob/879e5d587b84e6fc961829611999431778eb9f6a/django/db/backends/postgresql/psycopg_any.py#L21

Author: living180

debug_toolbar/panels/sql/tracking.py

@@ -43,16 +43,21 @@ def cursor(*args, **kwargs):
             # See:
             # https://github.com/jazzband/django-debug-toolbar/pull/615
             # https://github.com/jazzband/django-debug-toolbar/pull/896
+            cursor = connection._djdt_cursor(*args, **kwargs)
+            if connection._djdt_in_record:
+                return cursor
             if allow_sql.get():
                 wrapper = NormalCursorWrapper
             else:
                 wrapper = ExceptionCursorWrapper
-            return wrapper(connection._djdt_cursor(*args, **kwargs), connection, panel)
+            return wrapper(cursor, connection, panel)
 
         def chunked_cursor(*args, **kwargs):
             # prevent double wrapping
             # solves https://github.com/jazzband/django-debug-toolbar/issues/1239
             cursor = connection._djdt_chunked_cursor(*args, **kwargs)
+            if connection._djdt_in_record:
+                return cursor
             if not isinstance(cursor, BaseCursorWrapper):
                 if allow_sql.get():
                     wrapper = NormalCursorWrapper
@@ -63,6 +68,7 @@ def chunked_cursor(*args, **kwargs):
 
         connection.cursor = cursor
         connection.chunked_cursor = chunked_cursor
+        connection._djdt_in_record = False
 
 
 def unwrap_cursor(connection):
@@ -154,90 +160,96 @@ def _decode(self, param):
             return "(encoded string)"
 
     def _record(self, method, sql, params):
-        alias = self.db.alias
-        vendor = self.db.vendor
-
-        if vendor == "postgresql":
-            # The underlying DB connection (as opposed to Django's wrapper)
-            conn = self.db.connection
-            initial_conn_status = conn.info.transaction_status
-
-        start_time = time()
+        self.db._djdt_in_record = True
         try:
-            return method(sql, params)
-        finally:
-            stop_time = time()
-            duration = (stop_time - start_time) * 1000
-            _params = ""
-            try:
-                _params = json.dumps(self._decode(params))
-            except TypeError:
-                pass  # object not JSON serializable
-            template_info = get_template_info()
-
-            # Sql might be an object (such as psycopg Composed).
-            # For logging purposes, make sure it's str.
-            if vendor == "postgresql" and not isinstance(sql, str):
-                sql = sql.as_string(conn)
-            else:
-                sql = str(sql)
-
-            params = {
-                "vendor": vendor,
-                "alias": alias,
-                "sql": self.db.ops.last_executed_query(
-                    self.cursor, sql, self._quote_params(params)
-                ),
-                "duration": duration,
-                "raw_sql": sql,
-                "params": _params,
-                "raw_params": params,
-                "stacktrace": get_stack_trace(skip=2),
-                "start_time": start_time,
-                "stop_time": stop_time,
-                "is_slow": duration > dt_settings.get_config()["SQL_WARNING_THRESHOLD"],
-                "is_select": sql.lower().strip().startswith("select"),
-                "template_info": template_info,
-            }
+            alias = self.db.alias
+            vendor = self.db.vendor
 
             if vendor == "postgresql":
-                # If an erroneous query was ran on the connection, it might
-                # be in a state where checking isolation_level raises an
-                # exception.
+                # The underlying DB connection (as opposed to Django's wrapper)
+                conn = self.db.connection
+                initial_conn_status = conn.info.transaction_status
+
+            start_time = time()
+            try:
+                return method(sql, params)
+            finally:
+                stop_time = time()
+                duration = (stop_time - start_time) * 1000
+                _params = ""
                 try:
-                    iso_level = conn.isolation_level
-                except conn.InternalError:
-                    iso_level = "unknown"
-                # PostgreSQL does not expose any sort of transaction ID, so it is
-                # necessary to generate synthetic transaction IDs here.  If the
-                # connection was not in a transaction when the query started, and was
-                # after the query finished, a new transaction definitely started, so get
-                # a new transaction ID from logger.new_transaction_id().  If the query
-                # was in a transaction both before and after executing, make the
-                # assumption that it is the same transaction and get the current
-                # transaction ID from logger.current_transaction_id().  There is an edge
-                # 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.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)
-                    else:
-                        trans_id = self.logger.new_transaction_id(alias)
+                    _params = json.dumps(self._decode(params))
+                except TypeError:
+                    pass  # object not JSON serializable
+                template_info = get_template_info()
+
+                # Sql might be an object (such as psycopg Composed).
+                # For logging purposes, make sure it's str.
+                if vendor == "postgresql" and not isinstance(sql, str):
+                    sql = sql.as_string(conn)
                 else:
-                    trans_id = None
-
-                params.update(
-                    {
-                        "trans_id": trans_id,
-                        "trans_status": conn.info.transaction_status,
-                        "iso_level": iso_level,
-                    }
-                )
-
-            # We keep `sql` to maintain backwards compatibility
-            self.logger.record(**params)
+                    sql = str(sql)
+
+                params = {
+                    "vendor": vendor,
+                    "alias": alias,
+                    "sql": self.db.ops.last_executed_query(
+                        self.cursor, sql, self._quote_params(params)
+                    ),
+                    "duration": duration,
+                    "raw_sql": sql,
+                    "params": _params,
+                    "raw_params": params,
+                    "stacktrace": get_stack_trace(skip=2),
+                    "start_time": start_time,
+                    "stop_time": stop_time,
+                    "is_slow": (
+                        duration > dt_settings.get_config()["SQL_WARNING_THRESHOLD"]
+                    ),
+                    "is_select": sql.lower().strip().startswith("select"),
+                    "template_info": template_info,
+                }
+
+                if vendor == "postgresql":
+                    # If an erroneous query was ran on the connection, it might
+                    # be in a state where checking isolation_level raises an
+                    # exception.
+                    try:
+                        iso_level = conn.isolation_level
+                    except conn.InternalError:
+                        iso_level = "unknown"
+                    # PostgreSQL does not expose any sort of transaction ID, so it is
+                    # necessary to generate synthetic transaction IDs here.  If the
+                    # connection was not in a transaction when the query started, and was
+                    # after the query finished, a new transaction definitely started, so get
+                    # a new transaction ID from logger.new_transaction_id().  If the query
+                    # was in a transaction both before and after executing, make the
+                    # assumption that it is the same transaction and get the current
+                    # transaction ID from logger.current_transaction_id().  There is an edge
+                    # 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.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)
+                        else:
+                            trans_id = self.logger.new_transaction_id(alias)
+                    else:
+                        trans_id = None
+
+                    params.update(
+                        {
+                            "trans_id": trans_id,
+                            "trans_status": conn.info.transaction_status,
+                            "iso_level": iso_level,
+                        }
+                    )
+
+                # We keep `sql` to maintain backwards compatibility
+                self.logger.record(**params)
+        finally:
+            self.db._djdt_in_record = False
 
     def callproc(self, procname, params=None):
         return self._record(self.cursor.callproc, procname, params)