如何在Airflow中无需SSH Operator实现条件任务执行与运行时任务添加

我来帮你实现Airflow里这个不用SSH Operator的条件任务执行需求，还能搞定task1失败时动态添加任务的逻辑，下面是完整的方案：

要满足你的需求，我们主要用到Airflow的这几个组件：

• BranchPythonOperator：用来根据task1的执行结果做分支判断，决定后续跑task2还是task3
• PythonOperator：封装自定义逻辑，包括动态添加任务的操作
• Airflow的上下文(context)和TaskInstance：用来获取任务的执行状态，实现动态任务的添加

from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.operators.bash import BashOperator
from datetime import datetime, timedelta
from airflow.models import Variable, TaskInstance
from airflow.utils.state import State

# 默认DAG参数
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 0,
    'retry_delay': timedelta(minutes=5),
}

with DAG(
    'conditional_task_with_dynamic_add',
    default_args=default_args,
    description='Conditional task execution with dynamic task addition',
    schedule_interval=timedelta(days=1),
    catchup=False,
) as dag:

    # Task 1: 示例任务，这里可以替换成你的实际任务逻辑
    task1 = BashOperator(
        task_id='task1',
        bash_command='echo "Running Task 1"; exit 0',  # 改成exit 1可以测试失败分支
    )

    # 分支判断函数：根据task1的状态决定走哪个分支
    def branch_func(**context):
        ti = context['task_instance']
        # 获取task1的执行状态
        task1_state = ti.get_task_instance(task_id='task1').state
        if task1_state == State.SUCCESS:
            return 'task2'
        else:
            return 'task3'

    branch_task = BranchPythonOperator(
        task_id='branch_task',
        python_callable=branch_func,
        provide_context=True,
    )

    # Task 2: task1成功时执行的任务
    task2 = BashOperator(
        task_id='task2',
        bash_command='echo "Task 1 succeeded, running Task 2"',
    )

    # Task 3: task1失败时执行的任务，同时实现动态添加任务
    def handle_task1_failure(**context):
        print("Task 1 failed, running Task 3 and adding dynamic task")
        
        # 这里实现动态添加任务的逻辑
        # 示例：创建一个动态Bash任务
        dynamic_task = BashOperator(
            task_id=f'dynamic_task_{datetime.now().strftime("%Y%m%d%H%M%S")}',
            bash_command='echo "This is a dynamically added task after Task 1 failed"',
            dag=dag,
        )
        
        # 将动态任务与当前任务建立依赖（可选，根据你的流程需求）
        context['task_instance'].task.set_downstream(dynamic_task)
        
        # 也可以用Variable存储动态任务信息，后续流程使用（如果需要）
        Variable.set("dynamic_task_added", "True", serialize_json=True)

    task3 = PythonOperator(
        task_id='task3',
        python_callable=handle_task1_failure,
        provide_context=True,
    )

    # 设置任务依赖
    task1 >> branch_task >> [task2, task3]

1. 分支判断逻辑

   • branch_func通过上下文获取TaskInstance，然后拿到task1的执行状态，根据状态返回对应的task_id，BranchPythonOperator会自动跳转到对应的任务执行。

2. 动态添加任务

   • 在handle_task1_failure函数里，我们直接创建一个新的Operator（这里用BashOperator示例，你可以换成其他Operator），指定dag=dag把它加入当前DAG。
   • 可以通过set_downstream或者set_upstream来建立动态任务和现有任务的依赖关系，确保流程符合你的需求。
   • 如果需要后续流程引用动态任务，也可以用Airflow的Variable来存储相关信息，方便其他任务读取。

3. 状态判断细节

   • 我们用airflow.utils.state.State里的常量来判断任务状态，比直接写字符串更可靠，避免拼写错误。

内容的提问来源于stack exchange，提问作者Manmohan