Skip to content

Commit

Permalink
Rewrite recursion into iteration (apache#26175)
Browse files Browse the repository at this point in the history
This helps to avoid RecursionError when viewing the graph
view of a dag with many tasks
  • Loading branch information
ephraimbuddy authored Sep 6, 2022
1 parent 55928b9 commit 8538a72
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 4 deletions.
13 changes: 9 additions & 4 deletions airflow/www/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -597,11 +597,16 @@ def collect_edges(task_group):
edges = set()

def get_downstream(task):
for child in task.downstream_list:
edge = (task.task_id, child.task_id)
if edge not in edges:
tasks_to_trace = task.downstream_list
while tasks_to_trace:
tasks_to_trace_next: Set[str] = set()
for child in tasks_to_trace:
edge = (task.task_id, child.task_id)
if edge in edges:
continue
tasks_to_trace_next.update(child.downstream_list)
edges.add(edge)
get_downstream(child)
tasks_to_trace = tasks_to_trace_next

for root in dag.roots:
get_downstream(root)
Expand Down
21 changes: 21 additions & 0 deletions tests/www/views/test_views_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -970,3 +970,24 @@ def test_task_fail_duration(app, admin_client, dag_maker, session):
assert resp.status_code == 200
assert sorted(item["key"] for item in cumulative_chart) == ["fail", "success"]
assert sorted(item["key"] for item in line_chart) == ["fail", "success"]


def test_graph_view_doesnt_fail_on_recursion_error(app, dag_maker, admin_client):
"""Test that the graph view doesn't fail on a recursion error."""
from airflow.utils.helpers import chain

with dag_maker('test_fails_with_recursion') as dag:

tasks = [
BashOperator(
task_id=f"task_{i}",
bash_command="echo test",
)
for i in range(1, 1000 + 1)
]
chain(*tasks)
with unittest.mock.patch.object(app, 'dag_bag') as mocked_dag_bag:
mocked_dag_bag.get_dag.return_value = dag
url = f'/dags/{dag.dag_id}/graph'
resp = admin_client.get(url, follow_redirects=True)
assert resp.status_code == 200

0 comments on commit 8538a72

Please sign in to comment.