Airflow中实现任务回滚机制的方案咨询

不用Branch Operator，用**TriggerRule（触发规则）+ TaskGroup（任务组）**是更贴合场景的方案——因为你的需求是「失败触发补偿」，而非分支选择，原生触发规则比分支算子更简洁直接。

核心思路 #

• 正常流程：task1 >> task2 >> task3
• 回滚触发逻辑：
  1. task2失败 → 仅执行undo_task1
  2. task3失败 → 按undo_task2 >> undo_task1顺序执行回滚
• 用TriggerRule控制回滚任务的触发时机，用TaskGroup归置回滚逻辑，避免代码混乱

完整代码示例 #

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.utils.trigger_rule import TriggerRule
from airflow.utils.task_group import TaskGroup
from datetime import datetime

# 模拟业务任务
def run_task(task_name):
    print(f"Running business task: {task_name}")
    # 可取消注释模拟失败场景
    # if task_name == "task2":
    #     raise Exception(f"{task_name} failed intentionally")
    # if task_name == "task3":
    #     raise Exception(f"{task_name} failed intentionally")

# 模拟回滚任务
def run_undo_task(undo_task_name, **context):
    # 通过上下文获取上游任务状态，精准判断是否需要执行回滚
    dagrun = context["task_instance"].get_dagrun()
    task_states = {ti.task_id: ti.state for ti in dagrun.get_task_instances()}

    # 判断触发条件：要么task2失败，要么task3失败且undo_task2已完成
    need_rollback = False
    if task_states.get("task2") == "failed":
        need_rollback = True
    elif task_states.get("task3") == "failed" and task_states.get("undo_task2") == "success":
        need_rollback = True

    if need_rollback:
        print(f"Executing rollback: {undo_task_name}")
    else:
        print(f"Skipping {undo_task_name} - no rollback needed")
        context["task_instance"].state = "skipped"

with DAG(
    dag_id="business_rollback_dag",
    start_date=datetime(2024, 1, 1),
    schedule_interval=None,
    catchup=False,
    tags=["rollback", "business"]
) as dag:
    # 正常业务任务
    task1 = PythonOperator(
        task_id="task1",
        python_callable=run_task,
        op_kwargs={"task_name": "task1"}
    )

    task2 = PythonOperator(
        task_id="task2",
        python_callable=run_task,
        op_kwargs={"task_name": "task2"}
    )

    task3 = PythonOperator(
        task_id="task3",
        python_callable=run_task,
        op_kwargs={"task_name": "task3"}
    )

    # 回滚任务组：统一管理回滚逻辑
    with TaskGroup("rollback_group") as rollback_group:
        undo_task2 = PythonOperator(
            task_id="undo_task2",
            python_callable=run_undo_task,
            op_kwargs={"undo_task_name": "undo_task2"},
            # 仅当task3失败时触发（task2已成功，因为task3依赖task2）
            trigger_rule=TriggerRule.ONE_FAILED,
            provide_context=True
        )

        undo_task1 = PythonOperator(
            task_id="undo_task1",
            python_callable=run_undo_task,
            op_kwargs={"undo_task_name": "undo_task1"},
            # 不管上游状态如何先触发，再在函数内判断是否执行
            trigger_rule=TriggerRule.ALL_DONE,
            provide_context=True
        )

        # 强制回滚顺序：undo_task2执行完才允许执行undo_task1
        undo_task2 >> undo_task1

    # 设置依赖关系
    task1 >> task2 >> task3
    # 回滚任务的依赖绑定
    task3 >> undo_task2  # undo_task2仅依赖task3失败
    [task2, undo_task2] >> undo_task1  # undo_task1依赖task2状态或undo_task2执行结果
    task1 >> undo_task1  # 确保task1已执行才回滚

逻辑说明 #

1. 正常流程：task1→task2→task3全部成功时，回滚任务因触发规则和内部判断不会执行，直接跳过。
2. task2失败：task3不会启动，undo_task1的内部判断检测到task2失败，执行回滚；undo_task2因task3未执行，触发规则不满足，不会启动。
3. task3失败：undo_task2因task3失败触发，执行回滚；undo_task1检测到task3失败且undo_task2已完成，执行回滚，严格遵循undo_task2 >> undo_task1顺序。

为什么不用Branch Operator？ #

Branch Operator的核心是主动选择分支路径，而你的场景是被动触发补偿操作——用TriggerRule更贴合Airflow的原生状态依赖语义，代码更简洁，不需要额外的分支判断逻辑，维护成本更低。

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