如何在Airflow的Spark Steps中引入用户配置的日期参数？

核心问题说明 #

你用os.environ传递变量失败的原因是：Airflow任务运行在独立的worker进程中，触发DAG时传入的JSON参数不会自动注入进程环境变量，跨任务的环境变量传递也不具备可靠性。以下是两种落地可行的解决方案：

方案1：利用Airflow Jinja2模板渲染直接替换 #

EmrAddStepsOperator的steps参数原生支持Jinja2模板渲染，可直接引用DAG触发时传入的conf参数完成路径替换：

1. 定义带模板占位符的Spark步骤配置

SPARK_STEPS_ADHOC = [
    {
        "Name": "Adhoc Spark Job",
        "ActionOnFailure": "CONTINUE",
        "HadoopJarStep": {
            "Jar": "command-runner.jar",
            "Args": [
                "spark-submit",
                "--class", "com.example.AdhocJob",
                "s3://your-bucket/jobs/adhoc-job.jar",
                "--input-path", "s3://your-bucket/input/{{ dag_run.conf['start_date'] }}-{{ dag_run.conf['end_date'] }}",
                "--output-path", "s3://your-bucket/output/{{ dag_run.conf['start_date'] }}-{{ dag_run.conf['end_date'] }}"
            ]
        }
    }
]

2. 配置EmrAddStepsOperator
   直接使用上述配置，Airflow会在任务执行阶段自动替换模板中的变量：

from airflow.providers.amazon.aws.operators.emr_add_steps import EmrAddStepsOperator

add_emr_steps = EmrAddStepsOperator(
    task_id="add_adhoc_spark_step",
    job_flow_id="your-emr-cluster-id",  # 可通过Airflow Variable或其他方式传入
    steps=SPARK_STEPS_ADHOC,
    dag=dag
)

触发DAG时传入指定格式的JSON参数即可（示例：{"type":"adhoc", "start_date":"2022-01-01", "end_date":"2022-05-01"}）。

方案2：通过PythonOperator动态生成步骤配置 #

如果需要处理更复杂的参数逻辑，可先解析JSON参数，动态替换S3路径后再传递给EmrAddStepsOperator：

1. 定义步骤模板

import json
from airflow.models import Variable
from airflow.operators.python import PythonOperator
from airflow.providers.amazon.aws.operators.emr_add_steps import EmrAddStepsOperator

# 带占位符的基础步骤模板
SPARK_STEPS_TEMPLATE = [
    {
        "Name": "Adhoc Spark Job",
        "ActionOnFailure": "CONTINUE",
        "HadoopJarStep": {
            "Jar": "command-runner.jar",
            "Args": [
                "spark-submit",
                "--class", "com.example.AdhocJob",
                "s3://your-bucket/jobs/adhoc-job.jar",
                "--input-path", "s3://your-bucket/input/{STARTDATE}-{ENDDATE}",
                "--output-path", "s3://your-bucket/output/{STARTDATE}-{ENDDATE}"
            ]
        }
    }
]

2. 编写参数解析与步骤生成函数

def generate_spark_steps(**context):
    # 从DAG触发conf中提取JSON参数
    dag_conf = context["dag_run"].conf
    start_date = dag_conf["start_date"]
    end_date = dag_conf["end_date"]
    
    # 替换模板中的占位符
    steps_json = json.dumps(SPARK_STEPS_TEMPLATE)
    steps_json = steps_json.replace("{STARTDATE}", start_date).replace("{ENDDATE}", end_date)
    final_steps = json.loads(steps_json)
    
    # 将生成的步骤推送到XCom，供后续任务调用
    context["ti"].xcom_push(key="adhoc_spark_steps", value=final_steps)

3. 串联任务流程

# 生成步骤的Python任务
generate_steps = PythonOperator(
    task_id="generate_adhoc_spark_steps",
    python_callable=generate_spark_steps,
    provide_context=True,
    dag=dag
)

# EMR添加步骤任务
add_emr_steps = EmrAddStepsOperator(
    task_id="add_adhoc_spark_step",
    job_flow_id=Variable.get("emr_cluster_id"),
    steps="{{ ti.xcom_pull(key='adhoc_spark_steps') }}",
    dag=dag
)

generate_steps >> add_emr_steps

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