Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with
or
.
Download ZIP
Browse files

Add RunPython migration operation and tests

  • Loading branch information...
commit 3b810c5656774b4e1a89fd69ff26935e0a5833d2 1 parent 05656f2
@andrewgodwin andrewgodwin authored
View
6 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 @@ def unapply(self, project_state, schema_editor):
# 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))
View
2  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
View
3  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)
View
42 django/db/migrations/operations/special.py
@@ -1,5 +1,5 @@
import re
-
+import textwrap
from .base import Operation
@@ -59,6 +59,10 @@ def __init__(self, sql, reverse_sql=None, state_operations=None, multiple=False)
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 @@ def database_backwards(self, app_label, schema_editor, from_state, to_state):
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")
View
29 tests/migrations/test_operations.py
@@ -282,7 +282,7 @@ def test_alter_index_together(self):
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 @@ def test_run_sql(self):
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):
"""
Please sign in to comment.
Something went wrong with that request. Please try again.