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

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从 temp_staff_data 读取数据并更新 sys_user 表

逻辑：
 1. 对于 temp_staff_data 中每一行，使用 nickname = 姓名 AND account_type = '1' 在 sys_user 中查找唯一记录
 2. 若找到唯一记录：
     - 计算 user_no = zero_pad( sys_user.user_name, 6 )，并写回到 sys_user.user_no
     - 将 sys_user.user_name 和 sys_user.phonenumber 都更新为 temp_staff_data 中的手机号字段值
 3. 若未找到或找到多条记录，跳过并记录日志

用法：
    python update_sys_user_from_temp.py --db-config ../学校历史数据迁移/db_config.py --dry-run
"""

from __future__ import annotations
import os
import sys
import argparse
import importlib.util
from typing import Optional, Sequence
import logging

import pandas as pd
from sqlalchemy import create_engine, text

LOG = logging.getLogger("update_sys_user")
LOG.setLevel(logging.INFO)
ch = logging.StreamHandler()
ch.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
LOG.addHandler(ch)


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 pick_column_name(possible_names: Sequence[str], available: Sequence[str]) -> Optional[str]:
    """从可能的列名集合中挑选 temp_staff_data 中存在的一项"""
    for name in possible_names:
        if name in available:
            return name
    return None


def zero_pad_to_six(value) -> str:
    """将给定值转换为字符串并左侧补0到6位（保留原始字符串的数字内容优先）"""
    if value is None:
        return ""
    s = str(value).strip()
    # 尝试提取整数部分
    try:
        # 支持像 '00123' 或 '123' 的情况
        num = int(s)
        return str(num).zfill(6)
    except Exception:
        # 若不能转为整数，直接用原始字符串左补0（或裁切）
        return s.zfill(6)[-6:]


def main(db_config_path: str, table_name: str = "temp_staff_data", dry_run: bool = True):
    base_dir = os.path.dirname(os.path.abspath(__file__))
    if not os.path.isabs(db_config_path):
        # 相对路径从当前脚本目录解析
        db_config_path = os.path.join(base_dir, db_config_path)

    if not os.path.exists(db_config_path):
        LOG.error("找不到 db_config 文件: %s", db_config_path)
        sys.exit(1)

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

    # 使用事务（begin）确保更新会被提交
    with engine.begin() as conn:
        # 读取表头，确定列名映射
        res = conn.execute(text(f"SELECT * FROM `{table_name}` LIMIT 1"))
        cols = list(res.keys())
        LOG.info("temp_staff_data 可用列: %s", cols)

        name_col = pick_column_name(["姓名", "name"], cols)
        phone_col = pick_column_name(["手机号", "手机号_常用的一个_", "phone", "phonenumber"], cols)
        if not name_col or not phone_col:
            LOG.error("无法在 %s 中找到 必需 列 (姓名/手机号)。请检查表结构。", table_name)
            return

        # 读取所有数据（若数据量巨大可以改为分批处理）
        df = pd.read_sql(text(f"SELECT * FROM `{table_name}`"), conn)
        LOG.info("读取到 %d 条 temp_staff_data 记录", len(df))

        updated = 0
        skipped_no_user = 0
        skipped_multi = 0
        errors = 0

        for _, row in df.iterrows():
            staff_name = row.get(name_col)
            phone = row.get(phone_col)
            if pd.isna(staff_name) or staff_name is None:
                LOG.warning("跳过一行，姓名为空")
                skipped_no_user += 1
                continue

            staff_name_str = str(staff_name).strip()
            # 查找 sys_user
            try:
                users = conn.execute(
                    text("SELECT user_id, user_name FROM sys_user WHERE nick_name = :nick AND account_type = '1'"),
                    {"nick": staff_name_str},
                ).fetchall()
            except Exception as e:
                LOG.exception("查询 sys_user 失败: %s", e)
                errors += 1
                continue

            if len(users) == 0:
                LOG.info("未找到 sys_user (nickname=%s, account_type=1)，跳过", staff_name_str)
                skipped_no_user += 1
                continue
            if len(users) > 1:
                LOG.warning("找到多条 sys_user (nickname=%s, account_type=1)，跳过以避免歧义", staff_name_str)
                skipped_multi += 1
                continue

            user_row = users[0]
            user_id = user_row[0]
            current_user_name = user_row[1]
            # 计算 user_no 并更新 user_name & phonenumber
            user_no = zero_pad_to_six(current_user_name)
            new_user_name = "" if phone is None or (isinstance(phone, float) and pd.isna(phone)) else str(phone).strip()
            new_phonenumber = new_user_name

            LOG.info("准备更新 sys_user.id=%s: user_no=%s, user_name=%s, phonenumber=%s",
                     user_id, user_no, new_user_name, new_phonenumber)

            if dry_run:
                LOG.info("dry-run 模式，未执行写入")
                updated += 1
                continue

            try:
                conn.execute(
                    text("UPDATE sys_user SET user_no = :user_no, user_name = :user_name, phonenumber = :phone WHERE user_id = :id"),
                    {"user_no": user_no, "user_name": new_user_name, "phone": new_phonenumber, "id": user_id},
                )
                updated += 1
            except Exception as e:
                LOG.exception("更新 sys_user.id=%s 失败: %s", user_id, e)
                errors += 1

        LOG.info("处理完成，统计: updated=%d, skipped_no_user=%d, skipped_multi=%d, errors=%d",
                 updated, skipped_no_user, skipped_multi, errors)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="从 temp_staff_data 更新 sys_user")
    parser.add_argument("--db-config", type=str, default="../学校历史数据迁移/db_config.py", help="数据库配置文件路径（python模块含 DB_CONFIG）")
    parser.add_argument("--table", type=str, default="temp_staff_data", help="临时表名")
    parser.add_argument("--dry-run", action="store_true", help="仅打印将要执行的操作，不写入数据库")
    args = parser.parse_args()

    main(db_config_path=args.db_config, table_name=args.table, dry_run=args.dry_run)