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

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
04 - 专业信息入库脚本

说明:
- 从 `his_xssj` 表读取学院编号、专业编号、专业名称（去重）
- 根据学院编号映射到学院名称，然后在 `sys_dept`（或目标表）中查找父节点 parent_id
- 若专业在目标表中不存在，则生成 INSERT 语句（默认 dry-run，仅打印）
- 可使用 --execute 开关执行插入

用法示例:
  python3 "学校历史数据迁移/04-insert_majors.py" --host localhost --user root --password secret --database hzdyjs_xg
  python3 "学校历史数据迁移/04-insert_majors.py" --execute --host ...  (实际写入)

注意:
- 请提前安装依赖: pip install pymysql
"""
from __future__ import annotations

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

import pymysql

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


COLLEGE_MAP: Dict[str, str] = {
    "14": "智能控制学院",
    "15": "智能制造学院",
    "19": "人文学院",
    "18": "餐饮休闲学院",
    "28": "数字信息学院",
    "16": "医药健康学院",
    "17": "生态园林学院",
}


def load_local_db_config() -> Dict:
    """
    尝试从当前脚本目录的 db_config.py 加载 DB_CONFIG 字典并返回（失败则返回空字典）
    """
    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(host: str, port: int, user: str, password: str, database: str):
    """建立并返回 pymysql 连接（默认 autocommit False）"""
    return pymysql.connect(
        host=host,
        port=port,
        user=user,
        password=password,
        database=database,
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
        autocommit=False,
    )


def college_has_majors(conn, parent_id: int, target_table: str) -> bool:
    """
    检查指定 parent_id（学院）下是否已有专业类记录（dept_category 以 'zy-' 开头或任意子节点）
    若存在任一子记录则返回 True。
    """
    # note: percent signs must be escaped (%%) because PyMySQL uses Python %-formatting during mogrify
    # Only consider existing records that look like majors (dept_category starting with 'zy-').
    # Previously the condition included "OR dept_category IS NOT NULL" which caused false positives
    # and made the script skip entire colleges when any child had a non-null dept_category.
    sql = f"SELECT 1 FROM {target_table} WHERE parent_id = %s AND dept_category LIKE 'zy-%%' LIMIT 1"
    with conn.cursor() as cur:
        cur.execute(sql, (parent_id,))
        return cur.fetchone() is not None


def fetch_distinct_majors(conn) -> List[Tuple[str, str, str]]:
    """
    从 his_xssj 表读取 distinct (college_code, major_code, major_name)
    假设表名为 his_xssj，字段名可调整如 xybm（学院编码）、zybm（专业编码）、zymc（专业名称）
    """
    # use Chinese column names present in his_xssj table
    sql = """
SELECT DISTINCT
    COALESCE(`学院编号`, '') AS college_code,
    COALESCE(`专业编号`, '') AS major_code,
    COALESCE(`专业名称`, '') AS major_name
FROM his_xssj
WHERE `专业编号` IS NOT NULL AND `专业编号` <> '' AND `专业名称` IS NOT NULL AND `专业名称` <> ''
ORDER BY college_code, major_code;
"""
    with conn.cursor() as cur:
        cur.execute(sql)
        rows = cur.fetchall()
    return [(r["college_code"], r["major_code"], r["major_name"]) for r in rows]


def find_parent_dept(conn, college_name: str) -> Optional[Dict]:
    """
    根据学院名称查找 sys_dept 中的父节点（按 dept_name 精确匹配）
    返回整行记录（字典）或 None
    """
    sql = "SELECT * FROM sys_dept WHERE dept_name = %s LIMIT 1"
    with conn.cursor() as cur:
        cur.execute(sql, (college_name,))
        return cur.fetchone()


def exists_major(conn, parent_id: int, major_code: str, major_name: str, target_table: str) -> bool:
    """
    判断目标表中是否已存在该专业（通过 dept_category 或 dept_name 与 parent_id 共同判断）
    """
    # Consider a major existing only if BOTH dept_category and dept_name match under the same parent_id
    sql = f"SELECT 1 FROM {target_table} WHERE parent_id = %s AND dept_category = %s AND dept_name = %s LIMIT 1"
    with conn.cursor() as cur:
        cur.execute(sql, (parent_id, f"zy-{major_code}", major_name))
        return cur.fetchone() is not None


def gen_dept_id(counter: int) -> int:
    """
    生成一个大整数 dept_id:
    使用当前时间 (YYYYMMDDHHMMSS) + 3-digit counter 保证唯一性（可根据需要替换为雪花ID）
    """
    ts = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    return int(f"{ts}{counter:03d}")


def insert_major(conn, target_table: str, parent_row: Dict, major_code: str, major_name: str, dept_id: int):
    """
    插入一条 sys_dept 记录到 target_table
    - parent_row: 父节点整行（包含 dept_id, ancestors）
    """
    now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    parent_id = parent_row["dept_id"]
    parent_ancestors = parent_row.get("ancestors") or ""
    if parent_ancestors:
        ancestors = f"{parent_ancestors},{parent_id}"
    else:
        ancestors = str(parent_id)

    dept_category = f"zy-{major_code}"
    sql = f"""
INSERT INTO {target_table}
    (dept_id, tenant_id, parent_id, ancestors, dept_name, dept_category, order_num, leader, phone, email, status, del_flag, create_dept, create_by, create_time, update_by, update_time)
VALUES
    (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
"""
    params = (
        dept_id,
        "000000",  # tenant_id
        parent_id,
        ancestors,
        major_name,
        dept_category,
        1,  # order_num
        None,  # leader
        None,  # phone
        None,  # email
        "0",  # status
        "0",  # del_flag
        0,  # create_dept
        1,  # create_by
        now,
        1,  # update_by
        now,
    )
    with conn.cursor() as cur:
        cur.execute(sql, params)


def main(argv: Optional[List[str]] = None):
    # 尝试从同目录下的 db_config.py 读取默认数据库配置
    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"))
    target_table_default = os.getenv("TARGET_TABLE", "sys_dept")

    parser = argparse.ArgumentParser(description="从 his_xssj 导入专业到 sys_dept")
    parser.add_argument("--host", default=host_default)
    parser.add_argument("--port", type=int, default=port_default)
    parser.add_argument("--user", default=user_default)
    parser.add_argument("--password", default=password_default)
    parser.add_argument("--database", default=database_default)
    parser.add_argument("--target-table", default=target_table_default)
    parser.add_argument("--dry-run", action="store_true", default=False, help="仅打印将要执行的 SQL（默认执行写库）")
    parser.add_argument("--execute", action="store_true", default=True, help="执行插入（会在事务中执行并提交，默认开启）")
    parser.add_argument("--limit", type=int, default=0, help="仅处理前 N 条（用于测试）")
    args = parser.parse_args(argv)

    conn = get_db_conn(args.host, args.port, args.user, args.password, args.database)
    try:
        majors = fetch_distinct_majors(conn)
        logger.info("从 his_xssj 获取到 %d 条去重专业记录", len(majors))

        counter = 1
        inserted = 0
        skipped = 0
        sql_statements: List[str] = []
        # 记录对每个学院的整体决定（'skip' or 'insert'），避免对同一学院重复查询
        college_decision: Dict[str, str] = {}

        for idx, (college_code, major_code, major_name) in enumerate(majors, start=1):
            if args.limit and inserted >= args.limit:
                break

            college_name = COLLEGE_MAP.get(college_code)
            if not college_name:
                logger.warning("未在 COLLEGE_MAP 中找到学院编码 %s，跳过: %s", college_code, major_name)
                skipped += 1
                continue

            parent_row = find_parent_dept(conn, college_name)
            if not parent_row:
                logger.warning("未在 sys_dept 中找到学院父节点（dept_name=%s），请先确认父节点存在，跳过: %s", college_name, major_name)
                skipped += 1
                continue

            # 判重：在同一父节点下，若 dept_category 和 dept_name 同时匹配则视为已存在并跳过
            if exists_major(conn, parent_row["dept_id"], major_code, major_name, args.target_table):
                logger.info("专业已存在，跳过: %s (%s) under %s", major_name, major_code, college_name)
                skipped += 1
                continue

            dept_id = gen_dept_id(counter)
            counter += 1

            # Prepare SQL string for logging / dry-run
            now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            parent_id = parent_row["dept_id"]
            parent_ancestors = parent_row.get("ancestors") or ""
            ancestors = f"{parent_ancestors},{parent_id}" if parent_ancestors else str(parent_id)
            dept_category = f"zy-{major_code}"
            # escape single quotes in major_name for SQL literal
            safe_major_name = str(major_name).replace("'", "''")
            # ancestors is numeric-comma separated; safe to include directly
            insert_sql = (
                "INSERT INTO {table} (dept_id, tenant_id, parent_id, ancestors, dept_name, dept_category, order_num, leader, phone, email, status, del_flag, create_dept, create_by, create_time, update_by, update_time) "
                "VALUES ({did}, '000000', {pid}, '{ancs}', '{dname}', '{dcat}', 1, NULL, NULL, NULL, '0', '0', 0, 1, '{now}', 1, '{now}');"
            ).format(
                table=args.target_table,
                did=dept_id,
                pid=parent_id,
                ancs=ancestors,
                dname=safe_major_name,
                dcat=dept_category,
                now=now,
            )
            sql_statements.append(insert_sql)

            logger.info("准备插入专业: %s (%s) parent=%s new_dept_id=%s", major_name, major_code, college_name, dept_id)

            if args.execute:
                insert_major(conn, args.target_table, parent_row, major_code, major_name, dept_id)
                conn.commit()
                inserted += 1
            else:
                skipped += 1

        # 输出 dry-run 的 SQL
        if sql_statements:
            logger.info("共准备 %d 条 INSERT 语句 (dry-run=%s, execute=%s)", len(sql_statements), args.dry_run, args.execute)
            for s in sql_statements:
                print(s)
        else:
            logger.info("没有需要插入的专业记录")

        logger.info("完成: 插入 %d 条, 跳过 %d 条", inserted, skipped)
    finally:
        conn.close()


if __name__ == "__main__":
    main()