Airflow TaskFlow API: Cleaner DAGs with Decorators

Rewriting the DAG with the TaskFlow API

TaskFlow is a decorator-based API. Instead of wrapping functions in operators, you decorate them. Two decorators are all you need: @task turns a function into an Airflow task, and @dag turns a function into a DAG.

Creating the TaskFlow DAG file

In VS Code, go to the Airflow directory, then the DAG folder, and create a new file called taskflow.py.

Import only the DAG and task decorators, plus datetime to set the start date:

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

Defining the tasks with decorators

Define the first function using the @task decorator. It returns the amount of data (42 as an example) and prints something to the logs:

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

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

Setting up the DAG decorator

Use the @dag decorator on a function that will define the pipeline. The first part of the decorator sets the DAG ID, start date, schedule, and catchup:

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

my_pipeline()

Running the DAG

In the Airflow UI, go to the DAGs section and search for taskflow_dag. Trigger a single run. Both tasks should succeed.

Check the logs:

• For extract, the output is 42.
• For transform, the output is 84.