Python django.db.migrations.RunPython() Examples

def get_migrations_for_django_21_and_newer():
    return [
        # remove primary key information from 'key' field
        migrations.AlterField(
            model_name='resetpasswordtoken',
            name='key',
            field=models.CharField(db_index=True, primary_key=False, max_length=64, unique=True, verbose_name='Key'),
        ),
        # add a new id field
        migrations.AddField(
            model_name='resetpasswordtoken',
            name='id',
            field=models.AutoField(primary_key=True, serialize=False),
            preserve_default=False,
        ),
        migrations.RunPython(
            populate_auto_incrementing_pk_field,
            migrations.RunPython.noop
        ),
    ] 

def remove_duplicates(apps, schema_editor):
    Reference = apps.get_model('mangaki', 'Reference')
    db_alias = schema_editor.connection.alias

    unique_fields = ['work', 'source', 'identifier']

    duplicates = (Reference.objects.using(db_alias)
                                   .values(*unique_fields)
                                   .order_by()
                                   .annotate(max_id=models.Max('id'),
                                             count_id=models.Count('id'))
                                   .filter(count_id__gt=1)
                                   .iterator())

    for duplicate in duplicates:
        (Reference.objects.using(db_alias)
                          .filter(**{x: duplicate[x] for x in unique_fields})
                          .exclude(id=duplicate['max_id'])
                          .delete())

# migrations.RunPython.noop cause circular reference error… 

def add_project_member(apps, schema_editor):
    # Using historical versions as recommended for RunPython
    PublicDataAccess = apps.get_model("public_data", "PublicDataAccess")
    DataRequestProjectMember = apps.get_model(
        "private_sharing", "DataRequestProjectMember"
    )
    DataRequestProject = apps.get_model("private_sharing", "DataRequestProject")
    db_alias = schema_editor.connection.alias

    def id_label_to_project(id_label):
        match = re.match(r"direct-sharing-(?P<id>\d+)", id_label)
        if match:
            project = DataRequestProject.objects.using(db_alias).get(
                id=int(match.group("id"))
            )
            return project

    for pda in PublicDataAccess.objects.using(db_alias).filter(project_membership=None):
        project = id_label_to_project(pda.data_source)
        drpm = DataRequestProjectMember.objects.using(db_alias).get(
            project=project, member=pda.participant.member
        )
        pda.project_membership = drpm
        pda.save() 

def _test_run_python(self, app_label, should_run, hints=None):
        with override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter()]):
            project_state = self.set_up_test_model(app_label)

        # Create the operation
        def inner_method(models, schema_editor):
            Pony = models.get_model(app_label, "Pony")
            Pony.objects.create(pink=1, weight=3.55)
            Pony.objects.create(weight=5)

        operation = migrations.RunPython(inner_method, hints=hints or {})
        # Test the state alteration does nothing
        new_state = project_state.clone()
        operation.state_forwards(app_label, new_state)
        self.assertEqual(new_state, project_state)
        # Test the database alteration
        self.assertEqual(project_state.apps.get_model(app_label, "Pony").objects.count(), 0)
        with connection.schema_editor() as editor:
            operation.database_forwards(app_label, editor, project_state, new_state)
        Pony = project_state.apps.get_model(app_label, "Pony")
        if should_run:
            self.assertEqual(Pony.objects.count(), 2)
        else:
            self.assertEqual(Pony.objects.count(), 0) 

def analyse_data_migration(self, migration):
        errors = []
        ignored = []
        warnings = []

        for operation in migration.operations:
            if isinstance(operation, RunPython):
                op_errors, op_ignored, op_warnings = self.lint_runpython(operation)
                if op_errors:
                    errors += op_errors
                if op_ignored:
                    ignored += op_ignored
                if op_warnings:
                    warnings += op_warnings

        return errors, ignored, warnings 

def _test_run_python(self, app_label, should_run, hints=None):
        with override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter()]):
            project_state = self.set_up_test_model(app_label)

        # Create the operation
        def inner_method(models, schema_editor):
            Pony = models.get_model(app_label, "Pony")
            Pony.objects.create(pink=1, weight=3.55)
            Pony.objects.create(weight=5)

        operation = migrations.RunPython(inner_method, hints=hints or {})
        # Test the state alteration does nothing
        new_state = project_state.clone()
        operation.state_forwards(app_label, new_state)
        self.assertEqual(new_state, project_state)
        # Test the database alteration
        self.assertEqual(project_state.apps.get_model(app_label, "Pony").objects.count(), 0)
        with connection.schema_editor() as editor:
            operation.database_forwards(app_label, editor, project_state, new_state)
        Pony = project_state.apps.get_model(app_label, "Pony")
        if should_run:
            self.assertEqual(Pony.objects.count(), 2)
        else:
            self.assertEqual(Pony.objects.count(), 0) 

def _test_run_python(self, app_label, should_run, hints=None):
        with override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter()]):
            project_state = self.set_up_test_model(app_label)

        # Create the operation
        def inner_method(models, schema_editor):
            Pony = models.get_model(app_label, "Pony")
            Pony.objects.create(pink=1, weight=3.55)
            Pony.objects.create(weight=5)

        operation = migrations.RunPython(inner_method, hints=hints or {})
        # Test the state alteration does nothing
        new_state = project_state.clone()
        operation.state_forwards(app_label, new_state)
        self.assertEqual(new_state, project_state)
        # Test the database alteration
        self.assertEqual(project_state.apps.get_model(app_label, "Pony").objects.count(), 0)
        with connection.schema_editor() as editor:
            operation.database_forwards(app_label, editor, project_state, new_state)
        Pony = project_state.apps.get_model(app_label, "Pony")
        if should_run:
            self.assertEqual(Pony.objects.count(), 2)
        else:
            self.assertEqual(Pony.objects.count(), 0) 

def test_run_python_noop(self):
        """
        #24098 - Tests no-op RunPython operations.
        """
        project_state = ProjectState()
        new_state = project_state.clone()
        operation = migrations.RunPython(migrations.RunPython.noop, migrations.RunPython.noop)
        with connection.schema_editor() as editor:
            operation.database_forwards("test_runpython", editor, project_state, new_state)
            operation.database_backwards("test_runpython", editor, new_state, project_state) 

def get_migrations_for_django_before_21():
    return [
        # add a new id field (without primary key information)
        migrations.AddField(
            model_name='resetpasswordtoken',
            name='id',
            field=models.IntegerField(null=True),
            preserve_default=True,
        ),
        # fill the new pk field
        migrations.RunPython(
            populate_auto_incrementing_pk_field,
            migrations.RunPython.noop
        ),
        # add primary key information to id field
        migrations.AlterField(
            model_name='resetpasswordtoken',
            name='id',
            field=models.AutoField(primary_key=True, serialize=False)
        ),
        # remove primary key information from 'key' field
        migrations.AlterField(
            model_name='resetpasswordtoken',
            name='key',
            field=models.CharField(db_index=True, max_length=64, unique=True, verbose_name='Key'),
        ),
    ] 

def do_nothing(*args, **kwargs):
    ''' A no-op for reverse migration. Django 1.8 has RunPython.noop(), but
    1.7 does not. '''
    pass 

def set_data_source(apps, schema_editor):
    # Using historical versions as recommended for RunPython
    PublicDataAccess = apps.get_model("public_data", "PublicDataAccess")
    db_alias = schema_editor.connection.alias

    for pda in PublicDataAccess.objects.using(db_alias).filter(data_source=None):
        pda.data_source = "direct-sharing-{}".format(pda.project_membership.project.id)
        pda.save() 

def test_run_python_noop(self):
        """
        #24098 - Tests no-op RunPython operations.
        """
        project_state = ProjectState()
        new_state = project_state.clone()
        operation = migrations.RunPython(migrations.RunPython.noop, migrations.RunPython.noop)
        with connection.schema_editor() as editor:
            operation.database_forwards("test_runpython", editor, project_state, new_state)
            operation.database_backwards("test_runpython", editor, new_state, project_state) 

def test_run_python_noop(self):
        """
        #24098 - Tests no-op RunPython operations.
        """
        project_state = ProjectState()
        new_state = project_state.clone()
        operation = migrations.RunPython(migrations.RunPython.noop, migrations.RunPython.noop)
        with connection.schema_editor() as editor:
            operation.database_forwards("test_runpython", editor, project_state, new_state)
            operation.database_backwards("test_runpython", editor, new_state, project_state) 

def test_run_python_related_assignment(self):
        """
        #24282 - Model changes to a FK reverse side update the model
        on the FK side as well.
        """

        def inner_method(models, schema_editor):
            Author = models.get_model("test_authors", "Author")
            Book = models.get_model("test_books", "Book")
            author = Author.objects.create(name="Hemingway")
            Book.objects.create(title="Old Man and The Sea", author=author)

        create_author = migrations.CreateModel(
            "Author",
            [
                ("id", models.AutoField(primary_key=True)),
                ("name", models.CharField(max_length=100)),
            ],
            options={},
        )
        create_book = migrations.CreateModel(
            "Book",
            [
                ("id", models.AutoField(primary_key=True)),
                ("title", models.CharField(max_length=100)),
                ("author", models.ForeignKey("test_authors.Author", models.CASCADE))
            ],
            options={},
        )
        add_hometown = migrations.AddField(
            "Author",
            "hometown",
            models.CharField(max_length=100),
        )
        create_old_man = migrations.RunPython(inner_method, inner_method)

        project_state = ProjectState()
        new_state = project_state.clone()
        with connection.schema_editor() as editor:
            create_author.state_forwards("test_authors", new_state)
            create_author.database_forwards("test_authors", editor, project_state, new_state)
        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            create_book.state_forwards("test_books", new_state)
            create_book.database_forwards("test_books", editor, project_state, new_state)
        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            add_hometown.state_forwards("test_authors", new_state)
            add_hometown.database_forwards("test_authors", editor, project_state, new_state)
        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            create_old_man.state_forwards("test_books", new_state)
            create_old_man.database_forwards("test_books", editor, project_state, new_state) 

def test_model_with_bigautofield(self):
        """
        A model with BigAutoField can be created.
        """
        def create_data(models, schema_editor):
            Author = models.get_model("test_author", "Author")
            Book = models.get_model("test_book", "Book")
            author1 = Author.objects.create(name="Hemingway")
            Book.objects.create(title="Old Man and The Sea", author=author1)
            Book.objects.create(id=2 ** 33, title="A farewell to arms", author=author1)

            author2 = Author.objects.create(id=2 ** 33, name="Remarque")
            Book.objects.create(title="All quiet on the western front", author=author2)
            Book.objects.create(title="Arc de Triomphe", author=author2)

        create_author = migrations.CreateModel(
            "Author",
            [
                ("id", models.BigAutoField(primary_key=True)),
                ("name", models.CharField(max_length=100)),
            ],
            options={},
        )
        create_book = migrations.CreateModel(
            "Book",
            [
                ("id", models.BigAutoField(primary_key=True)),
                ("title", models.CharField(max_length=100)),
                ("author", models.ForeignKey(to="test_author.Author", on_delete=models.CASCADE))
            ],
            options={},
        )
        fill_data = migrations.RunPython(create_data)

        project_state = ProjectState()
        new_state = project_state.clone()
        with connection.schema_editor() as editor:
            create_author.state_forwards("test_author", new_state)
            create_author.database_forwards("test_author", editor, project_state, new_state)

        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            create_book.state_forwards("test_book", new_state)
            create_book.database_forwards("test_book", editor, project_state, new_state)

        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            fill_data.state_forwards("fill_data", new_state)
            fill_data.database_forwards("fill_data", editor, project_state, new_state) 

def test_model_with_bigautofield(self):
        """
        A model with BigAutoField can be created.
        """
        def create_data(models, schema_editor):
            Author = models.get_model("test_author", "Author")
            Book = models.get_model("test_book", "Book")
            author1 = Author.objects.create(name="Hemingway")
            Book.objects.create(title="Old Man and The Sea", author=author1)
            Book.objects.create(id=2 ** 33, title="A farewell to arms", author=author1)

            author2 = Author.objects.create(id=2 ** 33, name="Remarque")
            Book.objects.create(title="All quiet on the western front", author=author2)
            Book.objects.create(title="Arc de Triomphe", author=author2)

        create_author = migrations.CreateModel(
            "Author",
            [
                ("id", models.BigAutoField(primary_key=True)),
                ("name", models.CharField(max_length=100)),
            ],
            options={},
        )
        create_book = migrations.CreateModel(
            "Book",
            [
                ("id", models.BigAutoField(primary_key=True)),
                ("title", models.CharField(max_length=100)),
                ("author", models.ForeignKey(to="test_author.Author", on_delete=models.CASCADE))
            ],
            options={},
        )
        fill_data = migrations.RunPython(create_data)

        project_state = ProjectState()
        new_state = project_state.clone()
        with connection.schema_editor() as editor:
            create_author.state_forwards("test_author", new_state)
            create_author.database_forwards("test_author", editor, project_state, new_state)

        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            create_book.state_forwards("test_book", new_state)
            create_book.database_forwards("test_book", editor, project_state, new_state)

        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            fill_data.state_forwards("fill_data", new_state)
            fill_data.database_forwards("fill_data", editor, project_state, new_state) 

def test_run_python_related_assignment(self):
        """
        #24282 - Model changes to a FK reverse side update the model
        on the FK side as well.
        """

        def inner_method(models, schema_editor):
            Author = models.get_model("test_authors", "Author")
            Book = models.get_model("test_books", "Book")
            author = Author.objects.create(name="Hemingway")
            Book.objects.create(title="Old Man and The Sea", author=author)

        create_author = migrations.CreateModel(
            "Author",
            [
                ("id", models.AutoField(primary_key=True)),
                ("name", models.CharField(max_length=100)),
            ],
            options={},
        )
        create_book = migrations.CreateModel(
            "Book",
            [
                ("id", models.AutoField(primary_key=True)),
                ("title", models.CharField(max_length=100)),
                ("author", models.ForeignKey("test_authors.Author", models.CASCADE))
            ],
            options={},
        )
        add_hometown = migrations.AddField(
            "Author",
            "hometown",
            models.CharField(max_length=100),
        )
        create_old_man = migrations.RunPython(inner_method, inner_method)

        project_state = ProjectState()
        new_state = project_state.clone()
        with connection.schema_editor() as editor:
            create_author.state_forwards("test_authors", new_state)
            create_author.database_forwards("test_authors", editor, project_state, new_state)
        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            create_book.state_forwards("test_books", new_state)
            create_book.database_forwards("test_books", editor, project_state, new_state)
        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            add_hometown.state_forwards("test_authors", new_state)
            add_hometown.database_forwards("test_authors", editor, project_state, new_state)
        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            create_old_man.state_forwards("test_books", new_state)
            create_old_man.database_forwards("test_books", editor, project_state, new_state) 

def test_model_with_bigautofield(self):
        """
        A model with BigAutoField can be created.
        """
        def create_data(models, schema_editor):
            Author = models.get_model("test_author", "Author")
            Book = models.get_model("test_book", "Book")
            author1 = Author.objects.create(name="Hemingway")
            Book.objects.create(title="Old Man and The Sea", author=author1)
            Book.objects.create(id=2 ** 33, title="A farewell to arms", author=author1)

            author2 = Author.objects.create(id=2 ** 33, name="Remarque")
            Book.objects.create(title="All quiet on the western front", author=author2)
            Book.objects.create(title="Arc de Triomphe", author=author2)

        create_author = migrations.CreateModel(
            "Author",
            [
                ("id", models.BigAutoField(primary_key=True)),
                ("name", models.CharField(max_length=100)),
            ],
            options={},
        )
        create_book = migrations.CreateModel(
            "Book",
            [
                ("id", models.BigAutoField(primary_key=True)),
                ("title", models.CharField(max_length=100)),
                ("author", models.ForeignKey(to="test_author.Author", on_delete=models.CASCADE))
            ],
            options={},
        )
        fill_data = migrations.RunPython(create_data)

        project_state = ProjectState()
        new_state = project_state.clone()
        with connection.schema_editor() as editor:
            create_author.state_forwards("test_author", new_state)
            create_author.database_forwards("test_author", editor, project_state, new_state)

        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            create_book.state_forwards("test_book", new_state)
            create_book.database_forwards("test_book", editor, project_state, new_state)

        project_state = new_state
        new_state = new_state.clone()
        with connection.schema_editor() as editor:
            fill_data.state_forwards("fill_data", new_state)
            fill_data.database_forwards("fill_data", editor, project_state, new_state) 

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.apps.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.apps.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.apps.get_model("test_runpythonatomic", "Pony").objects.count(), 0)
        # Otherwise, the non-atomic operation should leave a row there
        else:
            self.assertEqual(project_state.apps.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.apps.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.apps.get_model("test_runpythonatomic", "Pony").objects.count(), 1)
        # And deconstruction
        definition = non_atomic_migration.operations[0].deconstruct()
        self.assertEqual(definition[0], "RunPython")
        self.assertEqual(definition[1], [])
        self.assertEqual(sorted(definition[2]), ["atomic", "code"]) 

def lint_runpython(self, runpython):
        function_name = runpython.code.__name__
        error = []
        ignored = []
        warning = []

        if not runpython.reversible:
            issue = {
                "code": "REVERSIBLE_DATA_MIGRATION",
                "msg": "'{}': RunPython data migration is not reversible".format(
                    function_name
                ),
            }
            if issue["code"] in self.exclude_migration_tests:
                ignored.append(issue)
            else:
                warning.append(issue)

        if PY2:
            args_spec = inspect.getargspec(runpython.code)
        else:
            args_spec = inspect.getfullargspec(runpython.code)

        if tuple(args_spec.args) != EXPECTED_DATA_MIGRATION_ARGS:
            issue = {
                "code": "NAMING_CONVENTION_RUNPYTHON_ARGS",
                "msg": (
                    "'{}': By convention, "
                    "RunPython names the two arguments: apps, schema_editor"
                ).format(function_name),
            }
            if issue["code"] in self.exclude_migration_tests:
                ignored.append(issue)
            else:
                warning.append(issue)

        # Detect wrong model imports
        # Forward
        issues = self.get_data_migration_model_import_issues(runpython.code)
        for issue in issues:
            if issue["code"] in self.exclude_migration_tests:
                ignored.append(issue)
            else:
                error.append(issue)

        # Backward
        if runpython.reversible:
            issues = self.get_data_migration_model_import_issues(runpython.reverse_code)
            for issue in issues:
                if issue and issue["code"] in self.exclude_migration_tests:
                    ignored.append(issue)
                else:
                    error.append(issue)
        return error, ignored, warning