Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/services/ui/structure.py: 9%
77 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.
18"""
19Private service for dag structure.
21:meta private:
22"""
24from __future__ import annotations
26from collections import defaultdict
28from sqlalchemy import select
29from sqlalchemy.orm import Session
31from airflow.models.asset import AssetAliasModel, AssetEvent
32from airflow.models.dag_version import DagVersion
33from airflow.models.dagrun import DagRun
34from airflow.models.serialized_dag import SerializedDagModel
37def get_upstream_assets(
38 asset_expression: dict, entry_node_ref: str, level: int = 0
39) -> tuple[list[dict], list[dict]]:
40 edges: list[dict] = []
41 nodes: list[dict] = []
42 asset_expression_type: str | None = None
44 # include assets, asset-alias, asset-name-refs, asset-uri-refs
45 assets_info: list[dict] = []
47 nested_expression: dict = {}
49 expr_key = ""
50 if asset_expression.keys() == {"any"}:
51 asset_expression_type = "or-gate"
52 expr_key = "any"
53 elif asset_expression.keys() == {"all"}:
54 asset_expression_type = "and-gate"
55 expr_key = "all"
57 if expr_key in asset_expression:
58 asset_exprs: list[dict] = asset_expression[expr_key]
59 for expr in asset_exprs:
60 nested_expr_key = next(iter(expr.keys()))
61 if nested_expr_key in ("any", "all"):
62 nested_expression = expr
63 elif nested_expr_key in ("asset", "alias", "asset-name-ref", "asset-uri-ref"):
64 asset_info = expr[nested_expr_key]
65 asset_info["type"] = nested_expr_key if nested_expr_key != "alias" else "asset-alias"
67 assets_info.append(asset_info)
68 else:
69 raise TypeError(f"Unsupported type: {expr.keys()}")
71 if asset_expression_type and assets_info:
72 asset_condition_id = f"{asset_expression_type}-{level}"
73 edges.append(
74 {
75 "source_id": asset_condition_id,
76 "target_id": entry_node_ref,
77 "is_source_asset": level == 0,
78 }
79 )
80 nodes.append(
81 {
82 "id": asset_condition_id,
83 "label": asset_condition_id,
84 "type": "asset-condition",
85 "asset_condition_type": asset_expression_type,
86 }
87 )
89 for asset in assets_info:
90 asset_type = asset["type"]
92 if asset_type == "asset":
93 source_id = str(asset["id"])
94 label = asset["name"]
95 elif asset_type == "asset-alias" or asset_type == "asset-name-ref":
96 source_id = asset["name"]
97 label = asset["name"]
98 elif asset_type == "asset-uri-ref":
99 source_id = asset["uri"]
100 label = asset["uri"]
101 else:
102 raise TypeError(f"Unsupported type: {asset_type}")
104 edges.append(
105 {
106 "source_id": source_id,
107 "target_id": asset_condition_id,
108 }
109 )
110 nodes.append(
111 {
112 "id": source_id,
113 "label": label,
114 "type": asset_type,
115 }
116 )
118 if nested_expression is not None:
119 n, e = get_upstream_assets(nested_expression, asset_condition_id, level=level + 1)
121 nodes = nodes + n
122 edges = edges + e
124 return nodes, edges
127def bind_output_assets_to_tasks(
128 edges: list[dict], serialized_dag: SerializedDagModel, version_number: int, session: Session
129) -> None:
130 """
131 Try to bind the downstream assets to the relevant task that produces them.
133 This function will mutate the `edges` in place.
134 """
135 # bind normal assets present in the `task_outlet_asset_references`
136 outlet_asset_references = serialized_dag.dag_model.task_outlet_asset_references
138 downstream_asset_edges = [
139 edge
140 for edge in edges
141 if edge["target_id"].startswith("asset:") and not edge.get("resolved_from_alias")
142 ]
144 for edge in downstream_asset_edges:
145 # Try to attach the outlet assets to the relevant tasks
146 asset_id = int(edge["target_id"].replace("asset:", "", 1))
147 outlet_asset_reference = next(
148 outlet_asset_reference
149 for outlet_asset_reference in outlet_asset_references
150 if outlet_asset_reference.asset_id == asset_id
151 )
152 edge["source_id"] = outlet_asset_reference.task_id
154 # bind assets resolved from aliases, they do not populate the `outlet_asset_references`
155 downstream_alias_resolved_edges = [
156 edge for edge in edges if edge["target_id"].startswith("asset:") and edge.get("resolved_from_alias")
157 ]
159 aliases_names = {edges["resolved_from_alias"] for edges in downstream_alias_resolved_edges}
161 result = session.scalars(
162 select(AssetEvent)
163 .join(AssetEvent.source_aliases)
164 .join(AssetEvent.source_dag_run)
165 # That's a simplification, instead doing `version_number` in `DagRun.dag_versions`.
166 .join(DagRun.created_dag_version)
167 .where(AssetEvent.source_aliases.any(AssetAliasModel.name.in_(aliases_names)))
168 .where(AssetEvent.source_dag_run.has(DagRun.dag_id == serialized_dag.dag_model.dag_id))
169 .where(DagVersion.version_number == version_number)
170 ).unique()
172 asset_id_to_task_ids = defaultdict(set)
173 for asset_event in result:
174 asset_id_to_task_ids[asset_event.asset_id].add(asset_event.source_task_id)
176 for edge in downstream_alias_resolved_edges:
177 asset_id = int(edge["target_id"].replace("asset:", "", 1))
178 task_ids = asset_id_to_task_ids.get(asset_id, set())
180 for index, task_id in enumerate(task_ids):
181 if index == 0:
182 edge["source_id"] = task_id
183 continue
184 edge_copy = {**edge, "source_id": task_id}
185 edges.append(edge_copy)