nex_docus/backend/scripts/init_db.py

"""
数据库初始化脚本
用于创建数据库表并插入基础数据
"""
import sys
import os
import asyncio
from pathlib import Path

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

from sqlalchemy import text
from app.core.database import engine, async_session
from app.models.user import User
from app.models.role import Role
from app.models.menu import SystemMenu
from app.models.project import Project, ProjectMember
from app.core.security import get_password_hash


async def init_tables():
    """创建所有数据库表"""
    print("正在创建数据库表...")
    
    # 导入所有模型以确保它们被注册
    from app.models import Base
    
    async with engine.begin() as conn:
        # 创建所有表
        await conn.run_sync(Base.metadata.create_all)
    
    print("✓ 数据库表创建成功")


async def init_roles():
    """初始化系统角色"""
    print("正在初始化系统角色...")
    
    async with async_session() as session:
        # 检查是否已存在角色
        result = await session.execute(text("SELECT COUNT(*) FROM roles"))
        count = result.scalar()
        
        if count > 0:
            print("  角色已存在，跳过初始化")
            return
        
        # 创建默认角色
        roles = [
            Role(
                role_name="超级管理员",
                role_code="super_admin",
                description="系统超级管理员，拥有所有权限",
                status=1
            ),
            Role(
                role_name="管理员",
                role_code="admin",
                description="系统管理员",
                status=1
            ),
            Role(
                role_name="普通用户",
                role_code="user",
                description="普通用户",
                status=1
            ),
        ]
        
        for role in roles:
            session.add(role)
        
        await session.commit()
        print("✓ 系统角色初始化成功")


async def init_menus():
    """初始化系统菜单"""
    print("正在初始化系统菜单...")

    async with async_session() as session:
        # 检查是否已存在菜单
        result = await session.execute(text("SELECT COUNT(*) FROM system_menus"))
        count = result.scalar()

        if count > 0:
            print("  菜单已存在，跳过初始化")
            return

        # 创建系统菜单
        menus = [
            # 一级菜单
            SystemMenu(
                id=1,
                parent_id=0,
                menu_name="项目空间",
                menu_code="projects",
                menu_type=0,
                path="/projects",
                icon="ProjectOutlined",
                sort_order=3,
                visible=1,
                status=1
            ),
            SystemMenu(
                id=8,
                parent_id=0,
                menu_name="个人桌面",
                menu_code="desktop",
                menu_type=1,
                path="/desktop",
                component="Desktop",
                icon="DesktopOutlined",
                sort_order=1,
                visible=1,
                status=1
            ),
            SystemMenu(
                id=9,
                parent_id=0,
                menu_name="管理面板",
                menu_code="dashboard",
                menu_type=1,
                path="/dashboard",
                component="Dashboard",
                icon="DashboardOutlined",
                sort_order=2,
                visible=1,
                status=1
            ),
            SystemMenu(
                id=10,
                parent_id=0,
                menu_name="知识库空间",
                menu_code="knowledge",
                menu_type=0,
                path="/knowledge",
                icon="BookOutlined",
                sort_order=4,
                visible=1,
                status=1
            ),
            SystemMenu(
                id=20,
                parent_id=0,
                menu_name="系统管理",
                menu_code="system",
                menu_type=0,
                path="/system",
                icon="SettingOutlined",
                sort_order=5,
                visible=1,
                status=1
            ),

            # 项目空间子菜单
            SystemMenu(
                id=2,
                parent_id=1,
                menu_name="我的项目",
                menu_code="projects:my",
                menu_type=1,
                path="/projects",
                component="MyProjects",
                icon="FolderOutlined",
                sort_order=1,
                visible=1,
                status=1
            ),

            # 我的项目子菜单（按钮权限）
            SystemMenu(
                id=3,
                parent_id=2,
                menu_name="创建项目",
                menu_code="projects:create",
                menu_type=2,
                permission="projects:create",
                sort_order=1,
                visible=1,
                status=1
            ),
            SystemMenu(
                id=4,
                parent_id=2,
                menu_name="编辑项目",
                menu_code="projects:edit",
                menu_type=2,
                permission="projects:edit",
                sort_order=2,
                visible=1,
                status=1
            ),
            SystemMenu(
                id=5,
                parent_id=2,
                menu_name="删除项目",
                menu_code="projects:delete",
                menu_type=2,
                permission="projects:delete",
                sort_order=3,
                visible=1,
                status=1
            ),

            # 知识库空间子菜单
            SystemMenu(
                id=11,
                parent_id=10,
                menu_name="我的知识库",
                menu_code="knowledge:my",
                menu_type=1,
                path="/knowledge",
                component="MyKnowledge",
                icon="ReadOutlined",
                sort_order=1,
                visible=1,
                status=1
            ),

            # 我的知识库子菜单（按钮权限）
            SystemMenu(
                id=12,
                parent_id=11,
                menu_name="编辑知识库",
                menu_code="knowledge:edit",
                menu_type=2,
                permission="knowledge:edit",
                sort_order=1,
                visible=1,
                status=1
            ),
            SystemMenu(
                id=13,
                parent_id=11,
                menu_name="删除知识库",
                menu_code="knowledge:delete",
                menu_type=2,
                permission="knowledge:delete",
                sort_order=2,
                visible=1,
                status=1
            ),

            # 系统管理子菜单
            SystemMenu(
                id=14,
                parent_id=20,
                menu_name="权限管理",
                menu_code="system:permissions",
                menu_type=1,
                path="/system/permissions",
                component="System/Permissions",
                icon="SafetyOutlined",
                sort_order=3,
                visible=1,
                status=1
            ),
            SystemMenu(
                id=15,
                parent_id=20,
                menu_name="用户管理",
                menu_code="system:users",
                menu_type=1,
                path="/system/users",
                component="System/Users",
                icon="UserOutlined",
                sort_order=1,
                visible=1,
                status=1
            ),
            SystemMenu(
                id=16,
                parent_id=20,
                menu_name="角色管理",
                menu_code="system:roles",
                menu_type=1,
                path="/system/roles",
                component="System/Roles",
                icon="TeamOutlined",
                sort_order=2,
                visible=1,
                status=1
            ),
        ]

        for menu in menus:
            session.add(menu)

        await session.commit()
        print("✓ 系统菜单初始化成功")


async def init_admin_user():
    """初始化管理员用户"""
    print("正在初始化管理员账号...")

    # 从环境变量获取管理员信息
    admin_username = os.getenv("ADMIN_USERNAME", "admin")
    admin_password = os.getenv("ADMIN_PASSWORD", "admin@123")
    admin_email = os.getenv("ADMIN_EMAIL", "admin@unisspace.com")
    admin_nickname = os.getenv("ADMIN_NICKNAME", "系统管理员")

    async with async_session() as session:
        # 检查管理员是否已存在
        result = await session.execute(
            text("SELECT COUNT(*) FROM users WHERE username = :username"),
            {"username": admin_username}
        )
        count = result.scalar()

        if count > 0:
            print(f"  管理员账号 {admin_username} 已存在，跳过初始化")
            return

        # 创建管理员用户
        admin_user = User(
            username=admin_username,
            password_hash=get_password_hash(admin_password),
            nickname=admin_nickname,
            email=admin_email,
            status=1,
            is_superuser=True
        )

        session.add(admin_user)
        await session.flush()

        # 获取超级管理员角色
        result = await session.execute(
            text("SELECT id FROM roles WHERE role_code = 'super_admin' LIMIT 1")
        )
        role_id = result.scalar()

        if role_id:
            # 分配角色
            await session.execute(
                text("INSERT INTO user_roles (user_id, role_id) VALUES (:user_id, :role_id)"),
                {"user_id": admin_user.id, "role_id": role_id}
            )

        # 为超级管理员角色分配菜单权限：ID 9、20 及其所有子菜单
        # ID 9: 管理面板
        # ID 20: 系统管理及其子菜单 (14, 15, 16)
        authorized_menu_ids = [9, 20, 14, 15, 16]

        if role_id:
            for menu_id in authorized_menu_ids:
                await session.execute(
                    text("INSERT INTO role_menus (role_id, menu_id) VALUES (:role_id, :menu_id)"),
                    {"role_id": role_id, "menu_id": menu_id}
                )

        await session.commit()

        print("✓ 管理员账号初始化成功")
        print(f"  用户名: {admin_username}")
        print(f"  密码: {admin_password}")
        print(f"  邮箱: {admin_email}")
        print(f"  已授权菜单: ID {', '.join(map(str, authorized_menu_ids))}")
        print("  ⚠️  请登录后及时修改默认密码！")


async def main():
    """主函数"""
    print("=" * 60)
    print("NEX Docus 数据库初始化")
    print("=" * 60)
    
    try:
        # 1. 创建数据库表
        await init_tables()
        
        # 2. 初始化角色
        await init_roles()
        
        # 3. 初始化菜单
        await init_menus()
        
        # 4. 初始化管理员
        await init_admin_user()
        
        print("\n" + "=" * 60)
        print("✓ 数据库初始化完成！")
        print("=" * 60)
        
    except Exception as e:
        print(f"\n✗ 初始化失败: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
    finally:
        await engine.dispose()


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