How to retrive data from python function and use it in a emr operator

Viewed 134

Airflow version :2.0.2 Trying to create Emr cluster, by retrying data from AWS secrets manager.

I am trying to write an airflow dag and, my task is to get data from this get_secret function and use it in Spark_steps

def get_secret():
    secret_name =  Variable.get("secret_name")
    region_name = Variable.get(region_name)
    # Create a Secrets Manager client
    session = boto3.session.Session()
    client = session.client(service_name='secretsmanager', region_name=region_name)
    account_id = boto3.client('sts').get_caller_identity().get('Account')
    
    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"]
            return  airflow_path 

... I need to use "airflow_path" return value in below spark_steps Spark_steps:

SPARK_STEPS = [

    {
        'Name': 'Spark-Submit Command',
        "ActionOnFailure": "CONTINUE",
        'HadoopJarStep': {
            "Jar": "command-runner.jar",
            "Args": [
                'spark-submit',
                '--py-files',
                's3://'+airflow_path+'-pyspark/pitchbook/config.zip,s3://'+airflow_path+'-pyspark/pitchbook/jobs.zip,s3://'+airflow_path+'-pyspark/pitchbook/DDL.zip',
                's3://'+airflow_path+'-pyspark/pitchbook/main.py'
                    ],
                        },
    },

I saw on the internet I need to use Xcom, is this right ?, and do I need to run this function in python operator first and then get the value. please provide an example as I am a newbie.

Thanks for your help. Xi

1 Answers
Related