init_db.py

"""
@Author: li
@Email: lijianqiao2906@live.com
@FileName: init_db.py
@DateTime: 2025年11月24日
@Docs: 数据库初始化脚本
"""

import asyncio
import logging
import tomllib
from pathlib import Path
from typing import Any

from loguru import logger
from sqlalchemy import select, text
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.config import settings
from app.core.security import get_password_hash
from app.db.base import Base
from app.db.session import AsyncSessionLocal, engine
from app.models import Department, DictData, DictType, Menu, Post, Role, SysConfig, User

# 屏蔽 SQLAlchemy 的 SQL 日志
logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)

# 配置文件目录
INIT_DATA_DIR = Path(__file__).parent / "init_data"


def load_toml(filename: str) -> dict[str, Any]:
    """加载 TOML 配置文件."""
    filepath = INIT_DATA_DIR / filename
    with open(filepath, "rb") as f:
        return tomllib.load(f)


# ==================== 菜单初始化 ====================
async def init_menus(db: AsyncSession) -> None:
    """初始化菜单数据."""
    stmt = select(Menu).where(Menu.name == "System")
    result = await db.execute(stmt)
    if result.scalars().first():
        logger.info("菜单已存在，跳过创建。")
        return

    logger.info("创建菜单...")
    data = load_toml("menus.toml")

    async def create_menu(menu_data: dict, parent_id: int | None = None) -> None:
        """递归创建菜单."""
        children = menu_data.pop("children", [])
        buttons = menu_data.pop("buttons", [])

        menu = Menu(
            title=menu_data["title"],
            name=menu_data["name"],
            component_name=menu_data.get("component_name"),
            icon=menu_data.get("icon"),
            path=menu_data.get("path"),
            component=menu_data.get("component"),
            sort=menu_data.get("sort", 0),
            type=menu_data.get("type", 1),
            perms=menu_data.get("perms"),
            parent_id=parent_id,
        )
        db.add(menu)
        await db.flush()

        # 创建按钮
        for btn in buttons:
            db.add(
                Menu(
                    title=btn["title"],
                    type=2,
                    parent_id=menu.id,
                    perms=btn["perms"],
                )
            )

        # 递归创建子菜单
        for child in children:
            await create_menu(child, menu.id)

    for menu_data in data.get("menus", []):
        await create_menu(menu_data)


# ==================== 角色初始化 ====================
async def init_roles(db: AsyncSession) -> Role | None:
    """初始化角色数据."""
    stmt = select(Role).where(Role.code == "admin")
    result = await db.execute(stmt)
    admin_role = result.scalars().first()

    if admin_role:
        logger.info("角色已存在，跳过创建。")
        return admin_role

    logger.info("创建角色...")
    data = load_toml("roles.toml")

    for role_data in data.get("roles", []):
        role = Role(
            name=role_data["name"],
            code=role_data["code"],
            description=role_data.get("description"),
        )
        db.add(role)
        if role_data["code"] == "admin":
            admin_role = role

    await db.flush()
    return admin_role


# ==================== 用户初始化 ====================
async def init_superuser(db: AsyncSession, admin_role: Role | None) -> None:
    """初始化超级管理员."""
    stmt = select(User).where(User.username == settings.FIRST_SUPERUSER)
    result = await db.execute(stmt)
    if result.scalars().first():
        logger.info("超级用户已存在，跳过创建。")
        return

    logger.info("创建超级用户...")
    user = User(
        username=settings.FIRST_SUPERUSER,
        password_hash=get_password_hash(settings.FIRST_SUPERUSER_PASSWORD),
        email=settings.FIRST_SUPERUSER_EMAIL,
        phone=settings.FIRST_SUPERUSER_PHONE,
        is_superuser=True,
        is_active=True,
        nickname="Admin",
    )
    if admin_role:
        user.roles.append(admin_role)
    db.add(user)


# ==================== 系统配置初始化 ====================
async def init_configs(db: AsyncSession) -> None:
    """初始化系统配置."""
    stmt = select(SysConfig).where(SysConfig.key == "site.name")
    result = await db.execute(stmt)
    if result.scalars().first():
        logger.info("系统配置已存在，跳过创建。")
        return

    logger.info("创建系统配置...")
    data = load_toml("configs.toml")

    for config_data in data.get("configs", []):
        db.add(
            SysConfig(
                name=config_data["name"],
                key=config_data["key"],
                value=config_data["value"],
                config_type=config_data.get("config_type", "system"),
                group=config_data.get("group", "site"),
                is_public=config_data.get("is_public", True),
                remark=config_data.get("remark"),
            )
        )


# ==================== 部门初始化 ====================
async def init_depts(db: AsyncSession) -> None:
    """初始化部门数据."""
    stmt = select(Department).where(Department.code == "ROOT")
    result = await db.execute(stmt)
    if result.scalars().first():
        logger.info("部门已存在，跳过创建。")
        return

    logger.info("创建部门...")
    data = load_toml("depts.toml")

    for dept_data in data.get("depts", []):
        children = dept_data.pop("children", [])

        dept = Department(
            name=dept_data["name"],
            code=dept_data["code"],
            sort=dept_data.get("sort", 0),
            leader=dept_data.get("leader"),
            status=True,
        )
        db.add(dept)
        await db.flush()

        for child in children:
            db.add(
                Department(
                    name=child["name"],
                    code=child["code"],
                    sort=child.get("sort", 0),
                    parent_id=dept.id,
                    status=True,
                )
            )


# ==================== 岗位初始化 ====================
async def init_posts(db: AsyncSession) -> None:
    """初始化岗位数据."""
    stmt = select(Post).where(Post.code == "CEO")
    result = await db.execute(stmt)
    if result.scalars().first():
        logger.info("岗位已存在，跳过创建。")
        return

    logger.info("创建岗位...")
    data = load_toml("posts.toml")

    for post_data in data.get("posts", []):
        db.add(
            Post(
                code=post_data["code"],
                name=post_data["name"],
                sort=post_data.get("sort", 0),
                status=True,
                remark=post_data.get("remark"),
            )
        )


# ==================== 字典初始化 ====================
async def init_dicts(db: AsyncSession) -> None:
    """初始化字典类型和字典数据."""
    stmt = select(DictType).where(DictType.code == "sys_common_status")
    result = await db.execute(stmt)
    if result.scalars().first():
        logger.info("字典已存在，跳过创建。")
        return

    logger.info("创建字典...")
    data = load_toml("dicts.toml")

    for dict_type_data in data.get("dict_types", []):
        items = dict_type_data.pop("items", [])

        dict_type = DictType(
            name=dict_type_data["name"],
            code=dict_type_data["code"],
            status=True,
            remark=dict_type_data.get("remark"),
        )
        db.add(dict_type)
        await db.flush()

        for item in items:
            db.add(
                DictData(
                    dict_type_id=dict_type.id,
                    label=item["label"],
                    value=item["value"],
                    sort=item.get("sort", 0),
                    css_class=item.get("css_class"),
                    list_class=item.get("list_class"),
                    status=True,
                )
            )


# ==================== 主初始化函数 ====================
async def init_db(db: AsyncSession) -> None:
    """初始化数据库数据."""
    # 1. 菜单
    await init_menus(db)

    # 2. 角色
    admin_role = await init_roles(db)

    # 3. 超级用户
    await init_superuser(db, admin_role)

    # 4. 系统配置
    await init_configs(db)

    # 5. 部门
    await init_depts(db)

    # 6. 岗位
    await init_posts(db)

    # 7. 字典
    await init_dicts(db)

    await db.commit()
    logger.info("初始化数据完成!")


# ==================== 数据库操作函数 ====================
async def reset_db() -> None:
    """重置数据库 (清空表数据)."""
    async with engine.begin() as conn:
        logger.info("重置数据库 (清空表数据)...")
        await conn.execute(
            text(
                "DO $$ DECLARE r RECORD; BEGIN FOR r IN "
                "(SELECT tablename FROM pg_tables WHERE schemaname = current_schema()) "
                "LOOP EXECUTE 'TRUNCATE TABLE ' || quote_ident(r.tablename) || ' CASCADE'; "
                "END LOOP; END $$;"
            )
        )


async def drop_db() -> None:
    """删除数据库 (删除表)."""
    async with engine.begin() as conn:
        logger.info("删除所有表...")
        await conn.run_sync(Base.metadata.drop_all)


async def create_tables() -> None:
    """创建数据库表."""
    async with engine.begin() as conn:
        logger.info("创建所有表...")
        await conn.run_sync(Base.metadata.create_all)


async def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(description="数据库初始化脚本")
    parser.add_argument("--reset", action="store_true", help="重置数据库 (清空表数据)")
    parser.add_argument("--drop", action="store_true", help="删除数据库 (删除表)")
    parser.add_argument("--init", action="store_true", help="初始化数据库 (创建基础数据)")

    args = parser.parse_args()

    if args.drop:
        await drop_db()
        logger.info("删除数据库完成。")

    if args.reset:
        await reset_db()
        logger.info("重置数据库完成。")

    if args.init:
        await create_tables()
        async with AsyncSessionLocal() as session:
            await init_db(session)
        logger.info("初始化数据库完成。")

    if not (args.reset or args.drop or args.init):
        parser.print_help()


if __name__ == "__main__":
    asyncio.run(main())