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

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""" 

18Database helpers for Airflow REST API. 

19 

20:meta private: 

21""" 

22 

23from __future__ import annotations 

24 

25from collections.abc import AsyncGenerator, Generator, Sequence 

26from typing import TYPE_CHECKING, Annotated, Literal, overload 

27 

28from fastapi import Depends 

29from sqlalchemy import func, select 

30from sqlalchemy.ext.asyncio import AsyncSession 

31from sqlalchemy.orm import Session 

32 

33from airflow.utils.db import get_query_count, get_query_count_async 

34from airflow.utils.session import NEW_SESSION, create_session, create_session_async, provide_session 

35 

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

37 from sqlalchemy.sql import Select 

38 

39 from airflow.api_fastapi.core_api.base import OrmClause 

40 

41# Rows a single scan reads. Result sets that fit are counted exactly; wider ones report a floor. 

42# Shared by the dashboard's historical metrics and by cursor-paginated listings, which surface an 

43# item count without counting every matching row — the point of cursor pagination on large tables. 

44EXACT_COUNT_LIMIT = 50_000 

45 

46 

47def _get_session() -> Generator[Session, None, None]: 

48 with create_session(scoped=False) as session: 

49 yield session 

50 

51 

52SessionDep = Annotated[Session, Depends(_get_session, scope="function")] 

53 

54 

55def apply_filters_to_select( 

56 *, statement: Select, filters: Sequence[OrmClause | None] | None = None 

57) -> Select: 

58 if filters is None: 

59 return statement 

60 for f in filters: 

61 if f is None: 

62 continue 

63 statement = f.to_orm(statement) 

64 

65 return statement 

66 

67 

68def bounded_total_entries( 

69 *, 

70 statement: Select, 

71 filters: Sequence[OrmClause | None] | None = None, 

72 session: Session, 

73) -> tuple[int, int]: 

74 """ 

75 Count the rows a cursor-paginated listing matches, reading at most ``EXACT_COUNT_LIMIT``. 

76 

77 Returns ``(total, limit)`` where ``total`` is ``min(actual_count, EXACT_COUNT_LIMIT)`` — a 

78 ``total`` equal to ``limit`` means only that at least that many rows match — and ``limit`` is 

79 the cap that was applied, for the caller to surface as ``total_entries_limit``. 

80 

81 The ``LIMIT`` sits inside the counted subquery so the database stops scanning once the cap is 

82 reached, keeping the count cheap on tables that cursor pagination exists to handle. ORDER BY is 

83 stripped for the same reason :func:`~airflow.utils.db.get_query_count` strips it: it cannot 

84 change a count and only constrains the planner. 

85 """ 

86 statement = apply_filters_to_select(statement=statement, filters=filters) 

87 bounded = statement.order_by(None).limit(EXACT_COUNT_LIMIT).subquery() 

88 total = session.scalar(select(func.count()).select_from(bounded)) or 0 

89 return total, EXACT_COUNT_LIMIT 

90 

91 

92async def _get_async_session() -> AsyncGenerator[AsyncSession, None]: 

93 async with create_session_async() as session: 

94 yield session 

95 

96 

97AsyncSessionDep = Annotated[AsyncSession, Depends(_get_async_session)] 

98 

99 

100@overload 

101async def paginated_select_async( 101 ↛ exitline 101 didn't return from function 'paginated_select_async' because

102 *, 

103 statement: Select, 

104 filters: Sequence[OrmClause] | None = None, 

105 order_by: OrmClause | None = None, 

106 offset: OrmClause | None = None, 

107 limit: OrmClause | None = None, 

108 session: AsyncSession, 

109 return_total_entries: Literal[True] = True, 

110) -> tuple[Select, int]: ... 

111 

112 

113@overload 

114async def paginated_select_async( 114 ↛ exitline 114 didn't return from function 'paginated_select_async' because

115 *, 

116 statement: Select, 

117 filters: Sequence[OrmClause] | None = None, 

118 order_by: OrmClause | None = None, 

119 offset: OrmClause | None = None, 

120 limit: OrmClause | None = None, 

121 session: AsyncSession, 

122 return_total_entries: Literal[False], 

123) -> tuple[Select, None]: ... 

124 

125 

126async def paginated_select_async( 

127 *, 

128 statement: Select, 

129 filters: Sequence[OrmClause | None] | None = None, 

130 order_by: OrmClause | None = None, 

131 offset: OrmClause | None = None, 

132 limit: OrmClause | None = None, 

133 session: AsyncSession, 

134 return_total_entries: bool = True, 

135) -> tuple[Select, int | None]: 

136 statement = apply_filters_to_select( 

137 statement=statement, 

138 filters=filters, 

139 ) 

140 

141 total_entries = None 

142 if return_total_entries: 

143 total_entries = await get_query_count_async(statement, session=session) 

144 

145 statement = apply_filters_to_select( 

146 statement=statement, 

147 filters=[order_by, offset, limit], 

148 ) 

149 

150 return statement, total_entries 

151 

152 

153@overload 

154def paginated_select( 154 ↛ exitline 154 didn't return from function 'paginated_select' because

155 *, 

156 statement: Select, 

157 filters: Sequence[OrmClause] | None = None, 

158 order_by: OrmClause | None = None, 

159 offset: OrmClause | None = None, 

160 limit: OrmClause | None = None, 

161 session: Session = NEW_SESSION, 

162 return_total_entries: Literal[True] = True, 

163) -> tuple[Select, int]: ... 

164 

165 

166@overload 

167def paginated_select( 167 ↛ exitline 167 didn't return from function 'paginated_select' because

168 *, 

169 statement: Select, 

170 filters: Sequence[OrmClause] | None = None, 

171 order_by: OrmClause | None = None, 

172 offset: OrmClause | None = None, 

173 limit: OrmClause | None = None, 

174 session: Session = NEW_SESSION, 

175 return_total_entries: Literal[False], 

176) -> tuple[Select, None]: ... 

177 

178 

179@provide_session 

180def paginated_select( 

181 *, 

182 statement: Select, 

183 filters: Sequence[OrmClause] | None = None, 

184 order_by: OrmClause | None = None, 

185 offset: OrmClause | None = None, 

186 limit: OrmClause | None = None, 

187 session: Session = NEW_SESSION, 

188 return_total_entries: bool = True, 

189) -> tuple[Select, int | None]: 

190 statement = apply_filters_to_select( 

191 statement=statement, 

192 filters=filters, 

193 ) 

194 

195 total_entries = None 

196 if return_total_entries: 

197 total_entries = get_query_count(statement, session=session) 

198 

199 statement = apply_filters_to_select(statement=statement, filters=[order_by, offset, limit]) 

200 

201 return statement, total_entries