Introduction to Apache Airflow® DAGs

1
"""
2
## Astronaut ETL example DAG
3
4
This DAG queries the list of astronauts currently in space from the
5
Open Notify API and prints each astronaut's name and flying craft.
6
7
There are three tasks, one to get the data from the API and save the results,
8
one to print the results and a final task to react. The first two tasks are
9
written in Python using Airflow's TaskFlow API, which allows you to easily turn
10
Python functions into Airflow tasks, and automatically infer dependencies and pass data.
11
12
The second task uses dynamic task mapping to create a copy of the task for
13
each Astronaut in the list retrieved from the API. This list will change
14
depending on how many Astronauts are in space, and the DAG will adjust
15
accordingly each time it runs.
16
17
The third task is defined using the BashOperator, which is a traditional
18
operator, allowing you to run a bash command.
19
"""
20
21
from airflow.decorators import dag, task
22
from airflow.operators.bash import BashOperator
23
from airflow.models.baseoperator import chain
24
from pendulum import datetime
25
import requests
26
27
28
# Define the basic parameters of the DAG, like schedule and start_date
29
@dag(
30
    start_date=datetime(2024, 1, 1),
31
    schedule="@daily",
32
    catchup=False,
33
    doc_md=__doc__,
34
    default_args={"owner": "Astro", "retries": 3},
35
    tags=["example"],
36
)
37
def example_astronauts_three_tasks():
38
    # Define tasks
39
    @task
40
    def get_astronauts(**context) -> list[dict]:
41
        """
42
        This task uses the requests library to retrieve a list of Astronauts
43
        currently in space. The results are pushed to XCom with a specific key
44
        so they can be used in a downstream pipeline. The task returns a list
45
        of Astronauts to be used in the next task.
46
        """
47
        r = requests.get("http://api.open-notify.org/astros.json")
48
        number_of_people_in_space = r.json()["number"]
49
        list_of_people_in_space = r.json()["people"]
50
51
        context["ti"].xcom_push(
52
            key="number_of_people_in_space", value=number_of_people_in_space
53
        )
54
        return list_of_people_in_space
55
56
    @task
57
    def print_astronaut_craft(greeting: str, person_in_space: dict) -> None:
58
        """
59
        This task creates a print statement with the name of an
60
        Astronaut in space and the craft they are flying on from
61
        the API request results of the previous task, along with a
62
        greeting which is hard-coded in this example.
63
        """
64
        craft = person_in_space["craft"]
65
        name = person_in_space["name"]
66
67
        print(f"{name} is currently in space flying on the {craft}! {greeting}")
68
69
    # define a task using a traditional operator
70
    # this task will run a bash command
71
    print_reaction = BashOperator(
72
        task_id="print_reaction",
73
        bash_command="echo This is awesome!",
74
    )
75
76
    # Use dynamic task mapping to run the print_astronaut_craft task for each
77
    # Astronaut in space
78
    print_astronaut_craft_obj = print_astronaut_craft.partial(
79
        greeting="Hello! :)"
80
    ).expand(
81
        person_in_space=get_astronauts()  # Define dependencies using TaskFlow API syntax
82
    )
83
84
    # set the dependency between the second and third task explicitly
85
    chain(print_astronaut_craft_obj, print_reaction)
86
87
88
# Instantiate the DAG
89
example_astronauts_three_tasks()

1
"""
2
## Example of a complex DAG structure
3
4
This DAG demonstrates how to set up a complex structures including:
5
- Branches with Labels
6
- Task Groups
7
- Dynamically mapped tasks
8
- Dynamically mapped task groups
9
10
The tasks themselves are empty or simple bash statements.
11
"""
12
13
from airflow.decorators import dag, task_group, task
14
from airflow.operators.empty import EmptyOperator
15
from airflow.operators.bash import BashOperator
16
from airflow.models.baseoperator import chain, chain_linear
17
from airflow.utils.edgemodifier import Label
18
from airflow.datasets import Dataset
19
from pendulum import datetime
20
21
22
# Define the DAG
23
@dag(
24
    start_date=datetime(2024, 1, 1),
25
    schedule=None,
26
    catchup=False,
27
)
28
def complex_dag_structure():
29
30
    start = EmptyOperator(task_id="start")
31
32
    # Dynamically mapped tasks: .partial() contains all parameters that stay the same
33
    # between mapped task instances. .expand() contains the parameters that change
34
    # one mapped task instance will be created for each value in the list provided to `bash_command`
35
    sales_data_extract = BashOperator.partial(task_id="sales_data_extract").expand(
36
        bash_command=["echo 1", "echo 2", "echo 3", "echo 4"]
37
    )
38
    internal_api_extract = BashOperator.partial(task_id="internal_api_extract").expand(
39
        bash_command=["echo 1", "echo 2", "echo 3", "echo 4"]
40
    )
41
42
    # Branch task that picks which branch to follow
43
    @task.branch
44
    def determine_load_type() -> str:
45
        """Randomly choose a branch. The return value is the task_id of the branch to follow."""
46
        import random
47
48
        if random.choice([True, False]):
49
            return "internal_api_load_full"
50
        return "internal_api_load_incremental"
51
52
    sales_data_transform = EmptyOperator(task_id="sales_data_transform")
53
54
    # When using the TaskFlow API it is common to assign the called task to an object
55
    # to use in several dependency definitions without creating several instances of the same task
56
    determine_load_type_obj = determine_load_type()
57
58
    sales_data_load = EmptyOperator(task_id="sales_data_load")
59
    internal_api_load_full = EmptyOperator(task_id="internal_api_load_full")
60
    internal_api_load_incremental = EmptyOperator(
61
        task_id="internal_api_load_incremental"
62
    )
63
64
    # defining a task group that task a parameter (a) to allow for dynamic task group mapping
65
    @task_group
66
    def sales_data_reporting(a):
67
        # the trigger_rule of the first task in the task group can be set to "all_done"
68
        # to ensure that it runs even if one of the upstream tasks (`internal_api_load_full`)
69
        # is skipped due to the branch task
70
        prepare_report = EmptyOperator(
71
            task_id="prepare_report", trigger_rule="all_done"
72
        )
73
        publish_report = EmptyOperator(task_id="publish_report")
74
75
        # setting dependencies within the task group
76
        chain(prepare_report, publish_report)
77
78
    # dynamically mapping the task group with `a` being the mapped parameter
79
    # one task group instance will be created for each value in the list provided to `a`
80
    sales_data_reporting_obj = sales_data_reporting.expand(a=[1, 2, 3, 4, 5, 6])
81
82
    # defining a task group that does not use any additional parameters
83
    @task_group
84
    def cre_integration():
85
        # the trigger_rule of the first task in the task group can be set to "all_done"
86
        # to ensure that it runs even if one of the upstream tasks (`internal_api_load_full`)
87
        # is skipped due to the branch task
88
        cre_extract = EmptyOperator(task_id="cre_extract", trigger_rule="all_done")
89
        cre_transform = EmptyOperator(task_id="cre_transform")
90
        cre_load = EmptyOperator(task_id="cre_load")
91
92
        # setting dependencies within the task group
93
        chain(cre_extract, cre_transform, cre_load)
94
95
    # calling the task group to instantiate it and assigning it to an object
96
    # to use in dependency definitions
97
    cre_integration_obj = cre_integration()
98
99
    @task_group
100
    def mlops():
101
        # the trigger_rule of the first task in the task group can be set to "all_done"
102
        # to ensure that it runs even if one of the upstream tasks (`internal_api_load_incremental`)
103
        # is skipped due to the branch task
104
        set_up_cluster = EmptyOperator(
105
            task_id="set_up_cluster", trigger_rule="all_done"
106
        )
107
        # the outlets parameter is used to define which datasets are updated by this task
108
        train_model = EmptyOperator(
109
            task_id="train_model", outlets=[Dataset("model_trained")]
110
        )
111
        tear_down_cluster = EmptyOperator(task_id="tear_down_cluster")
112
113
        # setting dependencies within the task group
114
        chain(set_up_cluster, train_model, tear_down_cluster)
115
116
        # turning the `tear_down_cluster`` task into a teardown task and set
117
        # the `set_up_cluster` task as the associated setup task
118
        tear_down_cluster.as_teardown(setups=set_up_cluster)
119
120
    mlops_obj = mlops()
121
122
    end = EmptyOperator(task_id="end")
123
124
    # --------------------- #
125
    # Defining dependencies #
126
    # --------------------- #
127
128
    chain(
129
        start,
130
        sales_data_extract,
131
        sales_data_transform,
132
        sales_data_load,
133
        [sales_data_reporting_obj, cre_integration_obj],
134
        end,
135
    )
136
    chain(
137
        start,
138
        internal_api_extract,
139
        determine_load_type_obj,
140
        [internal_api_load_full, internal_api_load_incremental],
141
        mlops_obj,
142
        end,
143
    )
144
145
    chain_linear(
146
        [sales_data_load, internal_api_load_full],
147
        [sales_data_reporting_obj, cre_integration_obj],
148
    )
149
150
    # Adding labels to two edges
151
    chain(
152
        determine_load_type_obj, Label("additional data"), internal_api_load_incremental
153
    )
154
    chain(
155
        determine_load_type_obj, Label("changed existing data"), internal_api_load_full
156
    )
157
158
159
# Calling the DAG function will instantiate the DAG
160
complex_dag_structure()

1
# Import all packages needed at the top level of the DAG
2
from airflow.sdk import dag, task, chain
3
from pendulum import datetime
4
5
6
# Define the DAG function a set of parameters
7
@dag(
8
    start_date=datetime(2025, 4, 1),
9
    schedule="@daily",
10
)
11
def taskflow_dag():
12
    # Define tasks within the DAG context
13
    @task
14
    def my_task_1():
15
        import time  # import packages only needed in the task function
16
17
        time.sleep(5)
18
        print(1)
19
20
    @task
21
    def my_task_2():
22
        print(2)
23
24
    # Define dependencies and call task functions
25
    chain(my_task_1(), my_task_2())
26
27
28
# Call the DAG function
29
taskflow_dag()

1
# Import all packages needed at the top level of the DAG
2
from airflow.sdk import DAG
3
from airflow.providers.standard.operators.python import PythonOperator
4
from pendulum import datetime
5
6
7
def my_task_1_func():
8
    import time  # import packages only needed in the task function
9
10
    time.sleep(5)
11
    print(1)
12
13
14
# Instantiate the DAG
15
with DAG(
16
    dag_id="traditional_syntax_dag",
17
    start_date=datetime(2025, 4, 1),
18
    schedule="@daily",
19
):
20
    # Instantiate tasks within the DAG context
21
    my_task_1 = PythonOperator(
22
        task_id="my_task_1",
23
        python_callable=my_task_1_func,
24
    )
25
26
    my_task_2 = PythonOperator(
27
        task_id="my_task_2",
28
        python_callable=lambda: print(2),
29
    )
30
31
    # Define dependencies
32
    my_task_1 >> my_task_2