Add test for PostgreSQL transaction tracking

Author: living180

tests/panels/test_sql.py

@@ -7,7 +7,7 @@
 import django
 from asgiref.sync import sync_to_async
 from django.contrib.auth.models import User
-from django.db import connection
+from django.db import connection, transaction
 from django.db.models import Count
 from django.db.utils import DatabaseError
 from django.shortcuts import render
@@ -527,3 +527,76 @@ def test_aliases(self):
 
         query = self.panel._queries[-1]
         self.assertEqual(query[0], "replica")
+
+    @unittest.skipUnless(
+        connection.vendor == "postgresql", "Test valid only on PostgreSQL"
+    )
+    def test_transaction_status(self):
+        """
+        Test case for tracking the transaction status is properly associated
+        with queries.
+        """
+        self.assertEqual(len(self.panel._queries), 0)
+
+        with transaction.atomic():
+            list(User.objects.all())
+            list(User.objects.using("replica").all())
+
+            with transaction.atomic(using="replica"):
+                list(User.objects.all())
+                list(User.objects.using("replica").all())
+
+        with transaction.atomic():
+            list(User.objects.all())
+
+        list(User.objects.using("replica").all())
+
+        response = self.panel.process_request(self.request)
+        self.panel.generate_stats(self.request, response)
+
+        self.assertEqual(len(self.panel._queries), 6)
+
+        query = self.panel._queries[0]
+        self.assertEqual(query[0], "default")
+        self.assertIsNotNone(query[1]["trans_id"])
+        self.assertTrue(query[1].get("starts_trans", False))
+        self.assertTrue(query[1].get("in_trans", False))
+        self.assertFalse(query[1].get("ends_trans", False))
+
+        query = self.panel._queries[-1]
+        self.assertEqual(query[0], "replica")
+        self.assertIsNone(query[1]["trans_id"])
+        self.assertFalse(query[1].get("starts_trans", False))
+        self.assertFalse(query[1].get("in_trans", False))
+        self.assertFalse(query[1].get("ends_trans", False))
+        self.assertFalse(query[1].get("ends_trans", False))
+
+        query = self.panel._queries[2]
+        self.assertEqual(query[0], "default")
+        self.assertEqual(query[1]["trans_id"], self.panel._queries[0][1]["trans_id"])
+        self.assertFalse(query[1].get("starts_trans", False))
+        self.assertTrue(query[1].get("in_trans", False))
+        self.assertTrue(query[1].get("ends_trans", False))
+
+        query = self.panel._queries[3]
+        self.assertEqual(query[0], "replica")
+        self.assertNotEqual(query[1]["trans_id"], self.panel._queries[0][1]["trans_id"])
+        self.assertTrue(query[1].get("starts_trans", False))
+        self.assertTrue(query[1].get("in_trans", False))
+        self.assertTrue(query[1].get("ends_trans", False))
+
+        query = self.panel._queries[4]
+        self.assertEqual(query[0], "default")
+        self.assertNotEqual(query[1]["trans_id"], self.panel._queries[0][1]["trans_id"])
+        self.assertNotEqual(query[1]["trans_id"], self.panel._queries[3][1]["trans_id"])
+        self.assertTrue(query[1].get("starts_trans", False))
+        self.assertTrue(query[1].get("in_trans", False))
+        self.assertTrue(query[1].get("ends_trans", False))
+
+        query = self.panel._queries[5]
+        self.assertEqual(query[0], "replica")
+        self.assertIsNone(query[1]["trans_id"])
+        self.assertFalse(query[1].get("starts_trans", False))
+        self.assertFalse(query[1].get("in_trans", False))
+        self.assertFalse(query[1].get("ends_trans", False))
+        self.assertFalse(query[1].get("ends_trans", False))