在 BeamRunPythonPipelineOperator 中使用 Airflow DAG 运行配置答案

【问题标题】：Using Airflow DAG run conf in BeamRunPythonPipelineOperator在 BeamRunPythonPipelineOperator 中使用 Airflow DAG 运行配置
【发布时间】：2021-12-26 07:31:55
【问题描述】：

背景

我正在尝试将 Apache Beam 管道作为 Airflow DAG 的一部分运行。 DAG 运行是使用user_id 参数手动触发的。通常，在我的 PythonOperator 中，我可以访问我的 DAG 运行 conf 作为 python_callable 的参数的一部分。例如：

def python_operator_example(ds, conf, dag_run):
    # Can easily access my user_id parameter
    user_id = dag_run.conf['user_id']

问题

现在，我创建了一个BeamRunPythonPipelineOperator，我也想在其中访问user_id。

process_data = BeamRunPythonPipelineOperator(
    dag=dag,
    task_id="process_data",
    pipeline_options={
        "input": "/Users/myuser/Desktop/test.txt",
        "output": "/Users/myuser/Desktop/foo/"
    },
    py_file='pipeline/process_data.py'
)

但是，我似乎找不到在 BeamRunPythonPipelineOperator 中访问 DAG 运行 conf 的方法。似乎没有办法在运行时将输入从 dag_run.conf 传递到 Beam 管道。理想情况下，我想做这样的事情：

## 1. Define a function to set pipeline inputs dynamically using dag_run

def create_beam_pipeline_options(ds, conf, dag_run):
    return {
        "user_id": dag_run.conf['user_id'],
        "input": "/Users/myuser/Desktop/test.txt",
        "output": "/Users/myuser/Desktop/foo/"
    }

## 2. Pass the function into operator to generate options when pipeline task is invoked

process_data = BeamRunPythonPipelineOperator(
    dag=dag,
    task_id="process_data",
    pipeline_options=create_beam_pipeline_options,
    py_file='pipeline/process_data.py'
)

问题

将像dag_run.conf 这样在运行时定义为管道选项的值传递到BeamRunPythonPipelineOperator 的正确方法是什么？

我已经尝试过的

我检查了管道的run 函数中的argv 是否包含dag_run.conf，但它没有。
我查看了是否可以使用 XCom 通过 BeamRunPythonPipelineOperator 传递来自先前 DAG 任务的数据，但没有找到任何东西。
我探索了一些解决方案，允许我为每个 user_id 动态创建一个新的 DAG，但这似乎是一种浪费和反模式。
我阅读了 docs 和 code 的建议

如果这是一个愚蠢的问题，我提前致以诚挚的歉意，我对使用 Beam / Airflow 很陌生！

【问题讨论】：

标签： python airflow apache-beam dataflow data-processing

【解决方案1】：

我通过使用PythonOperator调用一个@ 987654322实现了这一@。它看起来是这样的。 P>

from airflow.providers.apache.beam.operators.beam import (
    BeamHook,
    BeamRunnerType
)

beam_hook = BeamHook(BeamRunnerType.DirectRunner)

def process_data(ds, conf, dag_run):
    beam_hook.start_python_pipeline(
        variables={"user_id": dag_run.conf["user_id"]},
        py_file='dags/pipeline/pipline/process_data.py',
        py_options=[],
        py_requirements=['apache-beam'],
        py_interpreter='python3',
        py_system_site_packages=False,
    )

【讨论】：