Fix encoding error. (#75)

Refactoring migrate logic

.gitignore

@@ -143,4 +143,5 @@ cython_debug/
 .idea
 migrations
 aerich.ini
-src
+src
+.vscode

CHANGELOG.md

@@ -1,5 +1,32 @@
 # ChangeLog
 
+## 0.3
+
+### 0.3.3
+
+- Fix encoding error. (#75)
+- Support multiple databases. (#68)
+- Compatible with models file in directory. (#70)
+
+### 0.3.2
+
+- Fix migrate to new database error. (#62)
+
+### 0.3.1
+
+- Fix first version error.
+- Fix init error. (#61)
+
+### 0.3.0
+
+- Refactoring migrate logic, and this version is not compatible with previous version.
+- Now there don't need `old_models.py` and it store in database.
+- Upgrade steps:
+  1. Upgrade aerich version.
+  2. Drop aerich table in database.
+  3. Delete `migrations/{app}` folder and rerun `aerich init-db`.
+  4. Update model and `aerich migrate` normally.
+
 ## 0.2
 
 ### 0.2.5

Makefile

@@ -22,7 +22,7 @@ up:
 	@poetry update
 
 deps:
-	@poetry install -E dbdrivers --no-root
+	@poetry install -E dbdrivers
 
 style: deps
 	isort -src $(checkfiles)

README.md

@@ -7,12 +7,10 @@
 
 ## Introduction
 
-Tortoise-ORM is the best asyncio ORM now, but it lacks a database
-migrations tool like alembic for SQLAlchemy, or Django ORM with it\'s
-own migrations tool.
+Aerich is a database migrations tool for Tortoise-ORM, which like alembic for SQLAlchemy, or Django ORM with it\'s
+own migrations solution.
 
-This project aim to be a best migrations tool for Tortoise-ORM and which
-written by one of contributors of Tortoise-ORM.
+**Important: You can only use absolutely import in your `models.py` to make `aerich` work.**
 
 ## Install
 
@@ -25,7 +23,7 @@ Just install from pypi:
 ## Quick Start
 
 ```shell
-$ aerich -h
+> aerich -h
 
 Usage: aerich [OPTIONS] COMMAND [ARGS]...
 
@@ -37,7 +35,7 @@ Options:
   -h, --help         Show this message and exit.
 
 Commands:
-  downgrade  Downgrade to previous version.
+  downgrade  Downgrade to specified version.
   heads      Show current available heads in migrate location.
   history    List all migrate items.
   init       Init config file and generate root migrate location.
@@ -66,7 +64,7 @@ TORTOISE_ORM = {
 ### Initialization
 
 ```shell
-$ aerich init -h
+> aerich init -h
 
 Usage: aerich init [OPTIONS]
 
@@ -82,7 +80,7 @@ Options:
 Init config file and location:
 
 ```shell
-$ aerich init -t tests.backends.mysql.TORTOISE_ORM
+> aerich init -t tests.backends.mysql.TORTOISE_ORM
 
 Success create migrate location ./migrations
 Success generate config file aerich.ini
@@ -91,7 +89,7 @@ Success generate config file aerich.ini
 ### Init db
 
 ```shell
-$ aerich init-db
+> aerich init-db
 
 Success create app migrate location ./migrations/models
 Success generate schema for app "models"
@@ -103,7 +101,7 @@ If your Tortoise-ORM app is not default `models`, you must specify
 ### Update models and make migrate
 
 ```shell
-$ aerich migrate --name drop_column
+> aerich migrate --name drop_column
 
 Success migrate 1_202029051520102929_drop_column.json
 ```
@@ -118,27 +116,39 @@ If you use `MySQL`, only MySQL8.0+ support `rename..to` syntax.
 ### Upgrade to latest version
 
 ```shell
-$ aerich upgrade
+> aerich upgrade
 
 Success upgrade 1_202029051520102929_drop_column.json
 ```
 
 Now your db is migrated to latest.
 
-### Downgrade to previous version
+### Downgrade to specified version
 
 ```shell
-$ aerich downgrade
+> aerich init -h
+
+Usage: aerich downgrade [OPTIONS]
+
+  Downgrade to specified version.
+
+Options:
+  -v, --version INTEGER  Specified version, default to last.  [default: -1]
+  -h, --help             Show this message and exit.
+```
+
+```shell
+> aerich downgrade
 
 Success downgrade 1_202029051520102929_drop_column.json
 ```
 
-Now your db rollback to previous version.
+Now your db rollback to specified version.
 
 ### Show history
 
 ```shell
-$ aerich history
+> aerich history
 
 1_202029051520102929_drop_column.json
 ```
@@ -146,15 +156,29 @@ $ aerich history
 ### Show heads to be migrated
 
 ```shell
-$ aerich heads
+> aerich heads
 
 1_202029051520102929_drop_column.json
 ```
 
-## Support this project
+### Multiple databases
 
-- Just give a star!
-- Donation.
+```python
+tortoise_orm = {
+    "connections": {
+        "default": expand_db_url(db_url, True),
+        "second": expand_db_url(db_url_second, True),
+    },
+    "apps": {
+        "models": {"models": ["tests.models", "aerich.models"], "default_connection": "default"},
+        "models_second": {"models": ["tests.models_second"], "default_connection": "second",},
+    },
+}
+```
+
+You need only specify `aerich.models` in one app, and must specify `--app` when run `aerich migrate` and so on.
+
+## Support this project
 
 | AliPay                                                                                 | WeChatPay                                                                                 | PayPal                                                           |
 | -------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- | ---------------------------------------------------------------- |

aerich/__init__.py

@@ -1 +1 @@
-__version__ = "0.2.5"
+__version__ = "0.3.3"

aerich/cli.py

@@ -26,7 +26,11 @@ def coro(f):
     @wraps(f)
     def wrapper(*args, **kwargs):
         loop = asyncio.get_event_loop()
-        return loop.run_until_complete(f(*args, **kwargs))
+        ctx = args[0]
+        loop.run_until_complete(f(*args, **kwargs))
+        app = ctx.obj.get("app")
+        if app:
+            Migrate.remove_old_model_file(app, ctx.obj["location"])
 
     return wrapper
 
@@ -34,7 +38,11 @@ def coro(f):
 @click.group(context_settings={"help_option_names": ["-h", "--help"]})
 @click.version_option(__version__, "-V", "--version")
 @click.option(
-    "-c", "--config", default="aerich.ini", show_default=True, help="Config file.",
+    "-c",
+    "--config",
+    default="aerich.ini",
+    show_default=True,
+    help="Config file.",
 )
 @click.option("--app", required=False, help="Tortoise-ORM app name.")
 @click.option(
@@ -62,12 +70,10 @@ async def cli(ctx: Context, config, app, name):
 
         tortoise_config = get_tortoise_config(ctx, tortoise_orm)
         app = app or list(tortoise_config.get("apps").keys())[0]
-        if "aerich.models" not in tortoise_config.get("apps").get(app).get("models"):
-            raise UsageError("Check your tortoise config and add aerich.models to it.", ctx=ctx)
         ctx.obj["config"] = tortoise_config
         ctx.obj["location"] = location
         ctx.obj["app"] = app
-
+        Migrate.app = app
         if invoked_subcommand != "init-db":
             await Migrate.init_with_old_models(tortoise_config, app, location)
 
@@ -77,64 +83,85 @@ async def cli(ctx: Context, config, app, name):
 @click.pass_context
 @coro
 async def migrate(ctx: Context, name):
-    config = ctx.obj["config"]
-    location = ctx.obj["location"]
-    app = ctx.obj["app"]
     ret = await Migrate.migrate(name)
     if not ret:
         return click.secho("No changes detected", fg=Color.yellow)
-    Migrate.write_old_models(config, app, location)
     click.secho(f"Success migrate {ret}", fg=Color.green)
 
 
-@cli.command(help="Upgrade to latest version.")
+@cli.command(help="Upgrade to specified version.")
 @click.pass_context
 @coro
 async def upgrade(ctx: Context):
     config = ctx.obj["config"]
     app = ctx.obj["app"]
+    location = ctx.obj["location"]
     migrated = False
-    for version in Migrate.get_all_version_files():
+    for version_file in Migrate.get_all_version_files():
         try:
-            exists = await Aerich.exists(version=version, app=app)
+            exists = await Aerich.exists(version=version_file, app=app)
         except OperationalError:
             exists = False
         if not exists:
             async with in_transaction(get_app_connection_name(config, app)) as conn:
-                file_path = os.path.join(Migrate.migrate_location, version)
+                file_path = os.path.join(Migrate.migrate_location, version_file)
                 with open(file_path, "r", encoding="utf-8") as f:
                     content = json.load(f)
                     upgrade_query_list = content.get("upgrade")
                     for upgrade_query in upgrade_query_list:
                         await conn.execute_script(upgrade_query)
-                await Aerich.create(version=version, app=app)
-            click.secho(f"Success upgrade {version}", fg=Color.green)
+                await Aerich.create(
+                    version=version_file,
+                    app=app,
+                    content=Migrate.get_models_content(config, app, location),
+                )
+            click.secho(f"Success upgrade {version_file}", fg=Color.green)
             migrated = True
     if not migrated:
         click.secho("No migrate items", fg=Color.yellow)
 
 
-@cli.command(help="Downgrade to previous version.")
+@cli.command(help="Downgrade to specified version.")
+@click.option(
+    "-v",
+    "--version",
+    default=-1,
+    type=int,
+    show_default=True,
+    help="Specified version, default to last.",
+)
 @click.pass_context
+@click.confirmation_option(
+    prompt="Downgrade is dangerous, which maybe lose your data, are you sure?",
+)
 @coro
-async def downgrade(ctx: Context):
+async def downgrade(ctx: Context, version: int):
     app = ctx.obj["app"]
     config = ctx.obj["config"]
-    last_version = await Migrate.get_last_version()
-    if not last_version:
-        return click.secho("No last version found", fg=Color.yellow)
-    file = last_version.version
-    async with in_transaction(get_app_connection_name(config, app)) as conn:
-        file_path = os.path.join(Migrate.migrate_location, file)
-        with open(file_path, "r", encoding="utf-8") as f:
-            content = json.load(f)
-            downgrade_query_list = content.get("downgrade")
-            if not downgrade_query_list:
-                return click.secho("No downgrade item found", fg=Color.yellow)
-            for downgrade_query in downgrade_query_list:
-                await conn.execute_query(downgrade_query)
-            await last_version.delete()
-        return click.secho(f"Success downgrade {file}", fg=Color.green)
+    if version == -1:
+        specified_version = await Migrate.get_last_version()
+    else:
+        specified_version = await Aerich.filter(app=app, version__startswith=f"{version}_").first()
+    if not specified_version:
+        return click.secho("No specified version found", fg=Color.yellow)
+    if version == -1:
+        versions = [specified_version]
+    else:
+        versions = await Aerich.filter(app=app, pk__gte=specified_version.pk)
+    for version in versions:
+        file = version.version
+        async with in_transaction(get_app_connection_name(config, app)) as conn:
+            file_path = os.path.join(Migrate.migrate_location, file)
+            with open(file_path, "r", encoding="utf-8") as f:
+                content = json.load(f)
+                downgrade_query_list = content.get("downgrade")
+                if not downgrade_query_list:
+                    return click.secho("No downgrade item found", fg=Color.yellow)
+                for downgrade_query in downgrade_query_list:
+                    await conn.execute_query(downgrade_query)
+                await version.delete()
+            os.unlink(file_path)
+            click.secho(f"Success downgrade {file}", fg=Color.green)
 
 
 @cli.command(help="Show current available heads in migrate location.")
@@ -149,7 +176,7 @@ async def heads(ctx: Context):
             click.secho(version, fg=Color.green)
             is_heads = True
     if not is_heads:
-        click.secho("No available heads,try migrate", fg=Color.green)
+        click.secho("No available heads,try migrate first", fg=Color.green)
 
 
 @cli.command(help="List all migrate items.")
@@ -171,12 +198,17 @@ async def history(ctx: Context):
     help="Tortoise-ORM config module dict variable, like settings.TORTOISE_ORM.",
 )
 @click.option(
-    "--location", default="./migrations", show_default=True, help="Migrate store location."
+    "--location",
+    default="./migrations",
+    show_default=True,
+    help="Migrate store location.",
 )
 @click.pass_context
 @coro
 async def init(
-    ctx: Context, tortoise_orm, location,
+    ctx: Context,
+    tortoise_orm,
+    location,
 ):
     config_file = ctx.obj["config_file"]
     name = ctx.obj["name"]
@@ -219,8 +251,6 @@ async def init_db(ctx: Context, safe):
     else:
         return click.secho(f"Inited {app} already", fg=Color.yellow)
 
-    Migrate.write_old_models(config, app, location)
-
     await Tortoise.init(config=config)
     connection = get_app_connection(config, app)
     await generate_schema_for_client(connection, safe)
@@ -228,7 +258,11 @@ async def init_db(ctx: Context, safe):
     schema = get_schema_sql(connection, safe)
 
     version = await Migrate.generate_version()
-    await Aerich.create(version=version, app=app)
+    await Aerich.create(
+        version=version,
+        app=app,
+        content=Migrate.get_models_content(config, app, location),
+    )
     with open(os.path.join(dirname, version), "w", encoding="utf-8") as f:
         content = {
             "upgrade": [schema],
@@ -240,3 +274,7 @@ async def init_db(ctx: Context, safe):
 def main():
     sys.path.insert(0, ".")
     cli()
+
+
+if __name__ == "__main__":
+    main()

aerich/ddl/sqlite/__init__.py

@@ -13,7 +13,7 @@ class SqliteDDL(BaseDDL):
     DIALECT = SqliteSchemaGenerator.DIALECT
 
     def drop_column(self, model: "Type[Model]", column_name: str):
-        raise NotSupportError("Drop column is not support in SQLite.")
+        raise NotSupportError("Drop column is unsupported in SQLite.")
 
     def modify_column(self, model: "Type[Model]", field_object: Field):
-        raise NotSupportError("Modify column is not support in SQLite.")
+        raise NotSupportError("Modify column is unsupported in SQLite.")

aerich/migrate.py

@@ -1,9 +1,11 @@
+import inspect
 import json
 import os
 import re
 from datetime import datetime
 from importlib import import_module
-from typing import Dict, List, Tuple, Type
+from io import StringIO
+from typing import Dict, List, Optional, Tuple, Type
 
 import click
 from tortoise import (
@@ -14,6 +16,7 @@ from tortoise import (
     Model,
     Tortoise,
 )
+from tortoise.exceptions import OperationalError
 from tortoise.fields import Field
 
 from aerich.ddl import BaseDDL
@@ -41,8 +44,8 @@ class Migrate:
     dialect: str
 
     @classmethod
-    def get_old_model_file(cls):
-        return cls.old_models + ".py"
+    def get_old_model_file(cls, app: str, location: str):
+        return os.path.join(location, app, cls.old_models + ".py")
 
     @classmethod
     def get_all_version_files(cls) -> List[str]:
@@ -52,18 +55,33 @@ class Migrate:
         )
 
     @classmethod
-    async def get_last_version(cls) -> Aerich:
-        return await Aerich.filter(app=cls.app).first()
+    async def get_last_version(cls) -> Optional[Aerich]:
+        try:
+            return await Aerich.filter(app=cls.app).first()
+        except OperationalError:
+            pass
+
+    @classmethod
+    def remove_old_model_file(cls, app: str, location: str):
+        try:
+            os.unlink(cls.get_old_model_file(app, location))
+        except FileNotFoundError:
+            pass
 
     @classmethod
     async def init_with_old_models(cls, config: dict, app: str, location: str):
-        migrate_config = cls._get_migrate_config(config, app, location)
-
+        await Tortoise.init(config=config)
+        last_version = await cls.get_last_version()
         cls.app = app
-        cls.migrate_config = migrate_config
         cls.migrate_location = os.path.join(location, app)
+        if last_version:
+            content = last_version.content
+            with open(cls.get_old_model_file(app, location), "w", encoding="utf-8") as f:
+                f.write(content)
 
-        await Tortoise.init(config=migrate_config)
+            migrate_config = cls._get_migrate_config(config, app, location)
+            cls.migrate_config = migrate_config
+            await Tortoise.init(config=migrate_config)
 
         connection = get_app_connection(config, app)
         cls.dialect = connection.schema_generator.DIALECT
@@ -102,6 +120,10 @@ class Migrate:
     @classmethod
     async def _generate_diff_sql(cls, name):
         version = await cls.generate_version(name)
+        # delete if same version exists
+        for version_file in cls.get_all_version_files():
+            if version_file.startswith(version.split("_")[0]):
+                os.unlink(os.path.join(cls.migrate_location, version_file))
         content = {
             "upgrade": cls.upgrade_operators,
             "downgrade": cls.downgrade_operators,
@@ -151,26 +173,6 @@ class Migrate:
             else:
                 cls.downgrade_operators.append(operator)
 
-    @classmethod
-    def cp_models(
-        cls, app: str, model_files: List[str], old_model_file,
-    ):
-        """
-        cp currents models to old_model_files
-        :param app:
-        :param model_files:
-        :param old_model_file:
-        :return:
-        """
-        pattern = rf"(\n)?('|\")({app})(.\w+)('|\")"
-        for i, model_file in enumerate(model_files):
-            with open(model_file, "r", encoding="utf-8") as f:
-                content = f.read()
-            ret = re.sub(pattern, rf"\2{cls.diff_app}\4\5", content)
-            mode = "w" if i == 0 else "a"
-            with open(old_model_file, mode, encoding="utf-8") as f:
-                f.write(f"{ret}\n")
-
     @classmethod
     def _get_migrate_config(cls, config: dict, app: str, location: str):
         """
@@ -189,7 +191,7 @@ class Migrate:
         return config
 
     @classmethod
-    def write_old_models(cls, config: dict, app: str, location: str):
+    def get_models_content(cls, config: dict, app: str, location: str):
         """
         write new models to old models
         :param config:
@@ -197,14 +199,26 @@ class Migrate:
         :param location:
         :return:
         """
-        cls.app = app
-
         old_model_files = []
         models = config.get("apps").get(app).get("models")
         for model in models:
-            old_model_files.append(import_module(model).__file__)
-
-        cls.cp_models(app, old_model_files, os.path.join(location, app, cls.get_old_model_file()))
+            module = import_module(model)
+            possible_models = [getattr(module, attr_name) for attr_name in dir(module)]
+            for attr in filter(
+                lambda x: inspect.isclass(x) and issubclass(x, Model) and x is not Model,
+                possible_models,
+            ):
+                file = inspect.getfile(attr)
+                if file not in old_model_files:
+                    old_model_files.append(file)
+        pattern = rf"(\n)?('|\")({app})(.\w+)('|\")"
+        str_io = StringIO()
+        for i, model_file in enumerate(old_model_files):
+            with open(model_file, "r", encoding="utf-8") as f:
+                content = f.read()
+            ret = re.sub(pattern, rf"\2{cls.diff_app}\4\5", content)
+            str_io.write(f"{ret}\n")
+        return str_io.getvalue()
 
     @classmethod
     def diff_models(
@@ -289,12 +303,15 @@ class Migrate:
                             is_rename = diff_key in cls._rename_new
                         if is_rename:
                             cls._add_operator(
-                                cls._rename_field(new_model, old_field, new_field), upgrade,
+                                cls._rename_field(new_model, old_field, new_field),
+                                upgrade,
                             )
                             break
                 else:
                     cls._add_operator(
-                        cls._add_field(new_model, new_field), upgrade, cls._is_fk_m2m(new_field),
+                        cls._add_field(new_model, new_field),
+                        upgrade,
+                        cls._is_fk_m2m(new_field),
                     )
             else:
                 old_field = old_fields_map.get(new_key)
@@ -345,11 +362,15 @@ class Migrate:
                 if isinstance(new_field, ForeignKeyFieldInstance):
                     if old_field.db_constraint and not new_field.db_constraint:
                         cls._add_operator(
-                            cls._drop_fk(new_model, new_field), upgrade, True,
+                            cls._drop_fk(new_model, new_field),
+                            upgrade,
+                            True,
                         )
                     if new_field.db_constraint and not old_field.db_constraint:
                         cls._add_operator(
-                            cls._add_fk(new_model, new_field), upgrade, True,
+                            cls._add_fk(new_model, new_field),
+                            upgrade,
+                            True,
                         )
 
         for old_key in old_keys:
@@ -359,12 +380,20 @@ class Migrate:
                     not upgrade and old_key not in cls._rename_new
                 ):
                     cls._add_operator(
-                        cls._remove_field(old_model, field), upgrade, cls._is_fk_m2m(field),
+                        cls._remove_field(old_model, field),
+                        upgrade,
+                        cls._is_fk_m2m(field),
                     )
 
         for new_index in new_indexes:
             if new_index not in old_indexes:
-                cls._add_operator(cls._add_index(new_model, new_index,), upgrade)
+                cls._add_operator(
+                    cls._add_index(
+                        new_model,
+                        new_index,
+                    ),
+                    upgrade,
+                )
         for old_index in old_indexes:
             if old_index not in new_indexes:
                 cls._add_operator(cls._remove_index(old_model, old_index), upgrade)

aerich/models.py

@@ -6,6 +6,7 @@ MAX_VERSION_LENGTH = 255
 class Aerich(Model):
     version = fields.CharField(max_length=MAX_VERSION_LENGTH)
     app = fields.CharField(max_length=20)
+    content = fields.TextField()
 
     class Meta:
         ordering = ["-id"]

conftest.py

@@ -13,10 +13,15 @@ from aerich.ddl.sqlite import SqliteDDL
 from aerich.migrate import Migrate
 
 db_url = os.getenv("TEST_DB", "sqlite://:memory:")
+db_url_second = os.getenv("TEST_DB_SECOND", "sqlite://:memory:")
 tortoise_orm = {
-    "connections": {"default": expand_db_url(db_url, True)},
+    "connections": {
+        "default": expand_db_url(db_url, True),
+        "second": expand_db_url(db_url_second, True),
+    },
     "apps": {
         "models": {"models": ["tests.models", "aerich.models"], "default_connection": "default"},
+        "models_second": {"models": ["tests.models_second"], "default_connection": "second"},
     },
 }
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "aerich"
-version = "0.2.5"
+version = "0.3.3"
 description = "A database migrations tool for Tortoise ORM."
 authors = ["long2ice <long2ice@gmail.com>"]
 license = "Apache-2.0"
@@ -25,7 +25,7 @@ asyncpg = {version = "*", optional = true}
 [tool.poetry.dev-dependencies]
 flake8 = "*"
 isort = "*"
-black = "^19.10b0"
+black = "*"
 pytest = "*"
 pytest-xdist = "*"
 pytest-asyncio = "*"

tests/models_second.py

@@ -0,0 +1,63 @@
+import datetime
+from enum import IntEnum
+
+from tortoise import Model, fields
+
+
+class ProductType(IntEnum):
+    article = 1
+    page = 2
+
+
+class PermissionAction(IntEnum):
+    create = 1
+    delete = 2
+    update = 3
+    read = 4
+
+
+class Status(IntEnum):
+    on = 1
+    off = 0
+
+
+class User(Model):
+    username = fields.CharField(max_length=20, unique=True)
+    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")
+    is_superuser = fields.BooleanField(default=False, description="Is SuperUser")
+    avatar = fields.CharField(max_length=200, default="")
+    intro = fields.TextField(default="")
+
+
+class Email(Model):
+    email = fields.CharField(max_length=200)
+    is_primary = fields.BooleanField(default=False)
+    user = fields.ForeignKeyField("models_second.User", db_constraint=False)
+
+
+class Category(Model):
+    slug = fields.CharField(max_length=200)
+    name = fields.CharField(max_length=200)
+    user = fields.ForeignKeyField("models_second.User", description="User")
+    created_at = fields.DatetimeField(auto_now_add=True)
+
+
+class Product(Model):
+    categories = fields.ManyToManyField("models_second.Category")
+    name = fields.CharField(max_length=50)
+    view_num = fields.IntField(description="View Num")
+    sort = fields.IntField()
+    is_reviewed = fields.BooleanField(description="Is Reviewed")
+    type = fields.IntEnumField(ProductType, description="Product Type")
+    image = fields.CharField(max_length=200)
+    body = fields.TextField()
+    created_at = fields.DatetimeField(auto_now_add=True)
+
+
+class Config(Model):
+    label = fields.CharField(max_length=200)
+    key = fields.CharField(max_length=20)
+    value = fields.JSONField()
+    status: Status = fields.IntEnumField(Status, default=Status.on)