Ticket #17258: 17258.thread-local-connections.3.diff

django/db/__init__.py

@@ +1 @@
+from threading import local
 from django.conf import settings
 from django.core import signals
 from django.core.exceptions import ImproperlyConfigured
-from django.db.utils import (ConnectionHandler, ConnectionRouter,
-    load_backend, DEFAULT_DB_ALIAS, DatabaseError, IntegrityError)
+from django.db.utils import (ConnectionHandler, ConnectionRouter, load_backend,
+    DEFAULT_DB_ALIAS, DatabaseError, IntegrityError, SharedConnectionError)
 
 __all__ = ('backend', 'connection', 'connections', 'router', 'DatabaseError',
     'IntegrityError', 'DEFAULT_DB_ALIAS')
@@ -22 +23 @@
 # we manually create the dictionary from the settings, passing only the
 # settings that the database backends care about. Note that TIME_ZONE is used
 # by the PostgreSQL backends.
-# we load all these up for backwards compatibility, you should use
+# We load all these up for backwards compatibility, you should use
 # connections['default'] instead.
-connection = connections[DEFAULT_DB_ALIAS]
+class DefaultConnectionProxy(object):
+    """
+    Proxy for the thread-local default connection.
+    """
+    def __getattr__(self, item):
+        return getattr(connections[DEFAULT_DB_ALIAS], item)
+
+    def __setattr__(self, name, value):
+        return setattr(connections[DEFAULT_DB_ALIAS], name, value)
+
+connection = DefaultConnectionProxy()
 backend = load_backend(connection.settings_dict['ENGINE'])
 
 # Register an event that closes the database connection

django/db/backends/__init__.py

@@ -2 +2 @@
     import thread
 except ImportError:
     import dummy_thread as thread
-from threading import local
 from contextlib import contextmanager
 
 from django.conf import settings
-from django.db import DEFAULT_DB_ALIAS
+from django.db import DEFAULT_DB_ALIAS, SharedConnectionError
 from django.db.backends import util
 from django.db.transaction import TransactionManagementError
 from django.utils.importlib import import_module
 from django.utils.timezone import is_aware
 
 
-class BaseDatabaseWrapper(local):
+class BaseDatabaseWrapper(object):
     """
     Represents a database connection.
     """
@@ -34 +33 @@
         self.transaction_state = []
         self.savepoint_state = 0
         self._dirty = None
+        self.allow_thread_sharing = False
+        self._thread_ident = None
 
     def __eq__(self, other):
         return self.alias == other.alias
@@ -41 +42 @@
     def __ne__(self, other):
         return not self == other
 
+    def _check_thread_sharing(self):
+        if self.allow_thread_sharing:
+            return
+        if self._thread_ident is None:
+            self._thread_ident = thread.get_ident()
+            return
+        if self._thread_ident <> thread.get_ident():
+            raise SharedConnectionError(
+                "This connection seems to be shared between threads. This "
+                "is not allowed.")
+
     def _commit(self):
         if self.connection is not None:
             return self.connection.commit()

django/db/backends/postgresql_psycopg2/base.py

@@ -153 +153 @@
     pg_version = property(_get_pg_version)
 
     def _cursor(self):
+        self._check_thread_sharing()
         settings_dict = self.settings_dict
         if self.connection is None:
             if settings_dict['NAME'] == '':

django/db/backends/sqlite3/base.py

@@ -231 +231 @@
         self.validation = BaseDatabaseValidation(self)
 
     def _cursor(self):
+        self._check_thread_sharing()
         if self.connection is None:
             settings_dict = self.settings_dict
             if not settings_dict['NAME']:

django/db/utils.py

@@ -1 +1 @@
 import os
+from threading import local
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
@@ -46 +47 @@
 class ConnectionDoesNotExist(Exception):
     pass
 
+class SharedConnectionError(Exception):
+    pass
 
 class ConnectionHandler(object):
     def __init__(self, databases):
         self.databases = databases
-        self._connections = {}
+        self._connections = local()
 
     def ensure_defaults(self, alias):
         """
@@ -73 +76 @@
             conn.setdefault(setting, None)
 
     def __getitem__(self, alias):
-        if alias in self._connections:
-            return self._connections[alias]
+        if hasattr(self._connections, alias):
+            return getattr(self._connections, alias)
 
         self.ensure_defaults(alias)
         db = self.databases[alias]
         backend = load_backend(db['ENGINE'])
         conn = backend.DatabaseWrapper(db, alias)
-        self._connections[alias] = conn
+        setattr(self._connections, alias, conn)
         return conn
 
+    def __setitem__(self, key, value):
+        setattr(self._connections, key, value)
+
     def __iter__(self):
         return iter(self.databases)
 

tests/regressiontests/backends/tests.py

@@ -3 +3 @@
 from __future__ import with_statement, absolute_import
 
 import datetime
+import threading
 
 from django.conf import settings
 from django.core.management.color import no_style
 from django.db import (backend, connection, connections, DEFAULT_DB_ALIAS,
-    IntegrityError, transaction)
+    IntegrityError, transaction, SharedConnectionError)
 from django.db.backends.signals import connection_created
 from django.db.backends.postgresql_psycopg2 import version as pg_version
 from django.db.utils import ConnectionHandler, DatabaseError
@@ -283 +284 @@
         connection_created.connect(receiver)
         connection.close()
         cursor = connection.cursor()
-        self.assertTrue(data["connection"] is connection)
+        self.assertTrue(data["connection"].connection is connection.connection)
 
         connection_created.disconnect(receiver)
         data.clear()
@@ -446 +447 @@
                         connection.check_constraints()
             finally:
                 transaction.rollback()
+
+class ThreadTests(TestCase):
+
+    @unittest.skipIf(connection.vendor == 'sqlite',
+                     "SQLite doesn't allow connection sharing between threads")
+    def test_default_connection_thread_local(self):
+        """
+        Ensure that the default connection (i.e. django.db.connection) is
+        different for each thread.
+        Refs #17258.
+        """
+        connections_set = set()
+        connection.cursor()
+        connections_set.add(connection.connection)
+        def runner():
+            from django.db import connection
+            connection.cursor()
+            connections_set.add(connection.connection)
+        for x in xrange(2):
+            t = threading.Thread(target=runner)
+            t.start()
+            t.join()
+        self.assertEquals(len(connections_set), 3)
+        # Finish by closing the connections opened by the other threads (the
+        # connection opened in the main thread will automatically be closed on
+        # teardown).
+        for conn in connections_set:
+            if conn != connection.connection:
+                conn.close()
+
+    @unittest.skipIf(connection.vendor == 'sqlite',
+                     "SQLite doesn't allow connection sharing between threads")
+    def test_connections_thread_local(self):
+        """
+        Ensure that the connections are different for each thread.
+        Refs #17258.
+        """
+        connections_set = set()
+        for conn in connections.all():
+            connections_set.add(conn)
+        def runner():
+            from django.db import connections
+            for conn in connections.all():
+                connections_set.add(conn)
+        for x in xrange(2):
+            t = threading.Thread(target=runner)
+            t.start()
+            t.join()
+        self.assertEquals(len(connections_set), 6)
+        # Finish by closing the connections opened by the other threads (the
+        # connection opened in the main thread will automatically be closed on
+        # teardown).
+        for conn in connections_set:
+            if conn != connection:
+                conn.close()
+
+    @unittest.skipIf(connection.vendor == 'sqlite',
+                     "SQLite doesn't allow connection sharing between threads")
+    def test_connection_passing(self):
+        """
+        Using the same connection in two separate threads is not allowed
+        unless explicitly requested by the user. Refs #17258.
+        """
+        errors = []
+        def runner(connection):
+            # Usage is defined by taking a cursor.
+            try:
+                connection.cursor()
+                connection.close()
+            except SharedConnectionError, e:
+                errors.append(str(e))
+
+        connection = connections[DEFAULT_DB_ALIAS]
+        # Make sure the connection is claimed by this thread.
+        connection.cursor()
+        for x in xrange(2):
+            t = threading.Thread(target=runner, args=(connection,))
+            t.start()
+            t.join()
+        self.assertEquals(len(errors), 2)
+
+        errors = []
+        connection.allow_thread_sharing = True
+        for x in xrange(2):
+            t = threading.Thread(target=runner, args=(connection,))
+            t.start()
+            t.join()
+        self.assertEquals(len(errors), 0)