如何借助Beam与Airflow动态获取GCS文件路径并解决XCom报错

错误原因 #

直接调用collect_paths_job.xcom_pull()会触发TypeError，因为Task实例的xcom_pull()方法必须接收Airflow执行上下文（context）作为参数——该上下文包含TaskInstance、DAG运行时信息等，仅在任务执行阶段可用，无法在DAG定义阶段直接调用。

解决方案 #

推荐方案：使用Airflow 2.2+ 动态任务映射（Dynamic Task Mapping） #

Cloud Composer默认支持Airflow 2.x，动态任务映射可以直接基于上游任务的XCom输出自动生成并发任务，无需手动遍历路径列表，是最简洁的实现方式。

步骤1：完善路径采集任务

确保路径采集任务正确将结果推送到XCom（或Airflow Variable）：

from airflow import DAG
from airflow.operators.python import PythonOperator
from google.cloud import storage
from airflow.models import Variable
import json
from datetime import datetime

def collect_paths(**context):
    # 初始化GCS客户端
    client = storage.Client()
    bucket = client.get_bucket("staging")
    
    # 过滤目标文件：包含main_folder且以.txt结尾
    blobs = bucket.list_blobs()
    txt_paths = [
        f"gs://{bucket.name}/{blob.name}" 
        for blob in blobs 
        if blob.name.endswith(".txt") and "/main_folder/" in blob.name
    ]
    
    # 选项1：推送到XCom（供动态任务映射使用）
    context["ti"].xcom_push(key="txt_paths", value=txt_paths)
    
    # 选项2：导出到Airflow Variable（满足你的第二个目标）
    Variable.set("daily_txt_paths", json.dumps(txt_paths), serialize_json=True)
    
    return txt_paths

with DAG(
    dag_id="gcs_file_processor",
    schedule_interval="@daily",
    start_date=datetime(2024, 1, 1),
    catchup=False
) as dag:
    collect_paths_job = PythonOperator(
        task_id="collect_paths",
        python_callable=collect_paths,
        provide_context=True,
        do_xcom_push=True
    )

步骤2：动态生成文件处理任务

使用partial() + expand()实现动态任务映射，为每个文件生成独立的处理任务：

def process_single_file(file_path, **context):
    # 单文件处理逻辑：可调用Apache Beam/Dataflow或自定义处理逻辑
    print(f"Processing file: {file_path}")
    # 示例：启动Beam pipeline处理单个文件
    # from apache_beam.options.pipeline_options import PipelineOptions
    # pipeline_options = PipelineOptions([
    #     "--runner=DataflowRunner",
    #     "--project=your-gcp-project",
    #     "--temp_location=gs://staging/temp"
    # ])
    # with beam.Pipeline(options=pipeline_options) as p:
    #     p | beam.io.ReadFromText(file_path) | ...  # 后续处理步骤

# 动态生成任务：基于collect_paths_job的XCom输出
process_files = PythonOperator.partial(
    task_id="process_single_file",
    python_callable=process_single_file,
    provide_context=True
).expand(op_kwargs=[{"file_path": path} for path in collect_paths_job.output])

# 设置任务依赖
collect_paths_job >> process_files

备选方案：在Operator内部通过Context获取XCom（旧版Airflow兼容） #

如果使用Airflow 2.2以下版本，可在单独的PythonOperator中通过上下文获取XCom，再触发子任务（需结合SubDAG或TriggerDagRunOperator，复杂度较高）：

def generate_process_tasks(**context):
    ti = context["ti"]
    # 从XCom获取路径列表
    txt_paths = ti.xcom_pull(task_ids="collect_paths", key="txt_paths")
    
    # 此处可通过TriggerDagRunOperator触发多个子DAG，或调用SubDAG
    # 示例：遍历路径并触发单个任务
    for idx, path in enumerate(txt_paths):
        trigger_task = TriggerDagRunOperator(
            task_id=f"trigger_process_{idx}",
            trigger_dag_id="single_file_processor_dag",
            conf={"file_path": path},
            dag=context["dag"]
        )
        trigger_task.execute(context=context)

generate_tasks_job = PythonOperator(
    task_id="generate_process_tasks",
    python_callable=generate_process_tasks,
    provide_context=True,
    dag=dag
)

collect_paths_job >> generate_tasks_job

关键注意事项 #

1. Airflow Variable序列化：存储列表到Variable时需用json.dumps()序列化，读取时用json.loads()反序列化，避免格式错误。
2. 动态任务映射限制：每个DAG的动态任务数量受Airflow配置限制（默认最大1000个），若文件数量过大需调整配置或分批处理。
3. XCom存储限制：XCom默认存储大小有限（约48KB），若路径列表过大，建议优先使用Airflow Variable或GCS存储路径文件。

内容的提问来源于stack exchange，提问作者alex-mont