CODEWITHSUNDEEP
airflowairflow-scheduler
Ray
Ray

Reputation: 3225

Execute single task AFTER dynamically-generated tasks via for-loop

Suppose I have the follow DAG (basic placeholder functions), that uses a for-loop to dynamically generate tasks (from iterating over a list):

from airflow import DAG
from airflow.operators.python_operator import PythonOperator

default_args = {
    'owner': 'ETLUSER',
    'depends_on_past': False,
    'start_date': datetime(2019, 12, 16, 0, 0, 0),
    'email': ['xxx@xxx.com'],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5)
}

dag = DAG('xxx', catchup=False,
          default_args=default_args, schedule_interval='0 */4 * * *')


# Some dummy function
def StepOne(x):
  print(x)


def StepTwo():
  print("Okay, we finished all of Step 1.")


some_list = [1, 2, 3, 4, 5, 6]


for t in some_list:
  task_id = f'FirstStep_{t}'
  task = PythonOperator(
      task_id=task_id,
      python_callable=StepOne,
      provide_context=False,
      op_kwargs={'x': str(t)},
      dag=dag
  )
  task

I want to introduce some additional task that's simply:

task2 = PythonOperator(
    task_id="SecondStep",
    python_callable=StepTwo,
    provide_context=False,
    dag=dag
  )

That runs only after all the steps in the first have finished. Linearly, this would be task >> task2

How do I go about doing this?

Upvotes: 2

Views: 2014

Answers (1)

Emma
Emma

Reputation: 9363

You can have task dependencies with array.

Do taskC after both taskA and taskB finished.

[taskA, taskB] >> taskC

or

Do taskB and taskC in parallel after taskA finished.

taskA >> [taskB, taskC]

as long as 1 side of upstream or downstream are non-array.

Thus, for your example,

task1 = []
for t in some_list:
    task_id = f'FirstStep_{t}'
    task1.append(PythonOperator(
        task_id=task_id,
        python_callable=StepOne,
        provide_context=False,
        op_kwargs={'x': str(t)},
        dag=dag))

task2 = PythonOperator(
    task_id="SecondStep",
    python_callable=StepTwo,
    provide_context=False,
    dag=dag)

task1 >> task2

Upvotes: 3

Related Questions