Setup and Teardown Tasks

Concepts covered: Trigger rules and branching

A DAG that works in a scratch table runs 4 tasks in one chain. The first 2 create the table and load it; the last 2 roll it up and drop it. The chain has a hole in it. When load_orders fails, drop_tmp sits downstream of a failure and goes upstream_failed, so the table stays until a person notices. On a scratch cluster the same hole bills by the hour all night. The repair the trigger rules lesson suggests, trigger_rule="all_done" on the drop, opens a second hole: the run's state follows its leaf tasks, drop_tmp is the only leaf, and once it succeeds the run turns green with a failed load inside it, so the failure callback never fires.

Marking the pair

from airflow.sdk import dag, setup, task, teardown
from etl import warehouse


@dag(schedule=None, catchup=False)
def scratch_etl():
    @setup
    def create_tmp():
        warehouse.create("tmp_orders")

    @task
    def load_orders():
        warehouse.load("tmp_orders")

    @task
    def aggregate():
        warehouse.rollup("tmp_orders")

    @teardown
    def drop_tmp():
        warehouse.drop("tmp_orders")

    with create_tmp() >> drop_tmp():
        load_orders() >> aggregate()


scratch_etl()

@setup and @teardown replace @task on the 2 tasks that frame the work, and the with block wires everything created inside it after the setup and before the teardown, so the 4 tasks form the same chain as the plain version. The marking changes 3 things about drop_tmp. Its trigger rule becomes all_done_setup_success: it runs once every upstream task has finished, in any state, provided the setup it belongs to succeeded. Its own state stops counting toward the run's state, so a green teardown under a failed load leaves the run failed. And clearing any task between the pair clears the pair with it, so a rerun of aggregate recreates the table before it and drops it after.

Play the night below. A manual trigger at 02:02 creates the run; create_tmp runs for 2 minutes and succeeds; load_orders runs for 4 minutes and fails; aggregate goes upstream_failed without starting; and drop_tmp, with its setup green and every upstream cell settled, is scheduled, runs for 1 minute and succeeds. The run closes as failed with the teardown green inside it. Switch the variant to a night where the load succeeds and the only difference is the 2 cells in the middle.

from airflow.sdk import dag, setup, task, teardown

@dag(schedule=None, catchup=False)
def scratch_etl():
    @setup
    def create_tmp():
        print("create tmp_orders")
    @task
    def load_orders():
        print("load tmp_orders")
    @task
    def aggregate():
        print("rollup tmp_orders")
    @teardown
    def drop_tmp():
        print("drop tmp_orders")
    with create_tmp() >> drop_tmp():
        load_orders() >> aggregate()

scratch_etl()

The last 2 rows show the default for a teardown that fails on its own: the work succeeded, the run reports success, and the undropped table shows only as the teardown's red cell. For a cluster that bills until it is deleted, that default hides the cost. @teardown(on_failure_fail_dagrun=True) makes the teardown's failure fail the run, which routes it to the same on_failure_callback as any other failure. Choose per teardown: a temp table nobody pays for can stay a red cell, a cluster cannot.

Classic operators and several setups

Operators from providers take the same marking through 2 methods. create.as_setup() returns the task marked as a setup, and delete.as_teardown(setups=create) returns the teardown tied to it. setups takes a list when one teardown closes several setups, and the teardown then waits for all of its upstream tasks and runs when its setups succeeded. The work tasks between the pair need no marking at all; the chain through the teardown places them.

from airflow.providers.standard.operators.bash import BashOperator

create = BashOperator(task_id="create_cluster", bash_command="cluster-cli create etl-scratch")
run_job = BashOperator(task_id="run_job", bash_command="cluster-cli submit etl-scratch job.py")
delete = BashOperator(task_id="delete_cluster", bash_command="cluster-cli delete etl-scratch")

create.as_setup() >> run_job >> delete.as_teardown(setups=create, on_failure_fail_dagrun=True)

Clearing follows the pair as a unit. Clear run_job after fixing its bug and the grid clears create_cluster and delete_cluster with it, because the rerun needs a cluster that exists and must delete it afterwards. Clear delete_cluster alone and only the delete reruns, which is the move when the job finished and the delete was the task that failed. A clear of the setup alone also clears the work and the teardown, since work that ran on the old cluster cannot be kept once a new one is made.

About this lesson

This section is part of the Task Lifecycle: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.