Add RunPython migration operation and tests

This commit is contained in:
Andrew Godwin 2013-09-25 13:58:07 +01:00
parent 05656f2388
commit 3b810c5656
5 changed files with 79 additions and 3 deletions

View File

@ -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))

View File

@ -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 File

@ -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 File

@ -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")

View File

@ -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):
"""