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

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 

20import json 

21 

22from fastapi import HTTPException, status 

23from pydantic import ValidationError 

24from sqlalchemy import select 

25 

26from airflow._shared.secrets_masker import merge 

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.connections import ConnectionBody 

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

37from airflow.models.connection import Connection 

38 

39 

40def update_orm_from_pydantic( 

41 orm_conn: Connection, pydantic_conn: ConnectionBody, update_mask: list[str] | None = None 

42) -> None: 

43 """Update ORM object from Pydantic object.""" 

44 # Not all fields match and some need setters, therefore copy partly manually via setters 

45 non_update_fields = {"connection_id", "conn_id"} 

46 setter_fields = {"password", "extra"} 

47 fields_set = pydantic_conn.model_fields_set 

48 if "schema_" in fields_set: # Alias is not resolved correctly, need to patch 

49 fields_set.remove("schema_") 

50 fields_set.add("schema") 

51 fields_to_update = fields_set - non_update_fields - setter_fields 

52 if update_mask: 

53 fields_to_update = fields_to_update.intersection(update_mask) 

54 conn_data = pydantic_conn.model_dump(by_alias=True) 

55 for key, val in conn_data.items(): 

56 if key in fields_to_update: 

57 setattr(orm_conn, key, val) 

58 

59 if (not update_mask and "password" in pydantic_conn.model_fields_set) or ( 

60 update_mask and "password" in update_mask 

61 ): 

62 if pydantic_conn.password is None: 

63 orm_conn.set_password(pydantic_conn.password) 

64 else: 

65 merged_password = merge(pydantic_conn.password, orm_conn.password, "password") 

66 orm_conn.set_password(merged_password) 

67 if (not update_mask and "extra" in pydantic_conn.model_fields_set) or ( 

68 update_mask and "extra" in update_mask 

69 ): 

70 if pydantic_conn.extra is None or orm_conn.extra is None: 70 ↛ 73line 70 didn't jump to line 73 because the condition on line 70 was always true

71 orm_conn.set_extra(pydantic_conn.extra) 

72 return 

73 try: 

74 merged_extra = merge(json.loads(pydantic_conn.extra), json.loads(orm_conn.extra)) 

75 orm_conn.set_extra(json.dumps(merged_extra)) 

76 except json.JSONDecodeError: 

77 # We can't merge fields in an unstructured `extra` 

78 orm_conn.set_extra(pydantic_conn.extra) 

79 

80 

81class BulkConnectionService(BulkService[ConnectionBody]): 

82 """Service for handling bulk operations on connections.""" 

83 

84 def categorize_connections(self, connection_ids: set) -> tuple[dict, set, set]: 

85 """ 

86 Categorize the given connection_ids into matched_connection_ids and not_found_connection_ids based on existing connection_ids. 

87 

88 Existed connections are returned as a dict of {connection_id : Connection}. 

89 

90 :param connection_ids: set of connection_ids 

91 :return: tuple of dict of existed connections, set of matched connection_ids, set of not found connection_ids 

92 """ 

93 existed_connections = self.session.execute( 

94 select(Connection).filter(Connection.conn_id.in_(connection_ids)) 

95 ).scalars() 

96 existed_connections_dict = {conn.conn_id: conn for conn in existed_connections} 

97 matched_connection_ids = set(existed_connections_dict.keys()) 

98 not_found_connection_ids = connection_ids - matched_connection_ids 

99 return existed_connections_dict, matched_connection_ids, not_found_connection_ids 

100 

101 def handle_bulk_create( 

102 self, action: BulkCreateAction[ConnectionBody], results: BulkActionResponse 

103 ) -> None: 

104 """Bulk create connections.""" 

105 to_create_connection_ids = {connection.connection_id for connection in action.entities} 

106 existed_connections_dict, matched_connection_ids, not_found_connection_ids = ( 

107 self.categorize_connections(to_create_connection_ids) 

108 ) 

109 try: 

110 if action.action_on_existence == BulkActionOnExistence.FAIL and matched_connection_ids: 

111 raise HTTPException( 

112 status_code=status.HTTP_409_CONFLICT, 

113 detail=f"The connections with these connection_ids: {matched_connection_ids} already exist.", 

114 ) 

115 if action.action_on_existence == BulkActionOnExistence.SKIP: 

116 create_connection_ids = not_found_connection_ids 

117 else: 

118 create_connection_ids = to_create_connection_ids 

119 

120 for connection in action.entities: 

121 if connection.connection_id in create_connection_ids: 121 ↛ 120line 121 didn't jump to line 120 because the condition on line 121 was always true

122 if connection.connection_id in matched_connection_ids: 122 ↛ 123line 122 didn't jump to line 123 because the condition on line 122 was never true

123 existed_connection = existed_connections_dict[connection.connection_id] 

124 # Only overwrite fields the request actually provided (see pools.py for the 

125 # full rationale). Plain ``model_dump()`` resets omitted fields to their 

126 # defaults on the existing connection — e.g. silently nulling ``team_name`` 

127 # multi-team ownership. ``exclude_unset=True`` writes only the fields present 

128 # in the request body. 

129 for key, val in connection.model_dump(by_alias=True, exclude_unset=True).items(): 

130 setattr(existed_connection, key, val) 

131 else: 

132 self.session.add(Connection(**connection.model_dump(by_alias=True))) 

133 results.success.append(connection.connection_id) 

134 

135 except HTTPException as e: 

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

137 

138 def handle_bulk_update( 

139 self, action: BulkUpdateAction[ConnectionBody], results: BulkActionResponse 

140 ) -> None: 

141 """Bulk Update connections.""" 

142 to_update_connection_ids = {connection.connection_id for connection in action.entities} 

143 existed_connections_dict, matched_connection_ids, not_found_connection_ids = ( 

144 self.categorize_connections(to_update_connection_ids) 

145 ) 

146 

147 try: 

148 if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_connection_ids: 148 ↛ 149line 148 didn't jump to line 149 because the condition on line 148 was never true

149 raise HTTPException( 

150 status_code=status.HTTP_404_NOT_FOUND, 

151 detail=f"The connections with these connection_ids: {not_found_connection_ids} were not found.", 

152 ) 

153 if action.action_on_non_existence == BulkActionNotOnExistence.SKIP: 

154 update_connection_ids = matched_connection_ids 

155 else: 

156 update_connection_ids = to_update_connection_ids 

157 

158 for connection in action.entities: 

159 if connection.connection_id in update_connection_ids: 159 ↛ 158line 159 didn't jump to line 158 because the condition on line 159 was always true

160 old_connection = existed_connections_dict.get(connection.connection_id) 

161 if old_connection is None: 161 ↛ 162line 161 didn't jump to line 162 because the condition on line 161 was never true

162 raise ValidationError( 

163 f"The Connection with connection_id: `{connection.connection_id}` was not found" 

164 ) 

165 ConnectionBody(**connection.model_dump()) 

166 

167 update_orm_from_pydantic(old_connection, connection) 

168 results.success.append(connection.connection_id) 

169 

170 except HTTPException as e: 

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

172 

173 except ValidationError as e: 

174 results.errors.append({"error": f"{e.errors()}"}) 

175 

176 def handle_bulk_delete( 

177 self, action: BulkDeleteAction[ConnectionBody], results: BulkActionResponse 

178 ) -> None: 

179 """Bulk delete connections.""" 

180 to_delete_connection_ids = set(action.entities) 

181 existed_connections_dict, matched_connection_ids, not_found_connection_ids = ( 

182 self.categorize_connections(to_delete_connection_ids) 

183 ) 

184 

185 try: 

186 if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_connection_ids: 

187 raise HTTPException( 

188 status_code=status.HTTP_404_NOT_FOUND, 

189 detail=f"The connections with these connection_ids: {not_found_connection_ids} were not found.", 

190 ) 

191 if action.action_on_non_existence == BulkActionNotOnExistence.SKIP: 

192 delete_connection_ids = matched_connection_ids 

193 else: 

194 delete_connection_ids = to_delete_connection_ids 

195 

196 for connection_id in delete_connection_ids: 196 ↛ 197line 196 didn't jump to line 197 because the loop on line 196 never started

197 existing_connection = existed_connections_dict.get(connection_id) 

198 if existing_connection: 

199 self.session.delete(existing_connection) 

200 results.success.append(connection_id) 

201 

202 except HTTPException as e: 

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