226 lines
8.7 KiB
Python
226 lines
8.7 KiB
Python
#!/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())
|
||
|
||
|