Airflow 2.0.2中如何获取Python函数返回值供EMR Operator使用

XCom跨任务传值的方案完全可行，你确实需要先通过PythonOperator运行get_secret函数：Airflow中PythonOperator执行的函数如果有return返回值，会被自动推送到XCom中，后续创建EMR集群的任务只需要在执行阶段拉取该值，即可完成S3路径拼接。

注意两个原代码的易错点：

1. 原get_secret函数中Variable.get(region_name)写法错误，变量名需要加字符串引号，正确写法为Variable.get("region_name")
2. 原函数的try块没有异常捕获和兜底返回逻辑，建议补上避免任务无报错信息直接失败

import json
import boto3
from datetime import datetime
from airflow import DAG
from airflow.models import Variable
from airflow.operators.python import PythonOperator
from airflow.providers.amazon.aws.operators.emr import EmrCreateJobFlowOperator

# 修正后的密钥获取函数
def get_secret(**context):
    secret_name = Variable.get("secret_name")
    region_name = Variable.get("region_name")
    # 初始化Secrets Manager客户端
    session = boto3.session.Session()
    client = session.client(service_name='secretsmanager', region_name=region_name)
    
    try:
        get_secret_value_response = client.get_secret_value(SecretId=secret_name)
        if 'SecretString' in get_secret_value_response:
            secret_str = get_secret_value_response['SecretString']
            secret = json.loads(secret_str)            
            airflow_path = secret["airflow_path"]
            # 函数返回值会自动存入XCom，key为return_value
            return airflow_path
    except Exception as e:
        raise ValueError(f"获取Secrets Manager密钥失败: {str(e)}")

# 构造EMR Spark步骤的函数，从XCom拉取airflow_path
def build_spark_steps(**context):
    # 从XCom拉取上游get_secret任务的返回值，ti是task instance的缩写
    airflow_path = context['ti'].xcom_pull(task_ids='get_airflow_path_from_secret')
    spark_steps = [
        {
            'Name': 'Spark-Submit Command',
            "ActionOnFailure": "CONTINUE",
            'HadoopJarStep': {
                "Jar": "command-runner.jar",
                "Args": [
                    'spark-submit',
                    '--py-files',
                    f's3://{airflow_path}-pyspark/pitchbook/config.zip,s3://{airflow_path}-pyspark/pitchbook/jobs.zip,s3://{airflow_path}-pyspark/pitchbook/DDL.zip',
                    f's3://{airflow_path}-pyspark/pitchbook/main.py'
                ],
            },
        }
    ]
    return spark_steps

# DAG默认参数
default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 1, 1),
    'retries': 0
}

with DAG(
    dag_id='emr_cluster_create_demo',
    default_args=default_args,
    schedule_interval=None,
    catchup=False
) as dag:
    # 任务1：从Secrets Manager拉取airflow_path
    get_path_task = PythonOperator(
        task_id='get_airflow_path_from_secret',
        python_callable=get_secret,
        provide_context=True
    )

    # 任务2：构造Spark步骤配置
    build_steps_task = PythonOperator(
        task_id='build_spark_steps',
        python_callable=build_spark_steps,
        provide_context=True
    )

    # 任务3：创建EMR集群，从XCom拉取构造好的SPARK_STEPS
    create_emr_task = EmrCreateJobFlowOperator(
        task_id='create_emr_cluster',
        aws_conn_id='aws_default', # 提前在Airflow连接中配置好AWS凭证
        # 这里的job_flow_overrides只演示SPARK_STEPS部分，其余EMR配置（比如实例配置、EC2密钥等）按你的实际需求补全即可
        job_flow_overrides={
            'Name': 'pitchbook-spark-cluster',
            'ReleaseLabel': 'emr-6.10.0',
            'Instances': {
                'InstanceGroups': [], # 补全你的实例组配置
                'KeepJobFlowAliveWhenNoSteps': False,
                'TerminationProtected': False
            },
            # 直接用Jinja模板拉取上游build_spark_steps任务的XCom返回值
            'Steps': "{{ ti.xcom_pull(task_ids='build_spark_steps') }}",
            'JobFlowRole': 'EMR_EC2_DefaultRole',
            'ServiceRole': 'EMR_DefaultRole'
        }
    )

    # 定义任务依赖顺序
    get_path_task >> build_steps_task >> create_emr_task

• 不要在DAG顶层（with DAG块之外、或者任务定义的参数里直接写死拉XCom的逻辑）尝试拉取XCom值：DAG解析阶段任务还没运行，XCom里没有对应值，会直接报错。所有XCom拉取逻辑要么放在任务执行的函数里，要么用Jinja模板字符串让Airflow在任务执行时自动渲染。
• 确保你的Airflow环境已经安装了Amazon provider包，如果运行时报错找不到EmrCreateJobFlowOperator，执行pip install apache-airflow-providers-amazon安装对应依赖即可。
• 确保Airflow配置的AWS连接对应的账号，有Secrets Manager的读权限、EMR集群创建权限，否则会报权限错误。

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