complete InspectDb

tortoise · Dec 25, 2020 · 5ae8b9e · 5ae8b9e
1 parent 55a6d4b
commit 5ae8b9e
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -7,8 +7,8 @@
 
 ## Introduction
 
-Aerich is a database migrations tool for Tortoise-ORM, which like alembic for SQLAlchemy, or Django ORM with it\'s
-own migrations solution.
+Aerich is a database migrations tool for Tortoise-ORM, which like alembic for SQLAlchemy, or Django ORM with it\'s own
+migrations solution.
 
 **Important: You can only use absolutely import in your `models.py` to make `aerich` work.**
 
@@ -40,14 +40,14 @@ Commands:
   history    List all migrate items.
   init       Init config file and generate root migrate location.
   init-db    Generate schema and generate app migrate location.
+  inspectdb  Introspects the database tables to standard output as...
   migrate    Generate migrate changes file.
   upgrade    Upgrade to latest version.
 ```
 
 ## Usage
 
-You need add `aerich.models` to your `Tortoise-ORM` config first,
-example:
+You need add `aerich.models` to your `Tortoise-ORM` config first, example:
 
 ```python
 TORTOISE_ORM = {
@@ -109,7 +109,8 @@ Success migrate 1_202029051520102929_drop_column.sql
 Format of migrate filename is
 `{version_num}_{datetime}_{name|update}.sql`.
 
-And if `aerich` guess you are renaming a column, it will ask `Rename {old_column} to {new_column} [True]`, you can choice `True` to rename column without column drop, or choice `False` to drop column then create.
+And if `aerich` guess you are renaming a column, it will ask `Rename {old_column} to {new_column} [True]`, you can
+choice `True` to rename column without column drop, or choice `False` to drop column then create.
 
 ### Upgrade to latest version
 
@@ -163,6 +164,33 @@ Now your db rollback to specified version.
 1_202029051520102929_drop_column.sql
 ```
 
+### Inspect db tables to TortoiseORM model
+
+```shell
+Usage: aerich inspectdb [OPTIONS]
+
+  Introspects the database tables to standard output as TortoiseORM model.
+
+Options:
+  -t, --table TEXT  Which tables to inspect.
+  -h, --help        Show this message and exit.
+```
+
+Inspect all tables and print to console:
+
+```shell
+aerich --app models inspectdb -t user 
+```
+
+Inspect a specified table in default app and redirect to `models.py`:
+
+```shell
+aerich inspectdb -t user > models.py
+```
+
+Note that this command is restricted, which is not supported in some solutions, such as `IntEnumField`
+and `ForeignKeyField` and so on.
+
 ### Multiple databases
 
 ```python
@@ -173,7 +201,7 @@ tortoise_orm = {
     },
     "apps": {
         "models": {"models": ["tests.models", "aerich.models"], "default_connection": "default"},
-        "models_second": {"models": ["tests.models_second"], "default_connection": "second",},
+        "models_second": {"models": ["tests.models_second"], "default_connection": "second", },
     },
 }
 ```

diff --git a/aerich/inspectdb.py b/aerich/inspectdb.py
@@ -1,3 +1,4 @@
+import sys
 from typing import List, Optional
 
 from ddlparse import DdlParse
@@ -6,6 +7,17 @@
 
 
 class InspectDb:
+    _table_template = "class {table}(Model):\n"
+    _field_template_mapping = {
+        "INT": "    {field} = fields.IntField({pk}{unique}{comment})",
+        "SMALLINT": "    {field} = fields.IntField({pk}{unique}{comment})",
+        "TINYINT": "    {field} = fields.BooleanField({null}{default}{comment})",
+        "VARCHAR": "    {field} = fields.CharField({pk}{unique}{length}{null}{default}{comment})",
+        "LONGTEXT": "    {field} = fields.TextField({null}{default}{comment})",
+        "TEXT": "    {field} = fields.TextField({null}{default}{comment})",
+        "DATETIME": "    {field} = fields.DatetimeField({null}{default}{comment})",
+    }
+
     def __init__(self, conn: BaseDBAsyncClient, tables: Optional[List[str]] = None):
         self.conn = conn
         self.tables = tables
@@ -14,9 +26,9 @@ def __init__(self, conn: BaseDBAsyncClient, tables: Optional[List[str]] = None):
     async def show_create_tables(self):
         if self.DIALECT == MySQLSchemaGenerator.DIALECT:
             if not self.tables:
-                sql_tables = f"SELECT table_name FROM information_schema.tables WHERE table_schema = '{self.conn.database}';"
+                sql_tables = f"SELECT table_name FROM information_schema.tables WHERE table_schema = '{self.conn.database}';"  # nosec: B608
                 ret = await self.conn.execute_query(sql_tables)
-                self.tables = map(lambda x: x[0], ret)
+                self.tables = map(lambda x: x["TABLE_NAME"], ret[1])
             for table in self.tables:
                 sql_show_create_table = f"SHOW CREATE TABLE {table}"
                 ret = await self.conn.execute_query(sql_show_create_table)
@@ -26,7 +38,49 @@ async def show_create_tables(self):
 
     async def inspect(self):
         ddl_list = self.show_create_tables()
+        result = "from tortoise import Model, fields\n\n\n"
+        tables = []
         async for ddl in ddl_list:
             parser = DdlParse(ddl, DdlParse.DATABASE.mysql)
             table = parser.parse()
-            print(table)
+            name = table.name.title()
+            columns = table.columns
+            fields = []
+            model = self._table_template.format(table=name)
+            for column_name, column in columns.items():
+                comment = default = length = unique = null = pk = ""
+                if column.primary_key:
+                    pk = "pk=True, "
+                if column.unique:
+                    unique = "unique=True, "
+                if column.data_type == "VARCHAR":
+                    length = f"max_length={column.length}, "
+                if not column.not_null:
+                    null = "null=True, "
+                if column.default is not None:
+                    if column.data_type == "TINYINT":
+                        default = f"default={'True' if column.default == '1' else 'False'}, "
+                    elif column.data_type == "DATETIME":
+                        if "CURRENT_TIMESTAMP" in column.default:
+                            if "ON UPDATE CURRENT_TIMESTAMP" in ddl:
+                                default = "auto_now_add=True, "
+                            else:
+                                default = "auto_now=True, "
+                    else:
+                        default = f"default={column.default}, "
+
+                if column.comment:
+                    comment = f"description='{column.comment}', "
+
+                field = self._field_template_mapping[column.data_type].format(
+                    field=column_name,
+                    pk=pk,
+                    unique=unique,
+                    length=length,
+                    null=null,
+                    default=default,
+                    comment=comment,
+                )
+                fields.append(field)
+            tables.append(model + "\n".join(fields))
+        sys.stdout.write(result + "\n\n\n".join(tables))