Airflow set DAG options conditionally

Viewed 31

I'm implementing a python script to create a bunch of Airflow dag based on json config files. One json config file contains all the fields to be used in DAG(), and the last three fields are optional(will use global default if not set).

{
  "owner": "Mike",
  "start_date": "2022-04-10",
  "schedule_interval": "0 0 * * *",
  "on_failure_callback": "slack",
  "is_paused_upon_creation": False,
  "catchup": True
}

Now, my question is how to create the DAG conditionally? Since the last three option on_failure_callback, is_paused_upon_creation and catchup is optional, wonder what's the best way to use them in DAG()?

One solution_1 I tried is to use default_arg=optional_fields, and add optional fields into it with an if statement. However, it doesn't work. The DAG is not taking these three optional fields' values.

def create_dag(name, config):
  # config is a dict that generate from the json config file

  optional_fields = {
    'owner': config['owner']
  }
  if 'on_failure_callback' in config:
    optional_fields['on_failure_callback'] = partial(xxx(config['on_failure_callback']))
  if 'is_paused_upon_creation' in config:
    optional_fields['is_paused_upon_creation'] = config['is_paused_upon_creation']

  dag = DAG(
    dag_id=name,
    start_date=datetime.strptime(config['start_date'], '%Y-%m-%d'),
    schedule_interval=config['schedule_interval'],
    default_args=optional_fields
  )

Then, I tried solution_2 with **optional_fields, but got an error TypeError: __init__() got an unexpected keyword argument 'owner'

  dag = DAG(
    dag_id=name,
    start_date=datetime.strptime(config['start_date'], '%Y-%m-%d'),
    schedule_interval=config['schedule_interval'],
    **optional_fields
  )

Then solution_3 works as the following.

def create_dag(name, config):
  # config is a dict that generate from the json config file

  default_args = {
    'owner': config['owner']
  }

  optional_fields = {}
  if 'on_failure_callback' in config:
    optional_fields['on_failure_callback'] = partial(xxx(config['on_failure_callback']))
  if 'is_paused_upon_creation' in config:
    optional_fields['is_paused_upon_creation'] = config['is_paused_upon_creation']

  dag = DAG(
    dag_id=name,
    start_date=datetime.strptime(config['start_date'], '%Y-%m-%d'),
    schedule_interval=config['schedule_interval'],
    default_args=default_args
    **optional_fields
  )

However, I'm confused 1) which fields should be set in optional_fields vs default_args? 2) is there any other way to achieve it?

0 Answers
Related