Airflow自定义Operator返回值如何正确传递给下游Operator作为输入

前置说明 #

你当前开发的HelloWorldOperator已经符合返回值自动存入XCom的要求，绝大多数场景下不需要修改Operator核心逻辑，仅需调整DAG代码和少量Operator配置即可实现上下游参数传递。

注意你提供的代码存在参数不匹配问题：Operator初始化接收的参数名为some_input，但DAG中调用时传递的是name参数，需要先修正该问题，否则会触发运行错误。

实现方式（两种可选，均为业界规范写法） #

方案1：Jinja模板传参（兼容性最好，全Airflow版本支持） #

这是最常用的实现方式，仅需两步操作：

1. 给自定义Operator添加template_fields属性，将需要接收上游传值的字段加入列表，让Airflow知道该字段需要做模板渲染：

from airflow.models.baseoperator import BaseOperator

class HelloWorldOperator(BaseOperator):
    # 新增这一行，把需要支持模板传参的字段放进去
    template_fields = ("some_input",)

    def __init__(
            self,
            some_input: str,
            **kwargs) -> None:
        super().__init__(**kwargs)
        self.some_input = some_input

    def execute(self, context):
        # 原有业务逻辑无需修改
        token = "MyGeneratedToken"
        return token

2. 在DAG中给下游Operator的对应字段传入XCom拉取模板即可：

import logging
from airflow import DAG
from airflow.utils.dates import days_ago
from custom_operators.hello_world import HelloWorldOperator

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 0,
}

dag = DAG("hello_world_test",
    description='Testing out a operator',
    start_date=days_ago(1),
    schedule_interval=None,
    catchup=False,
    default_args=default_args)

get_token = HelloWorldOperator(
    task_id='hello_world_task_1',
    some_input='My input to generate a token',
    dag=dag
)

run_this = HelloWorldOperator(
    task_id='hello_world_task_2',
    # 直接用模板语法拉取上游任务返回的XCom值
    some_input="{{ ti.xcom_pull(task_ids='hello_world_task_1') }}",
    dag=dag
)

get_token >> run_this

方案2：TaskFlow API传参（Airflow 2.0+支持，写法更简洁） #

如果你的用户使用的是Airflow 2.0及以上版本，可以用TaskFlow的原生能力自动处理XCom推拉，不需要手动写模板语法：

from airflow.decorators import dag
from airflow.utils.dates import days_ago
from custom_operators.hello_world import HelloWorldOperator

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 0,
}

@dag(
    "hello_world_test",
    description='Testing out a operator',
    start_date=days_ago(1),
    schedule_interval=None,
    catchup=False,
    default_args=default_args
)
def hello_world_dag():
    get_token = HelloWorldOperator(
        task_id='hello_world_task_1',
        some_input='My input to generate a token',
    )

    run_this = HelloWorldOperator(
        task_id='hello_world_task_2',
        # 直接传入上游任务的output属性即可，Airflow自动处理XCom拉取
        some_input=get_token.output,
    )

    get_token >> run_this

dag = hello_world_dag()

常见问题说明 #

• 上下游是不同Operator也可以用完全相同的逻辑传值，只需要保证下游Operator要接收值的字段已经加入它自身的template_fields列表即可。
• 不要尝试在DAG解析阶段直接把上游返回值赋值给普通Python变量，DAG解析时任务还未运行，XCom值还未生成，必须用模板或者TaskFlow的懒加载机制实现运行时传值。
• 不需要修改Operator的execute方法逻辑，只要返回值会自动存入XCom，就可以用上述方式拉取。

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