Branching, Setup/Teardown, and Params
Conditional production DAGs need explicit joins, resource lifecycle tasks, and typed params instead of ad hoc conf dictionaries.
The foundations series taught the three moving parts of a forking DAG: @task.branch picks a path, everything on the roads not taken is set to skipped, and a join underneath needs trigger_rule="none_failed_min_one_success" or it skips itself. If any of that is fuzzy, Branching and Trigger Rules is the place to firm it up — this post assumes it. What’s left for a production DAG is the harder half: branches that nest, branches that feed a dynamic fan-out, resources that have to be created before the work and destroyed after it even when the work explodes, and manual runs that need real, validated inputs instead of a free-form conf dictionary someone typed by hand at 2 a.m. Those are three different tools — advanced branching, setup/teardown tasks, and the Param system — and each one closes a gap that the basics leave open.
Branching, past the basics
Two quick recaps so the rest stands on its own. @task.branch returns the task_id (or a list of them) of the path to run; @task.short_circuit returns a boolean and skips everything downstream on a falsy return. Both live in the Task SDK alongside @task, and the classic BranchPythonOperator still exists at airflow.providers.standard.operators.python if you meet it in an older DAG. Every join beneath a branch needs a trigger rule that tolerates a skipped parent — none_failed_min_one_success nine times out of ten. That’s the whole of what you carry forward. Now the parts the foundations post deliberately skipped.
Skipping without a branch. A @task.branch is overkill when you have exactly one task that sometimes shouldn’t run and there’s no alternate path to choose. Raise AirflowSkipException from inside an ordinary @task and that single instance goes to skipped — no branch node, no dummy task to steer toward:
from airflow.sdk import dag, task
from airflow.sdk.exceptions import AirflowSkipException
@task
def load_reviews(path: str) -> int:
rows = read_parquet(path)
if not rows:
# nothing landed today — skip, don't fail, don't fabricate an empty load
raise AirflowSkipException("no review file for this run")
return warehouse.load(rows, table="stg_reviews")
The distinction that matters: this task decides its own fate, mid-work, using information it only has once it’s running. A branch decides other tasks’ fate before any of them start. Reach for AirflowSkipException for “this specific step turned out to be a no-op today,” and for a branch when the shape of the downstream graph changes. The skip still cascades and still trips joins, so anything below a task that might raise it wants the same none_failed_min_one_success treatment.
Nested branches. Production decisions are rarely a single fork. The bookshop’s nightly run first asks is this a full refresh or an incremental one, and only on the incremental path does it then ask did any regions actually change. That’s a branch whose chosen path contains another branch, and the only real subtlety is that each join needs its own trigger rule — the discipline doesn’t compound, it just repeats at every fork:
@dag(schedule="@daily", catchup=False, tags=["bookshop"])
def bookshop_nightly():
@task.branch
def choose_refresh(params: dict) -> str:
return "full_reload" if params["full_refresh"] else "check_changes"
@task
def full_reload():
warehouse.rebuild_all()
@task.branch
def check_changes() -> str:
return "incremental_load" if changed_regions() else "no_changes"
@task
def incremental_load():
warehouse.merge_changed()
@task
def no_changes():
print("nothing changed since last run")
@task(trigger_rule="none_failed_min_one_success")
def write_run_report():
print("run summary written")
# Call each task EXACTLY once and hold the reference. Calling `full_reload()`
# twice does not reuse the task — it creates a second one (`full_reload__1`)
# with no upstream, which then runs on every single run, unconditionally.
refresh = choose_refresh()
inner = check_changes()
full = full_reload()
incremental = incremental_load()
nothing = no_changes()
report = write_run_report()
refresh >> [full, inner]
inner >> [incremental, nothing]
[full, incremental, nothing] >> report
bookshop_nightly()
Read the fan-in on the last line. write_run_report sits below the outer branch’s full_reload and below both arms of the inner branch. On a full-refresh run, full_reload succeeds and both inner arms are skipped — one success, two skips, no failures, so none_failed_min_one_success fires. On an incremental run with changes, full_reload is skipped, incremental_load succeeds, no_changes is skipped — again one success amid skips. The rule holds at the join no matter how deep the branching that feeds it, because it only ever asks the same question: did anything upstream fail, and did at least one thing actually work? Nesting branches doesn’t require a new idea, only that you remember to put the rule on every convergence, not just the outermost one.
Branch into a dynamic fan-out. The other combination worth knowing is a branch that steers into dynamic task mapping. The branch decides whether to process regions at all; if it says yes, a mapped task fans out over however many regions today’s run turned up. A branch task can name the first task of the mapped path as its target, and mapping proceeds from there as normal:
@task.branch
def gate(params: dict) -> str:
return "list_regions" if not params["skip_regions"] else "skip_all"
@task
def list_regions() -> list[str]:
return storage.list_regions() # runtime-decided width
@task
def load_region(region: str) -> int:
return warehouse.load_region(region)
@task
def skip_all():
print("regional load disabled for this run")
@task(trigger_rule="none_failed_min_one_success")
def consolidate(counts):
print(f"loaded {sum(c for c in counts if c)} rows")
g = gate()
regions = list_regions()
loads = load_region.expand(region=regions)
skipped = skip_all()
summary = consolidate(loads) # consuming the mapped output IS the dependency
g >> [regions, skipped]
skipped >> summary # so the skip path also reaches the join
The one thing to hold in your head: when gate chooses skip_all, the entire mapped load_region node is skipped as a unit — Airflow doesn’t expand it and then skip each instance, it never expands it at all, because its upstream list_regions was skipped and produced no list. The reduce step consolidate collects the mapped outputs when they exist and tolerates the all-skipped case through its trigger rule. Branch-then-map gives you a fan-out whose width is data-driven and whose existence is decision-driven, which is exactly the pipeline that only touches the warehouse when there’s genuinely something to touch.
Setup and teardown: resource lifecycle as first-class tasks
Here’s the problem the trigger rules can’t solve cleanly. The nightly dbt build wants an isolated place to work — a scratch schema, a temp cluster, a checked-out branch of a repo — created before the transform runs and destroyed after. The destroy step is non-negotiable: leave the schema behind and tomorrow’s run collides with today’s, or you leak storage until someone notices the bill. So the cleanup has to run whether the transform succeeded, failed, or was skipped. The foundations series’ answer to “always run” was trigger_rule="all_done", and it works, but it’s a blunt instrument — an all_done cleanup will also fire if the pipeline never got far enough to create the thing it’s cleaning up, and it doesn’t model the fact that the cleanup and the creation are paired. Airflow 3.x has a purpose-built construct for exactly this shape: setup and teardown tasks.
A setup task creates a resource. Its paired teardown destroys it. Airflow understands the relationship, and that understanding buys three things all_done can’t: the teardown runs even when the work between them fails, a failed teardown doesn’t automatically fail your DAG run (cleanup breaking is not the same as the pipeline breaking), and the teardown is excluded from the DAG’s success calculation, so a green run means the work succeeded, not that the cleanup happened to. You mark them with the @setup and @teardown decorators from the Task SDK:
import re
from airflow.sdk import dag, task, setup, teardown, get_current_context
from airflow.providers.postgres.hooks.postgres import PostgresHook
@dag(schedule="@daily", catchup=False, tags=["bookshop"])
def bookshop_dbt_nightly():
@setup
def create_scratch_schema() -> str:
# a schema unique to this run, so concurrent/backfill runs never collide
run_id = get_current_context()["run_id"]
suffix = re.sub(r"[^0-9a-zA-Z]+", "_", run_id).strip("_").lower()
schema = f"dbt_scratch_{suffix}"
hook = PostgresHook(postgres_conn_id="warehouse")
hook.run(f"CREATE SCHEMA IF NOT EXISTS {schema}")
return schema
@task
def run_dbt(schema: str) -> None:
# build the bookshop models into the throwaway schema
run_dbt_build(target_schema=schema)
@task
def test_dbt(schema: str) -> None:
run_dbt_tests(target_schema=schema)
@teardown
def drop_scratch_schema(schema: str) -> None:
hook = PostgresHook(postgres_conn_id="warehouse")
hook.run(f"DROP SCHEMA IF EXISTS {schema} CASCADE")
schema = create_scratch_schema()
built = run_dbt(schema)
tested = test_dbt(schema)
schema >> built >> tested >> drop_scratch_schema(schema)
bookshop_dbt_nightly()
Note how the schema name gets built, because it’s the mistake everyone makes once: the run id comes out of get_current_context(), in Python. It is tempting to write f"dbt_scratch_{{ run_id }}" instead — but that is two bugs stacked. An f-string treats {{ as an escaped literal brace, so you would get the characters { run_id } in your schema name; and even if you dodged that, Jinja does not render inside a @task body — templating happens in an operator’s template_fields, not in arbitrary Python. Inside a task, reach into the context object. (A real run_id looks like scheduled__2026-07-12T00:00:00+00:00, which is not a legal identifier, hence the re.sub — the colons and plus signs have to go.)
The wiring on the last line is ordinary >>, but because create_scratch_schema is a setup and drop_scratch_schema is a teardown, Airflow reads the tasks between them — run_dbt and test_dbt — as being in the setup’s scope. That scope is the whole point. If run_dbt fails, test_dbt is skipped as usual, but drop_scratch_schema still runs and drops the schema. If test_dbt fails, same story. The teardown fires on the failure path exactly as reliably as on the success path — which is what you always wanted from cleanup and what all_done only approximates. And crucially, the run’s overall state reflects the work: if run_dbt failed, the DAG run is failed even though the teardown succeeded, because a teardown’s success doesn’t paper over a failure in its scope.
When the teardown itself fails. By default, a teardown that throws does not fail the DAG run — the reasoning being that if the real work succeeded, a hiccup dropping a temp schema shouldn’t turn a green pipeline red and page someone. Sometimes that’s wrong: if leaking the resource is genuinely dangerous (an expensive cluster left running, a lock never released), you want a failed teardown to be loud. That’s the on_failure_fail_dagrun flag:
@teardown(on_failure_fail_dagrun=True)
def drop_scratch_schema(schema: str) -> None:
PostgresHook(postgres_conn_id="warehouse").run(f"DROP SCHEMA IF EXISTS {schema} CASCADE")
With that set, a DROP that fails marks the run failed, so the leak surfaces instead of hiding behind a green checkmark. Leave it at the default for cheap, idempotent cleanups; turn it on when the cost of a silent leak is real.
Teardown as a context manager. Wiring the scope with explicit arrows is clear for a couple of tasks, but for a scope with several steps it’s easy to forget one and leave a task dangling outside the teardown’s protection. The teardown task doubles as a context manager, and anything defined inside the block is automatically placed in scope between the setup and the teardown:
schema = create_scratch_schema()
drop = drop_scratch_schema(schema)
with drop:
built = run_dbt(schema)
tested = test_dbt(schema)
built >> tested
schema >> built
Every task created within with drop: is guaranteed to be inside the scope — you can’t accidentally strand one downstream of the teardown or upstream of the setup. For a scratch-schema build with two steps the arrow form reads fine; for a scope with a TaskGroup and a fan-out inside it, the context manager is the safer default because it makes “these tasks are protected by this cleanup” a structural fact rather than something you have to wire correctly by hand.
More than one setup, more than one teardown. A real build often borrows more than one resource. The bookshop nightly might need both a scratch schema and a temporary object-store prefix for staged extracts, each with its own creation and its own cleanup. You don’t nest scopes for this — you pair each setup with the teardown that undoes it, and the arrows describe which teardown belongs to which setup:
@setup
def create_scratch_schema() -> str: ...
@teardown
def drop_scratch_schema(schema: str) -> None: ...
@setup
def create_staging_prefix() -> str: ...
@teardown
def clear_staging_prefix(prefix: str) -> None: ...
schema = create_scratch_schema()
prefix = create_staging_prefix()
drop = drop_scratch_schema(schema)
clear = clear_staging_prefix(prefix)
# both setups feed the work; both teardowns close over it
[schema, prefix] >> run_dbt(schema, prefix) >> [drop, clear]
The rule Airflow applies is worth stating plainly: a task is in the scope of a setup when it sits downstream of that setup and upstream of the paired teardown. run_dbt sits downstream of both setups and upstream of both teardowns, so it’s in both scopes — if it fails, both teardowns still fire, and neither resource leaks. A teardown fires as soon as everything in its scope has finished, so if you also had a schema-only step that no longer needs the staging prefix, clear_staging_prefix could run while later schema work continues, freeing the resource the moment it’s genuinely done rather than at the end of the DAG. Pairing is per-resource; the arrows, not proximity in the file, decide what each teardown protects.
One more property that makes setup/teardown more than sugar over all_done: setups and teardowns are excluded from the DAG’s clear-and-rerun logic in the way you’d hope. Clear a failed run_dbt to retry it, and Airflow, by default, also reruns the setup — it knows the setup has to run again to recreate the scratch schema the retry needs, so the resource lifecycle travels with the work it wraps. When the resource is expensive and still healthy — a warm cluster you don’t want rebuilt just to re-run one downstream step — clear with the ignore setup/teardown option (the flag is exposed on the clear action in the UI and on airflow tasks clear), and Airflow leaves the already-created resource in place and reruns only the work. The default is the safe choice; the ignore option is the escape hatch for when recreation is the expensive part. Either way that’s the difference between modeling the relationship and merely bolting a cleanup task onto the end of the graph.
Params: typed inputs instead of a conf junk drawer
The last gap is input. A manually-triggered run — a backfill, a one-off full refresh, a re-run for a specific date — needs to be told what to do, and the lazy way is dag_run.conf, a free-form dictionary the operator fills in as raw JSON. Nothing validates it. Misspell a key, pass a string where the DAG expected a boolean, forget a field entirely, and the run launches anyway and fails deep inside a task, or worse, does the wrong thing silently. The Param system replaces that junk drawer with a typed, validated, self-documenting surface: you declare the inputs a DAG accepts, their types and constraints, and Airflow renders a real form in the UI and rejects invalid input before the run ever starts.
You declare params in the params= argument on @dag, each one a Param with a default and a schema:
from airflow.sdk import dag, task, Param
@dag(
schedule="@daily",
catchup=False,
tags=["bookshop"],
params={
"full_refresh": Param(
False,
type="boolean",
title="Full refresh",
description="Rebuild every model from scratch instead of merging changes.",
),
"run_date": Param(
None,
type=["null", "string"],
format="date",
title="Run date",
description="Override the logical date. Leave blank to use the run's own date.",
),
"batch_size": Param(
5000,
type="integer",
minimum=100,
maximum=50000,
title="Batch size",
),
"environment": Param(
"staging",
type="string",
enum=["staging", "production"],
title="Target environment",
),
},
)
def bookshop_configurable():
...
Each Param is a small JSON-Schema declaration. type="boolean" renders a checkbox; type="integer" with minimum/maximum renders a number field that the form refuses to submit out of range; enum=[...] renders a dropdown of exactly those choices, so environment can never be a typo’d "prod". type=["null", "string"] makes run_date genuinely optional — its default of None is a valid value — while format="date" gives the field a date picker. title and description are what the operator reads in the form; spend a sentence on each, because this is the documentation someone sees at the moment they’re deciding what to run.
Where params show up. Two places. In any templated field you reference them through the Jinja context as {{ params.x }}:
from airflow.providers.standard.operators.bash import BashOperator
run_models = BashOperator(
task_id="run_dbt",
bash_command=(
"dbt build --target {{ params.environment }} "
"{% if params.full_refresh %}--full-refresh{% endif %} "
"--vars '{run_date: {{ params.run_date or ds }}}'"
),
)
Note {{ params.run_date or ds }} — the param defaults to None, and when it’s left blank the template falls back to ds, the run’s own logical date. A scheduled run gets its date automatically; a manual or backfill run can override it by filling in the field. That’s the idiomatic way to make a date both automatic and overridable: a nullable param that falls through to the built-in macro.
Inside a @task, reach the params through the runtime context rather than Jinja:
from airflow.sdk import task, get_current_context
@task
def summarize_run() -> None:
context = get_current_context()
params = context["params"]
if params["full_refresh"]:
print(f"full rebuild into {params['environment']}, batch={params['batch_size']}")
else:
print(f"incremental into {params['environment']} for {params['run_date'] or context['ds']}")
get_current_context() (from airflow.sdk) hands back the same dict Jinja sees, so context["params"] is your typed, already-validated input. This is how the nested branch from earlier read params["full_refresh"] to choose its path — the branch is just a @task.branch that pulls the param and returns a task_id.
The form, and the validation. In the UI, the Trigger DAG w/ config action reads these param declarations and builds a form: a checkbox for full_refresh, a bounded number field for batch_size, a dropdown for environment, a date picker for run_date. The operator fills it in and Airflow validates before creating the run — a batch_size of 80 is rejected against minimum=100 with a message, not accepted and blown up on later. This is the payoff over conf: the invalid run never launches, the valid one carries structured values every task can trust, and the form itself documents what the DAG accepts. A newcomer triggering a backfill sees titled fields and descriptions instead of a blank JSON box, and can’t invent a key the DAG doesn’t understand.
Two sharp edges. First, params are validated at trigger time, but a param with no default and no null in its type makes the DAG un-runnable on a schedule — the scheduler has nothing to fill it with and there’s no form to prompt. If a param must be supplied every run, that DAG is manual-only by design; if it should schedule, give it a sensible default. Second, dag_run.conf still exists and still overrides params when a run is triggered with an explicit config payload (the API, airflow dags trigger --conf, an upstream TriggerDagRunOperator), so params and programmatic triggering coexist — the form is the human front door, conf is the machine one, and both land in the same validated context["params"].
The machine door is validated too, which is the part people miss. The same JSON-Schema that renders the form runs against a programmatic conf payload — a TriggerDagRunOperator or a POST to /dagRuns that supplies environment="prod" is rejected against enum=["staging", "production"] before the run is created, exactly as the dropdown would reject it, and the caller gets the validation error back rather than a run that fails halfway through. Airflow 3.x is strict about unexpected keys as well: by default a conf that includes a key you never declared as a Param is rejected rather than silently passed through, so a typo’d field name surfaces at trigger time instead of being read as None three tasks later. The upshot is that “validated input” isn’t a UI nicety bolted onto a still-loose backend; declaring params tightens every path into the DAG, human and programmatic alike, which is what lets a downstream task treat context["params"] as trustworthy without re-checking it.
Bringing it together
The three tools stack. The bookshop’s nightly DAG takes a full_refresh param that the operator flips for a manual rebuild; a @task.branch reads that param and forks between a full reload and an incremental path; the incremental path fans out over regions with dynamic mapping; and the whole transform runs inside a setup/teardown scope that spins up a scratch schema and drops it whether the build passes or fails. Each tool covers a distinct axis — the param is what to do, the branch is which shape the graph takes, the setup/teardown is what resources the work borrows and returns. None of them substitutes for another, and a production DAG of any real complexity ends up using all three at once, which is precisely why they share a chapter.
Final thoughts
The through-line here is making the implicit explicit. A branch makes a conditional visible in the graph instead of buried in an if inside one fat task. Setup/teardown makes a resource’s lifetime a pair of nodes Airflow tracks and guarantees, instead of cleanup hidden in a finally or bolted on with all_done and hope. Params make a run’s inputs a typed, validated, documented form instead of a dictionary nobody checks. In every case the win is the same: Airflow now understands something it used to merely execute, and because it understands it, it can show it in the UI, validate it before it runs, and do the right thing when a piece fails. That’s the whole thesis of the advanced series in miniature — the orchestrator earns its keep not by running your code but by knowing what your code is for. Push the conditionals, the lifecycles, and the inputs out of your task bodies and into the constructs built for them, and the DAG stops being a script that happens to run on a schedule and becomes something a stranger can read, trigger safely, and trust.
Next: Custom XCom Backends
Comments