all-project/学校历史数据迁移/03-insert_sys_user.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
03-插入老师数据到 sys_user 表

功能：
- 保留 user_id = 1 的用户，其他用户的 del_flag 都改成 '1'（虚拟删除，可交互确认）
- 从 teacher.xlsx 读取教师信息，插入到 sys_user
  - user_name: 工号
  - nick_name: 姓名
  - phonenumber: 手机号（取前11位数字）
  - sex: 男->'0', 女->'1'（其他为空）
  - password: 使用固定哈希值
  - account_type: '1'
  - dept_id: 通过部门名称关联 sys_dept.dept_id（找不到会写 NULL 并打印警告）
  - 跳过已存在相同 user_name 的记录（可选择覆盖）

用法：
  python 03-insert_sys_user.py           # 交互式确认虚拟删除
  python 03-insert_sys_user.py --yes     # 自动确认虚拟删除
  python 03-insert_sys_user.py --teacher-file /path/to/teacher.xlsx
  python 03-insert_sys_user.py --dry-run  # 仅打印将执行的操作，不写库
"""
from __future__ import annotations
import os
import sys
import re
import argparse
import datetime
from typing import List, Dict, Any, Optional
import pandas as pd
from sqlalchemy import create_engine, text
import importlib.util


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}"
    engine = create_engine(uri, pool_pre_ping=True)
    return engine


def find_column(df: pd.DataFrame, keywords: List[str]) -> Optional[str]:
    for col in df.columns:
        name = str(col).strip()
        low = name.lower()
        for k in keywords:
            if k in low:
                return name
    return None


def read_teacher_df_from_db(engine, table_name: str = "his_lssj") -> pd.DataFrame:
    """
    从数据库表 his_lssj 读取教师数据，返回 pandas DataFrame（列名使用表中的列名）。
    假定表中包含列：`工号`, `姓名`, `部门`, `性别`, `手机号码`
    """
    print(f"从数据库表 `{table_name}` 读取教师数据（假定列：工号, 姓名, 部门, 性别, 手机号码）")
    q = f"SELECT `工号`, `姓名`, `部门`, `性别`, `手机号码` FROM `{table_name}`"
    df = pd.read_sql_query(q, con=engine)
    return df


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 soft_delete_other_users(engine, yes: bool, dry_run: bool = False):
    # Renamed behavior: perform physical DELETE except user_id = 1
    if not yes:
        resp = input("将删除 sys_user 中除 user_id=1 外的所有用户，确认？(yes/no) [no]: ").strip().lower()
        if resp not in ("y", "yes"):
            print("取消删除步骤。")
            return
    q_count = text("SELECT COUNT(1) as c FROM sys_user WHERE user_id != 1")
    with engine.connect() as conn:
        cnt = conn.execute(q_count).scalar()
    print(f"将被删除的用户数：{cnt}")
    if dry_run:
        print("dry-run 模式，跳过删除。")
        return
    q = text("DELETE FROM sys_user WHERE user_id != 1")
    with engine.begin() as conn:
        conn.execute(q)
    print("已删除其他用户（保留 user_id = 1）。")


def generate_user_id(conn):
    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 get_dept_id(conn, dept_name: str) -> Optional[int]:
    if dept_name is None or str(dept_name).strip() == "":
        return None
    r = conn.execute(text("SELECT dept_id FROM sys_dept WHERE dept_name = :name LIMIT 1"), {"name": dept_name.strip()}).fetchone()
    return int(r[0]) if r and r[0] is not None else None


def user_exists(conn, user_name: str, dept_id: Optional[int]) -> bool:
    # skip only if both user_name and dept_id match an existing record
    q = text(
        "SELECT 1 FROM sys_user WHERE user_name = :u AND "
        "((dept_id = :did) OR (dept_id IS NULL AND :did IS NULL)) LIMIT 1"
    )
    r = conn.execute(q, {"u": user_name, "did": dept_id}).fetchone()
    return bool(r)


def dept_has_users(conn, dept_id: int) -> bool:
    """检查指定 dept_id 下是否已有用户存在（用于决定是否跳过该部门的批量插入）"""
    if dept_id is None:
        return False
    r = conn.execute(text("SELECT 1 FROM sys_user WHERE dept_id = :did LIMIT 1"), {"did": dept_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 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_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():
    base_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = os.path.join(base_dir, "db_config.py")
    if not os.path.exists(config_path):
        print(f"找不到配置文件: {config_path}")
        sys.exit(1)
    db_config = load_db_config_from_file(config_path)
    engine = build_engine_from_config(db_config)

    parser = argparse.ArgumentParser()
    parser.add_argument("--yes", action="store_true", help="自动确认虚拟删除现有用户（除 user_id=1）")
    parser.add_argument("--teacher-file", type=str, default=os.path.join(base_dir, "teacher.xlsx"))
    parser.add_argument("--dry-run", action="store_true", help="只打印将执行的操作，不写入数据库")
    args = parser.parse_args()

    # 1) 虚拟删除其他用户
    soft_delete_other_users(engine, args.yes, dry_run=args.dry_run)

    # 2) 从 his_lssj 表读取教师数据（更快）
    df = read_teacher_df_from_db(engine, table_name="his_lssj")
    # 直接使用约定列名（表中创建时已用这些列名）
    col_workno = "工号"
    col_name = "姓名"
    col_phone = "手机号码"
    col_dept = "部门"
    col_sex = "性别"
    # ensure columns exist
    for c in (col_workno, col_name):
        if c not in df.columns:
            print(f"必需列 {c} 在表 his_lssj 中未找到，请检查表结构或使用 Excel 源。")
            return

    password_hash = "$2a$10$FGldU5BGGq9ONuBmYhxLweFhmbkE45keMKASNZpjETXawE1oGbtZy"
    now_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    # iterate rows and insert
    inserted = 0
    skipped = 0
    warnings = 0
    # cache per-dept decision to avoid repeated checks: dept_id -> 'skip'|'insert'
    dept_decision: Dict[Optional[int], str] = {}
    with engine.begin() as conn:
        for idx, row in df.iterrows():
            user_name_val = str(row.get(col_workno)).strip() if row.get(col_workno) is not None else ""
            if not user_name_val:
                skipped += 1
                continue

            # determine department first so existence check can use dept_id
            dept_raw = str(row.get(col_dept)).strip() if col_dept and row.get(col_dept) is not None else ""
            dept_id = get_dept_id(conn, dept_raw) if dept_raw else None
            if dept_raw and not dept_id:
                print(f"警告: 第 {idx+1} 行，部门 '{dept_raw}' 在 sys_dept 中未找到，dept_id 将设为 NULL。")
                warnings += 1
            # If we've decided to skip this department because it already has users, skip adding users under it
            if dept_id in dept_decision and dept_decision[dept_id] == "skip":
                skipped += 1
                continue

            # First time encountering this dept: if dept already has users, mark skip
            if dept_id not in dept_decision:
                try:
                    if dept_has_users(conn, dept_id):
                        dept_decision[dept_id] = "skip"
                        skipped += 1
                        continue
                    else:
                        dept_decision[dept_id] = "insert"
                except Exception:
                    dept_decision[dept_id] = "insert"

            if user_exists(conn, user_name_val, dept_id):
                skipped += 1
                continue

            nick = str(row.get(col_name)).strip() if row.get(col_name) is not None else ""
            raw_phone = row.get(col_phone) if col_phone else None
            phone = normalize_phone(raw_phone)
            sex_val = map_sex(row.get(col_sex)) if col_sex else None

            # generate user_id
            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 or user_name_val,
                "user_type": "sys_user",
                "email": None,
                "phonenumber": phone,
                "sex": sex_val or "0",
                "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",
            }

            if args.dry_run:
                print("DRY:", user_record)
                inserted += 1
                continue

            insert_user(conn, user_record)
            # assign roles to the new user
            role_ids = [1938120323613728770, 4]
            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

    print(f"完成：插入 {inserted} 条，跳过 {skipped} 条，警告 {warnings} 条。")


if __name__ == "__main__":
    main()