BaseDatabaseSchemaEditor.__exit__ doesn't exit its atomic block when deferred SQL raises

[Gavin Wahl(vendor)]

BaseDatabaseSchemaEditor (normally accessed through connection.schema_editor) implements the context manager protocol and wraps transaction.atomic() in its __enter__ and __exit__. In __exit__, it runs the deferred SQL and then calls self.atomic.__exit__. If the deferred SQL raises, the second call never happens and the atomic block stays open.

After that the connection is stuck. in_atomic_block is True with no atomic block on the stack, so nothing will ever commit. Once anything sets needs_rollback you get TransactionManagementError on every query. close() doesn't help because inside an atomic block it keeps the dead connection object around and waits for an Atomic.__exit__ to reset it.

This doesn't affect manage.py migrate from a shell, since the process exits and the server rolls back. It matters if you catch the error and keep going, e.g. tests, call_command("migrate"), or a tool using schema_editor directly.

There is an existing test_migrations_not_applied_on_deferred_sql_failure but it has atomic = False so it doesn't hit this.

Minimal reproducer attached.

Related #29721

[Sarah Boyce]

Thank you, reproduced:

tests/schema/tests.py

@@ -6078 +6078 @@
             if connection.vendor == "postgresql":
                 with connection.cursor() as cursor:
                     cursor.execute("DROP COLLATION IF EXISTS case_insensitive")
+
+
+class SchemaEditorTransactionTests(TransactionTestCase):
+    available_apps = []
+
+    @skipUnlessDBFeature("can_rollback_ddl")
+    def test_deferred_sql_failure_exits_atomic_block(self):
+        editor = connection.schema_editor()
+        try:
+            with self.assertRaises(DatabaseError):
+                with editor:
+                    editor.deferred_sql.append(
+                        "SELECT * FROM django_nonexistent_table"
+                    )
+
+            self.assertFalse(connection.in_atomic_block)
+        finally:
+            # Keep the test database usable while the bug is present.
+            if editor.atomic in connection.atomic_blocks:
+                editor.atomic.__exit__(DatabaseError, DatabaseError(), None)