airflow_scheduler/dataops_scripts/execution_sql.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import logging
import sys
import os
from datetime import datetime
import psycopg2
import jinja2
# 修改导入方式，避免使用相对导入
import sys
import os
import traceback
import time
from datetime import datetime, timedelta
import pytz
import re

from script_utils import get_config_param
SCHEDULE_TABLE_SCHEMA = get_config_param("SCHEDULE_TABLE_SCHEMA")

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

logger = logging.getLogger("execution_sql")

# 将同级目录加入到Python搜索路径
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.append(current_dir)

# 尝试导入script_utils
try:
    import script_utils
    from script_utils import get_pg_config, logger as utils_logger
    logger.info("成功导入script_utils模块的get_pg_config方法")
except ImportError as e:
    logger.error(f"无法直接导入script_utils方法: {str(e)}")
    
    # 尝试备用方法1：完整路径导入
    # try:
    #     sys.path.append(os.path.dirname(current_dir))  # 添加父目录
    #     from dataops_scripts.script_utils import get_pg_config
    #     logger.info("使用完整路径成功导入script_utils模块的方法")
    # except ImportError as e2:
    #     logger.error(f"使用完整路径导入失败: {str(e2)}")
        
    #     # 尝试备用方法2：动态导入
    #     try:
    #         import importlib.util
    #         script_utils_path = os.path.join(current_dir, "script_utils.py")
    #         logger.info(f"尝试从路径动态导入: {script_utils_path}")
            
    #         spec = importlib.util.spec_from_file_location("script_utils", script_utils_path)
    #         script_utils = importlib.util.module_from_spec(spec)
    #         spec.loader.exec_module(script_utils)
    #         get_pg_config = script_utils.get_pg_config
    #         logger.info("通过动态导入成功加载script_utils模块的方法")
    #     except Exception as e3:
    #         logger.error(f"动态导入也失败: {str(e3)}")
    #         raise ImportError(f"无法导入script_utils模块的方法，所有方法都失败")

# 使用script_utils中的方法获取配置
try:
    # 获取PostgreSQL配置
    PG_CONFIG = get_pg_config()
    logger.info(f"通过script_utils.get_pg_config()获取PostgreSQL配置成功")
except Exception as e:
    logger.error(f"获取配置失败，使用默认值: {str(e)}")
    # 默认配置
    PG_CONFIG = {
        "host": "localhost",
        "port": 5432,
        "user": "postgres",
        "password": "postgres",
        "database": "dataops",
    }

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

def get_script_content(target_table, script_name):
    """
    从data_transform_scripts表中获取脚本内容和目标日期列
    
    参数:
        target_table (str): 目标表名
        script_name (str): 脚本名称
    
    返回:
        tuple: (script_content, target_dt_column) 脚本内容和目标日期列
    """
    logger.info(f"正在从data_transform_scripts表获取脚本内容，目标表: {target_table}，脚本名称: {script_name}")
    conn = None
    cursor = None
    try:
        conn = get_pg_conn()
        cursor = conn.cursor()
        
        query = """
            SELECT script_content, target_dt_column 
            FROM {}.data_transform_scripts 
            WHERE target_table = %s AND script_name = %s 
            LIMIT 1
        """.format(SCHEDULE_TABLE_SCHEMA)
        
        logger.info(f"执行SQL查询: {query}")
        logger.info(f"查询参数: target_table={target_table}, script_name={script_name}")
        
        cursor.execute(query, (target_table, script_name))
        result = cursor.fetchone()
        
        if result is None:
            logger.error(f"未找到目标表 '{target_table}' 和脚本名 '{script_name}' 对应的脚本")
            return None, None
        
        # 获取脚本内容和目标日期列    
        script_content = result[0]
        target_dt_column = result[1] if len(result) > 1 else None
        
        # 记录结果
        logger.info(f"目标日期列: {target_dt_column if target_dt_column else '未设置'}")
        
        # 记录脚本内容，但可能很长，只记录前500个字符和后100个字符
        if len(script_content) > 600:
            logger.info(f"成功获取脚本内容，总长度: {len(script_content)}字符")
            logger.info(f"脚本内容前500字符: \n{script_content[:500]}")
            logger.info(f"脚本内容后100字符: \n{script_content[-100:]}")
        else:
            logger.info(f"成功获取脚本内容，内容如下: \n{script_content}")
            
        return script_content, target_dt_column
        
    except Exception as e:
        logger.error(f"获取脚本内容时出错: {str(e)}", exc_info=True)
        return None, None
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()

def execute_sql(sql, params=None):
    """
    执行SQL语句
    
    参数:
        sql (str): SQL语句
        params (dict, optional): SQL参数
    
    返回:
        tuple: (成功标志, 影响的行数或结果)
    """
    conn = None
    cursor = None
    try:
        conn = get_pg_conn()
        cursor = conn.cursor()
        
        # 记录SQL（不包含敏感参数）
        # 由于SQL可能很长，只记录前200个字符
        if len(sql) > 200:
            logger.info(f"执行SQL (前200字符): {sql[:200]}...")
        else:
            logger.info(f"执行SQL: {sql}")
            
        if params:
            logger.info(f"SQL参数: {params}")
        
        # 执行SQL
        cursor.execute(sql, params)
        
        # 获取影响的行数
        if cursor.rowcount >= 0:
            affected_rows = cursor.rowcount
            logger.info(f"SQL执行成功，影响了 {affected_rows} 行")
        else:
            affected_rows = 0
            logger.info("SQL执行成功，但无法确定影响的行数")
        
        # 如果是SELECT语句，获取结果
        if sql.strip().upper().startswith("SELECT"):
            result = cursor.fetchall()
            logger.info(f"查询返回 {len(result)} 行结果")
            conn.commit()
            return True, {"affected_rows": affected_rows, "result": result}
        else:
            # 对于非SELECT语句，提交事务
            conn.commit()
            return True, {"affected_rows": affected_rows}
            
    except Exception as e:
        logger.error(f"执行SQL时出错: {str(e)}", exc_info=True)
        if conn:
            conn.rollback()
        return False, {"error": str(e)}
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()

def render_sql_template(sql_content, template_params):
    """
    使用Jinja2渲染SQL模板
    
    参数:
        sql_content (str): SQL模板内容
        template_params (dict): 模板参数
    
    返回:
        str: 渲染后的SQL
    """
    try:
        # 创建Jinja2环境
        env = jinja2.Environment(
            autoescape=False,
            undefined=jinja2.StrictUndefined  # 严格模式，未定义变量会抛出异常
        )
        
        # 创建模板并渲染
        template = env.from_string(sql_content)
        rendered_sql = template.render(**template_params)
        
        logger.info("SQL模板渲染成功")
        
        # 记录渲染后的SQL，只记录前200个字符和后100个字符
        if len(rendered_sql) > 300:
            logger.info(f"渲染后SQL前200字符: \n{rendered_sql[:200]}...")
            logger.info(f"渲染后SQL后100字符: \n...{rendered_sql[-100:]}")
        else:
            logger.info(f"渲染后SQL: \n{rendered_sql}")
            
        return rendered_sql
    except Exception as e:
        logger.error(f"渲染SQL模板时出错: {str(e)}", exc_info=True)
        raise

def run(script_type=None, target_table=None, script_name=None, exec_date=None, 
        schedule_frequency=None, **kwargs):
    """
    执行SQL脚本主入口函数
    
    参数:
        script_type (str): 脚本类型，必须为'sql'
        target_table (str): 目标表名
        script_name (str): 脚本名称
        exec_date (str): 执行日期，格式为YYYY-MM-DD
        schedule_frequency (str): 频率，可选值为 daily, weekly, monthly, quarterly, yearly
        **kwargs: 其他参数
    
    返回:
        bool: 是否执行成功
    """
    # 记录开始执行的时间
    start_time = datetime.now()
    logger.info(f"===== 开始执行 SQL 脚本 =====")
    logger.info(f"脚本类型: {script_type}")
    logger.info(f"目标表: {target_table}")
    logger.info(f"脚本名称: {script_name}")
    logger.info(f"执行日期: {exec_date}")
    logger.info(f"频率: {schedule_frequency}")
    
    # 记录其他参数
    for key, value in kwargs.items():
        logger.info(f"其他参数 - {key}: {value}")
    
    # 验证必要参数
    if not script_type or script_type.lower() != 'sql':
        logger.error(f"脚本类型必须为'sql'，当前为: {script_type}")
        return False
    
    if not target_table:
        logger.error("未提供目标表名")
        return False
    
    if not script_name:
        logger.error("未提供脚本名称")
        return False
    
    if not exec_date:
        logger.error("未提供执行日期")
        return False
    
    if not schedule_frequency:
        logger.error("未提供频率")
        return False
    
    try:
        # 获取脚本内容，直接从数据库获取，不检查文件是否存在
        sql_content, target_dt_column = get_script_content(target_table, script_name)
        if not sql_content:
            logger.error("无法获取脚本内容，执行失败")
            return False
        
        logger.info(f"成功获取脚本内容，长度: {len(sql_content)} 字符")
        
        # 检查是否开启ETL幂等性
        target_table_label = kwargs.get('target_table_label', '')
        script_exec_mode = kwargs.get('update_mode', 'append')  # 只使用update_mode
        is_manual_dag_trigger = kwargs.get('is_manual_dag_trigger', False)  # 新增：获取是否为手动触发的DAG
        
        logger.info(f"脚本更新模式: {script_exec_mode}")
        logger.info(f"是否手工DAG触发: {is_manual_dag_trigger}")
        
        # 导入config模块获取幂等性开关
        try:
            config = __import__('config')
            enable_idempotency = getattr(config, 'ENABLE_ETL_IDEMPOTENCY', False)
        except ImportError:
            logger.warning("无法导入config模块获取幂等性开关，默认为False")
            enable_idempotency = False
        
        logger.info(f"ETL幂等性开关状态: {enable_idempotency}")
        logger.info(f"目标表标签: {target_table_label}")
        
        # 如果开启了ETL幂等性且是SQL类型
        if script_type.lower() == 'sql' and enable_idempotency:
            
            # 处理append模式
            if script_exec_mode.lower() == 'append':
                logger.info("当前为append模式，开始考虑ETL幂等性处理")
                
                if is_manual_dag_trigger:
                    logger.info("manual dag被手工触发")
                    # 计算日期范围
                    try:
                        start_date, end_date = script_utils.get_date_range(exec_date, schedule_frequency)
                        logger.info(f"计算得到的日期范围: start_date={start_date}, end_date={end_date}")
                    except Exception as date_err:
                        logger.error(f"计算日期范围时出错: {str(date_err)}", exc_info=True)
                        return False
                    
                    # 检查是否有目标日期列
                    if target_dt_column:
                        logger.info(f"找到目标日期列 {target_dt_column}，将生成DELETE语句")
                        
                        # 生成DELETE语句
                        delete_sql = f"""DELETE FROM {target_table}
WHERE {target_dt_column} >= '{{{{ start_date }}}}'
  AND {target_dt_column} < '{{{{ end_date }}}}';"""
                        
                        logger.info(f"生成的DELETE语句: {delete_sql}")
                        
                        # 准备模板参数
                        template_params = {
                            'start_date': start_date,
                            'end_date': end_date,
                            # 可以添加更多默认参数
                        }
                        
                        # 渲染DELETE SQL
                        try:
                            rendered_delete_sql = render_sql_template(delete_sql, template_params)
                            logger.info("成功渲染清理SQL")
                        except Exception as render_err:
                            logger.error(f"渲染清理SQL时出错: {str(render_err)}", exc_info=True)
                            # 即使清理SQL失败，仍然继续执行后续SQL
                            logger.warning("继续执行原始SQL")
                        else:
                            # 执行DELETE SQL
                            logger.info("执行清理SQL以实现幂等性")
                            delete_success, delete_result = execute_sql(rendered_delete_sql)
                            
                            if delete_success:
                                if isinstance(delete_result, dict) and "affected_rows" in delete_result:
                                    logger.info(f"清理SQL执行成功，删除了 {delete_result['affected_rows']} 行数据")
                                else:
                                    logger.info("清理SQL执行成功")
                            else:
                                logger.error(f"清理SQL执行失败: {delete_result.get('error', '未知错误')}")
                                # 继续执行原始SQL
                                logger.warning("继续执行原始SQL")
                    else:
                        logger.warning(f"目标表 {target_table} 没有设置目标日期列(target_dt_column)，无法生成DELETE语句实现幂等性")
                        logger.warning("将直接执行原始SQL，可能导致数据重复")
                else:
                    logger.info(f"不是manual dag手工触发，对目标表 {target_table} create_time 进行数据清理")
                    
                    # 获取一天范围的日期
                    start_date, end_date = script_utils.get_one_day_range(exec_date)
                    logger.info(f"使用一天范围的日期: start_date={start_date}, end_date={end_date}")
                    
                    # 生成基于create_time的DELETE语句
                    delete_sql = f"""DELETE FROM {target_table}
WHERE create_time >= '{start_date}'
  AND create_time < '{end_date}';"""
                    
                    logger.info(f"生成的基于create_time的DELETE语句: {delete_sql}")
                    
                    # 执行DELETE SQL
                    logger.info("基于create_time执行清理SQL以实现幂等性")
                    delete_success, delete_result = execute_sql(delete_sql)
                    
                    if delete_success:
                        if isinstance(delete_result, dict) and "affected_rows" in delete_result:
                            logger.info(f"基于create_time,清理SQL执行成功，删除了 {delete_result['affected_rows']} 行数据")
                        else:
                            logger.info("基于create_time,清理SQL执行成功")
                    else:
                        logger.error(f"基于create_time,清理SQL执行失败: {delete_result.get('error', '未知错误')}")
                        # 继续执行原始SQL
                        logger.warning("继续执行原始SQL")
            
            # 处理full_refresh模式
            elif script_exec_mode.lower() == 'full_refresh':
                logger.info("当前为full_refresh模式，将执行TRUNCATE操作")
                
                # 构建TRUNCATE语句
                truncate_sql = f"TRUNCATE TABLE {target_table};"
                logger.info(f"生成的TRUNCATE SQL: {truncate_sql}")
                
                # 执行TRUNCATE操作
                truncate_success, truncate_result = execute_sql(truncate_sql)
                
                if truncate_success:
                    logger.info(f"TRUNCATE TABLE {target_table} 执行成功，表已清空")
                else:
                    error_msg = truncate_result.get("error", "未知错误")
                    logger.error(f"TRUNCATE TABLE执行失败: {error_msg}")
                    # 继续执行原始SQL
                    logger.warning("TRUNCATE失败，继续执行原始SQL")
            
            else:
                logger.info(f"当前更新模式 {script_exec_mode} 不是append或full_refresh，不执行幂等性处理")
        else:
            logger.info("未满足ETL幂等性处理条件，直接执行原始SQL")
        
        # 计算日期范围(如果之前未计算)
        if 'start_date' not in locals() or 'end_date' not in locals():
            try:
                start_date, end_date = script_utils.get_date_range(exec_date, schedule_frequency)
                logger.info(f"计算得到的日期范围: start_date={start_date}, end_date={end_date}")
            except Exception as date_err:
                logger.error(f"计算日期范围时出错: {str(date_err)}", exc_info=True)
                return False
        
        # 准备模板参数(如果之前未准备)
        if 'template_params' not in locals():
            template_params = {
                'start_date': start_date,
                'end_date': end_date,
                # 可以添加更多默认参数
            }
        
        # 渲染原始SQL模板
        try:
            rendered_sql = render_sql_template(sql_content, template_params)
        except Exception as render_err:
            logger.error(f"渲染SQL模板时出错: {str(render_err)}", exc_info=True)
            return False
        
        # 执行原始SQL
        success, result = execute_sql(rendered_sql)
        
        if success:
            # 记录执行结果
            if isinstance(result, dict):
                if "affected_rows" in result:
                    logger.info(f"SQL执行成功，影响了 {result['affected_rows']} 行数据")
                if "result" in result:
                    logger.info(f"SQL查询返回 {len(result['result'])} 行结果")
            else:
                logger.info("SQL执行成功")
            
            # 记录执行时间
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()
            logger.info(f"===== SQL脚本执行完成 (成功) =====")
            logger.info(f"总耗时: {duration:.2f}秒")
            
            return True
        else:
            # 记录错误信息
            error_msg = result.get("error", "未知错误")
            logger.error(f"SQL执行失败: {error_msg}")
            
            # 记录执行时间
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()
            logger.info(f"===== SQL脚本执行完成 (失败) =====")
            logger.info(f"总耗时: {duration:.2f}秒")
            
            return False
            
    except Exception as e:
        # 捕获所有未处理的异常
        logger.error(f"执行SQL脚本时发生未预期的错误: {str(e)}", exc_info=True)
        return False

if __name__ == "__main__":
    # 直接执行时的测试代码
    import argparse
    
    parser = argparse.ArgumentParser(description='执行SQL脚本')
    parser.add_argument('--target-table', type=str, required=True, help='目标表名')
    parser.add_argument('--script-name', type=str, required=True, help='脚本名称')
    parser.add_argument('--exec-date', type=str, required=True, help='执行日期 (YYYY-MM-DD)')
    parser.add_argument('--schedule_frequency', type=str, required=True, 
                        choices=['daily', 'weekly', 'monthly', 'quarterly', 'yearly'], 
                        help='频率: daily, weekly, monthly, quarterly, yearly')
    
    args = parser.parse_args()
    
    # 构造必要的 kwargs
    run_kwargs = {
        "script_type": "sql",
        "target_table": args.target_table,
        "script_name": args.script_name,
        "exec_date": args.exec_date,
        "schedule_frequency": args.schedule_frequency,
    }

    logger.info("命令行测试执行参数: " + str(run_kwargs))

    success = run(**run_kwargs)
    
    if success:
        logger.info("SQL脚本执行成功")
        sys.exit(0)
    else:
        logger.error("SQL脚本执行失败")
        sys.exit(1)