Add RunPython migration operation and tests

django/db/migrations/migration.py

@@ -32,6 +32,10 @@ class Migration(object):
     # are not applied.
     replaces = []
 
+    # Error class which is raised when a migration is irreversible
+    class IrreversibleError(RuntimeError):
+        pass
+
     def __init__(self, name, app_label):
         self.name = name
         self.app_label = app_label
@@ -91,6 +95,8 @@ class Migration(object):
         # We need to pre-calculate the stack of project states
         to_run = []
         for operation in self.operations:
+            if not operation.reversible:
+                raise Migration.IrreversibleError("Operation %s in %s is not reversible" % (operation, sekf))
             new_state = project_state.clone()
             operation.state_forwards(self.app_label, new_state)
             to_run.append((operation, project_state, new_state))

django/db/migrations/operations/__init__.py

@@ -1,3 +1,3 @@
 from .models import CreateModel, DeleteModel, AlterModelTable, AlterUniqueTogether, AlterIndexTogether
 from .fields import AddField, RemoveField, AlterField, RenameField
-from .special import SeparateDatabaseAndState, RunSQL
+from .special import SeparateDatabaseAndState, RunSQL, RunPython

django/db/migrations/operations/base.py

@@ -15,6 +15,9 @@ class Operation(object):
     # Some operations are impossible to reverse, like deleting data.
     reversible = True
 
+    # Can this migration be represented as SQL? (things like RunPython cannot)
+    reduces_to_sql = True
+
     def __new__(cls, *args, **kwargs):
         # We capture the arguments to make returning them trivial
         self = object.__new__(cls)

django/db/migrations/operations/special.py

@@ -1,5 +1,5 @@
 import re
-
+import textwrap
 from .base import Operation
 
 
@@ -59,6 +59,10 @@ class RunSQL(Operation):
         self.state_operations = state_operations or []
         self.multiple = multiple
 
+    @property
+    def reversible(self):
+        return self.reverse_sql is not None
+
     def state_forwards(self, app_label, state):
         for state_operation in self.state_operations:
             state_operation.state_forwards(app_label, state)
@@ -92,3 +96,39 @@ class RunSQL(Operation):
 
     def describe(self):
         return "Raw SQL operation"
+
+
+class RunPython(Operation):
+    """
+    Runs Python code in a context suitable for doing versioned ORM operations.
+    """
+
+    reduces_to_sql = False
+    reversible = False
+
+    def __init__(self, code):
+        # Trim any leading whitespace that is at the start of all code lines
+        # so users can nicely indent code in migration files
+        code = textwrap.dedent(code)
+        # Run the code through a parser first to make sure it's at least
+        # syntactically correct
+        self.code = compile(code, "<string>", "exec")
+
+    def state_forwards(self, app_label, state):
+        # RunPython objects have no state effect. To add some, combine this
+        # with SeparateDatabaseAndState.
+        pass
+
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
+        # We now execute the Python code in a context that contains a 'models'
+        # object, representing the versioned models as an AppCache.
+        # We could try to override the global cache, but then people will still
+        # use direct imports, so we go with a documentation approach instead.
+        context = {
+            "models": from_state.render(),
+            "schema_editor": schema_editor,
+        }
+        eval(self.code, context)
+
+    def database_backwards(self, app_label, schema_editor, from_state, to_state):
+        raise NotImplementedError("You cannot reverse this operation")

tests/migrations/test_operations.py

@@ -282,7 +282,7 @@ class OperationTests(MigrationTestBase):
 
     def test_run_sql(self):
         """
-        Tests the AlterIndexTogether operation.
+        Tests the RunSQL operation.
         """
         project_state = self.set_up_test_model("test_runsql")
         # Create the operation
@@ -306,6 +306,33 @@ class OperationTests(MigrationTestBase):
             operation.database_backwards("test_runsql", editor, new_state, project_state)
         self.assertTableNotExists("i_love_ponies")
 
+    def test_run_python(self):
+        """
+        Tests the RunPython operation
+        """
+
+        project_state = self.set_up_test_model("test_runpython")
+        # Create the operation
+        operation = migrations.RunPython(
+            """
+            Pony = models.get_model("test_runpython", "Pony")
+            Pony.objects.create(pink=2, weight=4.55)
+            Pony.objects.create(weight=1)
+            """,
+        )
+        # Test the state alteration does nothing
+        new_state = project_state.clone()
+        operation.state_forwards("test_runpython", new_state)
+        self.assertEqual(new_state, project_state)
+        # Test the database alteration
+        self.assertEqual(project_state.render().get_model("test_runpython", "Pony").objects.count(), 0)
+        with connection.schema_editor() as editor:
+            operation.database_forwards("test_runpython", editor, project_state, new_state)
+        self.assertEqual(project_state.render().get_model("test_runpython", "Pony").objects.count(), 2)
+        # And test reversal fails
+        with self.assertRaises(NotImplementedError):
+            operation.database_backwards("test_runpython", None, new_state, project_state)
+
 
 class MigrateNothingRouter(object):
     """