181 lines
7.1 KiB
Python
181 lines
7.1 KiB
Python
#!/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)
|
||
|