airflow_scheduler/dags/dag_dataops_model_scheduler.py

# dag_dataops_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 logging
from common import (
    get_pg_conn, execute_with_monitoring, get_datamodel_dependency_from_neo4j,
    generate_optimized_execution_order, get_today_date
)

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

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

def get_datamodel_tasks(exec_date):
    """从airflow_dag_schedule表获取DataModel任务"""
    conn = get_pg_conn()
    cursor = conn.cursor()
    try:
        cursor.execute("""
            SELECT target_table, script_name, script_exec_mode
            FROM airflow_dag_schedule 
            WHERE exec_date = %s AND target_table_label = 'DataModel'
        """, (exec_date,))
        
        results = cursor.fetchall()
        
        tasks = []
        for row in results:
            target_table, script_name, script_exec_mode = row
            tasks.append({
                "target_table": target_table,
                "script_name": script_name,
                "script_exec_mode": script_exec_mode or "append"  # 默认为append
            })
        
        logger.info(f"使用日期 {exec_date} 获取到 {len(tasks)} 个DataModel任务")
        return tasks
    except Exception as e:
        logger.error(f"获取DataModel任务时出错: {str(e)}")
        return []
    finally:
        cursor.close()
        conn.close()

# 创建DAG
with DAG(
    "dag_dataops_model_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)
    }
) as dag:
    
    # 等待resource DAG完成
    wait_for_resource = ExternalTaskSensor(
        task_id="wait_for_resource",
        external_dag_id="dag_dataops_resource_scheduler",
        external_task_id="resource_processing_completed",
        mode="poke",
        timeout=3600,
        poke_interval=30,
        dag=dag
    )
    
    # 处理完成标记
    model_processing_completed = EmptyOperator(
        task_id="model_processing_completed",
        dag=dag
    )
    
    try:
        # 获取最近的日期
        latest_date = get_latest_date_with_models()
        logger.info(f"使用最近的日期 {latest_date} 查询模型任务")
        
        # 获取所有DataModel任务
        model_tasks = get_datamodel_tasks(latest_date)
        
        if model_tasks:
            # 获取表名列表
            table_names = [task["target_table"] for task in model_tasks]
            
            # 获取依赖关系
            dependency_dict = get_datamodel_dependency_from_neo4j(table_names)
            
            # 生成优化的执行顺序
            execution_order = generate_optimized_execution_order(table_names, dependency_dict)
            logger.info(f"生成的优化执行顺序: {execution_order}")
            
            # 创建任务字典
            task_dict = {}
            
            # 为每个表创建处理任务
            for table_name in execution_order:
                # 查找表任务信息
                task_info = next((t for t in model_tasks if t["target_table"] == table_name), None)
                
                if task_info and task_info.get("script_name"):
                    process_task = PythonOperator(
                        task_id=f"process_model_{table_name.replace('.', '_')}",
                        python_callable=execute_with_monitoring,
                        op_kwargs={
                            "target_table": table_name,
                            "script_name": task_info["script_name"],
                            "script_exec_mode": task_info.get("script_exec_mode", "append"), 
                            "exec_date": latest_date  # 使用从数据库获取的最近日期
                        },
                        dag=dag
                    )
                    task_dict[table_name] = process_task
                    logger.info(f"创建处理任务: {table_name}")
                else:
                    logger.warning(f"表 {table_name} 没有script_name，跳过任务创建")
            
            # 设置任务间的依赖关系
            for target_table, task in task_dict.items():
                # 获取上游依赖
                upstream_tables = dependency_dict.get(target_table, [])
                upstream_tables = [t for t in upstream_tables if t in task_dict]
                
                if not upstream_tables:
                    # 如果没有上游依赖，直接连接到wait_for_resource
                    logger.info(f"表 {target_table} 没有上游依赖，连接到wait_for_resource")
                    wait_for_resource >> task
                else:
                    # 设置与上游表的依赖关系
                    for upstream_table in upstream_tables:
                        logger.info(f"设置依赖: {upstream_table} >> {target_table}")
                        task_dict[upstream_table] >> task
                
                # 检查是否是末端节点（没有下游任务）
                is_terminal = True
                for downstream, upstreams in dependency_dict.items():
                    if target_table in upstreams and downstream in task_dict:
                        is_terminal = False
                        break
                
                # 如果是末端节点，连接到model_processing_completed
                if is_terminal:
                    logger.info(f"表 {target_table} 是末端节点，连接到model_processing_completed")
                    task >> model_processing_completed
            
            # 处理特殊情况
            # 检查是否有任务连接到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 has_connection_to_completed:
                    break
            
            # 如果没有任务连接到model_processing_completed，连接所有任务到完成标记
            if not has_connection_to_completed and task_dict:
                logger.info("没有任务连接到model_processing_completed，连接所有任务到完成标记")
                for task in task_dict.values():
                    task >> model_processing_completed
            
            # 检查是否有任务连接到wait_for_resource
            has_connection_from_wait = False
            for task in task_dict.values():
                for upstream in task.upstream_list:
                    if upstream.task_id == wait_for_resource.task_id:
                        has_connection_from_wait = True
                        break
                
                if has_connection_from_wait:
                    break
            
            # 如果没有任务连接到wait_for_resource，连接wait_for_resource到所有任务
            if not has_connection_from_wait and task_dict:
                logger.info("没有任务连接到wait_for_resource，连接wait_for_resource到所有任务")
                for task in task_dict.values():
                    wait_for_resource >> task
        else:
            # 如果没有任务，直接将等待节点连接到完成
            wait_for_resource >> model_processing_completed
            logger.warning("没有找到DataModel任务，直接将等待节点连接到完成")
    except Exception as e:
        logger.error(f"创建模型处理DAG时出错: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        # 确保在出错时也有完整的执行流程
        wait_for_resource >> model_processing_completed