[1.7.x] Fixed #22496: Data migrations get transactions again!

5 files changed, 65 insertions, 4 deletions

diff --git a/django/db/migrations/migration.py b/django/db/migrations/migration.py
index 3d7f47f0fa..708de5393d 100644
--- a/django/db/migrations/migration.py
+++ b/django/db/migrations/migration.py
@@ -1,4 +1,5 @@
 from __future__ import unicode_literals
+from django.db.transaction import atomic
 
 
 class Migration(object):
@@ -97,7 +98,13 @@ class Migration(object):
             new_state = project_state.clone()
             operation.state_forwards(self.app_label, new_state)
             # Run the operation
-            operation.database_forwards(self.app_label, schema_editor, project_state, new_state)
+            if not schema_editor.connection.features.can_rollback_ddl and operation.atomic:
+                # We're forcing a transaction on a non-transactional-DDL backend
+                with atomic(schema_editor.connection.alias):
+                    operation.database_forwards(self.app_label, schema_editor, project_state, new_state)
+            else:
+                # Normal behaviour
+                operation.database_forwards(self.app_label, schema_editor, project_state, new_state)
             # Switch states
             project_state = new_state
         return project_state
@@ -129,7 +136,13 @@ class Migration(object):
         # Now run them in reverse
         to_run.reverse()
         for operation, to_state, from_state in to_run:
-            operation.database_backwards(self.app_label, schema_editor, from_state, to_state)
+            if not schema_editor.connection.features.can_rollback_ddl and operation.atomic:
+                # We're forcing a transaction on a non-transactional-DDL backend
+                with atomic(schema_editor.connection.alias):
+                    operation.database_backwards(self.app_label, schema_editor, from_state, to_state)
+            else:
+                # Normal behaviour
+                operation.database_backwards(self.app_label, schema_editor, from_state, to_state)
         return project_state
 
 
diff --git a/django/db/migrations/operations/base.py b/django/db/migrations/operations/base.py
index 7e93fb7836..d0ab0a39ad 100644
--- a/django/db/migrations/operations/base.py
+++ b/django/db/migrations/operations/base.py
@@ -24,6 +24,10 @@ class Operation(object):
     # Can this migration be represented as SQL? (things like RunPython cannot)
     reduces_to_sql = True
 
+    # Should this operation be forced as atomic even on backends with no
+    # DDL transaction support (i.e., does it have no DDL, like RunPython)
+    atomic = False
+
     serialization_expand_args = []
 
     def __new__(cls, *args, **kwargs):
diff --git a/django/db/migrations/operations/special.py b/django/db/migrations/operations/special.py
index a80096cb32..b7ffbbbd29 100644
--- a/django/db/migrations/operations/special.py
+++ b/django/db/migrations/operations/special.py
@@ -86,7 +86,8 @@ class RunPython(Operation):
 
     reduces_to_sql = False
 
-    def __init__(self, code, reverse_code=None):
+    def __init__(self, code, reverse_code=None, atomic=True):
+        self.atomic = atomic
         # Forwards code
         if not callable(code):
             raise ValueError("RunPython must be supplied with a callable")
diff --git a/docs/ref/migration-operations.txt b/docs/ref/migration-operations.txt
index 0b7e6a3912..9c7a5d497e 100644
--- a/docs/ref/migration-operations.txt
+++ b/docs/ref/migration-operations.txt
@@ -178,7 +178,7 @@ operation that adds that field and so will try to run it again).
 RunPython
 ---------
 
-.. class:: RunPython(code, reverse_code=None)
+.. class:: RunPython(code, reverse_code=None, atomic=True)
 
 Runs custom Python code in a historical context. ``code`` (and ``reverse_code``
 if supplied) should be callable objects that accept two arguments; the first is
@@ -230,6 +230,12 @@ or that you use :class:`SeparateDatabaseAndState` to add in operations that will
 reflect your changes to the model state - otherwise, the versioned ORM and
 the autodetector will stop working correctly.
 
+By default, ``RunPython`` will run its contents inside a transaction even
+on databases that do not support DDL transactions (for example, MySQL and
+Oracle). This should be safe, but may cause a crash if you attempt to use
+the ``schema_editor`` provided on these backends; in this case, please
+set ``atomic=False``.
+
 SeparateDatabaseAndState
 ------------------------
 
diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index d294bdcd26..32319f373c 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -752,6 +752,43 @@ class OperationTests(MigrationTestBase):
         self.assertEqual(project_state.render().get_model("test_runpython", "Pony").objects.count(), 6)
         self.assertEqual(project_state.render().get_model("test_runpython", "ShetlandPony").objects.count(), 2)
 
+    def test_run_python_atomic(self):
+        """
+        Tests the RunPython operation correctly handles the "atomic" keyword
+        """
+        project_state = self.set_up_test_model("test_runpythonatomic", mti_model=True)
+        def inner_method(models, schema_editor):
+            Pony = models.get_model("test_runpythonatomic", "Pony")
+            Pony.objects.create(pink=1, weight=3.55)
+            raise ValueError("Adrian hates ponies.")
+        atomic_migration = Migration("test", "test_runpythonatomic")
+        atomic_migration.operations = [migrations.RunPython(inner_method)]
+        non_atomic_migration = Migration("test", "test_runpythonatomic")
+        non_atomic_migration.operations = [migrations.RunPython(inner_method, atomic=False)]
+        # If we're a fully-transactional database, both versions should rollback
+        if connection.features.can_rollback_ddl:
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    atomic_migration.apply(project_state, editor)
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    non_atomic_migration.apply(project_state, editor)
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+        # Otherwise, the non-atomic operation should leave a row there
+        else:
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    atomic_migration.apply(project_state, editor)
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    non_atomic_migration.apply(project_state, editor)
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 1)
+
+
 
 class MigrateNothingRouter(object):
     """