You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2026/04/28 07:35:45 UTC
(airflow) branch main updated: Allow accessing a Dag's members via `[]` (#65586)
This is an automated email from the ASF dual-hosted git repository.
potiuk 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 d4d83acd60d Allow accessing a Dag's members via `[]` (#65586)
d4d83acd60d is described below
commit d4d83acd60db59aa75ecab2292721dd40dfaa42a
Author: Dev-iL <65...@users.noreply.github.com>
AuthorDate: Tue Apr 28 10:35:36 2026 +0300
Allow accessing a Dag's members via `[]` (#65586)
* Add `DAG.__getitem__` for subscript (`[]`) task access
- introduce `TaskItemNotFound` (a `KeyError` subtype) to the SDK exceptions module.
- Enable `dag[task_id]` syntax, raising `TaskItemNotFound` on miss.
- Add `SerializedDAG.__getitem__` for consistency.
* Extend `DAG.__getitem__` to return tasks and task groups
`dag[id]` now searches task_dict then task_group_dict, so both tasks and groups are reachable by their fully-qualified ID. SerializedDAG follows the same logic. Both chained access (`dag["group"]["task"]`) and qualified-id access (`dag["group.task"]`) are covered by new tests.
* Rename `TaskItemNotFound` to `NodeNotFound`
---
airflow-core/src/airflow/exceptions.py | 7 +++
.../src/airflow/serialization/definitions/dag.py | 10 ++++-
.../unit/serialization/test_serialized_objects.py | 51 ++++++++++++++++++++++
task-sdk/src/airflow/sdk/definitions/dag.py | 11 ++++-
task-sdk/src/airflow/sdk/definitions/taskgroup.py | 11 +++--
task-sdk/src/airflow/sdk/exceptions.py | 7 +++
task-sdk/tests/task_sdk/definitions/test_dag.py | 41 +++++++++++++++++
.../tests/task_sdk/definitions/test_taskgroup.py | 31 ++++++++++++-
task-sdk/tests/task_sdk/test_exceptions.py | 43 ++++++++++++++++++
9 files changed, 204 insertions(+), 8 deletions(-)
diff --git a/airflow-core/src/airflow/exceptions.py b/airflow-core/src/airflow/exceptions.py
index 64c4b8ee92a..f9a891472a3 100644
--- a/airflow-core/src/airflow/exceptions.py
+++ b/airflow-core/src/airflow/exceptions.py
@@ -37,6 +37,7 @@ try:
AirflowOptionalProviderFeatureException as AirflowOptionalProviderFeatureException,
AirflowRescheduleException as AirflowRescheduleException,
AirflowTimetableInvalid as AirflowTimetableInvalid,
+ NodeNotFound as NodeNotFound,
TaskNotFound as TaskNotFound,
)
except ModuleNotFoundError:
@@ -54,6 +55,12 @@ except ModuleNotFoundError:
class TaskNotFound(AirflowException): # type: ignore[no-redef]
"""Raise when a Task is not available in the system."""
+ class NodeNotFound(TaskNotFound, KeyError): # type: ignore[no-redef]
+ """Raise when attempting to access an invalid node (task or task group) using [] notation."""
+
+ def __str__(self) -> str:
+ return str(self.args[0]) if self.args else ""
+
class AirflowRescheduleException(AirflowException): # type: ignore[no-redef]
"""
Raise when the task should be re-scheduled at a later time.
diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py b/airflow-core/src/airflow/serialization/definitions/dag.py
index f9e5a22a91d..3fc4a8f6674 100644
--- a/airflow-core/src/airflow/serialization/definitions/dag.py
+++ b/airflow-core/src/airflow/serialization/definitions/dag.py
@@ -33,7 +33,7 @@ from sqlalchemy import func, or_, select, tuple_
from airflow._shared.observability.metrics.stats import Stats
from airflow._shared.timezones.timezone import coerce_datetime
from airflow.configuration import conf as airflow_conf
-from airflow.exceptions import AirflowException, TaskNotFound
+from airflow.exceptions import AirflowException, NodeNotFound, TaskNotFound
from airflow.models.dag import DagModel
from airflow.models.dag_version import DagVersion
from airflow.models.dagrun import DagRun
@@ -246,6 +246,14 @@ class SerializedDAG:
return self.task_dict[task_id]
raise TaskNotFound(f"Task {task_id} not found")
+ def __getitem__(self, node_id: str) -> SerializedOperator | SerializedTaskGroup:
+ """Return a task or task group by its fully-qualified ID."""
+ if (node := self.task_dict.get(node_id)) is not None:
+ return node
+ if (tg := self.task_group_dict.get(node_id)) is not None:
+ return tg
+ raise NodeNotFound(f"Task or group {node_id!r} not found")
+
@property
def task_group_dict(self):
return {k: v for k, v in self.task_group.get_task_group_dict().items() if k is not None}
diff --git a/airflow-core/tests/unit/serialization/test_serialized_objects.py b/airflow-core/tests/unit/serialization/test_serialized_objects.py
index f48afd9c002..c41b1670c6a 100644
--- a/airflow-core/tests/unit/serialization/test_serialized_objects.py
+++ b/airflow-core/tests/unit/serialization/test_serialized_objects.py
@@ -1037,3 +1037,54 @@ class TestKubernetesImportAvoidance:
result = _has_kubernetes()
assert result is True
+
+
+@pytest.mark.db_test
+def test_serialized_dag_getitem_returns_task(dag_maker):
+ from airflow.serialization.definitions.dag import SerializedDAG
+
+ with dag_maker() as dag:
+ BashOperator(task_id="my_task", bash_command="echo 1")
+
+ ser_dict = DagSerialization.to_dict(dag)
+ ser_dag = DagSerialization.from_dict(ser_dict)
+ assert isinstance(ser_dag, SerializedDAG)
+ assert ser_dag["my_task"].task_id == "my_task"
+
+
+@pytest.mark.db_test
+def test_serialized_dag_getitem_missing_raises_node_not_found(dag_maker):
+ from airflow.exceptions import NodeNotFound
+
+ with dag_maker() as dag:
+ BashOperator(task_id="my_task", bash_command="echo 1")
+
+ ser_dict = DagSerialization.to_dict(dag)
+ ser_dag = DagSerialization.from_dict(ser_dict)
+ with pytest.raises(NodeNotFound):
+ ser_dag["nonexistent"]
+
+
+@pytest.mark.db_test
+def test_serialized_dag_getitem_missing_is_key_error(dag_maker):
+ with dag_maker() as dag:
+ BashOperator(task_id="my_task", bash_command="echo 1")
+
+ ser_dict = DagSerialization.to_dict(dag)
+ ser_dag = DagSerialization.from_dict(ser_dict)
+ with pytest.raises(KeyError):
+ ser_dag["nonexistent"]
+
+
+@pytest.mark.db_test
+def test_serialized_dag_getitem_returns_task_group(dag_maker):
+ from airflow.serialization.definitions.dag import SerializedDAG
+
+ with dag_maker() as dag:
+ with TaskGroup(group_id="section") as tg:
+ BashOperator(task_id="t", bash_command="echo 1")
+
+ ser_dict = DagSerialization.to_dict(dag)
+ ser_dag = DagSerialization.from_dict(ser_dict)
+ assert isinstance(ser_dag, SerializedDAG)
+ assert ser_dag["section"].group_id == tg.group_id
diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py
index 3df36692513..ab85396335e 100644
--- a/task-sdk/src/airflow/sdk/definitions/dag.py
+++ b/task-sdk/src/airflow/sdk/definitions/dag.py
@@ -43,7 +43,7 @@ from airflow import settings
from airflow.sdk import TaskInstanceState, TriggerRule, XComArg
from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.bases.timetable import BaseTimetable
-from airflow.sdk.definitions._internal.node import validate_key
+from airflow.sdk.definitions._internal.node import DAGNode, validate_key
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set
from airflow.sdk.definitions.asset import AssetAll, BaseAsset
from airflow.sdk.definitions.context import Context
@@ -55,6 +55,7 @@ from airflow.sdk.exceptions import (
AirflowDagCycleException,
DuplicateTaskIdFound,
FailFastDagInvalidTriggerRule,
+ NodeNotFound,
ParamValidationError,
RemovedInAirflow4Warning,
TaskNotFound,
@@ -1032,6 +1033,14 @@ class DAG:
return self.task_dict[task_id]
raise TaskNotFound(f"Task {task_id} not found")
+ def __getitem__(self, node_id: str) -> DAGNode:
+ """Return a task or task group by its fully-qualified ID."""
+ if (node := self.task_dict.get(node_id)) is not None:
+ return node
+ if (tg := self.task_group_dict.get(node_id)) is not None:
+ return tg
+ raise NodeNotFound(f"Task or group {node_id!r} not found")
+
@property
def task(self) -> TaskDecoratorCollection:
from airflow.sdk.definitions.decorators import task
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
index 72779adb5e8..50527f6b43b 100644
--- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
+++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -33,6 +33,7 @@ from airflow.sdk.definitions._internal.node import DAGNode, validate_group_key
from airflow.sdk.exceptions import (
AirflowDagCycleException,
DuplicateTaskIdFound,
+ NodeNotFound,
TaskAlreadyInTaskGroup,
)
@@ -141,10 +142,6 @@ class TaskGroup(DAGNode):
if not dag:
raise RuntimeError("TaskGroup can only be used inside a dag")
- def __getitem__(self, key: str) -> DAGNode:
- """Return a child node by label via ``tg[label]`` syntax. See `get_child_by_label`."""
- return self.get_child_by_label(key)
-
def __attrs_post_init__(self):
# TODO: If attrs supported init only args we could use that here
# https://github.com/python-attrs/attrs/issues/342
@@ -496,6 +493,12 @@ class TaskGroup(DAGNode):
"""Get a child task/TaskGroup by its label (i.e. task_id/group_id without the group_id prefix)."""
return self.children[self.child_id(label)]
+ def __getitem__(self, label: str) -> DAGNode:
+ try:
+ return self.get_child_by_label(label)
+ except KeyError:
+ raise NodeNotFound(f"Task {label!r} not found")
+
def serialize_for_task_group(self) -> tuple[DagAttributeTypes, Any]:
"""Serialize task group; required by DagNode."""
from airflow.sdk.api.datamodels._generated import DagAttributeTypes
diff --git a/task-sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py
index b69abe62265..c7698ab75ff 100644
--- a/task-sdk/src/airflow/sdk/exceptions.py
+++ b/task-sdk/src/airflow/sdk/exceptions.py
@@ -330,6 +330,13 @@ class TaskNotFound(AirflowException):
"""Raise when a Task is not available in the system."""
+class NodeNotFound(TaskNotFound, KeyError):
+ """Raise when attempting to access an invalid node (task or task group) using [] notation."""
+
+ def __str__(self) -> str:
+ return str(self.args[0]) if self.args else ""
+
+
class TaskAlreadyRunningError(AirflowException):
"""Raised when a task is already running on another worker."""
diff --git a/task-sdk/tests/task_sdk/definitions/test_dag.py b/task-sdk/tests/task_sdk/definitions/test_dag.py
index 263259240d2..a7d1bfbd926 100644
--- a/task-sdk/tests/task_sdk/definitions/test_dag.py
+++ b/task-sdk/tests/task_sdk/definitions/test_dag.py
@@ -941,3 +941,44 @@ class TestCycleTester:
op1 >> Label("label") >> op2
assert not dag.check_cycle()
+
+
+class TestDagGetItem:
+ def test_getitem_returns_task(self):
+ dag = DAG("test_dag", schedule=None, start_date=DEFAULT_DATE)
+ with dag:
+ op = DoNothingOperator(task_id="my_task")
+ assert dag["my_task"] is op
+
+ def test_getitem_returns_task_group(self):
+ dag = DAG("test_dag", schedule=None, start_date=DEFAULT_DATE)
+ with dag:
+ with TaskGroup(group_id="section") as tg:
+ DoNothingOperator(task_id="t")
+ assert dag["section"] is tg
+
+ def test_getitem_nested_task_by_qualified_id(self):
+ dag = DAG("test_dag", schedule=None, start_date=DEFAULT_DATE)
+ with dag:
+ with TaskGroup(group_id="section"):
+ op = DoNothingOperator(task_id="t")
+ assert dag["section.t"] is op
+
+ def test_getitem_nested_task_via_chained_access(self):
+ dag = DAG("test_dag", schedule=None, start_date=DEFAULT_DATE)
+ with dag:
+ with TaskGroup(group_id="section"):
+ op = DoNothingOperator(task_id="t")
+ assert dag["section"]["t"] is op
+
+ def test_getitem_missing_raises_node_not_found(self):
+ from airflow.sdk.exceptions import NodeNotFound
+
+ dag = DAG("test_dag", schedule=None, start_date=DEFAULT_DATE)
+ with pytest.raises(NodeNotFound):
+ dag["nonexistent"]
+
+ def test_getitem_missing_is_key_error(self):
+ dag = DAG("test_dag", schedule=None, start_date=DEFAULT_DATE)
+ with pytest.raises(KeyError):
+ dag["nonexistent"]
diff --git a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
index 5eb69dddd4c..ed03e6c6773 100644
--- a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
+++ b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
@@ -243,8 +243,9 @@ def test_taskgroup_getitem_returns_child_by_label(prefix: bool):
assert group1["subgroup"] == subgroup
assert group1["subgroup"]["task2"] == task2
- # Missing label raises KeyError
- with pytest.raises(KeyError):
+ from airflow.sdk.exceptions import NodeNotFound
+
+ with pytest.raises(NodeNotFound):
group1["nonexistent"]
@@ -930,3 +931,29 @@ def test_build_task_group_with_operators():
# Testing Tasks downstream
assert dag.task_dict["task_start"].downstream_task_ids == {"section_1.task_1"}
assert dag.task_dict["section_1.task_3"].downstream_task_ids == {"task_end"}
+
+
+class TestTaskGroupGetItem:
+ def test_getitem_missing_raises_node_not_found(self):
+ import pendulum
+
+ from airflow.sdk.exceptions import NodeNotFound
+
+ start = pendulum.datetime(2016, 1, 1)
+ with DAG("test_dag", schedule=None, start_date=start):
+ with TaskGroup(group_id="section") as tg:
+ pass
+
+ with pytest.raises(NodeNotFound):
+ tg["nonexistent"]
+
+ def test_getitem_missing_is_key_error(self):
+ import pendulum
+
+ start = pendulum.datetime(2016, 1, 1)
+ with DAG("test_dag", schedule=None, start_date=start):
+ with TaskGroup(group_id="section") as tg:
+ pass
+
+ with pytest.raises(KeyError):
+ tg["nonexistent"]
diff --git a/task-sdk/tests/task_sdk/test_exceptions.py b/task-sdk/tests/task_sdk/test_exceptions.py
new file mode 100644
index 00000000000..2a1b5d5a683
--- /dev/null
+++ b/task-sdk/tests/task_sdk/test_exceptions.py
@@ -0,0 +1,43 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import pytest
+
+from airflow.sdk.exceptions import NodeNotFound, TaskNotFound
+
+
+def test_node_not_found_is_subclass_of_task_not_found():
+ assert issubclass(NodeNotFound, TaskNotFound)
+
+
+def test_node_not_found_is_subclass_of_key_error():
+ assert issubclass(NodeNotFound, KeyError)
+
+
+def test_node_not_found_caught_as_key_error():
+ with pytest.raises(KeyError):
+ raise NodeNotFound("missing_task")
+
+
+def test_node_not_found_caught_as_task_not_found():
+ with pytest.raises(TaskNotFound):
+ raise NodeNotFound("missing_task")
+
+
+def test_node_not_found_str_suppresses_key_error_repr():
+ assert str(NodeNotFound("missing")) == "missing"