Support migration for specified index. (#203)

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## 0.6
 
+### 0.6.2
+
+- Support migration for specified index. (#203)
+
 ### 0.6.1
 
 - Fix `pyproject.toml` not existing error. (#217)

README.md

@@ -7,7 +7,7 @@
 
 ## Introduction
 
-Aerich is a database migrations tool for Tortoise-ORM, which is like alembic for SQLAlchemy, or like Django ORM with
+Aerich is a database migrations tool for TortoiseORM, which is like alembic for SQLAlchemy, or like Django ORM with
 it\'s own migration solution.
 
 ## Install

aerich/cli.py

@@ -10,12 +10,11 @@
 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 @@ def drop_index(self, model: "Type[Model]", field_names: List[str], unique=False)
             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 @@ def _add_operator(cls, operator: str, upgrade=True, fk_m2m_index=False):
             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 @@ def diff_models(cls, old_models: Dict[str, dict], new_models: Dict[str, dict], u
                 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 @@ def _resolve_fk_fields_name(cls, model: Type[Model], fields_name: Tuple[str]):
         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"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "aerich"
-version = "0.6.1"
+version = "0.6.2"
 description = "A database migrations tool for Tortoise ORM."
 authors = ["long2ice <[email protected]>"]
 license = "Apache-2.0"
@@ -16,7 +16,7 @@ include = ["CHANGELOG.md", "LICENSE", "README.md"]
 
 [tool.poetry.dependencies]
 python = "^3.7"
-tortoise-orm = "*"
+tortoise-orm = { git = "https://github.com/tortoise/tortoise-orm.git", branch = "develop" }
 click = "*"
 asyncpg = { version = "*", optional = true }
 asyncmy = { version = "*", optional = true }