如何通过dag_run.conf获取DAG触发参数并传入SSHOperator命令

你的代码存在几个关键问题，导致无法正确获取触发参数并生成任务，下面先指出问题，再给出可行的实现方案：

原代码的核心错误 #

1. 函数未执行调用：你直接把run_config函数赋值给number和item变量，这只是拿到了函数对象，并没有执行函数获取实际的配置值。
2. DAG解析阶段访问dag_run非法：在DAG定义（解析）的时候，DAG还没被触发，dag_run对象根本不存在，此时直接访问dag_run.conf会报错。
3. task_id格式化语法错误：'run_'.format(number[i])这种写法不会替换变量，应该写成'run_{}'.format(number[i])或者用f-stringf"run_{number[i]}"。

正确实现方式 #

推荐方案：动态任务映射（Airflow 2.2+） #

Airflow 2.2及以上支持的动态任务映射可以完美适配这种基于触发参数生成批量任务的场景，无需手动写循环，代码简洁且符合Airflow最佳实践：

from airflow import DAG
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow.operators.python import PythonOperator
from airflow.utils.dates import days_ago

def fetch_trigger_config(**kwargs):
    # 从触发上下文获取配置参数
    number_list = kwargs['dag_run'].conf.get('number', [])
    items_list = kwargs['dag_run'].conf.get('items', [])
    
    # 校验两个列表长度一致，避免后续任务出错
    if len(number_list) != len(items_list):
        raise ValueError("触发参数中number和items列表长度必须一致")
    
    return number_list, items_list

with DAG(
    dag_id="dynamic_ssh_task",
    schedule_interval=None,  # 手动触发DAG
    start_date=days_ago(1),
    catchup=False
) as dag:
    # 第一步：获取触发时传入的配置参数
    get_config_task = PythonOperator(
        task_id="get_trigger_config",
        python_callable=fetch_trigger_config,
        provide_context=True  # 让函数能拿到Airflow上下文参数
    )

    # 第二步：基于配置动态生成SSHOperator任务
    process_tasks = SSHOperator.partial(
        task_id="process_task",
        ssh_conn_id="your_ssh_conn_id"  # 替换成你在Airflow中配置的SSH连接ID
    ).expand(
        # 根据配置生成每个任务的命令
        command=lambda conf: [f"run_something.sh {num} {item}" for num, item in zip(conf[0], conf[1])],
        # 根据number值生成唯一的task_id
        task_id=lambda conf: [f"run_{num}" for num in conf[0]],
        # 传入get_config_task的输出结果作为参数
        op_args=[get_config_task.output]
    )

    # 设置任务依赖
    get_config_task >> process_tasks

兼容旧版Airflow方案（2.2以下） #

如果你的Airflow版本低于2.2，无法使用动态任务映射，可以通过Airflow临时变量中转配置，再生成任务：

from airflow import DAG
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow.operators.python import PythonOperator
from airflow.utils.dates import days_ago
from airflow.models import Variable

def save_config_to_var(**kwargs):
    number_list = kwargs['dag_run'].conf.get('number', [])
    items_list = kwargs['dag_run'].conf.get('items', [])
    
    if len(number_list) != len(items_list):
        raise ValueError("number和items列表长度必须一致")
    
    # 把配置存入Airflow临时变量
    Variable.set("tmp_number_list", number_list, serialize_json=True)
    Variable.set("tmp_items_list", items_list, serialize_json=True)

def load_config_from_var():
    # 从变量中读取配置
    number_list = Variable.get("tmp_number_list", deserialize_json=True)
    items_list = Variable.get("tmp_items_list", deserialize_json=True)
    return number_list, items_list

with DAG(
    dag_id="ssh_task_with_conf",
    schedule_interval=None,
    start_date=days_ago(1),
    catchup=False
) as dag:
    # 先保存触发配置到变量
    save_config_task = PythonOperator(
        task_id="save_trigger_config",
        python_callable=save_config_to_var,
        provide_context=True
    )

    # 加载配置生成任务
    number_list, items_list = load_config_from_var()
    for num, item in zip(number_list, items_list):
        process_task = SSHOperator(
            task_id=f"run_{num}",
            ssh_conn_id="your_ssh_conn_id",
            command=f"run_something.sh {num} {item}"
        )
        save_config_task >> process_task

注意：旧版方案有局限性，需要确保DAG触发前变量已存在，或者在触发后重新解析DAG，实际使用中更推荐升级Airflow到2.2+使用动态任务映射。

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