Airflow动态任务映射：实现映射任务的顺序依赖控制

我有一个Airflow DAG，逻辑是查询BigQuery表获取待执行的存储过程列表，列表的数量和内容会随每次运行变化。我已经创建了动态任务，尝试通过设置max_active_tis_per_dag=1控制顺序执行，但需要实现：当某个映射任务失败时，后续所有映射任务直接终止执行并标记失败。

示例依赖要求：

• Task1（获取并排序任务列表）→ 映射Task2(0) → 映射Task2(1) → 映射Task2(2)
• 若Task2(0)失败，Task2(1)、Task2(2)不执行
• 若Task2(0)成功但Task2(1)失败，Task2(2)不执行

当前代码实现：

@dag(schedule=None, start_date=pendulum.datetime(2024, 5, 21, 12, 00, 00), catchup=False)
def get_execution_steps():
    # 1) 从BigQuery获取待执行的存储过程列表
    # 2) 按指定规则排序任务
    # 3) 扩展为动态任务，支持从失败点重跑，无需重新执行已完成的任务

    @task()
    def get_and_sort_tasks():
        client=bigquery.Client()
        unsorted_task_list = client.query(bq_steps_sql, location='US')
        logger.info(repr(unsorted_task_list.result()))
        sorted_steps_list = sorted(unsorted_task_list, key=lambda x: (int(x[0]), int(x[1])))
        logger.info(sorted_steps_list)
        sorted_task_list=[]
        for step in sorted_steps_list:
            execution_task_name = 'Execute_statement_task{}_subtask{}'.format(step[0], step[1])
            sql=step[2]
            logger.info([execution_task_name, sql])
            sorted_task_list.append([execution_task_name, sql])

    @task(max_active_tis_per_dag=1, task_id= 'execute_statement_task', map_index_template="{{task_name}}")
    def execute_statements(sorted_task_list):
        logger.info(sorted_task_list)
        context = get_current_context()
        context["task_name"] = sorted_task_list[0]
        logger.info('Running:{}'.format(sorted_task_list[1]))
        client=bigquery.Client()
        query=sorted_task_list[1]
        job = client.query(query, location='US')
        logger.info(repr(job.result()))

    tasks = execute_statements.partial().expand(sorted_task_list=get_and_sort_tasks())

尝试过的方法：
添加了下游任务并设置trigger_rule='one_failed'，但失败时后续映射任务仍会继续执行：

@task(task_id='task_exception', trigger_rule='one_failed')
def task_exception():
    raise AirflowFailException('Dynamic Statement Failed')

默认的expand生成的动态任务是并行调度的，max_active_tis_per_dag=1仅控制并发数，无法实现失败阻断+顺序依赖。要实现需求，需要让每个动态任务实例显式依赖前一个实例的成功状态，具体实现方式如下：

方法：链式构建任务依赖 #

通过遍历任务列表，逐个创建任务并设置依赖关系，确保每个任务仅在前一个任务成功时才执行：

from airflow.decorators import dag, task
from airflow.exceptions import AirflowFailException
import pendulum
from google.cloud import bigquery
import logging

logger = logging.getLogger(__name__)
bq_steps_sql = "YOUR_BIGQUERY_SQL_HERE"  # 替换为你的查询SQL

@dag(schedule=None, start_date=pendulum.datetime(2024, 5, 21, 12, 00, 00), catchup=False)
def get_execution_steps():
    @task()
    def get_and_sort_tasks():
        client = bigquery.Client()
        query_job = client.query(bq_steps_sql, location='US')
        unsorted_task_list = query_job.result()
        logger.info(repr(unsorted_task_list))
        # 按步骤编号排序
        sorted_steps_list = sorted(unsorted_task_list, key=lambda x: (int(x[0]), int(x[1])))
        sorted_task_list = []
        for step in sorted_steps_list:
            execution_task_name = f'Execute_statement_task{step[0]}_subtask{step[1]}'
            sql = step[2]
            sorted_task_list.append({"task_name": execution_task_name, "sql": sql})
        return sorted_task_list

    @task(task_id='execute_statement_task')
    def execute_statements(task_info):
        logger.info(f"Running task: {task_info['task_name']}")
        client = bigquery.Client()
        query = task_info['sql']
        job = client.query(query, location='US')
        # 等待查询完成，若失败会抛出异常
        job.result()
        logger.info(f"Completed task: {task_info['task_name']}")

    # 获取任务列表
    task_list = get_and_sort_tasks()
    # 链式构建依赖
    previous_task = None
    for idx, task_info in enumerate(task_list):
        current_task = execute_statements.override(task_id=f"execute_statement_task_{idx}")(task_info)
        if previous_task:
            # 当前任务依赖前一个任务成功完成
            current_task.set_upstream(previous_task)
        previous_task = current_task

get_execution_steps()

关键逻辑说明 #

1. 链式依赖：通过set_upstream(previous_task)让每个任务严格按顺序执行，前一个任务失败时，后续任务会自动处于upstream_failed状态，不会启动执行。
2. 唯一任务ID：通过override(task_id=...)为每个动态任务设置唯一ID，方便单独查看、重跑单个失败任务。
3. 异常传递：job.result()会在BigQuery查询失败时直接抛出异常，标记当前任务为失败，触发后续任务的阻断逻辑。

额外优化 #

• 支持失败点重跑：Airflow原生支持单独重跑某个失败的execute_statement_task_{idx}任务，无需重新执行前面的成功任务。
• 添加失败重试：可以给execute_statements任务增加重试配置，提升稳定性：

  @task(task_id='execute_statement_task', retries=2, retry_delay=pendulum.duration(minutes=5))
  def execute_statements(task_info):
      # ... 原有逻辑 ...
  

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