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():