Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/services/ui/dependencies.py: 5%
143 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 14:22 +0000
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.
18from __future__ import annotations
20from collections import defaultdict, deque
21from typing import TYPE_CHECKING
23from airflow.models.asset import AssetModel
24from airflow.models.dag import DagModel
26if TYPE_CHECKING: 26 ↛ 27line 26 didn't jump to line 27 because the condition on line 26 was never true
27 from sqlalchemy.orm import Session
30def _dfs_connected_components(
31 temp: list[str], node_id: str, visited: dict[str, bool], adjacency_matrix: dict[str, list[str]]
32) -> list[str]:
33 visited[node_id] = True
35 temp.append(node_id)
37 for adj_node_id in adjacency_matrix[node_id]:
38 if not visited[adj_node_id]:
39 temp = _dfs_connected_components(temp, adj_node_id, visited, adjacency_matrix)
41 return temp
44def extract_connected_components(adjacency_matrix: dict[str, list[str]]) -> list[list[str]]:
45 """Extract all connected components of a graph."""
46 visited: dict[str, bool] = {node_id: False for node_id in adjacency_matrix}
48 connected_components: list[list[str]] = []
50 for node_id in adjacency_matrix:
51 if visited[node_id] is False:
52 temp: list[str] = []
53 connected_components.append(_dfs_connected_components(temp, node_id, visited, adjacency_matrix))
54 return connected_components
57def extract_single_connected_component(
58 node_id: str, nodes: list[dict], edges: list[dict]
59) -> dict[str, list[dict]]:
60 """Find the connected component that contains the node with the id ``node_id``."""
61 adjacency_matrix: dict[str, list[str]] = defaultdict(list)
63 for edge in edges:
64 adjacency_matrix[edge["source_id"]].append(edge["target_id"])
65 adjacency_matrix[edge["target_id"]].append(edge["source_id"])
67 connected_components = extract_connected_components(adjacency_matrix)
69 filtered_connected_components = [cc for cc in connected_components if node_id in cc]
71 if len(filtered_connected_components) != 1:
72 raise ValueError(
73 f"Unique connected component not found, got {filtered_connected_components} for connected components of node {node_id}, expected only 1 connected component."
74 )
76 connected_component = filtered_connected_components[0]
78 nodes = [node for node in nodes if node["id"] in connected_component]
79 edges = [
80 edge
81 for edge in edges
82 if (edge["source_id"] in connected_component and edge["target_id"] in connected_component)
83 ]
85 return {"nodes": nodes, "edges": edges}
88def get_scheduling_dependencies(readable_dag_ids: set[str] | None = None) -> dict[str, list[dict]]:
89 """Get scheduling dependencies between Dags."""
90 from airflow.models.serialized_dag import SerializedDagModel
92 nodes_dict: dict[str, dict] = {}
93 edge_tuples: set[tuple[str, str]] = set()
95 dag_dependencies = SerializedDagModel.get_dag_dependencies()
96 for dag, dependencies in sorted(dag_dependencies.items()):
97 if readable_dag_ids is not None and dag not in readable_dag_ids:
98 continue
99 dag_node_id = f"dag:{dag}"
100 if dag_node_id not in nodes_dict:
101 for dep in dependencies:
102 # Skip dependency objects whose edge endpoints reference DAGs
103 # outside the caller's readable set. ``dep.node_id`` /
104 # ``dep.source`` / ``dep.target`` would otherwise embed those
105 # DAG ids in the response even when the top-level filter
106 # above hides the DAG itself.
107 if readable_dag_ids is not None:
108 referenced_dag_ids: set[str] = set()
109 if dep.source != dep.dependency_type and ":" not in dep.source:
110 referenced_dag_ids.add(dep.source)
111 if dep.target != dep.dependency_type and ":" not in dep.target:
112 referenced_dag_ids.add(dep.target)
113 if not referenced_dag_ids.issubset(readable_dag_ids):
114 continue
116 # Add nodes
117 nodes_dict[dag_node_id] = {"id": dag_node_id, "label": dag, "type": "dag"}
118 if dep.node_id not in nodes_dict:
119 nodes_dict[dep.node_id] = {
120 "id": dep.node_id,
121 "label": dep.label,
122 "type": dep.dependency_type,
123 }
125 # Add edges
126 # not start dep
127 if dep.source != dep.dependency_type:
128 source = dep.source if ":" in dep.source else f"dag:{dep.source}"
129 target = dep.node_id
130 edge_tuples.add((source, target))
132 # not end dep
133 if dep.target != dep.dependency_type:
134 source = dep.node_id
135 target = dep.target if ":" in dep.target else f"dag:{dep.target}"
136 edge_tuples.add((source, target))
138 # Create missing ``dag:`` nodes which may have been skipped by the loop above.
139 # A DAG referenced only as a trigger target or a sensor source may have no
140 # scheduling dependencies of its own. Without this loop, these DAGs will not be
141 # materialised and will result in dangling edges.
142 for source, target in edge_tuples:
143 for endpoint in (source, target):
144 if endpoint.startswith("dag:") and endpoint not in nodes_dict:
145 nodes_dict[endpoint] = {
146 "id": endpoint,
147 "label": endpoint.removeprefix("dag:"),
148 "type": "dag",
149 }
151 dag_ids = [node["label"] for node in nodes_dict.values() if node["type"] == "dag"]
152 if dag_ids:
153 dag_id_to_team = DagModel.get_dag_id_to_team_name_mapping(dag_ids)
154 for node in nodes_dict.values():
155 if node["type"] == "dag":
156 team_name = dag_id_to_team.get(node["label"])
157 if team_name:
158 node["team"] = team_name
160 return {
161 "nodes": list(nodes_dict.values()),
162 "edges": [{"source_id": source, "target_id": target} for source, target in sorted(edge_tuples)],
163 }
166def get_data_dependencies(
167 asset_id: int, session: Session, readable_dag_ids: set[str] | None = None
168) -> dict[str, list[dict]]:
169 """Get full task dependencies for an asset."""
170 from sqlalchemy import select, union_all
171 from sqlalchemy.orm import selectinload
173 from airflow.models.asset import (
174 DagScheduleAssetReference,
175 TaskInletAssetReference,
176 TaskOutletAssetReference,
177 )
179 SEPARATOR = "__SEPARATOR__"
181 # Hide the asset entirely if the user has no read access to any dag that produces,
182 # consumes, or is scheduled by it. Without this check, visiting the asset graph page
183 # for an unrelated asset would leak its existence and name (and of connected nodes
184 # reachable through other readable dags) even though the user has no legitimate
185 # lineage connection to it. A readable_dag_ids value of None means no filter is
186 # applied (the user has unrestricted dag read access).
187 if readable_dag_ids is not None:
188 connected_dag_ids_query = union_all(
189 select(TaskOutletAssetReference.dag_id).where(TaskOutletAssetReference.asset_id == asset_id),
190 select(TaskInletAssetReference.dag_id).where(TaskInletAssetReference.asset_id == asset_id),
191 select(DagScheduleAssetReference.dag_id).where(DagScheduleAssetReference.asset_id == asset_id),
192 )
193 connected_dag_ids = set(session.scalars(select(connected_dag_ids_query.subquery().c.dag_id)))
194 if not connected_dag_ids & readable_dag_ids:
195 return {"nodes": [], "edges": []}
197 nodes_dict: dict[str, dict] = {}
198 edge_set: set[tuple[str, str]] = set()
200 # BFS to trace full dependencies
201 assets_to_process: deque[int] = deque([asset_id])
202 processed_assets: set[int] = set()
203 processed_tasks: set[tuple[str, str]] = set() # (dag_id, task_id)
205 while assets_to_process:
206 current_asset_id = assets_to_process.popleft()
207 if current_asset_id in processed_assets:
208 continue
209 processed_assets.add(current_asset_id)
211 # Eagerload producing_tasks and consuming_tasks to avoid lazy queries
212 asset = session.scalar(
213 select(AssetModel)
214 .where(AssetModel.id == current_asset_id)
215 .options(
216 selectinload(AssetModel.producing_tasks),
217 selectinload(AssetModel.consuming_tasks),
218 )
219 )
220 if not asset:
221 continue
223 asset_node_id = f"asset:{current_asset_id}"
225 # Add asset node
226 if asset_node_id not in nodes_dict:
227 nodes_dict[asset_node_id] = {"id": asset_node_id, "label": asset.name, "type": "asset"}
229 # Process producing tasks (tasks that output this asset)
230 for ref in asset.producing_tasks:
231 # Filter out tasks from Dags the user doesn't have access to
232 if readable_dag_ids is not None and ref.dag_id not in readable_dag_ids:
233 continue
234 task_key = (ref.dag_id, ref.task_id)
235 task_node_id = f"task:{ref.dag_id}{SEPARATOR}{ref.task_id}"
237 # Add task node with dag_id.task_id label for disambiguation
238 if task_node_id not in nodes_dict:
239 nodes_dict[task_node_id] = {
240 "id": task_node_id,
241 "label": f"{ref.dag_id}.{ref.task_id}",
242 "type": "task",
243 }
245 # Add edge: task → asset
246 edge_set.add((task_node_id, asset_node_id))
248 # Find other assets this task consumes (inlets) to trace upstream
249 if task_key not in processed_tasks:
250 processed_tasks.add(task_key)
251 inlet_refs = session.scalars(
252 select(TaskInletAssetReference).where(
253 TaskInletAssetReference.dag_id == ref.dag_id,
254 TaskInletAssetReference.task_id == ref.task_id,
255 )
256 ).all()
257 for inlet_ref in inlet_refs:
258 if inlet_ref.asset_id not in processed_assets:
259 assets_to_process.append(inlet_ref.asset_id)
261 # Process consuming tasks (tasks that input this asset)
262 for ref in asset.consuming_tasks:
263 # Filter out tasks from Dags the user doesn't have access to
264 if readable_dag_ids is not None and ref.dag_id not in readable_dag_ids:
265 continue
266 task_key = (ref.dag_id, ref.task_id)
267 task_node_id = f"task:{ref.dag_id}{SEPARATOR}{ref.task_id}"
269 # Add task node with dag_id.task_id label for disambiguation
270 if task_node_id not in nodes_dict:
271 nodes_dict[task_node_id] = {
272 "id": task_node_id,
273 "label": f"{ref.dag_id}.{ref.task_id}",
274 "type": "task",
275 }
277 # Add edge: asset → task
278 edge_set.add((asset_node_id, task_node_id))
280 # Find other assets this task produces (outlets) to trace downstream
281 if task_key not in processed_tasks:
282 processed_tasks.add(task_key)
283 outlet_refs = session.scalars(
284 select(TaskOutletAssetReference).where(
285 TaskOutletAssetReference.dag_id == ref.dag_id,
286 TaskOutletAssetReference.task_id == ref.task_id,
287 )
288 ).all()
289 for outlet_ref in outlet_refs:
290 if outlet_ref.asset_id not in processed_assets:
291 assets_to_process.append(outlet_ref.asset_id)
293 all_dag_ids = list({dag_id for dag_id, _ in processed_tasks})
294 if all_dag_ids:
295 dag_id_to_team = DagModel.get_dag_id_to_team_name_mapping(all_dag_ids, session=session)
296 for node in nodes_dict.values():
297 if not node["id"].startswith("task:"):
298 continue
299 dag_id = node["id"].removeprefix("task:").split(SEPARATOR, 1)[0]
300 team_name = dag_id_to_team.get(dag_id)
301 if team_name:
302 node["team"] = team_name
304 return {
305 "nodes": list(nodes_dict.values()),
306 "edges": [{"source_id": source, "target_id": target} for source, target in edge_set],
307 }