Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/execution_api/versions/v2026_04_06.py: 59%

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 typing import Any 

21 

22from cadwyn import ResponseInfo, VersionChange, convert_response_to_previous_version_for, endpoint, schema 

23 

24from airflow.api_fastapi.common.types import UtcDateTime 

25from airflow.api_fastapi.execution_api.datamodels.asset_event import ( 

26 AssetEventResponse, 

27 AssetEventsResponse, 

28 DagRunAssetReference, 

29) 

30from airflow.api_fastapi.execution_api.datamodels.dagrun import TriggerDAGRunPayload 

31from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( 

32 DagRun, 

33 TIDeferredStatePayload, 

34 TIRunContext, 

35) 

36 

37 

38class AddPartitionKeyField(VersionChange): 

39 """Add ``partition_key`` to the Dag run, asset event, asset reference and trigger payload models.""" 

40 

41 description = __doc__ 

42 

43 instructions_to_migrate_to_previous_version = ( 

44 schema(DagRun).field("partition_key").didnt_exist, 

45 schema(AssetEventResponse).field("partition_key").didnt_exist, 

46 schema(TriggerDAGRunPayload).field("partition_key").didnt_exist, 

47 schema(DagRunAssetReference).field("partition_key").didnt_exist, 

48 ) 

49 

50 @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] 

51 def remove_partition_key_from_dag_run(response: ResponseInfo) -> None: # type: ignore[misc] 

52 """Remove the `partition_key` field from the dag_run object when converting to the previous version.""" 

53 dag_run = response.body.get("dag_run") 

54 if isinstance(dag_run, dict): 

55 dag_run.pop("partition_key", None) 

56 

57 @convert_response_to_previous_version_for(AssetEventsResponse) # type: ignore[arg-type] 

58 def remove_partition_key_from_asset_events(response: ResponseInfo) -> None: # type: ignore[misc] 

59 """Remove the `partition_key` field from the dag_run object when converting to the previous version.""" 

60 events = response.body["asset_events"] 

61 for elem in events: 

62 elem.pop("partition_key", None) 

63 

64 

65class MovePreviousRunEndpoint(VersionChange): 

66 """Add new previous-run endpoint and migrate old endpoint.""" 

67 

68 description = __doc__ 

69 

70 instructions_to_migrate_to_previous_version = ( 

71 endpoint("/dag-runs/previous", ["GET"]).didnt_exist, 

72 endpoint("/dag-runs/{dag_id}/previous", ["GET"]).existed, 

73 ) 

74 

75 

76class AddDagRunDetailEndpoint(VersionChange): 

77 """Add dag run detail endpoint.""" 

78 

79 description = __doc__ 

80 

81 instructions_to_migrate_to_previous_version = ( 

82 endpoint("/dag-runs/{dag_id}/{run_id}", ["GET"]).didnt_exist, 

83 ) 

84 

85 

86class MakeDagRunStartDateNullable(VersionChange): 

87 """Make DagRun.start_date field nullable for runs that haven't started yet.""" 

88 

89 description = __doc__ 

90 

91 instructions_to_migrate_to_previous_version = (schema(DagRun).field("start_date").had(type=UtcDateTime),) 

92 

93 @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] 

94 def ensure_start_date_in_ti_run_context(response: ResponseInfo) -> None: # type: ignore[misc] 

95 """ 

96 Ensure start_date is never None in DagRun for previous API versions. 

97 

98 Older Task SDK clients expect start_date to be non-nullable. When the 

99 DagRun hasn't started yet (e.g. queued), fall back to run_after. 

100 """ 

101 dag_run = response.body.get("dag_run") 

102 if isinstance(dag_run, dict) and dag_run.get("start_date") is None: 

103 dag_run["start_date"] = dag_run.get("run_after") 

104 

105 @convert_response_to_previous_version_for(DagRun) # type: ignore[arg-type] 

106 def ensure_start_date_in_dag_run(response: ResponseInfo) -> None: # type: ignore[misc] 

107 """Ensure start_date is never None in direct DagRun responses for previous API versions.""" 

108 if response.body.get("start_date") is None: 

109 response.body["start_date"] = response.body.get("run_after") 

110 

111 

112class ModifyDeferredTaskKwargsToJsonValue(VersionChange): 

113 """Change the types of `trigger_kwargs` and `next_kwargs` in TIDeferredStatePayload to JsonValue.""" 

114 

115 description = __doc__ 

116 

117 instructions_to_migrate_to_previous_version = ( 

118 schema(TIDeferredStatePayload).field("trigger_kwargs").had(type=dict[str, Any] | str), 

119 schema(TIDeferredStatePayload).field("next_kwargs").had(type=dict[str, Any]), 

120 ) 

121 

122 @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] 

123 def convert_next_kwargs_to_base_serialization(response: ResponseInfo) -> None: # type: ignore[misc] 

124 """ 

125 Convert next_kwargs from SDK serde format to BaseSerialization format for old workers. 

126 

127 Old workers (task-sdk < 1.2) only know BaseSerialization.deserialize(), which requires 

128 dicts wrapped as {"__type": "dict", "__var": {...}}. SDK serde produces plain dicts that 

129 BaseSerialization cannot parse, causing KeyError on __var. 

130 

131 We must deserialize SDK serde first to recover native Python objects (datetime, 

132 timedelta, etc.), then re-serialize with BaseSerialization so old workers get 

133 proper typed values instead of raw {"__classname__": ...} dicts. 

134 """ 

135 next_kwargs = response.body.get("next_kwargs") 

136 if next_kwargs is None: 

137 return 

138 

139 from airflow.sdk.serde import deserialize 

140 from airflow.serialization.serialized_objects import BaseSerialization 

141 

142 try: 

143 plain = deserialize(next_kwargs) 

144 except (ImportError, KeyError, AttributeError, TypeError): 

145 # Already in BaseSerialization format (rolling upgrade, old data in DB) 

146 return 

147 

148 response.body["next_kwargs"] = BaseSerialization.serialize(plain) 

149 

150 

151class RemoveUpstreamMapIndexesField(VersionChange): 

152 """Remove upstream_map_indexes field from TIRunContext - now computed by Task SDK.""" 

153 

154 description = __doc__ 

155 

156 instructions_to_migrate_to_previous_version = ( 

157 schema(TIRunContext) 

158 .field("upstream_map_indexes") 

159 .existed_as(type=dict[str, int | list[int] | None] | None), 

160 ) 

161 

162 @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] 

163 def add_upstream_map_indexes_field(response: ResponseInfo) -> None: # type: ignore[misc] 

164 """Add upstream_map_indexes field with None for older API versions.""" 

165 response.body["upstream_map_indexes"] = None 

166 

167 

168class AddNoteField(VersionChange): 

169 """Add note parameter to DagRun Model.""" 

170 

171 description = __doc__ 

172 

173 instructions_to_migrate_to_previous_version = ( 

174 schema(DagRun).field("note").didnt_exist, 

175 schema(TriggerDAGRunPayload).field("note").didnt_exist, 

176 ) 

177 

178 @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] 

179 def remove_note_field(response: ResponseInfo) -> None: # type: ignore[misc] 

180 """Remove note field for older API versions.""" 

181 if "dag_run" in response.body and isinstance(response.body["dag_run"], dict): 

182 response.body["dag_run"].pop("note", None) 

183 

184 

185class AddTaskInstanceStartDateField(VersionChange): 

186 """Add `start_date` field to TIRunContext.""" 

187 

188 description = __doc__ 

189 

190 instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("start_date").didnt_exist,) 

191 

192 @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] 

193 def remove_start_date_field(response: ResponseInfo) -> None: # type: ignore[misc] 

194 """Remove start_date field for older API versions.""" 

195 response.body.pop("start_date", None) 

196 

197 

198class AddDagEndpoint(VersionChange): 

199 """Add the `/dags/{dag_id}` endpoint.""" 

200 

201 description = __doc__ 

202 

203 instructions_to_migrate_to_previous_version = (endpoint("/dags/{dag_id}", ["GET"]).didnt_exist,) 

204 

205 

206class AddRunAfterField(VersionChange): 

207 """Add run_after parameter to TriggerDAGRunPayload Model.""" 

208 

209 description = __doc__ 

210 

211 instructions_to_migrate_to_previous_version = ( 

212 schema(TriggerDAGRunPayload).field("run_after").didnt_exist, 

213 )