Refactor all uses of thread locals to be more consistant and sane.

git-svn-id: http://code.djangoproject.com/svn/django/trunk@15232 bcc190cf-cafb-0310-a4f2-bffc1f526a37

django/db/transaction.py

@@ -25,6 +25,7 @@ except ImportError:
 from django.conf import settings
 from django.db import connections, DEFAULT_DB_ALIAS
 
+
 class TransactionManagementError(Exception):
     """
     This exception is thrown when something bad happens with transaction
@@ -32,19 +33,6 @@ class TransactionManagementError(Exception):
     """
     pass
 
-# The states are dictionaries of dictionaries of lists. The key to the outer
-# dict is the current thread, and the key to the inner dictionary is the
-# connection alias and the list is handled as a stack of values.
-state = {}
-savepoint_state = {}
-
-# The dirty flag is set by *_unless_managed functions to denote that the
-# code under transaction management has changed things to require a
-# database commit.
-# This is a dictionary mapping thread to a dictionary mapping connection
-# alias to a boolean.
-dirty = {}
-
 def enter_transaction_management(managed=True, using=None):
     """
     Enters transaction management for a running thread. It must be balanced with
@@ -58,15 +46,14 @@ def enter_transaction_management(managed=True, using=None):
     if using is None:
         using = DEFAULT_DB_ALIAS
     connection = connections[using]
-    thread_ident = thread.get_ident()
-    if thread_ident in state and state[thread_ident].get(using):
-        state[thread_ident][using].append(state[thread_ident][using][-1])
+
+    if connection.transaction_state:
+        connection.transaction_state.append(connection.transaction_state[-1])
     else:
-        state.setdefault(thread_ident, {})
-        state[thread_ident][using] = [settings.TRANSACTIONS_MANAGED]
-    if thread_ident not in dirty or using not in dirty[thread_ident]:
-        dirty.setdefault(thread_ident, {})
-        dirty[thread_ident][using] = False
+        connection.transaction_state.append(settings.TRANSACTIONS_MANAGED)
+
+    if connection.dirty is None:
+        connection.dirty = False
     connection._enter_transaction_management(managed)
 
 def leave_transaction_management(using=None):
@@ -78,16 +65,18 @@ def leave_transaction_management(using=None):
     if using is None:
         using = DEFAULT_DB_ALIAS
     connection = connections[using]
+
     connection._leave_transaction_management(is_managed(using=using))
-    thread_ident = thread.get_ident()
-    if thread_ident in state and state[thread_ident].get(using):
-        del state[thread_ident][using][-1]
+    if connection.transaction_state:
+        del connection.transaction_state[-1]
     else:
-        raise TransactionManagementError("This code isn't under transaction management")
-    if dirty.get(thread_ident, {}).get(using, False):
+        raise TransactionManagementError("This code isn't under transaction "
+            "management")
+    if connection.dirty:
         rollback(using=using)
-        raise TransactionManagementError("Transaction managed block ended with pending COMMIT/ROLLBACK")
-    dirty[thread_ident][using] = False
+        raise TransactionManagementError("Transaction managed block ended with "
+            "pending COMMIT/ROLLBACK")
+    connection.dirty = False
 
 def is_dirty(using=None):
     """
@@ -96,7 +85,9 @@ def is_dirty(using=None):
     """
     if using is None:
         using = DEFAULT_DB_ALIAS
-    return dirty.get(thread.get_ident(), {}).get(using, False)
+    connection = connections[using]
+
+    return connection.dirty
 
 def set_dirty(using=None):
     """
@@ -106,11 +97,13 @@ def set_dirty(using=None):
     """
     if using is None:
         using = DEFAULT_DB_ALIAS
-    thread_ident = thread.get_ident()
-    if thread_ident in dirty and using in dirty[thread_ident]:
-        dirty[thread_ident][using] = True
+    connection = connections[using]
+
+    if connection.dirty is not None:
+        connection.dirty = True
     else:
-        raise TransactionManagementError("This code isn't under transaction management")
+        raise TransactionManagementError("This code isn't under transaction "
+            "management")
 
 def set_clean(using=None):
     """
@@ -120,9 +113,10 @@ def set_clean(using=None):
     """
     if using is None:
         using = DEFAULT_DB_ALIAS
-    thread_ident = thread.get_ident()
-    if thread_ident in dirty and using in dirty[thread_ident]:
-        dirty[thread_ident][using] = False
+    connection = connections[using]
+
+    if connection.dirty is not None:
+        connection.dirty = False
     else:
         raise TransactionManagementError("This code isn't under transaction management")
     clean_savepoints(using=using)
@@ -130,9 +124,8 @@ def set_clean(using=None):
 def clean_savepoints(using=None):
     if using is None:
         using = DEFAULT_DB_ALIAS
-    thread_ident = thread.get_ident()
-    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
-        del savepoint_state[thread_ident][using]
+    connection = connections[using]
+    connection.savepoint_state = 0
 
 def is_managed(using=None):
     """
@@ -140,10 +133,9 @@ def is_managed(using=None):
     """
     if using is None:
         using = DEFAULT_DB_ALIAS
-    thread_ident = thread.get_ident()
-    if thread_ident in state and using in state[thread_ident]:
-        if state[thread_ident][using]:
-            return state[thread_ident][using][-1]
+    connection = connections[using]
+    if connection.transaction_state:
+        return connection.transaction_state[-1]
     return settings.TRANSACTIONS_MANAGED
 
 def managed(flag=True, using=None):
@@ -156,15 +148,16 @@ def managed(flag=True, using=None):
     if using is None:
         using = DEFAULT_DB_ALIAS
     connection = connections[using]
-    thread_ident = thread.get_ident()
-    top = state.get(thread_ident, {}).get(using, None)
+
+    top = connection.transaction_state
     if top:
         top[-1] = flag
         if not flag and is_dirty(using=using):
             connection._commit()
             set_clean(using=using)
     else:
-        raise TransactionManagementError("This code isn't under transaction management")
+        raise TransactionManagementError("This code isn't under transaction "
+            "management")
 
 def commit_unless_managed(using=None):
     """
@@ -221,13 +214,11 @@ def savepoint(using=None):
         using = DEFAULT_DB_ALIAS
     connection = connections[using]
     thread_ident = thread.get_ident()
-    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
-        savepoint_state[thread_ident][using].append(None)
-    else:
-        savepoint_state.setdefault(thread_ident, {})
-        savepoint_state[thread_ident][using] = [None]
+
+    connection.savepoint_state += 1
+
     tid = str(thread_ident).replace('-', '')
-    sid = "s%s_x%d" % (tid, len(savepoint_state[thread_ident][using]))
+    sid = "s%s_x%d" % (tid, connection.savepoint_state)
     connection._savepoint(sid)
     return sid
 
@@ -239,8 +230,8 @@ def savepoint_rollback(sid, using=None):
     if using is None:
         using = DEFAULT_DB_ALIAS
     connection = connections[using]
-    thread_ident = thread.get_ident()
-    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
+
+    if connection.savepoint_state:
         connection._savepoint_rollback(sid)
 
 def savepoint_commit(sid, using=None):
@@ -251,8 +242,8 @@ def savepoint_commit(sid, using=None):
     if using is None:
         using = DEFAULT_DB_ALIAS
     connection = connections[using]
-    thread_ident = thread.get_ident()
-    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
+
+    if connection.savepoint_state:
         connection._savepoint_commit(sid)
 
 ##############