Airflow自定义Operator中修改XcomArg默认XCom推送key的方法

你之前手动创建XComArg指定自定义key不生效，核心原因是execute方法中仅通过return返回值时，Airflow默认只会把返回结果推送到固定的return_value键，你没有显式往自定义的key推送数据，自然无法通过指定key拿到对应值。你可以通过以下两种方案实现需求：

方案1：修改自定义Operator内部逻辑（推荐，兼容原有写法） #

直接在你的FileToAzureBlobOperator类中做两处调整，既兼容历史对extract_list键的引用，也能保留.output属性的便捷用法：

from airflow.models.baseoperator import BaseOperator
from airflow.models.xcom_arg import XComArg

class FileToAzureBlobOperator(BaseOperator):
    # 原有__init__等逻辑不需要修改，仅新增覆写output属性的逻辑
    @property
    def output(self) -> XComArg:
        # 让.output默认返回key为extract_list的XCom值
        return XComArg(self, key="extract_list")
    
    def execute(self, context):
        # 你的业务逻辑，生成extract_list、最大ID、时间戳等数据
        extract_list = ["/test/file1", "/test/file2"]
        max_id = 1024
        exec_timestamp = "2024-05-01 12:00:00"
        
        # 显式推送所有需要的XCom值
        context["ti"].xcom_push(key="extract_list", value=extract_list)
        context["ti"].xcom_push(key="max_id", value=max_id)
        context["ti"].xcom_push(key="exec_timestamp", value=exec_timestamp)
        
        # 不需要return也可以，如果你还需要保留return_value的推送可以保留return语句
        return extract_list

修改完成后，你原来的DAG写法不需要做任何调整，extract.output会自动指向extract_list对应的XCom值，原有其他位置对extract_list键的引用也完全不受影响。

方案2：无需修改Operator，下游直接指定key拉取 #

如果你暂时不想调整自定义Operator的代码，也可以直接在DAG中显式指定XCom的key拉取对应值：

extract = FileToAzureBlobOperator(
    task_id="extract-test",
    remote_directories=["/input/test"],
    subfolders=["test", "raw"],
    params={
        "start": "{{ data_interval_start }}",
        "end": "{{ data_interval_end }}",
    },
)

transform = PrepareParquetOperator(
    task_id="transform-test",
    # 直接指定要拉取的XCom key，不需要用.output属性
    input_files=XComArg(operator=extract, key="extract_list"),
    output_folder="test/staging",
    custom_transform_script="scripts.common.test",
    partition_columns=["date_id"],
)

注意该方案需要你确保execute方法中已经显式调用xcom_push往extract_list键推送了对应值，否则会拉取到空值。

注意事项 #

• 如果你使用的是Airflow 2.0以下版本，对XComArg的支持不完善，可以直接用模板语法获取值：{{ ti.xcom_pull(task_ids='extract-test', key='extract_list') }}
• 推送大体积的列表类数据时，注意不要超过你使用的XCom后端的存储上限。

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