geMoldInsight/src/shared/database/init_db.py

import asyncio
import sys
from pathlib import Path
from sqlalchemy import text, select
from shared.database.database import db_manager
from shared.models.database import User, Role, Permission, UserRole, RolePermission
from shared.services.auth_service import get_password_hash
from shared.config.settings import settings
from shared.utils.logger import get_logger
from alembic.config import Config
from alembic import command

logger = get_logger(__name__)

_ALEMBIC_INI = Path(__file__).resolve().parents[3] / "alembic.ini"


def _alembic_stamp_head() -> None:
    """将当前 DB 标记为已到最新版本（基线既有 DB，不执行 SQL）"""
    cfg = Config(str(_ALEMBIC_INI))
    command.stamp(cfg, "head")


def _alembic_upgrade_head() -> None:
    """执行 alembic 迁移到最新版本"""
    cfg = Config(str(_ALEMBIC_INI))
    command.upgrade(cfg, "head")


async def _run_alembic_migrations() -> None:
    """以 alembic 管理 schema：既有未纳入管理的 DB 自动 stamp 基线，再 upgrade head

    - 全新 DB：upgrade head 执行初始迁移，创建全部表
    - 既有已纳入管理：upgrade head 为 no-op
    - 既有但无 alembic_version（历史 DB）：先 stamp head 基线，再 upgrade（no-op）
    """
    async with db_manager.engine.begin() as conn:
        has_alembic = await conn.execute(text("SELECT to_regclass('public.alembic_version')")).scalar()
        if not has_alembic:
            table_count = await conn.execute(
                text("SELECT count(*) FROM information_schema.tables WHERE table_schema='public' AND table_name <> 'alembic_version'")
            ).scalar()
            if table_count and table_count > 0:
                logger.info("检测到既有 DB 未纳入 alembic 管理，自动 stamp head 作为基线")
                await asyncio.to_thread(_alembic_stamp_head)
    await asyncio.to_thread(_alembic_upgrade_head)

DEFAULT_PERMISSIONS = [
    {"code": "view_dashboard", "name": "查看仪表盘", "module": "dashboard"},
    {"code": "view_moldinsight", "name": "使用模具分析", "module": "moldinsight"},
    {"code": "upload_file", "name": "上传文件", "module": "moldinsight"},
    {"code": "view_history", "name": "查看历史记录", "module": "moldinsight"},
    {"code": "view_inventory", "name": "查看库存", "module": "inventory"},
    {"code": "manage_inventory", "name": "管理库存", "module": "inventory"},
    {"code": "view_products", "name": "查看产品", "module": "inventory"},
    {"code": "manage_products", "name": "管理产品", "module": "inventory"},
    {"code": "view_suppliers", "name": "查看供应商", "module": "inventory"},
    {"code": "manage_suppliers", "name": "管理供应商", "module": "inventory"},
    {"code": "view_customers", "name": "查看客户", "module": "inventory"},
    {"code": "manage_customers", "name": "管理客户", "module": "inventory"},
    {"code": "view_finance", "name": "查看财务", "module": "finance"},
    {"code": "manage_receipts", "name": "管理收款", "module": "finance"},
    {"code": "manage_payments", "name": "管理付款", "module": "finance"},
    {"code": "void_finance_transaction", "name": "作废财务单据", "module": "finance"},
    {"code": "view_users", "name": "查看用户", "module": "admin"},
    {"code": "manage_users", "name": "管理用户", "module": "admin"},
    {"code": "manage_roles", "name": "管理角色", "module": "admin"},
]

DEFAULT_ROLES = [
    {"code": "admin", "name": "管理员", "description": "系统管理员，拥有所有权限", "is_system": True, "permissions": ["view_dashboard", "view_moldinsight", "upload_file", "view_history", "view_inventory", "manage_inventory", "view_products", "manage_products", "view_suppliers", "manage_suppliers", "view_customers", "manage_customers", "view_finance", "manage_receipts", "manage_payments", "void_finance_transaction", "view_users", "manage_users", "manage_roles"]},
    {"code": "user", "name": "普通用户", "description": "普通用户，可使用模具分析和查看库存", "is_system": False, "permissions": ["view_dashboard", "view_moldinsight", "upload_file", "view_history", "view_inventory", "view_products", "view_suppliers", "view_customers", "view_finance", "manage_receipts", "manage_payments"]},
    {"code": "viewer", "name": "只读用户", "description": "只读用户，只能查看数据", "is_system": False, "permissions": ["view_dashboard", "view_moldinsight", "view_history", "view_inventory", "view_products", "view_suppliers", "view_customers", "view_finance"]},
]


async def init_permissions(session):
    """初始化权限"""
    result = await session.execute(select(Permission))
    existing_perms = result.scalars().all()
    
    if existing_perms:
        logger.info("权限已初始化")
        return
    
    perm_map = {}
    for perm_data in DEFAULT_PERMISSIONS:
        perm = Permission(**perm_data)
        session.add(perm)
        await session.flush()
        perm_map[perm.code] = perm.id
    
    logger.info(f"创建了 {len(DEFAULT_PERMISSIONS)} 个权限")
    return perm_map


async def init_roles(session, perm_map):
    """初始化角色"""
    result = await session.execute(select(Role))
    existing_roles = result.scalars().all()
    
    if existing_roles:
        logger.info("角色已初始化")
        return
    
    for role_data in DEFAULT_ROLES:
        perm_ids = [perm_map[code] for code in role_data.pop("permissions")]
        role = Role(**role_data)
        session.add(role)
        await session.flush()
        
        for perm_id in perm_ids:
            rp = RolePermission(role_id=role.id, permission_id=perm_id)
            session.add(rp)
    
    logger.info(f"创建了 {len(DEFAULT_ROLES)} 个角色")


async def create_admin_user(session):
    """创建默认管理员"""
    result = await session.execute(select(User).where(User.username == settings.ADMIN_USERNAME))
    existing_admin = result.scalar_one_or_none()
    
    if existing_admin:
        logger.info("管理员账户已存在")
        return
    
    admin = User(
        username=settings.ADMIN_USERNAME,
        email=settings.ADMIN_EMAIL,
        hashed_password=get_password_hash(settings.ADMIN_PASSWORD),
        full_name=settings.ADMIN_FULL_NAME,
        is_active=True
    )
    session.add(admin)
    await session.flush()
    
    result = await session.execute(select(Role).where(Role.code == "admin"))
    admin_role = result.scalar_one_or_none()
    
    if admin_role:
        user_role = UserRole(user_id=admin.id, role_id=admin_role.id)
        session.add(user_role)
    
    await session.commit()
    logger.info(f"创建了管理员账户: {settings.ADMIN_USERNAME}")


async def init_database(keep_connected: bool = True):
    """初始化数据库"""
    try:
        await db_manager.connect()
        await _run_alembic_migrations()
        
        async with db_manager.session() as session:
            perm_map = await init_permissions(session)
            if perm_map is None:
                # Permissions already existed, fetch them from database
                result = await session.execute(select(Permission))
                perms = result.scalars().all()
                perm_map = {perm.code: perm.id for perm in perms}
            await init_roles(session, perm_map)
            await create_admin_user(session)


        logger.info("数据库初始化完成")
        print("=" * 60)
        print("数据库初始化成功！")
        print("=" * 60)
        print(f"管理员用户名: {settings.ADMIN_USERNAME}")
        print(f"管理员邮箱: {settings.ADMIN_EMAIL}")
        print("=" * 60)
        print("可以在 .env 文件中修改管理员配置:")
        print("  ADMIN_USERNAME")
        print("  ADMIN_EMAIL")
        print("  ADMIN_FULL_NAME")
        print("=" * 60)
        
        return True
        
    except Exception as e:
        logger.error(f"数据库初始化失败: {e}")
        print(f"数据库初始化失败: {e}")
        return False
    finally:
        if not keep_connected:
            await db_manager.disconnect()


if __name__ == "__main__":
    asyncio.run(init_database(keep_connected=False))