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

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

目的:
  将 `sys_dept` 表中，同一个学院（parent_id）下的专业记录（可通过 dept_category 判断）按照 `dept_name` 升序重新设置 `order_num`（从 1 开始递增）。

行为:
  - 默认 dry-run：打印将要执行的 UPDATE 语句，不做修改。
  - 传入 --apply 会在事务中执行更新。
  - 支持通过 --major-only 控制是否仅对 dept_category LIKE 'zy-%' 的记录进行排序（默认 True）。
  - 支持 --limit 控制处理的 parent_id 数量（用于测试）。

使用:
  python3 04-02-reorder_dept_order_num.py --host ... --user ... --database ... [--apply]
"""

from __future__ import annotations

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

import pymysql
from pymysql.connections import Connection

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


def load_local_db_config() -> dict:
    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 parse_args():
    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="Reorder sys_dept.order_num by dept_name under the same parent_id")
    p.add_argument("--host", default=host_default)
    p.add_argument("--port", type=int, default=port_default)
    p.add_argument("--user", default=user_default)
    p.add_argument("--password", default=password_default)
    p.add_argument("--database", default=database_default)
    p.add_argument("--charset", default="utf8mb4")
    p.add_argument("--apply", action="store_true", help="Apply changes (default dry-run)")
    p.add_argument("--limit", type=int, default=0, help="Limit number of parent groups to process (for testing)")
    p.add_argument("--major-only", action="store_true", default=True, help="Only reorder records where dept_category LIKE 'zy-%' (default True)")
    return p.parse_args()


def get_db_conn(args) -> Connection:
    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_parent_ids(conn: Connection, major_only: bool = True, limit: int = 0) -> List[int]:
    """
    Find parent_id values that have child departments to reorder.
    If major_only is True, only consider children with dept_category LIKE 'zy-%'.
    """
    if major_only:
        sql = """
        SELECT parent_id
        FROM sys_dept
        WHERE dept_category LIKE 'zy-%%' AND del_flag = '0'
        GROUP BY parent_id
        ORDER BY parent_id
        """
    else:
        sql = """
        SELECT parent_id
        FROM sys_dept
        WHERE parent_id IS NOT NULL AND del_flag = '0'
        GROUP BY parent_id
        ORDER BY parent_id
        """
    if limit and limit > 0:
        sql = sql + f" LIMIT {limit}"
    with conn.cursor() as cur:
        cur.execute(sql)
        rows = cur.fetchall()
    return [r["parent_id"] for r in rows if r["parent_id"] is not None]


def fetch_children(conn: Connection, parent_id: int, major_only: bool = True) -> List[Dict]:
    """
    Fetch child departments for a parent_id, ordered by dept_name asc then dept_id asc.
    """
    if major_only:
        sql = """
        SELECT dept_id, dept_name, order_num
        FROM sys_dept
        WHERE parent_id = %s AND dept_category LIKE 'zy-%%' AND del_flag = '0'
        ORDER BY dept_name ASC, dept_id ASC
        """
    else:
        sql = """
        SELECT dept_id, dept_name, order_num
        FROM sys_dept
        WHERE parent_id = %s AND del_flag = '0'
        ORDER BY dept_name ASC, dept_id ASC
        """
    with conn.cursor() as cur:
        cur.execute(sql, (parent_id,))
        return cur.fetchall()


def apply_updates(conn: Connection, updates: List[Tuple[int, int]], dry_run: bool):
    """
    updates: list of (dept_id, new_order_num)
    """
    if not updates:
        return
    logger.info("  %d updates to apply", len(updates))
    for dept_id, new_order in updates:
        sql = "UPDATE sys_dept SET order_num = %s WHERE dept_id = %s"
        logger.info("    %s -- params: (%s, %s)", sql, new_order, dept_id)
        if not dry_run:
            with conn.cursor() as cur:
                cur.execute(sql, (new_order, dept_id))


def process(conn: Connection, parent_ids: List[int], major_only: bool, apply: bool):
    dry_run = not apply
    summary = {"parents": 0, "updates": 0}
    for parent_id in parent_ids:
        children = fetch_children(conn, parent_id, major_only=major_only)
        if not children:
            logger.debug("Parent %s has no children to reorder", parent_id)
            continue
        updates = []
        for idx, child in enumerate(children, start=1):
            dept_id = child["dept_id"]
            current_order = child.get("order_num") or 0
            if current_order != idx:
                updates.append((dept_id, idx))
        if updates:
            logger.info("Parent %s: will update %d rows", parent_id, len(updates))
            apply_updates(conn, updates, dry_run=dry_run)
            summary["updates"] += len(updates)
        else:
            logger.info("Parent %s: no updates needed", parent_id)
        summary["parents"] += 1
    return summary


def main():
    args = parse_args()
    conn = get_db_conn(args)
    try:
        parent_ids = find_parent_ids(conn, major_only=args.major_only, limit=args.limit)
        if not parent_ids:
            logger.info("No parent groups found to process.")
            return 0
        logger.info("Found %d parent groups to process (limit=%s, major_only=%s).", len(parent_ids), args.limit, args.major_only)
        # show parents
        for pid in parent_ids:
            logger.info("  parent_id=%s", pid)

        if args.apply:
            logger.info("Applying updates inside a transaction...")
            try:
                summary = process(conn, parent_ids, args.major_only, apply=True)
                conn.commit()
                logger.info("Committed. Processed %d parents, total updates %d", summary["parents"], summary["updates"])
            except Exception as e:
                conn.rollback()
                logger.exception("Error during apply, rolled back. Error: %s", e)
                return 2
        else:
            logger.info("Dry-run mode: showing planned updates (no DB changes).")
            summary = process(conn, parent_ids, args.major_only, apply=False)
            logger.info("Dry-run finished. Processed %d parents, planned updates %d. Re-run with --apply to execute.", summary["parents"], summary["updates"])
        return 0
    finally:
        conn.close()


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