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

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 

21from collections.abc import Iterable, Mapping 

22from datetime import datetime 

23from typing import Annotated, Any 

24 

25from pydantic import Field, field_validator, model_validator 

26from pydantic_core.core_schema import ValidationInfo 

27 

28from airflow._shared.secrets_masker import redact, should_hide_value_for_key 

29from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel, make_partial_model 

30from airflow.configuration import conf 

31 

32 

33# Response Models 

34class ConnectionResponse(BaseModel): 

35 """Connection serializer for responses.""" 

36 

37 connection_id: str = Field(serialization_alias="connection_id", validation_alias="conn_id") 

38 conn_type: str 

39 description: str | None 

40 host: str | None 

41 login: str | None 

42 schema_: str | None = Field(alias="schema") 

43 port: int | None 

44 password: str | None 

45 extra: str | None 

46 team_name: str | None 

47 

48 @field_validator("password", mode="after") 

49 @classmethod 

50 def redact_password(cls, v: str | None, field_info: ValidationInfo) -> str | None: 

51 if v is None: 

52 return None 

53 return str(redact(v, field_info.field_name)) 

54 

55 @field_validator("extra", mode="before") 

56 @classmethod 

57 def redact_extra(cls, v: str | None) -> str | None: 

58 if v is None or v == "": 

59 return v 

60 try: 

61 extra_dict = json.loads(v) 

62 redacted_dict = redact(extra_dict) 

63 return json.dumps(redacted_dict) 

64 except json.JSONDecodeError: 

65 # Do not return un-redacted extra because this could cause sensitive information to be exposed. 

66 # This code path should never be hit as ``Connection._validate_extra`` makes sure that ``extra`` is 

67 # always a valid JSON string (if truthy). We add this safeguard just in case and to make the coupling 

68 # explicit. 

69 raise ValueError( 

70 "This code path should never happen as persisted Connections (DB layer) should always enforce `extra` as a JSON string." 

71 ) 

72 

73 

74class ConnectionCollectionResponse(BaseModel): 

75 """Connection Collection serializer for responses.""" 

76 

77 connections: Iterable[ConnectionResponse] 

78 total_entries: int 

79 

80 

81class ConnectionTestResponse(BaseModel): 

82 """Connection Test serializer for synchronous test responses.""" 

83 

84 status: bool 

85 message: str 

86 

87 

88class ConnectionTestQueuedResponse(BaseModel): 

89 """Response returned when a connection test has been enqueued for worker execution.""" 

90 

91 token: str 

92 connection_id: str 

93 state: str 

94 

95 

96class AsyncConnectionTestResponse(BaseModel): 

97 """Response returned when polling for the status of an enqueued connection test.""" 

98 

99 token: str 

100 connection_id: str 

101 state: str 

102 result_message: str | None = None 

103 created_at: datetime 

104 

105 

106class ConnectionHookFieldBehavior(BaseModel): 

107 """A class to store the behavior of each standard field of a Hook.""" 

108 

109 hidden: Annotated[ 

110 bool, 

111 Field(description="Flag if the form field should be hidden."), 

112 ] = False 

113 title: Annotated[ 

114 str | None, 

115 Field( 

116 description="Label / title for the field that should be displayed, if re-labelling is needed. Use `None` to display standard title." 

117 ), 

118 ] = None 

119 placeholder: Annotated[ 

120 str | None, 

121 Field(description="Placeholder text that should be populated to the form."), 

122 ] = None 

123 

124 

125class StandardHookFields(BaseModel): 

126 """Standard fields of a Hook that a form will render.""" 

127 

128 description: ConnectionHookFieldBehavior | None 

129 url_schema: ConnectionHookFieldBehavior | None 

130 host: ConnectionHookFieldBehavior | None 

131 port: ConnectionHookFieldBehavior | None 

132 login: ConnectionHookFieldBehavior | None 

133 password: ConnectionHookFieldBehavior | None 

134 

135 

136class ConnectionHookMetaData(BaseModel): 

137 """ 

138 Response model for Hook information == Connection type meta data. 

139 

140 It is used to transfer providers information loaded by providers_manager such that 

141 the API server/Web UI can use this data to render connection form UI. 

142 """ 

143 

144 connection_type: str | None 

145 hook_class_name: str | None 

146 default_conn_name: str | None 

147 hook_name: str 

148 standard_fields: StandardHookFields | None 

149 extra_fields: Mapping | None 

150 

151 @field_validator("extra_fields", mode="after") 

152 @classmethod 

153 def redact_extra_fields(cls, v: Mapping | None): 

154 if v is None: 

155 return None 

156 

157 # Check if extra_fields contains param spec structures (result of SerializedParam.dump()) 

158 # which have "value" and "schema" keys, or simple dictionary structures 

159 has_param_spec_structure = any( 

160 isinstance(field_spec, dict) and "value" in field_spec and "schema" in field_spec 

161 for field_spec in v.values() 

162 ) 

163 

164 if has_param_spec_structure: 

165 redacted_extra_fields: dict[str, Any] = {} 

166 for field_name, field_spec in v.items(): 

167 if isinstance(field_spec, dict) and "value" in field_spec and "schema" in field_spec: 

168 if should_hide_value_for_key(field_name) and field_spec.get("value") is not None: 

169 # Mask only the value, preserve everything else including schema.type 

170 redacted_extra_fields[field_name] = {**field_spec, "value": "***"} 

171 else: 

172 # Not sensitive or no value, keep as is 

173 redacted_extra_fields[field_name] = field_spec 

174 else: 

175 # Not a param spec structure, apply redact by default 

176 redacted_extra_fields[field_name] = redact(field_spec) 

177 

178 return redacted_extra_fields 

179 

180 # For simple dictionary structures, use the standard redact function 

181 return redact(v) 

182 

183 

184# Request Models 

185class ConnectionBody(StrictBaseModel): 

186 """Connection Serializer for requests body.""" 

187 

188 connection_id: str = Field(serialization_alias="conn_id", max_length=200, pattern=r"^[\w.-]+$") 

189 conn_type: str 

190 description: str | None = Field(default=None) 

191 host: str | None = Field(default=None) 

192 login: str | None = Field(default=None) 

193 schema_: str | None = Field(None, alias="schema") 

194 port: int | None = Field(default=None) 

195 password: str | None = Field(default=None) 

196 extra: str | None = Field(default=None) 

197 team_name: str | None = Field(max_length=50, default=None) 

198 

199 @field_validator("extra") 

200 @classmethod 

201 def validate_extra(cls, v: str | None) -> str | None: 

202 """ 

203 Validate that `extra` field is a JSON-encoded Python dict. 

204 

205 If `extra` field is not a valid JSON, it will be returned as is. 

206 """ 

207 if v is None: 

208 return v 

209 if v == "": 

210 return "{}" # Backward compatibility: treat "" as empty JSON object 

211 try: 

212 extra_dict = json.loads(v) 

213 if not isinstance(extra_dict, dict): 

214 raise ValueError("The `extra` field must be a valid JSON object (e.g., {'key': 'value'})") 

215 except json.JSONDecodeError: 

216 raise ValueError( 

217 "The `extra` field must be a valid JSON object (e.g., {'key': 'value'}), " 

218 "but encountered non-JSON in `extra` field" 

219 ) 

220 return v 

221 

222 @model_validator(mode="after") 

223 def validate_team_name(self) -> ConnectionBody: 

224 if self.team_name is not None and not conf.getboolean("core", "multi_team"): 

225 raise ValueError( 

226 "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." 

227 ) 

228 return self 

229 

230 

231ConnectionBodyPartial = make_partial_model(ConnectionBody) 

232 

233 

234class ConnectionTestRequestBody(ConnectionBody): 

235 """ 

236 Request body for enqueueing a connection test on a worker. 

237 

238 Inherits ``connection_id`` pattern, ``extra`` JSON validation, and 

239 ``team_name`` handling from ``ConnectionBody`` so tested connections share 

240 the same input contract as persisted ones. 

241 """ 

242 

243 commit_on_success: bool = Field( 

244 default=False, 

245 description="If True, save or update the connection in the connection table when the test succeeds.", 

246 ) 

247 executor: str | None = Field( 

248 default=None, 

249 description="Executor name to dispatch the connection test to.", 

250 ) 

251 queue: str | None = Field( 

252 default=None, 

253 description="Worker queue to route the connection test to (executor-dependent).", 

254 )