Support for custom code for non-trivial migrations

[merlinz01]

Good to see someone working on this!

In addition to that aerich migrations are db-specific, one of my reasons for opening tortoise/aerich#366 is that aerich doesn't support custom migration code. So I want to make sure that it is on the table for this project.

One real-life example of where custom code would be useful: replacing an "issue_status" enum field with an "issue_open" boolean field. Supposing we have a list of issues where some are open and some are closed. Merely removing the "issue_status" field and adding the "issue_open" field with a default results in data loss, because all the issue statuses are now the default for the added column.
We could improve this with some custom code like this:

class Migration:
    async def before_migrate(self, db):
        self.issues = await db.execute('SELECT id, issue_status FROM issues').all()

    async def after_migrate(self, db):
        for id, old_issue_status in self.issues:
            await db.execute('UPDATE issues SET issue_open=$1 WHERE id=$2', old_issue_status==IssueStatus.CLOSED, id)

Alternatively, (more flexible but maybe too complex):

class Migration:
    operations = [
        CustomFunction('get_old_issue_statuses'),
        RemoveField(...),
        AddField(...),
        CustomFunction('update_new_issue_statuses'),
    ]

    async def get_old_issue_statuses(self, db):
        self.issues = await db.execute('SELECT id, issue_status FROM issues').all()

    async def update_new_issue_statuses(self, db):
        for id, old_issue_status in self.issues:
            await db.execute('UPDATE issues SET issue_open=$1 WHERE id=$2', old_issue_status==IssueStatus.CLOSED, id)

Another thing to consider is optional support for downgrading with custom code, maybe with additional before/after_unmigrate methods or a second argument for CustomFunction.

I'm not really familiar with other migration tools so I don't know if there are any established paradigms for this, but I think we have the opportunity to take Tortoise ORM migrations to the next level.

[henadzit]

I like how Django approaches this:

from django.db import migrations


def forwards_func(apps, schema_editor):
    ...


def reverse_func(apps, schema_editor):
    ...


class Migration(migrations.Migration):
    dependencies = []

    operations = [
        migrations.RunPython(forwards_func, reverse_func),
    ]

Another approach would be to ask users to subclass RunPython and implement apply and revert (I probably like this one even more):

class SetFieldToTrue(RunPython):
  async def apply(self, connection):
    await Model.update(field=True)

  async def revert(self, connection):
    await Model.update(field=False)

With the approaches you suggested, I don't see how operations can be reverted but this is definitely something that the tool should support.

Another complexity here is that when executing queries in RunPython, you need to have the state of the model at the time of the migration execution, not the current state. Imagine the following migration history:

1. NewModel added
2. RunPython migration that sets NewModel.field to True
3. NewModel removed

If you pull the code and try to apply migration on a fresh database, the second migration will fail because there is no NewModel. Django has apps.get_model("myapp", "Country") that returns the state of the model at the time when the migration is executed. I'm guessing we will need to implement something similar but for now I guess we can just document that shortcoming.

[merlinz01]

Django's approach is almost like my second example but there's an aspect missing: It's important to have a clear way to save some kind of state before the schema changes and update the data based on it after the changes. I suppose you could do it with a global variable referenced from a second RunPython operation in the Django example, but it feels a bit hacky to resort to globals.

Here's an updated example showing how I think we could do this in a Pythonic manner and also support reverting custom migrations:

class Migration:
    operations = [
        Call(apply='get_old_issue_statuses', revert='revert_old_issue_statuses'),
        RemoveField(...),
        AddField(...),
        Call(apply='update_new_issue_statuses', revert='get_new_issue_statuses'),
    ]

    async def get_old_issue_statuses(self, db):
        # The class instance provides a simple way to preserve state from one point to another.
        self.issues = await db.execute('SELECT id, issue_status FROM issues').all()

    async def update_new_issue_statuses(self, db):
        for id, old_issue_status in self.issues:
            await db.execute('UPDATE issues SET issue_open=$1 WHERE id=$2', old_issue_status==IssueStatus.CLOSED, id)

    async def get_new_issue_statuses(self, db):
        self.issues_reverting = await db.execute('SELECT id, issue_open FROM issues').all()

    async def revert_old_issue_statuses(self, db):
        for id, new_issue_status in self.issues_reverting:
            status = IssueStatus.OPEN if new_issue_status else IssueStatus.CLOSED
            await db.execute('UPDATE issues SET issue_status=$1 WHERE id=$2', status, id)

As for the model/schema differences, to me it's not a problem that the models don't always match; that's an obvious side effect of migrations, and custom migrations should just use raw SQL (or generate it with pypika for cross-compatibility) if that is the case.