all-project/学校历史数据迁移/01-import_history.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
导入学校历史学生和老师数据到 MySQL 表

用法:
    python import_history.py

依赖:
    pip install -r requirements.txt

说明:
    - 从同目录下的 db_config.py 读取 DB_CONFIG
    - 将 `student.xls` 导入到表 `his_xssj`
    - 将 `teacher.xlsx` 导入到表 `his_lssj`
    - 自动创建表（所有列使用 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

# 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", "_")
        # 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 df_to_mysql_append(engine, df: pd.DataFrame, table_name: str, first_write: bool):
    # ensure all columns are strings (TEXT)
    df = df.where(pd.notnull(df), None)
    if first_write:
        # create table with TEXT columns by writing empty dataframe with dtype mapping
        empty = df.head(0)
        dtype_map = {col: sql_types.Text() for col in empty.columns}
        empty.to_sql(name=table_name, con=engine, index=False, if_exists="replace", dtype=dtype_map)
        # then append the first actual df
        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_xls(filepath: str, engine, table_name: str, chunk_size: int = 2000):
    if xlrd is None:
        raise RuntimeError("Missing xlrd library to read .xls files. Install from requirements.")
    book = xlrd.open_workbook(filepath, formatting_info=False)
    sheet = book.sheet_by_index(0)
    header = [str(c).strip() for c in sheet.row_values(0)]
    header = sanitize_columns(header)
    total = sheet.nrows - 1
    print(f"读取 .xls 文件 {filepath}, 行数(不含表头): {total}")
    batch: List[Dict[str, Any]] = []
    first_write = True
    for r in range(1, sheet.nrows):
        row = sheet.row_values(r)
        # convert nan floats to 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
            batch = []
            print(f"已写入 {chunk_size} 行到 {table_name}")
    if batch:
        df = pd.DataFrame(batch)
        df_to_mysql_append(engine, df, table_name, first_write)
        print(f"写入剩余 {len(batch)} 行到 {table_name}")


def process_xlsx(filepath: str, engine, table_name: str, chunk_size: int = 2000):
    if openpyxl is None:
        raise RuntimeError("Missing openpyxl library to read .xlsx files. Install from requirements.")
    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)
    batch: List[Dict[str, Any]] = []
    first_write = True
    written = 0
    for row in rows:
        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(filepath: str, engine, table_name: str, chunk_size: int = 2000):
    ext = os.path.splitext(filepath)[1].lower()
    if ext == ".xls":
        process_xls(filepath, engine, table_name, chunk_size=chunk_size)
    elif ext in (".xlsx",):
        process_xlsx(filepath, engine, table_name, chunk_size=chunk_size)
    else:
        # try pandas fallback (may be memory heavy)
        print(f"未知扩展名 {ext}，尝试使用 pandas 直接读取（可能占用大量内存）")
        df = pd.read_excel(filepath, engine="openpyxl" if ext == ".xlsx" else None)
        df.columns = sanitize_columns(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):
    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):
        print(f"找不到配置文件: {config_path}")
        sys.exit(1)
    db_config = load_db_config_from_file(config_path)
    engine = build_engine_from_config(db_config)
    # 在每次导入前询问是否清空目标表（或使用 --yes 自动确认）
    should_truncate = False
    if args and getattr(args, "yes", False):
        should_truncate = True
    else:
        resp = input("是否在导入前清空目标表 `his_xssj` 和 `his_lssj`？(yes/no) [no]: ").strip().lower()
        if resp in ("y", "yes"):
            should_truncate = True

    if should_truncate:
        print("尝试清空目标表: his_xssj, his_lssj ...")
        try:
            with engine.begin() as conn:
                # 关闭外键检查以确保能够 truncate 有外键依赖的表
                conn.exec_driver_sql("SET FOREIGN_KEY_CHECKS=0;")
                conn.exec_driver_sql("TRUNCATE TABLE `his_xssj`;")
                conn.exec_driver_sql("TRUNCATE TABLE `his_lssj`;")
                conn.exec_driver_sql("SET FOREIGN_KEY_CHECKS=1;")
            print("目标表已清空（如果表存在且有权限）")
        except Exception as e:
            print(f"清空表时发生异常（可能表不存在或无权限），继续执行导入。错误：{e}")
    else:
        print("跳过清空表，继续导入（若需要清空请重新运行并使用 --yes 参数）。")

    student_file = os.path.join(base_dir, "student1222.xls")
    teacher_file = os.path.join(base_dir, "teacher.xlsx")
    if not os.path.exists(student_file):
        print(f"找不到学生文件: {student_file}")
    else:
        print("开始导入学生数据到表 his_xssj ...")
        process_file(student_file, engine, "his_xssj", chunk_size=2000)
        print("学生数据导入完成 -> his_xssj")

    if not os.path.exists(teacher_file):
        print(f"找不到老师文件: {teacher_file}")
    else:
        print("开始导入老师数据到表 his_lssj ...")
        process_file(teacher_file, engine, "his_lssj", chunk_size=2000)
        print("老师数据导入完成 -> his_lssj")

    print("全部导入完成。")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="导入历史学生/老师数据或生成创建表SQL")
    parser.add_argument("--dump-create-sql", action="store_true", help="只读取表头并生成 CREATE TABLE SQL 文件（不写入数据库）")
    parser.add_argument("--student-file", type=str, default=None, help="学生文件路径，默认同目录下 student.xls")
    parser.add_argument("--teacher-file", type=str, default=None, help="老师文件路径，默认同目录下 teacher.xlsx")
    parser.add_argument("--yes", action="store_true", help="在导入前自动确认清空目标表（跳过交互式确认）")
    args = parser.parse_args()

    if args.dump_create_sql:
        # generate CREATE TABLE SQL based on excel headers
        def generate_create_table_sql(table_name: str, cols: List[str]) -> str:
            # base fields required by project
            lines: List[str] = []
            lines.append(f"CREATE TABLE `{table_name}` (")
            lines.append("  `id` BIGINT(20) NOT NULL AUTO_INCREMENT,")
            lines.append("  `tenant_id` BIGINT(20) DEFAULT NULL,")
            lines.append("  `deleted` TINYINT(1) NOT NULL DEFAULT 0,")
            lines.append("  `create_by` BIGINT(20) DEFAULT NULL,")
            lines.append("  `create_time` DATETIME DEFAULT NULL,")
            lines.append("  `update_by` BIGINT(20) DEFAULT NULL,")
            lines.append("  `update_time` DATETIME DEFAULT NULL,")
            # add excel columns as TEXT
            for col in cols:
                lines.append(f"  `{col}` TEXT DEFAULT NULL,")
            lines.append("  PRIMARY KEY (`id`)")
            lines.append(") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;")
            return "\n".join(lines)

        base_dir = os.path.dirname(os.path.abspath(__file__))
        student_path = args.student_file or os.path.join(base_dir, "student1222.xls")
        teacher_path = args.teacher_file or os.path.join(base_dir, "teacher.xlsx")

        for path, table in ((student_path, "his_xssj"), (teacher_path, "his_lssj")):
            if not os.path.exists(path):
                print(f"跳过，找不到文件: {path}")
                continue
            ext = os.path.splitext(path)[1].lower()
            try:
                # read header only
                if ext == ".xls":
                    df = pd.read_excel(path, nrows=0, engine="xlrd")
                else:
                    df = pd.read_excel(path, nrows=0, engine="openpyxl")
            except Exception as e:
                # fallback
                print(f"读取表头失败 ({path}): {e}")
                continue
            raw_cols = list(df.columns)
            safe_cols = sanitize_columns([str(c) for c in raw_cols])
            sql = generate_create_table_sql(table, safe_cols)
            out_file = os.path.join(base_dir, f"{table}_create.sql")
            with open(out_file, "w", encoding="utf-8") as fh:
                fh.write(sql)
            print(f"已生成 {out_file}，包含 {len(safe_cols)} 列（不含基础字段）")
    else:
        main(args)