Human-in-the-loop workflows with Airflow

1
from airflow.providers.standard.operators.hitl import HITLOperator
2
from airflow.sdk import dag, task, chain, Param
3
4
5
@dag
6
def HITLOperator_syntax_example():
7
8
    @task
9
    def upstream_task():
10
        return "Review expense report and approve vendor payment method."
11
12
    _upstream_task = upstream_task()
13
14
    _hitl_task = HITLOperator(
15
        task_id="hitl_task",
16
        subject="Expense Approval Required",  # templatable
17
        body="{{ ti.xcom_pull(task_ids='upstream_task') }}", # templatable
18
        options=["ACH Transfer", "Wire Transfer", "Corporate Check"],  # cannot be empty!
19
        defaults=["ACH Transfer"],
20
        multiple=False,  # default: False
21
        params={
22
            "expense_amount": Param(
23
                10000,
24
                type="number",
25
            )
26
        },
27
        execution_timeout=timedelta(minutes=5),  # default: None
28
    )
29
30
    @task
31
    def print_result(hitl_output):
32
        print(f"Expense amount: ${hitl_output['params_input']['expense_amount']}")
33
        print(f"Payment method: {hitl_output['chosen_options']}")
34
35
    _print_result = print_result(_hitl_task.output)
36
37
    chain(_upstream_task, _hitl_task)
38
39
40
HITLOperator_syntax_example()

1
from airflow.providers.standard.operators.hitl import ApprovalOperator
2
from airflow.sdk import dag, task, chain, Param
3
4
5
@dag
6
def ApprovalOperator_syntax_example():
7
8
    @task
9
    def upstream_task():
10
        return "Pineapple on pizza?"
11
12
    _upstream_task = upstream_task()
13
14
    _hitl_task = ApprovalOperator(
15
        task_id="approval_task",
16
        subject="Your task:",
17
        body="{{ ti.xcom_pull(task_ids='upstream_task') }}",
18
        defaults="Approve", # other option: "Reject"
19
        params={
20
            "second_topping": Param(
21
                "olives",
22
                type="string",
23
            )
24
        },
25
    )
26
27
    @task
28
    def print_result(hitl_output):
29
        print(f"Params input: {hitl_output['params_input']}")
30
        print(f"Chosen options: {hitl_output['chosen_options']}")
31
32
    _print_result = print_result(_hitl_task.output)
33
34
    chain(_upstream_task, _hitl_task)
35
36
37
ApprovalOperator_syntax_example()

1
from airflow.providers.standard.operators.hitl import HITLBranchOperator
2
from airflow.sdk import dag, task, chain
3
4
5
_budget_categories = ["marketing", "research_development", "facilities", "training", "technology"]
6
7
8
@dag
9
def HITLBranchOperator_syntax_example():
10
11
    @task
12
    def upstream_task():
13
        return {
14
            "total_budget": "$4B",
15
        }
16
17
    _upstream_task = upstream_task()
18
19
    _hitl_branch_task = HITLBranchOperator(
20
        task_id="hitl_branch_task",
21
        subject="Budget Category Approval",
22
        body="""**Total Budget Available:** {{ ti.xcom_pull(task_ids='upstream_task')['total_budget'] }}
23
24
Select the funding proposals to approve for this quarter.""",
25
        options=_budget_categories,
26
        defaults=["marketing", "research_development"],
27
        multiple=True,
28
    )
29
30
    for _category in _budget_categories:
31
32
        @task(
33
            task_id=f"{_category}",  # needs to match options in HITLBranchOperator
34
        )
35
        def category_task():
36
            print(f"Processing budget approval for {_category}")
37
38
        _category_task = category_task()
39
        chain(_hitl_branch_task, _category_task)
40
41
    chain(_upstream_task, _hitl_branch_task)
42
43
44
HITLBranchOperator_syntax_example()

1
from airflow.providers.standard.operators.hitl import HITLEntryOperator
2
from airflow.sdk import dag, task, chain, Param
3
4
5
@dag
6
def HITLEntryOperator_syntax_example():
7
8
    @task
9
    def upstream_task():
10
        return "How can I auto-pause a dag if it fails?"
11
12
    _upstream_task = upstream_task()
13
14
    _hitl_task = HITLEntryOperator(
15
        task_id="hitl_task",
16
        subject="Please respond to this ticket!",
17
        body="{{ ti.xcom_pull(task_ids='upstream_task') }}",
18
        params={
19
            "response": Param(
20
                "You can use the max_consecutive_failed_dag_runs parameter! :)",
21
                type="string",
22
            ),
23
            "urgency": Param(
24
                "p3",
25
                type="string",
26
            ),
27
        },
28
    )
29
30
    @task
31
    def print_result(hitl_output):
32
        print(f"Params input: {hitl_output['params_input']}")
33
        print(f"Chosen options: {hitl_output['chosen_options']}")
34
35
    _print_result = print_result(_hitl_task.output)
36
37
    chain(_upstream_task, _hitl_task)
38
39
40
HITLEntryOperator_syntax_example()

1
from airflow.sdk import BaseNotifier, Context, dag, task, Param
2
from airflow.providers.standard.operators.hitl import HITLOperator
3
from datetime import timedelta
4
5
_BASE_URL = "http://localhost:28080"
6
7
class MyNotifier(BaseNotifier):
8
    template_fields = ("message",)
9
10
    def __init__(self, message: str = "") -> None:
11
        self.message = message
12
13
    def notify(self, context: Context) -> None:
14
        
15
        task_state = context['ti'].state
16
        if task_state == "running":
17
18
            # this method generates a direct link to the UI page where the user can respond
19
            url = HITLOperator.generate_link_to_ui_from_context(
20
                context=context,
21
                base_url=_BASE_URL,
22
            )
23
24
            # placeholder code, you can send the URL to any service you want
25
            self.log.info(self.message)
26
            self.log.info("Url to respond %s", url)
27
        else:
28
            self.log.info("Task state: %s", task_state)
29
            self.log.info("No response needed!")
30
31
32
notifier_class = MyNotifier(
33
    message="""
34
Subject: {{ task.subject }}
35
Body: {{ task.body }}
36
Options: {{ task.options }}
37
"""
38
)
39
40
41
@dag
42
def notifier_example():
43
    HITLOperator(
44
        task_id="hitl_task",
45
        subject="Choose a number: ",
46
        options=["23", "19", "42"],
47
        notifiers=[notifier_class],
48
    )
49
50
51
notifier_example()

1
import requests
2
from typing import Any
3
4
_USERNAME = "admin"
5
_PASSWORD = "admin"
6
_HOST = "http://localhost:28080/"  # To learn how to send API requests to Airflow running on Astro see: https://www.astronomer.io/docs/astro/airflow-api/
7
8
_DAG_ID = "HITLOperator_syntax_example"
9
_TASK_ID = "hitl_task"
10
11
12
def _pick_option(options: list[str]):
13
    print("Available options: ", options)
14
    chosen_option = input("Enter the option you want to select: ")
15
    return chosen_option
16
17
18
def _pick_params(param: dict[str, Any]):
19
    print("Input for param:", param)
20
    param_input = input("Enter your value for the param: ")
21
    return param_input
22
23
24
def _get_jwt_token():
25
    token_url = f"{_HOST}/auth/token"
26
    payload = {"username": _USERNAME, "password": _PASSWORD}
27
    headers = {"Content-Type": "application/json"}
28
    response = requests.post(token_url, json=payload, headers=headers)
29
30
    token = response.json().get("access_token")
31
    return token
32
33
34
def _get_running_dagruns_for_dag(dag_id: str):
35
    url = f"{_HOST}/api/v2/dags/{dag_id}/dagRuns?state=running"
36
    headers = {"Authorization": f"Bearer {_get_jwt_token()}"}
37
    response = requests.get(url, headers=headers)
38
    return response.json()
39
40
41
def _get_hitl_details(dag_id: str, dag_run_id: str, task_id: str):
42
    url = f"{_HOST}/api/v2/hitlDetails/{dag_id}/{dag_run_id}/{task_id}"
43
    headers = {"Authorization": f"Bearer {_get_jwt_token()}"}
44
    response = requests.get(url, headers=headers)
45
    if response.status_code == 200:
46
        subject = response.json()["subject"]
47
        body = response.json()["body"]
48
        options = response.json()["options"]
49
        params = response.json()["params"]
50
        print("--------------------------------")
51
        print("Required Action found for: ", dag_id, "DAG Run: ", dag_run_id, "Task: ", task_id)
52
        print("Subject: ", subject)
53
        print("Body: ", body)
54
        print("Options: ", options)
55
        print("Params: ", params)
56
        print("--------------------------------")
57
        return {
58
            "subject": subject,
59
            "body": body,
60
            "options": options,
61
            "params": params,
62
        }
63
    elif response.status_code == 404:
64
        print("--------------------------------")
65
        print("404 - No required action found for: ", dag_id, "DAG Run: ", dag_run_id, "Task: ", task_id)
66
        print("Response: ", response.json())
67
        print("--------------------------------")
68
        return None
69
    else:
70
        print("--------------------------------")
71
        print("Error: ", response.status_code)
72
        print("Response: ", response.json())
73
        print("--------------------------------")
74
        return None
75
76
77
def _add_hitl_response(
78
    dag_id: str, dag_run_id: str, task_id: str, options: list[str], params: dict[str, Any]
79
):
80
    url = f"{_HOST}/api/v2/hitlDetails/{dag_id}/{dag_run_id}/{task_id}"
81
    headers = {"Authorization": f"Bearer {_get_jwt_token()}"}
82
    chosen_options = [_pick_option(options)]
83
    if params:
84
        params_input = {f"{param}": _pick_params(param) for param in params}
85
    else:
86
        params_input = {}
87
    response = requests.patch(
88
        url, headers=headers, json={"chosen_options": chosen_options, "params_input": params_input}
89
    )
90
    if response.status_code == 200:
91
        print("--------------------------------")
92
        print("Hitl response added for DAG: ", dag_id, "DAG Run: ", dag_run_id, "Task: ", task_id)
93
        print("Chosen options: ", chosen_options)
94
        print("Params input: ", params_input)
95
        print("Response status code: ", response.status_code)
96
        print("Response: ", response.json())
97
        print("--------------------------------")
98
    elif response.status_code == 409:
99
        print("--------------------------------")
100
        print("409 - Already updated action for: ", dag_id, "DAG Run: ", dag_run_id, "Task: ", task_id)
101
        print("Response: ", response.json())
102
        print("--------------------------------")
103
    else:
104
        print("--------------------------------")
105
        print("Error: ", response.status_code)
106
        print("Response: ", response.json())
107
        print("--------------------------------")
108
109
110
def main():
111
    dag_runs = _get_running_dagruns_for_dag(_DAG_ID)["dag_runs"]
112
    if not dag_runs:
113
        print("No running dag runs found for DAG: ", _DAG_ID)
114
    else:
115
        for dag_run in dag_runs:
116
            hitl_details = _get_hitl_details(_DAG_ID, dag_run["dag_run_id"], _TASK_ID)
117
            if hitl_details:
118
                _add_hitl_response(
119
                    _DAG_ID, dag_run["dag_run_id"], _TASK_ID, hitl_details["options"], hitl_details["params"]
120
                )
121
122
123
if __name__ == "__main__":
124
    main()