Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/common/db/dag_runs.py: 89%

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 defaultdict 

21from collections.abc import Sequence 

22from typing import TYPE_CHECKING 

23from uuid import UUID 

24 

25from sqlalchemy import func, select, tuple_, union_all 

26from sqlalchemy.orm import joinedload 

27from sqlalchemy.orm.interfaces import LoaderOption 

28 

29from airflow.models.dag import DagModel 

30from airflow.models.dag_version import DagVersion 

31from airflow.models.dagrun import DagRun 

32from airflow.models.taskinstance import TaskInstance 

33from airflow.models.taskinstancehistory import TaskInstanceHistory 

34 

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

36 from sqlalchemy.orm import Session 

37 

38# Use the hybrid_property for dag_display_name so the CASE WHEN fallback to dag_id 

39# is applied when dag_display_name is NULL. Using __table__.c would return raw NULLs 

40# and crash DagStatsResponse validation (https://github.com/apache/airflow/issues/64247). 

41dagruns_select_with_state_count = ( 

42 select( # type: ignore[call-overload] 

43 DagRun.__table__.c.dag_id, 

44 DagRun.__table__.c.state, 

45 DagModel.dag_display_name, 

46 func.count(DagRun.__table__.c.state).label("count"), 

47 ) 

48 .join(DagModel, DagRun.__table__.c.dag_id == DagModel.__table__.c.dag_id) 

49 .group_by(DagRun.__table__.c.dag_id, DagRun.__table__.c.state, DagModel.dag_display_name) 

50 .order_by(DagRun.__table__.c.dag_id) 

51) 

52 

53 

54def eager_load_dag_run_for_list() -> tuple[LoaderOption, ...]: 

55 """ 

56 Lightweight eager loading for the DagRun list endpoint. 

57 

58 Only loads the direct relationships needed for serialization (dag_model, 

59 dag_run_note, created_dag_version). The dag_versions property — which 

60 requires iterating every TI and TIH — is populated separately by 

61 :func:`attach_dag_versions_to_runs` using a single DISTINCT query. 

62 """ 

63 return ( 

64 joinedload(DagRun.dag_model), 

65 joinedload(DagRun.dag_run_note), 

66 joinedload(DagRun.created_dag_version).joinedload(DagVersion.bundle), 

67 ) 

68 

69 

70def attach_dag_versions_to_runs(dag_runs: Sequence[DagRun], *, session: Session) -> None: 

71 """ 

72 Prefetch distinct dag_version_ids for each DagRun via a lightweight query. 

73 

74 Instead of loading all TI and TIH rows (potentially thousands per run) 

75 through the ORM relationship just to extract distinct dag_version_ids, 

76 this issues a single query that returns only the distinct 

77 (dag_id, run_id, dag_version_id) tuples for the given runs. 

78 

79 The result is attached to each DagRun as ``_prefetched_dag_version_ids`` 

80 (a dict mapping version_id -> DagVersion), which the ``dag_versions`` 

81 property reads as an optimized substitute for traversing TI/TIH 

82 relationships. All business logic (bundle_version shortcut, sorting, 

83 deduplication) remains solely in ``DagRun.dag_versions``. 

84 """ 

85 if not dag_runs: 

86 return 

87 

88 # Only runs without a bundle_version need TI/TIH traversal; 

89 # runs with bundle_version use created_dag_version directly 

90 # (handled by the dag_versions property). 

91 runs_needing_versions = [dr for dr in dag_runs if not dr.bundle_version] 

92 if not runs_needing_versions: 92 ↛ 93line 92 didn't jump to line 93 because the condition on line 92 was never true

93 return 

94 

95 run_key_values = [(dr.dag_id, dr.run_id) for dr in runs_needing_versions] 

96 

97 ti_sub = ( 

98 select( 

99 TaskInstance.dag_id, 

100 TaskInstance.run_id, 

101 TaskInstance.dag_version_id, 

102 ) 

103 .where(TaskInstance.dag_version_id.isnot(None)) 

104 .where(tuple_(TaskInstance.dag_id, TaskInstance.run_id).in_(run_key_values)) 

105 .distinct() 

106 ) 

107 tih_sub = ( 

108 select( 

109 TaskInstanceHistory.dag_id, 

110 TaskInstanceHistory.run_id, 

111 TaskInstanceHistory.dag_version_id, 

112 ) 

113 .where(TaskInstanceHistory.dag_version_id.isnot(None)) 

114 .where(tuple_(TaskInstanceHistory.dag_id, TaskInstanceHistory.run_id).in_(run_key_values)) 

115 .distinct() 

116 ) 

117 combined = union_all(ti_sub, tih_sub).subquery() 

118 rows = session.execute( 

119 select(combined.c.dag_id, combined.c.run_id, combined.c.dag_version_id).distinct() 

120 ).all() 

121 

122 all_version_ids = {r.dag_version_id for r in rows} 

123 versions_by_id: dict[UUID, DagVersion] = {} 

124 if all_version_ids: 124 ↛ 132line 124 didn't jump to line 132 because the condition on line 124 was always true

125 dv_query = ( 

126 select(DagVersion) 

127 .where(DagVersion.id.in_(all_version_ids)) 

128 .options(joinedload(DagVersion.bundle)) 

129 ) 

130 versions_by_id = {dv.id: dv for dv in session.scalars(dv_query)} 

131 

132 versions_per_run: dict[tuple[str, str], dict[UUID, DagVersion]] = defaultdict(dict) 

133 for row in rows: 

134 dv = versions_by_id.get(row.dag_version_id) 

135 if dv: 135 ↛ 133line 135 didn't jump to line 133 because the condition on line 135 was always true

136 versions_per_run[(row.dag_id, row.run_id)][dv.id] = dv 

137 

138 for dr in runs_needing_versions: 

139 dr._prefetched_dag_version_ids = versions_per_run.get((dr.dag_id, dr.run_id), {})