This is an automated email from the ASF dual-hosted git repository.
weilee pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 94febf24170 Add options_mapping support to HITLBranchOperator (#55093)
94febf24170 is described below
commit 94febf24170771204b3effd0fcaad3164d9da854
Author: Constance Martineau <[email protected]>
AuthorDate: Mon Sep 1 10:39:52 2025 -0400
Add options_mapping support to HITLBranchOperator (#55093)
Co-authored-by: Wei Lee <[email protected]>
---
.../airflow/providers/standard/operators/hitl.py | 48 +++++++++
.../tests/unit/standard/operators/test_hitl.py | 110 ++++++++++++++++++++-
2 files changed, 157 insertions(+), 1 deletion(-)
diff --git
a/providers/standard/src/airflow/providers/standard/operators/hitl.py
b/providers/standard/src/airflow/providers/standard/operators/hitl.py
index 75588f8cbeb..440abf9a516 100644
--- a/providers/standard/src/airflow/providers/standard/operators/hitl.py
+++ b/providers/standard/src/airflow/providers/standard/operators/hitl.py
@@ -350,9 +350,57 @@ class HITLBranchOperator(HITLOperator, BranchMixIn):
inherits_from_skipmixin = True
+ def __init__(self, *, options_mapping: dict[str, str] | None = None,
**kwargs) -> None:
+ """
+ Initialize HITLBranchOperator.
+
+ Args:
+ options_mapping:
+ A dictionary mapping option labels (must match entries in
`self.options`)
+ to string values (e.g., task IDs). Defaults to an empty dict
if not provided.
+
+ Raises:
+ ValueError:
+ - If `options_mapping` contains keys not present in
`self.options`.
+ - If any value in `options_mapping` is not a string.
+ """
+ super().__init__(**kwargs)
+ self.options_mapping = options_mapping or {}
+ self.validate_options_mapping()
+
+ def validate_options_mapping(self) -> None:
+ """
+ Validate that `options_mapping` keys match `self.options` and all
values are strings.
+
+ Raises:
+ ValueError: If any key is not in `self.options` or any value is
not a string.
+ """
+ if not self.options_mapping:
+ return
+
+ # Validate that the choice options are keys in the mapping are the same
+ invalid_keys = set(self.options_mapping.keys()) - set(self.options)
+ if invalid_keys:
+ raise ValueError(
+ f"`options_mapping` contains keys that are not in `options`:
{sorted(invalid_keys)}"
+ )
+
+ # validate that all values are strings
+ invalid_entries = {
+ k: (v, type(v).__name__) for k, v in self.options_mapping.items()
if not isinstance(v, str)
+ }
+ if invalid_entries:
+ raise ValueError(
+ f"`options_mapping` values must be strings
(task_ids).\nInvalid entries: {invalid_entries}"
+ )
+
def execute_complete(self, context: Context, event: dict[str, Any]) -> Any:
+ """Execute the operator and branch based on chosen options."""
ret = super().execute_complete(context=context, event=event)
chosen_options = ret["chosen_options"]
+
+ # Map options to task IDs using the mapping, fallback to original
option
+ chosen_options = [self.options_mapping.get(option, option) for option
in chosen_options]
return self.do_branch(context=context,
branches_to_execute=chosen_options)
diff --git a/providers/standard/tests/unit/standard/operators/test_hitl.py
b/providers/standard/tests/unit/standard/operators/test_hitl.py
index 2257c4e6198..ca0b0ae9f4f 100644
--- a/providers/standard/tests/unit/standard/operators/test_hitl.py
+++ b/providers/standard/tests/unit/standard/operators/test_hitl.py
@@ -31,7 +31,7 @@ from urllib.parse import parse_qs, urlparse
import pytest
from sqlalchemy import select
-from airflow.exceptions import DownstreamTasksSkipped
+from airflow.exceptions import AirflowException, DownstreamTasksSkipped
from airflow.models import TaskInstance, Trigger
from airflow.models.hitl import HITLDetail
from airflow.providers.standard.operators.empty import EmptyOperator
@@ -505,3 +505,111 @@ class TestHITLBranchOperator:
},
)
assert set(exc_info.value.tasks) == set((f"branch_{i}", -1) for i in
range(4, 6))
+
+ def test_mapping_applies_for_single_choice(self, dag_maker):
+ # ["Approve"]; map -> "publish"
+ with dag_maker("hitl_map_dag", serialized=True):
+ op = HITLBranchOperator(
+ task_id="choose",
+ subject="S",
+ options=["Approve", "Reject"],
+ options_mapping={"Approve": "publish"},
+ )
+ op >> [EmptyOperator(task_id="publish"),
EmptyOperator(task_id="archive")]
+
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance("choose")
+
+ with pytest.raises(DownstreamTasksSkipped) as exc:
+ op.execute_complete(
+ context={"ti": ti, "task": ti.task},
+ event={"chosen_options": ["Approve"], "params_input": {}},
+ )
+ # checks to see that the "archive" task was skipped
+ assert set(exc.value.tasks) == {("archive", -1)}
+
+ def test_mapping_with_multiple_choices(self, dag_maker):
+ # multiple=True; mapping applied per option; no dedup implied
+ with dag_maker("hitl_map_dag", serialized=True):
+ op = HITLBranchOperator(
+ task_id="choose",
+ subject="S",
+ multiple=True,
+ options=["Approve", "KeepAsIs"],
+ options_mapping={"Approve": "publish", "KeepAsIs": "keep"},
+ )
+ op >> [
+ EmptyOperator(task_id="publish"),
+ EmptyOperator(task_id="keep"),
+ EmptyOperator(task_id="other"),
+ ]
+
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance("choose")
+
+ with pytest.raises(DownstreamTasksSkipped) as exc:
+ op.execute_complete(
+ context={"ti": ti, "task": ti.task},
+ event={"chosen_options": ["Approve", "KeepAsIs"],
"params_input": {}},
+ )
+ # publish + keep chosen → only "other" skipped
+ assert set(exc.value.tasks) == {("other", -1)}
+
+ def test_fallback_to_option_when_not_mapped(self, dag_maker):
+ # No mapping: option must match downstream task_id
+ with dag_maker("hitl_map_dag", serialized=True):
+ op = HITLBranchOperator(
+ task_id="choose",
+ subject="S",
+ options=["branch_1", "branch_2"], # no mapping for branch_2
+ )
+ op >> [EmptyOperator(task_id="branch_1"),
EmptyOperator(task_id="branch_2")]
+
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance("choose")
+
+ with pytest.raises(DownstreamTasksSkipped) as exc:
+ op.execute_complete(
+ context={"ti": ti, "task": ti.task},
+ event={"chosen_options": ["branch_2"], "params_input": {}},
+ )
+ assert set(exc.value.tasks) == {("branch_1", -1)}
+
+ def test_error_if_mapped_branch_not_direct_downstream(self, dag_maker):
+ # Don't add the mapped task downstream → expect a clean error
+ with dag_maker("hitl_map_dag", serialized=True):
+ op = HITLBranchOperator(
+ task_id="choose",
+ subject="S",
+ options=["Approve"],
+ options_mapping={"Approve": "not_a_downstream"},
+ )
+ # Intentionally no downstream "not_a_downstream"
+
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance("choose")
+
+ with pytest.raises(AirflowException, match="downstream|not found"):
+ op.execute_complete(
+ context={"ti": ti, "task": ti.task},
+ event={"chosen_options": ["Approve"], "params_input": {}},
+ )
+
+ @pytest.mark.parametrize("bad", [123, ["publish"], {"x": "y"}, b"publish"])
+ def test_options_mapping_non_string_value_raises(self, bad):
+ with pytest.raises(ValueError, match=r"values must be strings
\(task_ids\)"):
+ HITLBranchOperator(
+ task_id="choose",
+ subject="S",
+ options=["Approve"],
+ options_mapping={"Approve": bad},
+ )
+
+ def test_options_mapping_key_not_in_options_raises(self):
+ with pytest.raises(ValueError, match="contains keys that are not in
`options`"):
+ HITLBranchOperator(
+ task_id="choose",
+ subject="S",
+ options=["Approve", "Reject"],
+ options_mapping={"NotAnOption": "publish"},
+ )
