139 lines
4.9 KiB
Python
139 lines
4.9 KiB
Python
#!/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()
|
||
|