Commit 3b810c56 authored by Andrew Godwin's avatar Andrew Godwin
Browse files

Add RunPython migration operation and tests

parent 05656f23

django/db/migrations/migration.py

+6 −0
Original line number Diff line number Diff line
@@ -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 −1
Original line number Diff line number Diff line 
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

+3 −0
Original line number Diff line number Diff line
@@ -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

+41 −1
Original line number Diff line number Diff line 
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

+28 −1
Original line number Diff line number Diff line
@@ -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):

    """