Move postgresql connection logic to NormalCursorWrapper._record()

Previously, the logic for determining when to switch transaction IDs
was contained in SQLPanel.get_transaction_id().  However, this was not
sufficient, as it only looked at the transaction status after the query
was finished.  If two queries were issued consecutively but in separate
transactions, such as in the following:

    with transaction.atomic():
        list(User.objects.all())
    with transaction.atomic():
        list(Group.objects.all())

both queries would be given the same transaction ID since they would
both have the same transaction status after being executed
(TRANSACTION_STATUS_INTRANS).  Instead, move the logic to
NormalCursorWrapper._record() and compare the transaction status before
the query to the transaction status after the query to determine if a
new transaction was started.

Also, use the conn.status attribute for determining transaction status
instead of conn.get_transaction_status(), since the former is a local
connection variable, while the latter requires a call to the server.

Author: living180

debug_toolbar/panels/sql/panel.py

@@ -61,36 +61,29 @@ def __init__(self, *args, **kwargs):
         self._num_queries = 0
         self._queries = []
         self._databases = {}
-        self._transaction_status = {}
+        # synthetic transaction IDs, keyed by DB alias
         self._transaction_ids = {}
 
-    def get_transaction_id(self, alias):
-        if alias not in connections:
-            return
-        connection = connections[alias]
-        conn = connection.connection
-        if not conn:
-            return
-
-        if connection.vendor == "postgresql":
-            cur_status = conn.get_transaction_status()
-        else:
-            raise ValueError(connection.vendor)
-
-        last_status = self._transaction_status.get(alias)
-        self._transaction_status[alias] = cur_status
-
-        if not cur_status:
-            # No available state
-            return None
-
-        if cur_status != last_status:
-            if cur_status:
-                self._transaction_ids[alias] = uuid.uuid4().hex
-            else:
-                self._transaction_ids[alias] = None
-
-        return self._transaction_ids[alias]
+    def new_transaction_id(self, alias):
+        """
+        Generate and return a new synthetic transaction ID for the specified DB alias.
+        """
+        trans_id = uuid.uuid4().hex
+        self._transaction_ids[alias] = trans_id
+        return trans_id
+
+    def current_transaction_id(self, alias):
+        """
+        Return the current synthetic transaction ID for the specified DB alias.
+        """
+        trans_id = self._transaction_ids.get(alias)
+        # Sometimes it is not possible to detect the beginning of the first transaction,
+        # so current_transaction_id() will be called before new_transaction_id().  In
+        # that case there won't yet be a transaction ID. so it is necessary to generate
+        # one using new_transaction_id().
+        if trans_id is None:
+            trans_id = self.new_transaction_id(alias)
+        return trans_id
 
     def record(self, alias, **kwargs):
         self._queries.append((alias, kwargs))

debug_toolbar/panels/sql/tracking.py

@@ -10,8 +10,10 @@
 
 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
@@ -139,6 +141,14 @@ 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.status
+
         start_time = time()
         try:
             return method(sql, params)
@@ -156,10 +166,6 @@ def _record(self, method, sql, params):
                 pass  # object not JSON serializable
             template_info = get_template_info()
 
-            alias = self.db.alias
-            conn = self.db.connection
-            vendor = self.db.vendor
-
             # Sql might be an object (such as psycopg Composed).
             # For logging purposes, make sure it's str.
             sql = str(sql)
@@ -190,9 +196,29 @@ def _record(self, method, sql, params):
                     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.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": self.logger.get_transaction_id(alias),
+                        "trans_id": trans_id,
                         "trans_status": conn.get_transaction_status(),
                         "iso_level": iso_level,
                         "encoding": conn.encoding,