From a5a5de529b8c4d2d8f05635c1d887cfb0cb04fa1 Mon Sep 17 00:00:00 2001
From: long2ice <long2ice@gmail.com>
Date: Wed, 13 May 2020 18:59:24 +0800
Subject: [PATCH] add migrate and cli

---
 MANIFEST.in                          |  3 +
 alice/__init__.py                    | 10 ----
 alice/backends/__init__.py           |  2 +-
 alice/cli.py                         | 83 ++++++++++++++++++++++++++
 alice/cmd.py                         | 18 ------
 alice/diff.py                        |  0
 alice/migrate.py                     | 88 ++++++++++++++++++++++++++++
 alice/utils.py                       | 17 ++++++
 poetry.lock                          | 24 ++++++--
 pyproject.toml                       |  3 +-
 requirements.txt                     |  1 +
 setup.py                             | 43 ++++++++++++++
 tests/backends/mysql/__init__.py     |  7 ++-
 tests/backends/mysql/test_ddl.py     |  6 +-
 tests/backends/mysql/test_migrate.py |  7 +++
 tests/new_models.py                  | 69 ++++++++++++++++++++++
 16 files changed, 338 insertions(+), 43 deletions(-)
 create mode 100644 MANIFEST.in
 create mode 100644 alice/cli.py
 delete mode 100644 alice/cmd.py
 delete mode 100644 alice/diff.py
 create mode 100644 alice/migrate.py
 create mode 100644 alice/utils.py
 create mode 100644 tests/backends/mysql/test_migrate.py
 create mode 100644 tests/new_models.py

diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 0000000..5f9c2e9
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,3 @@
+include LICENSE
+include README.rst
+include requirements.txt
\ No newline at end of file
diff --git a/alice/__init__.py b/alice/__init__.py
index b6bb958..b794fd4 100644
--- a/alice/__init__.py
+++ b/alice/__init__.py
@@ -1,11 +1 @@
 __version__ = '0.1.0'
-
-from alice.cmd import CommandLine
-
-
-def main(argv):
-    command = CommandLine(argv)
-
-
-if __name__ == '__main__':
-    main(None)
diff --git a/alice/backends/__init__.py b/alice/backends/__init__.py
index 1837876..9ed6f93 100644
--- a/alice/backends/__init__.py
+++ b/alice/backends/__init__.py
@@ -63,7 +63,7 @@ class DDL:
                 )
                 if field_object.description else "",
                 is_primary_key=field_object.pk,
-                default=default,
+                default=default
             )
         )
 
diff --git a/alice/cli.py b/alice/cli.py
new file mode 100644
index 0000000..080cfc4
--- /dev/null
+++ b/alice/cli.py
@@ -0,0 +1,83 @@
+import importlib
+import os
+import sys
+from enum import Enum
+
+import click
+from click import BadParameter, ClickException
+
+sys.path.append(os.getcwd())
+
+
+class Color(str, Enum):
+    green = 'green'
+    red = 'red'
+
+
+@click.group(context_settings={'help_option_names': ['-h', '--help']})
+@click.option('-c', '--config', default='settings',
+              help='Tortoise-ORM config module, will read config variable from it, default is `settings`.')
+@click.option('-t', '--tortoise-orm', default='TORTOISE_ORM',
+              help='Tortoise-ORM config dict variable, default is `TORTOISE_ORM`.')
+@click.option('-l', '--location', default='./migrations',
+              help='Migrate store location, default is `./migrations`.')
+@click.option('--connection', default='default', help='Tortoise-ORM connection name, default is `default`.')
+@click.pass_context
+def cli(ctx, config, tortoise_orm, location, connection):
+    ctx.ensure_object(dict)
+    try:
+        config_module = importlib.import_module(config)
+        config = getattr(config_module, tortoise_orm, None)
+        if not config:
+            raise BadParameter(param_hint=['--config'],
+                               message=f'Can\'t get "{tortoise_orm}" from module "{config_module}"')
+        ctx.obj['config'] = config
+        ctx.obj['location'] = location
+        if connection not in config.get('connections').keys():
+            raise BadParameter(param_hint=['--connection'], message=f'No connection found in "{config}"')
+    except ModuleNotFoundError:
+        raise BadParameter(param_hint=['--tortoise-orm'], message=f'No module named "{config}"')
+
+
+@cli.command()
+@click.pass_context
+def migrate(ctx):
+    config = ctx.obj['config']
+
+
+@cli.command()
+@click.pass_context
+def upgrade():
+    pass
+
+
+@cli.command()
+@click.pass_context
+def downgrade():
+    pass
+
+
+@cli.command()
+@click.pass_context
+def initdb():
+    pass
+
+
+@cli.command()
+@click.option('--overwrite', type=bool, default=False, help='Overwrite old_models.py.')
+@click.pass_context
+def init(ctx, overwrite):
+    location = ctx.obj['location']
+    config = ctx.obj['config']
+    if not os.path.isdir(location) or overwrite:
+        os.mkdir(location)
+        connections = config.get('connections').keys()
+        for connection in connections:
+            dirname = os.path.join(location, connection)
+            if not os.path.isdir(dirname):
+                os.mkdir(dirname)
+                click.secho(f'Success create migrate location {dirname}', fg=Color.green)
+            if overwrite:
+                pass
+    else:
+        raise ClickException('Already inited')
diff --git a/alice/cmd.py b/alice/cmd.py
deleted file mode 100644
index 4d21f61..0000000
--- a/alice/cmd.py
+++ /dev/null
@@ -1,18 +0,0 @@
-class CommandLine:
-    def __init__(self, argv):
-        self.argv = argv
-
-    def migrate(self):
-        pass
-
-    def upgrade(self):
-        pass
-
-    def downgrade(self):
-        pass
-
-    def init_db(self):
-        pass
-
-    def init(self):
-        pass
diff --git a/alice/diff.py b/alice/diff.py
deleted file mode 100644
index e69de29..0000000
diff --git a/alice/migrate.py b/alice/migrate.py
new file mode 100644
index 0000000..c9b1330
--- /dev/null
+++ b/alice/migrate.py
@@ -0,0 +1,88 @@
+import importlib
+import inspect
+from typing import List, Type, Dict
+
+from tortoise import Model, ForeignKeyFieldInstance
+from tortoise.fields import Field
+
+from alice.backends import DDL
+
+
+class Migrate:
+    operators: List
+    ddl: DDL
+
+    def __init__(self, ddl: DDL):
+        self.operators = []
+        self.ddl = ddl
+
+    def diff_models_module(self, old_models_module, new_models_module):
+        old_module = importlib.import_module(old_models_module)
+        old_models = {}
+        new_models = {}
+        for name, obj in inspect.getmembers(old_module):
+            if inspect.isclass(obj) and issubclass(obj, Model):
+                old_models[obj.__name__] = obj
+
+        new_module = importlib.import_module(new_models_module)
+        for name, obj in inspect.getmembers(new_module):
+            if inspect.isclass(obj) and issubclass(obj, Model):
+                new_models[obj.__name__] = obj
+        self.diff_models(old_models, new_models)
+
+    def diff_models(self, old_models: Dict[str, Type[Model]], new_models: Dict[str, Type[Model]]):
+        for new_model_str, new_model in new_models.items():
+            if new_model_str not in old_models.keys():
+                self.add_model(new_model)
+            else:
+                self.diff_model(old_models.get(new_model_str), new_model)
+
+        for old_model in old_models:
+            if old_model not in new_models.keys():
+                self.remove_model(old_models.get(old_model))
+
+    def _add_operator(self, operator):
+        self.operators.append(operator)
+
+    def add_model(self, model: Type[Model]):
+        self._add_operator(self.ddl.create_table(model))
+
+    def remove_model(self, model: Type[Model]):
+        self._add_operator(self.ddl.drop_table(model))
+
+    def diff_model(self, old_model: Type[Model], new_model: Type[Model]):
+        old_fields_map = old_model._meta.fields_map
+        new_fields_map = new_model._meta.fields_map
+        old_keys = old_fields_map.keys()
+        new_keys = new_fields_map.keys()
+        for new_key in new_keys:
+            new_field = new_fields_map.get(new_key)
+            if new_key not in old_keys:
+                self._add_field(new_model, new_field)
+            else:
+                old_field = old_fields_map.get(new_key)
+                if old_field.index and not new_field.index:
+                    self._remove_index(old_model, old_field)
+                elif new_field.index and not old_field.index:
+                    self._add_index(new_model, new_field)
+        for old_key in old_keys:
+            if old_key not in new_keys:
+                field = old_fields_map.get(old_key)
+                self._remove_field(old_model, field)
+
+    def _remove_index(self, model: Type[Model], field: Field):
+        self._add_operator(self.ddl.drop_index(model, [field.model_field_name], field.unique))
+
+    def _add_index(self, model: Type[Model], field: Field):
+        self._add_operator(self.ddl.add_index(model, [field.model_field_name], field.unique))
+
+    def _add_field(self, model: Type[Model], field: Field):
+        if isinstance(field, ForeignKeyFieldInstance):
+            self._add_operator(self.ddl.add_fk(model, field))
+        else:
+            self._add_operator(self.ddl.add_column(model, field))
+
+    def _remove_field(self, model: Type[Model], field: Field):
+        if isinstance(field, ForeignKeyFieldInstance):
+            self._add_operator(self.ddl.drop_fk(model, field))
+        self._add_operator(self.ddl.drop_column(model, field.model_field_name))
diff --git a/alice/utils.py b/alice/utils.py
new file mode 100644
index 0000000..e5b9908
--- /dev/null
+++ b/alice/utils.py
@@ -0,0 +1,17 @@
+import re
+
+
+def cp_models(old_model_file, new_model_file, new_app):
+    """
+    cp models file to old_models.py and rename model app
+    :param old_app:
+    :param new_app:
+    :param old_model_file:
+    :param new_model_file:
+    :return:r
+    """
+    pattern = r'(ManyToManyField|ForeignKeyField|OneToOneField)\((model_name)?(\"|\')(?P<app>\w+).+\)'
+    with open(old_model_file, 'r') as f:
+        content = f.read()
+    ret = re.sub(pattern, rf'{new_app} \g<app>', content)
+    print(ret)
diff --git a/poetry.lock b/poetry.lock
index c69d758..8bce2a0 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -48,6 +48,14 @@ optional = false
 python-versions = "*"
 version = "2.1.3"
 
+[[package]]
+category = "main"
+description = "Composable command line interface toolkit"
+name = "click"
+optional = false
+python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
+version = "7.1.2"
+
 [[package]]
 category = "main"
 description = "cryptography is a package which provides cryptographic recipes and primitives to Python developers."
@@ -173,7 +181,7 @@ description = "Easy async ORM for python, built with relations in mind"
 name = "tortoise-orm"
 optional = false
 python-versions = "*"
-version = "0.16.10"
+version = "0.16.11"
 
 [package.dependencies]
 aiosqlite = ">=0.11.0"
@@ -182,6 +190,10 @@ iso8601 = ">=0.1.12"
 pypika = ">=0.36.5"
 typing-extensions = ">=3.7"
 
+[package.source]
+reference = "72f84f0848dc68041157f03e60cd1c92b0ee5137"
+type = "git"
+url = "https://github.com/tortoise/tortoise-orm.git"
 [[package]]
 category = "main"
 description = "Backported and Experimental Type Hints for Python 3.5+"
@@ -191,7 +203,7 @@ python-versions = "*"
 version = "3.7.4.2"
 
 [metadata]
-content-hash = "0b4c3c6eb6ed2e84b03745542ebc42dc53e6e5466bea0c4eee991bcd063a6ef3"
+content-hash = "708876857d4653fd45cb251e9a1689c4158966f42da2efca1e8167becef89837"
 python-versions = "^3.8"
 
 [metadata.files]
@@ -240,6 +252,10 @@ cffi = [
 ciso8601 = [
     {file = "ciso8601-2.1.3.tar.gz", hash = "sha256:bdbb5b366058b1c87735603b23060962c439ac9be66f1ae91e8c7dbd7d59e262"},
 ]
+click = [
+    {file = "click-7.1.2-py2.py3-none-any.whl", hash = "sha256:dacca89f4bfadd5de3d7489b7c8a566eee0d3676333fbb50030263894c38c0dc"},
+    {file = "click-7.1.2.tar.gz", hash = "sha256:d2b5255c7c6349bc1bd1e59e08cd12acbbd63ce649f2588755783aa94dfb6b1a"},
+]
 cryptography = [
     {file = "cryptography-2.9.2-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:daf54a4b07d67ad437ff239c8a4080cfd1cc7213df57d33c97de7b4738048d5e"},
     {file = "cryptography-2.9.2-cp27-cp27m-manylinux1_x86_64.whl", hash = "sha256:3b3eba865ea2754738616f87292b7f29448aec342a7c720956f8083d252bf28b"},
@@ -306,9 +322,7 @@ toml = [
     {file = "toml-0.10.0-py2.py3-none-any.whl", hash = "sha256:235682dd292d5899d361a811df37e04a8828a5b1da3115886b73cf81ebc9100e"},
     {file = "toml-0.10.0.tar.gz", hash = "sha256:229f81c57791a41d65e399fc06bf0848bab550a9dfd5ed66df18ce5f05e73d5c"},
 ]
-tortoise-orm = [
-    {file = "tortoise-orm-0.16.10.tar.gz", hash = "sha256:b3f4fdc9edabfc88413b7c5297b6cb9408420d1a97d9ad25051170b2b6228e02"},
-]
+tortoise-orm = []
 typing-extensions = [
     {file = "typing_extensions-3.7.4.2-py2-none-any.whl", hash = "sha256:f8d2bd89d25bc39dabe7d23df520442fa1d8969b82544370e03d88b5a591c392"},
     {file = "typing_extensions-3.7.4.2-py3-none-any.whl", hash = "sha256:6e95524d8a547a91e08f404ae485bbb71962de46967e1b71a0cb89af24e761c5"},
diff --git a/pyproject.toml b/pyproject.toml
index 6855f8e..de85af0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,8 +6,9 @@ authors = ["long2ice <long2ice@gmail.com>"]
 
 [tool.poetry.dependencies]
 python = "^3.8"
-tortoise-orm = "*"
+tortoise-orm = {git = "https://github.com/tortoise/tortoise-orm.git", branch = "develop"}
 aiomysql = "*"
+click = "*"
 
 [tool.poetry.dev-dependencies]
 taskipy = "*"
diff --git a/requirements.txt b/requirements.txt
index 093655f..9f16636 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,6 +2,7 @@ aiomysql==0.0.20
 aiosqlite==0.13.0
 cffi==1.14.0
 ciso8601==2.1.3; sys_platform != "win32" and implementation_name == "cpython"
+click==7.1.2
 cryptography==2.9.2
 iso8601==0.1.12; sys_platform == "win32" or implementation_name != "cpython"
 pycparser==2.20
diff --git a/setup.py b/setup.py
index e69de29..328cd46 100644
--- a/setup.py
+++ b/setup.py
@@ -0,0 +1,43 @@
+import os
+import re
+from setuptools import find_packages, setup
+
+
+def version():
+    ver_str_line = open('alice/__init__.py', 'rt').read()
+    mob = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", ver_str_line, re.M)
+    if not mob:
+        raise RuntimeError("Unable to find version string")
+    return mob.group(1)
+
+
+with open(os.path.join(os.path.dirname(__file__), 'README.rst')) as f:
+    long_description = f.read()
+
+
+def requirements():
+    return open('requirements.txt', 'rt').read().splitlines()
+
+
+setup(
+    name='alice',
+    version=version(),
+    description='A database migrations tool for Tortoise-ORM.',
+    author='long2ice',
+    long_description_content_type='text/x-rst',
+    long_description=long_description,
+    author_email='long2ice@gmail.com',
+    url='https://github.com/long2ice/alice',
+    license='MIT License',
+    packages=find_packages(include=['alice*']),
+    include_package_data=True,
+    zip_safe=True,
+    entry_points={
+        'console_scripts': ['alice = alice.cli:cli'],
+    },
+    platforms='any',
+    keywords=(
+        'migrations Tortoise-ORM mysql'
+    ),
+    install_requires=requirements(),
+)
diff --git a/tests/backends/mysql/__init__.py b/tests/backends/mysql/__init__.py
index cbd71a8..7d34ca5 100644
--- a/tests/backends/mysql/__init__.py
+++ b/tests/backends/mysql/__init__.py
@@ -2,17 +2,17 @@ from asynctest import TestCase
 from tortoise import Tortoise
 
 from alice.backends.mysql import MysqlDDL
-from tests.models import Category
+from alice.migrate import Migrate
 
 TORTOISE_ORM = {
     'connections': {
-        'default': 'mysql://root:123456@127.0.0.1:3306/test'
+        'default': 'mysql://root:123456@127.0.0.1:3306/test',
     },
     'apps': {
         'models': {
             'models': ['tests.models'],
             'default_connection': 'default',
-        }
+        },
     }
 }
 
@@ -22,6 +22,7 @@ class DBTestCase(TestCase):
         await Tortoise.init(config=TORTOISE_ORM)
         self.client = Tortoise.get_connection('default')
         self.ddl = MysqlDDL(self.client)
+        self.migrate = Migrate(ddl=self.ddl)
 
     async def tearDown(self) -> None:
         await Tortoise.close_connections()
diff --git a/tests/backends/mysql/test_ddl.py b/tests/backends/mysql/test_ddl.py
index 91932a7..82c3048 100644
--- a/tests/backends/mysql/test_ddl.py
+++ b/tests/backends/mysql/test_ddl.py
@@ -1,5 +1,5 @@
 from tests.backends.mysql import DBTestCase
-from tests.models import Category, User
+from tests.models import Category
 
 
 class TestDDL(DBTestCase):
@@ -47,7 +47,3 @@ class TestDDL(DBTestCase):
     def test_drop_fk(self):
         ret = self.ddl.drop_fk(Category, Category._meta.fields_map.get('user'))
         self.assertEqual(ret, "ALTER TABLE category DROP FOREIGN KEY fk_category_user_366ffa6f")
-
-    async def test_aa(self):
-        user = await User.get(username='test')
-        await user.save()
diff --git a/tests/backends/mysql/test_migrate.py b/tests/backends/mysql/test_migrate.py
new file mode 100644
index 0000000..204fc8f
--- /dev/null
+++ b/tests/backends/mysql/test_migrate.py
@@ -0,0 +1,7 @@
+from tests.backends.mysql import DBTestCase
+
+
+class TestMigrate(DBTestCase):
+    async def test_migrate(self):
+        self.migrate.diff_models_module('tests.models', 'tests.new_models')
+        print(self.migrate.operators)
diff --git a/tests/new_models.py b/tests/new_models.py
new file mode 100644
index 0000000..afcb683
--- /dev/null
+++ b/tests/new_models.py
@@ -0,0 +1,69 @@
+import datetime
+from enum import IntEnum
+from tortoise import fields, Model
+
+
+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='')
+
+
+    def __str__(self):
+        return f'{self.pk}#{self.username}'
+
+
+class Category(Model):
+    slug = fields.CharField(max_length=200)
+    name = fields.CharField(max_length=200)
+    user = fields.ForeignKeyField('new_models.User', description='User')
+    created_at = fields.DatetimeField(auto_now_add=True)
+
+    def __str__(self):
+        return f'{self.pk}#{self.name}'
+
+
+class Product(Model):
+    categories = fields.ManyToManyField('new_models.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)
+
+    def __str__(self):
+        return f'{self.pk}#{self.name}'
+
+
+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)
+
+    def __str__(self):
+        return f'{self.pk}#{self.label}'