提交
This commit is contained in:
zlwu committed
1 parent
3c629d2d58
commit
34c33ebc54
12 files changed
+1206
No files matched your search
@@ -0,0 +1,263 @@
|
||||
#!/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
|
||||
Reference in new issue
Block a user