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

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 

19import os 

20from typing import TYPE_CHECKING, Annotated 

21 

22from fastapi import Depends, Header, HTTPException, Query, status 

23from fastapi.exceptions import RequestValidationError 

24from pydantic import ValidationError 

25from sqlalchemy import select 

26from sqlalchemy.exc import IntegrityError 

27 

28from airflow.api_fastapi.app import get_auth_manager 

29from airflow.api_fastapi.auth.managers.models.resource_details import ConnectionDetails 

30from airflow.api_fastapi.common.db.common import SessionDep, paginated_select 

31from airflow.api_fastapi.common.parameters import ( 

32 QueryConnectionIdPatternSearch, 

33 QueryConnectionIdPrefixPatternSearch, 

34 QueryLimit, 

35 QueryOffset, 

36 SortParam, 

37) 

38from airflow.api_fastapi.common.router import AirflowRouter 

39from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT 

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

41 BulkBody, 

42 BulkResponse, 

43) 

44from airflow.api_fastapi.core_api.datamodels.connections import ( 

45 AsyncConnectionTestResponse, 

46 ConnectionBody, 

47 ConnectionBodyPartial, 

48 ConnectionCollectionResponse, 

49 ConnectionResponse, 

50 ConnectionTestQueuedResponse, 

51 ConnectionTestRequestBody, 

52 ConnectionTestResponse, 

53) 

54from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc 

55from airflow.api_fastapi.core_api.security import ( 

56 AuthManagerDep, 

57 GetUserDep, 

58 ReadableConnectionsFilterDep, 

59 requires_access_connection, 

60 requires_access_connection_bulk, 

61) 

62from airflow.api_fastapi.core_api.services.public.connections import ( 

63 BulkConnectionService, 

64 update_orm_from_pydantic, 

65) 

66from airflow.api_fastapi.logging.decorators import action_logging 

67from airflow.configuration import conf 

68from airflow.exceptions import AirflowNotFoundException 

69from airflow.executors.executor_loader import ExecutorLoader 

70from airflow.models import Connection 

71from airflow.models.connection_test import ConnectionTestRequest 

72from airflow.secrets.environment_variables import CONN_ENV_PREFIX 

73from airflow.utils.db import create_default_connections as db_create_default_connections 

74from airflow.utils.strings import get_random_string 

75 

76if TYPE_CHECKING: 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true

77 from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod 

78 

79connections_router = AirflowRouter(tags=["Connection"], prefix="/connections") 

80 

81 

82def _ensure_test_connection_enabled() -> None: 

83 """Raise 403 if connection testing is not enabled in the Airflow configuration.""" 

84 if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled": 84 ↛ exitline 84 didn't return from function '_ensure_test_connection_enabled' because the condition on line 84 was always true

85 raise HTTPException( 

86 status.HTTP_403_FORBIDDEN, 

87 "Testing connections is disabled in Airflow configuration. " 

88 "Contact your deployment admin to enable it.", 

89 ) 

90 

91 

92def _ensure_executor_is_configured(executor: str | None) -> None: 

93 """Raise 422 if the requested executor is not in the configured executors list.""" 

94 if executor is None: 

95 return 

96 configured = ExecutorLoader.get_executor_names(validate_teams=False) 

97 if not any( 

98 executor in (name.alias, name.module_path, name.module_path.split(".")[-1]) for name in configured 

99 ): 

100 raise HTTPException( 

101 HTTP_422_UNPROCESSABLE_CONTENT, 

102 f"Executor '{executor}' is not configured. " 

103 f"Configured executors: {[name.alias or name.module_path for name in configured]}", 

104 ) 

105 

106 

107_MASKED_CREDENTIAL_SENTINEL = "***" 

108 

109 

110def _same_endpoint(requested: str | int | None, stored: str | int | None) -> bool: 

111 """ 

112 Return True when request and stored host/port refer to the same destination. 

113 

114 The UI sends empty string for hidden unused host/port fields; the ORM stores 

115 those as NULL. Treat blank as unset so connection types that do not use 

116 host/port still reuse stored credentials. 

117 """ 

118 

119 def _norm(value: str | int | None) -> str | int | None: 

120 return None if value is None or value == "" else value 

121 

122 return _norm(requested) == _norm(stored) 

123 

124 

125def _supplies_own_credentials(test_body: ConnectionBody) -> bool: 

126 """ 

127 Return True when the request includes a real (non-masked) password. 

128 

129 The UI always posts the masked sentinel for unchanged secrets. That is not 

130 a caller-supplied credential and must not skip restoring stored extras. 

131 """ 

132 if "password" not in test_body.model_fields_set: 

133 return False 

134 password = test_body.password 

135 return bool(password) and password != _MASKED_CREDENTIAL_SENTINEL 

136 

137 

138@connections_router.delete( 

139 "/{connection_id}", 

140 status_code=status.HTTP_204_NO_CONTENT, 

141 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

142 dependencies=[Depends(requires_access_connection(method="DELETE")), Depends(action_logging())], 

143) 

144def delete_connection( 

145 connection_id: str, 

146 session: SessionDep, 

147): 

148 """Delete a connection entry.""" 

149 connection = session.scalar(select(Connection).filter_by(conn_id=connection_id)) 

150 

151 if connection is None: 

152 raise HTTPException( 

153 status.HTTP_404_NOT_FOUND, f"The Connection with connection_id: `{connection_id}` was not found" 

154 ) 

155 

156 session.delete(connection) 

157 

158 

159@connections_router.get( 

160 "/enqueue-test", 

161 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

162) 

163def get_connection_test( 

164 session: SessionDep, 

165 user: GetUserDep, 

166 connection_test_token: Annotated[str, Header(alias="Airflow-Connection-Test-Token")], 

167) -> AsyncConnectionTestResponse: 

168 """Poll for the status of an enqueued connection test by its token (passed as a header).""" 

169 connection_test = session.scalar(select(ConnectionTestRequest).filter_by(token=connection_test_token)) 

170 

171 if connection_test is None: 171 ↛ 177line 171 didn't jump to line 177 because the condition on line 171 was always true

172 raise HTTPException( 

173 status.HTTP_404_NOT_FOUND, 

174 f"No connection test found for token: `{connection_test_token}`", 

175 ) 

176 

177 if not get_auth_manager().is_authorized_connection( 

178 method="GET", 

179 details=ConnectionDetails(conn_id=connection_test.connection_id, team_name=connection_test.team_name), 

180 user=user, 

181 ): 

182 raise HTTPException( 

183 status.HTTP_404_NOT_FOUND, 

184 f"No connection test found for token: `{connection_test_token}`", 

185 ) 

186 

187 return AsyncConnectionTestResponse( 

188 token=connection_test.token, 

189 connection_id=connection_test.connection_id, 

190 state=connection_test.state, 

191 result_message=connection_test.result_message, 

192 created_at=connection_test.created_at, 

193 ) 

194 

195 

196@connections_router.get( 

197 "/{connection_id}", 

198 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

199 dependencies=[Depends(requires_access_connection(method="GET"))], 

200) 

201def get_connection( 

202 connection_id: str, 

203 session: SessionDep, 

204) -> ConnectionResponse: 

205 """Get a connection entry.""" 

206 connection = session.scalar(select(Connection).filter_by(conn_id=connection_id)) 

207 

208 if connection is None: 

209 raise HTTPException( 

210 status.HTTP_404_NOT_FOUND, f"The Connection with connection_id: `{connection_id}` was not found" 

211 ) 

212 

213 return connection 

214 

215 

216@connections_router.get( 

217 "", 

218 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

219 dependencies=[Depends(requires_access_connection(method="GET"))], 

220) 

221def get_connections( 

222 limit: QueryLimit, 

223 offset: QueryOffset, 

224 order_by: Annotated[ 

225 SortParam, 

226 Depends( 

227 SortParam( 

228 ["conn_id", "conn_type", "description", "host", "port", "id", "team_name"], 

229 Connection, 

230 {"connection_id": "conn_id"}, 

231 ).dynamic_depends() 

232 ), 

233 ], 

234 readable_connections_filter: ReadableConnectionsFilterDep, 

235 session: SessionDep, 

236 connection_id_pattern: QueryConnectionIdPatternSearch, 

237 connection_id_prefix_pattern: QueryConnectionIdPrefixPatternSearch, 

238) -> ConnectionCollectionResponse: 

239 """Get all connection entries.""" 

240 connection_select, total_entries = paginated_select( 

241 statement=select(Connection), 

242 filters=[connection_id_pattern, connection_id_prefix_pattern, readable_connections_filter], 

243 order_by=order_by, 

244 offset=offset, 

245 limit=limit, 

246 session=session, 

247 ) 

248 

249 connections = session.scalars(connection_select) 

250 

251 return ConnectionCollectionResponse( 

252 connections=connections, 

253 total_entries=total_entries, 

254 ) 

255 

256 

257@connections_router.post( 

258 "", 

259 status_code=status.HTTP_201_CREATED, 

260 responses=create_openapi_http_exception_doc( 

261 [status.HTTP_409_CONFLICT] 

262 ), # handled by global exception handler 

263 dependencies=[Depends(requires_access_connection(method="POST")), Depends(action_logging())], 

264) 

265def post_connection( 

266 post_body: ConnectionBody, 

267 session: SessionDep, 

268) -> ConnectionResponse: 

269 """Create connection entry.""" 

270 connection = Connection(**post_body.model_dump(by_alias=True)) 

271 session.add(connection) 

272 return connection 

273 

274 

275@connections_router.patch( 

276 "", dependencies=[Depends(requires_access_connection_bulk()), Depends(action_logging())] 

277) 

278def bulk_connections( 

279 request: BulkBody[ConnectionBody], 

280 session: SessionDep, 

281) -> BulkResponse: 

282 """Bulk create, update, and delete connections.""" 

283 return BulkConnectionService(session=session, request=request).handle_request() 

284 

285 

286@connections_router.patch( 

287 "/{connection_id}", 

288 responses=create_openapi_http_exception_doc( 

289 [ 

290 status.HTTP_400_BAD_REQUEST, 

291 status.HTTP_404_NOT_FOUND, 

292 ] 

293 ), 

294 dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())], 

295) 

296def patch_connection( 

297 connection_id: str, 

298 patch_body: ConnectionBody, 

299 session: SessionDep, 

300 update_mask: list[str] | None = Query(None), 

301) -> ConnectionResponse: 

302 """Update a connection entry.""" 

303 if patch_body.connection_id != connection_id: 

304 raise HTTPException( 

305 status.HTTP_400_BAD_REQUEST, 

306 "The connection_id in the request body does not match the URL parameter", 

307 ) 

308 

309 connection = session.scalar(select(Connection).filter_by(conn_id=connection_id).limit(1)) 

310 

311 if connection is None: 

312 raise HTTPException( 

313 status.HTTP_404_NOT_FOUND, f"The Connection with connection_id: `{connection_id}` was not found" 

314 ) 

315 

316 if update_mask: 

317 fields_to_update = patch_body.model_fields_set & set(update_mask) 

318 try: 

319 ConnectionBodyPartial(**patch_body.model_dump(include=fields_to_update)) 

320 except ValidationError as e: 

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

322 

323 update_orm_from_pydantic(connection, patch_body, update_mask) 

324 return connection 

325 

326 

327@connections_router.post("/test", dependencies=[Depends(requires_access_connection(method="POST"))]) 

328def test_connection( 

329 test_body: ConnectionBody, 

330 user: GetUserDep, 

331 auth_manager: AuthManagerDep, 

332) -> ConnectionTestResponse: 

333 """ 

334 Test an API connection. 

335 

336 This method first creates an in-memory transient conn_id & exports that to an env var, 

337 as some hook classes tries to find out the `conn` from their __init__ method & errors out if not found. 

338 It also deletes the conn id env connection after the test. 

339 """ 

340 _ensure_test_connection_enabled() 

341 

342 transient_conn_id = get_random_string() 

343 conn_env_var = f"{CONN_ENV_PREFIX}{transient_conn_id.upper()}" 

344 try: 

345 # Authorize read access on the requested ``connection_id`` *before* 

346 # touching the secrets backends. The route-level POST dependency only 

347 # verifies the caller can create connections; merging the existing 

348 # connection's hidden fields also requires read access to that 

349 # specific connection. Gating the backend lookup itself (rather than 

350 # the post-load merge) prevents an unauthorized caller from using 

351 # this endpoint to enumerate protected connection ids, generate 

352 # access-log entries in audited backends, or impose backend load for 

353 # arbitrary ids. ``get_team_name`` is a metadata-only DB lookup and 

354 # does not touch the configured secrets backends. 

355 # 

356 # When the connection has no metadata-DB row (e.g. it lives only in 

357 # a team-aware secrets backend like Vault or Kubernetes), fall back 

358 # to the request body's validated ``team_name`` so the GET 

359 # authorization and the secrets lookup both run in the right team 

360 # scope. ``ConnectionBody.validate_team_name`` already rejects 

361 # ``team_name`` from clients when ``[core] multi_team`` is off, so 

362 # a non-None body value here is always already gated by that 

363 # validator. 

364 team_name = Connection.get_team_name(test_body.connection_id) 

365 if team_name is None: 

366 team_name = test_body.team_name 

367 existing_conn: Connection | None = None 

368 if auth_manager.is_authorized_connection( 

369 method="GET", 

370 details=ConnectionDetails( 

371 conn_id=test_body.connection_id, 

372 team_name=team_name, 

373 ), 

374 user=user, 

375 ): 

376 try: 

377 existing_conn = Connection.get_connection_from_secrets( 

378 test_body.connection_id, team_name=team_name 

379 ) 

380 except AirflowNotFoundException: 

381 existing_conn = None 

382 

383 if existing_conn is not None: 

384 # Stored credentials are only reused to test the connection's own 

385 # host/port; testing a different destination must supply its own. 

386 fields_set = test_body.model_fields_set 

387 host_changed = "host" in fields_set and not _same_endpoint(test_body.host, existing_conn.host) 

388 port_changed = "port" in fields_set and not _same_endpoint(test_body.port, existing_conn.port) 

389 if host_changed or port_changed: 

390 if not _supplies_own_credentials(test_body): 

391 raise HTTPException( 

392 status.HTTP_400_BAD_REQUEST, 

393 "The host or port to test differs from the stored connection. " 

394 "Include the credentials to test in the request body.", 

395 ) 

396 existing_conn = None 

397 

398 if existing_conn is not None: 

399 existing_conn.conn_id = transient_conn_id 

400 update_orm_from_pydantic(existing_conn, test_body) 

401 conn = existing_conn 

402 else: 

403 data = test_body.model_dump(by_alias=True) 

404 data["conn_id"] = transient_conn_id 

405 conn = Connection(**data) 

406 

407 os.environ[conn_env_var] = conn.get_uri() 

408 test_status, test_message = conn.test_connection() 

409 return ConnectionTestResponse.model_validate({"status": test_status, "message": test_message}) 

410 finally: 

411 os.environ.pop(conn_env_var, None) 

412 

413 

414@connections_router.post( 

415 "/enqueue-test", 

416 status_code=status.HTTP_202_ACCEPTED, 

417 responses=create_openapi_http_exception_doc( 

418 [ 

419 status.HTTP_403_FORBIDDEN, 

420 status.HTTP_409_CONFLICT, 

421 HTTP_422_UNPROCESSABLE_CONTENT, 

422 ] 

423 ), 

424 dependencies=[Depends(action_logging())], 

425) 

426def enqueue_connection_test( 

427 test_body: ConnectionTestRequestBody, 

428 session: SessionDep, 

429 user: GetUserDep, 

430) -> ConnectionTestQueuedResponse: 

431 """Enqueue a connection test for deferred execution on a worker; returns a polling token.""" 

432 _ensure_test_connection_enabled() 

433 _ensure_executor_is_configured(test_body.executor) 

434 

435 existing = session.scalar(select(Connection).filter_by(conn_id=test_body.connection_id)) 

436 if existing is not None: 

437 effective_team = existing.team_name 

438 if test_body.team_name is not None and test_body.team_name != effective_team: 

439 raise HTTPException( 

440 status.HTTP_403_FORBIDDEN, 

441 f"team_name `{test_body.team_name}` does not match the team of connection " 

442 f"`{test_body.connection_id}`.", 

443 ) 

444 else: 

445 effective_team = test_body.team_name 

446 

447 auth_method: ResourceMethod = "PUT" if existing is not None and test_body.commit_on_success else "POST" 

448 if not get_auth_manager().is_authorized_connection( 

449 method=auth_method, 

450 details=ConnectionDetails(conn_id=test_body.connection_id, team_name=effective_team), 

451 user=user, 

452 ): 

453 raise HTTPException( 

454 status.HTTP_403_FORBIDDEN, 

455 f"You are not authorized to test connection `{test_body.connection_id}`.", 

456 ) 

457 

458 connection_test = ConnectionTestRequest( 

459 connection_id=test_body.connection_id, 

460 conn_type=test_body.conn_type, 

461 host=test_body.host, 

462 login=test_body.login, 

463 password=test_body.password, 

464 schema=test_body.schema_, 

465 port=test_body.port, 

466 extra=test_body.extra, 

467 commit_on_success=test_body.commit_on_success, 

468 executor=test_body.executor, 

469 queue=test_body.queue, 

470 team_name=effective_team, 

471 ) 

472 session.add(connection_test) 

473 try: 

474 session.flush() 

475 except IntegrityError: 

476 raise HTTPException( 

477 status.HTTP_409_CONFLICT, 

478 f"An active connection test already exists for connection_id `{test_body.connection_id}`.", 

479 ) 

480 

481 return ConnectionTestQueuedResponse( 

482 token=connection_test.token, 

483 connection_id=connection_test.connection_id, 

484 state=connection_test.state, 

485 ) 

486 

487 

488@connections_router.post( 

489 "/defaults", 

490 status_code=status.HTTP_204_NO_CONTENT, 

491 dependencies=[Depends(requires_access_connection(method="POST")), Depends(action_logging())], 

492) 

493def create_default_connections( 

494 session: SessionDep, 

495): 

496 """Create default connections.""" 

497 db_create_default_connections(session=session)