Airflow TaskFlow API: Cleaner DAGs with Decorators

The TaskFlow API

The TaskFlow API is Airflow's decorator-based way of writing DAGs. It replaces the classic pattern of wrapping functions in PythonOperator, passing values through the context dictionary, and manually defining dependencies at the end.

Two decorators are all you need:

• @task – turns a Python function into an Airflow task.
• @dag – turns a Python function into a DAG.

TaskFlow handles two things automatically that used to require manual code:

• XCom passing – the return value of one task is pushed to XCom.
• Dependencies – calling transform(count) inside the DAG function tells Airflow that transform depends on extract. No explicit >> line is required.

Creating the DAG File

In VS Code, navigate to the Airflow directory, then to the DAGs folder. Create a new file called taskflow.py.

Import only what is required:

from airflow import DAG
from airflow.decorators import task, dag
from datetime import datetime

• task and dag come from airflow.decorators. These are the TaskFlow API entry points.
• datetime sets the DAG's start_date.

Defining Tasks with @task

Decorate a function with @task to turn it into an Airflow task:

@task
def extract():
    amount = 42
    print(f"Extracted {amount} units of data")
    return amount

@task
def transform(count: int):
    result = count * 2
    print(f"Transformed result: {result}")
    return result

Defining the DAG with @dag

Decorate a Python function with @dag to define the DAG structure:

@dag(
    dag_id="taskflow_dag",
    start_date=datetime(2023, 1, 1),
    schedule="@daily",
    catchup=False
)
def my_pipeline():
    count = extract()
    transform(count)

Key points:

• dag_id – the name shown in the Airflow UI.
• schedule="@daily" – the DAG runs once per day.
• catchup=False – prevents backfilling of missed runs.
• Inside the pipeline function, calling transform(count) establishes the dependency automatically.

Call the DAG function at the bottom of the file to register it with Airflow:

my_pipeline()

Running the DAG

Open the Airflow UI, search for taskflow_dag, and trigger a single run.

Expected task output:

• extract – logs "Extracted 42 units of data" and returns 42.
• transform – logs "Transformed result: 84" after receiving 42 from extract.