Write and Run a Dag with the Task SDK in Apache Airflow 3

Table of Contents

1. Challenge

   From a script to a Dag

   Welcome to this lab on writing and running a Dag with the Apache Airflow 3 Task SDK. You play a data engineer at Globomantics. Every morning someone on the operations team runs a Python script by hand. It reads the latest orders export, keeps the orders that shipped, totals revenue per product, and writes a JSON summary that feeds a dashboard. The script works until someone forgets to run it, runs it twice, or runs it against the wrong region. The team is moving to Apache Airflow 3, and this script is the first pipeline to migrate.

   What the Task SDK gives you

   Airflow 3 introduced the Task SDK, a stable authoring interface that you import from airflow.sdk. It keeps your pipeline code separate from the scheduler internals. You define a Dag by decorating a Python function with @dag, you turn ordinary functions into tasks with @task, and Airflow moves data between tasks through XCom whenever a task returns a value. You declare dependencies by calling one task with the result of another, the same way you call functions.

   Airflow 3 spells it Dag in its own documentation and release notes, so this lab does too.

   Running a Dag without a scheduler

   dag.test() runs a whole Dag inside one Python process. No scheduler, no API server, and no web UI are involved. That makes it the fastest loop for developing a Dag: edit the file, run it, read the output, repeat. In this lab every run happens through dag.test(). The lab environment already initialized the local SQLite metadata database it needs.

   What you will build

   You start with dags/orders_pipeline.py, which holds the imports, two path constants, and an empty orders_pipeline() function that wraps three empty nested functions with TODO comments. Across the next three steps you turn that file into a working three-task Dag:

   • Step 2 defines the Dag with @dag and its metadata.
   • Step 3 implements extract, transform, and load as @task functions and wires them together.
   • Step 4 runs the Dag with dag.test(), passes in a region, and verifies the result from the returned DagRun.

   The data lives in data/orders.csv, 309 synthetic orders across four regions and four SKUs. The load task writes its summaries to outputs/. Each task ends with a check. Click Validate on the task to run it, and read the feedback if it fails.

   Note: If you get stuck at any point, the solutions/ folder contains the completed orders_pipeline.py for every task, and comments in each file name the task that adds each part.

2. Challenge

   Define a Dag with the @dag decorator

   In the Task SDK a Dag is a decorated Python function. The @dag decorator carries the metadata that Airflow needs: a unique dag_id, a schedule, a start_date, and whether Airflow should backfill missed runs with catchup. Calling the decorated function once at module level produces the Dag object, and that is what dag.test() runs later.

   This Dag has no schedule, because in this lab you always run it by hand. You still set catchup=False. It has no effect while schedule is None, but it is the habit to keep, because the moment you give a Dag a schedule, catchup decides whether Airflow runs every missed interval since start_date. The params you just declared are the hook for Step 4. A run can override any param through run_conf, and every task can read the merged values. That is how the same Dag will summarize all regions on one run and a single region on the next.

3. Challenge

   Create tasks with the @task decorator

   A @task function looks like a normal Python function, and that is the point. Whatever it returns becomes its XCom value, and whatever you pass into it from another task becomes a dependency. You never write set_upstream or >> for this pipeline. You write summary = transform(rows) and Airflow knows that transform runs after extract.

   One detail matters here. A decorated function does not become a task until you call it inside the Dag function. Decorating extract is not enough. The line rows = extract() inside orders_pipeline() is what registers the task on the Dag, so each task in this step adds both the decorator and the call.

   Note: XCom is for small handoffs. The 309 rows in this lab are about 15 KB of JSON, which is fine. For a real extract you would return a file path or an Asset and let the next task read the data itself. Declaring params as a named argument is how the Task SDK hands a task its run context. Airflow inspects the function signature and injects the matching context values. You could also call get_current_context() from airflow.sdk, but a named argument keeps the function easy to call and easy to test on its own, which is exactly what the check does. retries is operator configuration, so it lives on the decorator rather than inside the function. Airflow will rerun a failed load up to two more times in a scheduled deployment. In this lab the write succeeds on the first try, so you will not see a retry, but the setting travels with the Dag when it moves to production.

4. Challenge

   Run and verify the Dag with dag.test()

   You now have a complete Dag, but nothing has run it. dag.test() executes every task in dependency order inside the current Python process and returns a DagRun object. In this step you add the entry point that calls it, pass a region into a run, and then turn "it looks like it worked" into an assertion. The output is verbose because Airflow logs every state transition. Two lines matter. Each task reports Task instance state updated with new_state=success, and the run ends with DagRun Finished and state=success. The print calls inside your tasks appear between them. run_conf merges into the Dag's params for that run only. The default of all that you declared in Step 2 still applies whenever no region is passed, which is why the earlier summary_all.json came out the way it did. Compare states with the enums, not with str(state). TaskInstanceState and DagRunState are string enums, so state == TaskInstanceState.SUCCESS is true whether Airflow hands you the enum or the plain string success. str(TaskInstanceState.SUCCESS) is the text TaskInstanceState.SUCCESS, so a helper built on str(state) == "success" fails the moment a run object carries the enum, and Airflow 3.3 already returns the enum for dag_run.state. The helper you just wrote is small, and it is the difference between a run you watched and a run you verified.

5. Challenge

   Conclusion and next steps

   You rebuilt a script that someone ran by hand as an Airflow 3 Dag with the Task SDK. You defined the Dag with @dag and gave it an id, a start date, tags, params, and docs. You created three tasks with @task, passed data between them through XCom by returning values, declared the dependencies by calling the tasks, and put a retry policy on the step that writes to disk. Then you ran the Dag with dag.test(), parameterized a run with run_conf, and verified the outcome from the DagRun that came back, using the state enums that Airflow itself uses.

   Moving the Dag to a deployment

   Three things change when this file lands in a real Airflow 3 environment. The schedule becomes a cron string or a timetable instead of None, and catchup starts to matter. The file lives in the deployment's Dag bundle, where the scheduler parses it. The retries on load do real work when a disk or network hiccup fails a write. Nothing about the tasks, the XCom handoffs, or the dependencies has to change.

   Where to go next

   • Give the Dag a daily schedule, set catchup=True, and watch Airflow create a backfill run for every day since start_date.
   • Replace the single transform with dynamic task mapping over the four regions so each region becomes its own task instance.
   • Return an Asset from load so a downstream Dag can run whenever a new summary lands.
   • Keep dag.test() and verify_run in your workflow. They give you a local loop that takes seconds and catches most mistakes before the scheduler ever sees the file.