Airflow中如何在DAG任务组中参数化循环次数？

简而言之：如何用DAG params参数控制任务组的循环次数？ #

需要将DAG params中的input_number值，替换TaskGroup里循环的固定数值13：

with DAG(dag_id=v_dag_id,
     description=f'Ejecución SQL {v_project_name}',
     tags=['caba', '16860', 'kde063'],
     default_args=v_default_args,
     schedule_interval=None,
     max_active_tasks=12,
     max_active_runs=10,
     start_date=v_start_date,
     catchup=v_catchup,
     params = {"input_number": 13} # <------ 要引用这个数值
        ) as dag:

with TaskGroup(group_id='test_exadata') as test_exadata:

        a = []
        for i in range(0,13): # <------ 如何替换成input_number的值？
            nombre = 'test_exadata_'+str(i)
            a.append(
            CustomExadataOperator(
                task_id=nombre,
                repo_url= REPO_URL,
                repo_branch= REPO_BRANCH,
                sql_path= f'src/sql/test_query.sql',
                engine= SECRETS_ENGINE,
                secret= SECRETS_PATH,
                vault_token_id= 'AIRFLOW_VAULT_TOKEN',
                output_mode= 'log',
                dag= dag
                )
            )

test_exadata

详细场景说明 #

我在Airflow的DAG里配置了几个默认参数：上月日期（格式YYYYMM）、回溯月份数（整数）、表名前缀。比如给定month_input: 202412、months_hist_input:13、prefix_name_table_input:TEMP_XXX_，代码会从202412倒推，创建TEMP_XXX_202412、TEMP_XXX_202411……TEMP_XXX_202401这类表。

参数已经在DAG的params里定义：

with DAG(dag_id=v_dag_id,
         description=f'Ejecución SQL {v_project_name}',
         tags=['caba', '16860', 'kde063'],
         default_args=v_default_args,
         schedule_interval=None,
         max_active_tasks=12,
         max_active_runs=10,
         start_date=v_start_date,
         catchup=v_catchup,
         params = {
             "month_input": (date.today()+relativedelta(months=-1)).strftime('%Y%m'),
             "months_hist_input": 13,
             "prefix_name_table_input": "TEMP_XXX_"
             }
            ) as dag:

目前我已经能通过PythonOperator访问这些参数，生成日期和表名并推送到XCom：

def _create_dates_and_names(month,months_hist,prefix_name_table):
    
    months_hist = int(months_hist)
    dates = []
    names = []
    
    for i in range(0,months_hist):
        dates.append((datetime.strptime(month, '%Y%m') - relativedelta(months=i)).strftime('%Y%m'))
        names.append(f'{prefix_name_table}_{dates[i]}')

    context = get_current_context()
    ti = context["ti"]
    ti.xcom_push(key="dates", value=dates)
    ti.xcom_push(key="names", value=names)


create_dates_and_names = PythonOperator(
    task_id = 'create_dates_and_names',
    python_callable = _create_dates_and_names,
    op_kwargs = {
        'month' : '{{params["month_input"]}}',
        'months_hist': '{{params["months_hist_input"]}}',
        'prefix_name_table': '{{params["prefix_name_table_input"]}}'},
    provide_context = True
)

后续任务通过XCom获取这些值，但TaskGroup里的循环次数还是写死的13：

with TaskGroup(group_id='test_exadata') as test_exadata:

        a = []

        for i in range(0,13):  # <----------- 要替换成months_hist_input的值
            nombre = 'test_exadata_'+str(i)
            a.append(
            CustomExadataOperator(
                task_id=nombre,
                repo_url= REPO_URL,
                repo_branch= REPO_BRANCH,
                sql_path= f'src/sql/test_query.sql',
                sql_args = f'{{{{ti.xcom_pull(key="dates",task_ids="create_dates_and_names")[{i}]}}}} {{{{ti.xcom_pull(key="names",task_ids="create_dates_and_names")[{i}]}}}}',
                engine= SECRETS_ENGINE,
                secret= SECRETS_PATH,
                vault_token_id= 'AIRFLOW_VAULT_TOKEN',
                output_mode= 'log',
                trigger_rule = 'none_failed_min_one_success',
                dag= dag
                )
            )

create_dates_and_names >> test_exadata

现在代码能正常跑，但问题是没法把循环次数换成months_hist_input的值。直接写for i in range(0,int({{params["months_hist_input"]}})):会把模板字符串当成普通字符串，不生效。是不是得用自定义函数来实现任务组的参数化？

解决方案 #

核心问题是：DAG解析阶段（静态）无法直接访问运行时的params参数，所以不能直接在循环里用模板语法。以下是两种可行方案：

方案1：参数为静态默认值时直接引用 #

如果months_hist_input是固定默认值，或者在DAG解析时就能确定（比如从环境变量、配置文件读取），可以直接在DAG定义阶段提取参数值：

# 从DAG的params默认值中提取循环次数
months_hist = dag.params["months_hist_input"]

with TaskGroup(group_id='test_exadata') as test_exadata:
    a = []
    for i in range(0, months_hist):
        nombre = 'test_exadata_'+str(i)
        a.append(
            CustomExadataOperator(
                task_id=nombre,
                repo_url= REPO_URL,
                repo_branch= REPO_BRANCH,
                sql_path= f'src/sql/test_query.sql',
                sql_args = f'{{{{ti.xcom_pull(key="dates",task_ids="create_dates_and_names")[{i}]}}}} {{{{ti.xcom_pull(key="names",task_ids="create_dates_and_names")[{i}]}}}}',
                engine= SECRETS_ENGINE,
                secret= SECRETS_PATH,
                vault_token_id= 'AIRFLOW_VAULT_TOKEN',
                output_mode= 'log',
                trigger_rule = 'none_failed_min_one_success',
                dag= dag
            )
        )

方案2：支持运行时动态传参（Airflow 2.3+推荐） #

如果需要在触发DAG时手动修改months_hist_input，用**动态任务映射（Dynamic Task Mapping）**替代手动循环，这是Airflow 2.3+的原生特性：

1. 修改PythonOperator，返回日期和表名的元组列表：

def _create_dates_and_names(month,months_hist,prefix_name_table):
    months_hist = int(months_hist)
    dates = []
    names = []
    for i in range(0,months_hist):
        dates.append((datetime.strptime(month, '%Y%m') - relativedelta(months=i)).strftime('%Y%m'))
        names.append(f'{prefix_name_table}_{dates[i]}')
    # 返回元组列表，供动态映射使用
    return list(zip(dates, names))

create_dates_and_names = PythonOperator(
    task_id = 'create_dates_and_names',
    python_callable = _create_dates_and_names,
    op_kwargs = {
        'month' : '{{params["month_input"]}}',
        'months_hist': '{{params["months_hist_input"]}}',
        'prefix_name_table': '{{params["prefix_name_table_input"]}}'},
    provide_context = True,
    do_xcom_push=True  # 确保返回值推送到XCom
)

2. 用partial+expand实现动态任务生成：

test_exadata_task = CustomExadataOperator.partial(
    task_id='test_exadata',
    repo_url= REPO_URL,
    repo_branch= REPO_BRANCH,
    sql_path= f'src/sql/test_query.sql',
    engine= SECRETS_ENGINE,
    secret= SECRETS_PATH,
    vault_token_id= 'AIRFLOW_VAULT_TOKEN',
    output_mode= 'log',
    trigger_rule = 'none_failed_min_one_success',
    dag= dag
).expand(
    sql_args=create_dates_and_names.output.map(lambda item: f'{item[0]} {item[1]}')
)

create_dates_and_names >> test_exadata_task

Airflow会根据create_dates_and_names返回的列表长度自动生成对应数量的任务，完美适配运行时传入的参数值。

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