Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/routes/ui/dashboard.py: 44%

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. 

17from __future__ import annotations 

18 

19from decimal import ROUND_FLOOR, Context 

20from typing import TYPE_CHECKING, cast 

21 

22from fastapi import Depends, status 

23from sqlalchemy import func, select 

24from sqlalchemy.sql.expression import case, false 

25 

26from airflow._shared.timezones import timezone 

27from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity 

28from airflow.api_fastapi.common.db.common import EXACT_COUNT_LIMIT, SessionDep 

29from airflow.api_fastapi.common.parameters import DateTimeQuery, OptionalDateTimeQuery 

30from airflow.api_fastapi.common.router import AirflowRouter 

31from airflow.api_fastapi.core_api.datamodels.ui.dashboard import ( 

32 DashboardDagStatsResponse, 

33 HistoricalMetricDataResponse, 

34) 

35from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc 

36from airflow.api_fastapi.core_api.security import ReadableDagsFilterDep, requires_access_dag 

37from airflow.models.dag import DagModel 

38from airflow.models.dagrun import DagRun 

39from airflow.models.taskinstance import TaskInstance 

40from airflow.utils.state import DagRunState, TaskInstanceState 

41 

42if TYPE_CHECKING: 42 ↛ 43line 42 didn't jump to line 43 because the condition on line 42 was never true

43 from sqlalchemy.orm import Session 

44 

45dashboard_router = AirflowRouter(tags=["Dashboard"], prefix="/dashboard") 

46 

47 

48_ROUNDING = Context(prec=2, rounding=ROUND_FLOOR) 

49 

50 

51def _round_down(value: int) -> int: 

52 """Round to two significant digits, never upwards: that would claim uncounted rows.""" 

53 return int(_ROUNDING.create_decimal(value)) 

54 

55 

56def _compute_state_counts( 

57 model, filters, *, session: Session, join=None, null_label: str | None = None 

58) -> tuple[dict[str, int], bool]: 

59 """ 

60 Per-state counts for the window, and whether they are lower bounds rather than exact. 

61 

62 A scan that stopped early counted only some of the rows, but each count is still a floor 

63 on the real value. 

64 """ 

65 stmt = select(model.state.label("state")).select_from(model) 

66 if join is not None: 

67 stmt = stmt.join(join) 

68 window = stmt.where(*filters).limit(EXACT_COUNT_LIMIT + 1).subquery() 

69 rows = session.execute(select(window.c.state, func.count().label("cnt")).group_by(window.c.state)).all() 

70 are_lower_bounds = sum(row.cnt for row in rows) > EXACT_COUNT_LIMIT 

71 counts: dict[str, int] = {} 

72 for row in rows: 

73 label = row.state or null_label 

74 if label is None: 

75 continue 

76 counts[label] = _round_down(row.cnt) if are_lower_bounds else row.cnt 

77 return counts, are_lower_bounds 

78 

79 

80@dashboard_router.get( 

81 "/historical_metrics_data", 

82 responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST]), 

83 dependencies=[ 

84 Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE)), 

85 Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN)), 

86 ], 

87) 

88def historical_metrics( 

89 session: SessionDep, 

90 start_date: DateTimeQuery, 

91 readable_dags_filter: ReadableDagsFilterDep, 

92 end_date: OptionalDateTimeQuery = None, 

93) -> HistoricalMetricDataResponse: 

94 """Return cluster activity historical metrics.""" 

95 current_time = timezone.utcnow() 

96 permitted_dag_ids = cast("set[str]", readable_dags_filter.value) 

97 

98 end_bound = end_date if end_date is not None else current_time 

99 dag_run_filters = [ 

100 DagRun.run_after >= start_date, 

101 DagRun.run_after <= end_bound, 

102 DagRun.dag_id.in_(permitted_dag_ids), 

103 ] 

104 

105 # Judged separately: dag runs often fit when task instances do not. 

106 dag_run_states, dag_runs_are_lower_bounds = _compute_state_counts( 

107 DagRun, dag_run_filters, session=session 

108 ) 

109 task_instance_states, task_instances_are_lower_bounds = _compute_state_counts( 

110 TaskInstance, 

111 dag_run_filters, 

112 session=session, 

113 join=TaskInstance.dag_run, 

114 null_label="no_status", 

115 ) 

116 

117 return HistoricalMetricDataResponse.model_validate( 

118 { 

119 "dag_run_states": { 

120 **{dag_run_state.value: 0 for dag_run_state in DagRunState}, 

121 **dag_run_states, 

122 }, 

123 "task_instance_states": { 

124 "no_status": 0, 

125 **{ti_state.value: 0 for ti_state in TaskInstanceState}, 

126 **task_instance_states, 

127 }, 

128 "dag_run_counts_are_lower_bounds": dag_runs_are_lower_bounds, 

129 "task_instance_counts_are_lower_bounds": task_instances_are_lower_bounds, 

130 } 

131 ) 

132 

133 

134@dashboard_router.get( 

135 "/dag_stats", 

136 dependencies=[Depends(requires_access_dag(method="GET"))], 

137) 

138def dag_stats( 

139 session: SessionDep, 

140 readable_dags_filter: ReadableDagsFilterDep, 

141) -> DashboardDagStatsResponse: 

142 """Return basic Dag stats with counts of Dags in various states.""" 

143 permitted_dag_ids = cast("set[str]", readable_dags_filter.value) 

144 

145 latest_state = ( 

146 select(DagRun.state) 

147 .where(DagRun.dag_id == DagModel.dag_id, DagRun.logical_date.is_not(None)) 

148 .order_by(DagRun.logical_date.desc()) 

149 .limit(1) 

150 .correlate(DagModel) 

151 .scalar_subquery() 

152 ) 

153 dag_counts_query = ( 

154 select( 

155 func.coalesce(func.sum(case((DagModel.is_paused == false(), 1))), 0).label("active"), 

156 func.coalesce(func.sum(case((latest_state == DagRunState.FAILED, 1))), 0).label("failed"), 

157 ) 

158 .select_from(DagModel) 

159 .where(DagModel.is_stale == false()) 

160 .where(DagModel.dag_id.in_(permitted_dag_ids)) 

161 ) 

162 dag_counts = session.execute(dag_counts_query).one() 

163 

164 active_states_query = ( 

165 select(DagRun.state, func.count(func.distinct(DagRun.dag_id))) 

166 .join(DagModel, DagModel.dag_id == DagRun.dag_id) 

167 .where(DagModel.is_stale == false()) 

168 .where(DagRun.dag_id.in_(permitted_dag_ids)) 

169 .where(DagRun.state.in_([DagRunState.RUNNING, DagRunState.QUEUED])) 

170 .group_by(DagRun.state) 

171 ) 

172 active_counts = {state: count for state, count in session.execute(active_states_query)} 

173 

174 return DashboardDagStatsResponse( 

175 active_dag_count=dag_counts.active, 

176 failed_dag_count=dag_counts.failed, 

177 running_dag_count=active_counts.get(DagRunState.RUNNING, 0), 

178 queued_dag_count=active_counts.get(DagRunState.QUEUED, 0), 

179 )