From a624d1b43ba3f5c9e0e680e30de008499925f9fe Mon Sep 17 00:00:00 2001 From: Waket Zheng Date: Thu, 27 Feb 2025 22:23:26 +0800 Subject: [PATCH] fix: migrate does not recognise attribute changes for string primary key (#428) * refactor: show warning for unsupported pk field changes * fix: migrate does not recognise attribute changes for string primary key * docs: update changelog * refactor: reduce indents * chore: update docs --- CHANGELOG.md | 1 + aerich/migrate.py | 50 +++++++++++++++++++++++++++++++++++++------ tests/models.py | 2 ++ tests/old_models.py | 1 + tests/test_migrate.py | 50 ++++++++++++++++++++++++++----------------- 5 files changed, 77 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bf94d4d..ba431f3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### [0.8.2]**(Unreleased)** #### Added +- Support changes `max_length` or int type for primary key field. ([#428]) - feat: support psycopg. ([#425]) - Support run `poetry add aerich` in project that inited by poetry v2. ([#424]) - feat: support command `python -m aerich`. ([#417]) diff --git a/aerich/migrate.py b/aerich/migrate.py index 0ae61eb..529ecd2 100644 --- a/aerich/migrate.py +++ b/aerich/migrate.py @@ -16,6 +16,7 @@ from tortoise.indexes import Index from aerich.coder import load_index from aerich.ddl import BaseDDL +from aerich.enums import Color from aerich.models import MAX_VERSION_LENGTH, Aerich from aerich.utils import ( get_app_connection, @@ -424,14 +425,8 @@ class Migrate: ) old_indexes = cls._get_indexes(model, old_model_describe) new_indexes = cls._get_indexes(model, new_model_describe) - old_pk_field = old_model_describe.get("pk_field") - new_pk_field = new_model_describe.get("pk_field") # pk field - changes = diff(old_pk_field, new_pk_field) - for action, option, change in changes: - # current only support rename pk - if action == "change" and option == "name": - cls._add_operator(cls._rename_field(model, *change), upgrade) + cls._handle_pk_field_alter(model, old_model_describe, new_model_describe, upgrade) # fk fields args = (old_model_describe, new_model_describe, model, old_models, new_models) cls._handle_fk_fields(*args, upgrade=upgrade) @@ -605,6 +600,47 @@ class Migrate: continue cls._add_operator(cls.drop_model(old_models[old_model]["table"]), upgrade) + @classmethod + def _handle_pk_field_alter( + cls, + model: type[Model], + old_model_describe: dict[str, dict], + new_model_describe: dict[str, dict], + upgrade: bool, + ) -> None: + old_pk_field = old_model_describe.get("pk_field", {}) + new_pk_field = new_model_describe.get("pk_field", {}) + changes = cls._exclude_extra_field_types(diff(old_pk_field, new_pk_field)) + sqls: list[str] = [] + for action, option, change in changes: + if action != "change": + continue + if option == "db_column": + # rename pk + sql = cls._rename_field(model, *change) + elif option == "constraints.max_length": + sql = cls._modify_field(model, new_pk_field) + elif option == "field_type": + # Only support change field type between int fields, e.g.: IntField -> BigIntField + if not all(field_type.endswith("IntField") for field_type in change): + if upgrade: + model_name = model._meta.full_name.split(".")[-1] + field_name = new_pk_field.get("name", "") + msg = ( + f"Does not support change primary_key({model_name}.{field_name}) field type," + " you may need to do it manually." + ) + click.secho(msg, fg=Color.yellow) + return + sql = cls._modify_field(model, new_pk_field) + else: + # Skip option like 'constraints.ge', 'constraints.le', 'db_field_types.' + continue + sqls.append(sql) + for sql in sorted(sqls, key=lambda x: "RENAME" not in x): + # TODO: alter references field in m2m table + cls._add_operator(sql, upgrade) + @classmethod def _handle_field_changes( cls, diff --git a/tests/models.py b/tests/models.py index 4328bc9..a1f2725 100644 --- a/tests/models.py +++ b/tests/models.py @@ -79,6 +79,7 @@ class Category(Model): class Product(Model): + id = fields.BigIntField(primary_key=True) categories: fields.ManyToManyRelation[Category] = fields.ManyToManyField( "models.Category", null=False ) @@ -106,6 +107,7 @@ class Product(Model): class Config(Model): + slug = fields.CharField(primary_key=True, max_length=20) categories: fields.ManyToManyRelation[Category] = fields.ManyToManyField( "models.Category", through="config_category_map", related_name="category_set" ) diff --git a/tests/old_models.py b/tests/old_models.py index 6994194..7e46338 100644 --- a/tests/old_models.py +++ b/tests/old_models.py @@ -77,6 +77,7 @@ class Product(Model): class Config(Model): + slug = fields.CharField(primary_key=True, max_length=10) category: fields.ManyToManyRelation[Category] = fields.ManyToManyField("models.Category") categories: fields.ManyToManyRelation[Category] = fields.ManyToManyField( "models.Category", through="config_category_map", related_name="config_set" diff --git a/tests/test_migrate.py b/tests/test_migrate.py index e078904..e2f93ec 100644 --- a/tests/test_migrate.py +++ b/tests/test_migrate.py @@ -187,19 +187,19 @@ old_models_describe = { "unique_together": [], "indexes": [], "pk_field": { - "name": "id", - "field_type": "IntField", - "db_column": "id", - "python_type": "int", - "generated": True, + "name": "slug", + "field_type": "CharField", + "db_column": "slug", + "python_type": "str", + "generated": False, "nullable": False, "unique": True, "indexed": True, "default": None, "description": None, "docstring": None, - "constraints": {"ge": MIN_INT, "le": 2147483647}, - "db_field_types": {"": "INT"}, + "constraints": {"max_length": 10}, + "db_field_types": {"": "VARCHAR(10)"}, }, "data_fields": [ { @@ -939,6 +939,8 @@ def test_migrate(mocker: MockerFixture): """ models.py diff with old_models.py - change email pk: id -> email_id + - change product pk field type: IntField -> BigIntField + - change config pk field attribute: max_length=10 -> max_length=20 - add field: Email.address - add fk field: Config.user - drop fk field: Email.user @@ -991,8 +993,8 @@ def test_migrate(mocker: MockerFixture): "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", "ALTER TABLE `email` ADD `address` VARCHAR(200) NOT NULL", - "ALTER TABLE `email` ADD CONSTRAINT `fk_email_config_76a9dc71` FOREIGN KEY (`config_id`) REFERENCES `config` (`id`) ON DELETE CASCADE", - "ALTER TABLE `email` ADD `config_id` INT NOT NULL UNIQUE", + "ALTER TABLE `email` ADD CONSTRAINT `fk_email_config_88e28c1b` FOREIGN KEY (`config_id`) REFERENCES `config` (`slug`) ON DELETE CASCADE", + "ALTER TABLE `email` ADD `config_id` VARCHAR(20) NOT NULL UNIQUE", "ALTER TABLE `email` DROP INDEX `idx_email_company_1c9234`, ADD UNIQUE (`company`)", "ALTER TABLE `configs` RENAME TO `config`", "ALTER TABLE `product` DROP COLUMN `uuid`", @@ -1008,15 +1010,17 @@ def test_migrate(mocker: MockerFixture): "ALTER TABLE `product` ALTER COLUMN `view_num` SET DEFAULT 0", "ALTER TABLE `product` RENAME COLUMN `is_delete` TO `is_deleted`", "ALTER TABLE `product` RENAME COLUMN `is_review` TO `is_reviewed`", + "ALTER TABLE `product` MODIFY COLUMN `id` BIGINT NOT NULL", "ALTER TABLE `user` DROP COLUMN `avatar`", "ALTER TABLE `user` MODIFY COLUMN `password` VARCHAR(100) NOT NULL", "ALTER TABLE `user` MODIFY COLUMN `longitude` DECIMAL(10,8) NOT NULL", "ALTER TABLE `user` ADD UNIQUE INDEX `username` (`username`)", "CREATE TABLE `email_user` (\n `email_id` INT NOT NULL REFERENCES `email` (`email_id`) ON DELETE CASCADE,\n `user_id` INT NOT NULL REFERENCES `user` (`id`) ON DELETE CASCADE\n) CHARACTER SET utf8mb4", "CREATE TABLE IF NOT EXISTS `newmodel` (\n `id` INT NOT NULL PRIMARY KEY AUTO_INCREMENT,\n `name` VARCHAR(50) NOT NULL\n) CHARACTER SET utf8mb4", - "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", + "CREATE TABLE `product_user` (\n `product_id` BIGINT 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` VARCHAR(20) NOT NULL REFERENCES `config` (`slug`) ON DELETE CASCADE\n) CHARACTER SET utf8mb4", "DROP TABLE IF EXISTS `config_category`", + "ALTER TABLE `config` MODIFY COLUMN `slug` VARCHAR(20) NOT NULL", } upgrade_operators = set(Migrate.upgrade_operators) upgrade_more_than_expected = upgrade_operators - expected_upgrade_operators @@ -1037,11 +1041,12 @@ def test_migrate(mocker: MockerFixture): "ALTER TABLE `config` DROP FOREIGN KEY `fk_config_user_17daa970`", "ALTER TABLE `config` ALTER COLUMN `status` SET DEFAULT 1", "ALTER TABLE `config` DROP COLUMN `user_id`", + "ALTER TABLE `config` MODIFY COLUMN `slug` VARCHAR(10) NOT NULL", "ALTER TABLE `config` RENAME TO `configs`", "ALTER TABLE `email` ADD `user_id` INT NOT NULL", "ALTER TABLE `email` DROP COLUMN `address`", "ALTER TABLE `email` DROP COLUMN `config_id`", - "ALTER TABLE `email` DROP FOREIGN KEY `fk_email_config_76a9dc71`", + "ALTER TABLE `email` DROP FOREIGN KEY `fk_email_config_88e28c1b`", "ALTER TABLE `email` RENAME COLUMN `email_id` TO `id`", "ALTER TABLE `email` DROP INDEX `company`, ADD INDEX (`idx_email_company_1c9234`)", "ALTER TABLE `email` DROP INDEX `idx_email_email_4a1a33`", @@ -1056,6 +1061,7 @@ def test_migrate(mocker: MockerFixture): "ALTER TABLE `product` ALTER COLUMN `view_num` DROP DEFAULT", "ALTER TABLE `product` RENAME COLUMN `is_deleted` TO `is_delete`", "ALTER TABLE `product` RENAME COLUMN `is_reviewed` TO `is_review`", + "ALTER TABLE `product` MODIFY COLUMN `id` INT NOT NULL", "ALTER TABLE `user` ADD `avatar` VARCHAR(200) NOT NULL DEFAULT ''", "ALTER TABLE `user` DROP INDEX `username`", "ALTER TABLE `user` MODIFY COLUMN `password` VARCHAR(200) NOT NULL", @@ -1063,7 +1069,7 @@ def test_migrate(mocker: MockerFixture): "DROP TABLE IF EXISTS `newmodel`", "DROP TABLE IF EXISTS `product_user`", "ALTER TABLE `user` MODIFY COLUMN `longitude` DECIMAL(12,9) NOT NULL", - "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", + "CREATE TABLE `config_category` (\n `config_id` VARCHAR(20) NOT NULL REFERENCES `config` (`slug`) 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`", } downgrade_operators = set(Migrate.downgrade_operators) @@ -1081,17 +1087,18 @@ def test_migrate(mocker: MockerFixture): 'ALTER TABLE "category" ADD CONSTRAINT "fk_category_user_110d4c63" FOREIGN KEY ("owner_id") REFERENCES "user" ("id") ON DELETE CASCADE', 'CREATE INDEX IF NOT EXISTS "idx_category_slug_e9bcff" ON "category" USING HASH ("slug")', 'DROP INDEX IF EXISTS "idx_category_slug_e9bcff"', + 'ALTER TABLE "configs" RENAME TO "config"', 'ALTER TABLE "config" DROP COLUMN "name"', 'DROP INDEX IF EXISTS "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', - 'ALTER TABLE "configs" RENAME TO "config"', + 'ALTER TABLE "config" ALTER COLUMN "slug" TYPE VARCHAR(20) USING "slug"::VARCHAR(20)', + 'ALTER TABLE "email" ADD "config_id" VARCHAR(20) NOT NULL UNIQUE', 'ALTER TABLE "email" ADD "address" VARCHAR(200) NOT NULL', 'ALTER TABLE "email" RENAME COLUMN "id" TO "email_id"', 'ALTER TABLE "email" DROP COLUMN "user_id"', - 'ALTER TABLE "email" ADD CONSTRAINT "fk_email_config_76a9dc71" FOREIGN KEY ("config_id") REFERENCES "config" ("id") ON DELETE CASCADE', - 'ALTER TABLE "email" ADD "config_id" INT NOT NULL UNIQUE', + 'ALTER TABLE "email" ADD CONSTRAINT "fk_email_config_88e28c1b" FOREIGN KEY ("config_id") REFERENCES "config" ("slug") ON DELETE CASCADE', 'DROP INDEX IF EXISTS "idx_email_company_1c9234"', 'CREATE UNIQUE INDEX IF NOT EXISTS "uid_email_company_1c9234" ON "email" ("company")', 'DROP INDEX IF EXISTS "uid_product_uuid_d33c18"', @@ -1102,6 +1109,7 @@ def test_migrate(mocker: MockerFixture): 'ALTER TABLE "product" RENAME COLUMN "is_delete" TO "is_deleted"', 'ALTER TABLE "product" ADD "price" DOUBLE PRECISION', 'ALTER TABLE "product" ADD "no" UUID NOT NULL', + 'ALTER TABLE "product" ALTER COLUMN "id" TYPE BIGINT USING "id"::BIGINT', 'ALTER TABLE "user" ALTER COLUMN "password" TYPE VARCHAR(100) USING "password"::VARCHAR(100)', 'ALTER TABLE "user" DROP COLUMN "avatar"', 'ALTER TABLE "user" ALTER COLUMN "longitude" TYPE DECIMAL(10,8) USING "longitude"::DECIMAL(10,8)', @@ -1112,8 +1120,8 @@ def test_migrate(mocker: MockerFixture): 'CREATE TABLE IF NOT EXISTS "newmodel" (\n "id" SERIAL NOT NULL PRIMARY KEY,\n "name" VARCHAR(50) NOT NULL\n);\nCOMMENT ON COLUMN "config"."user_id" IS \'User\'', 'CREATE UNIQUE INDEX IF NOT EXISTS "uid_product_name_869427" ON "product" ("name", "type_db_alias")', 'CREATE UNIQUE INDEX IF NOT EXISTS "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)', + 'CREATE TABLE "product_user" (\n "product_id" BIGINT 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" VARCHAR(20) NOT NULL REFERENCES "config" ("slug") ON DELETE CASCADE\n)', 'DROP TABLE IF EXISTS "config_category"', } upgrade_operators = set(Migrate.upgrade_operators) @@ -1136,11 +1144,12 @@ def test_migrate(mocker: MockerFixture): 'ALTER TABLE "config" DROP CONSTRAINT IF EXISTS "fk_config_user_17daa970"', 'ALTER TABLE "config" RENAME TO "configs"', 'ALTER TABLE "config" DROP COLUMN "user_id"', + 'ALTER TABLE "config" ALTER COLUMN "slug" TYPE VARCHAR(10) USING "slug"::VARCHAR(10)', 'ALTER TABLE "email" ADD "user_id" INT NOT NULL', 'ALTER TABLE "email" DROP COLUMN "address"', 'ALTER TABLE "email" RENAME COLUMN "email_id" TO "id"', 'ALTER TABLE "email" DROP COLUMN "config_id"', - 'ALTER TABLE "email" DROP CONSTRAINT IF EXISTS "fk_email_config_76a9dc71"', + 'ALTER TABLE "email" DROP CONSTRAINT IF EXISTS "fk_email_config_88e28c1b"', 'CREATE INDEX IF NOT EXISTS "idx_email_company_1c9234" ON "email" ("company")', 'DROP INDEX IF EXISTS "uid_email_company_1c9234"', 'ALTER TABLE "product" ADD "uuid" INT NOT NULL UNIQUE', @@ -1151,6 +1160,7 @@ def test_migrate(mocker: MockerFixture): 'ALTER TABLE "product" RENAME COLUMN "is_reviewed" TO "is_review"', 'ALTER TABLE "product" DROP COLUMN "price"', 'ALTER TABLE "product" DROP COLUMN "no"', + 'ALTER TABLE "product" ALTER COLUMN "id" TYPE INT USING "id"::INT', 'ALTER TABLE "user" ADD "avatar" VARCHAR(200) NOT NULL DEFAULT \'\'', 'ALTER TABLE "user" ALTER COLUMN "password" TYPE VARCHAR(200) USING "password"::VARCHAR(200)', 'ALTER TABLE "user" ALTER COLUMN "longitude" TYPE DECIMAL(12,9) USING "longitude"::DECIMAL(12,9)', @@ -1162,7 +1172,7 @@ def test_migrate(mocker: MockerFixture): 'DROP INDEX IF EXISTS "idx_product_no_e4d701"', '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)', + 'CREATE TABLE "config_category" (\n "config_id" VARCHAR(20) NOT NULL REFERENCES "config" ("slug") ON DELETE CASCADE,\n "category_id" INT NOT NULL REFERENCES "category" ("id") ON DELETE CASCADE\n)', 'DROP TABLE IF EXISTS "config_category_map"', } downgrade_operators = set(Migrate.downgrade_operators)