How to use the values of input json configurations passed to trigger the airflow job inside the dag?

user_10495955 0
Viewed 203

I am new to airflow and working on a use case where I need to read the input json configs that are passed to an airflow dag and construct a string based on the configs read, this string will be used as the name of cluster we are creating in GCP data proc.

Ex : Input Json to dag

{ "x" : "data", "y": "engg", "z" : "usecase" }

I am expecting the cluster name to be "my-cluster-for-data-engg-usecase" with the code below but I am getting an error like it cannot recognize "dag_run". Any help would be appreciated.

from datetime import datetime
from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.operators.python_operator import PythonOperator
from airflow import models
from airflow.utils.dates import days_ago

CONN_ID = 'blah'
PROJECT_ID = 'xyz'
REGION = 'us-east4'
CLUSTER_NAME = "my-cluster-for-"+dag_run.conf['x']+dag_run.conf['y']+dag_run.conf['z']
with models.DAG('simple-python-dag', start_date=days_ago(1), schedule_interval=None) as dag:


    create_cluster_spark = XyzCreateClusterOperator(task_id='create_cluster_spark',
                                                    cluster_name=CLUSTER_NAME,
                                                    location=REGION,
                                                    gcp_conn_id=CONN_ID)

    create_cluster_spark
airflow google-cloud-platform python
1 Answers

The dag_run is only available in an active DAG run (i.e when you've triggered the DAG and it is running) - specifying dag_run.conf as an object will not work as the Airflow scheduler will parse your dags (every 30s by default) and no such object exists.

There are two ways access the dag_run:

  1. Use the templating system
CLUSTER_NAME = "my-cluster-for-{{ dag_run.conf['x'] }}{{ dag_run.conf['y'] }}{{ dag_run.conf['z'] }}"

This will only work if the cluster_name field is a templated field in the XyzCreateClusterOperator Operator. You can check if a field is templated by looking at the docstring, or looking at the Airflow documentation for operators.

If XyzCreateClusterOperator is a third party provider you can check the documentation for the provider here

You can see all available items in a template here (dag_run is one of many things you can access using this template syntax).

  1. You can access the dag_run inside a python operator
import pendulum
from airflow.decorators import dag, task

@dag(
    schedule_interval=None,
    start_date=pendulum.datetime(2022, 4, 18),
    catchup=False,
)
def example_dag():
    @task()
    def dag_run_example_task(**context):
        context["dag_run"].conf.get("x")
        context["dag_run"].conf.get("y")
        context["dag_run"].conf.get("z")

If cluster_name is not a templated field, you can use the Python Operator to wrap around your operator and substitute in:

import pendulum
from airflow.decorators import dag, task

@dag(
    schedule_interval=None,
    start_date=pendulum.datetime(2022, 4, 18),
    catchup=False,
)
def example_dag():
    @task()
    def dag_run_example_task(**context):
        cluster_name = "my-cluster-for-"+context["dag_run"].conf.get("x")+context["dag_run"].conf.get("y")+context["dag_run"].conf.get("z")

        create_cluster_spark = XyzCreateClusterOperator(
            task_id='create_cluster_spark',
            cluster_name=cluster_name,
            location=REGION,
            gcp_conn_id=CONN_ID,
        )

        create_cluster_spark.execute(context=context)

The same objects available in the templates here are available in the context object passed to a PythonOperator (and a few more as well).

user_131322710
Related