dingdanquanliucheng/backend/scripts/bootstrap_data.py

"""数据库初始化种子数据脚本。

首次部署或重建数据库时运行，用于插入初始数据：
- 角色（admin、manager、salesman、driver）
- 菜单与权限码
- 角色-菜单关联
- 系统配置项
- 产品分类
- 默认管理员用户（admin01 / 123456）
- 默认业务员用户（sales01 / 123456）

运行方式: python -m backend.scripts.bootstrap_data
"""

from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from sqlalchemy import select

from backend.app.core.security import hash_password
from backend.app.db import SessionLocal
from backend.app.models.business import ProductCategory
from backend.app.models.system import Menu, Role, RoleMenu, SystemConfig, User
from backend.app.services.bootstrap import (
    DEFAULT_ADMIN_USER,
    DEFAULT_CATEGORIES,
    DEFAULT_CONFIGS,
    DEFAULT_MENUS,
    DEFAULT_ROLE_MENU_CODES,
    DEFAULT_ROLES,
    DEFAULT_SALES_USER,
)


def main() -> None:
    """执行全量种子数据初始化，幂等操作（已存在则更新，不存在则插入）。"""
    session = SessionLocal()
    try:
        role_ids = seed_roles(session)
        menu_ids = seed_menus(session)
        seed_role_menu_relations(session, role_ids, menu_ids)
        seed_configs(session)
        seed_categories(session)
        seed_admin_user(session, role_ids)
        seed_sales_user(session, role_ids)
        session.commit()
        print("bootstrap finished")
        print(f"roles: {len(role_ids)}")
        print(f"menus: {len(menu_ids)}")
        print(f"configs: {len(DEFAULT_CONFIGS)}")
        print(f"categories: {len(DEFAULT_CATEGORIES)}")
        print(f"admin: {DEFAULT_ADMIN_USER['username']}")
        print(f"sales: {DEFAULT_SALES_USER['username']}")
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def seed_roles(session) -> dict[str, int]:
    """初始化角色数据。

    按 role_code 幂等插入或更新角色记录。

    Args:
        session: 数据库会话。

    Returns:
        dict[str, int]: 角色编码到角色 ID 的映射。
    """
    role_ids: dict[str, int] = {}
    for item in DEFAULT_ROLES:
        existed = session.execute(select(Role).where(Role.role_code == item["role_code"])).scalar_one_or_none()
        if existed is None:
            existed = Role(**item)
            session.add(existed)
            session.flush()
        else:
            existed.role_name = item["role_name"]
            existed.status = item["status"]
            existed.remark = item["remark"]
            session.add(existed)
            session.flush()
        role_ids[item["role_code"]] = existed.id
    return role_ids


def seed_menus(session) -> dict[str, int]:
    """初始化菜单数据。

    按 permission_code 幂等插入或更新菜单记录。

    Args:
        session: 数据库会话。

    Returns:
        dict[str, int]: 权限码到菜单 ID 的映射。
    """
    menu_ids: dict[str, int] = {}
    for item in DEFAULT_MENUS:
        existed = session.execute(select(Menu).where(Menu.permission_code == item["permission_code"])).scalar_one_or_none()
        if existed is None:
            data = {k: v for k, v in item.items() if k != "id"}
            existed = Menu(**data)
            session.add(existed)
            session.flush()
        else:
            existed.parent_id = item["parent_id"]
            existed.menu_name = item["menu_name"]
            existed.menu_path = item["menu_path"]
            existed.menu_type = item["menu_type"]
            existed.icon = item["icon"]
            existed.sort_no = item["sort_no"]
            existed.status = item["status"]
            session.add(existed)
            session.flush()
        menu_ids[item["permission_code"]] = existed.id
    return menu_ids


def seed_role_menu_relations(session, role_ids: dict[str, int], menu_ids: dict[str, int]) -> None:
    """同步角色-菜单关联关系。

    根据 DEFAULT_ROLE_MENU_CODES 定义，删除多余关联、补充缺失关联。

    Args:
        session: 数据库会话。
        role_ids: 角色编码到 ID 的映射。
        menu_ids: 权限码到菜单 ID 的映射。
    """
    for role_code, permission_codes in DEFAULT_ROLE_MENU_CODES.items():
        role_id = role_ids[role_code]
        expected_menu_ids = {menu_ids[code] for code in permission_codes if code in menu_ids}
        existed = list(session.execute(select(RoleMenu).where(RoleMenu.role_id == role_id)).scalars())
        existed_menu_ids = {item.menu_id for item in existed}

        # 删除不再需要的关联
        for relation in existed:
            if relation.menu_id not in expected_menu_ids:
                session.delete(relation)

        # 补充缺失的关联
        missing_menu_ids = expected_menu_ids - existed_menu_ids
        for menu_id in missing_menu_ids:
            session.add(RoleMenu(role_id=role_id, menu_id=menu_id))


def seed_configs(session) -> None:
    """初始化系统配置项（物流超时、沉默天数、欠款阈值等）。"""
    for item in DEFAULT_CONFIGS:
        existed = session.execute(select(SystemConfig).where(SystemConfig.config_key == item["config_key"])).scalar_one_or_none()
        if existed is None:
            session.add(SystemConfig(**item))
            continue
        existed.config_value = item["config_value"]
        existed.config_name = item["config_name"]
        existed.remark = item.get("remark")
        existed.status = item.get("status", 1)
        session.add(existed)


def seed_categories(session) -> None:
    """初始化产品分类数据。"""
    for item in DEFAULT_CATEGORIES:
        existed = session.execute(select(ProductCategory).where(ProductCategory.category_code == item["category_code"])).scalar_one_or_none()
        if existed is None:
            session.add(ProductCategory(**item))
            continue
        existed.category_name = item["category_name"]
        existed.sort_no = item["sort_no"]
        existed.status = item["status"]
        existed.remark = item.get("remark")
        session.add(existed)


def seed_admin_user(session, role_ids: dict[str, int]) -> None:
    """初始化默认管理员用户（admin）。

    如果已存在则更新基本信息，密码仅在为空时设置默认值 123456。
    """
    existed = session.execute(select(User).where(User.username == DEFAULT_ADMIN_USER["username"])).scalar_one_or_none()
    password_hash = hash_password("123456")
    if existed is None:
        user_data = {**DEFAULT_ADMIN_USER, "role_id": role_ids["admin"]}
        session.add(
            User(
                **user_data,
                password_hash=password_hash,
            )
        )
        return

    existed.real_name = DEFAULT_ADMIN_USER["real_name"]
    existed.mobile = DEFAULT_ADMIN_USER["mobile"]
    existed.role_id = role_ids["admin"]
    existed.status = DEFAULT_ADMIN_USER["status"]
    if not existed.password_hash:
        existed.password_hash = password_hash
    session.add(existed)


def seed_sales_user(session, role_ids: dict[str, int]) -> None:
    """初始化默认业务员用户（sales01）。

    如果已存在则更新基本信息，密码仅在为空时设置默认值 123456。
    """
    existed = session.execute(select(User).where(User.username == DEFAULT_SALES_USER["username"])).scalar_one_or_none()
    password_hash = hash_password("123456")
    if existed is None:
        user_data = {**DEFAULT_SALES_USER, "role_id": role_ids["salesman"]}
        session.add(
            User(
                **user_data,
                password_hash=password_hash,
            )
        )
        return

    existed.real_name = DEFAULT_SALES_USER["real_name"]
    existed.mobile = DEFAULT_SALES_USER["mobile"]
    existed.role_id = role_ids["salesman"]
    existed.status = DEFAULT_SALES_USER["status"]
    if not existed.password_hash:
        existed.password_hash = password_hash
    session.add(existed)


if __name__ == "__main__":
    main()