all-project/学校历史数据迁移/02-update_sys_dept.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
更新 sys_dept 表：
1) 打印 sys_dept 表结构供学习
2) 可删除 sys_dept 中除 dept_id 为 100/101/102 外的所有记录（需确认或使用 --yes）
3) 从 teacher.xlsx 中读取部门列，按规则插入到 sys_dept （部分学院放到 parent_id=101，其余放到 102）

用法：
  python update_sys_dept.py           # 交互式执行（会提示确认删除）
  python update_sys_dept.py --yes     # 自动确认删除
  python update_sys_dept.py --teacher-file /path/to/teacher.xlsx

注意：
  - 本脚本会尝试读取 `学校历史数据迁移/db_config.py` 中的 DB_CONFIG 连接数据库
  - 插入时会尽量推断并填充常用字段（ancestors, parent_id, dept_name, order_num, status, del_flag, create_time）
"""
from __future__ import annotations
import os
import sys
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 inspect_table(engine, table_name: str) -> List[Dict[str, Any]]:
    q = text(
        "SELECT COLUMN_NAME, IS_NULLABLE, COLUMN_DEFAULT, DATA_TYPE "
        "FROM information_schema.columns "
        "WHERE table_schema = :schema AND table_name = :table ORDER BY ORDINAL_POSITION"
    )
    schema = engine.url.database
    with engine.connect() as conn:
        result = conn.execute(q, {"schema": schema, "table": table_name})
        # Use mappings() so each row is a mapping of column_name->value
        rows = result.mappings().all()
    return [dict(r) for r in rows]


def find_dept_column(df: pd.DataFrame) -> Optional[str]:
    candidates = []
    for col in df.columns:
        name = str(col).strip()
        low = name.lower()
        if any(k in low for k in ("部门", "学院", "系", "组织", "院", "department", "dept", "office", "unit")):
            candidates.append(name)
    if candidates:
        return candidates[0]
    # fallback: ask user by listing columns
    return None


def parent_has_children(engine, parent_id: int) -> bool:
    """
    检查 parent_id 下是否已有子部门记录（用于决定是否跳过该 parent 的批量插入）
    """
    q = text("SELECT 1 FROM sys_dept WHERE parent_id = :pid LIMIT 1")
    with engine.connect() as conn:
        r = conn.execute(q, {"pid": parent_id}).fetchone()
    return bool(r)


def read_teacher_departments_from_db(engine, table_name: str = "his_lssj") -> List[str]:
    """
    从数据库表（默认 his_lssj）读取去重的部门名称（列名 '部门'）。
    返回去重、排序后的部门名称列表。
    """
    print(f"从数据库表 `{table_name}` 读取部门列（列名假定为 '部门'）")
    q = text(f"SELECT DISTINCT `部门` AS dept_name FROM `{table_name}` WHERE `部门` IS NOT NULL AND TRIM(`部门`) != ''")
    with engine.connect() as conn:
        rows = conn.execute(q).mappings().all()
    depts = sorted({str(r.get("dept_name")).strip() for r in rows if r.get("dept_name") is not None})
    print(f"发现 {len(depts)} 个去重部门名（示例前20）:")
    for item in depts[:20]:
        print("  -", item)
    return depts


def get_parent_info(engine, parent_ids: List[int]) -> Dict[int, Dict[str, Any]]:
    q = text("SELECT * FROM sys_dept WHERE dept_id IN :ids")
    with engine.connect() as conn:
        result = conn.execute(q, {"ids": tuple(parent_ids)})
        rows = result.mappings().all()
    out: Dict[int, Dict[str, Any]] = {}
    for r in rows:
        d = dict(r)
        out[int(d.get("dept_id"))] = d
    return out


def delete_except(engine, keep_ids: List[int], yes: bool):
    if not yes:
        resp = input(f"将删除 sys_dept 表中除 dept_id IN {keep_ids} 外的所有记录，确认继续？(yes/no): ").strip().lower()
        if resp not in ("y", "yes"):
            print("取消操作。")
            return False
    q = text("DELETE FROM sys_dept WHERE dept_id NOT IN :keep")
    with engine.begin() as conn:
        conn.execute(q, {"keep": tuple(keep_ids)})
    print("删除完成。")
    return True


def insert_departments(engine, dept_names: List[str], parent_mapping: Dict[str, int], parent_info: Dict[int, Dict[str, Any]]):
    """
    插入时遵循以下规则以保持与现有数据一致性：
    - dept_id: 若表中已有较大数值 id（如 2025xxxxxxxxxxxx），则以 max(dept_id)+1 继续；否则以当前时间 `%Y%m%d%H%M%S` + 4位序号生成
    - tenant_id: 写为 '000000'
    - order_num: 在同一 parent_id 下取 max(order_num)+1
    - ancestors: 从 parent_info 推断，若父项无 ancestors 则使用 "0,{parent_id}"
    - 其他字段填充 create_by/update_by/create_dept 等默认值以匹配样例
    """
    inserted = 0
    skipped = 0
    now_dt = datetime.datetime.now()
    now_str = now_dt.strftime("%Y-%m-%d %H:%M:%S")
    timestamp_prefix = now_dt.strftime("%Y%m%d%H%M%S")

    # helper to generate unique dept_id
    def generate_dept_id(conn, local_counter_start=1):
        # get current max dept_id
        r = conn.execute(text("SELECT MAX(dept_id) as m FROM sys_dept")).fetchone()
        max_id = int(r[0]) if r and r[0] is not None else 0
        if max_id > 10**14:
            return max_id + 1
        # else use timestamp-based id plus 4-digit counter
        counter = local_counter_start
        while True:
            candidate = int(f"{timestamp_prefix}{counter:04d}")
            exists = conn.execute(text("SELECT 1 FROM sys_dept WHERE dept_id = :did LIMIT 1"), {"did": candidate}).fetchone()
            if not exists:
                return candidate
            counter += 1

    # Cache per-parent decision to avoid repeated checks: 'skip' or 'insert'
    parent_decision: Dict[int, str] = {}
    with engine.begin() as conn:
        # fetch column info once
        col_info = inspect_table(engine, "sys_dept")
        col_names = [c["COLUMN_NAME"] for c in col_info]
        # prepare xy counter for parent_id = 101 based on existing dept_category values
        xy_counter = 1
        try:
            rows = conn.execute(text("SELECT dept_category FROM sys_dept WHERE parent_id = :pid AND dept_category LIKE 'xy-%'"), {"pid": 101}).mappings().all()
            nums = []
            for r in rows:
                val = r.get("dept_category")
                if not val:
                    continue
                parts = str(val).split("-")
                if len(parts) >= 2:
                    try:
                        nums.append(int(parts[-1]))
                    except Exception:
                        continue
            if nums:
                xy_counter = max(nums) + 1
        except Exception:
            # ignore and keep default
            xy_counter = 1

        for name in dept_names:
            pname = name.strip()
            if not pname:
                continue
            parent_id = parent_mapping.get(pname, None)
            if parent_id is None:
                parent_id = 102
            # If we've decided to skip inserting for this parent, continue
            if parent_id in parent_decision and parent_decision[parent_id] == "skip":
                skipped += 1
                continue

            # First time encountering this parent: if it already has children, skip inserting any under it
            if parent_id not in parent_decision:
                try:
                    has_children = conn.execute(text("SELECT 1 FROM sys_dept WHERE parent_id = :pid LIMIT 1"), {"pid": parent_id}).fetchone()
                    if has_children:
                        parent_decision[parent_id] = "skip"
                        logger = None  # no logger in this module-level function; just mark skip
                        skipped += 1
                        continue
                    else:
                        parent_decision[parent_id] = "insert"
                except Exception:
                    parent_decision[parent_id] = "insert"
            # skip if same dept_name and parent exists
            exists_q = text("SELECT count(1) as c FROM sys_dept WHERE dept_name = :name AND parent_id = :pid")
            res = conn.execute(exists_q, {"name": pname, "pid": parent_id}).fetchone()
            if res and int(res[0]) > 0:
                skipped += 1
                continue

            parent = parent_info.get(parent_id, {})
            parent_ancestors = parent.get("ancestors") if parent else None
            if parent_ancestors:
                ancestors = f"{parent_ancestors},{parent_id}"
            else:
                # ensure root '0' present
                if parent_id == 0:
                    ancestors = "0"
                else:
                    ancestors = f"0,{parent_id}"

            # compute next order_num under this parent
            r = conn.execute(text("SELECT COALESCE(MAX(order_num),0) as m FROM sys_dept WHERE parent_id = :pid"), {"pid": parent_id}).fetchone()
            next_order = int(r[0]) + 1 if r and r[0] is not None else 1

            # generate unique dept_id
            new_dept_id = generate_dept_id(conn)

            # prepare insert dict with default fields matching provided samples
            # set dept_category to 'xy-序号' for entries under parent_id == 101
            dept_category_val = None
            if parent_id == 101:
                dept_category_val = f"xy-{xy_counter}"
                xy_counter += 1

            insert_dict: Dict[str, Any] = {
                "dept_id": new_dept_id,
                "tenant_id": "000000",
                "parent_id": parent_id,
                "ancestors": ancestors,
                "dept_name": pname,
                "dept_category": dept_category_val,
                "order_num": next_order,
                "leader": None,
                "phone": None,
                "email": None,
                "status": "0",
                "del_flag": "0",
                "create_dept": 0,
                "create_by": 1,
                "create_time": now_str,
                "update_by": 1,
                "update_time": now_str,
            }

            # only keep keys that exist in the actual table
            use_cols = {k: v for k, v in insert_dict.items() if k in col_names}
            if "dept_name" not in use_cols:
                raise RuntimeError("目标表不包含 dept_name 列，无法继续")

            cols_part = ", ".join(f"`{k}`" for k in use_cols.keys())
            vals_part = ", ".join(f":{k}" for k in use_cols.keys())
            ins_q = text(f"INSERT INTO sys_dept ({cols_part}) VALUES ({vals_part})")
            conn.execute(ins_q, use_cols)
            inserted += 1

    print(f"插入完成：新增 {inserted} 条，跳过已存在 {skipped} 条。")


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)

    # 1. 学习表结构
    print("正在读取 sys_dept 表结构...")
    cols = inspect_table(engine, "sys_dept")
    for c in cols:
        print(f"  - {c['COLUMN_NAME']}  type={c['DATA_TYPE']}  nullable={c['IS_NULLABLE']}  default={c['COLUMN_DEFAULT']}")

    # 2. 删除除指定 dept_id 外的记录
    parser = argparse.ArgumentParser()
    parser.add_argument("--yes", action="store_true", help="自动确认删除")
    parser.add_argument("--teacher-file", type=str, default=os.path.join(base_dir, "teacher.xlsx"))
    args = parser.parse_args()

    keep_ids = [100, 101, 102]
    ok = delete_except(engine, keep_ids, args.yes)
    if not ok:
        print("用户取消，退出。")
        return

    # 3. 从 his_lssj 表导入部门（更快）
    dept_names = read_teacher_departments_from_db(engine, table_name="his_lssj")

    # mapping: specific colleges -> parent 101, others -> 102
    top_level_101 = {
        "人文学院", "餐饮休闲学院", "生态园林学院", "数字信息学院",
        "医药健康学院", "智能控制学院", "智能制造学院"
    }
    # build mapping for exact names
    parent_mapping = {name: 101 for name in top_level_101}

    # get parent info
    parent_info = get_parent_info(engine, [101, 102])
    if 101 not in parent_info or 102 not in parent_info:
        print("警告：未能在 sys_dept 中找到 parent_id 为 101 或 102 的记录，请确认这两个父组织是否存在。")
    insert_departments(engine, dept_names, parent_mapping, parent_info)


if __name__ == "__main__":
    main()