From 1b440477a2e729f6dcbeab6157c8786c5c41893d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Ciarcin=CC=81ski?= Date: Thu, 6 Aug 2020 17:45:56 +0200 Subject: [PATCH 1/4] PostgreSQL add/drop index/unique --- aerich/ddl/postgres/__init__.py | 27 ++++++++++++++++++++++++++- tests/diff_models.py | 2 +- tests/test_ddl.py | 12 ++++++++++++ 3 files changed, 39 insertions(+), 2 deletions(-) diff --git a/aerich/ddl/postgres/__init__.py b/aerich/ddl/postgres/__init__.py index fd10fbc..910c1c0 100644 --- a/aerich/ddl/postgres/__init__.py +++ b/aerich/ddl/postgres/__init__.py @@ -1,4 +1,4 @@ -from typing import Type +from typing import List, Type from tortoise import Model from tortoise.backends.asyncpg.schema_generator import AsyncpgSchemaGenerator @@ -10,6 +10,12 @@ from aerich.ddl import BaseDDL class PostgresDDL(BaseDDL): schema_generator_cls = AsyncpgSchemaGenerator DIALECT = AsyncpgSchemaGenerator.DIALECT + _ADD_INDEX_TEMPLATE = 'CREATE INDEX "{index_name}" ON "{table_name}" ({column_names})' + _ADD_UNIQUE_TEMPLATE = ( + 'ALTER TABLE "{table_name}" ADD CONSTRAINT "{index_name}" UNIQUE ({column_names})' + ) + _DROP_INDEX_TEMPLATE = 'DROP INDEX "{index_name}"' + _DROP_UNIQUE_TEMPLATE = 'ALTER TABLE "{table_name}" DROP CONSTRAINT "{index_name}"' _ALTER_DEFAULT_TEMPLATE = 'ALTER TABLE "{table_name}" ALTER COLUMN "{column}" {default}' _ALTER_NULL_TEMPLATE = 'ALTER TABLE "{table_name}" ALTER COLUMN "{column}" {set_drop} NOT NULL' _MODIFY_COLUMN_TEMPLATE = 'ALTER TABLE "{table_name}" ALTER COLUMN "{column}" TYPE {datatype}' @@ -41,6 +47,25 @@ class PostgresDDL(BaseDDL): datatype=field_object.get_for_dialect(self.DIALECT, "SQL_TYPE"), ) + def add_index(self, model: "Type[Model]", field_names: List[str], unique=False): + template = self._ADD_UNIQUE_TEMPLATE if unique else self._ADD_INDEX_TEMPLATE + return template.format( + index_name=self.schema_generator._generate_index_name( + "uid" if unique else "idx", model, field_names + ), + table_name=model._meta.db_table, + column_names=", ".join([self.schema_generator.quote(f) for f in field_names]), + ) + + def drop_index(self, model: "Type[Model]", field_names: List[str], unique=False): + template = self._DROP_UNIQUE_TEMPLATE if unique else self._DROP_INDEX_TEMPLATE + return template.format( + index_name=self.schema_generator._generate_index_name( + "uid" if unique else "idx", model, field_names + ), + table_name=model._meta.db_table, + ) + def set_comment(self, model: "Type[Model]", field_object: Field): db_table = model._meta.db_table return self._SET_COMMENT_TEMPLATE.format( diff --git a/tests/diff_models.py b/tests/diff_models.py index 7f5af93..02b644e 100644 --- a/tests/diff_models.py +++ b/tests/diff_models.py @@ -22,7 +22,7 @@ class Status(IntEnum): class User(Model): - username = fields.CharField(max_length=20,) + username = fields.CharField(max_length=20) password = fields.CharField(max_length=200) last_login = fields.DatetimeField(description="Last Login", default=datetime.datetime.now) is_active = fields.BooleanField(default=True, description="Is Active") diff --git a/tests/test_ddl.py b/tests/test_ddl.py index dcadbf7..3588f38 100644 --- a/tests/test_ddl.py +++ b/tests/test_ddl.py @@ -146,6 +146,12 @@ def test_add_index(): assert ( index_u == "ALTER TABLE `category` ADD UNIQUE INDEX `uid_category_name_8b0cb9` (`name`)" ) + elif isinstance(Migrate.ddl, PostgresDDL): + assert index == 'CREATE INDEX "idx_category_name_8b0cb9" ON "category" ("name")' + assert ( + index_u + == 'ALTER TABLE "category" ADD CONSTRAINT "uid_category_name_8b0cb9" UNIQUE ("name")' + ) else: assert index == 'ALTER TABLE "category" ADD INDEX "idx_category_name_8b0cb9" ("name")' assert ( @@ -155,10 +161,16 @@ def test_add_index(): def test_drop_index(): ret = Migrate.ddl.drop_index(Category, ["name"]) + ret_u = Migrate.ddl.drop_index(Category, ["name"], True) if isinstance(Migrate.ddl, MysqlDDL): assert ret == "ALTER TABLE `category` DROP INDEX `idx_category_name_8b0cb9`" + assert ret_u == "ALTER TABLE `category` DROP INDEX `uid_category_name_8b0cb9`" + elif isinstance(Migrate.ddl, PostgresDDL): + assert ret == 'DROP INDEX "idx_category_name_8b0cb9"' + assert ret_u == 'ALTER TABLE "category" DROP CONSTRAINT "uid_category_name_8b0cb9"' else: assert ret == 'ALTER TABLE "category" DROP INDEX "idx_category_name_8b0cb9"' + assert ret_u == 'ALTER TABLE "category" DROP INDEX "uid_category_name_8b0cb9"' def test_add_fk(): From 198e4e00327197b54b891777787ff0a25e581313 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Ciarcin=CC=81ski?= Date: Thu, 6 Aug 2020 17:49:26 +0200 Subject: [PATCH 2/4] Save 1 space --- aerich/ddl/postgres/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aerich/ddl/postgres/__init__.py b/aerich/ddl/postgres/__init__.py index 910c1c0..055b041 100644 --- a/aerich/ddl/postgres/__init__.py +++ b/aerich/ddl/postgres/__init__.py @@ -12,7 +12,7 @@ class PostgresDDL(BaseDDL): DIALECT = AsyncpgSchemaGenerator.DIALECT _ADD_INDEX_TEMPLATE = 'CREATE INDEX "{index_name}" ON "{table_name}" ({column_names})' _ADD_UNIQUE_TEMPLATE = ( - 'ALTER TABLE "{table_name}" ADD CONSTRAINT "{index_name}" UNIQUE ({column_names})' + 'ALTER TABLE "{table_name}" ADD CONSTRAINT "{index_name}" UNIQUE ({column_names})' ) _DROP_INDEX_TEMPLATE = 'DROP INDEX "{index_name}"' _DROP_UNIQUE_TEMPLATE = 'ALTER TABLE "{table_name}" DROP CONSTRAINT "{index_name}"' From cda9bd1c47f20c70928a6dc11305593446d285b6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Ciarcin=CC=81ski?= Date: Thu, 6 Aug 2020 17:50:28 +0200 Subject: [PATCH 3/4] Save 1 space in tests --- tests/test_ddl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_ddl.py b/tests/test_ddl.py index 3588f38..e53b289 100644 --- a/tests/test_ddl.py +++ b/tests/test_ddl.py @@ -150,7 +150,7 @@ def test_add_index(): assert index == 'CREATE INDEX "idx_category_name_8b0cb9" ON "category" ("name")' assert ( index_u - == 'ALTER TABLE "category" ADD CONSTRAINT "uid_category_name_8b0cb9" UNIQUE ("name")' + == 'ALTER TABLE "category" ADD CONSTRAINT "uid_category_name_8b0cb9" UNIQUE ("name")' ) else: assert index == 'ALTER TABLE "category" ADD INDEX "idx_category_name_8b0cb9" ("name")' From 8674142ba8a3e293f8aa998c7c70cd29834515a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Ciarcin=CC=81ski?= Date: Thu, 6 Aug 2020 17:59:12 +0200 Subject: [PATCH 4/4] Fix test_migrate --- tests/test_migrate.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_migrate.py b/tests/test_migrate.py index ed53778..3a51e1b 100644 --- a/tests/test_migrate.py +++ b/tests/test_migrate.py @@ -1,6 +1,7 @@ from tortoise import Tortoise from aerich.ddl.mysql import MysqlDDL +from aerich.ddl.postgres import PostgresDDL from aerich.migrate import Migrate @@ -19,6 +20,15 @@ def test_migrate(): "ALTER TABLE `category` DROP COLUMN `name`", "ALTER TABLE `user` DROP INDEX `uid_user_usernam_9987ab`", ] + elif isinstance(Migrate.ddl, PostgresDDL): + assert Migrate.upgrade_operators == [ + 'ALTER TABLE "category" ADD "name" VARCHAR(200) NOT NULL', + 'ALTER TABLE "user" ADD CONSTRAINT "uid_user_usernam_9987ab" UNIQUE ("username")', + ] + assert Migrate.downgrade_operators == [ + 'ALTER TABLE "category" DROP COLUMN "name"', + 'ALTER TABLE "user" DROP CONSTRAINT "uid_user_usernam_9987ab"', + ] else: assert Migrate.upgrade_operators == [ 'ALTER TABLE "category" ADD "name" VARCHAR(200) NOT NULL',