提交

2026-01-09 10:24:31 +08:00
1 parent 3c629d2d58
commit 34c33ebc54
@@ -0,0 +1,138 @@
+#!/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()
+