3b0374aa31
- Set cursor.fast_executemany = True (pyodbc 5.x moved it to cursor) - merge_sources() now returns source row counts, avoiding re-reading files - Single commit at end of batch_insert instead of per-batch commits - Pre-build all rows before SQL execution - Result: PG_2022 migration improved from ~5min to ~36s (7.6x faster) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
547 lines
17 KiB
Python
547 lines
17 KiB
Python
"""
|
||
Excel to SQL Server Migration Script
|
||
Migrates production execution card data from Excel (.xlsm) to SQL Server.
|
||
Configuration is loaded from config.yaml.
|
||
"""
|
||
|
||
import argparse
|
||
import logging
|
||
import os
|
||
import sys
|
||
from collections import OrderedDict
|
||
from datetime import datetime
|
||
|
||
import openpyxl
|
||
import pyodbc
|
||
import yaml
|
||
|
||
logger = logging.getLogger('app')
|
||
|
||
|
||
# ============================================================
|
||
# 配置加载
|
||
# ============================================================
|
||
|
||
def load_config(config_path=None):
|
||
"""加载 YAML 配置文件"""
|
||
if config_path is None:
|
||
config_path = os.path.join(
|
||
os.path.dirname(os.path.abspath(__file__)), 'config.yaml'
|
||
)
|
||
with open(config_path, 'r', encoding='utf-8') as f:
|
||
return yaml.safe_load(f)
|
||
|
||
|
||
def get_connection_string(cfg):
|
||
"""根据配置构建 pyodbc 连接字符串"""
|
||
sql = cfg['sql_server']
|
||
return (
|
||
f"DRIVER={{{sql['driver']}}};"
|
||
f"SERVER={sql['server']};"
|
||
f"DATABASE={sql['database']};"
|
||
f"UID={sql['username']};"
|
||
f"PWD={sql['password']};"
|
||
f"TrustServerCertificate={sql['TrustServerCertificate']};"
|
||
)
|
||
|
||
|
||
# ============================================================
|
||
# Excel 读取与数据合并
|
||
# ============================================================
|
||
|
||
def sanitize_headers(headers):
|
||
"""处理原始表头：
|
||
1. 空列名 → col_N
|
||
2. 重复列名 → 第二次出现加 _2 后缀
|
||
返回处理后的列名列表
|
||
"""
|
||
seen = {}
|
||
result = []
|
||
for i, h in enumerate(headers):
|
||
h = str(h).strip() if h else ''
|
||
if not h:
|
||
h = f'col_{i}'
|
||
if h in seen:
|
||
result.append(f'{h}_{seen[h] + 1}')
|
||
seen[h] += 1
|
||
else:
|
||
seen[h] = 1
|
||
result.append(h)
|
||
return result
|
||
|
||
|
||
def remap_headers(headers, mapping):
|
||
"""根据 config.yaml 中的 column_mapping 重映射列名"""
|
||
if not mapping:
|
||
return headers
|
||
return [mapping.get(h, h) for h in headers]
|
||
|
||
|
||
def resolve_pk_index(headers, primary_key):
|
||
"""找到主键列的索引。
|
||
优先精确匹配 primary_key，其次匹配含 'BW' 的列，最后回退到第 0 列。
|
||
"""
|
||
for idx, h in enumerate(headers):
|
||
if h == primary_key:
|
||
return idx
|
||
# 回退：首列含 BW（如 20BW序号）
|
||
for idx, h in enumerate(headers):
|
||
if 'BW' in h:
|
||
return idx
|
||
return 0
|
||
|
||
|
||
def read_sheet(file_path, sheet_name, primary_key, cfg):
|
||
"""读取单个工作表。
|
||
返回 (列名列表, 实际主键列名, 数据列表)。
|
||
每行数据为 dict {列名: 值}。
|
||
跳过主键列为空的行。
|
||
"""
|
||
wb = openpyxl.load_workbook(file_path, read_only=True, data_only=True)
|
||
ws = wb[sheet_name]
|
||
|
||
rows_iter = ws.iter_rows(values_only=True)
|
||
# 第一行为表头
|
||
raw_headers = next(rows_iter)
|
||
headers = sanitize_headers(raw_headers)
|
||
headers = remap_headers(headers, cfg.get('column_mapping'))
|
||
|
||
pk_idx = resolve_pk_index(headers, primary_key)
|
||
actual_pk_col = headers[pk_idx]
|
||
|
||
data = []
|
||
for row in rows_iter:
|
||
pk_val = row[pk_idx] if pk_idx < len(row) else None
|
||
if pk_val is None or str(pk_val).strip() == '':
|
||
continue
|
||
|
||
record = {}
|
||
for j, h in enumerate(headers):
|
||
val = row[j] if j < len(row) else None
|
||
record[h] = val
|
||
data.append(record)
|
||
|
||
wb.close()
|
||
return headers, actual_pk_col, data
|
||
|
||
|
||
def merge_sources(sources, excel_dir, primary_key, cfg):
|
||
"""合并多个工作表的数据。
|
||
- 列名取并集（按出现顺序）
|
||
- 数据按主键去重，后读覆盖先读
|
||
返回 (unified_headers, actual_pk_col, merged_data_ordered_dict, source_row_counts)
|
||
"""
|
||
all_headers = [] # 保持顺序的列名并集
|
||
header_set = set()
|
||
merged = OrderedDict() # key=主键值, value=dict
|
||
resolved_pk = primary_key # 实际使用的主键列名
|
||
source_row_counts = [] # 各源的行数 [(file, sheet, count), ...]
|
||
|
||
for src in sources:
|
||
file_path = os.path.join(excel_dir, src['file'])
|
||
sheet_name = src['sheet']
|
||
|
||
headers, actual_pk_col, data = read_sheet(file_path, sheet_name, primary_key, cfg)
|
||
resolved_pk = actual_pk_col
|
||
source_row_counts.append((src['file'], src['sheet'], len(data)))
|
||
|
||
# 扩展统一列名集
|
||
for h in headers:
|
||
if h not in header_set:
|
||
all_headers.append(h)
|
||
header_set.add(h)
|
||
|
||
# 合并数据（后读覆盖先读）
|
||
for record in data:
|
||
pk_val = str(record.get(actual_pk_col, '')).strip()
|
||
if not pk_val:
|
||
continue
|
||
if pk_val in merged:
|
||
merged[pk_val].update(record)
|
||
else:
|
||
merged[pk_val] = record
|
||
|
||
return all_headers, resolved_pk, merged, source_row_counts
|
||
|
||
|
||
# ============================================================
|
||
# SQL Server 操作
|
||
# ============================================================
|
||
|
||
def get_connection(cfg):
|
||
"""根据配置建立 SQL Server 连接"""
|
||
conn_str = get_connection_string(cfg)
|
||
conn = pyodbc.connect(conn_str, autocommit=False)
|
||
return conn
|
||
|
||
|
||
def create_schema(cursor, schema_name):
|
||
"""创建 schema（如不存在）"""
|
||
sql = f"""
|
||
IF NOT EXISTS (
|
||
SELECT 1 FROM sys.schemas WHERE name = ?
|
||
)
|
||
BEGIN
|
||
EXEC('CREATE SCHEMA [{schema_name}]')
|
||
END
|
||
"""
|
||
cursor.execute(sql, (schema_name,))
|
||
cursor.commit()
|
||
|
||
|
||
def infer_sql_type(col_name, type_rules, primary_key):
|
||
"""根据 config.yaml 中的 type_rules 推断 SQL Server 数据类型"""
|
||
name = col_name
|
||
|
||
# 主键列（含 TM_2022 的 20BW序号）
|
||
if name == primary_key or '20BW' in name:
|
||
return type_rules['primary_key_type']
|
||
|
||
# DATETIME
|
||
if any(kw in name for kw in type_rules.get('datetime', [])):
|
||
return 'DATETIME'
|
||
|
||
# INT（精确匹配）
|
||
for rule in type_rules.get('int', []):
|
||
if rule.get('exact') and rule['name'] == name:
|
||
return 'INT'
|
||
elif not rule.get('exact') and rule['name'] in name:
|
||
return 'INT'
|
||
|
||
# DECIMAL（精确匹配）
|
||
for rule in type_rules.get('decimal', []):
|
||
if rule.get('exact') and rule['name'] == name:
|
||
return 'DECIMAL(18,2)'
|
||
elif not rule.get('exact') and rule['name'] in name:
|
||
return 'DECIMAL(18,2)'
|
||
|
||
# NVARCHAR(MAX) 长文本
|
||
if any(kw in name for kw in type_rules.get('long_text', [])):
|
||
return 'NVARCHAR(MAX)'
|
||
|
||
return type_rules.get('default_type', 'NVARCHAR(500)')
|
||
|
||
|
||
def table_exists(cursor, schema, table_name):
|
||
"""检查表是否已存在"""
|
||
full_name = f'[{schema}].[{table_name}]'
|
||
cursor.execute(f"SELECT OBJECT_ID('{full_name}', 'U')")
|
||
row = cursor.fetchone()
|
||
return row[0] is not None
|
||
|
||
|
||
def get_existing_columns(cursor, schema, table_name):
|
||
"""查询现有表的列名集合（排除 id 自增列）"""
|
||
full_name = f'[{schema}].[{table_name}]'
|
||
cursor.execute(f"""
|
||
SELECT COLUMN_NAME
|
||
FROM INFORMATION_SCHEMA.COLUMNS
|
||
WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?
|
||
ORDER BY ORDINAL_POSITION
|
||
""", (schema, table_name))
|
||
cols = {row[0] for row in cursor.fetchall()}
|
||
cols.discard('id')
|
||
return cols
|
||
|
||
|
||
def ensure_table(cursor, table_name, headers, cfg, actual_pk, use_truncate=False):
|
||
"""确保目标表就绪。
|
||
- use_truncate=False（默认）→ 调用 create_table()，DROP + CREATE
|
||
- use_truncate=True →
|
||
1. 表不存在 → CREATE TABLE
|
||
2. 表存在且列名集合一致 → TRUNCATE TABLE
|
||
3. 表存在但列名有变化 → DROP + CREATE（兜底）
|
||
"""
|
||
if not use_truncate:
|
||
return create_table(cursor, table_name, headers, cfg, actual_pk)
|
||
|
||
schema = cfg['schema']
|
||
full_name = f'[{schema}].[{table_name}]'
|
||
|
||
if not table_exists(cursor, schema, table_name):
|
||
# 表不存在，走正常 CREATE
|
||
return create_table(cursor, table_name, headers, cfg, actual_pk)
|
||
|
||
# 比较列名集合
|
||
existing_cols = get_existing_columns(cursor, schema, table_name)
|
||
new_cols = set(headers)
|
||
# 主键列也应包含在比较中
|
||
new_cols.add(actual_pk)
|
||
|
||
if existing_cols == new_cols:
|
||
# 列名一致，TRUNCATE
|
||
cursor.execute(f"TRUNCATE TABLE {full_name}")
|
||
cursor.commit()
|
||
return f"TRUNCATE TABLE {full_name}"
|
||
else:
|
||
# 列名有变化，兜底 DROP + CREATE
|
||
return create_table(cursor, table_name, headers, cfg, actual_pk)
|
||
|
||
|
||
def create_table(cursor, table_name, headers, cfg, actual_pk):
|
||
"""动态创建表。如已存在则先 DROP。"""
|
||
schema = cfg['schema']
|
||
type_rules = cfg['type_rules']
|
||
|
||
full_name = f'[{schema}].[{table_name}]'
|
||
|
||
# DROP if exists
|
||
cursor.execute(f"IF OBJECT_ID('{full_name}', 'U') IS NOT NULL DROP TABLE {full_name}")
|
||
cursor.commit()
|
||
|
||
# 构建列定义
|
||
col_defs = ['[id] INT IDENTITY(1,1)']
|
||
for h in headers:
|
||
sql_type = infer_sql_type(h, type_rules, actual_pk)
|
||
col_defs.append(f'[{h}] {sql_type}')
|
||
|
||
# 主键约束
|
||
col_defs.append(f'CONSTRAINT [PK_{table_name}] PRIMARY KEY ([{actual_pk}])')
|
||
|
||
create_sql = f"CREATE TABLE {full_name} ({', '.join(col_defs)})"
|
||
cursor.execute(create_sql)
|
||
cursor.commit()
|
||
|
||
return create_sql
|
||
|
||
|
||
def build_typed_cols(headers, type_rules, actual_pk):
|
||
"""构建需要做类型转换的列映射 {列名: 'datetime'|'int'|'decimal'}"""
|
||
typed = {}
|
||
for h in headers:
|
||
if h == actual_pk or '20BW' in h:
|
||
continue
|
||
if any(kw in h for kw in type_rules.get('datetime', [])):
|
||
typed[h] = 'datetime'
|
||
for rule in type_rules.get('int', []):
|
||
match = (rule['name'] == h) if rule.get('exact') else (rule['name'] in h)
|
||
if match:
|
||
typed[h] = 'int'
|
||
for rule in type_rules.get('decimal', []):
|
||
match = (rule['name'] == h) if rule.get('exact') else (rule['name'] in h)
|
||
if match:
|
||
typed[h] = 'decimal'
|
||
return typed
|
||
|
||
|
||
def safe_datetime(val):
|
||
"""将值转换为 datetime，失败则返回 None"""
|
||
if val is None:
|
||
return None
|
||
if isinstance(val, str) and val.strip() == '':
|
||
return None
|
||
if isinstance(val, datetime):
|
||
if val.year < 1753:
|
||
return None
|
||
return val
|
||
if isinstance(val, str):
|
||
for fmt in ('%Y-%m-%d %H:%M:%S', '%Y-%m-%d %H:%M', '%Y-%m-%d',
|
||
'%Y/%m/%d %H:%M:%S', '%Y/%m/%d %H:%M', '%Y/%m/%d'):
|
||
try:
|
||
dt = datetime.strptime(val.strip(), fmt)
|
||
if dt.year >= 1753:
|
||
return dt
|
||
return None
|
||
except ValueError:
|
||
continue
|
||
return None
|
||
return None
|
||
|
||
|
||
def safe_int(val):
|
||
"""将值转换为 int，失败则返回 None"""
|
||
if val is None:
|
||
return None
|
||
if isinstance(val, str) and val.strip() == '':
|
||
return None
|
||
if isinstance(val, (int, float)):
|
||
return int(val)
|
||
try:
|
||
return int(float(str(val).strip()))
|
||
except (ValueError, TypeError):
|
||
return None
|
||
|
||
|
||
def safe_decimal(val):
|
||
"""将值转换为 float（pyodbc 会映射到 DECIMAL），失败则返回 None"""
|
||
if val is None:
|
||
return None
|
||
if isinstance(val, str) and val.strip() == '':
|
||
return None
|
||
if isinstance(val, (int, float)):
|
||
return float(val)
|
||
try:
|
||
return float(str(val).strip())
|
||
except (ValueError, TypeError):
|
||
return None
|
||
|
||
|
||
TYPE_CONVERTERS = {
|
||
'datetime': safe_datetime,
|
||
'int': safe_int,
|
||
'decimal': safe_decimal,
|
||
}
|
||
|
||
|
||
def batch_insert(cursor, table_name, headers, data_dict, cfg, actual_pk):
|
||
"""批量插入数据到 SQL Server。"""
|
||
schema = cfg['schema']
|
||
batch_size = cfg.get('batch_size', 5000)
|
||
type_rules = cfg['type_rules']
|
||
full_name = f'[{schema}].[{table_name}]'
|
||
|
||
typed_cols = build_typed_cols(headers, type_rules, actual_pk)
|
||
|
||
col_list = ', '.join(f'[{h}]' for h in headers)
|
||
placeholders = ', '.join(['?'] * len(headers))
|
||
insert_sql = f"INSERT INTO {full_name} ({col_list}) VALUES ({placeholders})"
|
||
|
||
# 预构建所有行
|
||
all_rows = []
|
||
for pk, record in data_dict.items():
|
||
row = []
|
||
for h in headers:
|
||
val = record.get(h)
|
||
if isinstance(val, str) and val.strip() == '':
|
||
val = None
|
||
elif h in typed_cols and val is not None:
|
||
converter = TYPE_CONVERTERS[typed_cols[h]]
|
||
val = converter(val)
|
||
row.append(val)
|
||
all_rows.append(row)
|
||
|
||
total = len(all_rows)
|
||
|
||
# 启用 fast_executemany（pyodbc 5.x 在 cursor 上设置）
|
||
cursor.fast_executemany = True
|
||
|
||
# 分批 executemany，最后统一 commit
|
||
for i in range(0, total, batch_size):
|
||
cursor.executemany(insert_sql, all_rows[i:i + batch_size])
|
||
|
||
cursor.commit()
|
||
return total
|
||
|
||
|
||
# ============================================================
|
||
# 主流程
|
||
# ============================================================
|
||
|
||
def migrate_table(cursor, table_config, cfg, dry_run=False, use_truncate=False):
|
||
"""迁移单个表"""
|
||
table_name = table_config['table']
|
||
sources = table_config['sources']
|
||
primary_key = cfg['primary_key']
|
||
excel_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'Excel')
|
||
|
||
logger.info(f"迁移 {table_name} ...")
|
||
|
||
# 1. 读取并合并数据
|
||
headers, actual_pk, merged, source_row_counts = merge_sources(
|
||
sources, excel_dir, primary_key, cfg
|
||
)
|
||
|
||
# 打印各源行数
|
||
for file_name, sheet_name, count in source_row_counts:
|
||
logger.info(f" 读取: {file_name} > {sheet_name} → {count:,} 行")
|
||
|
||
logger.info(f" 合并去重后: {len(merged):,} 行, {len(headers)} 列 (主键列: {actual_pk})")
|
||
|
||
if dry_run:
|
||
logger.info(f" [DRY-RUN] 建表 SQL:")
|
||
type_rules = cfg['type_rules']
|
||
col_defs = ['[id] INT IDENTITY(1,1)']
|
||
for h in headers:
|
||
sql_type = infer_sql_type(h, type_rules, actual_pk)
|
||
col_defs.append(f' [{h}] {sql_type}')
|
||
col_defs.append(f' CONSTRAINT [PK_{table_name}] PRIMARY KEY ([{actual_pk}])')
|
||
logger.info(f" CREATE TABLE [{cfg['schema']}].[{table_name}] (")
|
||
logger.info(',\n'.join(col_defs))
|
||
logger.info(" );")
|
||
return len(merged)
|
||
|
||
# 2. 建表
|
||
create_sql = ensure_table(cursor, table_name, headers, cfg, actual_pk, use_truncate=use_truncate)
|
||
action = "TRUNCATE" if use_truncate and "TRUNCATE" in (create_sql or "") else "建表"
|
||
logger.info(f" {action}: [{cfg['schema']}].[{table_name}] ({len(headers)} 列)")
|
||
|
||
# 3. 插入数据
|
||
total = batch_insert(cursor, table_name, headers, merged, cfg, actual_pk)
|
||
logger.info(f" 插入: {total:,} 行 done")
|
||
|
||
return total
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='Excel to SQL Server Migration')
|
||
parser.add_argument('--config', default=None, help='配置文件路径 (默认: config.yaml)')
|
||
parser.add_argument('--dry-run', action='store_true', help='仅预览，不写入数据库')
|
||
parser.add_argument('--truncate', action='store_true',
|
||
help='使用 TRUNCATE 模式（表存在时截断而非重建）')
|
||
parser.add_argument('--table', nargs='+', default=None,
|
||
help='只迁移指定的表，例如 --table PG_2022 TM_2025')
|
||
args = parser.parse_args()
|
||
|
||
# 加载配置
|
||
cfg = load_config(args.config)
|
||
|
||
# 初始化日志
|
||
from logger import setup_logger
|
||
log_cfg = cfg.get('logging', {})
|
||
setup_logger(
|
||
'app',
|
||
log_dir=log_cfg.get('log_dir'),
|
||
level=log_cfg.get('level', 'INFO'),
|
||
max_keep_days=log_cfg.get('max_keep_days', 30),
|
||
)
|
||
|
||
# 按参数筛选表
|
||
if args.table:
|
||
specified = set(args.table)
|
||
tables = [t for t in cfg['tables'] if t['table'] in specified]
|
||
missing = specified - {t['table'] for t in tables}
|
||
if missing:
|
||
logger.error(f"未找到表 {', '.join(missing)}")
|
||
available = ', '.join(t['table'] for t in cfg['tables'])
|
||
logger.info(f"可用表: {available}")
|
||
sys.exit(1)
|
||
else:
|
||
tables = cfg['tables']
|
||
|
||
total_tables = len(tables)
|
||
primary_key = cfg['primary_key']
|
||
|
||
logger.info("=" * 50)
|
||
logger.info(f"Excel → SQL Server 迁移工具")
|
||
logger.info(f" 目标: [{cfg['schema']}], {total_tables} 张表")
|
||
logger.info(f" 模式: {'DRY-RUN' if args.dry_run else 'LIVE'}")
|
||
logger.info("=" * 50)
|
||
|
||
conn = None
|
||
cursor = None
|
||
|
||
if not args.dry_run:
|
||
conn = get_connection(cfg)
|
||
cursor = conn.cursor()
|
||
create_schema(cursor, cfg['schema'])
|
||
logger.info(f"Schema [{cfg['schema']}] 已就绪")
|
||
|
||
total_rows = 0
|
||
for i, table_cfg in enumerate(tables, 1):
|
||
logger.info(f"[{i}/{total_tables}]")
|
||
rows = migrate_table(cursor, table_cfg, cfg, dry_run=args.dry_run,
|
||
use_truncate=args.truncate)
|
||
total_rows += rows
|
||
|
||
if not args.dry_run:
|
||
cursor.close()
|
||
conn.close()
|
||
|
||
logger.info("=" * 50)
|
||
logger.info(f"完成! 共迁移 {total_tables} 张表, {total_rows:,} 行")
|
||
logger.info("=" * 50)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|