Add Command class. (#148) (#141) (#123) (#106)

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## 0.5
 
+### 0.5.6
+
+- Add `Command` class. (#148) (#141) (#123) (#106)
+
 ### 0.5.5
 
 - Fix KeyError: 'src_folder' after upgrading aerich to 0.5.4. (#176)

aerich/__init__.py

@@ -1 +1,144 @@
 __version__ = "0.5.5"
+
+import os
+from pathlib import Path
+from typing import List
+
+from tortoise import Tortoise, generate_schema_for_client
+from tortoise.exceptions import OperationalError
+from tortoise.transactions import in_transaction
+from tortoise.utils import get_schema_sql
+
+from aerich.exceptions import DowngradeError
+from aerich.inspectdb import InspectDb
+from aerich.migrate import Migrate
+from aerich.models import Aerich
+from aerich.utils import (
+    add_src_path,
+    get_app_connection,
+    get_app_connection_name,
+    get_models_describe,
+    get_version_content_from_file,
+    write_version_file,
+)
+
+
+class Command:
+    def __init__(
+        self,
+        tortoise_config: dict,
+        app: str,
+        location: str = "./migrations",
+        src_folder: str = ".",
+    ):
+        self.tortoise_config = tortoise_config
+        self.app = app
+        self.location = location
+        self.src_folder = src_folder
+        Migrate.app = app
+        add_src_path(src_folder)
+
+    async def init(self):
+        await Migrate.init(self.tortoise_config, self.app, self.location)
+
+    async def upgrade(self):
+        migrated = []
+        for version_file in Migrate.get_all_version_files():
+            try:
+                exists = await Aerich.exists(version=version_file, app=self.app)
+            except OperationalError:
+                exists = False
+            if not exists:
+                async with in_transaction(
+                    get_app_connection_name(self.tortoise_config, self.app)
+                ) as conn:
+                    file_path = Path(Migrate.migrate_location, version_file)
+                    content = get_version_content_from_file(file_path)
+                    upgrade_query_list = content.get("upgrade")
+                    for upgrade_query in upgrade_query_list:
+                        await conn.execute_script(upgrade_query)
+                    await Aerich.create(
+                        version=version_file,
+                        app=self.app,
+                        content=get_models_describe(self.app),
+                    )
+                migrated.append(version_file)
+        return migrated
+
+    async def downgrade(self, version: int, delete: bool):
+        ret = []
+        if version == -1:
+            specified_version = await Migrate.get_last_version()
+        else:
+            specified_version = await Aerich.filter(
+                app=self.app, version__startswith=f"{version}_"
+            ).first()
+        if not specified_version:
+            raise DowngradeError("No specified version found")
+        if version == -1:
+            versions = [specified_version]
+        else:
+            versions = await Aerich.filter(app=self.app, pk__gte=specified_version.pk)
+        for version in versions:
+            file = version.version
+            async with in_transaction(
+                get_app_connection_name(self.tortoise_config, self.app)
+            ) as conn:
+                file_path = Path(Migrate.migrate_location, file)
+                content = get_version_content_from_file(file_path)
+                downgrade_query_list = content.get("downgrade")
+                if not downgrade_query_list:
+                    raise DowngradeError("No downgrade items found")
+                for downgrade_query in downgrade_query_list:
+                    await conn.execute_query(downgrade_query)
+                await version.delete()
+                if delete:
+                    os.unlink(file_path)
+                ret.append(file)
+        return ret
+
+    async def heads(self):
+        ret = []
+        versions = Migrate.get_all_version_files()
+        for version in versions:
+            if not await Aerich.exists(version=version, app=self.app):
+                ret.append(version)
+        return ret
+
+    async def history(self):
+        ret = []
+        versions = Migrate.get_all_version_files()
+        for version in versions:
+            ret.append(version)
+        return ret
+
+    async def inspectdb(self, tables: List[str]):
+        connection = get_app_connection(self.tortoise_config, self.app)
+        inspect = InspectDb(connection, tables)
+        await inspect.inspect()
+
+    async def migrate(self, name: str):
+        return await Migrate.migrate(name)
+
+    async def init_db(self, safe: bool):
+        location = self.location
+        app = self.app
+        dirname = Path(location, app)
+        dirname.mkdir(parents=True)
+
+        await Tortoise.init(config=self.tortoise_config)
+        connection = get_app_connection(self.tortoise_config, app)
+        await generate_schema_for_client(connection, safe)
+
+        schema = get_schema_sql(connection, safe)
+
+        version = await Migrate.generate_version()
+        await Aerich.create(
+            version=version,
+            app=app,
+            content=get_models_describe(app),
+        )
+        content = {
+            "upgrade": [schema],
+        }
+        write_version_file(Path(dirname, version), content)

aerich/cli.py

@@ -7,26 +7,13 @@ from typing import List
 
 import click
 from click import Context, UsageError
-from tortoise import Tortoise, generate_schema_for_client
+from tortoise import Tortoise
-from tortoise.exceptions import OperationalError
-from tortoise.transactions import in_transaction
-from tortoise.utils import get_schema_sql
 
-from aerich.inspectdb import InspectDb
+from aerich.exceptions import DowngradeError
-from aerich.migrate import Migrate
+from aerich.utils import add_src_path, get_tortoise_config
-from aerich.utils import (
-    add_src_path,
-    get_app_connection,
-    get_app_connection_name,
-    get_models_describe,
-    get_tortoise_config,
-    get_version_content_from_file,
-    write_version_file,
-)
 
-from . import __version__
+from . import Command, __version__
 from .enums import Color
-from .models import Aerich
 
 parser = ConfigParser()
 
@@ -53,7 +40,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(
@@ -79,20 +70,16 @@ async def cli(ctx: Context, config, app, name):
         location = parser[name]["location"]
         tortoise_orm = parser[name]["tortoise_orm"]
         src_folder = parser[name].get("src_folder", CONFIG_DEFAULT_VALUES["src_folder"])
-
-        # Add specified source folder to path
-        add_src_path(src_folder)
-
         tortoise_config = get_tortoise_config(ctx, tortoise_orm)
         app = app or list(tortoise_config.get("apps").keys())[0]
-        ctx.obj["config"] = tortoise_config
+        command = Command(
-        ctx.obj["location"] = location
+            tortoise_config=tortoise_config, app=app, location=location, src_folder=src_folder
-        ctx.obj["app"] = app
+        )
-        Migrate.app = app
+        ctx.obj["command"] = command
         if invoked_subcommand != "init-db":
             if not Path(location, app).exists():
                 raise UsageError("You must exec init-db first", ctx=ctx)
-            await Migrate.init(tortoise_config, app, location)
+            await command.init()
 
 
 @cli.command(help="Generate migrate changes file.")
@@ -100,7 +87,8 @@ async def cli(ctx: Context, config, app, name):
 @click.pass_context
 @coro
 async def migrate(ctx: Context, name):
-    ret = await Migrate.migrate(name)
+    command = ctx.obj["command"]
+    ret = await command.migrate(name)
     if not ret:
         return click.secho("No changes detected", fg=Color.yellow)
     click.secho(f"Success migrate {ret}", fg=Color.green)
@@ -110,28 +98,13 @@ async def migrate(ctx: Context, name):
 @click.pass_context
 @coro
 async def upgrade(ctx: Context):
-    config = ctx.obj["config"]
+    command = ctx.obj["command"]
-    app = ctx.obj["app"]
+    migrated = await command.upgrade()
-    migrated = False
-    for version_file in Migrate.get_all_version_files():
-        try:
-            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 = Path(Migrate.migrate_location, version_file)
-                content = get_version_content_from_file(file_path)
-                upgrade_query_list = content.get("upgrade")
-                for upgrade_query in upgrade_query_list:
-                    await conn.execute_script(upgrade_query)
-                await Aerich.create(
-                    version=version_file, app=app, content=get_models_describe(app),
-                )
-            click.secho(f"Success upgrade {version_file}", fg=Color.green)
-            migrated = True
     if not migrated:
         click.secho("No upgrade items found", fg=Color.yellow)
+    else:
+        for version_file in migrated:
+            click.secho(f"Success upgrade {version_file}", fg=Color.green)
 
 
 @cli.command(help="Downgrade to specified version.")
@@ -157,59 +130,37 @@ async def upgrade(ctx: Context):
 )
 @coro
 async def downgrade(ctx: Context, version: int, delete: bool):
-    app = ctx.obj["app"]
+    command = ctx.obj["command"]
-    config = ctx.obj["config"]
+    try:
-    if version == -1:
+        files = await command.downgrade(version, delete)
-        specified_version = await Migrate.get_last_version()
+    except DowngradeError as e:
-    else:
+        return click.secho(str(e), fg=Color.yellow)
-        specified_version = await Aerich.filter(app=app, version__startswith=f"{version}_").first()
+    for file in files:
-    if not specified_version:
+        click.secho(f"Success downgrade {file}", fg=Color.green)
-        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 = Path(Migrate.migrate_location, file)
-            content = get_version_content_from_file(file_path)
-            downgrade_query_list = content.get("downgrade")
-            if not downgrade_query_list:
-                click.secho("No downgrade items found", fg=Color.yellow)
-                return
-            for downgrade_query in downgrade_query_list:
-                await conn.execute_query(downgrade_query)
-            await version.delete()
-            if delete:
-                os.unlink(file_path)
-            click.secho(f"Success downgrade {file}", fg=Color.green)
 
 
 @cli.command(help="Show current available heads in migrate location.")
 @click.pass_context
 @coro
 async def heads(ctx: Context):
-    app = ctx.obj["app"]
+    command = ctx.obj["command"]
-    versions = Migrate.get_all_version_files()
+    head_list = await command.heads()
-    is_heads = False
+    if not head_list:
-    for version in versions:
+        return click.secho("No available heads, try migrate first", fg=Color.green)
-        if not await Aerich.exists(version=version, app=app):
+    for version in head_list:
-            click.secho(version, fg=Color.green)
+        click.secho(version, fg=Color.green)
-            is_heads = True
-    if not is_heads:
-        click.secho("No available heads,try migrate first", fg=Color.green)
 
 
 @cli.command(help="List all migrate items.")
 @click.pass_context
 @coro
 async def history(ctx: Context):
-    versions = Migrate.get_all_version_files()
+    command = ctx.obj["command"]
+    versions = await command.history()
+    if not versions:
+        return click.secho("No history, try migrate", fg=Color.green)
     for version in versions:
         click.secho(version, fg=Color.green)
-    if not versions:
-        click.secho("No history,try migrate", fg=Color.green)
 
 
 @cli.command(help="Init config file and generate root migrate location.")
@@ -220,7 +171,10 @@ 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.option(
     "-s",
@@ -272,49 +226,32 @@ async def init(ctx: Context, tortoise_orm, location, src_folder):
 @click.pass_context
 @coro
 async def init_db(ctx: Context, safe):
-    config = ctx.obj["config"]
+    command = ctx.obj["command"]
-    location = ctx.obj["location"]
+    app = command.app
-    app = ctx.obj["app"]
+    dirname = Path(command.location, app)
-
-    dirname = Path(location, app)
     try:
-        dirname.mkdir(parents=True)
+        await command.init_db(safe)
         click.secho(f"Success create app migrate location {dirname}", fg=Color.green)
+        click.secho(f'Success generate schema for app "{app}"', fg=Color.green)
     except FileExistsError:
         return click.secho(
             f"Inited {app} already, or delete {dirname} and try again.", fg=Color.yellow
         )
 
-    await Tortoise.init(config=config)
-    connection = get_app_connection(config, app)
-    await generate_schema_for_client(connection, safe)
-
-    schema = get_schema_sql(connection, safe)
-
-    version = await Migrate.generate_version()
-    await Aerich.create(
-        version=version, app=app, content=get_models_describe(app),
-    )
-    content = {
-        "upgrade": [schema],
-    }
-    write_version_file(Path(dirname, version), content)
-    click.secho(f'Success generate schema for app "{app}"', fg=Color.green)
-
 
 @cli.command(help="Introspects the database tables to standard output as TortoiseORM model.")
 @click.option(
-    "-t", "--table", help="Which tables to inspect.", multiple=True, required=False,
+    "-t",
+    "--table",
+    help="Which tables to inspect.",
+    multiple=True,
+    required=False,
 )
 @click.pass_context
 @coro
 async def inspectdb(ctx: Context, table: List[str]):
-    config = ctx.obj["config"]
+    command = ctx.obj["command"]
-    app = ctx.obj["app"]
+    await command.inspectdb(table)
-    connection = get_app_connection(config, app)
-
-    inspect = InspectDb(connection, table)
-    await inspect.inspect()
 
 
 def main():

aerich/ddl/__init__.py

@@ -78,11 +78,15 @@ class BaseDDL:
         auto_now_add = field_describe.get("auto_now_add", False)
         auto_now = field_describe.get("auto_now", False)
         if default is not None or auto_now_add:
-            if field_describe.get("field_type") in [
+            if (
-                "UUIDField",
+                field_describe.get("field_type")
-                "TextField",
+                in [
-                "JSONField",
+                    "UUIDField",
-            ] or is_default_function(default):
+                    "TextField",
+                    "JSONField",
+                ]
+                or is_default_function(default)
+            ):
                 default = ""
             else:
                 try:
@@ -115,7 +119,9 @@ class BaseDDL:
                 nullable="NOT NULL" if not field_describe.get("nullable") else "",
                 unique="UNIQUE" if field_describe.get("unique") else "",
                 comment=self.schema_generator._column_comment_generator(
-                    table=db_table, column=db_column, comment=field_describe.get("description"),
+                    table=db_table,
+                    column=db_column,
+                    comment=field_describe.get("description"),
                 )
                 if description
                 else "",

aerich/exceptions.py

@@ -2,3 +2,9 @@ class NotSupportError(Exception):
     """
     raise when features not support
     """
+
+
+class DowngradeError(Exception):
+    """
+    raise when downgrade error
+    """

aerich/migrate.py

@@ -264,7 +264,11 @@ class Migrate:
                             # rename field
                             if (
                                 changes[0]
-                                == ("change", "name", (old_data_field_name, new_data_field_name),)
+                                == (
+                                    "change",
+                                    "name",
+                                    (old_data_field_name, new_data_field_name),
+                                )
                                 and changes[1]
                                 == (
                                     "change",
@@ -302,11 +306,16 @@ class Migrate:
                                         )
                                     else:
                                         cls._add_operator(
-                                            cls._rename_field(model, *changes[1][2]), upgrade,
+                                            cls._rename_field(model, *changes[1][2]),
+                                            upgrade,
                                         )
                     if not is_rename:
                         cls._add_operator(
-                            cls._add_field(model, new_data_field,), upgrade,
+                            cls._add_field(
+                                model,
+                                new_data_field,
+                            ),
+                            upgrade,
                         )
                 # remove fields
                 for old_data_field_name in set(old_data_fields_name).difference(
@@ -406,7 +415,8 @@ class Migrate:
                         else:
                             # modify column
                             cls._add_operator(
-                                cls._modify_field(model, new_data_field), upgrade,
+                                cls._modify_field(model, new_data_field),
+                                upgrade,
                             )
 
         for old_model in old_models:

aerich/utils.py

@@ -3,7 +3,7 @@ import os
 import re
 import sys
 from pathlib import Path
-from typing import Dict
+from typing import Dict, Union
 
 from click import BadOptionUsage, ClickException, Context
 from tortoise import BaseDBAsyncClient, Tortoise
@@ -36,7 +36,8 @@ def get_app_connection_name(config, app_name: str) -> str:
     if app:
         return app.get("default_connection", "default")
     raise BadOptionUsage(
-        option_name="--app", message=f'Can\'t get app named "{app_name}"',
+        option_name="--app",
+        message=f'Can\'t get app named "{app_name}"',
     )
 
 
@@ -80,7 +81,7 @@ _UPGRADE = "-- upgrade --\n"
 _DOWNGRADE = "-- downgrade --\n"
 
 
-def get_version_content_from_file(version_file: str) -> Dict:
+def get_version_content_from_file(version_file: Union[str, Path]) -> Dict:
     """
     get version content
     :param version_file:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "aerich"
-version = "0.5.5"
+version = "0.5.6"
 description = "A database migrations tool for Tortoise ORM."
 authors = ["long2ice <long2ice@gmail.com>"]
 license = "Apache-2.0"