shanxia2.0/backend/scripts/setup_business_roles.py

#!/usr/bin/env python3
"""重置业务角色：保留超级管理员，创建 5 个销售体系角色并授权菜单、绑定用户。

用法：
  cd backend && ./venv/bin/python scripts/setup_business_roles.py --env-file .env.production
  ./venv/bin/python scripts/setup_business_roles.py --dry-run --env-file .env.production
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Set

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine

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

from app.core.utils import generate_uuid

TENANT_NO = "shanxia2.0"
ADMIN_USER_ID = "d3000000000000000000000000000001"
SUPER_ADMIN_ROLE_ID = "097bd2c6-4e72-11f1-9f47-00163e3bdd96"

# 四个一级业务模块根菜单
MODULE_MENU_ROOTS = [
    "CRM00000000000000000000000000001",  # CRM管理
    "BX00000000000000000000000000001",   # 报销管理
    "DB00000000000000000000000000001",   # 地磅管理
    "PL_MENU_ROOT",                      # 运输管理
]

NEW_ROLES: Dict[str, str] = {
    "管理员": "R65CAAE057914CCCAE2A81BE627A5001",
    "销售主管": "R65CAAE057914CCCAE2A81BE627A5002",
    "销售计划员": "R65CAAE057914CCCAE2A81BE627A5003",
    "销售实录员": "R65CAAE057914CCCAE2A81BE627A5004",
    "销售办证员": "R65CAAE057914CCCAE2A81BE627A5005",
}

USER_ROLE_ASSIGNMENTS: Dict[str, List[str]] = {
    "管理员": ["吴良萍", "吴九华", "汪璇"],
    "销售主管": ["方金强", "吴九华"],
    "销售计划员": ["程欢欢"],
    "销售实录员": ["郑丽娜"],
    "销售办证员": ["刘丹"],
}


def load_runtime_settings():
    from app.core.config import Settings

    return Settings()


def load_settings_env(env_file: str) -> None:
    if env_file:
        os.environ["ENV_FILE"] = env_file
    from dotenv import load_dotenv

    path = Path(env_file) if env_file else BACKEND_ROOT / ".env.production"
    if not path.is_absolute():
        path = BACKEND_ROOT / path
    if path.is_file():
        load_dotenv(path, override=True)


def collect_subtree_menu_ids(conn, root_ids: List[str], tenant_no: str) -> Set[str]:
    """BFS 收集模块根下全部菜单（含中间目录）。"""
    all_ids: Set[str] = set(root_ids)
    frontier = list(root_ids)
    while frontier:
        placeholders = ", ".join(f":p{i}" for i in range(len(frontier)))
        params = {f"p{i}": frontier[i] for i in range(len(frontier))}
        params["tenant"] = tenant_no
        rows = conn.execute(
            text(
                f"""
                SELECT id FROM sys_menu
                WHERE del_flag = 0 AND tenant_no = :tenant
                  AND parent_id IN ({placeholders})
                """
            ),
            params,
        ).fetchall()
        new_ids = [str(r[0]) for r in rows if str(r[0]) not in all_ids]
        if not new_ids:
            break
        all_ids.update(new_ids)
        frontier = new_ids
    return all_ids


def perm_menu_ids(conn, menu_ids: Set[str], tenant_no: str) -> List[str]:
    """持久化所有带 perms 的菜单（C 页面 :list + F 按钮），与前端 menuCode 过滤对齐。"""
    if not menu_ids:
        return []
    placeholders = ", ".join(f":m{i}" for i in range(len(menu_ids)))
    params = {f"m{i}": mid for i, mid in enumerate(sorted(menu_ids))}
    params["tenant"] = tenant_no
    rows = conn.execute(
        text(
            f"""
            SELECT id FROM sys_menu
            WHERE del_flag = 0 AND tenant_no = :tenant AND id IN ({placeholders})
              AND perms IS NOT NULL AND TRIM(perms) <> ''
            ORDER BY id
            """
        ),
        params,
    ).fetchall()
    return [str(r[0]) for r in rows if r and r[0]]


def leaf_menu_ids(conn, menu_ids: Set[str], tenant_no: str) -> List[str]:
    """与角色管理页一致：仅持久化叶子菜单（无子节点的菜单）。"""
    if not menu_ids:
        return []
    placeholders = ", ".join(f":m{i}" for i in range(len(menu_ids)))
    params = {f"m{i}": mid for i, mid in enumerate(sorted(menu_ids))}
    params["tenant"] = tenant_no
    rows = conn.execute(
        text(
            f"""
            SELECT id, parent_id FROM sys_menu
            WHERE del_flag = 0 AND tenant_no = :tenant AND id IN ({placeholders})
            """
        ),
        params,
    ).mappings().all()
    id_set = {str(r["id"]) for r in rows}
    parents_with_children: Set[str] = set()
    for r in rows:
        pid = str(r.get("parent_id") or "")
        if pid in id_set:
            parents_with_children.add(pid)
    leaves = sorted(mid for mid in id_set if mid not in parents_with_children)
    return leaves


def fetch_user_id_map(conn, tenant_no: str) -> Dict[str, str]:
    rows = conn.execute(
        text(
            """
            SELECT id, user_no, dept_id FROM sys_user
            WHERE tenant_no = :tenant AND del_flag = 0
            """
        ),
        {"tenant": tenant_no},
    ).mappings().all()
    return {str(r["user_no"]): {"id": str(r["id"]), "dept_id": str(r.get("dept_id") or "")} for r in rows}


def purge_non_super_roles(conn, tenant_no: str, dry_run: bool) -> int:
    cnt = conn.execute(
        text(
            """
            SELECT COUNT(*) FROM sys_role
            WHERE tenant_no = :tenant AND del_flag = 0 AND id <> :super_id
            """
        ),
        {"tenant": tenant_no, "super_id": SUPER_ADMIN_ROLE_ID},
    ).scalar()
    if dry_run:
        return int(cnt or 0)
    now = datetime.now()
    conn.execute(
        text(
            """
            UPDATE sys_role_menu
            SET del_flag = 1, update_time = :now
            WHERE tenant_no = :tenant AND role_id <> :super_id AND del_flag = 0
            """
        ),
        {"tenant": tenant_no, "super_id": SUPER_ADMIN_ROLE_ID, "now": now},
    )
    conn.execute(
        text(
            """
            DELETE FROM sys_role_menu
            WHERE tenant_no = :tenant AND role_id <> :super_id
            """
        ),
        {"tenant": tenant_no, "super_id": SUPER_ADMIN_ROLE_ID},
    )
    conn.execute(
        text(
            """
            UPDATE sys_role_data_scope_detail
            SET del_flag = 1, update_time = :now
            WHERE tenant_no = :tenant AND role_id <> :super_id AND del_flag = 0
            """
        ),
        {"tenant": tenant_no, "super_id": SUPER_ADMIN_ROLE_ID, "now": now},
    )
    conn.execute(
        text(
            """
            UPDATE sys_user_role
            SET del_flag = 1, update_time = :now
            WHERE tenant_no = :tenant AND role_id <> :super_id AND del_flag = 0
            """
        ),
        {"tenant": tenant_no, "super_id": SUPER_ADMIN_ROLE_ID, "now": now},
    )
    conn.execute(
        text(
            """
            UPDATE sys_role
            SET del_flag = 1, update_time = :now, update_by = :admin
            WHERE tenant_no = :tenant AND del_flag = 0 AND id <> :super_id
            """
        ),
        {"tenant": tenant_no, "super_id": SUPER_ADMIN_ROLE_ID, "now": now, "admin": ADMIN_USER_ID},
    )
    return int(cnt or 0)


def ensure_new_roles(conn, tenant_no: str, dry_run: bool) -> None:
    order = 10
    for role_name, role_id in NEW_ROLES.items():
        exists = conn.execute(
            text(
                """
                SELECT id FROM sys_role
                WHERE id = :id AND tenant_no = :tenant AND del_flag = 0
                LIMIT 1
                """
            ),
            {"id": role_id, "tenant": tenant_no},
        ).first()
        if exists:
            order += 10
            continue
        if not dry_run:
            conn.execute(
                text(
                    """
                    INSERT INTO sys_role (
                        id, role_name, data_scope, status, tenant_no, order_num, remark,
                        create_by, create_dept_id, create_time, update_by, update_time, del_flag
                    ) VALUES (
                        :id, :name, '1', '1', :tenant, :order_num,
                        '用户迁移后业务角色初始化', :admin, '', :now, :admin, :now, 0
                    )
                    """
                ),
                {
                    "id": role_id,
                    "name": role_name,
                    "tenant": tenant_no,
                    "order_num": order,
                    "admin": ADMIN_USER_ID,
                    "now": datetime.now(),
                },
            )
        order += 10


def grant_module_menus(conn, role_id: str, leaf_ids: List[str], tenant_no: str, dry_run: bool) -> int:
    if dry_run:
        return len(leaf_ids)
    conn.execute(
        text("DELETE FROM sys_role_menu WHERE role_id = :rid"),
        {"rid": role_id},
    )
    now = datetime.now()
    for mid in leaf_ids:
        conn.execute(
            text(
                """
                INSERT INTO sys_role_menu (
                    id, role_id, menu_id, tenant_no, order_num, remark,
                    create_by, create_dept_id, create_time, update_by, update_time, del_flag
                ) VALUES (
                    :id, :role_id, :menu_id, :tenant, 0, 'setup_business_roles',
                    :admin, '', :now, :admin, :now, 0
                )
                """
            ),
            {
                "id": generate_uuid(),
                "role_id": role_id,
                "menu_id": mid,
                "tenant": tenant_no,
                "admin": ADMIN_USER_ID,
                "now": now,
            },
        )
    return len(leaf_ids)


def assign_user_roles(
    conn,
    user_map: Dict[str, dict],
    tenant_no: str,
    dry_run: bool,
) -> List[dict]:
    rows_out: List[dict] = []
    for role_name, user_nos in USER_ROLE_ASSIGNMENTS.items():
        role_id = NEW_ROLES[role_name]
        for user_no in user_nos:
            info = user_map.get(user_no)
            if not info:
                raise ValueError(f"用户不存在: {user_no}")
            exists = conn.execute(
                text(
                    """
                    SELECT id FROM sys_user_role
                    WHERE user_id = :uid AND role_id = :rid AND tenant_no = :tenant AND del_flag = 0
                    LIMIT 1
                    """
                ),
                {"uid": info["id"], "rid": role_id, "tenant": tenant_no},
            ).first()
            if exists:
                rows_out.append({"user_no": user_no, "role_name": role_name, "status": "exists"})
                continue
            if not dry_run:
                conn.execute(
                    text(
                        """
                        INSERT INTO sys_user_role (
                            id, user_id, role_id, tenant_no, order_num, remark,
                            create_by, create_dept_id, create_time, update_by, update_time, del_flag
                        ) VALUES (
                            :id, :uid, :rid, :tenant, 0, 'setup_business_roles',
                            :admin, :dept_id, :now, :admin, :now, 0
                        )
                        """
                    ),
                    {
                        "id": generate_uuid(),
                        "uid": info["id"],
                        "rid": role_id,
                        "tenant": tenant_no,
                        "admin": ADMIN_USER_ID,
                        "dept_id": info["dept_id"],
                        "now": datetime.now(),
                    },
                )
            rows_out.append({"user_no": user_no, "role_name": role_name, "status": "inserted"})
    return rows_out


def main() -> int:
    parser = argparse.ArgumentParser(description="初始化业务角色与菜单权限")
    parser.add_argument("--env-file", default=".env.production")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument(
        "--report-out",
        default=str(BACKEND_ROOT.parent / "docs/data-migration/mappings/role_setup_result.json"),
    )
    args = parser.parse_args()

    load_settings_env(args.env_file)
    settings = load_runtime_settings()
    engine = create_engine(settings.DATABASE_URL, pool_pre_ping=True)

    print("=== setup_business_roles ===")
    print(f"模式: {'DRY-RUN' if args.dry_run else 'EXECUTE'}")
    print(f"租户: {TENANT_NO}")
    print()

    with engine.begin() as conn:
        removed = purge_non_super_roles(conn, TENANT_NO, args.dry_run)
        ensure_new_roles(conn, TENANT_NO, args.dry_run)
        subtree = collect_subtree_menu_ids(conn, MODULE_MENU_ROOTS, TENANT_NO)
        grant_ids = perm_menu_ids(conn, subtree, TENANT_NO)
        menu_grants = {}
        for role_name, role_id in NEW_ROLES.items():
            menu_grants[role_name] = grant_module_menus(conn, role_id, grant_ids, TENANT_NO, args.dry_run)
        user_map = fetch_user_id_map(conn, TENANT_NO)
        assignments = assign_user_roles(conn, user_map, TENANT_NO, args.dry_run)

    report = {
        "generated_at": datetime.now().isoformat(timespec="seconds"),
        "removed_old_roles": removed,
        "module_menu_roots": MODULE_MENU_ROOTS,
        "subtree_menu_count": len(subtree),
        "perm_menu_count": len(grant_ids),
        "menu_grants_per_role": menu_grants,
        "user_role_assignments": assignments,
        "new_roles": NEW_ROLES,
    }

    if not args.dry_run:
        Path(args.report_out).write_text(
            json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
        )

    print(f"清理旧角色（非超级管理员）: {removed}")
    print(f"模块子树菜单数: {len(subtree)}，带 perms 菜单数: {len(grant_ids)}")
    for role_name, cnt in menu_grants.items():
        print(f"  {role_name}: {cnt} 个菜单权限")
    print("用户-角色绑定:")
    for item in assignments:
        print(f"  {item['user_no']} -> {item['role_name']} ({item['status']})")
    if not args.dry_run:
        print(f"报告: {args.report_out}")
    print("完成.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())