kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3169602718
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -259,6 +284,87 @@ def test_connection(test_body: ConnectionBody) ->
ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
+ "/test-async",
+ status_code=status.HTTP_202_ACCEPTED,
+ responses=create_openapi_http_exception_doc(
+ [
+ status.HTTP_403_FORBIDDEN,
+ status.HTTP_409_CONFLICT,
+ status.HTTP_422_UNPROCESSABLE_ENTITY,
+ ]
+ ),
+ dependencies=[Depends(requires_access_connection(method="POST")),
Depends(action_logging())],
+)
+def test_connection_async(
+ test_body: ConnectionTestRequestBody,
+ session: SessionDep,
+) -> ConnectionTestQueuedResponse:
+ """
+ Queue an async connection test to be executed on a worker.
+
+ The connection data is stored in the test request table and the worker
+ reads from there. Returns a token to poll for the result via
+ GET /connections/test-async/{token}.
+ """
+ _ensure_test_connection_enabled()
+ _ensure_executor_is_configured(test_body.executor)
+
+ connection_test = ConnectionTestRequest(
+ connection_id=test_body.connection_id,
+ conn_type=test_body.conn_type,
+ host=test_body.host,
+ login=test_body.login,
+ password=test_body.password,
+ schema=test_body.schema_,
+ port=test_body.port,
+ extra=test_body.extra,
+ commit_on_success=test_body.commit_on_success,
+ executor=test_body.executor,
+ queue=test_body.queue,
+ )
+ session.add(connection_test)
+
+ return ConnectionTestQueuedResponse(
+ token=connection_test.token,
+ connection_id=connection_test.connection_id,
+ state=connection_test.state,
+ )
+
+
+@connections_router.get(
+ "/test-async/{connection_test_token}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
+ dependencies=[Depends(requires_access_connection(method="GET"))],
Review Comment:
I think there's an authorization bypass here. `requires_access_connection`
reads `connection_id` from `request.path_params`, but the path param on this
route is `connection_test_token`, not `connection_id`. So the auth callback
runs with `connection_id=None` and only checks generic 'can read connections'
rights, never the per-connection (or team-level) rights for the actual
`ct.connection_id`.
A user without rights on `prod_db` who obtains a token (log leak, captured
request, etc.) can poll the result and read `result_message`, which usually
carries hostnames, usernames, and error context. In multi-team setups this
circumvents team isolation.
Fix is to load the `ConnectionTestRequest` row first, then call
`is_authorized_connection(GET, ConnectionDetails(conn_id=ct.connection_id,
team_name=...))`. Returning 404 on unauthorized (rather than 403) avoids a
token-existence oracle.
##########
airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py:
##########
@@ -0,0 +1,43 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from pydantic import Field
+
+from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
+from airflow.models.connection_test import ConnectionTestState
+
+
+class ConnectionTestResultBody(StrictBaseModel):
+ """Payload sent by workers to report connection test results."""
+
+ state: ConnectionTestState
Review Comment:
Typing `state` as the full `ConnectionTestState` enum lets a worker PATCH
the test back to `PENDING` / `QUEUED` / `RUNNING` indefinitely. That bumps
`updated_at` on every call, so the reaper (which gates on `updated_at < now -
timeout - grace`) never fires, and `commit_on_success` is never triggered
either. Combined with the `active_connection_id` unique constraint, a stuck
test blocks all future tests for the same `connection_id`.
Consider constraining this to `Literal[ConnectionTestState.SUCCESS,
ConnectionTestState.FAILED]` since those are the only states a worker should be
reporting via this endpoint. The `RUNNING` transition the supervisor needs is
better expressed as a separate path (e.g. a dedicated `/start` endpoint) rather
than a free-form state setter.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py:
##########
@@ -0,0 +1,132 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Response, Security, status
+
+from airflow.api_fastapi.auth.tokens import JWTGenerator
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.connection_test import (
+ ConnectionTestConnectionResponse,
+ ConnectionTestResultBody,
+)
+from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+from airflow.api_fastapi.execution_api.deps import DepContainer
+from airflow.api_fastapi.execution_api.security import CurrentTIToken,
ExecutionAPIRoute, require_auth
+from airflow.models.connection_test import (
+ TERMINAL_STATES,
+ ConnectionTestRequest,
+ ConnectionTestState,
+)
+
+router = VersionedAPIRouter()
+
+ct_id_router = VersionedAPIRouter(
+ route_class=ExecutionAPIRoute,
+ dependencies=[
+ Security(require_auth, scopes=["ct:self"]),
+ ],
+)
+
+
+@ct_id_router.get(
+ "/{connection_test_id}/connection",
+ responses={
+ status.HTTP_404_NOT_FOUND: {"description": "Connection test not
found"},
+ },
+)
+def get_connection_test_connection(
+ connection_test_id: UUID,
+ session: SessionDep,
+) -> ConnectionTestConnectionResponse:
+ """Return the connection data stored in a test request (called by
workers)."""
+ ct = session.get(ConnectionTestRequest, connection_test_id)
+ if ct is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Connection test {connection_test_id} not found",
+ },
+ )
+
+ return ConnectionTestConnectionResponse(
+ conn_id=ct.connection_id,
+ conn_type=ct.conn_type,
+ host=ct.host,
+ login=ct.login,
+ password=ct.password,
+ schema=ct.schema,
+ port=ct.port,
+ extra=ct.extra,
+ )
+
+
+@ct_id_router.patch(
+ "/{connection_test_id}",
+ status_code=status.HTTP_204_NO_CONTENT,
+ dependencies=[Security(require_auth, scopes=["token:execution",
"token:workload"])],
+ responses={
+ status.HTTP_404_NOT_FOUND: {"description": "Connection test not
found"},
+ status.HTTP_409_CONFLICT: {"description": "Connection test already in
a terminal state"},
+ },
+)
+def patch_connection_test(
+ connection_test_id: UUID,
+ body: ConnectionTestResultBody,
+ response: Response,
+ session: SessionDep,
+ services=DepContainer,
+ token: TIToken = CurrentTIToken,
+) -> None:
+ """Update the result of a connection test (called by workers)."""
+ ct = session.get(ConnectionTestRequest, connection_test_id,
with_for_update=True)
+ if ct is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Connection test {connection_test_id} not found",
+ },
+ )
+
+ if ct.state in TERMINAL_STATES:
+ raise HTTPException(
+ status_code=status.HTTP_409_CONFLICT,
+ detail={
+ "reason": "conflict",
+ "message": f"Connection test {connection_test_id} is already
in terminal state: {ct.state}",
+ },
+ )
+
+ ct.state = body.state
Review Comment:
Companion to the body-typing comment: even if the model is tightened, this
handler still blindly assigns `ct.state = body.state` without validating the
transition. A `RUNNING -> PENDING` write would still succeed today. Worth
pairing the model change with an explicit allow-list here (`if body.state not
in {SUCCESS, FAILED, RUNNING}: raise 422`) or a small state machine, so the
invariant is enforced in one obvious place.
##########
airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py:
##########
@@ -78,12 +79,49 @@ class ConnectionCollectionResponse(BaseModel):
class ConnectionTestResponse(BaseModel):
- """Connection Test serializer for responses."""
+ """Connection Test serializer for synchronous test responses."""
status: bool
message: str
+class ConnectionTestRequestBody(StrictBaseModel):
+ """Request body for async connection test."""
+
+ connection_id: str
+ conn_type: str
+ host: str | None = None
+ login: str | None = None
+ schema_: str | None = Field(None, alias="schema")
+ port: int | None = None
+ password: str | None = None
+ extra: str | None = None
Review Comment:
`ConnectionBody.validate_extra` rejects malformed JSON at the API boundary,
but `ConnectionTestRequestBody` skips that validator. Today an invalid `extra`
is only caught when the worker constructs `Connection(...)` and
`_validate_extra` runs via `set_extra`, by which point the request is already
persisted and the user only learns about the failure via the round-tripped
`result_message`.
Same `field_validator("extra")` as `ConnectionBody` would reject upfront and
avoid storing junk.
##########
airflow-core/src/airflow/models/crypto.py:
##########
@@ -93,6 +96,60 @@ def rotate(self, msg: bytes | str) -> bytes:
return self._fernet.rotate(msg)
+class FernetFieldsMixin:
+ """Mixin providing Fernet-encrypted ``password`` and ``extra`` fields."""
+
+ _password: Mapped[str | None] = mapped_column("password", Text(),
nullable=True)
+ _extra: Mapped[str | None] = mapped_column("extra", Text(), nullable=True)
+ is_encrypted: Mapped[bool] = mapped_column(Boolean, unique=False,
default=False, nullable=False)
+ is_extra_encrypted: Mapped[bool] = mapped_column(Boolean, unique=False,
default=False, nullable=False)
+
+ def get_password(self) -> str | None:
+ """Decrypt and return password."""
+ if self._password and self.is_encrypted:
+ fernet = get_fernet()
+ if not fernet.is_encrypted:
+ raise ValueError("Can't decrypt encrypted password, FERNET_KEY
configuration is missing")
+ return fernet.decrypt(bytes(self._password, "utf-8")).decode()
+ return self._password
+
+ def set_password(self, value: str | None):
+ """Encrypt and store password."""
+ if value:
+ fernet = get_fernet()
+ self._password = fernet.encrypt(bytes(value, "utf-8")).decode()
Review Comment:
`set_password` only writes when `value` is truthy. Setting `obj.password =
None` or `""` leaves `_password` and `is_encrypted` unchanged, so an existing
encrypted secret stays in the DB. `set_extra` already handles the falsy case in
its `else` branch (lines 143-145). Worth mirroring that pattern here so 'clear
the password' actually clears it.
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py:
##########
@@ -1222,6 +1231,184 @@ def
test_should_test_new_connection_without_existing(self, test_client):
assert response.json()["status"] is True
+class TestAsyncConnectionTest(TestConnectionEndpoint):
+ """Tests for the async connection test endpoints (POST + GET polling)."""
+
+ TEST_REQUEST_BODY = {
+ "connection_id": TEST_CONN_ID,
+ "conn_type": TEST_CONN_TYPE,
+ "host": TEST_CONN_HOST,
+ }
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_post_should_respond_202(self, test_client, session):
+ """POST /connections/test-async returns 202 + token."""
+ response = test_client.post("/connections/test-async",
json=self.TEST_REQUEST_BODY)
+ assert response.status_code == 202
+ body = response.json()
+ assert "token" in body
+ assert body["connection_id"] == TEST_CONN_ID
+ assert body["state"] == "pending"
+ assert len(body["token"]) > 0
+
+ def test_should_respond_401(self, unauthenticated_test_client):
+ response = unauthenticated_test_client.post("/connections/test-async",
json=self.TEST_REQUEST_BODY)
+ assert response.status_code == 401
+
+ def test_should_respond_403(self, unauthorized_test_client):
+ response = unauthorized_test_client.post("/connections/test-async",
json=self.TEST_REQUEST_BODY)
+ assert response.status_code == 403
+
+ def test_should_respond_403_by_default(self, test_client):
+ """Connection testing is disabled by default."""
+ response = test_client.post("/connections/test-async",
json=self.TEST_REQUEST_BODY)
+ assert response.status_code == 403
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_post_creates_connection_test_request_row(self, test_client,
session):
+ """POST creates a ConnectionTestRequest row in PENDING state with
connection fields."""
+ response = test_client.post("/connections/test-async",
json=self.TEST_REQUEST_BODY)
+ assert response.status_code == 202
+ token = response.json()["token"]
+
+ ct =
session.scalar(select(ConnectionTestRequest).filter_by(token=token))
+ assert ct is not None
+ assert ct.connection_id == TEST_CONN_ID
+ assert ct.conn_type == TEST_CONN_TYPE
+ assert ct.host == TEST_CONN_HOST
+ assert ct.state == "pending"
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_post_passes_queue_parameter(self, test_client, session):
+ """POST /connections/test-async passes the queue parameter."""
+ body = {**self.TEST_REQUEST_BODY, "queue": "gpu_workers"}
+ response = test_client.post("/connections/test-async", json=body)
+ assert response.status_code == 202
+ token = response.json()["token"]
+
+ ct =
session.scalar(select(ConnectionTestRequest).filter_by(token=token))
+ assert ct is not None
+ assert ct.queue == "gpu_workers"
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_post_stores_commit_on_success(self, test_client, session):
+ """POST /connections/test-async stores the commit_on_success flag."""
+ body = {**self.TEST_REQUEST_BODY, "commit_on_success": True}
+ response = test_client.post("/connections/test-async", json=body)
+ assert response.status_code == 202
+ token = response.json()["token"]
+
+ ct =
session.scalar(select(ConnectionTestRequest).filter_by(token=token))
+ assert ct is not None
+ assert ct.commit_on_success is True
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_post_returns_409_for_duplicate_active_test(self, test_client,
session):
+ """POST returns 409 when there's already an active test for the same
connection_id."""
+ response = test_client.post("/connections/test-async",
json=self.TEST_REQUEST_BODY)
+ assert response.status_code == 202
+
+ response = test_client.post("/connections/test-async",
json=self.TEST_REQUEST_BODY)
+ assert response.status_code == 409
+ assert response.json()["detail"]["reason"] == "Unique constraint
violation"
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_post_rejects_unknown_executor_with_422(self, test_client,
session):
+ """POST returns 422 when the requested executor is not configured."""
+ body = {**self.TEST_REQUEST_BODY, "executor": "no_such_executor"}
+ response = test_client.post("/connections/test-async", json=body)
+ assert response.status_code == 422
+ assert "no_such_executor" in response.text
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_post_accepts_configured_executor(self, test_client, session):
+ """POST accepts an executor name that matches a configured executor."""
+ configured = ExecutorLoader.get_executor_names(validate_teams=False)
+ executor_name = configured[0].alias or configured[0].module_path
+ body = {**self.TEST_REQUEST_BODY, "executor": executor_name}
+ response = test_client.post("/connections/test-async", json=body)
+ assert response.status_code == 202
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_get_status_returns_pending(self, test_client, session):
+ """GET /connections/test-async/{token} returns current status."""
+ post_response = test_client.post("/connections/test-async",
json=self.TEST_REQUEST_BODY)
+ token = post_response.json()["token"]
+
+ response = test_client.get(f"/connections/test-async/{token}")
+ assert response.status_code == 200
+ body = response.json()
+ assert body["token"] == token
+ assert body["connection_id"] == TEST_CONN_ID
+ assert body["state"] == "pending"
+ assert body["result_message"] is None
+ assert "created_at" in body
+ assert "reverted" not in body
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_get_status_returns_completed_result(self, test_client, session):
+ """GET returns result after the worker has updated the test."""
+ post_response = test_client.post("/connections/test-async",
json=self.TEST_REQUEST_BODY)
+ token = post_response.json()["token"]
+
+ ct =
session.scalar(select(ConnectionTestRequest).filter_by(token=token))
+ ct.state = ConnectionTestState.SUCCESS
+ ct.result_message = "Connection successfully tested"
+ session.commit()
+
+ response = test_client.get(f"/connections/test-async/{token}")
+ assert response.status_code == 200
+ body = response.json()
+ assert body["state"] == "success"
+ assert body["result_message"] == "Connection successfully tested"
+
+ def test_get_status_returns_404_for_invalid_token(self, test_client):
Review Comment:
Tests cover unauthenticated/unauthorized for the POST and 404 for unknown
tokens, but nothing covers the cross-user case for GET-by-token: user A submits
a test for `connection_id=prod_db`, user B (no rights on `prod_db`) calls `GET
/connections/test-async/{token}` and reads the result. That's the path the
auth-bypass comment on the route is about, and a test here would have caught
it. Worth adding alongside `test_get_status_returns_404_for_invalid_token`.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3189,6 +3208,88 @@ def _activate_assets_generate_warnings() ->
Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ def _enqueue_connection_tests(self, *, session: Session) -> None:
+ """Enqueue pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("scheduler",
"max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("scheduler", "connection_test_timeout",
fallback=60)
+
+ num_occupied_slots = sum(executor.slots_occupied for executor in
self.executors)
+ parallelism_budget = conf.getint("core", "parallelism") -
num_occupied_slots
Review Comment:
`parallelism_budget` subtracts `slots_occupied` summed across all executors
from the global `core.parallelism`, but in multi-executor setups each executor
has its own parallelism. A task-heavy executor consuming all 'slots' here can
suppress connection-test dispatch even when the executor that would run the
test has plenty of free capacity.
The per-executor check inside the loop (via `executor.slots_available`) is
the more accurate gate. Could also just drop this aggregate check, since
`max_connection_test_concurrency` is the real safety knob for this feature.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3189,6 +3208,88 @@ def _activate_assets_generate_warnings() ->
Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ def _enqueue_connection_tests(self, *, session: Session) -> None:
+ """Enqueue pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("scheduler",
"max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("scheduler", "connection_test_timeout",
fallback=60)
+
+ num_occupied_slots = sum(executor.slots_occupied for executor in
self.executors)
+ parallelism_budget = conf.getint("core", "parallelism") -
num_occupied_slots
+ if parallelism_budget <= 0:
+ return
+
+ active_count = session.scalar(
Review Comment:
Two HA schedulers can both read `active_count = 0` here, both compute
`budget = max_concurrency`, and (because `with_row_locks(skip_locked=True)`
returns disjoint rows) each grabs up to `budget` PENDING tests. Net dispatched
on that tick is `N * budget`, which defeats the
`max_connection_test_concurrency` cap the PR description leans on as a
worker-protection knob.
Atomic alternatives: a single `UPDATE ... SET state='queued' ... WHERE
state='pending' LIMIT (max_concurrency - active_count) RETURNING ...`
(Postgres) or a sentinel row taken with `FOR UPDATE` to serialize the
read-modify-write. If the team is fine with N*budget per tick, please document
that explicitly in the docstring.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -338,10 +349,24 @@ def heartbeat(self) -> None:
self._emit_metrics(open_slots, num_running_workloads,
num_queued_workloads)
Review Comment:
`num_queued_workloads` passed in here is computed two lines up as
`len(self.queued_tasks) + len(self.queued_callbacks)` and does not include
`len(self.queued_connection_tests)`. So the `executor.queued_tasks` gauge
under-reports actual queue depth once connection tests are in flight, and
operators monitoring queue pressure won't see a backlog of tests until the
reaper starts failing them. (`slots_occupied` already counts them, so this is
purely a metric bug.)
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3189,6 +3208,88 @@ def _activate_assets_generate_warnings() ->
Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ def _enqueue_connection_tests(self, *, session: Session) -> None:
+ """Enqueue pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("scheduler",
"max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("scheduler", "connection_test_timeout",
fallback=60)
+
+ num_occupied_slots = sum(executor.slots_occupied for executor in
self.executors)
+ parallelism_budget = conf.getint("core", "parallelism") -
num_occupied_slots
+ if parallelism_budget <= 0:
+ return
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTestRequest.id)).where(
+ ConnectionTestRequest.state.in_(DISPATCHED_STATES)
+ )
+ )
+ concurrency_budget = max_concurrency - (active_count or 0)
+ budget = min(concurrency_budget, parallelism_budget)
+ if budget <= 0:
+ return
+
+ pending_stmt = (
+ select(ConnectionTestRequest)
+ .where(ConnectionTestRequest.state == ConnectionTestState.PENDING)
+ .order_by(ConnectionTestRequest.created_at)
+ .limit(budget)
+ )
+ pending_stmt = with_row_locks(pending_stmt, session,
of=ConnectionTestRequest, skip_locked=True)
+ pending_tests = session.scalars(pending_stmt).all()
+
+ if not pending_tests:
+ return
+
+ for ct in pending_tests:
+ executor = self._try_to_load_executor(ct, session)
+ if executor is None:
+ reason = f"No executor matches '{ct.executor}'"
+ ct.state = ConnectionTestState.FAILED
+ ct.result_message = reason
+ self.log.warning("Failing connection test %s: %s", ct.id,
reason)
+ continue
+ if not executor.supports_connection_test:
+ exec_name = executor.name
+ name = ct.executor or (exec_name and (exec_name.alias or
exec_name.module_path))
+ reason = f"Executor '{name}' does not support connection
testing"
+ ct.state = ConnectionTestState.FAILED
+ ct.result_message = reason
+ self.log.warning("Failing connection test %s: %s", ct.id,
reason)
+ continue
+
+ workload = workloads.TestConnection.make(
+ connection_test_id=ct.id,
+ connection_id=ct.connection_id,
+ timeout=timeout,
+ queue=ct.queue,
+ generator=executor.jwt_generator,
+ )
+ executor.queue_workload(workload, session=session)
+ ct.state = ConnectionTestState.QUEUED
+
+ session.flush()
+
+ @provide_session
+ def _reap_stale_connection_tests(self, *, session: Session = NEW_SESSION)
-> None:
+ """Mark connection tests that have exceeded their timeout as FAILED."""
+ timeout = conf.getint("scheduler", "connection_test_timeout",
fallback=60)
+ grace_period = max(30, timeout // 2)
+ cutoff = timezone.utcnow() - timedelta(seconds=timeout + grace_period)
+
+ stale_stmt = select(ConnectionTestRequest).where(
+ ConnectionTestRequest.state.in_(ACTIVE_STATES),
+ ConnectionTestRequest.updated_at < cutoff,
+ )
+ stale_stmt = with_row_locks(stale_stmt, session,
of=ConnectionTestRequest, skip_locked=True)
+ stale_tests = session.scalars(stale_stmt).all()
+
+ for ct in stale_tests:
+ ct.state = ConnectionTestState.FAILED
+ ct.result_message = f"Connection test timed out (exceeded
{timeout}s + {grace_period}s grace)"
+ self.log.warning("Reaped stale connection test %s", ct.id)
Review Comment:
When the reaper marks a test FAILED in the DB, the executor's in-memory
`queued_connection_tests` / `running` state is not touched. The next heartbeat
still hands the workload to the worker, which then PATCHes a SUCCESS/FAILED to
a row already in a terminal state and gets a 409 (handled by the endpoint at
line 110 of the execution-api route). Net effect is a wasted execution, a noisy
409 in worker logs, and a window where the connection's transient credentials
are exercised after the user has been told the test failed.
Worth either notifying executors to drop the key when reaping, or having the
worker re-check state before running the hook.
##########
task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py:
##########
@@ -0,0 +1,108 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Supervised execution of TestConnection workloads."""
+
+from __future__ import annotations
+
+import signal
+import uuid
+
+import structlog
+
+from airflow.sdk.api.client import Client
+from airflow.sdk.api.datamodels._generated import ConnectionTestState
+from airflow.sdk.definitions.connection import Connection as SDKConnection
+from airflow.sdk.execution_time.context import _preset_connections
+
+__all__ = ["supervise_connection_test"]
+
+log = structlog.get_logger(logger_name="connection_test_supervisor")
+
+
+def supervise_connection_test(
+ *,
+ connection_test_id: uuid.UUID,
+ connection_id: str,
+ timeout: int,
+ token: str,
+ server: str,
+) -> int:
+ """Execute a connection test on the worker and report the result via the
Execution API."""
+ from airflow.models.connection import Connection
Review Comment:
Importing `airflow.models.connection.Connection` from inside task-sdk
reverses task-sdk's dependency direction. task-sdk must not depend on
`airflow.models`, and the active work is to remove these edges, not add new
ones.
The import is also unnecessary. `Connection.test_connection()` in
`airflow-core/src/airflow/models/connection.py` just calls `self.get_hook()`
and then `hook.test_connection()`. The SDK `Connection`
(`task-sdk/src/airflow/sdk/definitions/connection.py:216`) already implements
`get_hook()` via `ProvidersManagerTaskRuntime`, and this function is already
constructing an `SDKConnection` (lines 69-78) with the same fields. The whole
`Connection(...)` / `conn.test_connection()` block (lines 59-68 and 81)
collapses to using `preset.get_hook()` and calling `hook.test_connection()`
directly. That drops the airflow-core import and the duplicate object.
##########
airflow-core/src/airflow/utils/db_cleanup.py:
##########
@@ -172,6 +172,7 @@ def readable_config(self):
),
_TableConfig(table_name="deadline", recency_column_name="deadline_time",
dag_id_column_name="dag_id"),
_TableConfig(table_name="revoked_token", recency_column_name="exp"),
+ _TableConfig(table_name="connection_test_request",
recency_column_name="created_at"),
Review Comment:
Cleanup uses `created_at` as the recency column with no state filter. With
long retention periods that's fine, but if an operator sets a short retention
(e.g. 1h for housekeeping), in-flight tests created near the cutoff can be
deleted while still PENDING/QUEUED/RUNNING. That orphans the executor's
in-memory entry and surfaces as a 404 on the worker's PATCH. Filtering on
`state in TERMINAL_STATES` (or using `updated_at` as the recency column) would
make this safe regardless of retention.
##########
airflow-core/src/airflow/migrations/versions/0113_3_3_0_add_connection_test_table.py:
##########
@@ -0,0 +1,84 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add connection_test_request table for async connection testing.
+
+Revision ID: a7e6d4c3b2f1
+Revises: fde9ed84d07b
+Create Date: 2026-02-22 00:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+from airflow.utils.sqlalchemy import UtcDateTime
+
+# revision identifiers, used by Alembic.
+revision = "a7e6d4c3b2f1"
+down_revision = "fde9ed84d07b"
+branch_labels = None
+depends_on = None
+airflow_version = "3.3.0"
+
+
+def upgrade():
+ """Create connection_test_request table."""
+ op.create_table(
+ "connection_test_request",
+ sa.Column("id", sa.Uuid(), nullable=False),
+ sa.Column("token", sa.String(64), nullable=False),
+ sa.Column("connection_id", sa.String(250), nullable=False),
+ sa.Column("state", sa.String(20), nullable=False),
+ sa.Column("result_message", sa.Text(), nullable=True),
+ sa.Column("created_at", UtcDateTime(timezone=True), nullable=False),
+ sa.Column("updated_at", UtcDateTime(timezone=True), nullable=False),
+ sa.Column("executor", sa.String(256), nullable=True),
+ sa.Column("queue", sa.String(256), nullable=True),
+ sa.Column("conn_type", sa.String(500), nullable=False),
+ sa.Column("host", sa.String(500), nullable=True),
+ sa.Column("login", sa.Text(), nullable=True),
+ sa.Column("password", sa.Text(), nullable=True),
+ sa.Column("schema", sa.String(500), nullable=True),
+ sa.Column("port", sa.Integer(), nullable=True),
+ sa.Column("extra", sa.Text(), nullable=True),
+ sa.Column("is_encrypted", sa.Boolean(), nullable=False),
Review Comment:
These two columns are NOT NULL but lack a `server_default`.
`commit_on_success` two lines below sets `server_default="0"` for exactly this
reason. Without a server default, any out-of-band insert (raw SQL, replication
tooling, or a future ORM path that forgets to set the flag) trips a NOT NULL
constraint instead of getting a sensible default. Cheap to add
`server_default="0"` here for consistency with the other Boolean column in the
same table.
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,210 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import secrets
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import (
+ Boolean,
+ Index,
+ Integer,
+ String,
+ Text,
+ UniqueConstraint,
+ Uuid,
+ select,
+)
+from sqlalchemy.orm import Mapped, mapped_column, validates
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import FernetFieldsMixin
+from airflow.utils.sqlalchemy import UtcDateTime
+
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
+log = structlog.get_logger(__name__)
+
+
+class ConnectionTestState(str, Enum):
+ """All possible states of a connection test."""
+
+ PENDING = "pending"
+ QUEUED = "queued"
+ RUNNING = "running"
+ SUCCESS = "success"
+ FAILED = "failed"
+
+ def __str__(self) -> str:
+ return self.value
+
+
+ACTIVE_STATES = frozenset(
+ (ConnectionTestState.PENDING, ConnectionTestState.QUEUED,
ConnectionTestState.RUNNING)
+)
+DISPATCHED_STATES = frozenset((ConnectionTestState.QUEUED,
ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS,
ConnectionTestState.FAILED))
+
+
+@dataclass(frozen=True, slots=True)
+class ConnectionTestKey:
+ """Typed key for connection-test workloads (wraps str(UUID))."""
+
+ id: str
+
+ def __str__(self) -> str:
+ return self.id
+
+
+class ConnectionTestRequest(Base, FernetFieldsMixin):
+ """
+ Tracks an async connection test request dispatched to a worker.
+
+ Stores the full connection details so the worker reads from this table
+ instead of the real ``connection`` table. The real ``connection`` table
+ is only modified if the test succeeds and ``commit_on_success`` is True.
+ """
+
+ __tablename__ = "connection_test_request"
+
+ id: Mapped[UUID] = mapped_column(Uuid(), primary_key=True,
default=uuid6.uuid7)
+ token: Mapped[str] = mapped_column(String(64), nullable=False, unique=True)
+ connection_id: Mapped[str] = mapped_column(String(250), nullable=False)
+ state: Mapped[str] = mapped_column(String(20), nullable=False,
default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
Review Comment:
`result_message` is `Text` with no length cap, and the API datamodel doesn't
bound it either. A buggy or hostile worker can write multi-MB messages, which
then get stored, returned on every poll, and counted against `db_cleanup` based
only on age. Worth a `String(N)` (or a `max_length` validator on
`ConnectionTestResultBody.result_message`) to keep this bounded.
##########
task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py:
##########
@@ -0,0 +1,108 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Supervised execution of TestConnection workloads."""
+
+from __future__ import annotations
+
+import signal
+import uuid
+
+import structlog
+
+from airflow.sdk.api.client import Client
+from airflow.sdk.api.datamodels._generated import ConnectionTestState
+from airflow.sdk.definitions.connection import Connection as SDKConnection
+from airflow.sdk.execution_time.context import _preset_connections
+
+__all__ = ["supervise_connection_test"]
+
+log = structlog.get_logger(logger_name="connection_test_supervisor")
+
+
+def supervise_connection_test(
+ *,
+ connection_test_id: uuid.UUID,
+ connection_id: str,
+ timeout: int,
+ token: str,
+ server: str,
+) -> int:
+ """Execute a connection test on the worker and report the result via the
Execution API."""
+ from airflow.models.connection import Connection
+
+ client = Client(base_url=server, token=token)
+
+ def _handle_timeout(signum, frame):
+ raise TimeoutError(f"Connection test timed out after {timeout}s")
+
+ signal.signal(signal.SIGALRM, _handle_timeout)
+ signal.alarm(timeout)
Review Comment:
`signal.SIGALRM` and `signal.alarm` are Unix-only. Today this is gated by
`LocalExecutor` having `supports_connection_test = True`, but if a future
executor that runs on Windows opts in (or someone runs the worker on Windows
directly), this raises `AttributeError` at import time. Either guard with `if
hasattr(signal, 'SIGALRM')` and fall back to a `threading.Timer`-based
watchdog, or document the platform constraint on `supports_connection_test`.
##########
airflow-core/src/airflow/executors/workloads/connection_test.py:
##########
@@ -0,0 +1,88 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Connection test workload schema for executor communication."""
+
+from __future__ import annotations
+
+import uuid
+from typing import TYPE_CHECKING, Literal
+
+from pydantic import Field
+
+from airflow.executors.workloads.base import BaseWorkloadSchema
+from airflow.models.connection_test import ConnectionTestKey
+
+if TYPE_CHECKING:
+ from airflow.api_fastapi.auth.tokens import JWTGenerator
+ from airflow.models.connection_test import ConnectionTestState
+
+
+class TestConnection(BaseWorkloadSchema):
+ """Execute a connection test on a worker."""
+
+ connection_test_id: uuid.UUID
+ connection_id: str
+ timeout: int
+ queue: str | None = None
+
+ type: Literal["TestConnection"] = Field(init=False,
default="TestConnection")
+
+ @property
+ def key(self) -> ConnectionTestKey:
+ """Return the connection-test key (str UUID) for this workload."""
+ return ConnectionTestKey(id=str(self.connection_test_id))
+
+ @property
+ def display_name(self) -> str:
+ """Return a human-readable name for logging and process titles."""
+ return f"connection-test {self.connection_id}"
+
+ @property
+ def success_state(self) -> ConnectionTestState:
+ from airflow.models.connection_test import ConnectionTestState
Review Comment:
These three properties each repeat `from airflow.models.connection_test
import ConnectionTestState` inside the function body, even though there's no
circular-import risk (line 27 already imports `ConnectionTestKey` from the same
module at module load). Move the import to the top alongside
`ConnectionTestKey` and drop the runtime imports on lines 56, 62, 68 (the
`TYPE_CHECKING` entry on line 31 also becomes redundant).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
