Airflow 2.x中ConsumeFromTopicOperator的XCOM_PUSH功能失效问题

问题现象 #

使用ConsumeFromTopicOperator调用的函数无法通过xcom_push推送值，其他任务无法获取：

• 尝试在函数内调用ti.xcom_push会触发运行时错误
• 直接返回值后，用ti.xcom_pull(task_ids='CONSUME_FROM_TOPIC_02')获取结果始终为None

问题原因 #

1. ConsumeFromTopicOperator的apply_function被设计为单独处理每条消息，且operator不会自动将函数返回值作为XCom推送
2. apply_function的执行上下文与operator主流程隔离，直接在其中调用xcom_push会因上下文缺失或实例访问错误触发异常

解决方法 #

利用operator提供的results_callback参数，在所有消息处理完成后统一推送XCom：

修改后的示例代码 #

from airflow_provider_kafka.operators.consume_from_topic import ConsumeFromTopicOperator
from airflow.utils.context import get_current_context

# 仅负责处理单条消息，返回需要传递的数据
def process_message(message):
    msg = message.value()
    return msg['my_cool_value']

# 回调函数：接收所有消息处理结果，统一推送XCom
def push_results_to_xcom(results):
    ctx = get_current_context()
    ti = ctx["ti"]
    # 根据需求推送单条结果或全部结果列表
    ti.xcom_push(key='data', value=results[0] if results else None)

ConsumeFromTopicOperator(
    dag=dag,
    task_id="CONSUME_FROM_TOPIC_02",
    topics=['my.Events'],
    apply_function=process_message,  # 直接传入函数对象，避免字符串路径的上下文问题
    results_callback=push_results_to_xcom,
    consumer_config=kafka_config,
    commit_cadence="end_of_batch",
    poll_timeout=20,
    max_messages=1,
    max_batch_size=100,
)

取值方式 #

在后续任务中通过以下代码获取XCom值：

ti.xcom_pull(task_ids='CONSUME_FROM_TOPIC_02', key='data')

说明 #

• results_callback在operator主执行流程中运行，拥有完整的Airflow上下文，可正常调用XCom相关方法
• apply_function专注于消息处理逻辑，返回的所有结果会被收集后传入results_callback
• 若需要传递多条消息的处理结果，可直接将results列表推送到XCom，无需取第一个元素

内容的提问来源于stack exchange，提问作者Mikell van der Laan