codegen_cli.py

"""
@Author: li
@Email: lijianqiao2906@live.com
@FileName: codegen_cli.py
@DateTime: 2025年11月27日
@Docs: 代码生成器命令行工具
"""

import argparse
import asyncio
import sys
from pathlib import Path

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))


def get_table_info_from_db(table_name: str, use_async: bool = True):
    """从数据库获取表信息.

    Args:
        table_name: 表名.
        use_async: 是否使用异步引擎.

    Returns:
        TableInfo 对象或 None.
    """
    from app.utils.codegen import TableInspector

    if use_async:
        from app.db.session import engine

        inspector = TableInspector(engine)
        return asyncio.run(inspector.inspect_table_async(table_name))
    else:
        # 创建同步引擎
        from sqlalchemy import create_engine

        from app.core.config import settings

        sync_url = settings.DATABASE_URL.replace("+asyncpg", "").replace("+aiosqlite", "")
        sync_engine = create_engine(sync_url)
        inspector = TableInspector(sync_engine)
        return inspector.inspect_table(table_name)


def get_table_info_from_metadata(table_name: str):
    """从 SQLAlchemy MetaData 获取表信息.

    Args:
        table_name: 表名.

    Returns:
        TableInfo 对象或 None.
    """
    from app.db.base import Base
    from app.models import (  # noqa: F401 - 导入所有模型以注册到 Base
        AuditLog,
        Department,
        DictData,
        DictType,
        Job,
        JobLog,
        LoginLog,
        Menu,
        Message,
        Notice,
        NoticeRead,
        OnlineUser,
        Post,
        Role,
        SysConfig,
        TokenBlacklist,
        User,
    )
    from app.utils.codegen import TableInspector

    return TableInspector.from_metadata(Base.metadata, table_name)


def list_tables(use_metadata: bool = True) -> list[str]:
    """列出所有表.

    Args:
        use_metadata: 是否从 MetaData 获取 (不需要数据库连接).

    Returns:
        表名列表.
    """
    if use_metadata:
        from app.db.base import Base
        from app.models import (  # noqa: F401 - 导入所有模型以注册到 Base
            AuditLog,
            Department,
            DictData,
            DictType,
            Job,
            JobLog,
            LoginLog,
            Menu,
            Message,
            Notice,
            NoticeRead,
            OnlineUser,
            Post,
            Role,
            SysConfig,
            TokenBlacklist,
            User,
        )

        return list(Base.metadata.tables.keys())
    else:
        from app.db.session import engine
        from app.utils.codegen import TableInspector

        inspector = TableInspector(engine)
        return asyncio.run(inspector.get_all_tables_async())


def main():
    """命令行入口."""
    parser = argparse.ArgumentParser(
        description="FastAPI Admin 代码生成器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 列出所有表
  python codegen_cli.py --list

  # 从模型定义生成代码 (推荐，不需要数据库连接)
  python codegen_cli.py --table sys_product

  # 从数据库读取表结构生成代码
  python codegen_cli.py --table sys_product --from-db

  # 只预览不生成文件
  python codegen_cli.py --table sys_product --preview

  # 覆盖已存在的文件
  python codegen_cli.py --table sys_product --overwrite

  # 仅生成 Model 和 Schema
  python codegen_cli.py --table sys_product --only model schema

  # 生成仅超级管理员可访问的 API
  python codegen_cli.py --table sys_product --superuser-only

  # 不使用权限控制
  python codegen_cli.py --table sys_product --no-permission
        """,
    )

    parser.add_argument("--table", "-t", help="要生成代码的表名")
    parser.add_argument("--list", "-l", action="store_true", help="列出所有可用的表")
    parser.add_argument("--from-db", action="store_true", help="从数据库读取表结构 (默认从模型定义)")
    parser.add_argument("--preview", "-p", action="store_true", help="预览生成的代码 (不保存文件)")
    parser.add_argument("--overwrite", "-o", action="store_true", help="覆盖已存在的文件")
    parser.add_argument(
        "--only",
        nargs="+",
        choices=["model", "schema", "service", "api", "test"],
        help="只生成指定类型的代码",
    )
    parser.add_argument("--superuser-only", action="store_true", help="API 仅超级管理员可访问")
    parser.add_argument("--no-permission", action="store_true", help="不使用权限控制")
    parser.add_argument("--no-soft-delete", action="store_true", help="不使用软删除")
    parser.add_argument("--no-timestamp", action="store_true", help="不使用时间戳")
    parser.add_argument("--author", default="li", help="作者名 (默认: li)")
    parser.add_argument("--email", default="lijianqiao2906@live.com", help="邮箱")
    parser.add_argument("--output", help="输出目录 (默认: 项目根目录)")

    args = parser.parse_args()

    # 列出所有表
    if args.list:
        print("\n可用的数据库表:")
        print("-" * 40)
        tables = list_tables(use_metadata=not args.from_db)
        for table in sorted(tables):
            print(f"  - {table}")
        print(f"\n共 {len(tables)} 张表")
        return

    # 检查必需参数
    if not args.table:
        parser.print_help()
        print("\n错误: 请指定表名 (--table)")
        sys.exit(1)

    # 获取表信息
    print(f"\n正在分析表: {args.table}")
    if args.from_db:
        table_info = get_table_info_from_db(args.table)
    else:
        table_info = get_table_info_from_metadata(args.table)

    if not table_info:
        print(f"错误: 表 '{args.table}' 不存在")
        sys.exit(1)

    print(f"表注释: {table_info.comment or '无'}")
    print(f"列数量: {len(table_info.columns)}")
    print(f"主键: {table_info.primary_key}")

    # 创建配置
    from app.utils.codegen import GeneratorConfig

    config = GeneratorConfig(
        project_root=Path(args.output) if args.output else Path.cwd(),
        overwrite=args.overwrite,
        use_soft_delete=not args.no_soft_delete,
        use_timestamp=not args.no_timestamp,
        use_permission=not args.no_permission and not args.superuser_only,
        superuser_only=args.superuser_only,
        author=args.author,
        email=args.email,
    )

    # 处理 --only 参数
    if args.only:
        config.generate_model = "model" in args.only
        config.generate_schema = "schema" in args.only
        config.generate_service = "service" in args.only
        config.generate_api = "api" in args.only
        config.generate_test = "test" in args.only

    # 创建生成器
    from app.utils.codegen import CodeGenerator

    generator = CodeGenerator(config)

    # 预览或生成
    if args.preview:
        generator.preview(table_info)
    else:
        print("\n开始生成代码...")
        saved_files = generator.save_all(table_info)
        print(f"\n生成完成! 共生成 {len(saved_files)} 个文件")

        # 提示后续操作
        if saved_files:
            print("\n后续步骤:")
            print("1. 检查生成的代码并根据需要修改")
            if "model" in saved_files or config.generate_model:
                print(
                    f"2. 在 app/models/__init__.py 中导入新模型: from app.models.{table_info.module_name} import {table_info.class_name}"
                )
            if "api" in saved_files or config.generate_api:
                print(
                    f"3. 在 app/api/v1/api.py 中注册路由: from app.api.v1.{table_info.module_name}s import router as {table_info.module_name}_router"
                )
                print(f"   api_router.include_router({table_info.module_name}_router)")
            print("4. 运行数据库迁移: alembic revision --autogenerate -m 'add {table_info.module_name}'")
            print("5. 应用迁移: alembic upgrade head")
            print("6. 运行测试: pytest tests/test_{table_info.module_name}s.py -v")


if __name__ == "__main__":
    main()