[1.7.x] Fixed #21239 -- Maintained atomicity when closing the connection.

aaugustin · aaugustin · commit 2e42c859da42 · 2014-04-10T23:18:35.000+02:00
Refs #15802 -- Reverted #7c657b24 as BaseDatabaseWrapper.close() now has a proper "finally" clause that may need to preserve self.connection. Backport of 2586009 from master.
diff --git a/django/db/backends/__init__.py b/django/db/backends/__init__.py
@@ -62,6 +62,7 @@ def __init__(self, settings_dict, alias=DEFAULT_DB_ALIAS,
 
         # Connection termination related attributes
         self.close_at = None
+        self.closed_in_transaction = False
         self.errors_occurred = False
 
         # Thread-safety related attributes
@@ -104,9 +105,11 @@ def connect(self):
         # In case the previous connection was closed while in an atomic block
         self.in_atomic_block = False
         self.savepoint_ids = []
+        self.needs_rollback = False
         # Reset parameters defining when to close the connection
         max_age = self.settings_dict['CONN_MAX_AGE']
         self.close_at = None if max_age is None else time.time() + max_age
+        self.closed_in_transaction = False
         self.errors_occurred = False
         # Establish the connection
         conn_params = self.get_connection_params()
@@ -188,7 +191,11 @@ def close(self):
         try:
             self._close()
         finally:
-            self.connection = None
+            if self.in_atomic_block:
+                self.closed_in_transaction = True
+                self.needs_rollback = True
+            else:
+                self.connection = None
         self.set_clean()
 
     ##### Backend-specific savepoint management methods #####
diff --git a/django/db/backends/postgresql_psycopg2/base.py b/django/db/backends/postgresql_psycopg2/base.py
@@ -3,7 +3,7 @@
 
 Requires psycopg 2: http://initd.org/projects/psycopg2
 """
-import logging
+
 import sys
 
 from django.conf import settings
@@ -36,8 +36,6 @@
 psycopg2.extensions.register_adapter(SafeBytes, psycopg2.extensions.QuotedString)
 psycopg2.extensions.register_adapter(SafeText, psycopg2.extensions.QuotedString)
 
-logger = logging.getLogger('django.db.backends')
-
 
 def utc_tzinfo_factory(offset):
     if offset != 0:
@@ -162,28 +160,6 @@ def create_cursor(self):
         cursor.tzinfo_factory = utc_tzinfo_factory if settings.USE_TZ else None
         return cursor
 
-    def close(self):
-        self.validate_thread_sharing()
-        if self.connection is None:
-            return
-
-        try:
-            self.connection.close()
-            self.connection = None
-        except Database.Error:
-            # In some cases (database restart, network connection lost etc...)
-            # the connection to the database is lost without giving Django a
-            # notification. If we don't set self.connection to None, the error
-            # will occur a every request.
-            self.connection = None
-            logger.warning(
-                'psycopg2 error while closing the connection.',
-                exc_info=sys.exc_info()
-            )
-            raise
-        finally:
-            self.set_clean()
-
     def _set_isolation_level(self, isolation_level):
         assert isolation_level in range(1, 5)     # Use set_autocommit for level = 0
         if self.psycopg2_version >= (2, 4, 2):
diff --git a/django/db/transaction.py b/django/db/transaction.py
@@ -313,7 +313,12 @@ def __exit__(self, exc_type, exc_value, traceback):
             connection.in_atomic_block = False
 
         try:
-            if exc_type is None and not connection.needs_rollback:
+            if connection.closed_in_transaction:
+                # The database will perform a rollback by itself.
+                # Wait until we exit the outermost block.
+                pass
+
+            elif exc_type is None and not connection.needs_rollback:
                 if connection.in_atomic_block:
                     # Release savepoint if there is one
                     if sid is not None:
@@ -353,13 +358,18 @@ def __exit__(self, exc_type, exc_value, traceback):
         finally:
             # Outermost block exit when autocommit was enabled.
             if not connection.in_atomic_block:
-                if connection.features.autocommits_when_autocommit_is_off:
+                if connection.closed_in_transaction:
+                    connection.connection = None
+                elif connection.features.autocommits_when_autocommit_is_off:
                     connection.autocommit = True
                 else:
                     connection.set_autocommit(True)
             # Outermost block exit when autocommit was disabled.
             elif not connection.savepoint_ids and not connection.commit_on_exit:
-                connection.in_atomic_block = False
+                if connection.closed_in_transaction:
+                    connection.connection = None
+                else:
+                    connection.in_atomic_block = False
 
     def __call__(self, func):
         @wraps(func, assigned=available_attrs(func))
diff --git a/tests/transactions/tests.py b/tests/transactions/tests.py
@@ -9,7 +9,7 @@
 from unittest import skipIf, skipUnless
 
 from django.db import (connection, transaction,
-    DatabaseError, IntegrityError, OperationalError)
+    DatabaseError, Error, IntegrityError, OperationalError)
 from django.test import TransactionTestCase, skipIfDBFeature, skipUnlessDBFeature
 from django.test.utils import IgnoreDeprecationWarningsMixin
 from django.utils import six
@@ -360,6 +360,20 @@ def test_atomic_allows_queries_after_fixing_transaction(self):
             r2.save(force_update=True)
         self.assertEqual(Reporter.objects.get(pk=r1.pk).last_name, "Calculus")
 
+    @skipUnlessDBFeature('test_db_allows_multiple_connections')
+    def test_atomic_prevents_queries_in_broken_transaction_after_client_close(self):
+        with transaction.atomic():
+            Reporter.objects.create(first_name="Archibald", last_name="Haddock")
+            connection.close()
+            # The connection is closed and the transaction is marked as
+            # needing rollback. This will raise an InterfaceError on databases
+            # that refuse to create cursors on closed connections (PostgreSQL)
+            # and a TransactionManagementError on other databases.
+            with self.assertRaises(Error):
+                Reporter.objects.create(first_name="Cuthbert", last_name="Calculus")
+        # The connection is usable again .
+        self.assertEqual(Reporter.objects.count(), 0)
+
 
 @skipUnless(connection.vendor == 'mysql', "MySQL-specific behaviors")
 class AtomicMySQLTests(TransactionTestCase):
