Run an airflow task after a task in a loop, not after all tasks in a loop

Question

Let's say we have these tasks:

for endpoint in ENDPOINTS:
    latest_only = LatestOnlyOperator(
        task_id=f'{endpoint.name}_latest_only',
    )

    s3 = SnowflakeQOperator(
        task_id=f'{endpoint.name}_to_S3',
        boostr_conn_id='boostr_default',
        s3_conn_id='aws_default',
        partition=endpoint.partition,
        endpoint=endpoint
    )

    short_circuit = ShortCircuitOperator(
        task_id=f"short_circuit_missing_{endpoint.name}",
        op_kwargs={'endpoint_to_check': endpoint, 'aws_conn_id': 'aws_default'},
        python_callable=check_file_exists,
        provide_context=True
    )

     s3 >> short_circuit

and let's say I want to add one task to run after nbc_to_s3 which is one of the '{endpoint.name}' task in the s3 task.

we're importing ENDPOINTS which contains several class, with the 'name' method:

@property
def name(self) -> str:
    return 'nbc'

I've tried to add it outside of the loop like this: nbc_to_s3 >> new_task but that doesn't work because 'nbc_to_s3' is not defined

user_15772377 · Answer

You could apply some logic within the loop to set a new dependency for new_task like so (apologies for the quick mockup):

from airflow.decorators import dag
from airflow.operators.dummy import DummyOperator

from datetime import datetime

ENDPOINTS = ["nbc", "cbs", "bravo", "espn"]
DEFAULT_ARGS = dict(owner="airflow", start_date=datetime(2021, 6, 9))
DAG_ARGS = dict(schedule_interval=None, default_args=DEFAULT_ARGS, catchup=False)


@dag(**DAG_ARGS)
def run_task_after_loop():
    for endpoint in ENDPOINTS:
        s3 = DummyOperator(
            task_id=f"{endpoint}_to_S3",
        )

        short_circuit = DummyOperator(
            task_id=f"short_circuit_missing_{endpoint}",
        )

        s3 >> short_circuit

        if endpoint == "nbc":
            new_task = DummyOperator(task_id=f"new_task_{endpoint}")

            s3 >> new_task


dag = run_task_after_loop()

Run an airflow task after a task in a loop, not after all tasks in a loop

1 Answers