fix: migrate drop the wrong m2m field when model have multi m2m fields

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### [0.8.1](Unreleased)
 
 #### Fixed
+- Migrate drop the wrong m2m field when model have multi m2m fields. (#376)
 - KeyError raised when removing or renaming an existing model (#386)
 - fix: error when there is `__init__.py` in the migration folder (#272)
 - Setting null=false on m2m field causes migration to fail. (#334)

aerich/migrate.py

@@ -13,7 +13,12 @@
 
 from aerich.ddl import BaseDDL
 from aerich.models import MAX_VERSION_LENGTH, Aerich
-from aerich.utils import get_app_connection, get_models_describe, is_default_function
+from aerich.utils import (
+    get_app_connection,
+    get_models_describe,
+    is_default_function,
+    reorder_m2m_fields,
+)
 
 MIGRATE_TEMPLATE = """from tortoise import BaseDBAsyncClient
 
@@ -279,10 +284,10 @@ def diff_models(
                 # m2m fields
                 old_m2m_fields = cast(List[dict], old_model_describe.get("m2m_fields"))
                 new_m2m_fields = cast(List[dict], new_model_describe.get("m2m_fields"))
-                if old_m2m_fields and len(new_m2m_fields) >= 2:
-                    length = len(old_m2m_fields)
-                    field_index = {f["name"]: i for i, f in enumerate(new_m2m_fields)}
-                    new_m2m_fields.sort(key=lambda field: field_index.get(field["name"], length))
+                if old_m2m_fields and new_m2m_fields:
+                    old_m2m_fields, new_m2m_fields = reorder_m2m_fields(
+                        old_m2m_fields, new_m2m_fields
+                    )
                 for action, option, change in diff(old_m2m_fields, new_m2m_fields):
                     if (option and option[-1] == "nullable") or change[0][0] == "db_constraint":
                         continue

aerich/utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import importlib.util
 import os
 import re
@@ -101,3 +103,112 @@ def import_py_file(file: Union[str, Path]) -> ModuleType:
     module = importlib.util.module_from_spec(spec)  # type:ignore[arg-type]
     spec.loader.exec_module(module)  # type:ignore[union-attr]
     return module
+
+
+def _pick_match_to_head(m2m_fields: list[dict], field_info: dict) -> list[dict]:
+    """
+    If there is a element in m2m_fields whose through or name is equal to field_info's
+    and this element is not at the first position, put it to the head then return the
+    new list, otherwise return the origin list.
+
+    Example::
+        >>> m2m_fields = [{'through': 'u1', 'name': 'u1'}, {'throught': 'u2', 'name': 'u2'}]
+        >>> _pick_match_to_head(m2m_fields, {'through': 'u2', 'name': 'u2'})
+        [{'through': 'u2', 'name': 'u2'}, {'through': 'u1', 'name': 'u1'}]
+
+    """
+    through = field_info["through"]
+    name = field_info["name"]
+    for index, field in enumerate(m2m_fields):
+        if field["through"] == through or field["name"] == name:
+            if index != 0:
+                m2m_fields = [field, *m2m_fields[:index], *m2m_fields[index + 1 :]]
+            break
+    return m2m_fields
+
+
+def reorder_m2m_fields(
+    old_m2m_fields: list[dict], new_m2m_fields: list[dict]
+) -> tuple[list[dict], list[dict]]:
+    """
+    Reorder m2m fields to help dictdiffer.diff generate more precise changes
+    :param old_m2m_fields: previous m2m field info list
+    :param new_m2m_fields: current m2m field info list
+    :return: ordered old/new m2m fields
+    """
+    length_old, length_new = len(old_m2m_fields), len(new_m2m_fields)
+    if length_old == length_new == 1:
+        # No need to change order if both of them have only one element
+        pass
+    elif length_old == 1:
+        # If any element of new fields match the one in old fields, put it to head
+        new_m2m_fields = _pick_match_to_head(new_m2m_fields, old_m2m_fields[0])
+    elif length_new == 1:
+        old_m2m_fields = _pick_match_to_head(old_m2m_fields, new_m2m_fields[0])
+    else:
+        old_table_names = [f["through"] for f in old_m2m_fields]
+        new_table_names = [f["through"] for f in new_m2m_fields]
+        old_field_names = [f["name"] for f in old_m2m_fields]
+        new_field_names = [f["name"] for f in new_m2m_fields]
+        if old_table_names == new_table_names:
+            pass
+        elif sorted(old_table_names) == sorted(new_table_names):
+            # If table name are the same but order not match,
+            # reorder new fields by through to match the order of old.
+
+            # Case like::
+            # old_m2m_fields = [
+            #     {'through': 'users_group', 'name': 'users',},
+            #     {'through': 'admins_group', 'name': 'admins'},
+            # ]
+            # new_m2m_fields = [
+            #     {'through': 'admins_group', 'name': 'admins_new'},
+            #     {'through': 'users_group', 'name': 'users_new',},
+            # ]
+            new_m2m_fields = sorted(
+                new_m2m_fields, key=lambda f: old_table_names.index(f["through"])
+            )
+        elif old_field_names == new_field_names:
+            pass
+        elif sorted(old_field_names) == sorted(new_field_names):
+            # Case like:
+            # old_m2m_fields = [
+            #     {'name': 'users', 'through': 'users_group'},
+            #     {'name': 'admins', 'through': 'admins_group'},
+            # ]
+            # new_m2m_fields = [
+            #     {'name': 'admins', 'through': 'admin_group_map'},
+            #     {'name': 'users', 'through': 'user_group_map'},
+            # ]
+            new_m2m_fields = sorted(new_m2m_fields, key=lambda f: old_field_names.index(f["name"]))
+        elif unchanged_table_names := set(old_table_names) & set(new_table_names):
+            # If old/new m2m field list have one or some unchanged table names, put them to head of list.
+
+            # Case like::
+            # old_m2m_fields = [
+            #     {'through': 'users_group', 'name': 'users',},
+            #     {'through': 'staffs_group', 'name': 'users',},
+            #     {'through': 'admins_group', 'name': 'admins'},
+            # ]
+            # new_m2m_fields = [
+            #     {'through': 'admins_group', 'name': 'admins_new'},
+            #     {'through': 'users_group', 'name': 'users_new',},
+            # ]
+            unchanged = sorted(unchanged_table_names, key=lambda name: old_table_names.index(name))
+            ordered_old_tables = unchanged + sorted(
+                set(old_table_names) - unchanged_table_names,
+                key=lambda name: old_table_names.index(name),
+            )
+            ordered_new_tables = unchanged + sorted(
+                set(new_table_names) - unchanged_table_names,
+                key=lambda name: new_table_names.index(name),
+            )
+            if ordered_old_tables != old_table_names:
+                old_m2m_fields = sorted(
+                    old_m2m_fields, key=lambda f: ordered_old_tables.index(f["through"])
+                )
+            if ordered_new_tables != new_table_names:
+                new_m2m_fields = sorted(
+                    new_m2m_fields, key=lambda f: ordered_new_tables.index(f["through"])
+                )
+    return old_m2m_fields, new_m2m_fields

tests/models.py

@@ -80,6 +80,9 @@ class Meta:
 
 
 class Config(Model):
+    categories: fields.ManyToManyRelation[Category] = fields.ManyToManyField(
+        "models.Category", through="config_category_map", related_name="category_set"
+    )
     label = fields.CharField(max_length=200)
     key = fields.CharField(max_length=20)
     value: dict = fields.JSONField()

tests/old_models.py

@@ -65,6 +65,10 @@ class Product(Model):
 
 
 class Config(Model):
+    category: fields.ManyToManyRelation[Category] = fields.ManyToManyField("models.Category")
+    categories: fields.ManyToManyRelation[Category] = fields.ManyToManyField(
+        "models.Category", through="config_category_map", related_name="config_set"
+    )
     name = fields.CharField(max_length=100, unique=True)
     label = fields.CharField(max_length=200)
     key = fields.CharField(max_length=20)

tests/test_migrate.py

@@ -270,7 +270,48 @@
         "backward_fk_fields": [],
         "o2o_fields": [],
         "backward_o2o_fields": [],
-        "m2m_fields": [],
+        "m2m_fields": [
+            {
+                "name": "category",
+                "field_type": "ManyToManyFieldInstance",
+                "python_type": "models.Category",
+                "generated": False,
+                "nullable": False,
+                "unique": False,
+                "indexed": False,
+                "default": None,
+                "description": None,
+                "docstring": None,
+                "constraints": {},
+                "model_name": "models.Category",
+                "related_name": "configs",
+                "forward_key": "category_id",
+                "backward_key": "config_id",
+                "through": "config_category",
+                "on_delete": "CASCADE",
+                "_generated": False,
+            },
+            {
+                "name": "categories",
+                "field_type": "ManyToManyFieldInstance",
+                "python_type": "models.Category",
+                "generated": False,
+                "nullable": False,
+                "unique": False,
+                "indexed": False,
+                "default": None,
+                "description": None,
+                "docstring": None,
+                "constraints": {},
+                "model_name": "models.Category",
+                "related_name": "config_set",
+                "forward_key": "category_id",
+                "backward_key": "config_id",
+                "through": "config_category_map",
+                "on_delete": "CASCADE",
+                "_generated": False,
+            },
+        ],
     },
     "models.Email": {
         "name": "models.Email",
@@ -898,6 +939,8 @@ def test_migrate(mocker: MockerFixture):
             "ALTER TABLE `product` MODIFY COLUMN `body` LONGTEXT NOT NULL",
             "ALTER TABLE `email` MODIFY COLUMN `is_primary` BOOL NOT NULL  DEFAULT 0",
             "CREATE TABLE `product_user` (\n    `product_id` INT NOT NULL REFERENCES `product` (`id`) ON DELETE CASCADE,\n    `user_id` INT NOT NULL REFERENCES `user` (`id`) ON DELETE CASCADE\n) CHARACTER SET utf8mb4",
+            "CREATE TABLE `config_category_map` (\n    `category_id` INT NOT NULL REFERENCES `category` (`id`) ON DELETE CASCADE,\n    `config_id` INT NOT NULL REFERENCES `config` (`id`) ON DELETE CASCADE\n) CHARACTER SET utf8mb4",
+            "DROP TABLE IF EXISTS `config_category`",
         }
         expected_downgrade_operators = {
             "ALTER TABLE `category` MODIFY COLUMN `name` VARCHAR(200) NOT NULL",
@@ -937,6 +980,8 @@ def test_migrate(mocker: MockerFixture):
             "ALTER TABLE `user` MODIFY COLUMN `longitude` DECIMAL(12,9) NOT NULL",
             "ALTER TABLE `product` MODIFY COLUMN `body` LONGTEXT NOT NULL",
             "ALTER TABLE `email` MODIFY COLUMN `is_primary` BOOL NOT NULL  DEFAULT 0",
+            "CREATE TABLE `config_category` (\n    `config_id` INT NOT NULL REFERENCES `config` (`id`) ON DELETE CASCADE,\n    `category_id` INT NOT NULL REFERENCES `category` (`id`) ON DELETE CASCADE\n) CHARACTER SET utf8mb4",
+            "DROP TABLE IF EXISTS `config_category_map`",
         }
         assert not set(Migrate.upgrade_operators).symmetric_difference(expected_upgrade_operators)
 
@@ -983,6 +1028,8 @@ def test_migrate(mocker: MockerFixture):
             'CREATE UNIQUE INDEX "uid_product_name_869427" ON "product" ("name", "type_db_alias")',
             'CREATE UNIQUE INDEX "uid_user_usernam_9987ab" ON "user" ("username")',
             'CREATE TABLE "product_user" (\n    "product_id" INT NOT NULL REFERENCES "product" ("id") ON DELETE CASCADE,\n    "user_id" INT NOT NULL REFERENCES "user" ("id") ON DELETE CASCADE\n)',
+            'CREATE TABLE "config_category_map" (\n    "category_id" INT NOT NULL REFERENCES "category" ("id") ON DELETE CASCADE,\n    "config_id" INT NOT NULL REFERENCES "config" ("id") ON DELETE CASCADE\n)',
+            'DROP TABLE IF EXISTS "config_category"',
         }
         expected_downgrade_operators = {
             'CREATE UNIQUE INDEX "uid_category_title_f7fc03" ON "category" ("title")',
@@ -1022,6 +1069,8 @@ def test_migrate(mocker: MockerFixture):
             'DROP INDEX IF EXISTS "uid_product_name_869427"',
             'DROP TABLE IF EXISTS "email_user"',
             'DROP TABLE IF EXISTS "newmodel"',
+            'CREATE TABLE "config_category" (\n    "config_id" INT NOT NULL REFERENCES "config" ("id") ON DELETE CASCADE,\n    "category_id" INT NOT NULL REFERENCES "category" ("id") ON DELETE CASCADE\n)',
+            'DROP TABLE IF EXISTS "config_category_map"',
         }
         assert not set(Migrate.upgrade_operators).symmetric_difference(expected_upgrade_operators)
         assert not set(Migrate.downgrade_operators).symmetric_difference(