Fixed #21239 -- Maintained atomicity when closing the connection.

Refs #15802 -- Reverted #7c657b24 as BaseDatabaseWrapper.close() now
has a proper "finally" clause that may need to preserve self.connection.

django/db/backends/__init__.py

@@ -59,6 +59,7 @@ class BaseDatabaseWrapper(object):
 
         # Connection termination related attributes.
         self.close_at = None
+        self.closed_in_transaction = False
         self.errors_occurred = False
 
         # Thread-safety related attributes.
@@ -101,9 +102,11 @@ class BaseDatabaseWrapper(object):
         # 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()
@@ -183,6 +186,10 @@ class BaseDatabaseWrapper(object):
         try:
             self._close()
         finally:
+            if self.in_atomic_block:
+                self.closed_in_transaction = True
+                self.needs_rollback = True
+            else:
                 self.connection = None
 
     ##### Backend-specific savepoint management methods #####

django/db/backends/postgresql_psycopg2/base.py

@@ -3,7 +3,7 @@ PostgreSQL database backend for Django.
 
 Requires psycopg 2: http://initd.org/projects/psycopg2
 """
-import logging
+
 import sys
 
 from django.conf import settings
@@ -36,8 +36,6 @@ psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
 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:
@@ -161,26 +159,6 @@ class DatabaseWrapper(BaseDatabaseWrapper):
         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
-
     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):

django/db/transaction.py

@@ -205,7 +205,12 @@ class Atomic(object):
             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:
@@ -245,12 +250,17 @@ class Atomic(object):
         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:
+                if connection.closed_in_transaction:
+                    connection.connection = None
+                else:
                     connection.in_atomic_block = False
 
     def __call__(self, func):

tests/transactions/tests.py

@@ -9,8 +9,8 @@ import time
 from unittest import skipIf, skipUnless
 
 from django.db import (connection, transaction,
-    DatabaseError, IntegrityError, OperationalError)
-from django.test import TransactionTestCase, skipIfDBFeature
+    DatabaseError, Error, IntegrityError, OperationalError)
+from django.test import TransactionTestCase, skipIfDBFeature, skipUnlessDBFeature
 from django.utils import six
 
 from .models import Reporter
@@ -338,6 +338,20 @@ class AtomicErrorsTests(TransactionTestCase):
             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):