如何将Airflow XCom返回值存入Python变量并生成S3Copy任务？

完全可以将Airflow的XCom值传入Python变量，但要注意代码执行时机：不能在DAG定义阶段直接调用task_instance.xcom_pull，因为此时任务还未执行，XCom数据尚未生成。需要在任务执行阶段获取XCom，或利用Airflow的动态任务映射功能来自动遍历返回值。

方法一：使用动态任务映射（Airflow 2.3+ 推荐） #

Airflow 2.3及以上版本支持动态任务映射，可直接基于上游任务返回的列表自动生成子任务，无需手动循环。

示例代码：

from airflow import DAG
from airflow.providers.amazon.aws.operators.s3 import S3ListPrefixesOperator, S3CopyObjectOperator
from datetime import datetime

with DAG(
    dag_id='s3_dynamic_copy',
    start_date=datetime(2024, 1, 1),
    schedule_interval=None,
    catchup=False
) as dag:
    # 1. 获取S3前缀列表
    list_prefixes = S3ListPrefixesOperator(
        task_id='list_s3_prefixes',
        bucket=get_bucket(),  # 替换为你的bucket获取逻辑
        prefix='my_key/',
        delimiter='/'
    )

    # 2. 动态生成S3Copy任务
    copy_objects = S3CopyObjectOperator.partial(
        task_id='copy_s3_objects',
        source_bucket_name=get_bucket(),
        dest_bucket_name='your_target_bucket',  # 替换为目标bucket
    ).expand(
        # 直接引用上游任务的XCom输出作为参数
        source_object_key=list_prefixes.output,
        # 自定义目标路径，这里给每个前缀添加copied/前缀
        dest_object_key=list_prefixes.output.map(lambda key: f"copied/{key}")
    )

    list_prefixes >> copy_objects

• partial()：设置所有子任务共享的固定参数
• expand()：接收上游返回的列表，每个元素对应一个独立的子任务
• list_prefixes.output：自动获取该任务的XCom返回值，无需手动调用xcom_pull

方法二：PythonOperator手动生成任务（Airflow 2.3以下版本） #

若Airflow版本较低，可通过PythonOperator在执行阶段获取XCom，再动态创建任务。

示例代码：

from airflow import DAG
from airflow.providers.amazon.aws.operators.s3 import S3ListPrefixesOperator, S3CopyObjectOperator
from airflow.operators.python import PythonOperator
from datetime import datetime

def create_copy_tasks(**context):
    # 从上下文获取TaskInstance，拉取上游XCom值
    ti = context['ti']
    bucket_keys = ti.xcom_pull(task_ids='list_s3_prefixes')
    
    # 遍历每个前缀，创建对应复制任务
    for key in bucket_keys:
        # 生成唯一的task_id，替换特殊字符避免报错
        task_id = f"copy_{key.replace('/', '_').strip('_')}"
        copy_task = S3CopyObjectOperator(
            task_id=task_id,
            source_bucket_name=get_bucket(),
            source_object_key=key,
            dest_bucket_name='your_target_bucket',
            dest_object_key=f"copied/{key}"
        )
        # 将当前Python任务与复制任务建立依赖
        context['task'] >> copy_task

with DAG(
    dag_id='s3_legacy_copy',
    start_date=datetime(2024, 1, 1),
    schedule_interval=None,
    catchup=False
) as dag:
    list_prefixes = S3ListPrefixesOperator(
        task_id='list_s3_prefixes',
        bucket=get_bucket(),
        prefix='my_key/',
        delimiter='/'
    )

    generate_copy_tasks = PythonOperator(
        task_id='generate_copy_tasks',
        python_callable=create_copy_tasks,
        provide_context=True  # 传递Airflow上下文，用于获取task和task_instance
    )

    list_prefixes >> generate_copy_tasks

• provide_context=True：将Airflow上下文（含TaskInstance、当前任务等）传递给Python函数
• 生成task_id时需替换特殊字符，确保ID唯一合法

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