BLD_sync/migration.py

import pandas as pd
import os
import shutil
import urllib
from sqlalchemy import create_engine, text
from config import DB_CONFIG, MIGRATION_TASKS, TEMP_DIR
import ntfy_utils  # 确保该文件在同一目录下

# ==========================================
# 1. 脚本配置 (Configuration)
# ==========================================

# 目标表配置
TARGET_DB_SCHEMA = "warehouseOutbound"
TARGET_TABLE_NAME = "customerProductType"
SQL_SOURCE_FILE_COL = "SourceFile"  # 你在SQL中新增的字段名

# 字段映射常量
SQL_COL_YEAR = "合同年份"
SQL_COL_WORKSHOP = "车间号"
SQL_COL_ORDER = "工令号"
SQL_COL_MODEL = "客户型号"

# 运行参数
FORCE_UPDATE = False  # 如果设为 True，则无视时间对比，强制更新所有文件

# ==========================================
# 2. 核心辅助函数
# ==========================================

def get_db_engine():
    params = urllib.parse.quote_plus(
        f"DRIVER={{{DB_CONFIG['driver']}}};"
        f"SERVER={DB_CONFIG['server']};"
        f"DATABASE={DB_CONFIG['database']};"
        f"UID={DB_CONFIG['username']};"
        f"PWD={DB_CONFIG['password']};"
        f"TrustServerCertificate=yes;"
    )
    # fast_executemany 极大提高写入速度
    return create_engine(f"mssql+pyodbc:///?odbc_connect={params}", fast_executemany=True)

def get_file_mtime(path):
    """获取文件最后修改时间戳"""
    try:
        return os.path.getmtime(path)
    except OSError:
        return 0

def delete_old_data(engine, filename):
    """根据 SourceFile 字段精确删除旧数据"""
    full_table = f"[{TARGET_DB_SCHEMA}].[{TARGET_TABLE_NAME}]"
    sql = text(f"DELETE FROM {full_table} WHERE [{SQL_SOURCE_FILE_COL}] = :fname")
    with engine.begin() as conn:
        conn.execute(sql, {"fname": filename})

# ==========================================
# 3. 迁移主逻辑
# ==========================================

def run_migration():
    # 初始化环境
    if not os.path.exists(TEMP_DIR):
        os.makedirs(TEMP_DIR)
    
    engine = get_db_engine()
    sync_count = 0
    error_count = 0

    print(f"🚀 开始增量同步任务 (强制更新={FORCE_UPDATE})")

    for task in MIGRATION_TASKS:
        remote_path = task['file_path']
        filename = os.path.basename(remote_path)
        local_path = os.path.join(TEMP_DIR, filename)
        
        # 1. 检查源文件
        if not os.path.exists(remote_path):
            msg = f"远程文件未找到: {remote_path}"
            print(f"❌ {msg}")
            ntfy_utils.send_error(msg)
            continue

        # 2. 增量判定
        remote_mtime = get_file_mtime(remote_path)
        local_mtime = get_file_mtime(local_path)

        if not FORCE_UPDATE and os.path.exists(local_path) and remote_mtime <= local_mtime:
            print(f"⏭️  跳过: {filename} (文件未变更)")
            continue

        print(f"🔄 正在处理: {filename} ...")

        try:
            # 3. 复制文件到本地 temp
            shutil.copy2(remote_path, local_path)
            
            # 4. 读取 Excel
            xls_dict = pd.read_excel(local_path, sheet_name=task['sheet_names'])
            if not isinstance(xls_dict, dict):
                xls_dict = {task['sheet_names'][0]: xls_dict}

            # 准备存放该文件所有 Sheet 的合并数据
            df_all_sheets = []

            for sheet_name, df in xls_dict.items():
                if df.empty: continue
                
                # 清洗与过滤
                df.columns = df.columns.astype(str).str.strip()
                source_cols = list(task['mapping'].keys())
                
                missing = [c for c in source_cols if c not in df.columns]
                if missing:
                    print(f"  ⚠️  Sheet[{sheet_name}] 缺失列: {missing}")
                    continue

                # 提取并重命名
                df_subset = df[source_cols].copy()
                df_subset.rename(columns=task['mapping'], inplace=True)
                
                # 注入年份和来源文件名
                df_subset[SQL_COL_YEAR] = task['year']
                df_subset[SQL_SOURCE_FILE_COL] = filename # 存入文件名，用于下次精准删除
                
                # 数据清洗
                subset_keys = [SQL_COL_YEAR, SQL_COL_WORKSHOP, SQL_COL_ORDER]
                df_subset.dropna(subset=subset_keys, inplace=True)
                df_subset.drop_duplicates(subset=subset_keys, keep='first', inplace=True)
                
                if not df_subset.empty:
                    df_all_sheets.append(df_subset)

            # 5. 写入数据库
            if df_all_sheets:
                final_df = pd.concat(df_all_sheets, ignore_index=True)
                
                # 执行删除并插入 (事务)
                with engine.begin() as conn:
                    # A. 删除旧记录
                    delete_sql = text(f"DELETE FROM [{TARGET_DB_SCHEMA}].[{TARGET_TABLE_NAME}] WHERE [{SQL_SOURCE_FILE_COL}] = :fname")
                    conn.execute(delete_sql, {"fname": filename})
                    
                    # B. 插入新记录
                    final_df.to_sql(
                        name=TARGET_TABLE_NAME,
                        schema=TARGET_DB_SCHEMA,
                        con=conn,
                        if_exists='append',
                        index=False,
                        chunksize=1000
                    )
                
                print(f"  ✅ 成功同步: {len(final_df)} 行记录")
                sync_count += 1
            else:
                print(f"  ⚠️  警告: 文件内容为空或格式不符")

        except Exception as e:
            error_msg = f"文件 [{filename}] 处理失败: {str(e)}"
            print(f"❌ {error_msg}")
            ntfy_utils.send_error(error_msg)
            error_count += 1

    # 结束汇总
    summary = f"同步完成: 成功 {sync_count} 个文件, 失败 {error_count} 个文件。"
    print(f"\n🏁 {summary}")
    if sync_count > 0:
        # 只有在有实际更新时才发送成功通知
        ntfy_utils.send_ntfy(summary, title="📊 数据迁移报告", tags=["package"])

if __name__ == "__main__":
    run_migration()