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
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 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()
        
        # 查询所有模型表任务
        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()
        
        # 整理资源表信息
        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
    }
    
    # 将执行计划保存到XCom，使用自定义序列化器处理日期对象
    kwargs['ti'].xcom_push(key='execution_plan', value=json.dumps(execution_plan, default=json_serial))
    logger.info(f"准备了执行计划，包含 {len(resource_tasks)} 个资源表任务和 {len(model_tasks)} 个模型表任务")
    
    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)}")
    
    return execute_with_monitoring(
        target_table=target_table,
        script_name=script_name,
        script_exec_mode=script_exec_mode,
        exec_date=exec_date
    )

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)}")
    
    return execute_with_monitoring(
        target_table=target_table,
        script_name=script_name,
        script_exec_mode=script_exec_mode,
        exec_date=exec_date
    )

# 创建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)
    }
) 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="poke",
        timeout=3600,
        poke_interval=30,
        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",
        dag=dag
    )
    
    # 设置初始任务依赖
    wait_for_prepare >> prepare_plan
    
    # 从执行计划JSON中获取信息
    execution_plan_json = '''{"exec_date": "2025-04-12", "resource_tasks": [], "model_tasks": [], "dependencies": {}}'''
    
    try:
        # 从文件或数据库中获取执行计划作为默认值
        import os
        plan_path = os.path.join(os.path.dirname(__file__), 'last_execution_plan.json')
        if os.path.exists(plan_path):
            with open(plan_path, 'r') as f:
                execution_plan_json = f.read()
    except Exception as e:
        logger.warning(f"读取执行计划默认值时出错: {str(e)}")
    
    # 解析执行计划获取任务信息
    try:
        execution_plan = json.loads(execution_plan_json)
        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", {})
        
        # 任务字典，用于设置依赖关系
        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": """{{ ti.xcom_pull(task_ids='prepare_execution_plan', key='execution_plan') }}"""
                },
                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
        
        # 创建有向图，用于检测模型表之间的依赖关系
        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]} 的依赖")
        
        # 生成拓扑排序，确定执行顺序
        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": """{{ ti.xcom_pull(task_ids='prepare_execution_plan', key='execution_plan') }}"""
                },
                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
                
                # 同时从所有资源表任务连接
                for resource_table in resource_tasks:
                    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}")
        
        # 找出所有终端任务（没有下游依赖的任务）
        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")
    
    except Exception as e:
        logger.error(f"构建任务DAG时出错: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        
        # 确保即使出错，DAG也能正常完成
        prepare_plan >> processing_completed