airflow_scheduler/dags/dag_data_model_scheduler.py

# dag_data_model_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
import pendulum
import logging
import networkx as nx
from utils import (
    get_enabled_tables,
    is_data_model_table,
    run_model_script,
    get_model_dependency_graph,
    check_script_exists,
    get_script_name_from_neo4j
)
from config import TASK_RETRY_CONFIG

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

def get_all_enabled_tables_for_today():
    """
    根据当前日期获取所有需要处理的表
    
    返回:
        list: 需要处理的表配置列表
    """
    today = pendulum.today()
    # 原始代码（注释）
    # is_monday = today.day_of_week == 0
    # is_first_day_of_month = today.day == 1
    # is_first_day_of_year = today.month == 1 and today.day == 1
    
    # 测试用：所有条件设为True
    is_monday = True
    is_first_day_of_month = True
    is_first_day_of_year = True
    
    logger.info(f"今日日期: {today.to_date_string()}")
    logger.info(f"日期特性: 是否周一={is_monday}, 是否月初={is_first_day_of_month}, 是否年初={is_first_day_of_year}")
    
    all_tables = []
    
    # 每天都处理daily表
    daily_tables = get_enabled_tables("daily")
    all_tables.extend(daily_tables)
    logger.info(f"添加daily表: {len(daily_tables)}个")
    
    # 周一处理weekly表
    if is_monday:
        weekly_tables = get_enabled_tables("weekly")
        all_tables.extend(weekly_tables)
        logger.info(f"今天是周一，添加weekly表: {len(weekly_tables)}个")
    
    # 月初处理monthly表
    if is_first_day_of_month:
        monthly_tables = get_enabled_tables("monthly")
        all_tables.extend(monthly_tables)
        logger.info(f"今天是月初，添加monthly表: {len(monthly_tables)}个")
    
    # 年初处理yearly表
    if is_first_day_of_year:
        yearly_tables = get_enabled_tables("yearly")
        all_tables.extend(yearly_tables)
        logger.info(f"今天是年初，添加yearly表: {len(yearly_tables)}个")
    
    # 去重
    unique_tables = {}
    for item in all_tables:
        table_name = item["table_name"]
        if table_name not in unique_tables:
            unique_tables[table_name] = item
        else:
            # 如果存在重复，优先保留execution_mode为full_refresh的配置
            if item["execution_mode"] == "full_refresh":
                unique_tables[table_name] = item
    
    result_tables = list(unique_tables.values())
    logger.info(f"去重后，共 {len(result_tables)} 个表需要处理")
    
    # 记录所有需要处理的表
    for idx, item in enumerate(result_tables, 1):
        logger.info(f"表[{idx}]: {item['table_name']}, 执行模式: {item['execution_mode']}")
    
    return result_tables

def optimize_execution_plan(tables):
    """
    优化表的执行计划
    
    参数:
        tables (list): 表配置列表
        
    返回:
        tuple: (优化后的表执行顺序, 依赖关系图)
    """
    logger.info("开始优化执行计划...")
    
    # 筛选出DataModel类型的表
    model_tables = []
    for table in tables:
        table_name = table["table_name"]
        if is_data_model_table(table_name):
            model_tables.append(table)
    
    logger.info(f"筛选出 {len(model_tables)} 个DataModel类型的表")
    
    if not model_tables:
        logger.warning("没有找到DataModel类型的表，无需优化执行计划")
        return [], {}
    
    # 获取表名列表
    table_names = [t["table_name"] for t in model_tables]
    
    # 创建有向图
    G = nx.DiGraph()
    
    # 添加所有节点
    for table_name in table_names:
        G.add_node(table_name)
    
    # 获取依赖关系
    dependency_dict = get_model_dependency_graph(table_names)
    logger.info(f"获取到 {len(dependency_dict)} 个表的依赖关系")
    
    # 添加依赖边
    edge_count = 0
    for target, upstreams in dependency_dict.items():
        for upstream in upstreams:
            if upstream in table_names:  # 确保只考虑当前处理的表
                G.add_edge(upstream, target)  # 从上游指向下游
                edge_count += 1
    
    logger.info(f"依赖图中添加了 {edge_count} 条边")
    
    # 检测循环依赖
    cycles = list(nx.simple_cycles(G))
    if cycles:
        logger.warning(f"检测到 {len(cycles)} 个循环依赖，将尝试打破循环")
        # 打破循环依赖（简单策略：移除每个循环中的最后一条边）
        for cycle in cycles:
            # 移除循环中的最后一条边
            G.remove_edge(cycle[-1], cycle[0])
            logger.info(f"打破循环依赖: 移除 {cycle[-1]} -> {cycle[0]} 的依赖")
    
    # 生成拓扑排序
    try:
        # 拓扑排序会从没有入边的节点开始，这些节点是上游节点，因此排序结果是从上游到下游
        execution_order = list(nx.topological_sort(G))
        logger.info(f"成功生成执行顺序，按从上游到下游顺序共 {len(execution_order)} 个表")
        
        # 创建结果依赖字典，包含所有表（即使没有依赖）
        result_dependency_dict = {name: [] for name in table_names}
        
        # 添加实际依赖关系
        for target, upstreams in dependency_dict.items():
            if target in table_names:  # 确保只考虑当前处理的表
                result_dependency_dict[target] = [u for u in upstreams if u in table_names]
        
        return execution_order, result_dependency_dict
    except Exception as e:
        logger.error(f"生成执行顺序失败: {str(e)}")
        # 如果拓扑排序失败，返回原始表名列表和空依赖图
        return table_names, {name: [] for name in table_names}

with DAG(
    "dag_data_model_scheduler", 
    start_date=datetime(2024, 1, 1), 
    schedule_interval="@daily", 
    catchup=False
) as dag:
    logger.info("初始化 dag_data_model_scheduler DAG")
    
    # 等待资源表 DAG 完成
    wait_for_resource = ExternalTaskSensor(
        task_id="wait_for_resource_loading",
        external_dag_id="dag_data_resource_scheduler",
        external_task_id="resource_loading_completed",
        mode="poke",
        timeout=3600,
        poke_interval=30
    )
    logger.info("创建资源表等待任务 - wait_for_resource_loading")

    # 创建一个完成标记任务
    model_processing_completed = EmptyOperator(
        task_id="model_processing_completed",
        dag=dag
    )
    logger.info("创建模型处理完成标记 - model_processing_completed")

    try:
        # 获取今日需要处理的所有表
        all_enabled_tables = get_all_enabled_tables_for_today()
        
        if not all_enabled_tables:
            logger.info("今天没有需要处理的表，直接连接开始和结束任务")
            wait_for_resource >> model_processing_completed
        else:
            # 优化执行计划
            execution_order, dependency_dict = optimize_execution_plan(all_enabled_tables)
            
            if not execution_order:
                logger.info("执行计划为空，直接连接开始和结束任务")
                wait_for_resource >> model_processing_completed
            else:
                # 创建任务字典
                task_dict = {}
                
                # 为每个表创建处理任务
                for table_name in execution_order:
                    # 查找表配置
                    table_config = next((t for t in all_enabled_tables if t["table_name"] == table_name), None)
                    
                    if table_config:
                        logger.info(f"为表 {table_name} 创建处理任务，执行模式: {table_config['execution_mode']}")
                        
                        # 创建任务
                        task = PythonOperator(
                            task_id=f"process_{table_name}",
                            python_callable=run_model_script,
                            op_kwargs={
                                "table_name": table_name,
                                "execution_mode": table_config["execution_mode"]
                            },
                            retries=TASK_RETRY_CONFIG["retries"],
                            retry_delay=timedelta(minutes=TASK_RETRY_CONFIG["retry_delay_minutes"]),
                            dag=dag
                        )
                        
                        # 将任务添加到字典
                        task_dict[table_name] = task
                
                # 设置任务间的依赖关系
                for table_name, task in task_dict.items():
                    # 获取上游依赖
                    upstream_tables = dependency_dict.get(table_name, [])
                    
                    if not upstream_tables:
                        # 如果没有上游依赖，直接连接到资源表等待任务
                        logger.info(f"表 {table_name} 没有上游依赖，连接到资源表等待任务")
                        wait_for_resource >> task
                    else:
                        # 设置与上游表的依赖关系
                        for upstream_table in upstream_tables:
                            if upstream_table in task_dict:
                                logger.info(f"设置依赖: {upstream_table} >> {table_name}")
                                task_dict[upstream_table] >> task
                    
                    # 检查是否是末端节点（没有下游节点）
                    is_terminal = True
                    for target, upstreams in dependency_dict.items():
                        if table_name in upstreams:
                            is_terminal = False
                            break
                    
                    # 如果是末端节点，连接到模型处理完成标记
                    if is_terminal:
                        logger.info(f"表 {table_name} 是末端节点，连接到模型处理完成标记")
                        task >> model_processing_completed
                
                # 处理特殊情况：检查是否有任务连接到完成标记
                has_connection_to_completed = False
                for task in task_dict.values():
                    for downstream in task.downstream_list:
                        if downstream.task_id == model_processing_completed.task_id:
                            has_connection_to_completed = True
                            break
                
                # 如果没有任务连接到完成标记，连接所有任务到完成标记
                if not has_connection_to_completed and task_dict:
                    logger.info("没有发现连接到完成标记的任务，连接所有任务到完成标记")
                    for task in task_dict.values():
                        task >> model_processing_completed
                
                # 处理特殊情况：如果资源等待任务没有下游任务，直接连接到完成标记
                if not wait_for_resource.downstream_list:
                    logger.info("资源等待任务没有下游任务，直接连接到完成标记")
                    wait_for_resource >> model_processing_completed
    
    except Exception as e:
        logger.error(f"构建DAG时出错: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        # 确保出错时也有完整的执行流
        wait_for_resource >> model_processing_completed