Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/datamodels/task_instances.py: 98%

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 collections.abc import Iterable 

20from datetime import datetime 

21from typing import Annotated, Any 

22from uuid import UUID 

23 

24from pydantic import ( 

25 AliasPath, 

26 AwareDatetime, 

27 BeforeValidator, 

28 Discriminator, 

29 Field, 

30 NonNegativeInt, 

31 StringConstraints, 

32 Tag, 

33 field_validator, 

34 model_validator, 

35) 

36 

37from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel 

38from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse 

39from airflow.api_fastapi.core_api.datamodels.job import JobResponse 

40from airflow.api_fastapi.core_api.datamodels.trigger import TriggerResponse 

41from airflow.utils.state import TaskInstanceState 

42 

43 

44class NewTaskResponse(BaseModel): 

45 """Lightweight response for new tasks that don't have TaskInstances yet.""" 

46 

47 task_id: str 

48 task_display_name: str 

49 

50 

51class TaskInstanceResponse(BaseModel): 

52 """TaskInstance serializer for responses.""" 

53 

54 id: UUID 

55 task_id: str 

56 dag_id: str 

57 run_id: str = Field(alias="dag_run_id") 

58 map_index: int 

59 logical_date: datetime | None 

60 run_after: datetime 

61 start_date: datetime | None 

62 end_date: datetime | None 

63 duration: float | None 

64 state: TaskInstanceState | None 

65 try_number: int 

66 max_tries: int 

67 task_display_name: str 

68 dag_display_name: str = Field(validation_alias=AliasPath("dag_run", "dag_model", "dag_display_name")) 

69 hostname: str | None 

70 unixname: str | None 

71 pool: str 

72 pool_slots: int 

73 queue: str | None 

74 priority_weight: int | None 

75 operator: str | None 

76 operator_name: str | None 

77 queued_dttm: datetime | None = Field(alias="queued_when") 

78 scheduled_dttm: datetime | None = Field(alias="scheduled_when") 

79 pid: int | None 

80 executor: str | None 

81 executor_config: Annotated[str, BeforeValidator(str)] 

82 note: str | None 

83 rendered_map_index: str | None 

84 rendered_fields: dict = Field( 

85 validation_alias=AliasPath("rendered_task_instance_fields", "rendered_fields"), 

86 default_factory=dict, 

87 ) 

88 trigger: TriggerResponse | None 

89 queued_by_job: JobResponse | None = Field(alias="triggerer_job") 

90 dag_version: DagVersionResponse | None 

91 

92 

93class TaskInstanceCollectionResponse(BaseModel): 

94 """ 

95 Task instance collection response supporting both offset and cursor pagination. 

96 

97 A single flat model is used instead of a discriminated union 

98 (``Annotated[Offset | Cursor, Field(discriminator=...)]``) because 

99 the OpenAPI ``oneOf`` + ``discriminator`` construct is not handled 

100 correctly by ``@hey-api/openapi-ts`` / ``@7nohe/openapi-react-query-codegen``: 

101 return types degrade to ``unknown`` in JSDoc and can produce 

102 incorrect TypeScript types (see hey-api/openapi-ts#1613, #3270). 

103 """ 

104 

105 task_instances: Iterable[TaskInstanceResponse] 

106 total_entries: int | None = Field( 

107 default=None, 

108 description="Number of matching items. For offset pagination this is the exact total. " 

109 "For cursor pagination it is capped at ``total_entries_limit``; a value equal to that " 

110 "limit means at least that many items match.", 

111 ) 

112 total_entries_limit: int | None = Field( 

113 default=None, 

114 description="Cap applied to ``total_entries`` under cursor pagination. ``null`` for offset " 

115 "pagination, where ``total_entries`` is exact.", 

116 ) 

117 next_cursor: str | None = Field( 

118 default=None, 

119 description="Token pointing to the next page. Populated for cursor pagination, " 

120 "``null`` when using offset pagination or when there is no next page.", 

121 ) 

122 previous_cursor: str | None = Field( 

123 default=None, 

124 description="Token pointing to the previous page. Populated for cursor pagination, " 

125 "``null`` when using offset pagination or when on the first page.", 

126 ) 

127 

128 

129def _task_instance_discriminator(v: Any) -> str: 

130 """Discriminate between TaskInstanceResponse and NewTaskResponse in the union.""" 

131 if isinstance(v, NewTaskResponse): 131 ↛ 132line 131 didn't jump to line 132 because the condition on line 131 was never true

132 return "new" 

133 if isinstance(v, dict): 133 ↛ 134line 133 didn't jump to line 134 because the condition on line 133 was never true

134 return "new" if "id" not in v else "full" 

135 # ORM objects and TaskInstanceResponse instances 

136 return "full" 

137 

138 

139class ClearTaskInstanceCollectionResponse(BaseModel): 

140 """Response for clear dag run dry run, which may contain new tasks without full TaskInstance data.""" 

141 

142 task_instances: Iterable[ 

143 Annotated[ 

144 Annotated[TaskInstanceResponse, Tag("full")] | Annotated[NewTaskResponse, Tag("new")], 

145 Discriminator(_task_instance_discriminator), 

146 ] 

147 ] 

148 total_entries: int 

149 

150 

151class TaskDependencyResponse(BaseModel): 

152 """Task Dependency serializer for responses.""" 

153 

154 name: str 

155 reason: str 

156 

157 

158class TaskDependencyCollectionResponse(BaseModel): 

159 """Task scheduling dependencies collection serializer for responses.""" 

160 

161 dependencies: list[TaskDependencyResponse] 

162 

163 

164class TaskInstancesBatchBody(StrictBaseModel): 

165 """Task Instance body for get batch.""" 

166 

167 dag_ids: list[str] | None = None 

168 dag_run_ids: list[str] | None = None 

169 task_ids: list[str] | None = None 

170 state: list[TaskInstanceState | None] | None = None 

171 

172 run_after_gte: AwareDatetime | None = None 

173 run_after_gt: AwareDatetime | None = None 

174 run_after_lte: AwareDatetime | None = None 

175 run_after_lt: AwareDatetime | None = None 

176 

177 logical_date_gte: AwareDatetime | None = None 

178 logical_date_gt: AwareDatetime | None = None 

179 logical_date_lte: AwareDatetime | None = None 

180 logical_date_lt: AwareDatetime | None = None 

181 

182 start_date_gte: AwareDatetime | None = None 

183 start_date_gt: AwareDatetime | None = None 

184 start_date_lte: AwareDatetime | None = None 

185 start_date_lt: AwareDatetime | None = None 

186 

187 end_date_gte: AwareDatetime | None = None 

188 end_date_gt: AwareDatetime | None = None 

189 end_date_lte: AwareDatetime | None = None 

190 end_date_lt: AwareDatetime | None = None 

191 

192 duration_gte: float | None = None 

193 duration_gt: float | None = None 

194 duration_lte: float | None = None 

195 duration_lt: float | None = None 

196 

197 pool: list[str] | None = None 

198 queue: list[str] | None = None 

199 executor: list[str] | None = None 

200 page_offset: NonNegativeInt = 0 

201 page_limit: NonNegativeInt = 100 

202 order_by: str | None = None 

203 

204 

205class ClearTaskInstancesBody(StrictBaseModel): 

206 """Request body for Clear Task Instances endpoint.""" 

207 

208 dry_run: bool = True 

209 start_date: AwareDatetime | None = None 

210 end_date: AwareDatetime | None = None 

211 only_failed: bool = True 

212 only_running: bool = False 

213 reset_dag_runs: bool = True 

214 task_ids: list[str | tuple[str, int]] | None = Field( 

215 default=None, 

216 description="A list of `task_id` or [`task_id`, `map_index`]. " 

217 "If only the `task_id` is provided for a mapped task, all of its map indices will be targeted.", 

218 ) 

219 dag_run_id: str | None = None 

220 include_upstream: bool = False 

221 include_downstream: bool = False 

222 include_future: bool = False 

223 include_past: bool = False 

224 run_on_latest_version: bool | None = Field( 

225 default=None, 

226 description="(Experimental) Run on the latest bundle version of the dag after " 

227 "clearing the task instances. " 

228 "If not specified, falls back to the DAG-level ``rerun_with_latest_version`` parameter, " 

229 "then the ``[core] rerun_with_latest_version`` config option, " 

230 "and finally ``False`` (the historical default for clear/rerun).", 

231 ) 

232 prevent_running_task: bool = False 

233 note: Annotated[str, StringConstraints(max_length=1000)] | None = None 

234 

235 @model_validator(mode="before") 

236 @classmethod 

237 def validate_model(cls, data: Any) -> Any: 

238 """Validate clear task instance form.""" 

239 if data.get("only_failed") and data.get("only_running"): 

240 raise ValueError("only_failed and only_running both are set to True") 

241 if data.get("start_date") and data.get("end_date"): 

242 if data.get("start_date") > data.get("end_date"): 

243 raise ValueError("end_date is sooner than start_date") 

244 if data.get("start_date") and data.get("end_date") and data.get("dag_run_id"): 

245 raise ValueError("Exactly one of dag_run_id or (start_date and end_date) must be provided") 

246 if data.get("start_date") and data.get("dag_run_id"): 

247 raise ValueError("Exactly one of dag_run_id or start_date must be provided") 

248 if data.get("end_date") and data.get("dag_run_id"): 

249 raise ValueError("Exactly one of dag_run_id or end_date must be provided") 

250 if isinstance(data.get("task_ids"), list) and len(data.get("task_ids")) < 1: 

251 raise ValueError("task_ids list should have at least 1 element.") 

252 return data 

253 

254 

255class PatchTaskInstanceBody(StrictBaseModel): 

256 """Request body for patching task instance state.""" 

257 

258 new_state: TaskInstanceState | None = None 

259 note: Annotated[str, StringConstraints(max_length=1000)] | None = None 

260 include_upstream: bool = False 

261 include_downstream: bool = False 

262 include_future: bool = False 

263 include_past: bool = False 

264 

265 @field_validator("new_state", mode="before") 

266 @classmethod 

267 def validate_new_state(cls, ns: str | None) -> str: 

268 """Validate new_state.""" 

269 valid_states = [ 

270 vs.name.lower() 

271 for vs in (TaskInstanceState.SUCCESS, TaskInstanceState.FAILED, TaskInstanceState.SKIPPED) 

272 ] 

273 if ns is None: 

274 raise ValueError("'new_state' should not be empty") 

275 ns = ns.lower() 

276 if ns not in valid_states: 

277 raise ValueError(f"'{ns}' is not one of {valid_states}") 

278 return ns 

279 

280 

281class BulkTaskInstanceBody(PatchTaskInstanceBody, StrictBaseModel): 

282 """Request body for bulk update, and delete task instances.""" 

283 

284 task_id: str 

285 map_index: int | None = None 

286 dag_id: str | None = None 

287 dag_run_id: str | None = None