From 99e72583f7a57a2b3490c9a5e6d7c2e5e95de719 Mon Sep 17 00:00:00 2001 From: Waket Zheng Date: Sun, 22 Dec 2024 00:23:11 +0800 Subject: [PATCH 1/3] refactor: remove puzzle vars --- aerich/migrate.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/aerich/migrate.py b/aerich/migrate.py index edca5de..e12cc7b 100644 --- a/aerich/migrate.py +++ b/aerich/migrate.py @@ -294,7 +294,7 @@ def diff_models( _aerich = f"{cls.app}.{cls._aerich}" old_models.pop(_aerich, None) new_models.pop(_aerich, None) - models_with_rename_field: Set[str] = set() + models_with_rename_field: Set[str] = set() # models that trigger the click.prompt for new_model_str, new_model_describe in new_models.items(): model = cls._get_model(new_model_describe["name"].split(".")[1]) @@ -367,12 +367,10 @@ def diff_models( new_data_fields_name = cast(List[str], [i.get("name") for i in new_data_fields]) # add fields or rename fields - rename_fields: Dict[str, str] = {} for new_data_field_name in set(new_data_fields_name).difference( set(old_data_fields_name) ): new_data_field = cls.get_field_by_name(new_data_field_name, new_data_fields) - model_rename_fields = cls._rename_fields.get(new_model_str) is_rename = False field_type = new_data_field.get("field_type") db_column = new_data_field.get("db_column") @@ -397,8 +395,11 @@ def diff_models( and old_data_field_name not in new_data_fields_name ): if upgrade: - if old_data_field_name in rename_fields or ( - new_data_field_name in rename_fields.values() + if ( + rename_fields := cls._rename_fields.get(new_model_str) + ) and ( + old_data_field_name in rename_fields + or new_data_field_name in rename_fields.values() ): continue prefix = f"({new_model_str}) " @@ -415,22 +416,18 @@ def diff_models( show_choices=True, ) if is_rename: + if rename_fields is None: + rename_fields = cls._rename_fields[new_model_str] = {} rename_fields[old_data_field_name] = new_data_field_name else: is_rename = False - if model_rename_fields and ( - rename_to := model_rename_fields.get(new_data_field_name) + if rename_to := cls._rename_fields.get(new_model_str, {}).get( + new_data_field_name ): is_rename = True if rename_to != old_data_field_name: continue if is_rename: - if upgrade: - if new_model_str not in cls._rename_fields: - cls._rename_fields[new_model_str] = {} - cls._rename_fields[new_model_str][ - old_data_field_name - ] = new_data_field_name # only MySQL8+ has rename syntax if ( cls.dialect == "mysql" From f93faa8afb03a7eab46f952bef718c9eda0dee53 Mon Sep 17 00:00:00 2001 From: Waket Zheng Date: Sun, 22 Dec 2024 00:23:47 +0800 Subject: [PATCH 2/3] fix: add o2o field does not create constraint when migrating (#396) * fix: add o2o field does not create constraint when migrating * Add testcase and update changelog * docs: update migrating list * refactor: use `_handle_o2o_fields` instead of `is_o2o=True` * Remove unused line --- CHANGELOG.md | 1 + aerich/migrate.py | 111 +++++++++++++++++++++++++++++------------- tests/models.py | 3 ++ tests/test_migrate.py | 9 ++++ 4 files changed, 90 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2e985ce..26d42eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### [0.8.1](Unreleased) #### Fixed +- fix: add o2o field does not create constraint when migrating. (#396) - fix: intermediate table for m2m relation not created. (#394) - Migrate add m2m field with custom through generate duplicated table. (#393) - Migrate drop the wrong m2m field when model have multi m2m fields. (#376) diff --git a/aerich/migrate.py b/aerich/migrate.py index a24b72f..65a48da 100644 --- a/aerich/migrate.py +++ b/aerich/migrate.py @@ -281,6 +281,68 @@ def _handle_m2m_fields( if add: cls._add_operator(cls.drop_m2m(table), upgrade, True) + @classmethod + def _handle_relational( + cls, + key: str, + old_model_describe: Dict, + new_model_describe: Dict, + model: Type[Model], + old_models: Dict, + new_models: Dict, + upgrade=True, + ) -> None: + old_fk_fields = cast(List[dict], old_model_describe.get(key)) + new_fk_fields = cast(List[dict], new_model_describe.get(key)) + + old_fk_fields_name: List[str] = [i.get("name", "") for i in old_fk_fields] + new_fk_fields_name: List[str] = [i.get("name", "") for i in new_fk_fields] + + # add + for new_fk_field_name in set(new_fk_fields_name).difference(set(old_fk_fields_name)): + fk_field = cls.get_field_by_name(new_fk_field_name, new_fk_fields) + if fk_field.get("db_constraint"): + ref_describe = cast(dict, new_models[fk_field["python_type"]]) + sql = cls._add_fk(model, fk_field, ref_describe) + cls._add_operator(sql, upgrade, fk_m2m_index=True) + # drop + for old_fk_field_name in set(old_fk_fields_name).difference(set(new_fk_fields_name)): + old_fk_field = cls.get_field_by_name(old_fk_field_name, cast(List[dict], old_fk_fields)) + if old_fk_field.get("db_constraint"): + ref_describe = cast(dict, old_models[old_fk_field["python_type"]]) + sql = cls._drop_fk(model, old_fk_field, ref_describe) + cls._add_operator(sql, upgrade, fk_m2m_index=True) + + @classmethod + def _handle_fk_fields( + cls, + old_model_describe: Dict, + new_model_describe: Dict, + model: Type[Model], + old_models: Dict, + new_models: Dict, + upgrade=True, + ) -> None: + key = "fk_fields" + cls._handle_relational( + key, old_model_describe, new_model_describe, model, old_models, new_models, upgrade + ) + + @classmethod + def _handle_o2o_fields( + cls, + old_model_describe: Dict, + new_model_describe: Dict, + model: Type[Model], + old_models: Dict, + new_models: Dict, + upgrade=True, + ) -> None: + key = "o2o_fields" + cls._handle_relational( + key, old_model_describe, new_model_describe, model, old_models, new_models, upgrade + ) + @classmethod def diff_models( cls, old_models: Dict[str, dict], new_models: Dict[str, dict], upgrade=True @@ -334,6 +396,13 @@ def diff_models( # current only support rename pk if action == "change" and option == "name": cls._add_operator(cls._rename_field(model, *change), upgrade) + # fk fields + args = (old_model_describe, new_model_describe, model, old_models, new_models) + cls._handle_fk_fields(*args, upgrade=upgrade) + # o2o fields + cls._handle_o2o_fields(*args, upgrade=upgrade) + old_o2o_columns = [i["raw_field"] for i in old_model_describe.get("o2o_fields", [])] + new_o2o_columns = [i["raw_field"] for i in new_model_describe.get("o2o_fields", [])] # m2m fields cls._handle_m2m_fields( old_model_describe, new_model_describe, model, new_models, upgrade @@ -424,7 +493,10 @@ def diff_models( ), upgrade, ) - if new_data_field["indexed"]: + if ( + new_data_field["indexed"] + and new_data_field["db_column"] not in new_o2o_columns + ): cls._add_operator( cls._add_index( model, (new_data_field["db_column"],), new_data_field["unique"] @@ -447,7 +519,10 @@ def diff_models( cls._remove_field(model, db_column), upgrade, ) - if old_data_field["indexed"]: + if ( + old_data_field["indexed"] + and old_data_field["db_column"] not in old_o2o_columns + ): is_unique_field = old_data_field.get("unique") cls._add_operator( cls._drop_index(model, {db_column}, is_unique_field), @@ -455,38 +530,6 @@ def diff_models( True, ) - old_fk_fields = cast(List[dict], old_model_describe.get("fk_fields")) - new_fk_fields = cast(List[dict], new_model_describe.get("fk_fields")) - - old_fk_fields_name: List[str] = [i.get("name", "") for i in old_fk_fields] - new_fk_fields_name: List[str] = [i.get("name", "") for i in new_fk_fields] - - # add fk - for new_fk_field_name in set(new_fk_fields_name).difference( - set(old_fk_fields_name) - ): - fk_field = cls.get_field_by_name(new_fk_field_name, new_fk_fields) - if fk_field.get("db_constraint"): - ref_describe = cast(dict, new_models[fk_field["python_type"]]) - cls._add_operator( - cls._add_fk(model, fk_field, ref_describe), - upgrade, - fk_m2m_index=True, - ) - # drop fk - for old_fk_field_name in set(old_fk_fields_name).difference( - set(new_fk_fields_name) - ): - old_fk_field = cls.get_field_by_name( - old_fk_field_name, cast(List[dict], old_fk_fields) - ) - if old_fk_field.get("db_constraint"): - ref_describe = cast(dict, old_models[old_fk_field["python_type"]]) - cls._add_operator( - cls._drop_fk(model, old_fk_field, ref_describe), - upgrade, - fk_m2m_index=True, - ) # change fields for field_name in set(new_data_fields_name).intersection(set(old_data_fields_name)): old_data_field = cls.get_field_by_name(field_name, old_data_fields) diff --git a/tests/models.py b/tests/models.py index e4ea87c..a9e0a15 100644 --- a/tests/models.py +++ b/tests/models.py @@ -40,6 +40,7 @@ class Email(Model): is_primary = fields.BooleanField(default=False) address = fields.CharField(max_length=200) users: fields.ManyToManyRelation[User] = fields.ManyToManyField("models.User") + config: fields.OneToOneRelation["Config"] = fields.OneToOneField("models.Config") def default_name(): @@ -91,6 +92,8 @@ class Config(Model): "models.User", description="User" ) + email: fields.OneToOneRelation["Email"] + class NewModel(Model): name = fields.CharField(max_length=50) diff --git a/tests/test_migrate.py b/tests/test_migrate.py index cab0fa9..7d324db 100644 --- a/tests/test_migrate.py +++ b/tests/test_migrate.py @@ -875,6 +875,7 @@ def test_migrate(mocker: MockerFixture): - drop field: User.avatar - add index: Email.email - add many to many: Email.users + - add one to one: Email.config - remove unique: Category.title - add unique: User.username - change column: length User.password @@ -914,6 +915,8 @@ def test_migrate(mocker: MockerFixture): "ALTER TABLE `config` ALTER COLUMN `status` DROP DEFAULT", "ALTER TABLE `config` MODIFY COLUMN `value` JSON NOT NULL", "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 `configs` RENAME TO `config`", "ALTER TABLE `product` DROP COLUMN `uuid`", "ALTER TABLE `product` DROP INDEX `uuid`", @@ -954,6 +957,8 @@ def test_migrate(mocker: MockerFixture): "ALTER TABLE `config` ALTER COLUMN `status` SET DEFAULT 1", "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 `config` RENAME TO `configs`", "ALTER TABLE `product` RENAME COLUMN `pic` TO `image`", "ALTER TABLE `email` RENAME COLUMN `email_id` TO `id`", @@ -1007,6 +1012,8 @@ def test_migrate(mocker: MockerFixture): 'ALTER TABLE "email" ADD "address" VARCHAR(200) NOT NULL', 'ALTER TABLE "email" RENAME COLUMN "id" TO "email_id"', 'ALTER TABLE "email" ALTER COLUMN "is_primary" TYPE BOOL USING "is_primary"::BOOL', + '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', 'DROP INDEX IF EXISTS "uid_product_uuid_d33c18"', 'ALTER TABLE "product" DROP COLUMN "uuid"', 'ALTER TABLE "product" ALTER COLUMN "view_num" SET DEFAULT 0', @@ -1048,6 +1055,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 "email" DROP COLUMN "config_id"', + 'ALTER TABLE "email" DROP CONSTRAINT IF EXISTS "fk_email_config_76a9dc71"', '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', From 7d22518c74b91c805e64b56051b58670013abf95 Mon Sep 17 00:00:00 2001 From: Waket Zheng Date: Sun, 22 Dec 2024 00:24:18 +0800 Subject: [PATCH 3/3] Fix create/drop indexes in every migration (#377) * Add `__eq__` method for `Index`instances * tests: add Index test case * refactor: compare index instances before set hash and eq func to class * fix: sort fields when generating index hash * docs: update changlog * fix style issue * refactor: use CustomIndex instead of postgres special HashIndex * Check tortoise version before patch Index * Add comment * Add comment for why > work --------- Co-authored-by: dbf --- CHANGELOG.md | 1 + aerich/migrate.py | 38 ++++++++++++++++++++++---------------- tests/indexes.py | 7 +++++++ tests/models.py | 7 +++++++ tests/old_models.py | 6 ++++++ tests/test_migrate.py | 4 +++- 6 files changed, 46 insertions(+), 17 deletions(-) create mode 100644 tests/indexes.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 26d42eb..1bc3bdb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ - Fix configuration file reading error when containing Chinese characters. (#286) - sqlite: failed to create/drop index. (#302) - PostgreSQL: Cannot drop constraint after deleting or rename FK on a model. (#378) +- Fix create/drop indexes in every migration. (#377) - Sort m2m fields before comparing them with diff. (#271) #### Changed diff --git a/aerich/migrate.py b/aerich/migrate.py index 65a48da..2adbac9 100644 --- a/aerich/migrate.py +++ b/aerich/migrate.py @@ -1,4 +1,5 @@ -import hashlib +from __future__ import annotations + import importlib import os from datetime import datetime @@ -6,6 +7,7 @@ from typing import Dict, Iterable, List, Optional, Set, Tuple, Type, Union, cast import asyncclick as click +import tortoise from dictdiffer import diff from tortoise import BaseDBAsyncClient, Model, Tortoise from tortoise.exceptions import OperationalError @@ -202,21 +204,25 @@ def _add_operator(cls, operator: str, upgrade=True, fk_m2m_index=False) -> None: @classmethod def _handle_indexes(cls, model: Type[Model], indexes: List[Union[Tuple[str], Index]]) -> list: - ret: list = [] - - def index_hash(self) -> str: - h = hashlib.new("MD5", usedforsecurity=False) # type:ignore[call-arg] - h.update( - self.index_name(cls.ddl.schema_generator, model).encode() - + self.__class__.__name__.encode() - ) - return h.hexdigest() - - for index in indexes: - if isinstance(index, Index): - index.__hash__ = index_hash # type:ignore[method-assign,assignment] - ret.append(index) - return ret + if tortoise.__version__ > "0.22.2": + # The min version of tortoise is '0.11.0', so we can compare it by a `>`, + # tortoise>0.22.2 have __eq__/__hash__ with Index class since 313ee76. + return indexes + if index_classes := set(index.__class__ for index in indexes if isinstance(index, Index)): + # Leave magic patch here to compare with older version of tortoise-orm + # TODO: limit tortoise>0.22.2 in pyproject.toml and remove this function when v0.9.0 released + for index_cls in index_classes: + if index_cls(fields=("id",)) != index_cls(fields=("id",)): + + def _hash(self) -> int: + return hash((tuple(sorted(self.fields)), self.name, self.expressions)) + + def _eq(self, other) -> bool: + return type(self) is type(other) and self.__dict__ == other.__dict__ + + setattr(index_cls, "__hash__", _hash) + setattr(index_cls, "__eq__", _eq) + return indexes @classmethod def _get_indexes(cls, model, model_describe: dict) -> Set[Union[Index, Tuple[str, ...]]]: diff --git a/tests/indexes.py b/tests/indexes.py new file mode 100644 index 0000000..83959f8 --- /dev/null +++ b/tests/indexes.py @@ -0,0 +1,7 @@ +from tortoise.indexes import Index + + +class CustomIndex(Index): + def __init__(self, *args, **kw) -> None: + super().__init__(*args, **kw) + self._foo = "" diff --git a/tests/models.py b/tests/models.py index a9e0a15..1316f84 100644 --- a/tests/models.py +++ b/tests/models.py @@ -3,6 +3,9 @@ from enum import IntEnum from tortoise import Model, fields +from tortoise.indexes import Index + +from tests.indexes import CustomIndex class ProductType(IntEnum): @@ -33,6 +36,10 @@ class User(Model): products: fields.ManyToManyRelation["Product"] + class Meta: + # reverse indexes elements + indexes = [CustomIndex(fields=("is_superuser",)), Index(fields=("username", "is_active"))] + class Email(Model): email_id = fields.IntField(primary_key=True) diff --git a/tests/old_models.py b/tests/old_models.py index 5444232..5e73020 100644 --- a/tests/old_models.py +++ b/tests/old_models.py @@ -2,6 +2,9 @@ from enum import IntEnum from tortoise import Model, fields +from tortoise.indexes import Index + +from tests.indexes import CustomIndex class ProductType(IntEnum): @@ -31,6 +34,9 @@ class User(Model): intro = fields.TextField(default="") longitude = fields.DecimalField(max_digits=12, decimal_places=9) + class Meta: + indexes = [Index(fields=("username", "is_active")), CustomIndex(fields=("is_superuser",))] + class Email(Model): email = fields.CharField(max_length=200) diff --git a/tests/test_migrate.py b/tests/test_migrate.py index 7d324db..8e6a8fa 100644 --- a/tests/test_migrate.py +++ b/tests/test_migrate.py @@ -3,6 +3,7 @@ import pytest import tortoise from pytest_mock import MockerFixture +from tortoise.indexes import Index from aerich.ddl.mysql import MysqlDDL from aerich.ddl.postgres import PostgresDDL @@ -10,6 +11,7 @@ from aerich.exceptions import NotSupportError from aerich.migrate import MIGRATE_TEMPLATE, Migrate from aerich.utils import get_models_describe +from tests.indexes import CustomIndex # tortoise-orm>=0.21 changes IntField constraints # from {"ge": 1, "le": 2147483647} to {"ge": -2147483648, "le": 2147483647} @@ -608,7 +610,7 @@ "description": None, "docstring": None, "unique_together": [], - "indexes": [], + "indexes": [Index(fields=("username", "is_active")), CustomIndex(fields=("is_superuser",))], "pk_field": { "name": "id", "field_type": "IntField",