Move Dag cycle detection from the Task SDK to shared

shared/dagnode/src/airflow_shared/dagnode/cycle.py

@@ -0,0 +1,73 @@
+# 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
+
+from collections import defaultdict, deque
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+# default of int is 0 which corresponds to _NEW
+_NEW = 0
+_IN_PROGRESS = 1
+_DONE = 2
+
+__all__ = ["detect_cycle"]
+
+
+def detect_cycle(node_ids: Iterable[str], downstream_of: Callable[[str], Iterable[str]]) -> str | None:
+    """
+    Search the graph for a cycle, following downstream edges only.
+
+    The graph is described by callbacks rather than by node objects so that callers
+    holding a serialized Dag can answer from raw data without hydrating it.
+
+    :param node_ids: Every node in the graph, so that disconnected components are covered too.
+    :param downstream_of: Returns the ids directly downstream of the given node id.
+    :return: The id of the node whose downstream edge closes a cycle, or None when acyclic.
+    """
+    visited: dict[str, int] = defaultdict(int)
+    path_stack: deque[str] = deque()
+
+    for node_id in node_ids:
+        if visited[node_id] == _DONE:
+            continue
+        path_stack.append(node_id)
+        # Iterative rather than recursive: Python's stack depth limit is easily
+        # reached by a graph with long chains.
+        while path_stack:
+            current_id = path_stack[-1]
+            if visited[current_id] == _NEW:
+                visited[current_id] = _IN_PROGRESS
+
+            child_to_check = None
+            for downstream_id in downstream_of(current_id):
+                if visited[downstream_id] == _IN_PROGRESS:
+                    return current_id
+                if visited[downstream_id] == _NEW:
+                    child_to_check = downstream_id
+                    break
+
+            if child_to_check is None:
+                visited[current_id] = _DONE
+                path_stack.pop()
+            else:
+                path_stack.append(child_to_check)
+
+    return None

shared/dagnode/tests/dagnode/test_cycle.py

@@ -0,0 +1,79 @@
+# 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 sys
+
+import pytest
+
+from airflow_shared.dagnode.cycle import detect_cycle
+
+
+def edges_of(graph: dict[str, list[str]]):
+    return lambda node_id: graph[node_id]
+
+
+@pytest.mark.parametrize(
+    "graph",
+    [
+        pytest.param({}, id="empty"),
+        pytest.param({"a": []}, id="single-node"),
+        pytest.param({"a": ["b"], "b": ["c"], "c": []}, id="chain"),
+        pytest.param({"a": ["b", "c"], "b": ["d"], "c": ["d"], "d": []}, id="diamond"),
+        pytest.param({"a": ["b"], "b": [], "c": ["d"], "d": []}, id="disconnected-components"),
+    ],
+)
+def test_returns_none_when_acyclic(graph):
+    assert detect_cycle(graph, edges_of(graph)) is None
+
+
+@pytest.mark.parametrize(
+    ("graph", "expected"),
+    [
+        pytest.param({"a": ["a"]}, "a", id="self-loop"),
+        pytest.param({"a": ["b"], "b": ["a"]}, "b", id="two-node-cycle"),
+        pytest.param({"a": ["b"], "b": ["c"], "c": ["a"]}, "c", id="three-node-cycle"),
+        pytest.param({"a": ["b"], "b": ["c"], "c": ["b"]}, "c", id="cycle-below-an-acyclic-root"),
+        pytest.param({"a": [], "b": ["c"], "c": ["b"]}, "c", id="cycle-outside-the-first-component"),
+        pytest.param({"a": [""], "": ["a"]}, "", id="empty-task-id"),
+    ],
+)
+def test_returns_the_node_whose_edge_closes_the_cycle(graph, expected):
+    assert detect_cycle(graph, edges_of(graph)) == expected
+
+
+def test_reads_the_graph_only_through_the_callback():
+    """Callers may answer from any shape, such as raw serialized Dag data."""
+    serialized = {
+        "tasks": [
+            {"task_id": "a", "downstream_task_ids": ["b"]},
+            {"task_id": "b", "downstream_task_ids": ["a"]},
+        ]
+    }
+    by_id = {task["task_id"]: task for task in serialized["tasks"]}
+
+    assert detect_cycle(by_id, lambda node_id: by_id[node_id]["downstream_task_ids"]) == "b"
+
+
+@pytest.mark.parametrize("cyclic", [False, True])
+def test_handles_a_graph_deeper_than_the_recursion_limit(cyclic):
+    depth = sys.getrecursionlimit() * 2
+    graph = {str(i): [str(i + 1)] for i in range(depth)}
+    graph[str(depth)] = ["0"] if cyclic else []
+
+    assert detect_cycle(graph, edges_of(graph)) == (str(depth) if cyclic else None)

task-sdk/src/airflow/sdk/definitions/dag.py

@@ -27,7 +27,7 @@
 import sys
 import warnings
 import weakref
-from collections import abc, defaultdict, deque
+from collections import abc
 from collections.abc import Callable, Collection, Iterable, MutableSet
 from datetime import datetime, timedelta
 from inspect import signature
@@ -41,6 +41,7 @@
 
 from airflow import settings
 from airflow.sdk import TaskInstanceState, TriggerRule, XComArg
+from airflow.sdk._shared.dagnode.cycle import detect_cycle
 from airflow.sdk.bases.operator import BaseOperator
 from airflow.sdk.bases.timetable import BaseTimetable
 from airflow.sdk.definitions._internal.node import DAGNode, validate_key
@@ -1142,40 +1143,12 @@ def check_cycle(self) -> None:
 
         :raises AirflowDagCycleException: If cycle is found in the Dag.
         """
-        # default of int is 0 which corresponds to CYCLE_NEW
-        CYCLE_NEW = 0
-        CYCLE_IN_PROGRESS = 1
-        CYCLE_DONE = 2
-
-        visited: dict[str, int] = defaultdict(int)
-        path_stack: deque[str] = deque()
         task_dict = self.task_dict
-
-        def _check_adjacent_tasks(task_id, current_task):
-            """Return first untraversed child task, else None if all tasks traversed."""
-            for adjacent_task in current_task.get_direct_relative_ids():
-                if visited[adjacent_task] == CYCLE_IN_PROGRESS:
-                    msg = f"Cycle detected in Dag: {self.dag_id}. Faulty task: {task_id}"
-                    raise AirflowDagCycleException(msg)
-                if visited[adjacent_task] == CYCLE_NEW:
-                    return adjacent_task
-            return None
-
-        for dag_task_id in self.task_dict.keys():
-            if visited[dag_task_id] == CYCLE_DONE:
-                continue
-            path_stack.append(dag_task_id)
-            while path_stack:
-                current_task_id = path_stack[-1]
-                if visited[current_task_id] == CYCLE_NEW:
-                    visited[current_task_id] = CYCLE_IN_PROGRESS
-                task = task_dict[current_task_id]
-                child_to_check = _check_adjacent_tasks(current_task_id, task)
-                if not child_to_check:
-                    visited[current_task_id] = CYCLE_DONE
-                    path_stack.pop()
-                else:
-                    path_stack.append(child_to_check)
+        faulty_task_id = detect_cycle(task_dict, lambda tid: task_dict[tid].get_direct_relative_ids())
+        if faulty_task_id is not None:
+            raise AirflowDagCycleException(
+                f"Cycle detected in Dag: {self.dag_id}. Faulty task: {faulty_task_id}"
+            )
 
     def cli(self):
         """Exposes a CLI specific to this Dag."""