all-project/学校历史数据迁移/04-01-dedupe_sys_dept.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
04-01-dedupe_sys_dept.py

目的:
  在 sys_dept 表中，对于同一学院(parent_id)下的相同专业名称(dept_name)只保留一条记录，删除重复记录。

行为:
  - 默认 dry-run：打印将要执行的 SQL 操作，不做修改。
  - 传入 --apply 会执行修改（在事务中）。
  - 对于重复记录，会保留 dept_id 最小的那条作为主记录（keep_id），
    并将重复记录的子部门(parent_id 指向被删除 id) 统一改为 keep_id，
    同时将 sys_user 表中引用的 dept_id 更新为 keep_id。
  - 会打印处理摘要。

使用:
  python3 04-01-dedupe_sys_dept.py --host ... --user ... --password ... --database ... [--apply]

注意:
  在运行前请先备份数据库。对非数字或特殊 dept_name 不做额外过滤，脚本基于 parent_id + dept_name 进行分组。
"""

from __future__ import annotations

import argparse
import getpass
import logging
import sys
import os
import importlib.util
from typing import List, Tuple

import pymysql
from pymysql.connections import Connection

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("dedupe_sys_dept")


def parse_args():
    """
    Parse arguments. If a local db_config.py exists in the same directory and defines DB_CONFIG dict,
    its values are used as defaults (same behavior as 04-insert_majors.py).
    """
    # try load local db_config.py for defaults
    local_db = load_local_db_config()
    host_default = local_db.get("host", os.getenv("DB_HOST", "localhost"))
    port_default = int(local_db.get("port", os.getenv("DB_PORT", 3306)))
    user_default = local_db.get("user", os.getenv("DB_USER", "root"))
    password_default = local_db.get("password", os.getenv("DB_PASSWORD", ""))
    database_default = local_db.get("database", os.getenv("DB_NAME", "hzdyjs_xg"))

    p = argparse.ArgumentParser(description="Deduplicate sys_dept entries by (parent_id, dept_name)")
    p.add_argument("--host", default=host_default, help="DB host")
    p.add_argument("--port", type=int, default=port_default, help="DB port")
    p.add_argument("--user", default=user_default, help="DB user")
    p.add_argument("--password", default=password_default, help="DB password (if omitted, prompted)")
    p.add_argument("--database", default=database_default, help="DB name")
    p.add_argument("--charset", default="utf8mb4", help="DB charset")
    p.add_argument("--apply", action="store_true", help="Actually apply deletions and updates. Default is dry-run.")
    p.add_argument("--limit", type=int, default=0, help="If >0, limit number of groups processed (for testing)")
    return p.parse_args()


def load_local_db_config() -> dict:
    """
    Attempt to load DB_CONFIG dict from db_config.py in the current script directory.
    Returns empty dict if not present or on error.
    """
    cfg_path = os.path.join(os.path.dirname(__file__), "db_config.py")
    if not os.path.exists(cfg_path):
        return {}
    try:
        spec = importlib.util.spec_from_file_location("local_db_config", cfg_path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)  # type: ignore
        return getattr(module, "DB_CONFIG", {}) or {}
    except Exception as e:
        logger.warning("加载本地 db_config.py 失败: %s", e)
        return {}


def get_db_conn(args) -> Connection:
    """
    Return a pymysql connection. Naming aligned with 04-insert_majors.py (`get_db_conn`).
    """
    pwd = args.password or getpass.getpass("DB password: ")
    return pymysql.connect(
        host=args.host,
        port=args.port,
        user=args.user,
        password=pwd,
        database=args.database,
        charset=args.charset,
        cursorclass=pymysql.cursors.DictCursor,
        autocommit=False,
    )


def find_duplicate_groups(conn: Connection, limit: int = 0) -> List[Tuple[int, str, List[int]]]:
    """
    返回列表，每项为 (parent_id, dept_name, [dept_id1, dept_id2, ...])，且 dept_id 列表长度 >=2
    """
    sql = """
    SELECT parent_id, dept_name, GROUP_CONCAT(dept_id ORDER BY dept_id ASC) AS ids, COUNT(*) AS cnt
    FROM sys_dept
    GROUP BY parent_id, dept_name
    HAVING cnt > 1
    ORDER BY parent_id, dept_name
    """
    if limit and limit > 0:
        sql = sql + f" LIMIT {limit}"
    with conn.cursor() as cur:
        cur.execute(sql)
        rows = cur.fetchall()
    groups = []
    for r in rows:
        ids = [int(x) for x in str(r["ids"]).split(",") if x.strip()]
        groups.append((r["parent_id"], r["dept_name"], ids))
    return groups


def reparent_children_and_users(conn: Connection, keep_id: int, delete_ids: List[int], dry_run: bool):
    if not delete_ids:
        return
    # 1) Reparent child depts whose parent_id in delete_ids -> set parent_id = keep_id
    sql_children = "SELECT dept_id FROM sys_dept WHERE parent_id IN ({ids})".format(
        ids=",".join(str(i) for i in delete_ids)
    )
    with conn.cursor() as cur:
        cur.execute(sql_children)
        children = cur.fetchall()
    if children:
        children_ids = [r["dept_id"] for r in children]
        logger.info("  Found %d child departments to reparent to %s: %s", len(children_ids), keep_id, children_ids)
        if not dry_run:
            sql_upd = "UPDATE sys_dept SET parent_id = %s WHERE parent_id IN ({ids})".format(
                ids=",".join(str(i) for i in delete_ids)
            )
            with conn.cursor() as cur:
                cur.execute(sql_upd, (keep_id,))
    else:
        logger.debug("  No child departments referencing deleted ids.")

    # 2) Update sys_user.dept_id referencing deleted ids -> keep_id
    sql_users = "SELECT count(*) as cnt FROM sys_user WHERE dept_id IN ({ids})".format(
        ids=",".join(str(i) for i in delete_ids)
    )
    with conn.cursor() as cur:
        cur.execute(sql_users)
        cnt_row = cur.fetchone()
        user_cnt = int(cnt_row["cnt"])
    if user_cnt > 0:
        logger.info("  Found %d users referencing deleted dept ids -> will update to %s", user_cnt, keep_id)
        if not dry_run:
            sql_upd_users = "UPDATE sys_user SET dept_id = %s WHERE dept_id IN ({ids})".format(
                ids=",".join(str(i) for i in delete_ids)
            )
            with conn.cursor() as cur:
                cur.execute(sql_upd_users, (keep_id,))


def delete_dept_ids(conn: Connection, delete_ids: List[int], dry_run: bool):
    if not delete_ids:
        return
    logger.info("  Will delete dept ids: %s", delete_ids)
    if dry_run:
        return
    sql = "DELETE FROM sys_dept WHERE dept_id IN ({ids})".format(ids=",".join(str(i) for i in delete_ids))
    with conn.cursor() as cur:
        cur.execute(sql)


def process_groups(conn: Connection, groups, apply: bool):
    dry_run = not apply
    summary = {"groups": 0, "deleted_ids_total": 0}
    for parent_id, dept_name, ids in groups:
        keep_id = min(ids)
        delete_ids = [i for i in ids if i != keep_id]
        summary["groups"] += 1
        summary["deleted_ids_total"] += len(delete_ids)
        logger.info("Group parent_id=%s dept_name=%r => keep=%s delete=%s", parent_id, dept_name, keep_id, delete_ids)
        # Reparent children and users first
        reparent_children_and_users(conn, keep_id, delete_ids, dry_run=dry_run)
        # Then delete
        delete_dept_ids(conn, delete_ids, dry_run=dry_run)
    return summary


def main():
    args = parse_args()
    conn = get_db_conn(args)
    try:
        groups = find_duplicate_groups(conn, limit=args.limit)
        if not groups:
            logger.info("No duplicate groups found.")
            return 0
        logger.info("Found %d duplicate groups (parent_id, dept_name with count>1).", len(groups))
        for idx, (parent_id, dept_name, ids) in enumerate(groups, start=1):
            logger.info(" %d) parent_id=%s dept_name=%r ids=%s", idx, parent_id, dept_name, ids)

        if args.apply:
            logger.info("Applying changes inside a transaction...")
            try:
                summary = process_groups(conn, groups, apply=True)
                conn.commit()
                logger.info("Committed. Processed %d groups, total deleted ids %d", summary["groups"], summary["deleted_ids_total"])
            except Exception as e:
                conn.rollback()
                logger.exception("Error during apply, rolled back. Error: %s", e)
                return 2
        else:
            logger.info("Dry-run mode (no changes will be made). Showing what would be executed:")
            _ = process_groups(conn, groups, apply=False)
            logger.info("Dry-run finished. Re-run with --apply to execute.")
        return 0
    finally:
        conn.close()


if __name__ == "__main__":
    sys.exit(main())