Airflow如何向python_callable同时传递XCom上下文和自定义参数

前置依赖导入 #

首先确保你在DAG文件头部导入了所需模块：

import os
import random
import glob
from airflow.operators.python import PythonOperator

1. 支持XCom推送的callable函数 #

你注释掉的XCom推送逻辑是正确的，直接放开即可，完整函数如下：

def _select_random_data(**kwargs):
    data_dir = kwargs.get('data_path')
    logger = kwargs.get('logger')

    if not os.path.exists(data_dir):
        raise RuntimeError('data directory does not exist')

    pattern = f'{data_dir}/*.csv'
    csv_files = [x for x in glob.glob(pattern) if os.path.isfile(x)]
    if not csv_files:
        raise RuntimeError(f'no csv file found in directory {data_dir}')
    
    random_filename = random.choice(csv_files)
    logger.info(f'random file name: {random_filename}')

    # 方式1：显式调用xcom_push推送值，自定义key
    task_instance = kwargs['task_instance']
    task_instance.xcom_push(key='selected_csv_file', value=random_filename)

    # 方式2：直接return值，会自动推送到XCom的return_value键，无需显式调用xcom_push
    # return random_filename

注：额外补充了空CSV列表的异常判断，避免目录下没有CSV文件时random.choice报错。

2. PythonOperator任务定义 #

你原有任务定义基本正确，Airflow 2.x版本可以省略provide_context=True参数，框架会自动传递上下文，兼容写法如下：

select_file_task = PythonOperator(
    task_id='select_file',
    python_callable=_select_random_data,
    # Airflow 1.x 需要保留provide_context=True，2.x可省略
    provide_context=True,
    op_kwargs={
        'data_path': 'path1',
        'logger': logger,
    }
)

后续任务获取XCom值示例 #

如果后续是PythonOperator的任务，可以通过以下两种方式获取推送的CSV路径：

• 对应上面显式push的自定义key的获取方式：

def downstream_task(**kwargs):
    selected_file = kwargs['task_instance'].xcom_pull(key='selected_csv_file', task_ids='select_file')
    # 后续处理逻辑

• 对应上面return方式的获取方式：

def downstream_task(**kwargs):
    selected_file = kwargs['task_instance'].xcom_pull(task_ids='select_file')
    # 后续处理逻辑

如果是在模板语法中使用（比如BashOperator的bash_command参数），可以直接用模板变量：{{ ti.xcom_pull(task_ids='select_file', key='selected_csv_file') }}

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