|
|
@@ -0,0 +1,266 @@
|
|
|
+#!/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)
|
wangxq