基于Airflow EMR模板，如何传入POST JSON中的变量？

你的问题核心在于全局定义的SPARK_TEST_STEPS在DAG加载阶段就会被解析，此时还没有dag_run上下文，无法获取POST传入的配置变量。要解决这个问题，我们可以利用Airflow的模板化特性，将步骤定义延迟到任务执行阶段（此时dag_run已经存在）来渲染变量。

具体解决方案 #

Airflow的EmrAddStepsOperator的steps参数支持Jinja2模板渲染，我们可以直接在步骤的Args中用Jinja语法引用dag_run.conf中的变量。下面是修改后的完整代码：

from datetime import timedelta
import airflow
from airflow import DAG
from airflow.contrib.operators.emr_create_job_flow_operator import EmrCreateJobFlowOperator
from airflow.contrib.operators.emr_add_steps_operator import EmrAddStepsOperator
from airflow.contrib.sensors.emr_step_sensor import EmrStepSensor
from airflow.contrib.operators.emr_terminate_job_flow_operator import EmrTerminateJobFlowOperator

DEFAULT_ARGS = {
    'owner': 'Airflow',
    'depends_on_past': False,
    'start_date': airflow.utils.dates.days_ago(2),
    'email': ['airflow@example.com'],
    'email_on_failure': False,
    'email_on_retry': False
}

JOB_FLOW_OVERRIDES = {
    'Name': 'PiCalc'
}

dag = DAG(
    'emr_job_flow_manual_steps_dag',
    default_args=DEFAULT_ARGS,
    dagrun_timeout=timedelta(hours=2),
    schedule_interval='0 3 * * *'
)

cluster_creator = EmrCreateJobFlowOperator(
    task_id='create_job_flow',
    job_flow_overrides=JOB_FLOW_OVERRIDES,
    aws_conn_id='aws_default',
    emr_conn_id='emr_default',
    dag=dag
)

# 直接在EmrAddStepsOperator中使用Jinja模板渲染步骤，动态获取dag_run.conf中的变量
step_adder = EmrAddStepsOperator(
    task_id='add_steps',
    job_flow_id="{{ task_instance.xcom_pull('create_job_flow', key='return_value') }}",
    aws_conn_id='aws_default',
    steps=[
        {
            'Name': 'calculate_pi',
            'ActionOnFailure': 'CONTINUE',
            'HadoopJarStep': {
                'Jar': 'command-runner.jar',
                'Args': [
                    '/usr/lib/spark/bin/run-example',
                    'SparkPi',
                    "{{ dag_run.conf.get('var_1', 'default_var1') }}",  # 用get方法避免变量不存在报错，可设置默认值
                    "{{ dag_run.conf.get('var_2', 'default_var2') }}",
                    "{{ dag_run.conf.get('var_3', 'default_var3') }}",
                    '10'
                ]
            }
        }
    ],
    dag=dag
)

step_checker = EmrStepSensor(
    task_id='watch_step',
    job_flow_id="{{ task_instance.xcom_pull('create_job_flow', key='return_value') }}",
    step_id="{{ task_instance.xcom_pull('add_steps', key='return_value')[0] }}",
    aws_conn_id='aws_default',
    dag=dag
)

cluster_remover = EmrTerminateJobFlowOperator(
    task_id='remove_cluster',
    job_flow_id="{{ task_instance.xcom_pull('create_job_flow', key='return_value') }}",
    aws_conn_id='aws_default',
    dag=dag
)

cluster_creator.set_downstream(step_adder)
step_adder.set_downstream(step_checker)
step_checker.set_downstream(cluster_remover)

关键说明 #

1. 模板化渲染时机：Jinja表达式{{ dag_run.conf.get('var_x') }}会在任务执行阶段（而非DAG加载阶段）被解析，此时dag_run已经包含了POST请求传入的配置信息。
2. 默认值处理：使用dag_run.conf.get('var_x', 'default_value')可以避免因POST请求未传入某个变量导致的报错，你可以根据实际需求设置合适的默认值。
3. 触发DAG的POST请求格式：触发时需要将变量放在conf字段中，示例JSON如下：

   {
     "conf": {
       "var_1": "100",
       "var_2": "200",
       "var_3": "300"
     }
   }
   

为什么原来的方式不可行？ #

你最初的代码中SPARK_TEST_STEPS是全局变量，Airflow在加载DAG文件时就会执行这段代码，此时还没有任何dag_run实例存在，自然无法获取kwargs['dag_run'].conf中的变量，会直接抛出KeyError或NameError。

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