Fix KeyError when deleting a field with unqiue=True (#365)

* Fix KeyError when deleting a field with unqiue=True

* refactor: rename `old_data_unique` to `is_unique_field`

* Add testcases for remove unique field

---------

Co-authored-by: gongchangku <gongchangku@anban.tech>
Co-authored-by: Waket Zheng <waketzheng@gmail.com>

CHANGELOG.md

@@ -7,7 +7,7 @@
 - Correct the click import. (#360)
 - Improve CLI help text and output. (#355)
 - Fix mysql drop unique index raises OperationalError. (#346)
-
+- Fix KeyError when deleting a field with unqiue=True. (#364)
   **Upgrade note:**
     1. Use column name as unique key name for mysql
     2. Drop support for Python3.7

aerich/migrate.py

@@ -425,8 +425,9 @@ class Migrate:
                         upgrade,
                     )
                     if old_data_field["indexed"]:
+                        is_unique_field = old_data_field.get("unique")
                         cls._add_operator(
-                            cls._drop_index(model, {db_column}),
+                            cls._drop_index(model, {db_column}, is_unique_field),
                             upgrade,
                             True,
                         )
@@ -548,13 +549,17 @@ class Migrate:
     def _resolve_fk_fields_name(cls, model: Type[Model], fields_name: Iterable[str]) -> List[str]:
         ret = []
         for field_name in fields_name:
-            field = model._meta.fields_map[field_name]
+            try:
-            if field.source_field:
+                field = model._meta.fields_map[field_name]
-                ret.append(field.source_field)
+            except KeyError:
-            elif field_name in model._meta.fk_fields:
+                # field dropped or to be add
-                ret.append(field_name + "_id")
+                pass
             else:
-                ret.append(field_name)
+                if field.source_field:
+                    field_name = field.source_field
+                elif field_name in model._meta.fk_fields:
+                    field_name += "_id"
+            ret.append(field_name)
         return ret
 
     @classmethod

tests/old_models.py

@@ -51,6 +51,7 @@ class Category(Model):
 
 class Product(Model):
     categories: fields.ManyToManyRelation[Category] = fields.ManyToManyField("models.Category")
+    uid = fields.IntField(source_field="uuid", unique=True)
     name = fields.CharField(max_length=50)
     view_num = fields.IntField(description="View Num")
     sort = fields.IntField()
@@ -64,6 +65,7 @@ class Product(Model):
 
 
 class Config(Model):
+    name = fields.CharField(max_length=100, unique=True)
     label = fields.CharField(max_length=200)
     key = fields.CharField(max_length=20)
     value: dict = fields.JSONField()

tests/test_migrate.py

@@ -188,6 +188,21 @@ old_models_describe = {
             "db_field_types": {"": "INT"},
         },
         "data_fields": [
+            {
+                "name": "name",
+                "field_type": "CharField",
+                "db_column": "name",
+                "python_type": "str",
+                "generated": False,
+                "nullable": False,
+                "unique": True,
+                "indexed": True,
+                "default": None,
+                "description": None,
+                "docstring": None,
+                "constraints": {"max_length": 100},
+                "db_field_types": {"": "VARCHAR(100)"},
+            },
             {
                 "name": "label",
                 "field_type": "CharField",
@@ -388,6 +403,21 @@ old_models_describe = {
                 "constraints": {"max_length": 50},
                 "db_field_types": {"": "VARCHAR(50)"},
             },
+            {
+                "name": "uid",
+                "field_type": "IntField",
+                "db_column": "uuid",
+                "python_type": "int",
+                "generated": False,
+                "nullable": False,
+                "unique": True,
+                "indexed": True,
+                "default": None,
+                "description": None,
+                "docstring": None,
+                "constraints": {"ge": -2147483648, "le": 2147483647},
+                "db_field_types": {"": "INT"},
+            },
             {
                 "name": "view_num",
                 "field_type": "IntField",
@@ -816,6 +846,7 @@ def test_migrate(mocker: MockerFixture):
     - add unique: User.username
     - change column: length User.password
     - add unique_together: (name,type) of Product
+    - drop unique field: Config.name
     - alter default: Config.status
     - rename column: Product.image -> Product.pic
     """
@@ -837,6 +868,8 @@ def test_migrate(mocker: MockerFixture):
             "ALTER TABLE `category` MODIFY COLUMN `name` VARCHAR(200)",
             "ALTER TABLE `category` MODIFY COLUMN `slug` VARCHAR(100) NOT NULL",
             "ALTER TABLE `category` DROP INDEX `title`",
+            "ALTER TABLE `config` DROP COLUMN `name`",
+            "ALTER TABLE `config` DROP INDEX `name`",
             "ALTER TABLE `config` ADD `user_id` INT NOT NULL  COMMENT 'User'",
             "ALTER TABLE `config` ADD CONSTRAINT `fk_config_user_17daa970` FOREIGN KEY (`user_id`) REFERENCES `user` (`id`) ON DELETE CASCADE",
             "ALTER TABLE `config` ALTER COLUMN `status` DROP DEFAULT",
@@ -844,6 +877,8 @@ def test_migrate(mocker: MockerFixture):
             "ALTER TABLE `email` ADD `address` VARCHAR(200) NOT NULL",
             "ALTER TABLE `email` DROP COLUMN `user_id`",
             "ALTER TABLE `configs` RENAME TO `config`",
+            "ALTER TABLE `product` DROP COLUMN `uuid`",
+            "ALTER TABLE `product` DROP INDEX `uuid`",
             "ALTER TABLE `product` RENAME COLUMN `image` TO `pic`",
             "ALTER TABLE `email` RENAME COLUMN `id` TO `email_id`",
             "ALTER TABLE `product` ADD INDEX `idx_product_name_869427` (`name`, `type_db_alias`)",
@@ -870,6 +905,8 @@ def test_migrate(mocker: MockerFixture):
             "ALTER TABLE `category` MODIFY COLUMN `name` VARCHAR(200) NOT NULL",
             "ALTER TABLE `category` MODIFY COLUMN `slug` VARCHAR(200) NOT NULL",
             "ALTER TABLE `category` ADD UNIQUE INDEX `title` (`title`)",
+            "ALTER TABLE `config` ADD `name` VARCHAR(100) NOT NULL UNIQUE",
+            "ALTER TABLE `config` ADD UNIQUE INDEX `name` (`name`)",
             "ALTER TABLE `config` DROP COLUMN `user_id`",
             "ALTER TABLE `config` DROP FOREIGN KEY `fk_config_user_17daa970`",
             "ALTER TABLE `config` ALTER COLUMN `status` SET DEFAULT 1",
@@ -878,6 +915,8 @@ def test_migrate(mocker: MockerFixture):
             "ALTER TABLE `config` RENAME TO `configs`",
             "ALTER TABLE `product` RENAME COLUMN `pic` TO `image`",
             "ALTER TABLE `email` RENAME COLUMN `email_id` TO `id`",
+            "ALTER TABLE `product` ADD `uuid` INT NOT NULL UNIQUE",
+            "ALTER TABLE `product` ADD UNIQUE INDEX `uuid` (`uuid`)",
             "ALTER TABLE `product` DROP INDEX `idx_product_name_869427`",
             "ALTER TABLE `email` DROP INDEX `idx_email_email_4a1a33`",
             "ALTER TABLE `product` DROP INDEX `uid_product_name_869427`",
@@ -915,6 +954,8 @@ def test_migrate(mocker: MockerFixture):
             'ALTER TABLE "category" ALTER COLUMN "name" DROP NOT NULL',
             'ALTER TABLE "category" ALTER COLUMN "slug" TYPE VARCHAR(100) USING "slug"::VARCHAR(100)',
             'ALTER TABLE "category" ALTER COLUMN "created_at" TYPE TIMESTAMPTZ USING "created_at"::TIMESTAMPTZ',
+            'ALTER TABLE "config" DROP COLUMN "name"',
+            'DROP INDEX "uid_config_name_2c83c8"',
             'ALTER TABLE "config" ADD "user_id" INT NOT NULL',
             'ALTER TABLE "config" ADD CONSTRAINT "fk_config_user_17daa970" FOREIGN KEY ("user_id") REFERENCES "user" ("id") ON DELETE CASCADE',
             'ALTER TABLE "config" ALTER COLUMN "status" DROP DEFAULT',
@@ -924,6 +965,8 @@ def test_migrate(mocker: MockerFixture):
             'ALTER TABLE "email" DROP COLUMN "user_id"',
             'ALTER TABLE "email" RENAME COLUMN "id" TO "email_id"',
             'ALTER TABLE "email" ALTER COLUMN "is_primary" TYPE BOOL USING "is_primary"::BOOL',
+            'DROP INDEX "uid_product_uuid_d33c18"',
+            'ALTER TABLE "product" DROP COLUMN "uuid"',
             'ALTER TABLE "product" ALTER COLUMN "view_num" SET DEFAULT 0',
             'ALTER TABLE "product" RENAME COLUMN "image" TO "pic"',
             'ALTER TABLE "product" ALTER COLUMN "is_reviewed" TYPE BOOL USING "is_reviewed"::BOOL',
@@ -948,6 +991,8 @@ def test_migrate(mocker: MockerFixture):
             'ALTER TABLE "category" ALTER COLUMN "name" SET NOT NULL',
             'ALTER TABLE "category" ALTER COLUMN "slug" TYPE VARCHAR(200) USING "slug"::VARCHAR(200)',
             'ALTER TABLE "category" ALTER COLUMN "created_at" TYPE TIMESTAMPTZ USING "created_at"::TIMESTAMPTZ',
+            'ALTER TABLE "config" ADD "name" VARCHAR(100) NOT NULL UNIQUE',
+            'CREATE UNIQUE INDEX "uid_config_name_2c83c8" ON "config" ("name")',
             'ALTER TABLE "config" ALTER COLUMN "status" SET DEFAULT 1',
             'ALTER TABLE "config" DROP COLUMN "user_id"',
             'ALTER TABLE "config" DROP CONSTRAINT "fk_config_user_17daa970"',
@@ -957,6 +1002,8 @@ def test_migrate(mocker: MockerFixture):
             'ALTER TABLE "email" DROP COLUMN "address"',
             'ALTER TABLE "email" RENAME COLUMN "email_id" TO "id"',
             'ALTER TABLE "email" ALTER COLUMN "is_primary" TYPE BOOL USING "is_primary"::BOOL',
+            'ALTER TABLE "product" ADD "uuid" INT NOT NULL UNIQUE',
+            'CREATE UNIQUE INDEX "uid_product_uuid_d33c18" ON "product" ("uuid")',
             'ALTER TABLE "product" ALTER COLUMN "view_num" DROP DEFAULT',
             'ALTER TABLE "product" RENAME COLUMN "pic" TO "image"',
             'ALTER TABLE "user" ADD "avatar" VARCHAR(200) NOT NULL  DEFAULT \'\'',