zkxw/backend/scripts/migrate_user_role_to_post.py

#!/usr/bin/env python3
"""将「用户-角色」迁移为「用户-职位 + 职位-角色」，并清理直接绑定。

步骤（与产品约定一致）：
1. 软删除清空 sys_post / sys_user_post / sys_post_role（指定租户）
2. 按 sys_role（未删除）逐条创建职位：post_name=角色名，post_code=名称拼音首字母缩写；
   并为职位绑定该角色（sys_post_role 一对一）
3. 按 sys_user_role 为用户绑定对应职位（sys_user_post）
4. 软删除清空 sys_user_role（指定租户）

用法（在 backend 目录）：

  # 务必指向测试库，勿连生产
  export ENV_FILE=./.env.test   # 或确认 .env 已是测试库

  source venv/bin/activate
  pip install pypinyin          # 若尚未安装

  # 只预览，不写库
  python scripts/migrate_user_role_to_post.py --dry-run

  # 真正执行（需同时传 --execute 与库名确认）
  python scripts/migrate_user_role_to_post.py --execute --confirm-database zkxw-sc

注意：
- 菜单/数据权限运行时已走有效角色并集，迁移后可由职位继承。
- 仍有少量代码只查 sys_user_role（如「实验专员」步骤负责人校验、
  质检金额授权列表、工作流按角色选人）。若业务依赖这些入口，需另改代码或
  保留对应直接绑定（可用 --preserve-role-ids）。
"""
from __future__ import annotations

import argparse
import os
import re
import sys
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple
from urllib.parse import urlparse

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

try:
    from pypinyin import Style, lazy_pinyin
except ImportError as exc:  # pragma: no cover
    raise SystemExit(
        "缺少依赖 pypinyin。请先执行: pip install pypinyin\n" + str(exc)
    ) from exc

from sqlalchemy import and_

from app.core.config import settings
from app.core.database import SessionLocal
from app.core.utils import generate_uuid
from app.models.models import SysPost, SysPostRole, SysRole, SysUserPost, SysUserRole

# 质检物料金额查看：授权管理模块仍读写 sys_user_role；默认迁移后一并清理，
# 若需保留其直接绑定，执行时加：
#   --preserve-role-ids ROLE000000000000000000000000000011
DEFAULT_PRESERVE_ROLE_IDS: Tuple[str, ...] = ()


def _mask_database_url(url: str) -> str:
    try:
        u = urlparse(url)
        user = u.username or ""
        host = u.hostname or ""
        port = f":{u.port}" if u.port else ""
        db = (u.path or "").lstrip("/")
        return f"{u.scheme}://{user}:***@{host}{port}/{db}"
    except Exception:
        return "***"


def _database_name(url: str) -> str:
    try:
        return (urlparse(url).path or "").lstrip("/").split("?")[0]
    except Exception:
        return ""


def role_name_to_post_code(role_name: str, used: Set[str]) -> str:
    """职位编码 = 角色名拼音首字母缩写（大写），冲突则追加序号。"""
    name = (role_name or "").strip() or "POST"
    letters = "".join(lazy_pinyin(name, style=Style.FIRST_LETTER))
    base = re.sub(r"[^A-Za-z0-9]", "", letters).upper()
    if not base:
        base = "POST"
    # 编码不宜过长
    base = base[:48]
    code = base
    n = 2
    while code in used:
        suffix = f"-{n}"
        code = (base[: 48 - len(suffix)] + suffix).upper()
        n += 1
    used.add(code)
    return code


def soft_clear_posts(db, tenant_no: str, now: datetime, dry_run: bool) -> Dict[str, int]:
    counts = {}
    for model, key in (
        (SysPostRole, "sys_post_role"),
        (SysUserPost, "sys_user_post"),
        (SysPost, "sys_post"),
    ):
        q = db.query(model).filter(model.tenant_no == tenant_no, model.del_flag == 0)
        n = q.count()
        counts[key] = n
        if not dry_run and n:
            q.update(
                {"del_flag": 1, "update_time": now},
                synchronize_session=False,
            )
    return counts


def load_roles(db, tenant_no: str) -> List[SysRole]:
    return (
        db.query(SysRole)
        .filter(SysRole.tenant_no == tenant_no, SysRole.del_flag == 0)
        .order_by(SysRole.order_num.asc(), SysRole.create_time.asc())
        .all()
    )


def create_posts_for_roles(
    db,
    roles: List[SysRole],
    tenant_no: str,
    operator_id: Optional[str],
    now: datetime,
    dry_run: bool,
) -> Dict[str, str]:
    """返回 role_id -> post_id。"""
    used_codes: Set[str] = set()
    role_to_post: Dict[str, str] = {}
    for idx, role in enumerate(roles, start=1):
        post_id = generate_uuid()
        code = role_name_to_post_code(role.role_name, used_codes)
        status = (role.status or "1").strip() or "1"
        print(
            f"  [{idx}/{len(roles)}] 角色「{role.role_name}」({role.id}) "
            f"→ 职位 {code} / {role.role_name}"
        )
        role_to_post[role.id] = post_id
        if dry_run:
            continue
        db.add(
            SysPost(
                id=post_id,
                post_name=role.role_name,
                post_code=code,
                status=status,
                tenant_no=tenant_no,
                order_num=role.order_num if role.order_num is not None else idx,
                remark=f"由角色迁移生成，role_id={role.id}",
                create_by=operator_id,
                create_dept_id="",
                create_time=now,
                update_by=operator_id,
                update_time=now,
                del_flag=0,
            )
        )
        db.add(
            SysPostRole(
                id=generate_uuid(),
                post_id=post_id,
                role_id=role.id,
                tenant_no=tenant_no,
                order_num=0,
                remark="角色一对一迁移",
                create_by=operator_id,
                create_dept_id="",
                create_time=now,
                update_by=operator_id,
                update_time=now,
                del_flag=0,
            )
        )
    return role_to_post


def migrate_user_posts(
    db,
    tenant_no: str,
    role_to_post: Dict[str, str],
    operator_id: Optional[str],
    now: datetime,
    dry_run: bool,
) -> Tuple[int, int]:
    """按用户-角色生成用户-职位。返回 (将写入条数, 跳过条数)。"""
    rows = (
        db.query(SysUserRole)
        .filter(SysUserRole.tenant_no == tenant_no, SysUserRole.del_flag == 0)
        .all()
    )
    # 同一用户同一职位只绑一次
    seen: Set[Tuple[str, str]] = set()
    written = 0
    skipped = 0
    for ur in rows:
        post_id = role_to_post.get(ur.role_id)
        if not post_id:
            skipped += 1
            continue
        key = (ur.user_id, post_id)
        if key in seen:
            skipped += 1
            continue
        seen.add(key)
        written += 1
        if dry_run:
            continue
        db.add(
            SysUserPost(
                id=generate_uuid(),
                user_id=ur.user_id,
                post_id=post_id,
                tenant_no=tenant_no,
                order_num=ur.order_num if ur.order_num is not None else 0,
                remark=f"由用户角色迁移，role_id={ur.role_id}",
                create_by=operator_id,
                create_dept_id="",
                create_time=now,
                update_by=operator_id,
                update_time=now,
                del_flag=0,
            )
        )
    return written, skipped


def clear_user_roles(
    db,
    tenant_no: str,
    preserve_role_ids: Set[str],
    now: datetime,
    dry_run: bool,
) -> Tuple[int, int]:
    q = db.query(SysUserRole).filter(
        SysUserRole.tenant_no == tenant_no, SysUserRole.del_flag == 0
    )
    if preserve_role_ids:
        q = q.filter(~SysUserRole.role_id.in_(list(preserve_role_ids)))
    n = q.count()
    preserved = 0
    if preserve_role_ids:
        preserved = (
            db.query(SysUserRole)
            .filter(
                SysUserRole.tenant_no == tenant_no,
                SysUserRole.del_flag == 0,
                SysUserRole.role_id.in_(list(preserve_role_ids)),
            )
            .count()
        )
    if not dry_run and n:
        q.update({"del_flag": 1, "update_time": now}, synchronize_session=False)
    return n, preserved


def run(
    *,
    tenant_no: str,
    dry_run: bool,
    execute: bool,
    confirm_database: Optional[str],
    operator_id: Optional[str],
    preserve_role_ids: Set[str],
) -> None:
    db_url = settings.DATABASE_URL
    db_name = _database_name(db_url)
    print("=== 用户角色 → 职位 迁移 ===")
    print(f"ENV_FILE     : {os.getenv('ENV_FILE', '(default backend/.env)')}")
    print(f"DATABASE_URL : {_mask_database_url(db_url)}")
    print(f"database     : {db_name}")
    print(f"tenant_no    : {tenant_no}")
    print(f"mode         : {'DRY-RUN' if dry_run or not execute else 'EXECUTE'}")
    if preserve_role_ids:
        print(f"preserve UR  : {', '.join(sorted(preserve_role_ids))}")
    print()

    if execute and not dry_run:
        if not confirm_database:
            raise SystemExit(
                "执行写库必须传 --confirm-database <库名>，"
                f"当前检测到库名为「{db_name}」。"
            )
        if confirm_database != db_name:
            raise SystemExit(
                f"--confirm-database={confirm_database!r} 与当前连接库名 "
                f"{db_name!r} 不一致，已中止。"
            )

    db = SessionLocal()
    now = datetime.now()
    try:
        roles = load_roles(db, tenant_no)
        user_role_count = (
            db.query(SysUserRole)
            .filter(SysUserRole.tenant_no == tenant_no, SysUserRole.del_flag == 0)
            .count()
        )
        print(f"当前有效角色数     : {len(roles)}")
        print(f"当前用户-角色绑定数: {user_role_count}")
        print()

        print("步骤1：清空职位相关表（软删除）…")
        cleared = soft_clear_posts(db, tenant_no, now, dry_run=dry_run or not execute)
        for k, v in cleared.items():
            print(f"  {k}: {v}")

        print("步骤2：按角色创建职位并绑定角色…")
        role_to_post = create_posts_for_roles(
            db,
            roles,
            tenant_no,
            operator_id,
            now,
            dry_run=dry_run or not execute,
        )

        print("步骤3：按用户-角色生成用户-职位…")
        written, skipped = migrate_user_posts(
            db,
            tenant_no,
            role_to_post,
            operator_id,
            now,
            dry_run=dry_run or not execute,
        )
        print(f"  将写入用户-职位: {written}，跳过: {skipped}")

        print("步骤4：清理用户-角色直接绑定（软删除）…")
        cleared_ur, preserved_ur = clear_user_roles(
            db,
            tenant_no,
            preserve_role_ids,
            now,
            dry_run=dry_run or not execute,
        )
        print(f"  将清理用户-角色: {cleared_ur}，保留: {preserved_ur}")

        if dry_run or not execute:
            db.rollback()
            print()
            print("DRY-RUN 完成，未写入数据库。确认无误后执行：")
            print(
                f"  ENV_FILE=... python scripts/migrate_user_role_to_post.py "
                f"--execute --confirm-database {db_name}"
            )
        else:
            db.commit()
            print()
            print("迁移已提交。")
            print(
                "请在测试环境登录验证：菜单权限、数据权限、职位管理、用户查看中的有效角色。"
            )
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="将用户-角色迁移为用户-职位（测试库专用）"
    )
    parser.add_argument("--tenant-no", default="zkxw", help="租户编号，默认 zkxw")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="仅预览（默认若不带 --execute 也等同预览）",
    )
    parser.add_argument(
        "--execute",
        action="store_true",
        help="真正写库（必须同时传 --confirm-database）",
    )
    parser.add_argument(
        "--confirm-database",
        default=None,
        help="当前连接的数据库名（与 DATABASE_URL 路径段一致）",
    )
    parser.add_argument(
        "--operator-id",
        default=None,
        help="create_by / update_by，可选",
    )
    parser.add_argument(
        "--preserve-role-ids",
        default="",
        help="清理 sys_user_role 时保留的角色 ID，逗号分隔",
    )
    args = parser.parse_args()

    preserve = set(DEFAULT_PRESERVE_ROLE_IDS)
    for part in (args.preserve_role_ids or "").split(","):
        pid = part.strip()
        if pid:
            preserve.add(pid)

    # 未显式 --execute 时一律 dry-run
    dry_run = bool(args.dry_run) or not bool(args.execute)
    run(
        tenant_no=args.tenant_no,
        dry_run=dry_run,
        execute=bool(args.execute),
        confirm_database=args.confirm_database,
        operator_id=args.operator_id,
        preserve_role_ids=preserve,
    )


if __name__ == "__main__":
    main()