Support migration for specified index. (#203)

aerich/cli.py

@@ -10,12 +10,11 @@ from click import Context, UsageError
 from tomlkit.exceptions import NonExistentKey
 from tortoise import Tortoise
 
+from aerich import Command
+from aerich.enums import Color
 from aerich.exceptions import DowngradeError
 from aerich.utils import add_src_path, get_tortoise_config
-
-from . import Command
-from .enums import Color
-from .version import __version__
+from aerich.version import __version__
 
 CONFIG_DEFAULT_VALUES = {
     "src_folder": ".",

aerich/coder.py

@@ -0,0 +1,31 @@
+import base64
+import json
+import pickle  # nosec: B301
+
+from tortoise.indexes import Index
+
+
+class JsonEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, Index):
+            return {
+                "type": "index",
+                "val": base64.b64encode(pickle.dumps(obj)).decode(),
+            }  # nosec: B301
+        else:
+            return super().default(obj)
+
+
+def object_hook(obj):
+    _type = obj.get("type")
+    if not _type:
+        return obj
+    return pickle.loads(base64.b64decode(obj["val"]))  # nosec: B301
+
+
+def encoder(obj: dict):
+    return json.dumps(obj, cls=JsonEncoder)
+
+
+def decoder(obj: str):
+    return json.loads(obj, object_hook=object_hook)

aerich/ddl/__init__.py

@@ -195,6 +195,12 @@ class BaseDDL:
             table_name=model._meta.db_table,
         )
 
+    def drop_index_by_name(self, model: "Type[Model]", index_name: str):
+        return self._DROP_INDEX_TEMPLATE.format(
+            index_name=index_name,
+            table_name=model._meta.db_table,
+        )
+
     def add_fk(self, model: "Type[Model]", field_describe: dict, reference_table_describe: dict):
         db_table = model._meta.db_table
 

aerich/migrate.py

@@ -1,12 +1,14 @@
 import os
 from datetime import datetime
+from hashlib import md5
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Type
+from typing import Dict, List, Optional, Tuple, Type, Union
 
 import click
 from dictdiffer import diff
 from tortoise import BaseDBAsyncClient, Model, Tortoise
 from tortoise.exceptions import OperationalError
+from tortoise.indexes import Index
 
 from aerich.ddl import BaseDDL
 from aerich.models import MAX_VERSION_LENGTH, Aerich
@@ -32,7 +34,7 @@ class Migrate:
     ddl: BaseDDL
     _last_version_content: Optional[dict] = None
     app: str
-    migrate_location: str
+    migrate_location: Path
     dialect: str
     _db_version: Optional[str] = None
 
@@ -157,6 +159,18 @@ class Migrate:
             else:
                 cls.downgrade_operators.append(operator)
 
+    @classmethod
+    def _handle_indexes(cls, model: Type[Model], indexes: List[Union[Tuple[str], Index]]):
+        ret = []
+        for index in indexes:
+            if isinstance(index, Index):
+                index.__hash__ = lambda self: md5(  # nosec: B303
+                    self.index_name(cls.ddl.schema_generator, model).encode()
+                    + self.__class__.__name__.encode()
+                ).hexdigest()
+            ret.append(index)
+        return ret
+
     @classmethod
     def diff_models(cls, old_models: Dict[str, dict], new_models: Dict[str, dict], upgrade=True):
         """
@@ -192,8 +206,18 @@ class Migrate:
                 new_unique_together = set(
                     map(lambda x: tuple(x), new_model_describe.get("unique_together"))
                 )
-                old_indexes = set(map(lambda x: tuple(x), old_model_describe.get("indexes", [])))
-                new_indexes = set(map(lambda x: tuple(x), new_model_describe.get("indexes", [])))
+                old_indexes = set(
+                    map(
+                        lambda x: x if isinstance(x, Index) else tuple(x),
+                        cls._handle_indexes(model, old_model_describe.get("indexes", [])),
+                    )
+                )
+                new_indexes = set(
+                    map(
+                        lambda x: x if isinstance(x, Index) else tuple(x),
+                        cls._handle_indexes(model, new_model_describe.get("indexes", [])),
+                    )
+                )
                 old_pk_field = old_model_describe.get("pk_field")
                 new_pk_field = new_model_describe.get("pk_field")
                 # pk field
@@ -463,12 +487,18 @@ class Migrate:
         return ret
 
     @classmethod
-    def _drop_index(cls, model: Type[Model], fields_name: Tuple[str], unique=False):
+    def _drop_index(cls, model: Type[Model], fields_name: Union[Tuple[str], Index], unique=False):
+        if isinstance(fields_name, Index):
+            return cls.ddl.drop_index_by_name(
+                model, fields_name.index_name(cls.ddl.schema_generator, model)
+            )
         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: Tuple[str], unique=False):
+    def _add_index(cls, model: Type[Model], fields_name: Union[Tuple[str], Index], unique=False):
+        if isinstance(fields_name, Index):
+            return fields_name.get_sql(cls.ddl.schema_generator, model, False)
         fields_name = cls._resolve_fk_fields_name(model, fields_name)
         return cls.ddl.add_index(model, fields_name, unique)
 

aerich/models.py

@@ -1,12 +1,14 @@
 from tortoise import Model, fields
 
+from aerich.coder import decoder, encoder
+
 MAX_VERSION_LENGTH = 255
 
 
 class Aerich(Model):
     version = fields.CharField(max_length=MAX_VERSION_LENGTH)
     app = fields.CharField(max_length=20)
-    content = fields.JSONField()
+    content = fields.JSONField(encoder=encoder, decoder=decoder)
 
     class Meta:
         ordering = ["-id"]

aerich/version.py

@@ -1 +1 @@
-__version__ = "0.6.1"
+__version__ = "0.6.2"