all-project/教职工数据迁移/insert_external_staff_to_sys_user.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将 temp_external_staff 的外聘职工数据插入到 sys_user，并为每个新用户分配角色

规则：
 - sys_user.user_name 使用 temp_external_staff.`手机号`
 - sys_user.user_no 使用 temp_external_staff.`序号`，格式为 'W' + zero_pad(序号,5)（例如序号=1 -> W00001；序号=209 -> W00209）
 - sys_user.nick_name 使用 temp_external_staff.`姓名`
 - 性别映射：'男' -> '0', '女' -> '1'，其他或空 -> NULL 或默认 '0'（脚本中采用 '0'）
 - 为每个插入的用户在 sys_user_role 中插入两个角色：4 和 1938120323613728770

注意：脚本默认不执行（dry-run），只有当去掉 --dry-run 参数并确认后才会写库。

用法示例（不写库）：
    python insert_external_staff_to_sys_user.py --dry-run

用法示例（写库 — 请先确认并备份）：
    python insert_external_staff_to_sys_user.py
"""

from __future__ import annotations
import os
import sys
import argparse
import importlib.util
import datetime
import re
from typing import Any, Dict, Optional

import pandas as pd
from sqlalchemy import create_engine, text


def load_db_config_from_file(config_path: str) -> dict:
    spec = importlib.util.spec_from_file_location("db_config", config_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)  # type: ignore
    if not hasattr(module, "DB_CONFIG"):
        raise RuntimeError(f"No DB_CONFIG found in {config_path}")
    return getattr(module, "DB_CONFIG")


def build_engine_from_config(db_config: dict):
    user = db_config.get("user")
    password = db_config.get("password")
    host = db_config.get("host", "localhost")
    port = db_config.get("port", 3306)
    database = db_config.get("database")
    charset = db_config.get("charset", "utf8mb4")
    uri = f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}?charset={charset}"
    return create_engine(uri, pool_pre_ping=True)


def normalize_phone(raw: Optional[str]) -> Optional[str]:
    if raw is None:
        return None
    s = str(raw)
    digits = "".join(re.findall(r"\d", s))
    return digits[:11] if digits else None


def map_sex(raw: Optional[str]) -> Optional[str]:
    if raw is None:
        return None
    s = str(raw).strip()
    if s in ("男", "M", "m", "male", "Male"):
        return "0"
    if s in ("女", "F", "f", "female", "Female"):
        return "1"
    return None


def zero_pad_seq_w_prefix(seq_value: Optional[str]) -> str:
    """将序号转换为 'W' + 5位零填充字符串；若无法解析为整数，则尝试取数字子串或返回 'W00000'"""
    if seq_value is None:
        return "W00000"
    s = str(seq_value).strip()
    # 提取整数部分
    m = re.search(r"\d+", s)
    if not m:
        return "W00000"
    num = int(m.group(0))
    return "W" + str(num).zfill(5)


def generate_user_id(conn) -> int:
    r = conn.execute(text("SELECT MAX(user_id) as m FROM sys_user")).fetchone()
    max_id = int(r[0]) if r and r[0] is not None else 0
    return max_id + 1


def user_exists(conn, user_name: str) -> bool:
    r = conn.execute(text("SELECT 1 FROM sys_user WHERE user_name = :u LIMIT 1"), {"u": user_name}).fetchone()
    return bool(r)


def user_has_role(conn, user_id: int, role_id: int) -> bool:
    r = conn.execute(text("SELECT 1 FROM sys_user_role WHERE user_id = :uid AND role_id = :rid LIMIT 1"), {"uid": user_id, "rid": role_id}).fetchone()
    return bool(r)


def insert_user(conn, data: Dict[str, Any]):
    cols = ", ".join(f"`{k}`" for k in data.keys())
    vals = ", ".join(f":{k}" for k in data.keys())
    ins = text(f"INSERT INTO sys_user ({cols}) VALUES ({vals})")
    conn.execute(ins, data)


def insert_user_role(conn, user_id: int, role_id: int):
    ins = text("INSERT INTO sys_user_role (`user_id`, `role_id`) VALUES (:uid, :rid)")
    conn.execute(ins, {"uid": user_id, "rid": role_id})


def main():
    parser = argparse.ArgumentParser(description="将 temp_external_staff 的数据插入 sys_user")
    parser.add_argument("--db-config", type=str, default="../学校历史数据迁移/db_config.py", help="db_config.py 路径")
    parser.add_argument("--table", type=str, default="temp_external_staff", help="源临时表名")
    parser.add_argument("--dry-run", action="store_true", help="仅打印将要执行的操作，不写入数据库")
    args = parser.parse_args()

    base_dir = os.path.dirname(os.path.abspath(__file__))
    db_config_path = args.db_config if os.path.isabs(args.db_config) else os.path.join(base_dir, args.db_config)
    if not os.path.exists(db_config_path):
        # 尝试上级目录中的学校历史数据迁移配置
        alt = os.path.join(os.path.dirname(base_dir), "学校历史数据迁移", "db_config.py")
        if os.path.exists(alt):
            db_config_path = alt
        else:
            print(f"找不到 db_config 文件: {db_config_path}")
            sys.exit(1)

    db_config = load_db_config_from_file(db_config_path)
    engine = build_engine_from_config(db_config)

    # 读取源数据
    with engine.connect() as conn:
        df = pd.read_sql(text(f"SELECT * FROM `{args.table}`"), conn)

    # 识别列名
    possible_seq_cols = ["序号", "序号_1", "col_1"]
    possible_name_cols = ["姓名", "name"]
    possible_phone_cols = ["手机号", "手机号_常用的一个_", "phone"]
    possible_sex_cols = ["性别", "sex"]
    possible_dept_cols = ["部门", "dept", "department"]

    cols = list(df.columns)
    def pick(names):
        for n in names:
            if n in cols:
                return n
        return None

    seq_col = pick(possible_seq_cols) or cols[1] if len(cols) > 1 else None
    name_col = pick(possible_name_cols) or cols[2] if len(cols) > 2 else None
    phone_col = pick(possible_phone_cols) or cols[6] if len(cols) > 6 else None
    sex_col = pick(possible_sex_cols) or cols[3] if len(cols) > 3 else None
    dept_col = pick(possible_dept_cols) or cols[5] if len(cols) > 5 else None

    print(f"使用列映射: 序号={seq_col}, 姓名={name_col}, 性别={sex_col}, 手机号={phone_col}")

    inserted = 0
    skipped_exists = 0
    errors = 0
    role_ids = [4, 1938120323613728770]
    password_hash = "$2a$10$FGldU5BGGq9ONuBmYhxLweFhmbkE45keMKASNZpjETXawE1oGbtZy"
    now_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    # iterate and insert
    with engine.begin() as conn:
        for idx, row in df.iterrows():
            seq_val = row.get(seq_col) if seq_col in row else None
            name_val = row.get(name_col) if name_col in row else None
            raw_phone = row.get(phone_col) if phone_col in row else None
            dept_raw = row.get(dept_col) if dept_col in row else None
            sex_raw = row.get(sex_col) if sex_col in row else None

            phone = normalize_phone(raw_phone)
            if not phone:
                print(f"跳过第{idx+1}行：手机号无效或为空（原值: {raw_phone}）")
                errors += 1
                continue

            user_name_val = phone
            if user_exists(conn, user_name_val):
                skipped_exists += 1
                continue

            user_no_val = zero_pad_seq_w_prefix(seq_val)
            nick = str(name_val).strip() if name_val is not None else user_name_val
            sex_val = map_sex(sex_raw) or "0"

            # resolve dept_id from 部门名称
            dept_id = None
            if dept_raw is not None and str(dept_raw).strip() != "":
                try:
                    r = conn.execute(text("SELECT dept_id FROM sys_dept WHERE dept_name = :name LIMIT 1"), {"name": str(dept_raw).strip()}).fetchone()
                    dept_id = int(r[0]) if r and r[0] is not None else None
                    if dept_id is None:
                        print(f"警告：第{idx+1}行，部门 '{dept_raw}' 在 sys_dept 中未找到，dept_id 将设为 NULL。")
                except Exception as e:
                    print(f"查找部门时出错（行 {idx+1}）: {e}")
                    dept_id = None

            new_user_id = generate_user_id(conn)
            user_record = {
                "user_id": new_user_id,
                "tenant_id": "000000",
                "dept_id": dept_id,
                "user_name": user_name_val,
                "nick_name": nick,
                "user_type": "sys_user",
                "email": None,
                "phonenumber": phone,
                "sex": sex_val,
                "avatar": None,
                "password": password_hash,
                "status": "0",
                "del_flag": "0",
                "login_ip": "",
                "login_date": None,
                "create_dept": 103,
                "create_by": 1,
                "create_time": now_str,
                "update_by": 1,
                "update_time": now_str,
                "remark": None,
                "account_type": "1",
                "user_no": user_no_val,
            }

            print(f"准备插入 user_id={new_user_id}, user_name={user_name_val}, user_no={user_no_val}, nick={nick}")
            if args.dry_run:
                inserted += 1
                continue

            try:
                insert_user(conn, user_record)
                for rid in role_ids:
                    if not user_has_role(conn, new_user_id, rid):
                        insert_user_role(conn, new_user_id, rid)
                inserted += 1
            except Exception as e:
                print(f"插入失败（行 {idx+1}）: {e}")
                errors += 1

    print(f"完成：插入 {inserted} 条，已存在跳过 {skipped_exists} 条，错误 {errors} 条（dry-run={args.dry_run})）")


if __name__ == "__main__":
    main()