Airflow开发：任务完成后将状态与task_id存入XCom字典

方案思路 #

借助Airflow的任务回调机制，在任务成功或失败时自动将task_id和对应状态写入XCom，后续可汇总成目标格式的字典。

具体实现代码 #

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

def task_18(ti):
    print("TASK 18 COMPLETE")

def task_19(ti):
    print("TASK 19 COMPLETE")

# 回调函数：将任务状态与task_id存入XCom
def task_state_callback(context):
    ti = context['ti']
    # 用task_id作为XCom的key，任务状态作为value存入
    ti.xcom_push(key=ti.task_id, value=ti.state)

# 定义DAG
with DAG(
    dag_id='task_state_xcom_dag',
    start_date=datetime(2024, 1, 1),
    schedule_interval=None,
    catchup=False
) as dag:

    opr_task_18 = PythonOperator(
        task_id='task.id.number.18',
        python_callable=task_18,
        on_success_callback=task_state_callback,
        on_failure_callback=task_state_callback
    )

    opr_task_19 = PythonOperator(
        task_id='task.id.number.19',
        python_callable=task_19,
        on_success_callback=task_state_callback,
        on_failure_callback=task_state_callback
    )

    # 任务依赖（按需调整）
    opr_task_18 >> opr_task_19

读取汇总状态字典 #

如果需要在后续任务中获取目标格式的状态字典，可添加以下任务：

def collect_task_states(ti):
    # 拉取当前DAG下所有任务的XCom数据
    xcom_data = ti.xcom_pull(dag_id=ti.dag_id, task_ids=None, key=None)
    # 整理成指定格式的字典
    state_dict = {k: v for k, v in xcom_data.items() if k in ['task.id.number.18', 'task.id.number.19']}
    print(state_dict)

collect_task = PythonOperator(
    task_id='collect_task_states',
    python_callable=collect_task_states
)

opr_task_19 >> collect_task

说明 #

• 回调函数task_state_callback通过context获取当前任务实例，自动捕获任务的成功/失败状态，并存入XCom。
• 后续任务通过xcom_pull拉取所有XCom数据，过滤出目标任务的状态后即可整理成你需要的字典格式。

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