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

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
导入教职工数据到临时表

用法:
    python import_staff_data.py

依赖:
    pip install pandas sqlalchemy pymysql openpyxl xlrd

说明:
    - 从同目录下的 db_config.py 读取 DB_CONFIG
    - 将 `人员清单列表1@正式员工.xlsx` 导入到临时表 `temp_staff_data`
    - 自动创建临时表（所有列使用 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, types as sql_types, text

# Optional, used for xls/xlsx low-level reading
try:
    import xlrd  # for .xls
except Exception:
    xlrd = None

try:
    import openpyxl  # for .xlsx
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 sanitize_columns(header: List[str]) -> List[str]:
    """清理和标准化列名"""
    seen: Dict[str, int] = {}
    out: List[str] = []
    for i, col in enumerate(header):
        if col is None:
            name = f"col_{i+1}"
        else:
            name = str(col).strip()
            if name == "":
                name = f"col_{i+1}"
        # replace problematic characters
        name = name.replace(" ", "_").replace("\n", "_").replace("\r", "_").replace("(", "_").replace(")", "_")
        # ensure starts with letter or underscore
        if not name[0].isalpha() and name[0] != "_":
            name = f"c_{name}"
        base = name
        cnt = seen.get(base, 0)
        if cnt:
            name = f"{base}_{cnt}"
        seen[base] = cnt + 1
        out.append(name)
    return out


def create_temp_table(engine, table_name: str, columns: List[str]):
    """创建临时表"""
    print(f"创建临时表: {table_name}")

    # 构建建表SQL
    lines: List[str] = []
    lines.append(f"DROP TABLE IF EXISTS `{table_name}`;")
    lines.append(f"CREATE TABLE `{table_name}` (")
    lines.append("  `id` BIGINT(20) NOT NULL AUTO_INCREMENT,")

    # 添加Excel列（TEXT类型）
    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;")

    create_sql = "\n".join(lines)

    try:
        with engine.begin() as conn:
            conn.execute(text(create_sql))
        print(f"临时表 {table_name} 创建成功")
    except Exception as e:
        print(f"创建临时表失败: {e}")
        raise


def df_to_mysql_append(engine, df: pd.DataFrame, table_name: str, first_write: bool):
    """将DataFrame追加到MySQL表"""
    # ensure all columns are strings (TEXT)
    df = df.where(pd.notnull(df), None)
    if first_write:
        # 追加数据，不创建表（表已预先创建）
        df.to_sql(name=table_name, con=engine, index=False, if_exists="append", chunksize=1000, method="multi")
    else:
        df.to_sql(name=table_name, con=engine, index=False, if_exists="append", chunksize=1000, method="multi")


def process_xlsx_staff(filepath: str, engine, table_name: str, chunk_size: int = 2000):
    """处理教职工Excel文件（.xlsx格式）"""
    if openpyxl is None:
        raise RuntimeError("Missing openpyxl library to read .xlsx files. Install from requirements.")

    print(f"开始读取教职工数据文件: {filepath}")
    wb = openpyxl.load_workbook(filename=filepath, read_only=True, data_only=True)
    sheet = wb.active
    rows = sheet.iter_rows(values_only=True)

    try:
        header_row = next(rows)
    except StopIteration:
        print(f"{filepath} 是空文件")
        return

    # 处理表头
    header = [str(c).strip() if c is not None else "" for c in header_row]
    header = sanitize_columns(header)
    print(f"表头列: {header}")

    # 创建临时表
    create_temp_table(engine, table_name, header)

    # 处理数据行
    batch: List[Dict[str, Any]] = []
    first_write = True
    written = 0

    for row in rows:
        # 转换NaN值为None
        row = [None if (isinstance(x, float) and math.isnan(x)) else x for x in row]
        row_dict = {col: row[idx] if idx < len(row) else None for idx, col in enumerate(header)}
        batch.append(row_dict)

        if len(batch) >= chunk_size:
            df = pd.DataFrame(batch)
            df_to_mysql_append(engine, df, table_name, first_write)
            first_write = False
            written += len(batch)
            batch = []
            print(f"已写入 {written} 行到 {table_name}")

    if batch:
        df = pd.DataFrame(batch)
        df_to_mysql_append(engine, df, table_name, first_write)
        written += len(batch)
        print(f"写入剩余 {len(batch)} 行到 {table_name}，总计写入 {written} 行")


def process_file_staff(filepath: str, engine, table_name: str, chunk_size: int = 2000):
    """处理教职工数据文件"""
    ext = os.path.splitext(filepath)[1].lower()
    if ext in (".xlsx",):
        process_xlsx_staff(filepath, engine, table_name, chunk_size=chunk_size)
    else:
        # 尝试使用pandas fallback
        print(f"未知扩展名 {ext}，尝试使用 pandas 直接读取（可能占用大量内存）")
        df = pd.read_excel(filepath, engine="openpyxl" if ext == ".xlsx" else None)
        df.columns = sanitize_columns(list(df.columns))

        # 创建临时表
        create_temp_table(engine, table_name, list(df.columns))

        # 导入数据
        df_to_mysql_append(engine, df, table_name, first_write=True)


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


def main(args=None):
    """主函数"""
    base_dir = os.path.dirname(os.path.abspath(__file__))

    # 查找数据库配置文件
    config_path = os.path.join(base_dir, "db_config.py")
    if not os.path.exists(config_path):
        # 尝试从上级目录查找
        config_path = os.path.join(os.path.dirname(base_dir), "学校历史数据迁移", "db_config.py")
        if not os.path.exists(config_path):
            print(f"找不到配置文件: {config_path}")
            sys.exit(1)

    print(f"使用配置文件: {config_path}")
    db_config = load_db_config_from_file(config_path)
    engine = build_engine_from_config(db_config)

    # 教职工数据文件路径
    staff_file = os.path.join(base_dir, "人员清单列表1@正式员工.xlsx")
    if args and args.file:
        staff_file = args.file

    table_name = "temp_staff_data"
    if args and args.table:
        table_name = args.table

    if not os.path.exists(staff_file):
        print(f"找不到教职工数据文件: {staff_file}")
        sys.exit(1)

    print(f"开始导入教职工数据到临时表 {table_name} ...")
    try:
        process_file_staff(staff_file, engine, table_name, chunk_size=1000)
        print(f"教职工数据导入完成 -> {table_name}")

        # 验证导入结果
        with engine.begin() as conn:
            result = conn.execute(text(f"SELECT COUNT(*) as count FROM {table_name}"))
            count = result.fetchone()[0]
            print(f"临时表 {table_name} 共导入 {count} 条记录")

            # 显示前几行数据
            result = conn.execute(text(f"SELECT * FROM {table_name} LIMIT 5"))
            columns = result.keys()
            print(f"\n前5行数据预览:")
            print("-" * 80)
            print(f"{' | '.join(columns)}")
            print("-" * 80)
            for row in result:
                print(f"{' | '.join(str(cell) if cell is not None else '' for cell in row)}")

    except Exception as e:
        print(f"导入过程中发生错误: {e}")
        sys.exit(1)

    print("\n导入完成！")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="导入教职工数据到临时表")
    parser.add_argument("--file", type=str, help="教职工数据文件路径")
    parser.add_argument("--table", type=str, default="temp_staff_data", help="临时表名")
    args = parser.parse_args()

    main(args)</contents>
</xai:function_call name="todo_write">
<parameter name="merge">true