airflow_scheduler/dataops/scripts/load_file.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import logging
import sys
import os
import pandas as pd
import psycopg2
from datetime import datetime
import csv
from dags.config import PG_CONFIG

# 配置日志记录器
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)

logger = logging.getLogger("load_file")

def get_pg_conn():
    """获取PostgreSQL连接"""
    return psycopg2.connect(**PG_CONFIG)

def get_table_columns(table_name):
    """
    获取表的列信息，包括列名和注释
    
    返回:
        dict: {列名: 列注释} 的字典
    """
    conn = get_pg_conn()
    cursor = conn.cursor()
    try:
        # 查询表列信息
        cursor.execute("""
            SELECT 
                column_name, 
                col_description((table_schema || '.' || table_name)::regclass::oid, ordinal_position) as column_comment
            FROM 
                information_schema.columns
            WHERE 
                table_name = %s
            ORDER BY 
                ordinal_position
        """, (table_name,))
        
        columns = {}
        for row in cursor.fetchall():
            col_name = row[0]
            col_comment = row[1] if row[1] else col_name  # 如果注释为空，使用列名
            columns[col_name] = col_comment
            
        return columns
    except Exception as e:
        logger.error(f"获取表 {table_name} 的列信息时出错: {str(e)}")
        return {}
    finally:
        cursor.close()
        conn.close()

def match_csv_columns(csv_headers, table_columns):
    """
    匹配CSV列名与表列名
    
    策略:
    1. 尝试通过表字段注释匹配CSV列名
    2. 尝试通过名称直接匹配
    
    参数:
        csv_headers (list): CSV文件的列名列表
        table_columns (dict): {列名: 列注释} 的字典
    
    返回:
        dict: {CSV列名: 表列名} 的映射字典
    """
    mapping = {}
    
    # 通过注释匹配
    comment_to_column = {comment: col for col, comment in table_columns.items()}
    for header in csv_headers:
        if header in comment_to_column:
            mapping[header] = comment_to_column[header]
            continue
        
        # 尝试直接名称匹配
        if header in table_columns:
            mapping[header] = header
    
    return mapping

def load_csv_to_table(csv_file, table_name, execution_mode='append'):
    """
    将CSV文件数据加载到目标表
    
    参数:
        csv_file (str): CSV文件路径
        table_name (str): 目标表名
        execution_mode (str): 执行模式，'append'或'full_refresh'
    
    返回:
        bool: 成功返回True，失败返回False
    """
    conn = None
    try:
        # 读取CSV文件，尝试自动检测编码
        try:
            df = pd.read_csv(csv_file, encoding='utf-8')
        except UnicodeDecodeError:
            try:
                df = pd.read_csv(csv_file, encoding='gbk')
            except UnicodeDecodeError:
                df = pd.read_csv(csv_file, encoding='latin1')
        
        logger.info(f"成功读取CSV文件: {csv_file}, 共 {len(df)} 行")
        
        # 获取CSV列名
        csv_headers = df.columns.tolist()
        logger.info(f"CSV列名: {csv_headers}")
        
        # 获取表结构
        table_columns = get_table_columns(table_name)
        if not table_columns:
            logger.error(f"无法获取表 {table_name} 的列信息")
            return False
        
        logger.info(f"表 {table_name} 的列信息: {table_columns}")
        
        # 匹配CSV列与表列
        column_mapping = match_csv_columns(csv_headers, table_columns)
        logger.info(f"列映射关系: {column_mapping}")
        
        if not column_mapping:
            logger.error(f"无法建立CSV列与表列的映射关系")
            return False
        
        # 筛选和重命名列
        df_mapped = df[list(column_mapping.keys())].rename(columns=column_mapping)
        
        # 连接数据库
        conn = get_pg_conn()
        cursor = conn.cursor()
        
        # 根据执行模式确定操作
        if execution_mode == 'full_refresh':
            # 如果是全量刷新，先清空表
            logger.info(f"执行全量刷新，清空表 {table_name}")
            cursor.execute(f"TRUNCATE TABLE {table_name}")
        
        # 构建INSERT语句
        columns = ', '.join(df_mapped.columns)
        placeholders = ', '.join(['%s'] * len(df_mapped.columns))
        insert_sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
        
        # 批量插入数据
        rows = [tuple(row) for row in df_mapped.values]
        cursor.executemany(insert_sql, rows)
        
        # 提交事务
        conn.commit()
        logger.info(f"成功插入 {len(rows)} 行数据到表 {table_name}")
        
        return True
    except Exception as e:
        logger.error(f"加载CSV数据到表时出错: {str(e)}")
        if conn:
            conn.rollback()
        return False
    finally:
        if conn:
            conn.close()

def run(table_name, execution_mode='append', exec_date=None, target_type=None, 
        storage_location=None, frequency=None, **kwargs):
    """
    统一入口函数，符合Airflow动态脚本调用规范
    
    参数:
        table_name (str): 要处理的表名
        execution_mode (str): 执行模式 (append/full_refresh)
        exec_date: 执行日期
        target_type: 目标类型，对于CSV文件应为'structure'
        storage_location: CSV文件路径
        frequency: 更新频率
        **kwargs: 其他可能的参数
    
    返回:
        bool: 执行成功返回True，否则返回False
    """
    logger.info(f"===== 开始执行CSV文件加载 =====")
    logger.info(f"表名: {table_name}")
    logger.info(f"执行模式: {execution_mode}")
    logger.info(f"执行日期: {exec_date}")
    logger.info(f"目标类型: {target_type}")
    logger.info(f"文件路径: {storage_location}")
    logger.info(f"更新频率: {frequency}")
    
    # 记录其他参数
    for key, value in kwargs.items():
        logger.info(f"其他参数 - {key}: {value}")
    
    # 检查必要参数
    if not storage_location:
        logger.error("未提供CSV文件路径")
        return False
    
    # 检查文件是否存在
    if not os.path.exists(storage_location):
        logger.error(f"CSV文件不存在: {storage_location}")
        return False
    
    # 记录执行开始时间
    start_time = datetime.now()
    
    try:
        # 加载CSV数据到表
        result = load_csv_to_table(storage_location, table_name, execution_mode)
        
        # 记录执行结束时间
        end_time = datetime.now()
        duration = (end_time - start_time).total_seconds()
        
        if result:
            logger.info(f"CSV文件加载成功，耗时: {duration:.2f}秒")
        else:
            logger.error(f"CSV文件加载失败，耗时: {duration:.2f}秒")
        
        return result
    except Exception as e:
        # 记录执行结束时间
        end_time = datetime.now()
        duration = (end_time - start_time).total_seconds()
        
        logger.error(f"CSV文件加载过程中出错: {str(e)}")
        logger.error(f"CSV文件加载失败，耗时: {duration:.2f}秒")
        
        return False
    finally:
        logger.info(f"===== CSV文件加载执行完成 =====")

if __name__ == "__main__":
    # 直接执行时的测试代码
    import argparse
    
    parser = argparse.ArgumentParser(description='从CSV文件加载数据到表')
    parser.add_argument('--table', type=str, required=True, help='目标表名')
    parser.add_argument('--file', type=str, required=True, help='CSV文件路径')
    parser.add_argument('--mode', type=str, default='append', help='执行模式: append或full_refresh')
    
    args = parser.parse_args()
    
    success = run(
        table_name=args.table,
        execution_mode=args.mode,
        storage_location=args.file,
        target_type='structure'
    )
    
    if success:
        print("CSV文件加载成功")
        sys.exit(0)
    else:
        print("CSV文件加载失败")
        sys.exit(1)