首次提交

paste/db/gen_models.py

@@ -0,0 +1,51 @@
+"""
+生成数据模型代码，注意这里显式排除了在配置文件 RBAC 中配置的数据表。
+如果不需要排除，可直接使用 sqlacodegen 自带的命令。
+"""
+import subprocess
+from os import path
+
+import pandas as pd
+
+from paste.core import config
+from paste.core.logging import echo_log
+from paste.db.basetable import BaseTable
+
+exclude_tables = list(config.get_config('rbac.table').values())
+"""
+需要排除的数据表，默认排除 RBAC 数据表，这部分表格在 RBAC 模块中已经配置好了。
+"""
+
+
+def db_engin():
+    return config.get_config('db_engine.engine')
+
+
+async def sqlacodegen(is_exclude_rbac_table: bool = True):
+    """
+    生成代码文件。
+
+    :param is_exclude_rbac_table: 是否排除 RBAC 相关的数据表，默认排除不生成数据模型
+    """
+    _table_names = await BaseTable.tables_in_db()
+
+    # 剔除 RBAC 数据表
+    if _table_names and is_exclude_rbac_table:
+        # 转换为 DataFrame
+        _tables_df = pd.DataFrame(_table_names)
+        # 剔除不要包含的表
+        _name_df: pd.DataFrame = _tables_df.loc[~_tables_df.iloc[:, 0].isin(exclude_tables)]
+        # 转换剩余数据为'表名'字符串列表
+        _table_names = _name_df[0].tolist()
+
+    if len(_table_names) == 0:
+        return
+
+    echo_log(f"将为以下表生成数据模型：{_table_names}")
+    _tables = f"--tables={','.join(_table_names)}"
+    # 默认创建在当前目录的 models 目录中
+    _outfile = f"--outfile={path.join(path.curdir, 'models', 'db_models.py')}"
+
+    _engin = db_engin()
+    subprocess.call(['sqlacodegen', _engin, _tables, _outfile])
+    echo_log(f"生成完成.")