airflow_scheduler/dags/dag_dataops_unified_data_scheduler.py

# dag_dataops_unified_data_scheduler.py
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.empty import EmptyOperator
from airflow.sensors.external_task import ExternalTaskSensor
from datetime import datetime, timedelta, date
import logging
import networkx as nx
import json
import os
from common import (
    get_pg_conn, 
    get_neo4j_driver,
    execute_with_monitoring,
    get_today_date
)
from config import TASK_RETRY_CONFIG

# 创建日志记录器
logger = logging.getLogger(__name__)

def get_latest_date():
    """获取数据库中包含记录的最近日期"""
    conn = get_pg_conn()
    cursor = conn.cursor()
    try:
        cursor.execute("""
            SELECT DISTINCT exec_date
            FROM airflow_dag_schedule 
            ORDER BY exec_date DESC
            LIMIT 1
        """)
        result = cursor.fetchone()
        if result:
            latest_date = result[0]
            logger.info(f"找到最近的包含记录的日期: {latest_date}")
            return latest_date
        else:
            logger.warning("未找到包含记录的日期，将使用当前日期")
            return get_today_date()
    except Exception as e:
        logger.error(f"查找最近日期时出错: {str(e)}")
        return get_today_date()
    finally:
        cursor.close()
        conn.close()

def get_all_tasks(exec_date):
    """获取所有需要执行的任务（DataResource和DataModel）"""
    conn = get_pg_conn()
    cursor = conn.cursor()
    try:
        # 查询数据表中记录总数
        cursor.execute("""
            SELECT COUNT(*) 
            FROM airflow_dag_schedule 
            WHERE exec_date = %s
        """, (exec_date,))
        total_count = cursor.fetchone()[0]
        logger.info(f"执行日期 {exec_date} 在airflow_dag_schedule表中共有 {total_count} 条记录")
        
        # 查询所有资源表任务
        cursor.execute("""
            SELECT source_table, target_table, target_table_label, script_name, script_exec_mode
            FROM airflow_dag_schedule 
            WHERE exec_date = %s AND target_table_label = 'DataResource' AND script_name IS NOT NULL
        """, (exec_date,))
        resource_results = cursor.fetchall()
        logger.info(f"查询到 {len(resource_results)} 条DataResource记录")
        
        # 查询所有模型表任务
        cursor.execute("""
            SELECT source_table, target_table, target_table_label, script_name, script_exec_mode
            FROM airflow_dag_schedule 
            WHERE exec_date = %s AND target_table_label = 'DataModel' AND script_name IS NOT NULL
        """, (exec_date,))
        model_results = cursor.fetchall()
        logger.info(f"查询到 {len(model_results)} 条DataModel记录")
        
        # 整理资源表信息
        resource_tasks = []
        for row in resource_results:
            source_table, target_table, target_table_label, script_name, script_exec_mode = row
            if script_name:  # 确保脚本名称不为空
                resource_tasks.append({
                    "source_table": source_table,
                    "target_table": target_table,
                    "target_table_label": target_table_label,
                    "script_name": script_name,
                    "script_exec_mode": script_exec_mode or "append"
                })
        
        # 整理模型表信息
        model_tasks = []
        for row in model_results:
            source_table, target_table, target_table_label, script_name, script_exec_mode = row
            if script_name:  # 确保脚本名称不为空
                model_tasks.append({
                    "source_table": source_table,
                    "target_table": target_table,
                    "target_table_label": target_table_label,
                    "script_name": script_name,
                    "script_exec_mode": script_exec_mode or "append"
                })
        
        logger.info(f"获取到 {len(resource_tasks)} 个资源表任务和 {len(model_tasks)} 个模型表任务")
        return resource_tasks, model_tasks
    except Exception as e:
        logger.error(f"获取任务信息时出错: {str(e)}")
        return [], []
    finally:
        cursor.close()
        conn.close()

def get_table_dependencies(table_names):
    """获取表之间的依赖关系"""
    driver = get_neo4j_driver()
    dependency_dict = {name: [] for name in table_names}
    
    try:
        with driver.session() as session:
            # 获取所有模型表之间的依赖关系
            query = """
                MATCH (source:DataModel)-[:DERIVED_FROM]->(target)
                WHERE source.en_name IN $table_names
                RETURN source.en_name AS source, target.en_name AS target, labels(target) AS target_labels
            """
            result = session.run(query, table_names=table_names)
            
            for record in result:
                source = record.get("source")
                target = record.get("target")
                target_labels = record.get("target_labels", [])
                
                if source and target:
                    # 将目标表添加到源表的依赖列表中
                    dependency_dict[source].append({
                        "table_name": target,
                        "table_type": next((label for label in target_labels if label in ["DataModel", "DataResource"]), None)
                    })
                    logger.debug(f"依赖关系: {source} 依赖于 {target}")
    except Exception as e:
        logger.error(f"从Neo4j获取依赖关系时出错: {str(e)}")
    finally:
        driver.close()
    
    return dependency_dict

def json_serial(obj):
    """将日期对象序列化为ISO格式字符串的JSON序列化器"""
    if isinstance(obj, (datetime, date)):
        return obj.isoformat()
    raise TypeError(f"类型 {type(obj)} 不能被序列化为JSON")

def prepare_unified_execution_plan(**kwargs):
    """准备统一执行计划的主函数"""
    # 获取执行日期
    exec_date = get_latest_date()
    logger.info(f"使用执行日期: {exec_date}")
    
    # 获取所有任务
    resource_tasks, model_tasks = get_all_tasks(exec_date)
    
    if not resource_tasks and not model_tasks:
        logger.warning(f"执行日期 {exec_date} 没有找到任务")
        return 0
    
    # 为所有模型表获取依赖关系
    model_table_names = [task["target_table"] for task in model_tasks]
    dependencies = get_table_dependencies(model_table_names)
    
    # 创建执行计划
    execution_plan = {
        "exec_date": exec_date,
        "resource_tasks": resource_tasks,
        "model_tasks": model_tasks,
        "dependencies": dependencies
    }
    
    # 记录资源任务和模型任务的名称，便于调试
    resource_names = [task["target_table"] for task in resource_tasks]
    model_names = [task["target_table"] for task in model_tasks]
    logger.info(f"资源表任务: {resource_names}")
    logger.info(f"模型表任务: {model_names}")
    
    # 将执行计划保存到XCom，使用自定义序列化器处理日期对象
    try:
        kwargs['ti'].xcom_push(key='execution_plan', value=json.dumps(execution_plan, default=json_serial))
        logger.info(f"准备了执行计划，包含 {len(resource_tasks)} 个资源表任务和 {len(model_tasks)} 个模型表任务")
    except Exception as e:
        logger.error(f"将执行计划保存到XCom时出错: {str(e)}")
        
        # 保存执行计划到文件以备后用
        try:
            plan_path = os.path.join(os.path.dirname(__file__), 'last_execution_plan.json')
            with open(plan_path, 'w') as f:
                json.dump(execution_plan, f, default=json_serial)
            logger.info(f"将执行计划保存到文件: {plan_path}")
        except Exception as file_e:
            logger.error(f"保存执行计划到文件时出错: {str(file_e)}")
    
    return len(resource_tasks) + len(model_tasks)

def process_resource(target_table, script_name, script_exec_mode, exec_date):
    """处理单个资源表"""
    logger.info(f"执行资源表 {target_table} 的脚本 {script_name}")
    # 检查exec_date是否是JSON字符串
    if isinstance(exec_date, str) and exec_date.startswith('{'):
        try:
            # 尝试解析JSON字符串
            exec_date_data = json.loads(exec_date)
            exec_date = exec_date_data.get("exec_date")
            logger.info(f"从JSON中提取执行日期: {exec_date}")
        except Exception as e:
            logger.error(f"解析exec_date JSON时出错: {str(e)}")
    
    try:
        # 正常调用执行监控函数
        result = execute_with_monitoring(
            target_table=target_table,
            script_name=script_name,
            script_exec_mode=script_exec_mode,
            exec_date=exec_date
        )
        logger.info(f"资源表 {target_table} 处理完成，结果: {result}")
        return result
    except Exception as e:
        logger.error(f"处理资源表 {target_table} 时出错: {str(e)}")
        # 确保即使出错也返回结果，不会阻塞DAG
        return False

def process_model(target_table, script_name, script_exec_mode, exec_date):
    """处理单个模型表"""
    logger.info(f"执行模型表 {target_table} 的脚本 {script_name}")
    # 检查exec_date是否是JSON字符串
    if isinstance(exec_date, str) and exec_date.startswith('{'):
        try:
            # 尝试解析JSON字符串
            exec_date_data = json.loads(exec_date)
            exec_date = exec_date_data.get("exec_date")
            logger.info(f"从JSON中提取执行日期: {exec_date}")
        except Exception as e:
            logger.error(f"解析exec_date JSON时出错: {str(e)}")
    
    try:
        result = execute_with_monitoring(
            target_table=target_table,
            script_name=script_name,
            script_exec_mode=script_exec_mode,
            exec_date=exec_date
        )
        logger.info(f"模型表 {target_table} 处理完成，结果: {result}")
        return result
    except Exception as e:
        logger.error(f"处理模型表 {target_table} 时出错: {str(e)}")
        # 确保即使出错也返回结果，不会阻塞DAG
        return False

# 预先加载数据以创建任务
try:
    logger.info("预先加载执行计划数据用于构建DAG")
    plan_path = os.path.join(os.path.dirname(__file__), 'last_execution_plan.json')
    execution_plan = {"exec_date": get_today_date(), "resource_tasks": [], "model_tasks": [], "dependencies": {}}
    
    if os.path.exists(plan_path):
        try:
            with open(plan_path, 'r') as f:
                execution_plan_json = f.read()
                execution_plan = json.loads(execution_plan_json)
                logger.info(f"从文件加载执行计划: {plan_path}")
        except Exception as e:
            logger.warning(f"读取执行计划文件出错: {str(e)}")
    else:
        logger.warning(f"执行计划文件不存在: {plan_path}，将创建基础DAG结构")
        
    # 提取信息
    exec_date = execution_plan.get("exec_date", get_today_date())
    resource_tasks = execution_plan.get("resource_tasks", [])
    model_tasks = execution_plan.get("model_tasks", [])
    dependencies = execution_plan.get("dependencies", {})
    
    logger.info(f"预加载执行计划: exec_date={exec_date}, resource_tasks数量={len(resource_tasks)}, model_tasks数量={len(model_tasks)}")
except Exception as e:
    logger.error(f"预加载执行计划数据时出错: {str(e)}")
    exec_date = get_today_date()
    resource_tasks = []
    model_tasks = []
    dependencies = {}

# 定义处理DAG失败的回调函数
def handle_dag_failure(context):
    logger.error(f"DAG执行失败: {context.get('exception')}")

# 创建DAG
with DAG(
    "dag_dataops_unified_data_scheduler", 
    start_date=datetime(2024, 1, 1), 
    schedule_interval="@daily", 
    catchup=False,
    default_args={
        'owner': 'airflow',
        'depends_on_past': False,
        'email_on_failure': False,
        'email_on_retry': False,
        'retries': 1,
        'retry_delay': timedelta(minutes=5)
    },
    on_failure_callback=handle_dag_failure  # 在这里设置回调函数
) as dag:
    
    # 等待准备DAG完成
    wait_for_prepare = ExternalTaskSensor(
        task_id="wait_for_prepare",
        external_dag_id="dag_dataops_unified_prepare_scheduler",
        external_task_id="preparation_completed",
        mode="reschedule",  # 改为reschedule模式，减少资源占用
        timeout=3600,
        poke_interval=30,
        execution_timeout=timedelta(hours=1),
        soft_fail=True,
        allowed_states=["success", "skipped"],
        dag=dag
    )
    
    # 准备执行计划
    prepare_plan = PythonOperator(
        task_id="prepare_execution_plan",
        python_callable=prepare_unified_execution_plan,
        provide_context=True,
        dag=dag
    )
    
    # 处理完成标记
    processing_completed = EmptyOperator(
        task_id="processing_completed",
        trigger_rule="none_failed_min_one_success",  # 只要有一个任务成功且没有失败的任务就标记为完成
        dag=dag
    )
    
    # 设置初始任务依赖
    wait_for_prepare >> prepare_plan
    
    # 任务字典，用于设置依赖关系
    task_dict = {}
    
    # 1. 预先创建资源表任务
    for task_info in resource_tasks:
        table_name = task_info["target_table"]
        script_name = task_info["script_name"]
        exec_mode = task_info.get("script_exec_mode", "append")
        
        # 创建安全的任务ID
        safe_table_name = table_name.replace(".", "_").replace("-", "_")
        task_id = f"resource_{safe_table_name}"
        
        resource_task = PythonOperator(
            task_id=task_id,
            python_callable=process_resource,
            op_kwargs={
                "target_table": table_name,
                "script_name": script_name,
                "script_exec_mode": exec_mode,
                "exec_date": exec_date  # 直接使用解析出的exec_date，不使用XCom
            },
            retries=TASK_RETRY_CONFIG["retries"],
            retry_delay=timedelta(minutes=TASK_RETRY_CONFIG["retry_delay_minutes"]),
            dag=dag
        )
        
        # 将任务添加到字典
        task_dict[table_name] = resource_task
        
        # 设置与prepare_plan的依赖 - 直接依赖，不需要其他条件
        prepare_plan >> resource_task
        logger.info(f"预先设置基本依赖: prepare_plan >> {task_id}")
    
    # 创建有向图，用于检测模型表之间的依赖关系
    G = nx.DiGraph()
    
    # 将所有模型表添加为节点
    for task_info in model_tasks:
        table_name = task_info["target_table"]
        G.add_node(table_name)
    
    # 添加模型表之间的依赖边
    for source, deps in dependencies.items():
        for dep in deps:
            if dep.get("table_type") == "DataModel" and dep.get("table_name") in G.nodes():
                G.add_edge(dep.get("table_name"), source)  # 依赖方向：依赖项 -> 目标
    
    # 检测循环依赖并处理
    cycles = list(nx.simple_cycles(G))
    if cycles:
        logger.warning(f"检测到循环依赖: {cycles}")
        for cycle in cycles:
            G.remove_edge(cycle[-1], cycle[0])
            logger.info(f"打破循环依赖: 移除 {cycle[-1]} -> {cycle[0]} 的依赖")
    
    # 生成拓扑排序，确定执行顺序
    execution_order = []
    try:
        execution_order = list(nx.topological_sort(G))
        logger.info(f"预加载计算的执行顺序: {execution_order}")
    except Exception as e:
        logger.error(f"生成拓扑排序失败: {str(e)}, 使用原始顺序")
        execution_order = [task_info["target_table"] for task_info in model_tasks]
    
    # 2. 预先创建模型表任务
    for table_name in execution_order:
        task_info = next((t for t in model_tasks if t["target_table"] == table_name), None)
        if not task_info:
            continue
            
        script_name = task_info["script_name"]
        exec_mode = task_info.get("script_exec_mode", "append")
        
        # 创建安全的任务ID
        safe_table_name = table_name.replace(".", "_").replace("-", "_")
        task_id = f"model_{safe_table_name}"
        
        model_task = PythonOperator(
            task_id=task_id,
            python_callable=process_model,
            op_kwargs={
                "target_table": table_name,
                "script_name": script_name,
                "script_exec_mode": exec_mode,
                "exec_date": exec_date  # 直接使用解析出的exec_date，不使用XCom
            },
            retries=TASK_RETRY_CONFIG["retries"],
            retry_delay=timedelta(minutes=TASK_RETRY_CONFIG["retry_delay_minutes"]),
            dag=dag
        )
        
        # 将任务添加到字典
        task_dict[table_name] = model_task
        
        # 设置依赖关系
        deps = dependencies.get(table_name, [])
        has_dependency = False
        
        # 处理模型表之间的依赖
        for dep in deps:
            dep_table = dep.get("table_name")
            dep_type = dep.get("table_type")
            
            if dep_table in task_dict:
                task_dict[dep_table] >> model_task
                has_dependency = True
                logger.info(f"预先设置依赖: {dep_table} >> {table_name}")
        
        # 如果没有依赖，则依赖于准备任务和所有资源表任务
        if not has_dependency:
            # 从prepare_plan任务直接连接
            prepare_plan >> model_task
            logger.info(f"预先设置基本依赖: prepare_plan >> {task_id}")
            
            # 同时从所有资源表任务连接 - 限制每个模型表最多依赖5个资源表，避免过度复杂的依赖关系
            resource_count = 0
            for resource_table in resource_tasks:
                if resource_count >= 5:
                    break
                
                resource_name = resource_table["target_table"]
                if resource_name in task_dict:
                    task_dict[resource_name] >> model_task
                    logger.info(f"预先设置资源依赖: {resource_name} >> {table_name}")
                    resource_count += 1
    
    # 找出所有终端任务（没有下游依赖的任务）
    terminal_tasks = []
    
    # 检查所有模型表任务
    for table_name in execution_order:
        # 检查是否有下游任务
        has_downstream = False
        for source, deps in dependencies.items():
            if source == table_name:  # 跳过自身
                continue
            for dep in deps:
                if dep.get("table_name") == table_name:
                    has_downstream = True
                    break
            if has_downstream:
                break
        
        # 如果没有下游任务，添加到终端任务列表
        if not has_downstream and table_name in task_dict:
            terminal_tasks.append(table_name)
    
    # 如果没有模型表任务，将所有资源表任务视为终端任务
    if not model_tasks and resource_tasks:
        terminal_tasks = [task["target_table"] for task in resource_tasks]
    
    # 如果既没有模型表任务也没有资源表任务，直接连接准备任务到完成标记
    if not terminal_tasks:
        prepare_plan >> processing_completed
        logger.warning("未找到任何任务，直接连接准备任务到完成标记")
    else:
        # 将所有终端任务连接到完成标记
        for table_name in terminal_tasks:
            if table_name in task_dict:
                task_dict[table_name] >> processing_completed
                logger.info(f"设置终端任务: {table_name} >> processing_completed")

logger.info(f"DAG dag_dataops_unified_data_scheduler 定义完成，预创建了 {len(task_dict)} 个任务")