dingdanquanliucheng/backend/scripts/bootstrap_data.py

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,
)


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)
        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']}")
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def seed_roles(session) -> dict[str, int]:
    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]:
    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:
            existed = Menu(**item)
            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:
    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:
    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)


if __name__ == "__main__":
    main()