Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/services/public/pools.py: 77%

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 cast 

21 

22from fastapi import HTTPException, status 

23from fastapi.exceptions import RequestValidationError 

24from pydantic import ValidationError 

25from sqlalchemy import select 

26 

27from airflow.api_fastapi.core_api.datamodels.common import ( 

28 BulkActionNotOnExistence, 

29 BulkActionOnExistence, 

30 BulkActionResponse, 

31 BulkCreateAction, 

32 BulkDeleteAction, 

33 BulkUpdateAction, 

34) 

35from airflow.api_fastapi.core_api.datamodels.pools import ( 

36 BasePool, 

37 PoolBody, 

38 PoolPatchBody, 

39) 

40from airflow.api_fastapi.core_api.services.public.common import BulkService 

41from airflow.models.pool import Pool 

42 

43 

44def update_orm_from_pydantic( 

45 pool: Pool, 

46 patch_body: PoolBody | PoolPatchBody, 

47 update_mask: list[str] | None, 

48) -> Pool: 

49 """ 

50 Update an existing pool. 

51 

52 :param pool: The existing Pool ORM object to update. 

53 :param patch_body: Pydantic model containing the fields to update. 

54 :param update_mask: Specific fields to update. If None, all provided fields will be considered. 

55 :return: The updated Pool instance. 

56 :raises HTTPException: If attempting to update disallowed fields on ``default_pool``. 

57 """ 

58 # Special restriction: default pool only allows limited fields to be patched 

59 pool_name = pool.pool 

60 if pool_name == Pool.DEFAULT_POOL_NAME: 

61 if update_mask and all(mask.strip() in {"slots", "include_deferred"} for mask in update_mask): 61 ↛ 63line 61 didn't jump to line 63 because the condition on line 61 was never true

62 # Validate only slots/include_deferred 

63 try: 

64 patch_body_subset = patch_body.model_dump( 

65 include={"slots", "include_deferred"}, exclude_unset=True, by_alias=True 

66 ) 

67 # Re-run validation with BasePool but only on allowed fields 

68 PoolPatchBody.model_validate(patch_body_subset) 

69 except ValidationError as e: 

70 raise RequestValidationError(errors=e.errors()) 

71 else: 

72 raise HTTPException( 

73 status.HTTP_400_BAD_REQUEST, 

74 "Only slots and included_deferred can be modified on Default Pool", 

75 ) 

76 else: 

77 fields_to_update = patch_body.model_fields_set 

78 try: 

79 # Dump with both input + output aliases handled 

80 body_dict = patch_body.model_dump( 

81 include=fields_to_update, 

82 by_alias=True, # ensures we get the API-facing alias keys 

83 ) 

84 

85 # Normalize keys for BasePool (expects "pool") 

86 if "name" in body_dict and "pool" not in body_dict: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true

87 body_dict["pool"] = body_dict.pop("name") 

88 

89 BasePool.model_validate(body_dict) 

90 

91 except ValidationError as e: 

92 raise RequestValidationError(errors=e.errors()) 

93 

94 # Delegate patch application to the common utility 

95 return cast( 

96 "Pool", 

97 BulkService.apply_patch_with_update_mask( 

98 model=pool, 

99 patch_body=patch_body, 

100 update_mask=update_mask, 

101 non_update_fields=None, 

102 ), 

103 ) 

104 

105 

106class BulkPoolService(BulkService[PoolBody]): 

107 """Service for handling bulk operations on pools.""" 

108 

109 def categorize_pools(self, pool_names: set) -> tuple[dict, set, set]: 

110 """ 

111 Categorize the given pool_names into matched_pool_names and not_found_pool_names based on existing pool_names. 

112 

113 Existing pools are returned as a dict of {pool_name : Pool}. 

114 

115 :param pool_names: set of pool_names 

116 :return: tuple of dict of existed pools, set of matched pool_names, set of not found pool_names 

117 """ 

118 existed_pools = self.session.execute(select(Pool).filter(Pool.pool.in_(pool_names))).scalars() 

119 existing_pools_dict = {pool.pool: pool for pool in existed_pools} 

120 matched_pool_names = set(existing_pools_dict.keys()) 

121 not_found_pool_names = pool_names - matched_pool_names 

122 return existing_pools_dict, matched_pool_names, not_found_pool_names 

123 

124 def handle_bulk_create(self, action: BulkCreateAction[PoolBody], results: BulkActionResponse) -> None: 

125 """Bulk create pools.""" 

126 to_create_pool_names = {pool.pool for pool in action.entities} 

127 existing_pools_dict, matched_pool_names, not_found_pool_names = self.categorize_pools( 

128 to_create_pool_names 

129 ) 

130 try: 

131 if action.action_on_existence == BulkActionOnExistence.FAIL and matched_pool_names: 

132 raise HTTPException( 

133 status_code=status.HTTP_409_CONFLICT, 

134 detail=f"The pools with these pool names: {matched_pool_names} already exist.", 

135 ) 

136 if action.action_on_existence == BulkActionOnExistence.SKIP: 

137 create_pool_names = not_found_pool_names 

138 else: 

139 create_pool_names = to_create_pool_names 

140 

141 for pool in action.entities: 

142 if pool.pool in create_pool_names: 

143 if pool.pool in matched_pool_names: 143 ↛ 144line 143 didn't jump to line 144 because the condition on line 143 was never true

144 existed_pool = existing_pools_dict[pool.pool] 

145 # Only overwrite fields the request actually provided. Plain ``model_dump()`` 

146 # emits every field at its default, so an overwrite that omits e.g. 

147 # ``team_name``/``description``/``include_deferred`` silently resets them on the 

148 # existing pool — most damagingly nulling its multi-team ``team_name`` ownership. 

149 # ``exclude_unset=True`` writes only fields present in the request body, so 

150 # omitted fields keep their current value while an explicitly-set field (even 

151 # ``None``) is still applied. 

152 for key, val in pool.model_dump(exclude_unset=True).items(): 

153 setattr(existed_pool, key, val) 

154 else: 

155 self.session.add(Pool(**pool.model_dump())) 

156 results.success.append(pool.pool) 

157 

158 self.session.flush() 

159 

160 except HTTPException as e: 

161 results.errors.append({"error": f"{e.detail}", "status_code": e.status_code}) 

162 

163 def handle_bulk_update(self, action: BulkUpdateAction[PoolBody], results: BulkActionResponse) -> None: 

164 """Bulk Update pools.""" 

165 to_update_pool_names = {pool.pool for pool in action.entities} 

166 existing_pools_dict, matched_pool_names, not_found_pool_names = self.categorize_pools( 

167 to_update_pool_names 

168 ) 

169 try: 

170 if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_pool_names: 

171 raise HTTPException( 

172 status_code=status.HTTP_404_NOT_FOUND, 

173 detail=f"The pools with these pool names: {not_found_pool_names} were not found.", 

174 ) 

175 if action.action_on_non_existence == BulkActionNotOnExistence.SKIP: 

176 update_pool_names = matched_pool_names 

177 else: 

178 update_pool_names = to_update_pool_names 

179 for pool in action.entities: 179 ↛ 180line 179 didn't jump to line 180 because the loop on line 179 never started

180 if pool.pool not in update_pool_names: 

181 continue 

182 

183 updated_pool = update_orm_from_pydantic( 

184 existing_pools_dict[pool.pool], pool, action.update_mask 

185 ) 

186 

187 results.success.append(str(updated_pool.pool)) # use request field, always consistent 

188 

189 except HTTPException as e: 

190 results.errors.append({"error": f"{e.detail}", "status_code": e.status_code}) 

191 

192 def handle_bulk_delete(self, action: BulkDeleteAction[PoolBody], results: BulkActionResponse) -> None: 

193 """Bulk delete pools.""" 

194 to_delete_pool_names = set(action.entities) 

195 existing_pools_dict, matched_pool_names, not_found_pool_names = self.categorize_pools( 

196 to_delete_pool_names 

197 ) 

198 

199 try: 

200 if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_pool_names: 200 ↛ 201line 200 didn't jump to line 201 because the condition on line 200 was never true

201 raise HTTPException( 

202 status_code=status.HTTP_404_NOT_FOUND, 

203 detail=f"The pools with these pool names: {not_found_pool_names} were not found.", 

204 ) 

205 

206 for pool_name in matched_pool_names: 

207 self.session.delete(existing_pools_dict[pool_name]) 

208 results.success.append(pool_name) 

209 

210 except HTTPException as e: 

211 results.errors.append({"error": f"{e.detail}", "status_code": e.status_code})