all-project/教职工数据迁移/import_external_staff.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将 `人员清单列表@外聘.xlsx` 导入到临时表 temp_external_staff

用法:
    python import_external_staff.py

依赖:
    pandas sqlalchemy pymysql openpyxl

说明:
    - 从同目录或上级目录读取 db_config.py 中的 DB_CONFIG
    - 自动根据 Excel 表头创建临时表（所有列使用 TEXT）
    - 支持批量写入
"""

from __future__ import annotations
import os
import sys
import math
import argparse
import importlib.util
from typing import List, Dict, Any

import pandas as pd
from sqlalchemy import create_engine, text

try:
    import openpyxl
except Exception:
    openpyxl = None


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 clean_column_name(col: str) -> str:
    if col is None:
        return ""
    s = str(col).strip()
    s = s.replace("（", "_").replace("）", "_").replace("(", "_").replace(")", "_")
    s = s.replace("、", "_").replace(" ", "_").replace("\n", "_").replace("\r", "_")
    # ensure valid identifier start
    if not s:
        s = "col"
    if not (s[0].isalpha() or s[0] == "_"):
        s = f"c_{s}"
    return s


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}"
    return create_engine(uri, pool_pre_ping=True)


def create_table(conn, table_name: str, columns: List[str]):
    # 先单独执行 DROP，再执行 CREATE（避免多语句执行限制导致语法错误）
    conn.execute(text(f"DROP TABLE IF EXISTS `{table_name}`"))
    lines: List[str] = []
    lines.append(f"CREATE TABLE `{table_name}` (")
    lines.append("  `id` BIGINT(20) NOT NULL AUTO_INCREMENT,")
    for col in columns:
        lines.append(f"  `{col}` TEXT DEFAULT NULL,")
    lines.append("  `import_time` DATETIME DEFAULT CURRENT_TIMESTAMP,")
    lines.append("  PRIMARY KEY (`id`)")
    lines.append(") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;")
    sql = "\n".join(lines)
    conn.execute(text(sql))


def import_excel_to_table(engine, excel_path: str, table_name: str, chunk_size: int = 1000):
    if not os.path.exists(excel_path):
        raise FileNotFoundError(f"找不到文件: {excel_path}")
    # 读取表头和全部数据使用 pandas（适合中小文件）
    df = pd.read_excel(excel_path, engine="openpyxl")
    raw_cols = list(df.columns)
    clean_cols = [clean_column_name(c) for c in raw_cols]
    # 重命名 DataFrame 列为清理后的列名
    df.columns = clean_cols

    with engine.begin() as conn:
        # 创建表
        create_table(conn, table_name, clean_cols)
        # 分块写入
        total = 0
        for start in range(0, len(df), chunk_size):
            chunk = df.iloc[start:start+chunk_size].where(pd.notnull(df), None)
            chunk.to_sql(name=table_name, con=conn, index=False, if_exists="append", method="multi", chunksize=chunk_size)
            total += len(chunk)
    return total


def main():
    parser = argparse.ArgumentParser(description="将人员清单（外聘）导入临时表")
    parser.add_argument("--file", type=str, default="人员清单列表@外聘.xlsx", help="Excel 文件名或路径")
    parser.add_argument("--table", type=str, default="temp_external_staff", help="临时表名")
    parser.add_argument("--db-config", type=str, default="../学校历史数据迁移/db_config.py", help="db_config.py 路径")
    args = parser.parse_args()

    base_dir = os.path.dirname(os.path.abspath(__file__))
    excel_path = args.file
    if not os.path.isabs(excel_path):
        excel_path = os.path.join(base_dir, excel_path)

    db_config_path = args.db_config
    if not os.path.isabs(db_config_path):
        db_config_path = os.path.join(base_dir, db_config_path)
    if not os.path.exists(db_config_path):
        # try parent folder
        db_config_path = os.path.join(os.path.dirname(base_dir), "学校历史数据迁移", "db_config.py")
        if not os.path.exists(db_config_path):
            print(f"找不到 db_config 文件: {db_config_path}")
            sys.exit(1)

    db_config = load_db_config_from_file(db_config_path)
    engine = build_engine_from_config(db_config)

    print(f"开始导入文件: {excel_path} 到表 {args.table}")
    total = import_excel_to_table(engine, excel_path, args.table, chunk_size=1000)
    print(f"导入完成，共插入 {total} 条记录到 {args.table}")


if __name__ == "__main__":
    main()