Merge pull request #34 from moubctez/postgresql_unique

PostgreSQL add/drop index/unique

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(

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")

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

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',