Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/airflow/utils/dag_cycle_tester.py: 28%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements. See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership. The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License. You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied. See the License for the
15# specific language governing permissions and limitations
16# under the License.
17"""DAG Cycle tester."""
19from __future__ import annotations
21from collections import defaultdict, deque
22from typing import TYPE_CHECKING
24from airflow.exceptions import AirflowDagCycleException, RemovedInAirflow3Warning
26if TYPE_CHECKING:
27 from airflow.models.dag import DAG
29CYCLE_NEW = 0
30CYCLE_IN_PROGRESS = 1
31CYCLE_DONE = 2
34def test_cycle(dag: DAG) -> None:
35 """
36 A wrapper function of `check_cycle` for backward compatibility purpose.
38 New code should use `check_cycle` instead since this function name `test_cycle` starts
39 with 'test_' and will be considered as a unit test by pytest, resulting in failure.
40 """
41 from warnings import warn
43 warn(
44 "Deprecated, please use `check_cycle` at the same module instead.",
45 RemovedInAirflow3Warning,
46 stacklevel=2,
47 )
48 return check_cycle(dag)
51def check_cycle(dag: DAG) -> None:
52 """Check to see if there are any cycles in the DAG.
54 :raises AirflowDagCycleException: If cycle is found in the DAG.
55 """
56 # default of int is 0 which corresponds to CYCLE_NEW
57 visited: dict[str, int] = defaultdict(int)
58 path_stack: deque[str] = deque()
59 task_dict = dag.task_dict
61 def _check_adjacent_tasks(task_id, current_task):
62 """Return first untraversed child task, else None if all tasks traversed."""
63 for adjacent_task in current_task.get_direct_relative_ids():
64 if visited[adjacent_task] == CYCLE_IN_PROGRESS:
65 msg = f"Cycle detected in DAG: {dag.dag_id}. Faulty task: {task_id}"
66 raise AirflowDagCycleException(msg)
67 elif visited[adjacent_task] == CYCLE_NEW:
68 return adjacent_task
69 return None
71 for dag_task_id in dag.task_dict.keys():
72 if visited[dag_task_id] == CYCLE_DONE:
73 continue
74 path_stack.append(dag_task_id)
75 while path_stack:
76 current_task_id = path_stack[-1]
77 if visited[current_task_id] == CYCLE_NEW:
78 visited[current_task_id] = CYCLE_IN_PROGRESS
79 task = task_dict[current_task_id]
80 child_to_check = _check_adjacent_tasks(current_task_id, task)
81 if not child_to_check:
82 visited[current_task_id] = CYCLE_DONE
83 path_stack.pop()
84 else:
85 path_stack.append(child_to_check)