From 559d3707e538f72f04589b796083a99bfe0e97e4 Mon Sep 17 00:00:00 2001 From: long2ice Date: Mon, 18 May 2020 11:15:11 +0800 Subject: [PATCH] Support indexes and unique_together --- CHANGELOG.rst | 3 ++ aerich/__init__.py | 2 +- aerich/cli.py | 9 ++++-- aerich/migrate.py | 74 ++++++++++++++++++++++++++++++++++++---------- 4 files changed, 68 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index cdc3089..bde3321 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -4,6 +4,9 @@ ChangeLog 0.1 === +0.1.3 +----- +- Support indexes and unique_together. 0.1.2 ----- - Now aerich support m2m. diff --git a/aerich/__init__.py b/aerich/__init__.py index b3f4756..ae73625 100644 --- a/aerich/__init__.py +++ b/aerich/__init__.py @@ -1 +1 @@ -__version__ = "0.1.2" +__version__ = "0.1.3" diff --git a/aerich/cli.py b/aerich/cli.py index a6e6778..dd7bdc9 100644 --- a/aerich/cli.py +++ b/aerich/cli.py @@ -6,7 +6,7 @@ from enum import Enum import asyncclick as click from asyncclick import Context, UsageError -from tortoise import Tortoise, generate_schema_for_client +from tortoise import Tortoise, generate_schema_for_client, ConfigurationError from aerich.migrate import Migrate from aerich.utils import get_app_connection, get_tortoise_config @@ -54,7 +54,10 @@ async def cli(ctx: Context, config, app, name): ctx.obj["app"] = app if invoked_subcommand != "init-db": - await Migrate.init_with_old_models(tortoise_config, app, location) + try: + await Migrate.init_with_old_models(tortoise_config, app, location) + except ConfigurationError: + raise UsageError(ctx=ctx, message='You must exec ini-db first') @cli.command(help="Generate migrate changes file.") @@ -149,7 +152,7 @@ def history(ctx): ) @click.pass_context async def init( - ctx: Context, tortoise_orm, location, + ctx: Context, tortoise_orm, location, ): config = ctx.obj["config"] name = ctx.obj["name"] diff --git a/aerich/migrate.py b/aerich/migrate.py index 74feae6..43fb942 100644 --- a/aerich/migrate.py +++ b/aerich/migrate.py @@ -25,8 +25,8 @@ from aerich.utils import get_app_connection class Migrate: upgrade_operators: List[str] = [] downgrade_operators: List[str] = [] - _upgrade_fk_m2m_operators: List[str] = [] - _downgrade_fk_m2m_operators: List[str] = [] + _upgrade_fk_m2m_index_operators: List[str] = [] + _downgrade_fk_m2m_index_operators: List[str] = [] _upgrade_m2m: List[str] = [] _downgrade_m2m: List[str] = [] @@ -126,18 +126,18 @@ class Migrate: """ if upgrade: if fk: - cls._upgrade_fk_m2m_operators.append(operator) + cls._upgrade_fk_m2m_index_operators.append(operator) else: cls.upgrade_operators.append(operator) else: if fk: - cls._downgrade_fk_m2m_operators.append(operator) + cls._downgrade_fk_m2m_index_operators.append(operator) else: cls.downgrade_operators.append(operator) @classmethod def cp_models( - cls, app: str, model_files: List[str], old_model_file, + cls, app: str, model_files: List[str], old_model_file, ): """ cp currents models to old_model_files @@ -187,7 +187,7 @@ class Migrate: @classmethod def _diff_models( - cls, old_models: Dict[str, Type[Model]], new_models: Dict[str, Type[Model]], upgrade=True + cls, old_models: Dict[str, Type[Model]], new_models: Dict[str, Type[Model]], upgrade=True ): """ diff models and add operators @@ -223,8 +223,15 @@ class Migrate: :param upgrade: :return: """ + old_indexes = old_model._meta.indexes + new_indexes = new_model._meta.indexes + + old_unique_together = old_model._meta.unique_together + new_unique_together = new_model._meta.unique_together + old_fields_map = old_model._meta.fields_map new_fields_map = new_model._meta.fields_map + old_keys = old_fields_map.keys() new_keys = new_fields_map.keys() for new_key in new_keys: @@ -241,13 +248,13 @@ class Migrate: old_field = old_fields_map.get(new_key) if old_field.index and not new_field.index: cls._add_operator( - cls._remove_index(old_model, old_field), + cls._remove_index(old_model, [old_field.model_field_name], old_field.unique), upgrade, isinstance(old_field, (ForeignKeyFieldInstance, ManyToManyFieldInstance)), ) elif new_field.index and not old_field.index: cls._add_operator( - cls._add_index(new_model, new_field), + cls._add_index(new_model, [new_field.model_field_name], new_field.unique), upgrade, isinstance(new_field, (ForeignKeyFieldInstance, ManyToManyFieldInstance)), ) @@ -261,13 +268,48 @@ class Migrate: isinstance(field, (ForeignKeyFieldInstance, ManyToManyFieldInstance)), ) - @classmethod - def _remove_index(cls, model: Type[Model], field: Field): - return cls.ddl.drop_index(model, [field.model_field_name], field.unique) + for new_index in new_indexes: + if new_index not in old_indexes: + cls._add_operator( + cls._add_index(new_model, new_index, ), upgrade + ) + for old_index in old_indexes: + if old_index not in new_indexes: + cls._add_operator( + cls._remove_index(old_model, old_index), upgrade + ) + + for new_unique in new_unique_together: + if new_unique not in old_unique_together: + cls._add_operator( + cls._add_index(new_model, new_unique, unique=True), upgrade + ) + + for old_unique in old_unique_together: + if old_unique not in new_unique_together: + cls._add_operator( + cls._remove_index(old_model, old_unique, unique=True), upgrade + ) @classmethod - def _add_index(cls, model: Type[Model], field: Field): - return cls.ddl.add_index(model, [field.model_field_name], field.unique) + def _resolve_fk_fields_name(cls, model: Type[Model], fields_name: List[str]): + ret = [] + for field_name in fields_name: + if field_name in model._meta.fk_fields: + ret.append(field_name + '_id') + else: + ret.append(field_name) + return ret + + @classmethod + def _remove_index(cls, model: Type[Model], fields_name: List[str], unique=False): + fields_name = cls._resolve_fk_fields_name(model, fields_name) + return cls.ddl.drop_index(model, fields_name, unique) + + @classmethod + def _add_index(cls, model: Type[Model], fields_name: List[str], unique=False): + fields_name = cls._resolve_fk_fields_name(model, fields_name) + return cls.ddl.add_index(model, fields_name, unique) @classmethod def _exclude_field(cls, field: Field, upgrade=False): @@ -331,16 +373,16 @@ class Migrate: @classmethod def _merge_operators(cls): """ - fk/m2m must be last when add,first when drop + fk/m2m/index must be last when add,first when drop :return: """ - for _upgrade_fk_m2m_operator in cls._upgrade_fk_m2m_operators: + for _upgrade_fk_m2m_operator in cls._upgrade_fk_m2m_index_operators: if "ADD" in _upgrade_fk_m2m_operator or "CREATE" in _upgrade_fk_m2m_operator: cls.upgrade_operators.append(_upgrade_fk_m2m_operator) else: cls.upgrade_operators.insert(0, _upgrade_fk_m2m_operator) - for _downgrade_fk_m2m_operator in cls._downgrade_fk_m2m_operators: + for _downgrade_fk_m2m_operator in cls._downgrade_fk_m2m_index_operators: if "ADD" in _downgrade_fk_m2m_operator or "CREATE" in _downgrade_fk_m2m_operator: cls.downgrade_operators.append(_downgrade_fk_m2m_operator) else: