Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/services/ui/grid.py: 21%

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 

18from __future__ import annotations 

19 

20from collections import Counter 

21from collections.abc import Iterable, Mapping 

22from dataclasses import dataclass, field 

23from datetime import datetime 

24from typing import Any 

25 

26import structlog 

27 

28from airflow.api_fastapi.common.parameters import state_priority 

29from airflow.api_fastapi.core_api.services.ui.task_group import get_task_group_children_getter 

30from airflow.models.taskmap import TaskMap 

31from airflow.serialization.definitions.baseoperator import SerializedBaseOperator 

32from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator 

33from airflow.serialization.definitions.taskgroup import SerializedTaskGroup 

34 

35log = structlog.get_logger(logger_name=__name__) 

36 

37 

38@dataclass 

39class GridNodeAgg: 

40 """Compact task instance summary used to aggregate grid state without keeping TI details.""" 

41 

42 child_states: Counter[Any] = field(default_factory=Counter) 

43 min_start_date: datetime | None = None 

44 max_end_date: datetime | None = None 

45 dag_version_number: int | None = None 

46 has_note: bool = False 

47 

48 def add_ti( 

49 self, 

50 *, 

51 state: Any, 

52 start_date: datetime | None, 

53 end_date: datetime | None, 

54 dag_version_number: int | None, 

55 has_note: bool = False, 

56 ) -> None: 

57 """Merge one task instance row into the summary.""" 

58 self.child_states[state] += 1 

59 if start_date is not None and (self.min_start_date is None or start_date < self.min_start_date): 

60 self.min_start_date = start_date 

61 if end_date is not None and (self.max_end_date is None or end_date > self.max_end_date): 

62 self.max_end_date = end_date 

63 if dag_version_number is not None and ( 

64 self.dag_version_number is None or dag_version_number > self.dag_version_number 

65 ): 

66 self.dag_version_number = dag_version_number 

67 self.has_note = self.has_note or has_note 

68 

69 def merge(self, other: GridNodeAgg) -> None: 

70 """Merge another summary into this one.""" 

71 self.child_states.update(other.child_states) 

72 if other.min_start_date is not None and ( 

73 self.min_start_date is None or other.min_start_date < self.min_start_date 

74 ): 

75 self.min_start_date = other.min_start_date 

76 if other.max_end_date is not None and ( 

77 self.max_end_date is None or other.max_end_date > self.max_end_date 

78 ): 

79 self.max_end_date = other.max_end_date 

80 if other.dag_version_number is not None and ( 

81 self.dag_version_number is None or other.dag_version_number > self.dag_version_number 

82 ): 

83 self.dag_version_number = other.dag_version_number 

84 self.has_note = self.has_note or other.has_note 

85 

86 def with_placeholder_state(self) -> GridNodeAgg: 

87 """Represent mapped tasks without rows as a single no-status square in the grid.""" 

88 if self.child_states: 

89 return self 

90 placeholder = GridNodeAgg(dag_version_number=self.dag_version_number) 

91 placeholder.add_ti( 

92 state=None, 

93 start_date=None, 

94 end_date=None, 

95 dag_version_number=self.dag_version_number, 

96 ) 

97 return placeholder 

98 

99 

100def _merge_node_dicts(current: list[dict[str, Any]], new: list[dict[str, Any]] | None) -> None: 

101 """Merge node dictionaries from different Dag versions, handling structure changes.""" 

102 # Handle None case - can occur when merging old Dag versions 

103 # where a TaskGroup was converted to a task or vice versa 

104 if new is None: 

105 return 

106 

107 current_nodes_by_id = {node["id"]: node for node in current} 

108 for node in new: 

109 node_id = node["id"] 

110 current_node = current_nodes_by_id.get(node_id) 

111 if current_node is not None: 

112 # Only merge children if current node already has children 

113 # This preserves the structure of the latest Dag version 

114 if current_node.get("children") is not None: 

115 _merge_node_dicts(current_node["children"], node.get("children")) 

116 else: 

117 current.append(node) 

118 current_nodes_by_id[node_id] = node 

119 

120 

121def agg_state(states): 

122 state_counts = states if isinstance(states, Counter) else Counter(states) 

123 for state in state_priority: 

124 if state in state_counts: 

125 return state 

126 return None 

127 

128 

129def _serialize_child_states(child_states: Counter[Any]) -> dict[str, int]: 

130 return {state if state is not None else "none": count for state, count in child_states.items()} 

131 

132 

133def _get_aggs_for_node(summary: GridNodeAgg) -> dict[str, Any]: 

134 return { 

135 "state": agg_state(summary.child_states), 

136 "min_start_date": summary.min_start_date, 

137 "max_end_date": summary.max_end_date, 

138 "child_states": _serialize_child_states(summary.child_states), 

139 "dag_version_number": summary.dag_version_number, 

140 "has_note": summary.has_note, 

141 } 

142 

143 

144def _find_aggregates( 

145 node: SerializedTaskGroup | SerializedBaseOperator | TaskMap, 

146 parent_node: SerializedTaskGroup | SerializedBaseOperator | TaskMap | None, 

147 ti_details: Mapping[str, GridNodeAgg], 

148 group_dict: dict[str | None, SerializedTaskGroup] | None = None, 

149) -> Iterable[tuple[dict[str, Any], GridNodeAgg]]: 

150 """Recursively fill the Task Group Map.""" 

151 node_id = node.node_id 

152 parent_id = parent_node.node_id if parent_node else None 

153 # Do not mutate ti_details by accidental key creation 

154 summary = ti_details.get(node_id) 

155 if summary is None: 

156 summary = GridNodeAgg() 

157 

158 if node is None: 

159 return 

160 if isinstance(node, SerializedMappedOperator): 

161 mapped_summary = summary.with_placeholder_state() 

162 yield ( 

163 { 

164 "task_id": node_id, 

165 "task_display_name": node.task_display_name, 

166 "type": "mapped_task", 

167 "parent_id": parent_id, 

168 **_get_aggs_for_node(mapped_summary), 

169 }, 

170 mapped_summary, 

171 ) 

172 

173 return 

174 if isinstance(node, SerializedTaskGroup): 

175 if group_dict is None: 

176 group_dict = node.dag.task_group.get_task_group_dict() 

177 children_summary = GridNodeAgg() 

178 for child in get_task_group_children_getter()(node, group_dict): 

179 for child_node, child_summary in _find_aggregates( 

180 node=child, parent_node=node, ti_details=ti_details, group_dict=group_dict 

181 ): 

182 if child_node["parent_id"] == node_id: 

183 children_summary.merge(child_summary) 

184 yield child_node, child_summary 

185 if node_id: 

186 yield ( 

187 { 

188 "task_id": node_id, 

189 "task_display_name": node_id, 

190 "type": "group", 

191 "parent_id": parent_id, 

192 **_get_aggs_for_node(children_summary), 

193 }, 

194 children_summary, 

195 ) 

196 return 

197 if isinstance(node, SerializedBaseOperator): 

198 yield ( 

199 { 

200 "task_id": node_id, 

201 "task_display_name": node.task_display_name, 

202 "type": "task", 

203 "parent_id": parent_id, 

204 **_get_aggs_for_node(summary), 

205 }, 

206 summary, 

207 ) 

208 return