286 lines
12 KiB
Python
286 lines
12 KiB
Python
#!/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)
|
||
|
||
|