Generate fake transaction ids so we can match up transaction end state correctly

1 files changed, 10 insertions, 16 deletions

diff --git a/debug_toolbar/utils/tracking/db.py b/debug_toolbar/utils/tracking/db.py
index 234247a..87f4550 100644
--- a/debug_toolbar/utils/tracking/db.py
+++ b/debug_toolbar/utils/tracking/db.py
@@ -29,17 +29,7 @@ class CursorWrapper(object):
         self.logger = logger
 
     def execute(self, sql, params=()):
-        alias = getattr(self, 'alias', 'default')
-        conn = connections[alias].connection
-        # HACK: avoid imports
-        if conn:
-            engine = conn.__class__.__module__.split('.', 1)[0]
-        else:
-            engine = 'unknown'
-
         start = datetime.now()
-        if engine == 'psycopg2':
-            trans_status = self.logger.get_transaction_status(alias)
         try:
             return self.cursor.execute(sql, params)
         finally:
@@ -66,6 +56,14 @@ class CursorWrapper(object):
                 pass
             del cur_frame
 
+            alias = getattr(self, 'alias', 'default')
+            conn = connections[alias].connection
+            # HACK: avoid imports
+            if conn:
+                engine = conn.__class__.__module__.split('.', 1)[0]
+            else:
+                engine = 'unknown'
+
             params = {
                 'engine': engine,
                 'alias': alias,
@@ -83,13 +81,9 @@ class CursorWrapper(object):
             }
 
             if engine == 'psycopg2':
-                import psycopg2.extensions
-                cur_trans_status = self.logger.get_transaction_status(alias, True)
                 params.update({
-                    'starts_trans': cur_trans_status > trans_status,
-                    'ends_trans': cur_trans_status < trans_status,
-                    'in_trans': cur_trans_status == psycopg2.extensions.TRANSACTION_STATUS_INTRANS,
-                    'trans_status': cur_trans_status,
+                    'trans_id': self.logger.get_transaction_id(alias),
+                    'trans_status': conn.get_transaction_status(),
                     'iso_level': conn.isolation_level,
                     'encoding': conn.encoding,
                 })