You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "anishgirianish (via GitHub)" <gi...@apache.org> on 2026/02/23 06:33:27 UTC
[PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
anishgirianish opened a new pull request, #62343:
URL: https://github.com/apache/airflow/pull/62343
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
## Summary
Follows the direction proposed by @potiuk in #59643 to move connection testing off the API server and onto workers.
Connection testing has been disabled by default since Airflow 2.7.0 because executing user-supplied driver code (ODBC/JDBC) on the API server poses security risks, and workers typically have network access to external systems that API servers don't — making API-server results misleading.
This PR adds async connection testing that dispatches tests to workers via the ExecutorCallback infrastructure:
- **`POST /connections/test-async`** : queues a connection test, returns a crypto-random polling token (`secrets.token_urlsafe`). Connection must already be saved; only `connection_id` is sent, workers fetch credentials
via `get_connection_from_secrets`.
- **`GET /connections/test-async/{token}`** : polls for results. Token knowledge serves as authorization.
- **`PATCH /execution/connection-tests/{id}`** : execution API endpoint for workers to report results back.
- **`ConnectionTest` model + migration** : state tracking (`PENDING` → `QUEUED` → `RUNNING` → `SUCCESS`/`FAILED`), callback FK, and `db clean` support via `created_at`.
The scheduler dispatch (picking up QUEUED ExecutorCallbacks and sending them to executors) depends on #61153. Once that merges `ExecuteCallback.make()` will be decoupled from `DagRun`, connection tests will execute end-to-end.
## References
- Dev mailing list discussion: [[DISCUSS] Move connection testing to workers](https://www.mail archive.com/dev@airflow.apache.org/msg21118.html)
- Builds on Callback infrastructure from #54796
- Depends on executor dispatch from #61153 for end-to-end flow
## Test plan
- [x] Unit tests for `ConnectionTest` model (token generation, state transitions)
- [x] Unit tests for `run_connection_test` worker function (success, failure, exception)
- [x] Unit tests for core API endpoints (POST 202, GET polling, 403 disabled, 404 not found)
- [x] Unit tests for execution API endpoint (PATCH 204, 404, 409 conflict, 422 validation)
- [x] Manual testing against PostgreSQL via Breeze (both core API and execution API)
- [x] Migration verified (`0105_3_2_0_add_connection_test_table`)
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177829001
##########
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:
Done, supervisor no longer imports airflow.models.connection; SDK Connection got test_connection(). Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177835113
##########
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
+ ct.result_message = body.result_message
+
+ if body.state == ConnectionTestState.SUCCESS and ct.commit_on_success:
+ ct.commit_to_connection_table(session=session)
+
+ # JWTReissueMiddleware also writes Refreshed-API-Token but skips workload tokens, so we set it here for the workload→execution swap.
+ if token.claims.scope == "workload":
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub": str(connection_test_id), "scope": "execution"})
+ response.headers["Refreshed-API-Token"] = execution_token
Review Comment:
Agreed, removed, by completion-PATCH time, the swap is dead. Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177835903
##########
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"]),
+ ],
+)
Review Comment:
Done, collapsed to a single router thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177836802
##########
airflow-core/src/airflow/api_fastapi/execution_api/security.py:
##########
@@ -190,6 +190,14 @@ async def require_auth(
detail="Token subject does not match task instance ID",
)
+ if "ct:self" in security_scopes.scopes:
Review Comment:
Done updated if→elif.
##########
airflow-core/src/airflow/api_fastapi/execution_api/security.py:
##########
@@ -190,6 +190,14 @@ async def require_auth(
detail="Token subject does not match task instance ID",
)
+ if "ct:self" in security_scopes.scopes:
Review Comment:
Done updated if→elif. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971260042
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via a TestConnection workload."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True)
+ reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0")
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(
+ self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return f"<ConnectionTest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+
+def run_connection_test(*, conn: Connection) -> tuple[bool, str]:
+ """
+ Worker-side function to execute a connection test.
+
+ Returns a (success, message) tuple. The caller is responsible for
+ reporting the result back via the Execution API.
+ """
+ try:
+ return conn.test_connection()
+ except Exception as e:
+ log.exception("Connection test failed", connection_id=conn.conn_id)
+ return False, str(e)
+
+
+_SNAPSHOT_FIELDS = (
+ "conn_type",
+ "description",
+ "host",
+ "login",
+ "_password",
+ "schema",
+ "port",
+ "_extra",
+ "is_encrypted",
+ "is_extra_encrypted",
+)
+
+
+def snapshot_connection(conn: Connection) -> dict:
+ """
+ Capture raw DB column values from a Connection for later restore.
+
+ Encrypted fields (``_password``, ``_extra``) are stored as ciphertext
+ so they can be written directly back without re-encryption.
+ """
+ return {field: getattr(conn, field) for field in _SNAPSHOT_FIELDS}
+
+
+def _revert_connection(conn: Connection, snapshot: dict) -> None:
+ """
+ Restore a Connection's columns from a snapshot dict.
+
+ Writes directly to ``_password`` and ``_extra`` (bypassing the
+ encrypting property setters) so the stored ciphertext is preserved.
+ """
+ for field, value in snapshot.items():
+ setattr(conn, field, value)
+
+
+def _decrypt_snapshot_field(snapshot: dict, field: str) -> str | None:
+ """Decrypt a single encrypted field from a snapshot dict using Fernet."""
+ raw = snapshot.get(field)
+ if raw is None:
+ return None
+ encrypted_flag = "is_encrypted" if field == "_password" else "is_extra_encrypted"
+ if not snapshot.get(encrypted_flag, False):
+ return raw
+ fernet = get_fernet()
+ return fernet.decrypt(bytes(raw, "utf-8")).decode()
+
+
+def _can_safely_revert(conn: Connection, post_snapshot: dict) -> bool:
+ """
+ Check whether the connection's current state matches the post-edit snapshot.
+
+ Compares **decrypted** values for encrypted fields and direct values for
+ non-encrypted fields. Returns ``False`` if any field differs, indicating
+ a concurrent edit has occurred and the revert should be skipped.
+ """
+ for field in _SNAPSHOT_FIELDS:
+ if field in ("is_encrypted", "is_extra_encrypted"):
+ continue
+
+ if field == "_password":
+ current_val = conn.password
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_password")
+ elif field == "_extra":
+ current_val = conn.extra
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_extra")
+ else:
+ current_val = getattr(conn, field)
+ snapshot_val = post_snapshot.get(field)
+
+ if current_val != snapshot_val:
+ return False
+ return True
+
+
+def attempt_revert(ct: ConnectionTest, *, session: Session) -> None:
+ """Revert a connection to its pre-edit values if no concurrent edit has occurred."""
+ if not ct.connection_snapshot:
+ log.warning("attempt_revert called without snapshot", connection_test_id=ct.id)
+ return
+
+ pre_snapshot = ct.connection_snapshot["pre"]
+ post_snapshot = ct.connection_snapshot["post"]
+
+ ct.connection_snapshot = None
+
+ conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id))
+ if conn is None:
+ ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists."
+ log.warning("Revert skipped: connection no longer exists", connection_id=ct.connection_id)
+ return
+
+ if not _can_safely_revert(conn, post_snapshot):
+ ct.result_message = (
+ ct.result_message or ""
Review Comment:
The snapshot mechanism has been removed now, and now its using request-buffer mechanism
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via a TestConnection workload."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True)
+ reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0")
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(
+ self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return f"<ConnectionTest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+
+def run_connection_test(*, conn: Connection) -> tuple[bool, str]:
+ """
+ Worker-side function to execute a connection test.
+
+ Returns a (success, message) tuple. The caller is responsible for
+ reporting the result back via the Execution API.
+ """
+ try:
+ return conn.test_connection()
+ except Exception as e:
+ log.exception("Connection test failed", connection_id=conn.conn_id)
+ return False, str(e)
+
+
+_SNAPSHOT_FIELDS = (
+ "conn_type",
+ "description",
+ "host",
+ "login",
+ "_password",
+ "schema",
+ "port",
+ "_extra",
+ "is_encrypted",
+ "is_extra_encrypted",
+)
+
+
+def snapshot_connection(conn: Connection) -> dict:
+ """
+ Capture raw DB column values from a Connection for later restore.
+
+ Encrypted fields (``_password``, ``_extra``) are stored as ciphertext
+ so they can be written directly back without re-encryption.
+ """
+ return {field: getattr(conn, field) for field in _SNAPSHOT_FIELDS}
+
+
+def _revert_connection(conn: Connection, snapshot: dict) -> None:
+ """
+ Restore a Connection's columns from a snapshot dict.
+
+ Writes directly to ``_password`` and ``_extra`` (bypassing the
+ encrypting property setters) so the stored ciphertext is preserved.
+ """
+ for field, value in snapshot.items():
+ setattr(conn, field, value)
+
+
+def _decrypt_snapshot_field(snapshot: dict, field: str) -> str | None:
+ """Decrypt a single encrypted field from a snapshot dict using Fernet."""
+ raw = snapshot.get(field)
+ if raw is None:
+ return None
+ encrypted_flag = "is_encrypted" if field == "_password" else "is_extra_encrypted"
+ if not snapshot.get(encrypted_flag, False):
+ return raw
+ fernet = get_fernet()
+ return fernet.decrypt(bytes(raw, "utf-8")).decode()
+
+
+def _can_safely_revert(conn: Connection, post_snapshot: dict) -> bool:
+ """
+ Check whether the connection's current state matches the post-edit snapshot.
+
+ Compares **decrypted** values for encrypted fields and direct values for
+ non-encrypted fields. Returns ``False`` if any field differs, indicating
+ a concurrent edit has occurred and the revert should be skipped.
+ """
+ for field in _SNAPSHOT_FIELDS:
+ if field in ("is_encrypted", "is_extra_encrypted"):
+ continue
+
+ if field == "_password":
+ current_val = conn.password
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_password")
+ elif field == "_extra":
+ current_val = conn.extra
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_extra")
+ else:
+ current_val = getattr(conn, field)
+ snapshot_val = post_snapshot.get(field)
+
+ if current_val != snapshot_val:
+ return False
+ return True
+
+
+def attempt_revert(ct: ConnectionTest, *, session: Session) -> None:
+ """Revert a connection to its pre-edit values if no concurrent edit has occurred."""
+ if not ct.connection_snapshot:
+ log.warning("attempt_revert called without snapshot", connection_test_id=ct.id)
+ return
+
+ pre_snapshot = ct.connection_snapshot["pre"]
+ post_snapshot = ct.connection_snapshot["post"]
+
+ ct.connection_snapshot = None
+
+ conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id))
+ if conn is None:
+ ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists."
+ log.warning("Revert skipped: connection no longer exists", connection_id=ct.connection_id)
+ return
+
+ if not _can_safely_revert(conn, post_snapshot):
Review Comment:
The snapshot mechanism has been removed now, and now its using request-buffer mechanism
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3025670031
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -305,10 +315,24 @@ def heartbeat(self) -> None:
self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads)
self.trigger_tasks(open_slots)
+ self.trigger_connection_tests()
+
# Calling child class sync method
self.log.debug("Calling the %s sync method", self.__class__)
self.sync()
+ def trigger_connection_tests(self) -> None:
+ """Process queued connection tests, respecting available slot capacity."""
+ if not self.supports_connection_test or not self.queued_connection_tests:
+ return
+
+ available = self.slots_available
+ if available <= 0:
+ return
+
+ tests_to_run = list(self.queued_connection_tests.values())[:available]
+ self._process_workloads(tests_to_run)
Review Comment:
This is not a deadlock. slots_available is the capacity model — it correctly accounts for all queued work (tasks, callbacks, connection tests). When slots are full, the function returns early, which is the intended behavior. Connection tests get processed when slots free up after tasks complete. The dispatch method in the scheduler (_dispatch_connection_tests) computes its own budget from parallelism - len(running) and picks up PENDING rows independently.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -305,10 +315,24 @@ def heartbeat(self) -> None:
self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads)
self.trigger_tasks(open_slots)
+ self.trigger_connection_tests()
+
# Calling child class sync method
self.log.debug("Calling the %s sync method", self.__class__)
self.sync()
+ def trigger_connection_tests(self) -> None:
+ """Process queued connection tests, respecting available slot capacity."""
+ if not self.supports_connection_test or not self.queued_connection_tests:
+ return
+
+ available = self.slots_available
+ if available <= 0:
+ return
+
+ tests_to_run = list(self.queued_connection_tests.values())[:available]
+ self._process_workloads(tests_to_run)
Review Comment:
This is not a deadlock. slots_available is the capacity model it correctly accounts for all queued work (tasks, callbacks, connection tests). When slots are full, the function returns early, which is the intended behavior. Connection tests get processed when slots free up after tasks complete. The dispatch method in the scheduler (_dispatch_connection_tests) computes its own budget from parallelism - len(running) and picks up PENDING rows independently.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4008414198
Hey @o-nikolas , @choo121600, thanks for the thorough review! I've addressed all the feedback:
- Switched to _try_to_load_executor for routing and added a supports_connection_test post-check.
- Connection tests now respect the parallelism budget and are counted in executor slot accounting.
- Moved dispatch into the main scheduler loop.
- Fixed the team name / executor name mixup in the test.
- Added tests for slot budget exhaustion and unsupported executor cases.
- Snapshot is now cleared (SQL NULL) after terminal state.
Would like to request you for another look when you get a chance!
Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2897389432
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Maybe the queue_type should also include something to indicate queue priority and whether that queue is intended to be used FIFO or prioritized?
Something vaguely like
If we add a NamedTuple called `WorkloadQueueDef(scheduling_tier: int, sort_key: int)` into the base workload and they are defined like this in the individual workloads:
```python
class TestConnection:
@property
def queue_def(self) ->WorkloadQueueDef :
return WorkloadQueueDef(
scheduling_tier = -1, # connection tests are super fast so let's always get them out of the way first
sort_key = 0, # FIFO order
)
class ExecuteCallback:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
scheduling_tier = 0, # always finish a callback before starting a new task, but don't block the much faster fast connection tests for a callback
sort_key = 0, # FIFO order
)
class ExecuteTask:
@property
def WorkloadQueueDef(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
# after all connection tests and callbacks are running, take the task with the highest priority
scheduling_tier = 1,
sort_key = self.ti.priority_weight,
)
```
and replace all of the "supports_callbacks", "supports_connection_test", etc in each executor can be replaced by a `supported_workload_types: set`. in LocalExecutr it might look like `supported_workload_types = {"ExecuteTask", "ExecuteCallback", "TestConnection"}`
then the entire chunk of code in _get_workloads_to_schedule that sorts out which jobs to pick up would be something real simple like
```python
def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]:
# create a meta_queue which consists of each sub-queue, sorted by that sub-queue's defined sort_key in the order of the scheduling tiers. Then take the top `open_slots` number of jobs off the list.
all_workloads: list[ExecutorWorkload] = [
workload
for queue in self.executor_queues.values()
for workload in queue.values()
]
# Sort by scheduling tier first, then by sort_key within each tier.
all_workloads.sort(key=lambda workload: workload.queue_def)
return [(workload.key, workload) for workload in all_workloads[:open_slots]]
```
As a bonus, this also assures that future workload types honor slots by default; if some future workload type wanted to break that, it would have to be done intentionally. It also means `slots_available` and `slots_occupied` turn into simple comprehensions along the lines of `sum(len(q) for q in self.executor_queues.values())` that don't have to be updated with every new workload. The code overall gets much cleaner.
I'm not sure, I haven't thought that through fully, but as we add more workload types we may need to get that sorted out before long. I have some rough ideas but they were all overkill for only having two types at the time so i figured this was a good place to get those thoughts "on paper". I'll resolve this comment so nobody sees it as blocking, but feel free to message me on Slack and/or tag me in the PR if you get to it.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Maybe the queue_type should also include something to indicate queue priority and whether that queue is intended to be used FIFO or prioritized?
Something vaguely like
If we add a NamedTuple called `WorkloadQueueDef(scheduling_tier: int, sort_key: int)` into the base workload and they are defined like this in the individual workloads:
```python
class TestConnection:
@property
def queue_def(self) -> WorkloadQueueDef :
return WorkloadQueueDef(
scheduling_tier = -1, # connection tests are super fast so let's always get them out of the way first
sort_key = 0, # FIFO order
)
class ExecuteCallback:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
scheduling_tier = 0, # always finish a callback before starting a new task, but don't block the much faster fast connection tests for a callback
sort_key = 0, # FIFO order
)
class ExecuteTask:
@property
def WorkloadQueueDef(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
# after all connection tests and callbacks are running, take the task with the highest priority
scheduling_tier = 1,
sort_key = self.ti.priority_weight,
)
```
and replace all of the "supports_callbacks", "supports_connection_test", etc in each executor can be replaced by a `supported_workload_types: set`. in LocalExecutr it might look like `supported_workload_types = {"ExecuteTask", "ExecuteCallback", "TestConnection"}`
then the entire chunk of code in _get_workloads_to_schedule that sorts out which jobs to pick up would be something real simple like
```python
def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]:
# create a meta_queue which consists of each sub-queue, sorted by that sub-queue's defined sort_key in the order of the scheduling tiers. Then take the top `open_slots` number of jobs off the list.
all_workloads: list[ExecutorWorkload] = [
workload
for queue in self.executor_queues.values()
for workload in queue.values()
]
# Sort by scheduling tier first, then by sort_key within each tier.
all_workloads.sort(key=lambda workload: workload.queue_def)
return [(workload.key, workload) for workload in all_workloads[:open_slots]]
```
As a bonus, this also assures that future workload types honor slots by default; if some future workload type wanted to break that, it would have to be done intentionally. It also means `slots_available` and `slots_occupied` turn into simple comprehensions along the lines of `sum(len(q) for q in self.executor_queues.values())` that don't have to be updated with every new workload. The code overall gets much cleaner.
I'm not sure, I haven't thought that through fully, but as we add more workload types we may need to get that sorted out before long. I have some rough ideas but they were all overkill for only having two types at the time so i figured this was a good place to get those thoughts "on paper". I'll resolve this comment so nobody sees it as blocking, but feel free to message me on Slack and/or tag me in the PR if you get to it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3302011936
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3245,6 +3261,94 @@ def _cleanup_orphaned_asset_state(*, session: Session) -> None:
)
session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+ def _enqueue_connection_tests(self, *, session: Session) -> None:
+ """
+ Enqueue pending connection tests to executors that support them.
+
+ ``max_concurrency`` is per-scheduler, not global: with N HA schedulers
+ the worst-case per-tick dispatch is ``N * max_concurrency``. Connection
+ tests are user-initiated and rare, so the overshoot self-corrects via
+ the reaper. For a true global cap, wrap the budget+claim below in a
+ sentinel-row ``SELECT ... FOR UPDATE``.
+ """
+ max_concurrency = conf.getint("connection_test", "max_concurrency", fallback=4)
+ timeout = conf.getint("connection_test", "timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTestRequest.id)).where(
+ ConnectionTestRequest.state.in_(DISPATCHED_STATES)
+ )
+ )
+ budget = max_concurrency - (active_count or 0)
+ 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)
Review Comment:
Done thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3316041342
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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(String(2000), nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+
+ conn_type: Mapped[str] = mapped_column(String(500), nullable=False)
+ host: Mapped[str | None] = mapped_column(String(500), nullable=True)
+ login: Mapped[str | None] = mapped_column(Text, nullable=True)
+ schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True)
+ port: Mapped[int | None] = mapped_column(Integer, nullable=True)
+ commit_on_success: Mapped[bool] = mapped_column(
+ Boolean, nullable=False, default=False, server_default="0"
+ )
+ is_encrypted: Mapped[bool] = mapped_column(
+ Boolean, unique=False, default=False, nullable=False, server_default="0"
+ )
+ is_extra_encrypted: Mapped[bool] = mapped_column(
+ Boolean, unique=False, default=False, nullable=False, server_default="0"
+ )
+
+ active_connection_id: Mapped[str | None] = mapped_column(String(250), nullable=True)
+
+ __table_args__ = (
+ Index("idx_connection_test_request_state_created_at", state, created_at),
+ UniqueConstraint(
+ "active_connection_id",
+ name="uq_connection_test_request_active_conn",
+ ),
+ )
+
+ def __init__(
+ self,
+ *,
+ connection_id: str,
+ conn_type: str,
+ host: str | None = None,
+ login: str | None = None,
+ password: str | None = None,
+ schema: str | None = None,
+ port: int | None = None,
+ extra: str | None = None,
+ commit_on_success: bool = False,
+ executor: str | None = None,
+ queue: str | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.conn_type = conn_type
+ self.host = host
+ self.login = login
+ self.password = password
+ self.schema = schema
+ self.port = port
+ self.extra = extra
+ self.commit_on_success = commit_on_success
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ @validates("state")
+ def _sync_active_connection_id(self, _key: str, value: str) -> str:
Review Comment:
Updated thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3316039628
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3247,6 +3264,97 @@ def _cleanup_orphaned_asset_state(*, session: Session) -> None:
)
session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+ def _enqueue_connection_tests(self, *, session: Session) -> None:
+ """
+ Enqueue pending connection tests to executors that support them.
+
+ ``max_concurrency`` is per-scheduler, not global: with N HA schedulers
+ the worst-case per-tick dispatch is ``N * max_concurrency``. Connection
+ tests are user-initiated and rare, so the overshoot self-corrects via
+ the reaper. For a true global cap, wrap the budget+claim below in a
+ sentinel-row ``SELECT ... FOR UPDATE``.
+ """
+ max_concurrency = conf.getint("connection_test", "max_concurrency", fallback=4)
+ timeout = conf.getint("connection_test", "timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTestRequest.id)).where(
+ ConnectionTestRequest.state.in_(DISPATCHED_STATES)
+ )
+ )
+ budget = max_concurrency - (active_count or 0)
+ 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:
+ team_name = (
+ Connection.get_team_name(ct.connection_id, session=session) if self._multi_team else None
+ )
+ executor = self._try_to_load_executor(ct, session, team_name=team_name)
+ 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("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_(CONNECTION_TEST_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)"
Review Comment:
Done thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2898741263
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Really appreciate you laying this out, Dennis. Planning to start on this right after #62343 lands. Will tag
you on the PR once I get it out. Thank you so much
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2893499858
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Love this pattern, would simplify things a lot. Excited to try this out as a follow-up!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2970099580
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
Review Comment:
So the implementation approach in this PR means that all (=_ALL_) executors must be touched to implement this feature.
As it is quite generic, would it be possible to somehow make it to some base class that all executors inherit and then not all need to implement this?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971266892
##########
airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml:
##########
@@ -1631,6 +1631,103 @@ paths:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
+ /api/v2/connections/{connection_id}/test:
Review Comment:
updated thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971856706
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
Review Comment:
Tha is sufficient if you have thought about it. Then some small adjustments are probably needed on all executors.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3169346979
##########
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:
The supervisor imports `airflow.models.connection.Connection` solely to call `test_connection()` on it. That method dispatches to the provider hook via `ProvidersManager` (core), which adds a new ORM dependency to task-sdk and is the wrong direction — task-sdk should not be pulling in core ORM models.
The right fix, which belongs in this PR since connection testing is its stated purpose:
**Add `test_connection()` to `airflow.sdk.definitions.connection.Connection`**, dispatching via `ProvidersManagerTaskRuntime` (already in task-sdk for exactly this kind of hook dispatch):
```python
def test_connection(self) -> tuple[bool, str]:
from airflow.sdk.execution_time.context import _preset_connections
from airflow.sdk.providers_manager_runtime import ProvidersManagerTaskRuntime
from airflow.sdk._shared.module_loading import import_string
hook_info = ProvidersManagerTaskRuntime().hooks.get(self.conn_type)
if not hook_info:
return False, f"Unknown conn_type: {self.conn_type!r}"
outer = _preset_connections.get() or {}
token = _preset_connections.set({**outer, self.conn_id: self})
try:
hook = import_string(hook_info.hook_class_name)(conn_id=self.conn_id)
if not hasattr(hook, "test_connection"):
return False, f"Hook {type(hook).__name__} doesn't implement test_connection"
return hook.test_connection()
finally:
_preset_connections.reset(token)
```
`_preset_connections` is the right injection point (checked before all secrets backends, including user-configured custom ones), but it should be encapsulated inside `test_connection()` rather than wired manually in the supervisor.
That reduces the supervisor to:
```python
conn = SDKConnection(conn_id=r.conn_id, conn_type=r.conn_type, ...)
with TimeoutPosix(seconds=timeout, error_message=f"Connection test timed out after {timeout}s"):
success, message = conn.test_connection()
```
Which also fixes the second issue here: `signal.signal(signal.SIGALRM, ...)` + `signal.alarm()` should use `TimeoutPosix` from `airflow.sdk.execution_time.timeout` — it already exists, handles non-POSIX gracefully, and uses `signal.setitimer` (sub-second capable) rather than integer-only `signal.alarm`.
---
Drafted-by: Claude Code (claude-sonnet-4-6); reviewed by @ashb before posting
##########
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)
Review Comment:
`session.add(connection_test)` will raise an `IntegrityError` if there is already an active test for the same `connection_id` (the `uq_connection_test_request_active_conn` unique constraint on `active_connection_id`). The endpoint documents a 409 response but nothing catches the DB error — the caller gets a 500 instead.
##########
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
+ commit_on_success: bool = Field(
Review Comment:
`ConnectionTestRequestBody` duplicates almost all of `ConnectionBody` (same `connection_id`, `conn_type`, `host`, `login`, `schema_`, `port`, `password`, `extra` fields). Use `ConnectionBody` as a component and extend it:
```python
class ConnectionTestRequestBody(ConnectionBody):
commit_on_success: bool = Field(default=False, description="...")
executor: str | None = None
queue: str | None = None
```
`ConnectionBody` already has the `extra` JSON validator, the `connection_id` pattern constraint, and the standard field definitions. Duplicating those creates drift — e.g. `ConnectionBody.connection_id` has `pattern=r"^[\w.-]+$"` but `ConnectionTestRequestBody` drops it.
We should also exclude `commit_on_success` from this schema as it makes no sense on the TaskSDK side. I think we've done something like this elsewhere, there's a pydantic setting/flag we can set in `Field()`
##########
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:
🚨 We cannot introduce new dependencies in TaskSDK to core airflow. We will need to do something else, like remove `test_connection()` from models.connection and move it in to the SDK version.
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -872,6 +874,25 @@ def get_detail_response(self, ti_id: uuid.UUID) -> HITLDetailResponse:
return HITLDetailResponse.model_validate_json(resp.read())
+class ConnectionTestOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def get_connection(self, connection_test_id: uuid.UUID) -> ConnectionResponse:
+ """Fetch connection data for a test request from the API server."""
+ resp = self.client.get(f"connection-tests/{connection_test_id}/connection")
+ return ConnectionResponse.model_validate_json(resp.read())
Review Comment:
Return type is `ConnectionResponse` but the endpoint (`GET /connection-tests/{id}/connection`) returns `ConnectionTestConnectionResponse`. Both models happen to have the same fields today so it works at runtime, but the annotation is wrong and will silently break if the two models ever diverge.
```suggestion
def get_connection(self, connection_test_id: uuid.UUID) -> ConnectionTestConnectionResponse:
"""Fetch connection data for a test request from the API server."""
resp = self.client.get(f"connection-tests/{connection_test_id}/connection")
return ConnectionTestConnectionResponse.model_validate_json(resp.read())
```
---
Drafted-by: Claude Code (claude-sonnet-4-6); reviewed by @ashb before posting
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -144,7 +145,16 @@ def _convert_variable_result_to_variable(var_result: VariableResult, deserialize
return Variable(**var_result.model_dump(exclude={"type"}))
+_preset_connections: ContextVar[dict[str, Connection]] = ContextVar("_preset_connections", default={})
Review Comment:
`default={}` is a single shared mutable object — every context that hasn't called `.set()` gets back the exact same dict from `.get()`. The current code never mutates it, so it doesn't blow up today, but it's a trap for anyone who writes `_preset_connections.get()[conn_id] = conn` instead of using `.set()`.
Change to `default=None` and guard at call sites:
```python
_preset_connections: ContextVar[dict[str, Connection] | None] = ContextVar("_preset_connections", default=None)
# in _get_connection:
preset = _preset_connections.get()
if preset and conn_id in preset:
...
```
Also: `.set({preset.conn_id: preset})` replaces the entire dict, shadowing any outer preset for the duration of the call. If `test_connection()` is ever called from within a task that has its own preset set (e.g. test tooling), the outer entries disappear until `.reset()`. Merge instead of replace:
```python
outer = _preset_connections.get() or {}
token = _preset_connections.set({**outer, self.conn_id: self})
```
---
Drafted-by: Claude Code (claude-sonnet-4-6); reviewed by @ashb before posting
##########
airflow-core/src/airflow/api_fastapi/execution_api/security.py:
##########
@@ -190,6 +190,14 @@ async def require_auth(
detail="Token subject does not match task instance ID",
)
+ if "ct:self" in security_scopes.scopes:
Review Comment:
```suggestion
elsif "ct:self" in security_scopes.scopes:
```
I think? ti:self. and ct:self can't co-exist?
##########
airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml:
##########
@@ -1844,6 +1844,115 @@ paths:
security:
- OAuth2PasswordBearer: []
- HTTPBearer: []
+ /api/v2/connections/test-async:
+ post:
+ tags:
+ - Connection
+ summary: Test Connection Async
+ description: 'Queue an async connection test to be executed on a worker.
Review Comment:
Yes this async, but also no this isn't async 😁
It's not hitting async python code.
I think lets make the URL be `/api/v2/connections/enqueue-test` and just remove "async" from the desc?
##########
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"]),
+ ],
+)
Review Comment:
Two routers seems needless?
```suggestion
router = VersionedAPIRouter(
route_class=ExecutionAPIRoute,
dependencies=[
Security(require_auth, scopes=["ct:self"]),
],
)
```
etc.
##########
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",
Review Comment:
I'm a little bit worried about this endpoint exposing or leaking the password somehow. I wonder if we should allow each ct to only be fetched once?
##########
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
+ ct.result_message = body.result_message
+
+ if body.state == ConnectionTestState.SUCCESS and ct.commit_on_success:
+ ct.commit_to_connection_table(session=session)
+
+ # JWTReissueMiddleware also writes Refreshed-API-Token but skips workload tokens, so we set it here for the workload→execution swap.
+ if token.claims.scope == "workload":
+ generator: JWTGenerator = services.get(JWTGenerator)
+ execution_token = generator.generate(extras={"sub": str(connection_test_id), "scope": "execution"})
+ response.headers["Refreshed-API-Token"] = execution_token
Review Comment:
Why do we need this? Isn't this in practice only ever used when the conn test is finishing, so the token doesn't matter anymore
##########
airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py:
##########
Review Comment:
This version is already released. We can't ret-conn it. You'll need to add a new migration instead.
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2602,6 +2602,33 @@ scheduler:
type: float
example: ~
default: "120.0"
+ connection_test_timeout:
+ description: |
+ Maximum number of seconds an async connection test is allowed to run
+ before it is considered timed out. The scheduler reaper uses this value
+ plus a grace period to mark stale tests as failed.
+ version_added: 3.3.0
+ type: integer
+ example: ~
+ default: "60"
+ max_connection_test_concurrency:
Review Comment:
Ditto.
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2602,6 +2602,33 @@ scheduler:
type: float
example: ~
default: "120.0"
+ connection_test_timeout:
Review Comment:
This does not belong under the `scheduler` section.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -216,6 +218,7 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None)
self.team_name: str | None = team_name
self.queued_tasks: dict[TaskInstanceKey, workloads.ExecuteTask] = {}
self.queued_callbacks: dict[str, workloads.ExecuteCallback] = {}
+ self.queued_connection_tests: dict[ConnectionTestKey, workloads.TestConnection] = {}
Review Comment:
Hmmmmmmm I'm really not liking the growth of this pattern. I wonder if there should be a single `queue` dict?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3292563283
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -259,6 +283,94 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
+ "/enqueue-test",
+ 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 enqueue_connection_test(
Review Comment:
**Cross-team privilege escalation via enqueue.**
The POST endpoint authorizes against the `team_name` in `test_body`, but the existing `Connection` row identified by `connection_id` may belong to a different team. Combined with `commit_on_success=True`, a user with POST rights on team A can submit a test for a team B `connection_id` with new credentials, and on a successful test the row gets overwritten with team A payload.
Resolve the team from the row before authz:
```python
actual_team = Connection.get_team_name(test_body.connection_id, session)
requires_access_connection("POST", team_name=actual_team)
```
and reject if `test_body.team_name` is set and doesn t match.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3302018592
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -259,6 +283,94 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
+ "/enqueue-test",
+ 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 enqueue_connection_test(
Review Comment:
Done thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3029224580
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,234 @@
+# 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 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, Uuid, select
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, synonym
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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))
+
+
+class ConnectionTestRequest(Base):
+ """
+ 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.
+ """
+
Review Comment:
Added thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3029229694
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -491,6 +491,24 @@ core:
type: string
example: ~
default: "Disabled"
+ connection_test_timeout:
+ description: |
+ Maximum number of seconds an async connection test is allowed to run
+ before it is considered timed out. The scheduler reaper uses this value
+ plus a grace period to mark stale tests as failed.
+ version_added: 3.2.0
+ type: integer
+ example: ~
+ default: "60"
+ max_connection_test_concurrency:
+ description: |
+ Maximum number of connection tests that can be active (QUEUED + RUNNING)
+ at the same time. Excess tests will remain in PENDING state until slots
+ become available.
+ version_added: 3.2.0
Review Comment:
moved it to scheduler section
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3029225910
##########
airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py:
##########
@@ -77,12 +78,46 @@ 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
+ commit_on_success: bool = False
Review Comment:
added docs thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3073694070
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,203 @@
+# 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 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, Uuid, select, text
+from sqlalchemy.orm import Mapped, mapped_column
+
+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))
+
+
+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)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+
+ # Connection fields — password and extra are Fernet-encrypted via FernetFieldsMixin.
+ conn_type: Mapped[str] = mapped_column(String(500), nullable=False)
+ host: Mapped[str | None] = mapped_column(String(500), nullable=True)
+ login: Mapped[str | None] = mapped_column(Text, nullable=True)
+ schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True)
+ port: Mapped[int | None] = mapped_column(Integer, nullable=True)
+ commit_on_success: Mapped[bool] = mapped_column(
+ Boolean, nullable=False, default=False, server_default="0"
+ )
+
+ __table_args__ = (
+ Index("idx_connection_test_request_state_created_at", state, created_at),
+ # since mysql lacks filtered/partial indices, this creates a
+ # duplicate index on mysql. Not the end of the world
+ Index(
+ "idx_connection_test_request_active_conn",
+ "connection_id",
+ unique=True,
+ postgresql_where=text("state IN ('pending', 'queued', 'running')"),
+ sqlite_where=text("state IN ('pending', 'queued', 'running')"),
+ ),
Review Comment:
updated thanks
##########
airflow-core/src/airflow/migrations/versions/0111_3_3_0_add_connection_test_table.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.
+
+"""
+Add connection_test_request table for async connection testing.
+
+Revision ID: a7e6d4c3b2f1
+Revises: a4c2d171ae18
+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 = "a4c2d171ae18"
+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("commit_on_success", sa.Boolean(), nullable=False, server_default="0"),
+ sa.PrimaryKeyConstraint("id", name=op.f("connection_test_request_pkey")),
+ sa.UniqueConstraint("token", name=op.f("connection_test_request_token_uq")),
+ )
+ op.create_index(
+ op.f("idx_connection_test_request_state_created_at"),
+ "connection_test_request",
+ ["state", "created_at"],
+ )
+ # since mysql lacks filtered/partial indices, this creates a
+ # duplicate index on mysql
+ op.create_index(
+ op.f("idx_connection_test_request_active_conn"),
+ "connection_test_request",
+ ["connection_id"],
+ unique=True,
+ postgresql_where=sa.text("state IN ('pending', 'queued', 'running')"),
+ sqlite_where=sa.text("state IN ('pending', 'queued', 'running')"),
+ )
+
+
Review Comment:
updated thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2901196801
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Hi @ferruzzi I've pulled out the workload queue generalization you suggested into a separate PR #[New Pr point to this branch](https://github.com/anishgirianish/airflow/pull/1) taking out branch from this pr. Would like to request you for your review on this. I will point it to main once connection testing pr is merged.
Thank you.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4009170592
Thank you so much @ferruzzi, for the review! Really appreciate it. Both future work ideas are great, excited to pick them up as follow-ups after this lands. Will definitely follow up on the supervised process as well, #62645 is a great reference. Thank you!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2932356638
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -174,6 +189,70 @@ def bulk_connections(
return BulkConnectionService(session=session, request=request).handle_request()
+@connections_router.patch(
+ "/{connection_id}/test",
+ responses=create_openapi_http_exception_doc(
+ [
+ status.HTTP_400_BAD_REQUEST,
+ status.HTTP_403_FORBIDDEN,
+ status.HTTP_404_NOT_FOUND,
+ ]
+ ),
+ dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())],
+)
+def patch_connection_and_test(
Review Comment:
I believe this should be `patch_connection_test` only or even removed, it's confusing. A single `async_test_connection` endpoint should be enough. (We can deduce if it's an edit or insert by fetching the db)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971266284
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +331,79 @@ 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_404_NOT_FOUND]),
+ 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 must already be saved. Returns a token that can be used
+ to poll for the test result via GET /connections/test-async/{token}.
+ """
+ _ensure_test_connection_enabled()
+
+ try:
+ Connection.get_connection_from_secrets(test_body.connection_id)
+ except AirflowNotFoundException:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ f"The Connection with connection_id: `{test_body.connection_id}` was not found. "
+ "Connection must be saved before testing.",
+ )
+
+ connection_test = ConnectionTest(
+ connection_id=test_body.connection_id, executor=test_body.executor, queue=test_body.queue
+ )
+ session.add(connection_test)
+ session.flush()
+
+ 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"))],
+)
+def get_connection_test_status(
Review Comment:
updated thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177847962
##########
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:
Done, falsy values now normalize to None. Thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177841400
##########
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:
Done, handler enforces state in TERMINAL_STATES → 409 and state not in ACTIVE_STATES → 409 before mutating. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177839784
##########
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:
Done as part of the previous comment thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177831286
##########
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)
Review Comment:
Done, wrapped session.flush() in try/except IntegrityError → 409 with conn-id-specific message. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177842203
##########
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:
Took your option-3 path: documented N * max_concurrency per-tick ceiling in the docstring, kept max_concurrency as per-scheduler. Sentinel-row breadcrumb in the docstring if you'd prefer the global cap later. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177845258
##########
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:
Done implicitly via the ConnectionTestRequestBody(ConnectionBody) inheritance, validate_extra is now applied at the API boundary. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177830554
##########
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
+ commit_on_success: bool = Field(
Review Comment:
Done, ConnectionTestRequestBody now extends ConnectionBody; only adds commit_on_success/executor/queue. Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177837496
##########
airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py:
##########
Review Comment:
Done, moved AddConnectionTestEndpoint into a new v2026_05_02.py; left v2026_04_06.py alone. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3029232752
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +286,84 @@ 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]),
+ 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()
+
+ # Only one active test per connection_id at a time.
+ _check_no_active_test(test_body.connection_id, session)
+
+ 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)
+ session.flush()
Review Comment:
removed thanks
##########
airflow-ctl/src/airflowctl/api/datamodels/generated.py:
##########
@@ -245,15 +245,58 @@ class ConnectionResponse(BaseModel):
team_name: Annotated[str | None, Field(title="Team Name")] = None
+class ConnectionTestQueuedResponse(BaseModel):
+ """
+ Response returned when an async connection test is queued.
+ """
+
+ token: Annotated[str, Field(title="Token")]
+ connection_id: Annotated[str, Field(title="Connection Id")]
+ state: Annotated[str, Field(title="State")]
+
+
+class ConnectionTestRequestBody(BaseModel):
+ """
+ Request body for async connection test.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+ connection_id: Annotated[str, Field(title="Connection Id")]
+ conn_type: Annotated[str, Field(title="Conn Type")]
+ host: Annotated[str | None, Field(title="Host")] = None
+ login: Annotated[str | None, Field(title="Login")] = None
+ schema_: Annotated[str | None, Field(alias="schema", title="Schema")] = None
+ port: Annotated[int | None, Field(title="Port")] = None
+ password: Annotated[str | None, Field(title="Password")] = None
+ extra: Annotated[str | None, Field(title="Extra")] = None
+ commit_on_success: Annotated[bool | None, Field(title="Commit On Success")] = False
+ executor: Annotated[str | None, Field(title="Executor")] = None
+ queue: Annotated[str | None, Field(title="Queue")] = None
+
+
class ConnectionTestResponse(BaseModel):
"""
- Connection Test serializer for responses.
+ Connection Test serializer for synchronous test responses.
"""
status: Annotated[bool, Field(title="Status")]
message: Annotated[str, Field(title="Message")]
+class ConnectionTestStatusResponse(BaseModel):
Review Comment:
updated thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3128342837
##########
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)
Review Comment:
Thanks for raising. I see SIGALRM is already used unconditionally in a few other places (task-sdk/.../timeout.py, utils/db.py... so I followed the existing pattern here. Happy to change if you'd prefer. Please let me know
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3292565580
##########
task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py:
##########
@@ -0,0 +1,87 @@
+# 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 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.exceptions import AirflowTaskTimeout
+from airflow.sdk.execution_time.timeout import TimeoutPosix
+
+__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."""
+ client = Client(base_url=server, token=token)
+
+ try:
+ r = client.connection_tests.get_connection(connection_test_id)
+
+ conn = SDKConnection(
+ conn_id=r.conn_id,
+ conn_type=r.conn_type,
+ host=r.host,
+ login=r.login,
+ password=r.password,
+ schema=r.schema_,
+ port=r.port,
+ extra=r.extra,
+ )
+ with TimeoutPosix(
Review Comment:
**Follow-up on the prior SIGALRM/Windows concern.**
Renaming to `TimeoutPosix` makes the platform constraint explicit, thanks. But the import of `signal.SIGALRM` inside `timeout.py` is still unconditional, so a Windows worker process raises `AttributeError` at import time before it can even decide whether to call the timeout. A `hasattr(signal, "SIGALRM")` gate at the import site, or a `threading.Timer` fallback, closes the loop.
If Windows is not a supported worker platform for this feature, please document that on `supports_connection_test` (and ideally fail fast at supervisor start rather than at the first connection test) so users understand the boundary.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-3960392203
@pierrejeambrun The dev list discussion ([thread](https://lists.apache.org/thread/xd7zmyp95y77cw36mb5wjp17dyynz100)) reached consensus — Jarek and Jens both supported the approach. The implementation follows the agreed design (store connection_id only, crypto-random token for polling, dedicated table with timestamp cleanup). Is there anything specific from the discussion you'd like to see addressed?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2932383533
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +331,79 @@ 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_404_NOT_FOUND]),
+ 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 must already be saved. Returns a token that can be used
+ to poll for the test result via GET /connections/test-async/{token}.
+ """
+ _ensure_test_connection_enabled()
+
+ try:
+ Connection.get_connection_from_secrets(test_body.connection_id)
+ except AirflowNotFoundException:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ f"The Connection with connection_id: `{test_body.connection_id}` was not found. "
+ "Connection must be saved before testing.",
+ )
+
+ connection_test = ConnectionTest(
+ connection_id=test_body.connection_id, executor=test_body.executor, queue=test_body.queue
+ )
+ session.add(connection_test)
+ session.flush()
+
+ 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"))],
+)
+def get_connection_test_status(
Review Comment:
get_connection_test
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2963468979
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -174,6 +189,70 @@ def bulk_connections(
return BulkConnectionService(session=session, request=request).handle_request()
+@connections_router.patch(
+ "/{connection_id}/test",
+ responses=create_openapi_http_exception_doc(
+ [
+ status.HTTP_400_BAD_REQUEST,
+ status.HTTP_403_FORBIDDEN,
+ status.HTTP_404_NOT_FOUND,
+ ]
+ ),
+ dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())],
+)
+def patch_connection_and_test(
+ connection_id: str,
+ patch_body: ConnectionBody,
+ session: SessionDep,
+ update_mask: list[str] | None = Query(None),
+ executor: str | None = Query(None, description="Executor to route the connection test to"),
+ queue: str | None = Query(None, description="Queue to route the connection test to"),
+) -> ConnectionSaveAndTestResponse:
+ """
+ Update a connection and queue an async test with revert-on-failure.
+
+ Atomically saves the edit and creates a ConnectionTest with snapshots of the
+ pre-edit and post-edit state. If the test fails, the connection is automatically
+ reverted to its pre-edit values.
+ """
+ _ensure_test_connection_enabled()
+
+ if patch_body.connection_id != connection_id:
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST,
+ "The connection_id in the request body does not match the URL parameter",
+ )
+
+ connection = session.scalar(select(Connection).filter_by(conn_id=connection_id).limit(1))
+ if connection is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ f"The Connection with connection_id: `{connection_id}` was not found",
+ )
+
+ try:
+ ConnectionBody(**patch_body.model_dump())
Review Comment:
`patch_body` is already a `ConnectionBody` instance — FastAPI validated it on parameter injection. Re-constructing `ConnectionBody(**patch_body.model_dump())` just validates the same data again. This block is dead code.
##########
airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py:
##########
@@ -0,0 +1,71 @@
+#
+# 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 table for async connection testing.
+
+Revision ID: a7e6d4c3b2f1
+Revises: 1d6611b6ab7c
+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 = "1d6611b6ab7c"
+branch_labels = None
+depends_on = None
+airflow_version = "3.2.0"
+
+
+def upgrade():
+ """Create connection_test table."""
+ op.create_table(
+ "connection_test",
+ 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(10), 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),
Review Comment:
Longest current value is `"pending"` (7 chars). `String(10)` leaves almost no room — `"timed_out"` (9) or `"cancelled"` (9) would need a new migration. Use `String(20)` for breathing room.
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -168,6 +177,84 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
+ """
+ Execute a connection test workload.
+
+ Constructs an SDK ``Client``, fetches the connection via the Execution API,
+ enforces a timeout via ``signal.alarm``, and reports all outcomes back
+ through the Execution API.
+
+ :param log: Logger instance
+ :param workload: The TestConnection workload to execute
+ :param team_conf: Team-specific executor configuration
+ """
+ # Lazy import: SDK modules must not be loaded at module level to avoid
+ # coupling core (scheduler-loaded) code to the SDK.
+ from airflow.sdk.api.client import Client
+ from airflow.sdk.execution_time.comms import ErrorResponse
+
+ setproctitle(
+ f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}",
+ log,
+ )
+
+ base_url = team_conf.get("api", "base_url", fallback="/")
+ if base_url.startswith("/"):
+ base_url = f"http://localhost:8080{base_url}"
+ default_execution_api_server = f"{base_url.rstrip('/')}/execution/"
+ server = team_conf.get("core", "execution_api_server_url", fallback=default_execution_api_server)
+
+ client = Client(base_url=server, token=workload.token)
+
+ def _handle_timeout(signum, frame):
+ raise TimeoutError(f"Connection test timed out after {workload.timeout}s")
+
+ signal.signal(signal.SIGALRM, _handle_timeout)
+ signal.alarm(workload.timeout)
+ try:
+ client.connection_tests.update_state(workload.connection_test_id, ConnectionTestState.RUNNING)
+
+ conn_response = client.connections.get(workload.connection_id)
+ if isinstance(conn_response, ErrorResponse):
+ raise RuntimeError(f"Connection '{workload.connection_id}' not found via Execution API")
+
+ conn = Connection(
+ conn_id=conn_response.conn_id,
+ conn_type=conn_response.conn_type,
+ host=conn_response.host,
+ login=conn_response.login,
+ password=conn_response.password,
+ schema=conn_response.schema_,
+ port=conn_response.port,
+ extra=conn_response.extra,
+ )
+ success, message = run_connection_test(conn=conn)
+
+ state = ConnectionTestState.SUCCESS if success else ConnectionTestState.FAILED
+ client.connection_tests.update_state(workload.connection_test_id, state, message)
+ except TimeoutError:
+ log.error(
+ "Connection test timed out after %ds",
+ workload.timeout,
+ connection_id=workload.connection_id,
+ )
+ client.connection_tests.update_state(
+ workload.connection_test_id,
+ ConnectionTestState.FAILED,
+ f"Connection test timed out after {workload.timeout}s",
+ )
+ except Exception as e:
+ log.exception("Connection test failed unexpectedly", connection_id=workload.connection_id)
+ client.connection_tests.update_state(
+ workload.connection_test_id,
+ ConnectionTestState.FAILED,
Review Comment:
`f"Connection test failed unexpectedly: {e}"[:500]` stores raw exception text. Some DB drivers include connection strings (with credentials) in error messages. This gets stored in `result_message` and exposed via the polling API.
Consider sanitizing or using a generic message with the exception type only:
```python
f"Connection test failed unexpectedly: {type(e).__name__}"
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971259042
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
Review Comment:
updated thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971262692
##########
airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py:
##########
@@ -0,0 +1,71 @@
+#
+# 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 table for async connection testing.
+
+Revision ID: a7e6d4c3b2f1
+Revises: 1d6611b6ab7c
+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 = "1d6611b6ab7c"
+branch_labels = None
+depends_on = None
+airflow_version = "3.2.0"
+
+
+def upgrade():
+ """Create connection_test table."""
+ op.create_table(
+ "connection_test",
+ 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(10), 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),
Review Comment:
updated thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2897389432
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Maybe the queue_type should also include something to indicate queue priority and whether that queue is intended to be used FIFO or prioritized?
Something vaguely like
If we add a NamedTuple called `WorkloadQueueDef(scheduling_tier: int, sort_key: int)` into the base workload and they are defined like this in the individual workloads:
```python
class TestConnection:
@property
def queue_def(self) -> WorkloadQueueDef :
return WorkloadQueueDef(
scheduling_tier = -1, # connection tests are super fast so let's always get them out of the way first
sort_key = 0, # FIFO order
)
class ExecuteCallback:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
scheduling_tier = 0, # always finish a callback before starting a new task, but don't block the much faster fast connection tests for a callback
sort_key = 0, # FIFO order
)
class ExecuteTask:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
# after all connection tests and callbacks are running, take the task with the highest priority
scheduling_tier = 1,
sort_key = self.ti.priority_weight,
)
```
((IN fact, we can set them to both default to 0 and that gets even more simplified.... but you get the idea))
and replace all of the "supports_callbacks", "supports_connection_test", etc in each executor can be replaced by a `supported_workload_types: set`. in LocalExecutr it might look like `supported_workload_types = {"ExecuteTask", "ExecuteCallback", "TestConnection"}`
then the entire chunk of code in _get_workloads_to_schedule that sorts out which jobs to pick up would be something real simple like
```python
def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]:
# create a meta_queue which consists of each sub-queue, sorted by that sub-queue's defined sort_key in the order of the scheduling tiers. Then take the top `open_slots` number of jobs off the list.
all_workloads: list[ExecutorWorkload] = [
workload
for queue in self.executor_queues.values()
for workload in queue.values()
]
# Sort by scheduling tier first, then by sort_key within each tier.
all_workloads.sort(key=lambda workload: workload.queue_def)
return [(workload.key, workload) for workload in all_workloads[:open_slots]]
```
As a bonus, this also assures that future workload types honor slots by default; if some future workload type wanted to break that, it would have to be done intentionally. It also means `slots_available` and `slots_occupied` turn into simple comprehensions along the lines of `sum(len(q) for q in self.executor_queues.values())` that don't have to be updated with every new workload. The code overall gets much cleaner.
I'm not sure, I haven't thought that through fully, but as we add more workload types we may need to get that sorted out before long. I have some rough ideas but they were all overkill for only having two types at the time so i figured this was a good place to get those thoughts "on paper". Feel free to resolve this comment once you have seen it so nobody sees it as blocking, but feel free to message me on Slack and/or tag me in the PR if you get to it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2897389432
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Maybe the queue_type should also include something to indicate queue priority and whether that queue is intended to be used FIFO or prioritized?
Something vaguely like
If we add a named tuple called `WorkloadQueueDef(scheduling_tier: int, sort_key: int)` into the base workload and they are defined like this in the individual workloads:
```python
class TestConnection:
@property
def WorkloadQueueDef(scheduling_tier: int, sort_key: int):
return (
scheduling_tier = -1 # connection tests are super fast so let's always get them out of the way first
sort_key = 0 # FIFO order
)
class ExecuteCallback:
@property
def WorkloadQueueDef(scheduling_tier: int, sort_key: int):
return (
scheduling_tier = 0 # always finish a callback before starting a new task, but don't block the much faster fast connection tests for a callback
sort_key = 0 # FIFO order
)
class ExecuteTask:
@property
def WorkloadQueueDef(scheduling_tier: int, sort_key: int):
return (
# after all connection tests and callbacks are running, take the task with the highest priority
scheduling_tier = 1
sort_key = self.ti.priority_weight
)
```
and replace all of the "supports_callbacks", "supports_connection_test", etc in each executor can be replaced by a `supported_workload_types: set`. in LocalExecutr it might look like `supported_workload_types = {"ExecuteTask", "ExecuteCallback", "TestConnection"}`
then the entire chunk of code in _get_workloads_to_schedule that sorts out which jobs to pick up would be something real simple like
```python
def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]:
# create a meta_queue which consists of each sub-queue, sorted by that sub-queue's defined sort_key in the order of the scheduling tiers. Then take the top `open_slots` number of jobs off the list.
all_workloads: list[ExecutorWorkload] = [
workload
for queue in self.executor_queues.values()
for workload in queue.values()
]
# Sort by scheduling tier first, then by sort_key within each tier.
all_workloads.sort(key=lambda workload: workload.queue_def)
return [(workload.key, workload) for workload in all_workloads[:open_slots]]
```
As a bonus, this also assures that future workload types honor slots by default; if some future workload type wanted to break that, it would have to be done intentionally. It also means `slots_available` and `slots_occupied` turn into simple comprehensions along the lines of `sum(len(q) for q in self.executor_queues.values())` that don't have to be updated with every new workload. The code overall gets much cleaner.
I'm not sure, I haven't thought that through fully, but as we add more workload types we may need to get that sorted out before long. I have some rough ideas but they were all overkill for only having two types at the time so i figured this was a good place to get those thoughts "on paper". I'll resolve this comment so nobody sees it as blocking, but feel free to message me on Slack and/or tag me in the PR if you get to it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2893084458
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -240,10 +242,18 @@ def queue_workload(self, workload: workloads.All, session: Session) -> None:
f"See LocalExecutor or CeleryExecutor for reference implementation."
)
self.queued_callbacks[workload.callback.id] = workload
+ elif isinstance(workload, workloads.TestConnection):
+ if not self.supports_connection_test:
+ raise NotImplementedError(
+ f"{type(self).__name__} does not support TestConnection workloads. "
+ f"Set supports_connection_test = True and implement connection test handling "
+ f"in _process_workloads(). See LocalExecutor for reference implementation."
+ )
+ self.queued_connection_tests[str(workload.connection_test_id)] = workload
else:
raise ValueError(
f"Un-handled workload type {type(workload).__name__!r} in {type(self).__name__}. "
- f"Workload must be one of: ExecuteTask, ExecuteCallback."
+ f"Workload must be one of: ExecuteTask, ExecuteCallback, TestConnection."
Review Comment:
Non-blocking future work unless you are already making more changes here. We'll eventually want to generate this list by unpacking ExecutorWorkload or something instead of manually adding new workload types I guess.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Non-blocking future work idea: I considered this when i was in here adding the callback workload but thought it wasn't worth the refactor at the time; maybe we're getting close to that time. I wonder if we should consider (again, not now) an alternative data object for the queues.
If we did something like `executor_queue: dict[workload_type, workload_queue_type]` and add the `workload_queue_type` field and type hint to each workload the same as the `type` field, then we could do `executor_queue[workload.type].update(new_entry)` and these slots_available and slots_occupied get simplified down to iterating the executor_queue and summing the lengths.
Anyway, like I said, absolutely not blocking but I wanted to get the idea out since it was starting to look relevant.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2970076291
##########
airflow-core/newsfragments/62343.feature.rst:
##########
Review Comment:
I am not sure if we really need a newsfragment for this. For 90% of users this is no difference and it is not a breaking change?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3073697229
##########
airflow-core/tests/unit/models/test_connection_test.py:
##########
@@ -0,0 +1,228 @@
+# 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 unittest import mock
+
+import pytest
+
+from airflow.models.connection import Connection
+from airflow.models.connection_test import (
+ ConnectionTestRequest,
+ ConnectionTestState,
+ run_connection_test,
+)
+
+from tests_common.test_utils.db import clear_db_connection_tests, clear_db_connections
+
+pytestmark = pytest.mark.db_test
+
+
+class TestConnectionTestRequestModel:
+ def test_token_is_generated(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.token is not None
+ assert len(ct.token) > 0
+
+ def test_initial_state_is_pending(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.state == ConnectionTestState.PENDING
+
+ def test_tokens_are_unique(self):
+ ct1 = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ ct2 = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct1.token != ct2.token
+
+ def test_repr(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ r = repr(ct)
+ assert "test_conn" in r
+ assert "pending" in r
+
+ def test_executor_parameter(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", executor="my_executor")
+ assert ct.executor == "my_executor"
+
+ def test_executor_defaults_to_none(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.executor is None
+
+ def test_queue_parameter(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", queue="my_queue")
+ assert ct.queue == "my_queue"
+
+ def test_queue_defaults_to_none(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.queue is None
+
+ def test_connection_fields_stored(self):
+ ct = ConnectionTestRequest(
+ connection_id="test_conn",
+ conn_type="postgres",
+ host="db.example.com",
+ login="user",
+ password="secret",
+ schema="mydb",
+ port=5432,
+ extra='{"key": "value"}',
+ )
+ assert ct.conn_type == "postgres"
+ assert ct.host == "db.example.com"
+ assert ct.login == "user"
+ assert ct.password == "secret"
+ assert ct.schema == "mydb"
+ assert ct.port == 5432
+ assert ct.extra == '{"key": "value"}'
+
+ def test_password_is_encrypted(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", password="secret")
+ assert ct._password is not None
+ assert ct._password != "secret"
+ assert ct.password == "secret"
+
+ def test_extra_is_encrypted(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", extra='{"key": "val"}')
+ assert ct._extra is not None
+ assert ct._extra != '{"key": "val"}'
+ assert ct.extra == '{"key": "val"}'
+
+ def test_null_password_and_extra(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="http")
+ assert ct._password is None
+ assert ct._extra is None
+
+ def test_commit_on_success_default(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.commit_on_success is False
+
+ def test_commit_on_success_true(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", commit_on_success=True)
+ assert ct.commit_on_success is True
+
+
+class TestToConnection:
+ def test_to_connection_returns_transient_connection(self):
+ ct = ConnectionTestRequest(
+ connection_id="test_conn",
+ conn_type="postgres",
+ host="db.example.com",
+ login="user",
+ password="secret",
+ schema="mydb",
+ port=5432,
+ extra='{"key": "value"}',
+ )
+ conn = ct.to_connection()
+ assert isinstance(conn, Connection)
+ assert conn.conn_id == "test_conn"
+ assert conn.conn_type == "postgres"
+ assert conn.host == "db.example.com"
+ assert conn.login == "user"
+ assert conn.password == "secret"
+ assert conn.schema == "mydb"
+ assert conn.port == 5432
+ assert conn.extra == '{"key": "value"}'
+
+
+class TestCommitToConnectionTable:
+ @pytest.fixture(autouse=True)
+ def setup_teardown(self):
+ clear_db_connections(add_default_connections_back=False)
+ clear_db_connection_tests()
+ yield
+ clear_db_connections(add_default_connections_back=False)
+ clear_db_connection_tests()
+
+ def test_creates_new_connection(self, session):
+ ct = ConnectionTestRequest(
+ connection_id="new_conn",
+ conn_type="postgres",
+ host="db.example.com",
+ login="user",
+ password="secret",
+ schema="mydb",
+ port=5432,
+ )
+ session.add(ct)
+ session.flush()
+
+ ct.commit_to_connection_table(session=session)
+ session.flush()
+
+ from sqlalchemy import select
Review Comment:
updated thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3307402411
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -259,6 +321,78 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
+ "/enqueue-test",
+ 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(action_logging())],
+)
+def enqueue_connection_test(
+ test_body: ConnectionTestRequestBody,
+ session: SessionDep,
+ user: GetUserDep,
+) -> ConnectionTestQueuedResponse:
+ """Enqueue a connection test for deferred execution on a worker; returns a polling token."""
+ _ensure_test_connection_enabled()
+ _ensure_executor_is_configured(test_body.executor)
+
+ existing = session.scalar(select(Connection).filter_by(conn_id=test_body.connection_id))
+ if existing is not None:
+ effective_team = existing.team_name
+ if test_body.team_name is not None and test_body.team_name != effective_team:
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ f"team_name `{test_body.team_name}` does not match the team of connection "
+ f"`{test_body.connection_id}`.",
+ )
+ else:
+ effective_team = test_body.team_name
Review Comment:
`effective_team` from `test_body.team_name` is used to authorize POST when there's no existing connection row (the `commit_on_success=True` flow for a brand-new connection), but it's never persisted on `ConnectionTestRequest`. The polling endpoint at line 141 then re-resolves team via `Connection.get_team_name(connection_test.connection_id)`, which returns `None` because the connection still doesn't exist, and authz collapses to the global "read connections" right. The scheduler dispatch at `scheduler_job_runner.py:3302-3304` does the same lookup and routes the workload to the global executor instead of the team-scoped one.
Cross-team isolation is broken end-to-end for the new-connection path in multi-team mode. Add a `team_name` column to `connection_test_request`, persist `effective_team` here, and read from the row in both the poll endpoint and the scheduler dispatch path.
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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(String(2000), nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+
+ conn_type: Mapped[str] = mapped_column(String(500), nullable=False)
+ host: Mapped[str | None] = mapped_column(String(500), nullable=True)
+ login: Mapped[str | None] = mapped_column(Text, nullable=True)
+ schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True)
+ port: Mapped[int | None] = mapped_column(Integer, nullable=True)
+ commit_on_success: Mapped[bool] = mapped_column(
+ Boolean, nullable=False, default=False, server_default="0"
+ )
+ is_encrypted: Mapped[bool] = mapped_column(
Review Comment:
`is_encrypted` / `is_extra_encrypted` flags are inherited from `FernetFieldsMixin`, but `airflow rotate-fernet-key` (`cli/commands/rotate_fernet_key_command.py`) only walks `Connection`, `Variable`, `Trigger`. After rotation, any in-flight or db-cleanup-retained `connection_test_request` row with `is_encrypted=True` becomes undecryptable.
Either add a `rotate_fernet_key()` method here and include it in the CLI walk, or document that connection-test rows must be drained before rotation.
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -1047,6 +1051,28 @@ def get_detail_response(self, ti_id: uuid.UUID) -> HITLDetailResponse:
return HITLDetailResponse.model_validate_json(resp.read())
+class ConnectionTestOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def get_connection(self, connection_test_id: uuid.UUID) -> ConnectionTestConnectionResponse:
+ """Fetch connection data for a test request from the API server."""
+ resp = self.client.get(f"connection-tests/{connection_test_id}/connection")
+ return ConnectionTestConnectionResponse.model_validate_json(resp.read())
+
+ def update_state(
+ self, id: uuid.UUID, state: ConnectionTestState, result_message: str | None = None
+ ) -> None:
+ """Report the state of a connection test to the API server."""
+ body = ConnectionTestResultBody(
+ state=state,
+ result_message=ResultMessage(result_message) if result_message is not None else None,
Review Comment:
`ResultMessage(result_message)` runs Pydantic validation with `max_length=2000`. In the happy path at `connection_test_supervisor.py:87`, `message` comes from `hook.test_connection()` and is user-controlled / unbounded. A 3000-char message raises `ValidationError`, the supervisor exits without PATCHing, and the row sits in RUNNING until the reaper labels it "timed out" with no signal that the test actually completed.
Truncate inside `update_state` before constructing `ResultMessage` (e.g. `result_message[:2000]`).
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3247,6 +3264,97 @@ def _cleanup_orphaned_asset_state(*, session: Session) -> None:
)
session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+ def _enqueue_connection_tests(self, *, session: Session) -> None:
+ """
+ Enqueue pending connection tests to executors that support them.
+
+ ``max_concurrency`` is per-scheduler, not global: with N HA schedulers
+ the worst-case per-tick dispatch is ``N * max_concurrency``. Connection
+ tests are user-initiated and rare, so the overshoot self-corrects via
+ the reaper. For a true global cap, wrap the budget+claim below in a
+ sentinel-row ``SELECT ... FOR UPDATE``.
+ """
+ max_concurrency = conf.getint("connection_test", "max_concurrency", fallback=4)
+ timeout = conf.getint("connection_test", "timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTestRequest.id)).where(
+ ConnectionTestRequest.state.in_(DISPATCHED_STATES)
+ )
+ )
+ budget = max_concurrency - (active_count or 0)
+ 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:
+ team_name = (
+ Connection.get_team_name(ct.connection_id, session=session) if self._multi_team else None
+ )
+ executor = self._try_to_load_executor(ct, session, team_name=team_name)
+ 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("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_(CONNECTION_TEST_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)"
Review Comment:
`CONNECTION_TEST_ACTIVE_STATES` covers PENDING/QUEUED/RUNNING and the message is the same for all three. A PENDING row reaped here was never dispatched (no executor picked it up before the cutoff); a RUNNING row exceeded its timeout. Same `result_message="timed out"` is misleading for the former.
Branch on `ct.state` before composing the message.
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2699,6 +2699,42 @@ scheduler:
type: boolean
example: ~
default: "False"
+connection_test:
+ description: |
+ Configuration for the deferred connection-test workflow that dispatches
+ test requests to workers via the executor (instead of running them
+ in-process on the API server).
+ options:
+ timeout:
+ description: |
+ Maximum number of seconds a worker-dispatched connection test is
+ allowed to run before it is considered timed out. The scheduler
+ reaper uses this value plus a grace period to mark stale tests as
+ failed.
+ version_added: 3.3.0
+ type: integer
+ example: ~
+ default: "60"
+ max_concurrency:
+ description: |
+ Maximum number of connection tests that can be active
+ (QUEUED + RUNNING) at the same time. Excess tests will remain in
+ PENDING state until slots become available. With N HA schedulers and
+ the dispatch lock seeded by the migration, this is enforced as a
Review Comment:
This description promises "with N HA schedulers and the dispatch lock seeded by the migration, this is enforced as a true global cap". The migration only creates `connection_test_request`; there's no sentinel / dispatch-lock row, and `scheduler_job_runner.py:3271-3275` explicitly documents the opposite ("per-scheduler, not global; worst-case `N * max_concurrency`").
Either align this description with the scheduler docstring, or implement the sentinel-row `SELECT ... FOR UPDATE` the docstring suggests.
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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(String(2000), nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+
+ conn_type: Mapped[str] = mapped_column(String(500), nullable=False)
+ host: Mapped[str | None] = mapped_column(String(500), nullable=True)
+ login: Mapped[str | None] = mapped_column(Text, nullable=True)
+ schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True)
+ port: Mapped[int | None] = mapped_column(Integer, nullable=True)
+ commit_on_success: Mapped[bool] = mapped_column(
+ Boolean, nullable=False, default=False, server_default="0"
+ )
+ is_encrypted: Mapped[bool] = mapped_column(
+ Boolean, unique=False, default=False, nullable=False, server_default="0"
+ )
+ is_extra_encrypted: Mapped[bool] = mapped_column(
+ Boolean, unique=False, default=False, nullable=False, server_default="0"
+ )
+
+ active_connection_id: Mapped[str | None] = mapped_column(String(250), nullable=True)
+
+ __table_args__ = (
+ Index("idx_connection_test_request_state_created_at", state, created_at),
+ UniqueConstraint(
+ "active_connection_id",
+ name="uq_connection_test_request_active_conn",
+ ),
+ )
+
+ def __init__(
+ self,
+ *,
+ connection_id: str,
+ conn_type: str,
+ host: str | None = None,
+ login: str | None = None,
+ password: str | None = None,
+ schema: str | None = None,
+ port: int | None = None,
+ extra: str | None = None,
+ commit_on_success: bool = False,
+ executor: str | None = None,
+ queue: str | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.conn_type = conn_type
+ self.host = host
+ self.login = login
+ self.password = password
+ self.schema = schema
+ self.port = port
+ self.extra = extra
+ self.commit_on_success = commit_on_success
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ @validates("state")
+ def _sync_active_connection_id(self, _key: str, value: str) -> str:
Review Comment:
`@validates("state")` runs on every `state` assignment, including the enum assignments in the route (`ct.state = body.state`) and the supervisor flow. Annotation should be `str | ConnectionTestState` (or just the enum) so callers / IDE / type checkers see the real contract.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3302010522
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -259,6 +283,94 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
+ "/enqueue-test",
+ 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 enqueue_connection_test(
+ test_body: ConnectionTestRequestBody,
+ session: SessionDep,
+) -> ConnectionTestQueuedResponse:
+ """Enqueue a connection test for deferred execution on a worker; returns a polling 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)
+ try:
+ session.flush()
+ except IntegrityError:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ f"An active connection test already exists for connection_id `{test_body.connection_id}`.",
+ )
+
+ return ConnectionTestQueuedResponse(
+ token=connection_test.token,
+ connection_id=connection_test.connection_id,
+ state=connection_test.state,
+ )
+
+
+@connections_router.get(
+ "/enqueue-test/{connection_test_token}",
Review Comment:
Went with option1. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3302014853
##########
airflow-core/src/airflow/models/crypto.py:
##########
@@ -93,6 +96,65 @@ 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()
+ self.is_encrypted = fernet.is_encrypted
+ else:
Review Comment:
Done. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177849992
##########
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:
Added, test_get_status_unauthorized_user_does_not_leak_row thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177832693
##########
airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml:
##########
@@ -1844,6 +1844,115 @@ paths:
security:
- OAuth2PasswordBearer: []
- HTTPBearer: []
+ /api/v2/connections/test-async:
+ post:
+ tags:
+ - Connection
+ summary: Test Connection Async
+ description: 'Queue an async connection test to be executed on a worker.
Review Comment:
Done, renamed to /api/v2/connections/enqueue-test, dropped "async" from descriptions. Thank you :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3041344867
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -169,6 +178,81 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
+ """
+ Execute a connection test workload.
+
+ Constructs an SDK ``Client``, fetches the connection via the Execution API,
+ enforces a timeout via ``signal.alarm``, and reports all outcomes back
+ through the Execution API.
+
+ :param log: Logger instance
+ :param workload: The TestConnection workload to execute
+ :param team_conf: Team-specific executor configuration
+ """
+ # Lazy import: SDK modules must not be loaded at module level to avoid
+ # coupling core (scheduler-loaded) code to the SDK.
+ from airflow.sdk.api.client import Client
+
+ setproctitle(
+ f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}",
+ log,
+ )
+
+ base_url = team_conf.get("api", "base_url", fallback="/")
+ if base_url.startswith("/"):
+ base_url = f"http://localhost:8080{base_url}"
Review Comment:
Yeah, for better or worse, that's existing code [here](https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/executors/local_executor.py#L136-L140) which I am moving into a new helper named `local_executor::_get_execution_api_server_url()` in https://github.com/apache/airflow/pull/62645 ([here](https://github.com/apache/airflow/pull/62645/changes#diff-e6c8f3d778a4f3366b23d8d46a1c9650608ccbf2184e7292f0f7bbf4df37dbebR53-R64))
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3128342837
##########
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)
Review Comment:
Thanks for raising. I see SIGALRM is already used unconditionally in a few other places (task-sdk/.../timeout.py, utils/db.py... so I followed the existing pattern here. Happy to change if you'd prefer. Please let me know. Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2891424498
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1759,6 +1760,15 @@ def _run_scheduler_loop(self) -> None:
action=bundle_cleanup_mgr.remove_stale_bundle_versions,
)
+ timers.call_regular_interval(
+ delay=conf.getfloat("scheduler", "connection_test_dispatch_interval", fallback=2.0),
+ action=self._dispatch_connection_tests,
Review Comment:
This ultimately schedules workloads to the executors, I believe this should be in the scheduler loop below.
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py:
##########
@@ -1140,6 +1148,275 @@ 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)."""
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_post_should_respond_202(self, test_client, session):
+ """POST /connections/test-async with a saved connection returns 202 + token."""
+ self.create_connection()
+ response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID})
+ 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={"connection_id": TEST_CONN_ID}
+ )
+ assert response.status_code == 401
+
+ def test_should_respond_403(self, unauthorized_test_client):
+ response = unauthorized_test_client.post(
+ "/connections/test-async", json={"connection_id": TEST_CONN_ID}
+ )
+ 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={"connection_id": TEST_CONN_ID})
+ assert response.status_code == 403
+ assert response.json() == {
+ "detail": "Testing connections is disabled in Airflow configuration. "
+ "Contact your deployment admin to enable it."
+ }
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_should_respond_404_for_nonexistent_connection(self, test_client):
+ """Connection must be saved before testing."""
+ response = test_client.post("/connections/test-async", json={"connection_id": "nonexistent"})
+ assert response.status_code == 404
+ assert "was not found" in response.json()["detail"]
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_post_creates_connection_test_row(self, test_client, session):
+ """POST creates a ConnectionTest row in PENDING state."""
+ self.create_connection()
+ response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID})
+ assert response.status_code == 202
+ token = response.json()["token"]
+
+ ct = session.scalar(select(ConnectionTest).filter_by(token=token))
+ assert ct is not None
+ assert ct.connection_id == TEST_CONN_ID
+ 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 to the ConnectionTest."""
+ self.create_connection()
+ response = test_client.post(
+ "/connections/test-async",
+ json={"connection_id": TEST_CONN_ID, "queue": "gpu_workers"},
+ )
+ assert response.status_code == 202
+ token = response.json()["token"]
+
+ ct = session.scalar(select(ConnectionTest).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_get_status_returns_pending(self, test_client, session):
+ """GET /connections/test-async/{token} returns current status (pending before scheduler dispatch)."""
+ self.create_connection()
+ post_response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID})
+ 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
+
+ @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 ConnectionTest."""
+ self.create_connection()
+ post_response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID})
+ token = post_response.json()["token"]
+
+ ct = session.scalar(select(ConnectionTest).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):
+ """GET with an unknown token returns 404."""
+ response = test_client.get("/connections/test-async/nonexistent-token")
+ assert response.status_code == 404
+
+
+class TestSaveAndTest(TestConnectionEndpoint):
+ """Tests for the combined PATCH /{connection_id}/test endpoint."""
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_save_and_test_returns_200_with_token(self, test_client, session):
+ """PATCH /{connection_id}/test updates the connection and returns a test token."""
+ self.create_connection()
+ response = test_client.patch(
+ f"/connections/{TEST_CONN_ID}/test",
+ json={
+ "connection_id": TEST_CONN_ID,
+ "conn_type": TEST_CONN_TYPE,
+ "host": "updated-host.example.com",
+ },
+ )
+ assert response.status_code == 200
+ body = response.json()
+ assert body["test_token"]
+ assert body["test_state"] == "pending"
+ assert body["connection"]["host"] == "updated-host.example.com"
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_save_and_test_creates_snapshot(self, test_client, session):
+ """PATCH /{connection_id}/test creates a ConnectionTest with a connection_snapshot."""
+ self.create_connection()
+ response = test_client.patch(
+ f"/connections/{TEST_CONN_ID}/test",
+ json={
+ "connection_id": TEST_CONN_ID,
+ "conn_type": TEST_CONN_TYPE,
+ "host": "new-host.example.com",
+ },
+ )
+ assert response.status_code == 200
+ token = response.json()["test_token"]
+
+ ct = session.scalar(select(ConnectionTest).filter_by(token=token))
+ assert ct is not None
+ assert ct.connection_snapshot is not None
+ snapshot = ct.connection_snapshot
+ assert "pre" in snapshot
+ assert "post" in snapshot
+ assert snapshot["pre"]["host"] == TEST_CONN_HOST
+ assert snapshot["post"]["host"] == "new-host.example.com"
+
+ @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
+ def test_save_and_test_passes_executor_parameter(self, test_client, session):
+ """PATCH /{connection_id}/test passes the executor parameter to the ConnectionTest."""
+ self.create_connection()
+ response = test_client.patch(
+ f"/connections/{TEST_CONN_ID}/test?executor=team_a",
Review Comment:
Why is the executor being set to a team name?
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3243,6 +3253,100 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ @provide_session
+ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None:
+ """Dispatch pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("core", "max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("core", "connection_test_timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTest.id)).where(ConnectionTest.state.in_(ACTIVE_STATES))
+ )
+ budget = max_concurrency - (active_count or 0)
+ if budget <= 0:
+ return
+
+ pending_stmt = (
+ select(ConnectionTest)
+ .where(ConnectionTest.state == ConnectionTestState.PENDING)
+ .order_by(ConnectionTest.created_at)
+ .limit(budget)
+ )
+ pending_stmt = with_row_locks(pending_stmt, session, of=ConnectionTest, skip_locked=True)
+ pending_tests = session.scalars(pending_stmt).all()
+
+ if not pending_tests:
+ return
+
+ for ct in pending_tests:
+ executor = self._find_executor_for_connection_test(ct.executor)
+ if executor is None:
+ reason = (
+ f"No executor matches '{ct.executor}'"
+ if ct.executor
+ else "No executor supports 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("core", "connection_test_timeout", fallback=60)
+ grace_period = max(30, timeout // 2)
+ cutoff = timezone.utcnow() - timedelta(seconds=timeout + grace_period)
+
+ stale_stmt = select(ConnectionTest).where(
+ ConnectionTest.state.in_(ACTIVE_STATES),
+ ConnectionTest.updated_at < cutoff,
+ )
+ stale_stmt = with_row_locks(stale_stmt, session, of=ConnectionTest, 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)
+ if ct.connection_snapshot:
+ attempt_revert(ct, session=session)
+
+ session.flush()
+
+ def _find_executor_for_connection_test(self, executor_name: str | None) -> BaseExecutor | None:
Review Comment:
This executor loading duplicates existing code. It already has the problem that happens when you duplicate code, it's drifted from the other implementation that all other workload use, it's missing team support and it's missing a case executors match on, class name.
Please use `_try_to_load_executor()` that accepts a workload (which your new test connection workload type should be a sub type of, if not, please update that). This is the single canonical location for executor lookups in the scheduler job.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3243,6 +3253,100 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ @provide_session
+ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None:
+ """Dispatch pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("core", "max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("core", "connection_test_timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTest.id)).where(ConnectionTest.state.in_(ACTIVE_STATES))
+ )
+ budget = max_concurrency - (active_count or 0)
Review Comment:
This must take into account available executor slots. It is the responsibility of the scheduler to ensure we don't over subscribe all executors across all teams.
Please look at the other workload types for reference `_enqueue_executor_callbacks` and `_critical_section_enqueue_task_instances`. These both contain logic to assess the number of available executor slots and to only submit/queue as many workloads as there are slots available.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2963468969
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via a TestConnection workload."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True)
+ reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0")
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(
+ self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return f"<ConnectionTest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+
+def run_connection_test(*, conn: Connection) -> tuple[bool, str]:
+ """
+ Worker-side function to execute a connection test.
+
+ Returns a (success, message) tuple. The caller is responsible for
+ reporting the result back via the Execution API.
+ """
+ try:
+ return conn.test_connection()
+ except Exception as e:
+ log.exception("Connection test failed", connection_id=conn.conn_id)
+ return False, str(e)
+
+
+_SNAPSHOT_FIELDS = (
+ "conn_type",
+ "description",
+ "host",
+ "login",
+ "_password",
+ "schema",
+ "port",
+ "_extra",
+ "is_encrypted",
+ "is_extra_encrypted",
+)
+
+
+def snapshot_connection(conn: Connection) -> dict:
+ """
+ Capture raw DB column values from a Connection for later restore.
+
+ Encrypted fields (``_password``, ``_extra``) are stored as ciphertext
+ so they can be written directly back without re-encryption.
+ """
+ return {field: getattr(conn, field) for field in _SNAPSHOT_FIELDS}
+
+
+def _revert_connection(conn: Connection, snapshot: dict) -> None:
+ """
+ Restore a Connection's columns from a snapshot dict.
+
+ Writes directly to ``_password`` and ``_extra`` (bypassing the
+ encrypting property setters) so the stored ciphertext is preserved.
+ """
+ for field, value in snapshot.items():
+ setattr(conn, field, value)
+
+
+def _decrypt_snapshot_field(snapshot: dict, field: str) -> str | None:
+ """Decrypt a single encrypted field from a snapshot dict using Fernet."""
+ raw = snapshot.get(field)
+ if raw is None:
+ return None
+ encrypted_flag = "is_encrypted" if field == "_password" else "is_extra_encrypted"
+ if not snapshot.get(encrypted_flag, False):
+ return raw
+ fernet = get_fernet()
+ return fernet.decrypt(bytes(raw, "utf-8")).decode()
+
+
+def _can_safely_revert(conn: Connection, post_snapshot: dict) -> bool:
+ """
+ Check whether the connection's current state matches the post-edit snapshot.
+
+ Compares **decrypted** values for encrypted fields and direct values for
+ non-encrypted fields. Returns ``False`` if any field differs, indicating
+ a concurrent edit has occurred and the revert should be skipped.
+ """
+ for field in _SNAPSHOT_FIELDS:
+ if field in ("is_encrypted", "is_extra_encrypted"):
+ continue
+
+ if field == "_password":
+ current_val = conn.password
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_password")
+ elif field == "_extra":
+ current_val = conn.extra
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_extra")
+ else:
+ current_val = getattr(conn, field)
+ snapshot_val = post_snapshot.get(field)
+
+ if current_val != snapshot_val:
+ return False
+ return True
+
+
+def attempt_revert(ct: ConnectionTest, *, session: Session) -> None:
+ """Revert a connection to its pre-edit values if no concurrent edit has occurred."""
+ if not ct.connection_snapshot:
+ log.warning("attempt_revert called without snapshot", connection_test_id=ct.id)
+ return
+
+ pre_snapshot = ct.connection_snapshot["pre"]
+ post_snapshot = ct.connection_snapshot["post"]
+
+ ct.connection_snapshot = None
+
+ conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id))
+ if conn is None:
+ ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists."
+ log.warning("Revert skipped: connection no longer exists", connection_id=ct.connection_id)
+ return
+
+ if not _can_safely_revert(conn, post_snapshot):
+ ct.result_message = (
+ ct.result_message or ""
Review Comment:
This `select(Connection)` reads without `with_for_update()`. A concurrent edit can land between the `_can_safely_revert` check and `_revert_connection`, causing the revert to silently overwrite the concurrent change — the exact scenario `_can_safely_revert` is meant to prevent.
```python
conn = session.scalar(
select(Connection).filter_by(conn_id=ct.connection_id).with_for_update()
)
```
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via a TestConnection workload."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True)
+ reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0")
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(
+ self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return f"<ConnectionTest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+
+def run_connection_test(*, conn: Connection) -> tuple[bool, str]:
+ """
+ Worker-side function to execute a connection test.
+
+ Returns a (success, message) tuple. The caller is responsible for
+ reporting the result back via the Execution API.
+ """
+ try:
+ return conn.test_connection()
+ except Exception as e:
+ log.exception("Connection test failed", connection_id=conn.conn_id)
+ return False, str(e)
+
+
+_SNAPSHOT_FIELDS = (
+ "conn_type",
+ "description",
+ "host",
+ "login",
+ "_password",
+ "schema",
+ "port",
+ "_extra",
+ "is_encrypted",
+ "is_extra_encrypted",
+)
+
+
+def snapshot_connection(conn: Connection) -> dict:
+ """
+ Capture raw DB column values from a Connection for later restore.
+
+ Encrypted fields (``_password``, ``_extra``) are stored as ciphertext
+ so they can be written directly back without re-encryption.
+ """
+ return {field: getattr(conn, field) for field in _SNAPSHOT_FIELDS}
+
+
+def _revert_connection(conn: Connection, snapshot: dict) -> None:
+ """
+ Restore a Connection's columns from a snapshot dict.
+
+ Writes directly to ``_password`` and ``_extra`` (bypassing the
+ encrypting property setters) so the stored ciphertext is preserved.
+ """
+ for field, value in snapshot.items():
+ setattr(conn, field, value)
+
+
+def _decrypt_snapshot_field(snapshot: dict, field: str) -> str | None:
+ """Decrypt a single encrypted field from a snapshot dict using Fernet."""
+ raw = snapshot.get(field)
+ if raw is None:
+ return None
+ encrypted_flag = "is_encrypted" if field == "_password" else "is_extra_encrypted"
+ if not snapshot.get(encrypted_flag, False):
+ return raw
+ fernet = get_fernet()
+ return fernet.decrypt(bytes(raw, "utf-8")).decode()
+
+
+def _can_safely_revert(conn: Connection, post_snapshot: dict) -> bool:
+ """
+ Check whether the connection's current state matches the post-edit snapshot.
+
+ Compares **decrypted** values for encrypted fields and direct values for
+ non-encrypted fields. Returns ``False`` if any field differs, indicating
+ a concurrent edit has occurred and the revert should be skipped.
+ """
+ for field in _SNAPSHOT_FIELDS:
+ if field in ("is_encrypted", "is_extra_encrypted"):
+ continue
+
+ if field == "_password":
+ current_val = conn.password
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_password")
+ elif field == "_extra":
+ current_val = conn.extra
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_extra")
+ else:
+ current_val = getattr(conn, field)
+ snapshot_val = post_snapshot.get(field)
+
+ if current_val != snapshot_val:
+ return False
+ return True
+
+
+def attempt_revert(ct: ConnectionTest, *, session: Session) -> None:
+ """Revert a connection to its pre-edit values if no concurrent edit has occurred."""
+ if not ct.connection_snapshot:
+ log.warning("attempt_revert called without snapshot", connection_test_id=ct.id)
+ return
+
+ pre_snapshot = ct.connection_snapshot["pre"]
+ post_snapshot = ct.connection_snapshot["post"]
+
+ ct.connection_snapshot = None
+
+ conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id))
+ if conn is None:
+ ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists."
+ log.warning("Revert skipped: connection no longer exists", connection_id=ct.connection_id)
+ return
+
+ if not _can_safely_revert(conn, post_snapshot):
Review Comment:
`ct.connection_snapshot = None` is set here, before the actual revert logic runs. If the subsequent `select(Connection)` query or `_revert_connection` call fails (DB hiccup, connection dropped), the snapshot is lost and the revert can never be retried.
Move `ct.connection_snapshot = None` to the end of the function, after the revert succeeds or is definitively skipped.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -305,10 +315,19 @@ def heartbeat(self) -> None:
self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads)
self.trigger_tasks(open_slots)
+ self.trigger_connection_tests()
+
# Calling child class sync method
self.log.debug("Calling the %s sync method", self.__class__)
self.sync()
+ def trigger_connection_tests(self) -> None:
+ """Process queued connection tests."""
+ if not self.supports_connection_test or not self.queued_connection_tests:
+ return
+
+ self._process_workloads(list(self.queued_connection_tests.values()))
Review Comment:
This is called from `heartbeat()` after `trigger_tasks(open_slots)` has already consumed slots. `trigger_connection_tests` then dispatches *all* queued connection tests via `_process_workloads` without checking remaining capacity.
The scheduler's `_enqueue_connection_tests` does slot accounting, but slots may have been consumed between enqueue time and executor heartbeat. With `max_connection_test_concurrency=4` the overshoot is bounded, but it'd be cleaner to pass remaining slots through or check `self.slots_available` here.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2967728694
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -168,6 +177,84 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
+ """
+ Execute a connection test workload.
+
+ Constructs an SDK ``Client``, fetches the connection via the Execution API,
+ enforces a timeout via ``signal.alarm``, and reports all outcomes back
+ through the Execution API.
+
+ :param log: Logger instance
+ :param workload: The TestConnection workload to execute
+ :param team_conf: Team-specific executor configuration
+ """
+ # Lazy import: SDK modules must not be loaded at module level to avoid
+ # coupling core (scheduler-loaded) code to the SDK.
+ from airflow.sdk.api.client import Client
+ from airflow.sdk.execution_time.comms import ErrorResponse
+
+ setproctitle(
+ f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}",
+ log,
+ )
+
+ base_url = team_conf.get("api", "base_url", fallback="/")
+ if base_url.startswith("/"):
+ base_url = f"http://localhost:8080{base_url}"
+ default_execution_api_server = f"{base_url.rstrip('/')}/execution/"
+ server = team_conf.get("core", "execution_api_server_url", fallback=default_execution_api_server)
+
+ client = Client(base_url=server, token=workload.token)
+
+ def _handle_timeout(signum, frame):
+ raise TimeoutError(f"Connection test timed out after {workload.timeout}s")
+
+ signal.signal(signal.SIGALRM, _handle_timeout)
+ signal.alarm(workload.timeout)
+ try:
+ client.connection_tests.update_state(workload.connection_test_id, ConnectionTestState.RUNNING)
+
+ conn_response = client.connections.get(workload.connection_id)
+ if isinstance(conn_response, ErrorResponse):
+ raise RuntimeError(f"Connection '{workload.connection_id}' not found via Execution API")
+
+ conn = Connection(
+ conn_id=conn_response.conn_id,
+ conn_type=conn_response.conn_type,
+ host=conn_response.host,
+ login=conn_response.login,
+ password=conn_response.password,
+ schema=conn_response.schema_,
+ port=conn_response.port,
+ extra=conn_response.extra,
+ )
+ success, message = run_connection_test(conn=conn)
+
+ state = ConnectionTestState.SUCCESS if success else ConnectionTestState.FAILED
+ client.connection_tests.update_state(workload.connection_test_id, state, message)
+ except TimeoutError:
+ log.error(
+ "Connection test timed out after %ds",
+ workload.timeout,
+ connection_id=workload.connection_id,
+ )
+ client.connection_tests.update_state(
+ workload.connection_test_id,
+ ConnectionTestState.FAILED,
+ f"Connection test timed out after {workload.timeout}s",
+ )
+ except Exception as e:
+ log.exception("Connection test failed unexpectedly", connection_id=workload.connection_id)
+ client.connection_tests.update_state(
+ workload.connection_test_id,
+ ConnectionTestState.FAILED,
Review Comment:
This one looks alright to me, he's embedding the connection id, not the error message.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971838499
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
Review Comment:
> LocalExecutor is "one" of the executors means it can serve as template for other executors. But including K8s, Celery, Edge, AWS at least 4 additional exist in the repo. So the logic would need to be added to each.
>
> Technically - w/o looking into the code - I see 2 technical options:
>
> * Implement it in the base class (BaseExecutor) and all individual implementations can inherit
> * Implement via a MixIn, then all Executors "just" need to use it... still all executor classes need to be adjusted
>
> Had only time to skim over the PR and have no answer, else I'd directly propose an alternative. But I worry a bit that complexity grows with this approach in each executor while I see in devlist additional features like https://lists.apache.org/thread/o0z8v01v9qq26r6qmvx8zwbkmho1fnbg might need to take the same route.
Hi @jscheffl thank you very much the feedback. I just realzied, this was discussed with (@ferruzzi @jason810496 ) a bit before in this very pr disscussion as well and was actually planned as fast follow up to be using the mechanism following the pattern in this pr https://github.com/apache/airflow/pull/62645 to make it more generic.
Please let me know your thoughts on this. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971264843
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -174,6 +189,70 @@ def bulk_connections(
return BulkConnectionService(session=session, request=request).handle_request()
+@connections_router.patch(
+ "/{connection_id}/test",
+ responses=create_openapi_http_exception_doc(
+ [
+ status.HTTP_400_BAD_REQUEST,
+ status.HTTP_403_FORBIDDEN,
+ status.HTTP_404_NOT_FOUND,
+ ]
+ ),
+ dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())],
+)
+def patch_connection_and_test(
+ connection_id: str,
+ patch_body: ConnectionBody,
+ session: SessionDep,
+ update_mask: list[str] | None = Query(None),
+ executor: str | None = Query(None, description="Executor to route the connection test to"),
+ queue: str | None = Query(None, description="Queue to route the connection test to"),
+) -> ConnectionSaveAndTestResponse:
+ """
+ Update a connection and queue an async test with revert-on-failure.
+
+ Atomically saves the edit and creates a ConnectionTest with snapshots of the
+ pre-edit and post-edit state. If the test fails, the connection is automatically
+ reverted to its pre-edit values.
+ """
+ _ensure_test_connection_enabled()
+
+ if patch_body.connection_id != connection_id:
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST,
+ "The connection_id in the request body does not match the URL parameter",
+ )
+
+ connection = session.scalar(select(Connection).filter_by(conn_id=connection_id).limit(1))
+ if connection is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ f"The Connection with connection_id: `{connection_id}` was not found",
+ )
+
+ try:
+ ConnectionBody(**patch_body.model_dump())
Review Comment:
Addressed, this entire endpoint is removed in the redesign.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971272801
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +286,83 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
Review Comment:
LocalExecutor is "one" of the executors means it can serve as template for other executors. But including K8s, Celery, Edge, AWS at least 4 additional exist in the repo. So the logic would need to be added to each.
Technically - w/o looking into the code - I see 2 technical options:
- Implement it in the base class (BaseExecutor) and all individual implementations can inherit
- Implement via a MixIn, then all Executors "just" need to use it... still all executor classes need to be adjusted
Had only time to skim over the PR and have no answer, else I'd directly propose an alternative. But I worry a bit that complexity grows with this approach in each executor while I see in devlist additional features like https://lists.apache.org/thread/o0z8v01v9qq26r6qmvx8zwbkmho1fnbg might need to take the same route.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3025675838
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -867,6 +869,25 @@ def get_detail_response(self, ti_id: uuid.UUID) -> HITLDetailResponse:
return HITLDetailResponse.model_validate_json(resp.read())
+class ConnectionTestOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def get_connection(self, connection_test_id: uuid.UUID) -> ConnectionResponse:
+ """Fetch connection data for a test request from the API server."""
+ resp = self.client.get(f"connection-tests/{connection_test_id}/connection")
+ return ConnectionResponse.model_validate_json(resp.read())
+
Review Comment:
The fields are identical today so this works correctly. The semantic distinction is minor can be addressed in a follow-up if the types diverge.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3029221476
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -491,6 +491,24 @@ core:
type: string
example: ~
default: "Disabled"
+ connection_test_timeout:
+ description: |
+ Maximum number of seconds an async connection test is allowed to run
+ before it is considered timed out. The scheduler reaper uses this value
+ plus a grace period to mark stale tests as failed.
+ version_added: 3.2.0
+ type: integer
+ example: ~
+ default: "60"
+ max_connection_test_concurrency:
+ description: |
+ Maximum number of connection tests that can be active (QUEUED + RUNNING)
+ at the same time. Excess tests will remain in PENDING state until slots
+ become available.
+ version_added: 3.2.0
Review Comment:
Thank you. Updated
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3029221476
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -491,6 +491,24 @@ core:
type: string
example: ~
default: "Disabled"
+ connection_test_timeout:
+ description: |
+ Maximum number of seconds an async connection test is allowed to run
+ before it is considered timed out. The scheduler reaper uses this value
+ plus a grace period to mark stale tests as failed.
+ version_added: 3.2.0
+ type: integer
+ example: ~
+ default: "60"
+ max_connection_test_concurrency:
+ description: |
+ Maximum number of connection tests that can be active (QUEUED + RUNNING)
+ at the same time. Excess tests will remain in PENDING state until slots
+ become available.
+ version_added: 3.2.0
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4339976836
> Feel free to resolve comments you have addressed so we know if works remain to be done before the next review.
Hi @pierrejeambrun thank you so much for the review. I have resolved feedback in the latest push, would like to request you for your re-review whenever you get a chance. Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177830554
##########
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
+ commit_on_success: bool = Field(
Review Comment:
Done, ConnectionTestRequestBody now extends ConnectionBody; only adds commit_on_success/executor/queue
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177849132
##########
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:
Switched to TimeoutPosix; documented the POSIX requirement on BaseExecutor.supports_connection_test.
##########
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:
Switched to TimeoutPosix; documented the POSIX requirement on BaseExecutor.supports_connection_test. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177835903
##########
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"]),
+ ],
+)
Review Comment:
Done,collapsed to a single router; scope tightened to ["ct:self", "token:workload"]. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177831794
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -872,6 +874,25 @@ def get_detail_response(self, ti_id: uuid.UUID) -> HITLDetailResponse:
return HITLDetailResponse.model_validate_json(resp.read())
+class ConnectionTestOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def get_connection(self, connection_test_id: uuid.UUID) -> ConnectionResponse:
+ """Fetch connection data for a test request from the API server."""
+ resp = self.client.get(f"connection-tests/{connection_test_id}/connection")
+ return ConnectionResponse.model_validate_json(resp.read())
Review Comment:
Done, return type is now ConnectionTestConnectionResponse. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3302016094
##########
task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py:
##########
@@ -0,0 +1,87 @@
+# 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 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.exceptions import AirflowTaskTimeout
+from airflow.sdk.execution_time.timeout import TimeoutPosix
+
+__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."""
+ client = Client(base_url=server, token=token)
+
+ try:
+ r = client.connection_tests.get_connection(connection_test_id)
+
+ conn = SDKConnection(
+ conn_id=r.conn_id,
+ conn_type=r.conn_type,
+ host=r.host,
+ login=r.login,
+ password=r.password,
+ schema=r.schema_,
+ port=r.port,
+ extra=r.extra,
+ )
+ with TimeoutPosix(
Review Comment:
Updated thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3028331645
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,234 @@
+# 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 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, Uuid, select
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, synonym
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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))
+
+
+class ConnectionTestRequest(Base):
+ """
+ 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.
+ """
+
Review Comment:
Can you add a db composite unique constraint on `connection_id` where state is in `('pending', 'queued', 'running')`.
To make sure that there cannot be two ConnectionTestRequest for the same ID.
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -491,6 +491,24 @@ core:
type: string
example: ~
default: "Disabled"
+ connection_test_timeout:
+ description: |
+ Maximum number of seconds an async connection test is allowed to run
+ before it is considered timed out. The scheduler reaper uses this value
+ plus a grace period to mark stale tests as failed.
+ version_added: 3.2.0
+ type: integer
+ example: ~
+ default: "60"
+ max_connection_test_concurrency:
+ description: |
+ Maximum number of connection tests that can be active (QUEUED + RUNNING)
+ at the same time. Excess tests will remain in PENDING state until slots
+ become available.
+ version_added: 3.2.0
Review Comment:
version_added isn't right.
##########
airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py:
##########
@@ -77,12 +78,46 @@ 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
+ commit_on_success: bool = False
Review Comment:
We should probably add a documentation on what `commit_on_success` is doing. So the generated openapi spec mention that piece of documentation.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +286,84 @@ 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]),
+ 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()
+
+ # Only one active test per connection_id at a time.
+ _check_no_active_test(test_body.connection_id, session)
+
+ 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)
+ session.flush()
Review Comment:
I don't think we need that manual flush.
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -169,6 +178,81 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
+ """
+ Execute a connection test workload.
+
+ Constructs an SDK ``Client``, fetches the connection via the Execution API,
+ enforces a timeout via ``signal.alarm``, and reports all outcomes back
+ through the Execution API.
+
+ :param log: Logger instance
+ :param workload: The TestConnection workload to execute
+ :param team_conf: Team-specific executor configuration
+ """
+ # Lazy import: SDK modules must not be loaded at module level to avoid
+ # coupling core (scheduler-loaded) code to the SDK.
+ from airflow.sdk.api.client import Client
+
+ setproctitle(
+ f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}",
+ log,
+ )
+
+ base_url = team_conf.get("api", "base_url", fallback="/")
+ if base_url.startswith("/"):
+ base_url = f"http://localhost:8080{base_url}"
Review Comment:
not sure that hardcoding `http://localhost:8080` makes sense.
##########
airflow-ctl/src/airflowctl/api/datamodels/generated.py:
##########
@@ -245,15 +245,58 @@ class ConnectionResponse(BaseModel):
team_name: Annotated[str | None, Field(title="Team Name")] = None
+class ConnectionTestQueuedResponse(BaseModel):
+ """
+ Response returned when an async connection test is queued.
+ """
+
+ token: Annotated[str, Field(title="Token")]
+ connection_id: Annotated[str, Field(title="Connection Id")]
+ state: Annotated[str, Field(title="State")]
+
+
+class ConnectionTestRequestBody(BaseModel):
+ """
+ Request body for async connection test.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+ connection_id: Annotated[str, Field(title="Connection Id")]
+ conn_type: Annotated[str, Field(title="Conn Type")]
+ host: Annotated[str | None, Field(title="Host")] = None
+ login: Annotated[str | None, Field(title="Login")] = None
+ schema_: Annotated[str | None, Field(alias="schema", title="Schema")] = None
+ port: Annotated[int | None, Field(title="Port")] = None
+ password: Annotated[str | None, Field(title="Password")] = None
+ extra: Annotated[str | None, Field(title="Extra")] = None
+ commit_on_success: Annotated[bool | None, Field(title="Commit On Success")] = False
+ executor: Annotated[str | None, Field(title="Executor")] = None
+ queue: Annotated[str | None, Field(title="Queue")] = None
+
+
class ConnectionTestResponse(BaseModel):
"""
- Connection Test serializer for responses.
+ Connection Test serializer for synchronous test responses.
"""
status: Annotated[bool, Field(title="Status")]
message: Annotated[str, Field(title="Message")]
+class ConnectionTestStatusResponse(BaseModel):
Review Comment:
```suggestion
class AsyncConnectionTestResponse(BaseModel):
```
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -491,6 +491,24 @@ core:
type: string
example: ~
default: "Disabled"
+ connection_test_timeout:
+ description: |
+ Maximum number of seconds an async connection test is allowed to run
+ before it is considered timed out. The scheduler reaper uses this value
+ plus a grace period to mark stale tests as failed.
+ version_added: 3.2.0
+ type: integer
+ example: ~
+ default: "60"
+ max_connection_test_concurrency:
+ description: |
+ Maximum number of connection tests that can be active (QUEUED + RUNNING)
+ at the same time. Excess tests will remain in PENDING state until slots
+ become available.
+ version_added: 3.2.0
Review Comment:
Config keys should be consistently in [scheduler] section
`connection_test_timeout` and `max_connection_test_concurrency` are in [core] but `connection_test_reaper_interval` is in [scheduler]. All three seems only consumed by the scheduler. Moving them all to [scheduler] would be consistent with how other scheduler-only settings are organized.
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,234 @@
+# 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 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, Uuid, select
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, synonym
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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))
+
+
+class ConnectionTestRequest(Base):
+ """
+ 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)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+
+ # Connection fields — password and extra are Fernet-encrypted.
+ conn_type: Mapped[str] = mapped_column(String(500), nullable=False)
+ host: Mapped[str | None] = mapped_column(String(500), nullable=True)
+ login: Mapped[str | None] = mapped_column(Text, nullable=True)
+ _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True)
+ schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True)
+ port: Mapped[int | None] = mapped_column(Integer, nullable=True)
+ _extra: Mapped[str | None] = mapped_column("extra", Text(), nullable=True)
+ commit_on_success: Mapped[bool] = mapped_column(
+ Boolean, nullable=False, default=False, server_default="0"
+ )
+
+ __table_args__ = (Index("idx_connection_test_request_state_created_at", state, created_at),)
+
+ def __init__(
+ self,
+ *,
+ connection_id: str,
+ conn_type: str,
+ host: str | None = None,
+ login: str | None = None,
+ password: str | None = None,
+ schema: str | None = None,
+ port: int | None = None,
+ extra: str | None = None,
+ commit_on_success: bool = False,
+ executor: str | None = None,
+ queue: str | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.conn_type = conn_type
+ self.host = host
+ self.login = login
+ self.password = password
+ self.schema = schema
+ self.port = port
+ self.extra = extra
+ self.commit_on_success = commit_on_success
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return (
+ f"<ConnectionTestRequest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+ )
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+ def get_password(self) -> str | None:
+ if self._password:
+ fernet = get_fernet()
+ if not fernet.is_encrypted:
+ return self._password
+ return fernet.decrypt(bytes(self._password, "utf-8")).decode()
+ return self._password
+
+ def set_password(self, value: str | None):
+ if value:
+ fernet = get_fernet()
+ self._password = fernet.encrypt(bytes(value, "utf-8")).decode()
+ else:
+ self._password = value
+
+ @declared_attr
+ def password(cls):
+ """Password. The value is decrypted/encrypted when reading/setting the value."""
+ return synonym("_password", descriptor=property(cls.get_password, cls.set_password))
+
+ def get_extra(self) -> str | None:
+ if self._extra:
+ fernet = get_fernet()
+ if not fernet.is_encrypted:
+ return self._extra
+ return fernet.decrypt(bytes(self._extra, "utf-8")).decode()
+ return self._extra
+
+ def set_extra(self, value: str | None):
+ if value:
+ fernet = get_fernet()
+ self._extra = fernet.encrypt(bytes(value, "utf-8")).decode()
+ else:
+ self._extra = value
+
+ @declared_attr
+ def extra(cls):
+ """Extra data. The value is decrypted/encrypted when reading/setting the value."""
+ return synonym("_extra", descriptor=property(cls.get_extra, cls.set_extra))
+
Review Comment:
I think all that is duplicated from `Connection` model. Can we consider factorizing this into a common mixin (ConnectionMixin, FernetFieldMixin) or shared utility.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3025670882
##########
airflow-core/src/airflow/utils/db.py:
##########
@@ -115,7 +115,7 @@ class MappedClassProtocol(Protocol):
"3.0.3": "fe199e1abd77",
"3.1.0": "cc92b33c6709",
"3.1.8": "509b94a1042d",
- "3.2.0": "1d6611b6ab7c",
+ "3.2.0": "a7e6d4c3b2f1",
"3.3.0": "a4c2d171ae18",
Review Comment:
Done updated
##########
airflow-core/src/airflow/migrations/versions/0111_3_2_0_add_connection_test_table.py:
##########
@@ -0,0 +1,77 @@
+#
+# 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: a4c2d171ae18
+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 = "a4c2d171ae18"
+branch_labels = None
+depends_on = None
+airflow_version = "3.2.0"
Review Comment:
updated
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3066488476
##########
airflow-core/tests/unit/models/test_connection_test.py:
##########
@@ -0,0 +1,228 @@
+# 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 unittest import mock
+
+import pytest
+
+from airflow.models.connection import Connection
+from airflow.models.connection_test import (
+ ConnectionTestRequest,
+ ConnectionTestState,
+ run_connection_test,
+)
+
+from tests_common.test_utils.db import clear_db_connection_tests, clear_db_connections
+
+pytestmark = pytest.mark.db_test
+
+
+class TestConnectionTestRequestModel:
+ def test_token_is_generated(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.token is not None
+ assert len(ct.token) > 0
+
+ def test_initial_state_is_pending(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.state == ConnectionTestState.PENDING
+
+ def test_tokens_are_unique(self):
+ ct1 = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ ct2 = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct1.token != ct2.token
+
+ def test_repr(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ r = repr(ct)
+ assert "test_conn" in r
+ assert "pending" in r
+
+ def test_executor_parameter(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", executor="my_executor")
+ assert ct.executor == "my_executor"
+
+ def test_executor_defaults_to_none(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.executor is None
+
+ def test_queue_parameter(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", queue="my_queue")
+ assert ct.queue == "my_queue"
+
+ def test_queue_defaults_to_none(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.queue is None
+
+ def test_connection_fields_stored(self):
+ ct = ConnectionTestRequest(
+ connection_id="test_conn",
+ conn_type="postgres",
+ host="db.example.com",
+ login="user",
+ password="secret",
+ schema="mydb",
+ port=5432,
+ extra='{"key": "value"}',
+ )
+ assert ct.conn_type == "postgres"
+ assert ct.host == "db.example.com"
+ assert ct.login == "user"
+ assert ct.password == "secret"
+ assert ct.schema == "mydb"
+ assert ct.port == 5432
+ assert ct.extra == '{"key": "value"}'
+
+ def test_password_is_encrypted(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", password="secret")
+ assert ct._password is not None
+ assert ct._password != "secret"
+ assert ct.password == "secret"
+
+ def test_extra_is_encrypted(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", extra='{"key": "val"}')
+ assert ct._extra is not None
+ assert ct._extra != '{"key": "val"}'
+ assert ct.extra == '{"key": "val"}'
+
+ def test_null_password_and_extra(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="http")
+ assert ct._password is None
+ assert ct._extra is None
+
+ def test_commit_on_success_default(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres")
+ assert ct.commit_on_success is False
+
+ def test_commit_on_success_true(self):
+ ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", commit_on_success=True)
+ assert ct.commit_on_success is True
+
+
+class TestToConnection:
+ def test_to_connection_returns_transient_connection(self):
+ ct = ConnectionTestRequest(
+ connection_id="test_conn",
+ conn_type="postgres",
+ host="db.example.com",
+ login="user",
+ password="secret",
+ schema="mydb",
+ port=5432,
+ extra='{"key": "value"}',
+ )
+ conn = ct.to_connection()
+ assert isinstance(conn, Connection)
+ assert conn.conn_id == "test_conn"
+ assert conn.conn_type == "postgres"
+ assert conn.host == "db.example.com"
+ assert conn.login == "user"
+ assert conn.password == "secret"
+ assert conn.schema == "mydb"
+ assert conn.port == 5432
+ assert conn.extra == '{"key": "value"}'
+
+
+class TestCommitToConnectionTable:
+ @pytest.fixture(autouse=True)
+ def setup_teardown(self):
+ clear_db_connections(add_default_connections_back=False)
+ clear_db_connection_tests()
+ yield
+ clear_db_connections(add_default_connections_back=False)
+ clear_db_connection_tests()
+
+ def test_creates_new_connection(self, session):
+ ct = ConnectionTestRequest(
+ connection_id="new_conn",
+ conn_type="postgres",
+ host="db.example.com",
+ login="user",
+ password="secret",
+ schema="mydb",
+ port=5432,
+ )
+ session.add(ct)
+ session.flush()
+
+ ct.commit_to_connection_table(session=session)
+ session.flush()
+
+ from sqlalchemy import select
Review Comment:
Avoid imports inside test functions; move `from sqlalchemy import select` to the module import section so the test file has a consistent import style and avoids per-test import side effects.
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2567,6 +2567,33 @@ scheduler:
type: float
example: ~
default: "120.0"
+ connection_test_timeout:
+ description: |
+ Maximum number of seconds an async connection test is allowed to run
+ before it is considered timed out. The scheduler reaper uses this value
+ plus a grace period to mark stale tests as failed.
+ version_added: 3.3.0
+ type: integer
+ example: ~
+ default: "60"
+ max_connection_test_concurrency:
+ description: |
+ Maximum number of connection tests that can be active (QUEUED + RUNNING)
+ at the same time. Excess tests will remain in PENDING state until slots
+ become available.
+ version_added: 3.3.0
+ type: integer
+ example: ~
+ default: "4"
Review Comment:
PR description/config section mentions `[core] connection_test_timeout` and `[core] max_connection_test_concurrency`, but the actual settings added here are under `[scheduler]` (and the scheduler code reads them from `[scheduler]`). Please align the documented config location (or move the settings) so operators aren’t misled about which section to configure.
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,203 @@
+# 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 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, Uuid, select, text
+from sqlalchemy.orm import Mapped, mapped_column
+
+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))
+
+
+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)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+
+ # Connection fields — password and extra are Fernet-encrypted via FernetFieldsMixin.
+ conn_type: Mapped[str] = mapped_column(String(500), nullable=False)
+ host: Mapped[str | None] = mapped_column(String(500), nullable=True)
+ login: Mapped[str | None] = mapped_column(Text, nullable=True)
+ schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True)
+ port: Mapped[int | None] = mapped_column(Integer, nullable=True)
+ commit_on_success: Mapped[bool] = mapped_column(
+ Boolean, nullable=False, default=False, server_default="0"
+ )
+
+ __table_args__ = (
+ Index("idx_connection_test_request_state_created_at", state, created_at),
+ # since mysql lacks filtered/partial indices, this creates a
+ # duplicate index on mysql. Not the end of the world
+ Index(
+ "idx_connection_test_request_active_conn",
+ "connection_id",
+ unique=True,
+ postgresql_where=text("state IN ('pending', 'queued', 'running')"),
+ sqlite_where=text("state IN ('pending', 'queued', 'running')"),
+ ),
Review Comment:
The `idx_connection_test_request_active_conn` index is declared `unique=True` with `postgresql_where`/`sqlite_where`. On MySQL (which doesn't support partial indexes), this will compile as an unconditional UNIQUE index on `connection_id`, preventing *any* subsequent connection tests for the same `connection_id` even after earlier tests reach terminal states. Consider making this non-unique on MySQL (conditional DDL in the migration) or using a MySQL-compatible pattern (e.g., generated/nullable “active_connection_id” column with a unique index) to enforce uniqueness only for active states.
##########
airflow-core/src/airflow/migrations/versions/0111_3_3_0_add_connection_test_table.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.
+
+"""
+Add connection_test_request table for async connection testing.
+
+Revision ID: a7e6d4c3b2f1
+Revises: a4c2d171ae18
+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 = "a4c2d171ae18"
+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("commit_on_success", sa.Boolean(), nullable=False, server_default="0"),
+ sa.PrimaryKeyConstraint("id", name=op.f("connection_test_request_pkey")),
+ sa.UniqueConstraint("token", name=op.f("connection_test_request_token_uq")),
+ )
+ op.create_index(
+ op.f("idx_connection_test_request_state_created_at"),
+ "connection_test_request",
+ ["state", "created_at"],
+ )
+ # since mysql lacks filtered/partial indices, this creates a
+ # duplicate index on mysql
+ op.create_index(
+ op.f("idx_connection_test_request_active_conn"),
+ "connection_test_request",
+ ["connection_id"],
+ unique=True,
+ postgresql_where=sa.text("state IN ('pending', 'queued', 'running')"),
+ sqlite_where=sa.text("state IN ('pending', 'queued', 'running')"),
+ )
+
+
Review Comment:
This migration creates `idx_connection_test_request_active_conn` as `unique=True` with only `postgresql_where`/`sqlite_where`. On MySQL, that becomes an unconditional UNIQUE index on `connection_id`, which will block creating a second test request for the same connection after the first one completes. Please adjust the migration to avoid an unconditional unique index on MySQL (dialect-conditional index creation or a MySQL-compatible partial-uniqueness workaround).
```suggestion
dialect_name = op.get_bind().dialect.name
if dialect_name in {"postgresql", "sqlite"}:
op.create_index(
op.f("idx_connection_test_request_active_conn"),
"connection_test_request",
["connection_id"],
unique=True,
postgresql_where=sa.text("state IN ('pending', 'queued', 'running')"),
sqlite_where=sa.text("state IN ('pending', 'queued', 'running')"),
)
else:
# MySQL does not support filtered/partial indexes. Create a plain
# non-unique index instead of an unconditional unique index.
op.create_index(
op.f("idx_connection_test_request_active_conn"),
"connection_test_request",
["connection_id"],
unique=False,
)
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2901196801
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Hi @ferruzzi I've pulled out the workload queue generalization you suggested into a separate PR #[New Pr that points to this branch](https://github.com/anishgirianish/airflow/pull/1) taking out branch from this pr. Would like to request you for your review on this. I will point it to main once connection testing pr is merged.
Thank you.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2897389432
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Maybe the queue_type should also include something to indicate queue priority and whether that queue is intended to be used FIFO or prioritized?
Kinda typing this as I think, so assume the examples are just rough pseudocode, the "plan" is only half-formed, etc... but....
If we add a NamedTuple called `WorkloadQueueDef(scheduling_tier: int, sort_key: int)` into the base workload and they are defined like this in the individual workloads:
```python
class TestConnection:
@property
def queue_def(self) -> WorkloadQueueDef :
return WorkloadQueueDef(
scheduling_tier = -1, # connection tests are super fast so let's always get them out of the way first
sort_key = 0, # FIFO order
)
class ExecuteCallback:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
scheduling_tier = 0, # always finish a callback before starting a new task, but don't block the much faster fast connection tests for a callback
sort_key = 0, # FIFO order
)
class ExecuteTask:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
# after all connection tests and callbacks are running, take the task with the highest priority
scheduling_tier = 1,
sort_key = self.ti.priority_weight,
)
```
((IN fact, we can set them to both default to 0 and that gets even more simplified.... but you get the idea))
and replace all of the "supports_callbacks", "supports_connection_test", etc in each executor can be replaced by a `supported_workload_types: set`. in LocalExecutr it might look like `supported_workload_types = {"ExecuteTask", "ExecuteCallback", "TestConnection"}`
then the entire chunk of code in _get_workloads_to_schedule that sorts out which jobs to pick up would be something real simple like
```python
def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]:
# create a meta_queue which consists of each sub-queue, sorted by that sub-queue's defined sort_key in the order of the scheduling tiers. Then take the top `open_slots` number of jobs off the list.
all_workloads: list[ExecutorWorkload] = [
workload
for queue in self.executor_queues.values()
for workload in queue.values()
]
# Sort by scheduling tier first, then by sort_key within each tier.
all_workloads.sort(key=lambda workload: workload.queue_def)
return [(workload.key, workload) for workload in all_workloads[:open_slots]]
```
As a bonus, this also assures that future workload types honor slots by default; if some future workload type wanted to break that, it would have to be done intentionally. It also means `slots_available` and `slots_occupied` turn into simple comprehensions along the lines of `sum(len(q) for q in self.executor_queues.values())` that don't have to be updated with every new workload. The code overall gets much cleaner.
I'm not sure, I haven't thought that through fully, but as we add more workload types we may need to get that sorted out before long. I have some rough ideas but they were all overkill for only having two types at the time so i figured this was a good place to get those thoughts "on paper". Feel free to resolve this comment once you have seen it so nobody sees it as blocking, but feel free to message me on Slack and/or tag me in the PR if you get to it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4015530671
Hi, @kaxil @pierrejeambrun, I would like to request your review whenever you get a chance.
Thank you so much.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2897389432
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Maybe the queue_type should also include something to indicate queue priority and whether that queue is intended to be used FIFO or prioritized?
Something vaguely like
If we add a NamedTuple called `WorkloadQueueDef(scheduling_tier: int, sort_key: int)` into the base workload and they are defined like this in the individual workloads:
```python
class TestConnection:
@property
def queue_def(self) -> WorkloadQueueDef :
return WorkloadQueueDef(
scheduling_tier = -1, # connection tests are super fast so let's always get them out of the way first
sort_key = 0, # FIFO order
)
class ExecuteCallback:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
scheduling_tier = 0, # always finish a callback before starting a new task, but don't block the much faster fast connection tests for a callback
sort_key = 0, # FIFO order
)
class ExecuteTask:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
# after all connection tests and callbacks are running, take the task with the highest priority
scheduling_tier = 1,
sort_key = self.ti.priority_weight,
)
```
and replace all of the "supports_callbacks", "supports_connection_test", etc in each executor can be replaced by a `supported_workload_types: set`. in LocalExecutr it might look like `supported_workload_types = {"ExecuteTask", "ExecuteCallback", "TestConnection"}`
then the entire chunk of code in _get_workloads_to_schedule that sorts out which jobs to pick up would be something real simple like
```python
def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]:
# create a meta_queue which consists of each sub-queue, sorted by that sub-queue's defined sort_key in the order of the scheduling tiers. Then take the top `open_slots` number of jobs off the list.
all_workloads: list[ExecutorWorkload] = [
workload
for queue in self.executor_queues.values()
for workload in queue.values()
]
# Sort by scheduling tier first, then by sort_key within each tier.
all_workloads.sort(key=lambda workload: workload.queue_def)
return [(workload.key, workload) for workload in all_workloads[:open_slots]]
```
As a bonus, this also assures that future workload types honor slots by default; if some future workload type wanted to break that, it would have to be done intentionally. It also means `slots_available` and `slots_occupied` turn into simple comprehensions along the lines of `sum(len(q) for q in self.executor_queues.values())` that don't have to be updated with every new workload. The code overall gets much cleaner.
I'm not sure, I haven't thought that through fully, but as we add more workload types we may need to get that sorted out before long. I have some rough ideas but they were all overkill for only having two types at the time so i figured this was a good place to get those thoughts "on paper". I'll resolve this comment so nobody sees it as blocking, but feel free to message me on Slack and/or tag me in the PR if you get to it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2963468962
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via a TestConnection workload."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True)
+ reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0")
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(
+ self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return f"<ConnectionTest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+
+def run_connection_test(*, conn: Connection) -> tuple[bool, str]:
+ """
+ Worker-side function to execute a connection test.
+
+ Returns a (success, message) tuple. The caller is responsible for
+ reporting the result back via the Execution API.
+ """
+ try:
+ return conn.test_connection()
+ except Exception as e:
+ log.exception("Connection test failed", connection_id=conn.conn_id)
+ return False, str(e)
+
+
+_SNAPSHOT_FIELDS = (
+ "conn_type",
+ "description",
+ "host",
+ "login",
+ "_password",
+ "schema",
+ "port",
+ "_extra",
+ "is_encrypted",
+ "is_extra_encrypted",
+)
+
+
+def snapshot_connection(conn: Connection) -> dict:
Review Comment:
**[Critical] Secrets in `connection_snapshot` JSON column**
`snapshot_connection()` stores Fernet-encrypted `_password` and `_extra` ciphertext in a JSON column. This expands the attack surface — anyone with read access to the `connection_test` table gets credential ciphertext that was previously confined to the `connection` table.
Two additional problems:
1. If Fernet keys rotate, stored snapshots become undecryptable and reverts will write garbage.
2. DB-level audit/masking rules on `connection._password` won't cover this secondary storage.
If the snapshot approach survives (vs. Pierre's request-buffer redesign), consider storing only non-secret fields and re-fetching credentials at revert time, or encrypting the entire snapshot blob with a short-lived key.
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
Review Comment:
**[High] PENDING tests never get reaped**
`ACTIVE_STATES` includes QUEUED and RUNNING but not PENDING. The reaper (`_reap_stale_connection_tests`) only checks `ACTIVE_STATES`, so tests stuck in PENDING — e.g., scheduler restart between creation and dispatch, or no supporting executor found — accumulate forever.
If those PENDING tests have snapshots (from save-and-test), the connection stays in its post-edit state with no resolution path.
```suggestion
ACTIVE_STATES = frozenset((ConnectionTestState.PENDING, ConnectionTestState.QUEUED, ConnectionTestState.RUNNING))
```
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via a TestConnection workload."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True)
+ reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0")
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(
+ self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return f"<ConnectionTest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+
+def run_connection_test(*, conn: Connection) -> tuple[bool, str]:
+ """
+ Worker-side function to execute a connection test.
+
+ Returns a (success, message) tuple. The caller is responsible for
+ reporting the result back via the Execution API.
+ """
+ try:
+ return conn.test_connection()
+ except Exception as e:
+ log.exception("Connection test failed", connection_id=conn.conn_id)
+ return False, str(e)
+
+
+_SNAPSHOT_FIELDS = (
+ "conn_type",
+ "description",
+ "host",
+ "login",
+ "_password",
+ "schema",
+ "port",
+ "_extra",
+ "is_encrypted",
+ "is_extra_encrypted",
+)
+
+
+def snapshot_connection(conn: Connection) -> dict:
+ """
+ Capture raw DB column values from a Connection for later restore.
+
+ Encrypted fields (``_password``, ``_extra``) are stored as ciphertext
+ so they can be written directly back without re-encryption.
+ """
+ return {field: getattr(conn, field) for field in _SNAPSHOT_FIELDS}
+
+
+def _revert_connection(conn: Connection, snapshot: dict) -> None:
+ """
+ Restore a Connection's columns from a snapshot dict.
+
+ Writes directly to ``_password`` and ``_extra`` (bypassing the
+ encrypting property setters) so the stored ciphertext is preserved.
+ """
+ for field, value in snapshot.items():
+ setattr(conn, field, value)
+
+
+def _decrypt_snapshot_field(snapshot: dict, field: str) -> str | None:
+ """Decrypt a single encrypted field from a snapshot dict using Fernet."""
+ raw = snapshot.get(field)
+ if raw is None:
+ return None
+ encrypted_flag = "is_encrypted" if field == "_password" else "is_extra_encrypted"
+ if not snapshot.get(encrypted_flag, False):
+ return raw
+ fernet = get_fernet()
+ return fernet.decrypt(bytes(raw, "utf-8")).decode()
+
+
+def _can_safely_revert(conn: Connection, post_snapshot: dict) -> bool:
+ """
+ Check whether the connection's current state matches the post-edit snapshot.
+
+ Compares **decrypted** values for encrypted fields and direct values for
+ non-encrypted fields. Returns ``False`` if any field differs, indicating
+ a concurrent edit has occurred and the revert should be skipped.
+ """
+ for field in _SNAPSHOT_FIELDS:
+ if field in ("is_encrypted", "is_extra_encrypted"):
+ continue
+
+ if field == "_password":
+ current_val = conn.password
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_password")
+ elif field == "_extra":
+ current_val = conn.extra
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_extra")
+ else:
+ current_val = getattr(conn, field)
+ snapshot_val = post_snapshot.get(field)
+
+ if current_val != snapshot_val:
+ return False
+ return True
+
+
+def attempt_revert(ct: ConnectionTest, *, session: Session) -> None:
+ """Revert a connection to its pre-edit values if no concurrent edit has occurred."""
+ if not ct.connection_snapshot:
+ log.warning("attempt_revert called without snapshot", connection_test_id=ct.id)
+ return
+
+ pre_snapshot = ct.connection_snapshot["pre"]
+ post_snapshot = ct.connection_snapshot["post"]
+
+ ct.connection_snapshot = None
+
+ conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id))
+ if conn is None:
+ ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists."
+ log.warning("Revert skipped: connection no longer exists", connection_id=ct.connection_id)
+ return
+
+ if not _can_safely_revert(conn, post_snapshot):
+ ct.result_message = (
+ ct.result_message or ""
Review Comment:
**[High] `attempt_revert` doesn't lock the Connection row**
This `select(Connection)` reads without `with_for_update()`. A concurrent edit can land between the `_can_safely_revert` check and `_revert_connection`, causing the revert to silently overwrite the concurrent change — the exact scenario `_can_safely_revert` is meant to prevent.
```python
conn = session.scalar(
select(Connection).filter_by(conn_id=ct.connection_id).with_for_update()
)
```
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via a TestConnection workload."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True)
+ reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0")
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(
+ self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return f"<ConnectionTest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+
+def run_connection_test(*, conn: Connection) -> tuple[bool, str]:
+ """
+ Worker-side function to execute a connection test.
+
+ Returns a (success, message) tuple. The caller is responsible for
+ reporting the result back via the Execution API.
+ """
+ try:
+ return conn.test_connection()
+ except Exception as e:
+ log.exception("Connection test failed", connection_id=conn.conn_id)
+ return False, str(e)
+
+
+_SNAPSHOT_FIELDS = (
+ "conn_type",
+ "description",
+ "host",
+ "login",
+ "_password",
+ "schema",
+ "port",
+ "_extra",
+ "is_encrypted",
+ "is_extra_encrypted",
+)
+
+
+def snapshot_connection(conn: Connection) -> dict:
+ """
+ Capture raw DB column values from a Connection for later restore.
+
+ Encrypted fields (``_password``, ``_extra``) are stored as ciphertext
+ so they can be written directly back without re-encryption.
+ """
+ return {field: getattr(conn, field) for field in _SNAPSHOT_FIELDS}
+
+
+def _revert_connection(conn: Connection, snapshot: dict) -> None:
+ """
+ Restore a Connection's columns from a snapshot dict.
+
+ Writes directly to ``_password`` and ``_extra`` (bypassing the
+ encrypting property setters) so the stored ciphertext is preserved.
+ """
+ for field, value in snapshot.items():
+ setattr(conn, field, value)
+
+
+def _decrypt_snapshot_field(snapshot: dict, field: str) -> str | None:
+ """Decrypt a single encrypted field from a snapshot dict using Fernet."""
+ raw = snapshot.get(field)
+ if raw is None:
+ return None
+ encrypted_flag = "is_encrypted" if field == "_password" else "is_extra_encrypted"
+ if not snapshot.get(encrypted_flag, False):
+ return raw
+ fernet = get_fernet()
+ return fernet.decrypt(bytes(raw, "utf-8")).decode()
+
+
+def _can_safely_revert(conn: Connection, post_snapshot: dict) -> bool:
+ """
+ Check whether the connection's current state matches the post-edit snapshot.
+
+ Compares **decrypted** values for encrypted fields and direct values for
+ non-encrypted fields. Returns ``False`` if any field differs, indicating
+ a concurrent edit has occurred and the revert should be skipped.
+ """
+ for field in _SNAPSHOT_FIELDS:
+ if field in ("is_encrypted", "is_extra_encrypted"):
+ continue
+
+ if field == "_password":
+ current_val = conn.password
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_password")
+ elif field == "_extra":
+ current_val = conn.extra
+ snapshot_val = _decrypt_snapshot_field(post_snapshot, "_extra")
+ else:
+ current_val = getattr(conn, field)
+ snapshot_val = post_snapshot.get(field)
+
+ if current_val != snapshot_val:
+ return False
+ return True
+
+
+def attempt_revert(ct: ConnectionTest, *, session: Session) -> None:
+ """Revert a connection to its pre-edit values if no concurrent edit has occurred."""
+ if not ct.connection_snapshot:
+ log.warning("attempt_revert called without snapshot", connection_test_id=ct.id)
+ return
+
+ pre_snapshot = ct.connection_snapshot["pre"]
+ post_snapshot = ct.connection_snapshot["post"]
+
+ ct.connection_snapshot = None
+
+ conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id))
+ if conn is None:
+ ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists."
+ log.warning("Revert skipped: connection no longer exists", connection_id=ct.connection_id)
+ return
+
+ if not _can_safely_revert(conn, post_snapshot):
Review Comment:
**[High] Snapshot cleared before revert completes**
`ct.connection_snapshot = None` is set here, before the actual revert logic runs. If the subsequent `select(Connection)` query or `_revert_connection` call fails (DB hiccup, connection dropped), the snapshot is lost and the revert can never be retried.
Move `ct.connection_snapshot = None` to the end of the function, after the revert succeeds or is definitively skipped.
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -168,6 +177,84 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
+ """
+ Execute a connection test workload.
+
+ Constructs an SDK ``Client``, fetches the connection via the Execution API,
+ enforces a timeout via ``signal.alarm``, and reports all outcomes back
+ through the Execution API.
+
+ :param log: Logger instance
+ :param workload: The TestConnection workload to execute
+ :param team_conf: Team-specific executor configuration
+ """
+ # Lazy import: SDK modules must not be loaded at module level to avoid
+ # coupling core (scheduler-loaded) code to the SDK.
+ from airflow.sdk.api.client import Client
+ from airflow.sdk.execution_time.comms import ErrorResponse
+
+ setproctitle(
+ f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}",
+ log,
+ )
+
+ base_url = team_conf.get("api", "base_url", fallback="/")
+ if base_url.startswith("/"):
+ base_url = f"http://localhost:8080{base_url}"
+ default_execution_api_server = f"{base_url.rstrip('/')}/execution/"
+ server = team_conf.get("core", "execution_api_server_url", fallback=default_execution_api_server)
+
+ client = Client(base_url=server, token=workload.token)
+
+ def _handle_timeout(signum, frame):
+ raise TimeoutError(f"Connection test timed out after {workload.timeout}s")
+
+ signal.signal(signal.SIGALRM, _handle_timeout)
+ signal.alarm(workload.timeout)
+ try:
+ client.connection_tests.update_state(workload.connection_test_id, ConnectionTestState.RUNNING)
+
+ conn_response = client.connections.get(workload.connection_id)
+ if isinstance(conn_response, ErrorResponse):
+ raise RuntimeError(f"Connection '{workload.connection_id}' not found via Execution API")
+
+ conn = Connection(
+ conn_id=conn_response.conn_id,
+ conn_type=conn_response.conn_type,
+ host=conn_response.host,
+ login=conn_response.login,
+ password=conn_response.password,
+ schema=conn_response.schema_,
+ port=conn_response.port,
+ extra=conn_response.extra,
+ )
+ success, message = run_connection_test(conn=conn)
+
+ state = ConnectionTestState.SUCCESS if success else ConnectionTestState.FAILED
+ client.connection_tests.update_state(workload.connection_test_id, state, message)
+ except TimeoutError:
+ log.error(
+ "Connection test timed out after %ds",
+ workload.timeout,
+ connection_id=workload.connection_id,
+ )
+ client.connection_tests.update_state(
+ workload.connection_test_id,
+ ConnectionTestState.FAILED,
+ f"Connection test timed out after {workload.timeout}s",
+ )
+ except Exception as e:
+ log.exception("Connection test failed unexpectedly", connection_id=workload.connection_id)
+ client.connection_tests.update_state(
+ workload.connection_test_id,
+ ConnectionTestState.FAILED,
Review Comment:
**[Low] Exception messages in `result_message` may leak connection details**
`f"Connection test failed unexpectedly: {e}"[:500]` stores raw exception text. Some DB drivers include connection strings (with credentials) in error messages. This gets stored in `result_message` and exposed via the polling API.
Consider sanitizing or using a generic message with the exception type only:
```python
f"Connection test failed unexpectedly: {type(e).__name__}"
```
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -305,10 +315,19 @@ def heartbeat(self) -> None:
self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads)
self.trigger_tasks(open_slots)
+ self.trigger_connection_tests()
+
# Calling child class sync method
self.log.debug("Calling the %s sync method", self.__class__)
self.sync()
+ def trigger_connection_tests(self) -> None:
+ """Process queued connection tests."""
+ if not self.supports_connection_test or not self.queued_connection_tests:
+ return
+
+ self._process_workloads(list(self.queued_connection_tests.values()))
Review Comment:
**[Medium] `trigger_connection_tests` dispatches all queued tests without rechecking slots**
This is called from `heartbeat()` after `trigger_tasks(open_slots)` has already consumed slots. `trigger_connection_tests` then dispatches *all* queued connection tests via `_process_workloads` without checking remaining capacity.
The scheduler's `_enqueue_connection_tests` does slot accounting, but slots may have been consumed between enqueue time and executor heartbeat. With `max_connection_test_concurrency=4` the overshoot is bounded, but it'd be cleaner to pass remaining slots through or check `self.slots_available` here.
##########
airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py:
##########
@@ -0,0 +1,71 @@
+#
+# 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 table for async connection testing.
+
+Revision ID: a7e6d4c3b2f1
+Revises: 1d6611b6ab7c
+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 = "1d6611b6ab7c"
+branch_labels = None
+depends_on = None
+airflow_version = "3.2.0"
+
+
+def upgrade():
+ """Create connection_test table."""
+ op.create_table(
+ "connection_test",
+ 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(10), 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),
Review Comment:
**[Medium] `state` column `String(10)` is tight**
Longest current value is `"pending"` (7 chars). `String(10)` leaves almost no room — `"timed_out"` (9) or `"cancelled"` (9) would need a new migration. Use `String(20)` for breathing room.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -174,6 +189,70 @@ def bulk_connections(
return BulkConnectionService(session=session, request=request).handle_request()
+@connections_router.patch(
+ "/{connection_id}/test",
+ responses=create_openapi_http_exception_doc(
+ [
+ status.HTTP_400_BAD_REQUEST,
+ status.HTTP_403_FORBIDDEN,
+ status.HTTP_404_NOT_FOUND,
+ ]
+ ),
+ dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())],
+)
+def patch_connection_and_test(
+ connection_id: str,
+ patch_body: ConnectionBody,
+ session: SessionDep,
+ update_mask: list[str] | None = Query(None),
+ executor: str | None = Query(None, description="Executor to route the connection test to"),
+ queue: str | None = Query(None, description="Queue to route the connection test to"),
+) -> ConnectionSaveAndTestResponse:
+ """
+ Update a connection and queue an async test with revert-on-failure.
+
+ Atomically saves the edit and creates a ConnectionTest with snapshots of the
+ pre-edit and post-edit state. If the test fails, the connection is automatically
+ reverted to its pre-edit values.
+ """
+ _ensure_test_connection_enabled()
+
+ if patch_body.connection_id != connection_id:
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST,
+ "The connection_id in the request body does not match the URL parameter",
+ )
+
+ connection = session.scalar(select(Connection).filter_by(conn_id=connection_id).limit(1))
+ if connection is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ f"The Connection with connection_id: `{connection_id}` was not found",
+ )
+
+ try:
+ ConnectionBody(**patch_body.model_dump())
Review Comment:
**[Medium] Redundant validation**
`patch_body` is already a `ConnectionBody` instance — FastAPI validated it on parameter injection. Re-constructing `ConnectionBody(**patch_body.model_dump())` just validates the same data again. This block is dead code.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3103384659
##########
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)
Review Comment:
Pretty sure signal.SIGALRM is *nix only and won't work on other operating systems. Do we care about that here?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-3985326131
Hi @kaxil @pierrejeambrun @jason810496 I have updated the pr based on recent discussion we had in the thread. I would like to request you for your re-review whenever you get chance. Thank you so much.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2876755045
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -168,6 +177,84 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
Review Comment:
Agreed, for sure will follow up with this change after this, or more than happy to include in this pr if it gets merged first, thanks.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177847292
##########
airflow-core/src/airflow/migrations/versions/0114_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:
Done, added server_default="0" to both columns (in migration + model override on ConnectionTestRequest to keep alembic autogen happy). thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177833542
##########
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",
Review Comment:
Done, now GET locks the row, transitions PENDING/QUEUED→RUNNING server-side, second fetch returns 409. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2842160078
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -249,6 +277,99 @@ 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_404_NOT_FOUND]),
+ 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 must already be saved. Returns a token that can be used
+ to poll for the test result via GET /connections/test-async/{token}.
+ """
+ if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ "Testing connections is disabled in Airflow configuration. "
+ "Contact your deployment admin to enable it.",
+ )
+
+ try:
+ Connection.get_connection_from_secrets(test_body.connection_id)
+ except AirflowNotFoundException:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ f"The Connection with connection_id: `{test_body.connection_id}` was not found. "
+ "Connection must be saved before testing.",
+ )
+
+ connection_test = ConnectionTest(connection_id=test_body.connection_id)
+ session.add(connection_test)
+ session.flush()
+
+ # ExecutorCallback requires an object satisfying ImportPathExecutorCallbackDefProtocol,
+ # but the only concrete impl (SyncCallback) lives in task-sdk which core API cannot
+ # import. This adapter will be replaced by ExecuteCallback.make() once #61153 merges.
+ callback_def = _ImportPathCallbackDef(
+ path=RUN_CONNECTION_TEST_PATH,
+ kwargs={
+ "connection_id": test_body.connection_id,
+ "connection_test_id": str(connection_test.id),
+ },
+ )
+ callback = ExecutorCallback(callback_def, fetch_method=CallbackFetchMethod.IMPORT_PATH)
Review Comment:
moved it to connection_test.py and added a create_callback() factory method on the model. The route now just calls connection_test.create_callback(). Much cleaner, thank you!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2842187524
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -56,13 +59,38 @@
from airflow.configuration import conf
from airflow.exceptions import AirflowNotFoundException
from airflow.models import Connection
+from airflow.models.callback import CallbackFetchMethod, ExecutorCallback
+from airflow.models.connection_test import (
+ RUN_CONNECTION_TEST_PATH,
+ ConnectionTest,
+ ConnectionTestState,
+)
from airflow.secrets.environment_variables import CONN_ENV_PREFIX
from airflow.utils.db import create_default_connections as db_create_default_connections
from airflow.utils.strings import get_random_string
connections_router = AirflowRouter(tags=["Connection"], prefix="/connections")
Review Comment:
thanks for the pointer! Would be great to refactor the sync endpoint to internally queue and stream/wait using the same approach. Happy to explore that as a follow-up once this lands.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2839476949
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,125 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import Boolean, ForeignKey, Index, String, Text, Uuid
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+if TYPE_CHECKING:
+ from airflow.models.callback import Callback
+
+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
+
+
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+# Path used by ExecutorCallback to locate the worker function.
+RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test"
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via ExecutorCallback."""
+
+ __tablename__ = "connection_test"
+
+ id: Mapped[UUID] = mapped_column(Uuid(), primary_key=True, default=uuid6.uuid7)
+ token: Mapped[str] = mapped_column(String(64), nullable=False, unique=True)
Review Comment:
May I ask why not just using `ConnectionTest.id` instead of `ConnectionTest.token` as identifier to poll the status of `ConnectionTest`?
Since the authentication is already handled by the JWT layer, the `token` here will still be served as _another_ identifier for retrieving the corresponding `ConnectionTest` model.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -249,6 +277,99 @@ 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_404_NOT_FOUND]),
+ 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 must already be saved. Returns a token that can be used
+ to poll for the test result via GET /connections/test-async/{token}.
+ """
+ if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ "Testing connections is disabled in Airflow configuration. "
+ "Contact your deployment admin to enable it.",
+ )
+
+ try:
+ Connection.get_connection_from_secrets(test_body.connection_id)
+ except AirflowNotFoundException:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ f"The Connection with connection_id: `{test_body.connection_id}` was not found. "
+ "Connection must be saved before testing.",
+ )
+
+ connection_test = ConnectionTest(connection_id=test_body.connection_id)
+ session.add(connection_test)
+ session.flush()
+
+ # ExecutorCallback requires an object satisfying ImportPathExecutorCallbackDefProtocol,
+ # but the only concrete impl (SyncCallback) lives in task-sdk which core API cannot
+ # import. This adapter will be replaced by ExecuteCallback.make() once #61153 merges.
+ callback_def = _ImportPathCallbackDef(
+ path=RUN_CONNECTION_TEST_PATH,
+ kwargs={
+ "connection_id": test_body.connection_id,
+ "connection_test_id": str(connection_test.id),
+ },
+ )
+ callback = ExecutorCallback(callback_def, fetch_method=CallbackFetchMethod.IMPORT_PATH)
Review Comment:
How above moving `_ImportPathCallbackDef` to `models.test_connection` module and add a factory method for constructing `ExecutorCallback`?
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -56,13 +59,38 @@
from airflow.configuration import conf
from airflow.exceptions import AirflowNotFoundException
from airflow.models import Connection
+from airflow.models.callback import CallbackFetchMethod, ExecutorCallback
+from airflow.models.connection_test import (
+ RUN_CONNECTION_TEST_PATH,
+ ConnectionTest,
+ ConnectionTestState,
+)
from airflow.secrets.environment_variables import CONN_ENV_PREFIX
from airflow.utils.db import create_default_connections as db_create_default_connections
from airflow.utils.strings import get_random_string
connections_router = AirflowRouter(tags=["Connection"], prefix="/connections")
Review Comment:
Just FYI that if rest of the reviewers also agree with current approach, we could refactor the existing synchronous test connection route like
https://github.com/apache/airflow/blob/01318970138e873deb5b6e5883cb1a33b0d15a04/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py#L527-L537
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,125 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import Boolean, ForeignKey, Index, String, Text, Uuid
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+if TYPE_CHECKING:
+ from airflow.models.callback import Callback
+
+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
+
+
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+# Path used by ExecutorCallback to locate the worker function.
+RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test"
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via ExecutorCallback."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_status: Mapped[bool | None] = mapped_column(Boolean, nullable=True)
Review Comment:
The `state` column should be sufficient, do we still need `result_status`?
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py:
##########
@@ -0,0 +1,65 @@
+# 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 fastapi import APIRouter, HTTPException, status
+
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.connection_test import ConnectionTestResultBody
+from airflow.models.connection_test import TERMINAL_STATES, ConnectionTest
+
+router = APIRouter()
+
+
Review Comment:
This route is not currently used in TaskSDK, but I think it's fine to keep it in current PR.
Since the synchronous callback _might_ be executed under TaskSDK instead directly being executed under Scheduler/Executor process in the long-term.
Will confirm in https://github.com/apache/airflow/pull/61153#discussion_r2839816002
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-3961599489
> @pierrejeambrun thank you very much for the review. The dev list discussion ([thread](https://lists.apache.org/thread/xd7zmyp95y77cw36mb5wjp17dyynz100)) received positive feedback , and the implementation follows the suggestions outlined in the thread. Happy to wait for a more formal conclusion on the thread. Would love to know if there's anything you'd like me to revisit or improve in the meantime.
I do plan to review proposal, PR and reply but occupied today. So tomorrow or Friday.
There isn't any hurry either since earliest this would go would be 3.2 anyway and there is some time to go for it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4001226267
> Hi @anishgirianish, thank you for the update.
>
> However, the main problem mentioned in the last review has not yet been resolved. We’re still mixing the `executor` field from the Test Connection API with the `team_name` field in Executor.
>
> Additionally, after the last fix, the `queue` semantics mentioned in the dev list are missing from the current PR.
>
> Could you please revise the semantics of `executor`, `team_name`, and `queue` when you have time? Thanks!
Hi @jason810496 Thank you so much for the re-review. Addressed all review feedback:
- Executor matching uses executor.name.alias/module_path instead of
team_name
- Restored queue field (model, migration, API, workload)
- Moved supports_connection_test guard into trigger_connection_tests()
- Defined ACTIVE_STATES frozenset, used in scheduler
Would like to request you for your re-review thanks.
E2e verified via Breeze (screencast attached): async test flow + save-and-test with revert-on-failure both working.
https://github.com/user-attachments/assets/a2c2396e-e1be-4803-be77-7a1f9c3c6662
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4336352723
Feel free to resolve comments you have addressed so we know if works remain to be done before the next review.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3292563762
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -259,6 +283,94 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
+ "/enqueue-test",
+ 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 enqueue_connection_test(
+ test_body: ConnectionTestRequestBody,
+ session: SessionDep,
+) -> ConnectionTestQueuedResponse:
+ """Enqueue a connection test for deferred execution on a worker; returns a polling 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)
+ try:
+ session.flush()
+ except IntegrityError:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ f"An active connection test already exists for connection_id `{test_body.connection_id}`.",
+ )
+
+ return ConnectionTestQueuedResponse(
+ token=connection_test.token,
+ connection_id=connection_test.connection_id,
+ state=connection_test.state,
+ )
+
+
+@connections_router.get(
+ "/enqueue-test/{connection_test_token}",
Review Comment:
**Bearer token in URL path leaks via access logs.**
Putting the worker callback token in the URL path means it ends up in API server access logs, reverse-proxy logs, browser history, and any intermediate observability that captures request URIs. A token sitting in a URL is treated as low-sensitivity by every layer between the worker and the API, but it grants the same authority as a JWT in an `Authorization` header.
Two options:
1. Move the token to an `Authorization: Bearer ...` header (or a POST body field) and resolve the request by `request_id` alone.
2. Keep the URL stable but store only a SHA-256 hash of the token, then accept the plaintext via header and compare hashes server-side.
Option 2 is friendlier if the worker has to retry without re-issuing tokens, but option 1 is simpler.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3245,6 +3261,94 @@ def _cleanup_orphaned_asset_state(*, session: Session) -> None:
)
session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+ def _enqueue_connection_tests(self, *, session: Session) -> None:
+ """
+ Enqueue pending connection tests to executors that support them.
+
+ ``max_concurrency`` is per-scheduler, not global: with N HA schedulers
+ the worst-case per-tick dispatch is ``N * max_concurrency``. Connection
+ tests are user-initiated and rare, so the overshoot self-corrects via
+ the reaper. For a true global cap, wrap the budget+claim below in a
+ sentinel-row ``SELECT ... FOR UPDATE``.
+ """
+ max_concurrency = conf.getint("connection_test", "max_concurrency", fallback=4)
+ timeout = conf.getint("connection_test", "timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTestRequest.id)).where(
+ ConnectionTestRequest.state.in_(DISPATCHED_STATES)
+ )
+ )
+ budget = max_concurrency - (active_count or 0)
+ 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)
Review Comment:
**Multi-team executor routing is broken at dispatch.**
`_get_workload_team_name` calls `ct.get_dag_id()`, which returns `None` for `ConnectionTestRequest` because there is no DAG behind a connection test. The `team_name` resolution then collapses to `None`, and every connection-test workload routes through the global executor regardless of which team the connection belongs to. The team-scoped executor pool the rest of the PR sets up is bypassed for this workload type.
Resolve the team directly from the row:
```python
team_name = Connection.get_team_name(ct.connection_id, session)
executor = self._try_to_load_executor(executor_name, team_name=team_name)
```
Ideally a regression test that spins up two team-scoped executors and asserts a `ConnectionTest` workload for team B lands on team B executor, not the global one.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3157909440
##########
airflow-core/src/airflow/models/connection.py:
##########
@@ -131,7 +131,7 @@ class Connection(Base, LoggingMixin):
host: Mapped[str | None] = mapped_column(String(500), nullable=True)
schema: Mapped[str | None] = mapped_column(String(500), nullable=True)
login: Mapped[str | None] = mapped_column(Text(), nullable=True)
- _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True)
+ # _password and _extra columns inherited from FernetFieldsMixin
Review Comment:
Removed. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3029282065
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,234 @@
+# 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 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, Uuid, select
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, synonym
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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))
+
+
+class ConnectionTestRequest(Base):
+ """
+ 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)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+
+ # Connection fields — password and extra are Fernet-encrypted.
+ conn_type: Mapped[str] = mapped_column(String(500), nullable=False)
+ host: Mapped[str | None] = mapped_column(String(500), nullable=True)
+ login: Mapped[str | None] = mapped_column(Text, nullable=True)
+ _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True)
+ schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True)
+ port: Mapped[int | None] = mapped_column(Integer, nullable=True)
+ _extra: Mapped[str | None] = mapped_column("extra", Text(), nullable=True)
+ commit_on_success: Mapped[bool] = mapped_column(
+ Boolean, nullable=False, default=False, server_default="0"
+ )
+
+ __table_args__ = (Index("idx_connection_test_request_state_created_at", state, created_at),)
+
+ def __init__(
+ self,
+ *,
+ connection_id: str,
+ conn_type: str,
+ host: str | None = None,
+ login: str | None = None,
+ password: str | None = None,
+ schema: str | None = None,
+ port: int | None = None,
+ extra: str | None = None,
+ commit_on_success: bool = False,
+ executor: str | None = None,
+ queue: str | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.conn_type = conn_type
+ self.host = host
+ self.login = login
+ self.password = password
+ self.schema = schema
+ self.port = port
+ self.extra = extra
+ self.commit_on_success = commit_on_success
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return (
+ f"<ConnectionTestRequest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+ )
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+ def get_password(self) -> str | None:
+ if self._password:
+ fernet = get_fernet()
+ if not fernet.is_encrypted:
+ return self._password
+ return fernet.decrypt(bytes(self._password, "utf-8")).decode()
+ return self._password
+
+ def set_password(self, value: str | None):
+ if value:
+ fernet = get_fernet()
+ self._password = fernet.encrypt(bytes(value, "utf-8")).decode()
+ else:
+ self._password = value
+
+ @declared_attr
+ def password(cls):
+ """Password. The value is decrypted/encrypted when reading/setting the value."""
+ return synonym("_password", descriptor=property(cls.get_password, cls.set_password))
+
+ def get_extra(self) -> str | None:
+ if self._extra:
+ fernet = get_fernet()
+ if not fernet.is_encrypted:
+ return self._extra
+ return fernet.decrypt(bytes(self._extra, "utf-8")).decode()
+ return self._extra
+
+ def set_extra(self, value: str | None):
+ if value:
+ fernet = get_fernet()
+ self._extra = fernet.encrypt(bytes(value, "utf-8")).decode()
+ else:
+ self._extra = value
+
+ @declared_attr
+ def extra(cls):
+ """Extra data. The value is decrypted/encrypted when reading/setting the value."""
+ return synonym("_extra", descriptor=property(cls.get_extra, cls.set_extra))
+
Review Comment:
Done, extracted FernetFieldsMixin into models/crypto.py with shared encrypted _password/_extra columns and synonyms. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3025380679
##########
airflow-core/src/airflow/utils/db.py:
##########
@@ -115,7 +115,7 @@ class MappedClassProtocol(Protocol):
"3.0.3": "fe199e1abd77",
"3.1.0": "cc92b33c6709",
"3.1.8": "509b94a1042d",
- "3.2.0": "1d6611b6ab7c",
+ "3.2.0": "a7e6d4c3b2f1",
"3.3.0": "a4c2d171ae18",
Review Comment:
The `_REVISION_HEADS_MAP` entry for `3.2.0` now points to revision `a7e6d4c3b2f1`, but that revision’s migration file declares `down_revision = "a4c2d171ae18"` (a `3.3.0` migration). This makes the mapping non-monotonic (`3.2.0` > `3.3.0`) and will break logic that relies on version→head ordering. The new migration should be mapped to the *next* Airflow version after `3.3.0` (and `3.2.0` should continue to map to `1d6611b6ab7c`).
```suggestion
"3.2.0": "1d6611b6ab7c",
"3.3.0": "a7e6d4c3b2f1",
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3025679592
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +286,83 @@ 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]),
+ 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()
+
+ # Only one active test per connection_id at a time.
+ _check_no_active_test(test_body.connection_id, session)
+
+ 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)
+ session.flush()
Review Comment:
Theoretically valid but practically unlikely this is triggered by a user manually clicking "test connection" in the UI. Two concurrent requests for the same connection_id within the same transaction window is an extreme edge case. Adding a partial unique index or locking would require extra migration complexity for marginal benefit. Can revisit as a follow-up if it ever surfaces in practice.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177830013
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -144,7 +145,16 @@ def _convert_variable_result_to_variable(var_result: VariableResult, deserialize
return Variable(**var_result.model_dump(exclude={"type"}))
+_preset_connections: ContextVar[dict[str, Connection]] = ContextVar("_preset_connections", default={})
Review Comment:
Done, default=None + guards at both read sites + merge-not-replace inside test_connection().
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2932356638
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -174,6 +189,70 @@ def bulk_connections(
return BulkConnectionService(session=session, request=request).handle_request()
+@connections_router.patch(
+ "/{connection_id}/test",
+ responses=create_openapi_http_exception_doc(
+ [
+ status.HTTP_400_BAD_REQUEST,
+ status.HTTP_403_FORBIDDEN,
+ status.HTTP_404_NOT_FOUND,
+ ]
+ ),
+ dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())],
+)
+def patch_connection_and_test(
Review Comment:
I believe this should be `patch_connection_test` only.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971257991
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via a TestConnection workload."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True)
+ reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0")
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(
+ self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return f"<ConnectionTest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+
+def run_connection_test(*, conn: Connection) -> tuple[bool, str]:
+ """
+ Worker-side function to execute a connection test.
+
+ Returns a (success, message) tuple. The caller is responsible for
+ reporting the result back via the Execution API.
+ """
+ try:
+ return conn.test_connection()
+ except Exception as e:
+ log.exception("Connection test failed", connection_id=conn.conn_id)
+ return False, str(e)
+
+
+_SNAPSHOT_FIELDS = (
+ "conn_type",
+ "description",
+ "host",
+ "login",
+ "_password",
+ "schema",
+ "port",
+ "_extra",
+ "is_encrypted",
+ "is_extra_encrypted",
+)
+
+
+def snapshot_connection(conn: Connection) -> dict:
Review Comment:
The snapshot mechanism has been removed now, and now its using request-buffer mechanism
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971262343
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -168,6 +177,84 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
+ """
+ Execute a connection test workload.
+
+ Constructs an SDK ``Client``, fetches the connection via the Execution API,
+ enforces a timeout via ``signal.alarm``, and reports all outcomes back
+ through the Execution API.
+
+ :param log: Logger instance
+ :param workload: The TestConnection workload to execute
+ :param team_conf: Team-specific executor configuration
+ """
+ # Lazy import: SDK modules must not be loaded at module level to avoid
+ # coupling core (scheduler-loaded) code to the SDK.
+ from airflow.sdk.api.client import Client
+ from airflow.sdk.execution_time.comms import ErrorResponse
+
+ setproctitle(
+ f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}",
+ log,
+ )
+
+ base_url = team_conf.get("api", "base_url", fallback="/")
+ if base_url.startswith("/"):
+ base_url = f"http://localhost:8080{base_url}"
+ default_execution_api_server = f"{base_url.rstrip('/')}/execution/"
+ server = team_conf.get("core", "execution_api_server_url", fallback=default_execution_api_server)
+
+ client = Client(base_url=server, token=workload.token)
+
+ def _handle_timeout(signum, frame):
+ raise TimeoutError(f"Connection test timed out after {workload.timeout}s")
+
+ signal.signal(signal.SIGALRM, _handle_timeout)
+ signal.alarm(workload.timeout)
+ try:
+ client.connection_tests.update_state(workload.connection_test_id, ConnectionTestState.RUNNING)
+
+ conn_response = client.connections.get(workload.connection_id)
+ if isinstance(conn_response, ErrorResponse):
+ raise RuntimeError(f"Connection '{workload.connection_id}' not found via Execution API")
+
+ conn = Connection(
+ conn_id=conn_response.conn_id,
+ conn_type=conn_response.conn_type,
+ host=conn_response.host,
+ login=conn_response.login,
+ password=conn_response.password,
+ schema=conn_response.schema_,
+ port=conn_response.port,
+ extra=conn_response.extra,
+ )
+ success, message = run_connection_test(conn=conn)
+
+ state = ConnectionTestState.SUCCESS if success else ConnectionTestState.FAILED
+ client.connection_tests.update_state(workload.connection_test_id, state, message)
+ except TimeoutError:
+ log.error(
+ "Connection test timed out after %ds",
+ workload.timeout,
+ connection_id=workload.connection_id,
+ )
+ client.connection_tests.update_state(
+ workload.connection_test_id,
+ ConnectionTestState.FAILED,
+ f"Connection test timed out after {workload.timeout}s",
+ )
+ except Exception as e:
+ log.exception("Connection test failed unexpectedly", connection_id=workload.connection_id)
+ client.connection_tests.update_state(
+ workload.connection_test_id,
+ ConnectionTestState.FAILED,
Review Comment:
updated the log thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2970084882
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +286,83 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
Review Comment:
This is the API that is submitted by UI to queue a test task? I am not sure if we should add this to "public API" because I think we should rather expose this to UI only... and keep it a private API. Once reason is that I assume we don't need to have this for external integrations (If I automatically add connections most probably I have tested them before...) as well as (2) this would be a documented exposure that somebody uses this as a door to DDos workload?
Also this would make it harder to change it if it is public... then it is a commiment to keep it stable.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2851604976
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,155 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import ForeignKey, Index, String, Text, Uuid
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+if TYPE_CHECKING:
+ from airflow.models.callback import Callback, ExecutorCallback
+
+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
+
+
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+# Path used by ExecutorCallback to locate the worker function.
+RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test"
+
+
+class _ImportPathCallbackDef:
+ """
+ Minimal implementation of ImportPathExecutorCallbackDefProtocol.
+
+ ExecutorCallback.__init__ expects an object satisfying this protocol, but no concrete
+ implementation exists in airflow-core — the only one (SyncCallback) lives in the task-sdk.
+ Once #61153 lands and ExecuteCallback.make() is decoupled from DagRun, this adapter can
+ be replaced with the proper factory method.
+ """
+
+ def __init__(self, path: str, kwargs: dict, executor: str | None = None):
+ self.path = path
+ self.kwargs = kwargs
+ self.executor = executor
+
+ def serialize(self) -> dict:
+ return {"path": self.path, "kwargs": self.kwargs, "executor": self.executor}
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via ExecutorCallback."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+
+ callback_id: Mapped[UUID | None] = mapped_column(
+ Uuid(), ForeignKey("callback.id", ondelete="SET NULL"), nullable=True
+ )
+ callback: Mapped[Callback | None] = relationship("Callback", uselist=False)
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(self, *, connection_id: str, **kwargs):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return (
+ f"<ConnectionTest token={self.token!r} connection_id={self.connection_id!r} state={self.state}>"
Review Comment:
Yep, for sure added. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2851603185
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -249,6 +256,88 @@ 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_404_NOT_FOUND]),
+ 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 must already be saved. Returns a token that can be used
+ to poll for the test result via GET /connections/test-async/{token}.
+ """
+ if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
Review Comment:
Extracted it to a util. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3316030532
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -259,6 +321,78 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
+ "/enqueue-test",
+ 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(action_logging())],
+)
+def enqueue_connection_test(
+ test_body: ConnectionTestRequestBody,
+ session: SessionDep,
+ user: GetUserDep,
+) -> ConnectionTestQueuedResponse:
+ """Enqueue a connection test for deferred execution on a worker; returns a polling token."""
+ _ensure_test_connection_enabled()
+ _ensure_executor_is_configured(test_body.executor)
+
+ existing = session.scalar(select(Connection).filter_by(conn_id=test_body.connection_id))
+ if existing is not None:
+ effective_team = existing.team_name
+ if test_body.team_name is not None and test_body.team_name != effective_team:
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ f"team_name `{test_body.team_name}` does not match the team of connection "
+ f"`{test_body.connection_id}`.",
+ )
+ else:
+ effective_team = test_body.team_name
Review Comment:
Done thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3316033055
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2699,6 +2699,42 @@ scheduler:
type: boolean
example: ~
default: "False"
+connection_test:
+ description: |
+ Configuration for the deferred connection-test workflow that dispatches
+ test requests to workers via the executor (instead of running them
+ in-process on the API server).
+ options:
+ timeout:
+ description: |
+ Maximum number of seconds a worker-dispatched connection test is
+ allowed to run before it is considered timed out. The scheduler
+ reaper uses this value plus a grace period to mark stale tests as
+ failed.
+ version_added: 3.3.0
+ type: integer
+ example: ~
+ default: "60"
+ max_concurrency:
+ description: |
+ Maximum number of connection tests that can be active
+ (QUEUED + RUNNING) at the same time. Excess tests will remain in
+ PENDING state until slots become available. With N HA schedulers and
+ the dispatch lock seeded by the migration, this is enforced as a
Review Comment:
Updated thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3316037103
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -1047,6 +1051,28 @@ def get_detail_response(self, ti_id: uuid.UUID) -> HITLDetailResponse:
return HITLDetailResponse.model_validate_json(resp.read())
+class ConnectionTestOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def get_connection(self, connection_test_id: uuid.UUID) -> ConnectionTestConnectionResponse:
+ """Fetch connection data for a test request from the API server."""
+ resp = self.client.get(f"connection-tests/{connection_test_id}/connection")
+ return ConnectionTestConnectionResponse.model_validate_json(resp.read())
+
+ def update_state(
+ self, id: uuid.UUID, state: ConnectionTestState, result_message: str | None = None
+ ) -> None:
+ """Report the state of a connection test to the API server."""
+ body = ConnectionTestResultBody(
+ state=state,
+ result_message=ResultMessage(result_message) if result_message is not None else None,
Review Comment:
Done thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3302013192
##########
airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_05_02.py:
##########
@@ -0,0 +1,31 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
Done thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-3989390075
> Thank you for the update!
>
> I’m still a bit concerned about mixing the concepts of `queue`, `executor`, and `team_name`, as this might create ambiguity for users in the future. It may be better to align this with Task‑level semantics, since running a connection test on a worker is effectively another type of workload. We should ensure that the way users specify the runtime environment is clear and consistent.
>
> Thanks!
Thanks @jason810496 ! Addressed all your comments in the latest push, renamed queue to executor across the board to avoid ambiguity with task-level queue semantics. Also raised NotImplementedError in base_executor and added skip_locked to the reaper. Would like to request you for a re-review whenever you get a chance!
thank you!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers [airflow]
Posted by "choo121600 (via GitHub)" <gi...@apache.org>.
choo121600 commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2889535133
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py:
##########
@@ -0,0 +1,73 @@
+# 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, status
+
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.connection_test import ConnectionTestResultBody
+from airflow.models.connection_test import (
+ TERMINAL_STATES,
+ ConnectionTest,
+ ConnectionTestState,
+ attempt_revert,
+)
+
+router = VersionedAPIRouter()
+
+
+@router.patch(
+ "/{connection_test_id}",
+ status_code=status.HTTP_204_NO_CONTENT,
+ 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,
+ session: SessionDep,
+) -> None:
+ """Update the result of a connection test (called by workers)."""
+ ct = session.get(ConnectionTest, 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
+ ct.result_message = body.result_message
+
+ if body.state == ConnectionTestState.FAILED and ct.connection_snapshot:
+ attempt_revert(ct, session=session)
Review Comment:
I might be missing something, but it looks like `connection_snapshot` (Fernet-encrypted `_password` / `_extra`) remains in the DB even after the revert decision is made.
Since the snapshot doesn’t seem to be needed after that point, would it make sense to clear it once the terminal state is reached?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2883213524
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3243,6 +3253,92 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ @provide_session
+ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None:
+ """Dispatch pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("core", "max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("core", "connection_test_timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTest.id)).where(
+ ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING])
+ )
+ )
+ budget = max_concurrency - (active_count or 0)
+ if budget <= 0:
+ return
+
+ pending_stmt = (
+ select(ConnectionTest)
+ .where(ConnectionTest.state == ConnectionTestState.PENDING)
+ .order_by(ConnectionTest.created_at)
+ .limit(budget)
+ )
+ pending_stmt = with_row_locks(pending_stmt, session, of=ConnectionTest, skip_locked=True)
+ pending_tests = session.scalars(pending_stmt).all()
+
+ if not pending_tests:
+ return
+
+ for ct in pending_tests:
+ executor = self._find_executor_for_connection_test(ct.executor)
+ if executor is None:
+ reason = (
+ f"No executor matches '{ct.executor}'"
+ if ct.executor
+ else "No executor supports 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,
+ 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("core", "connection_test_timeout", fallback=60)
+ grace_period = max(30, timeout // 2)
+ cutoff = timezone.utcnow() - timedelta(seconds=timeout + grace_period)
+
+ stale_stmt = select(ConnectionTest).where(
+ ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]),
+ ConnectionTest.updated_at < cutoff,
+ )
+ stale_stmt = with_row_locks(stale_stmt, session, of=ConnectionTest, 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)
+ if ct.connection_snapshot:
+ attempt_revert(ct, session=session)
+
+ session.flush()
+
+ def _find_executor_for_connection_test(self, executor_name: str | None) -> BaseExecutor | None:
+ """Find an executor that supports connection testing, optionally matching by team name."""
+ if executor_name is not None:
+ for executor in self.executors:
+ if executor.supports_connection_test and executor.team_name == executor_name:
+ return executor
+ return None
+ for executor in self.executors:
+ if executor.supports_connection_test:
+ return executor
+ return None
Review Comment:
The main problems pointed out in last review doesn't seem to be resolved yet.
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -168,6 +177,84 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
Review Comment:
Additionally, it would be nice to verify the e2e flow with real system setup (perhaps via Breeze) when you have time, thanks!
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -314,10 +324,20 @@ def heartbeat(self) -> None:
self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads)
self.trigger_tasks(open_slots)
+ if self.supports_connection_test and self.queued_connection_tests:
Review Comment:
Could we move the check into `trigger_connection_tests` function to encapsulate the logic?
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3243,6 +3253,92 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ @provide_session
+ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None:
+ """Dispatch pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("core", "max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("core", "connection_test_timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTest.id)).where(
+ ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING])
Review Comment:
How about defining as frozenset in `models.connection_test`?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2851045539
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -249,6 +256,88 @@ 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_404_NOT_FOUND]),
+ 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 must already be saved. Returns a token that can be used
+ to poll for the test result via GET /connections/test-async/{token}.
+ """
+ if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
Review Comment:
Perhaps we could have a small util for the validation here as it's same as `test_connection` route.
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,125 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import Boolean, ForeignKey, Index, String, Text, Uuid
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+if TYPE_CHECKING:
+ from airflow.models.callback import Callback
+
+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
+
+
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+# Path used by ExecutorCallback to locate the worker function.
+RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test"
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via ExecutorCallback."""
+
+ __tablename__ = "connection_test"
+
+ id: Mapped[UUID] = mapped_column(Uuid(), primary_key=True, default=uuid6.uuid7)
+ token: Mapped[str] = mapped_column(String(64), nullable=False, unique=True)
Review Comment:
I see. I just checked the dev mailing list for more context and you’re right. Thanks for the clarification!
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,155 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import ForeignKey, Index, String, Text, Uuid
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+if TYPE_CHECKING:
+ from airflow.models.callback import Callback, ExecutorCallback
+
+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
+
+
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+# Path used by ExecutorCallback to locate the worker function.
+RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test"
+
+
+class _ImportPathCallbackDef:
+ """
+ Minimal implementation of ImportPathExecutorCallbackDefProtocol.
+
+ ExecutorCallback.__init__ expects an object satisfying this protocol, but no concrete
+ implementation exists in airflow-core — the only one (SyncCallback) lives in the task-sdk.
+ Once #61153 lands and ExecuteCallback.make() is decoupled from DagRun, this adapter can
+ be replaced with the proper factory method.
+ """
+
+ def __init__(self, path: str, kwargs: dict, executor: str | None = None):
+ self.path = path
+ self.kwargs = kwargs
+ self.executor = executor
+
+ def serialize(self) -> dict:
+ return {"path": self.path, "kwargs": self.kwargs, "executor": self.executor}
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via ExecutorCallback."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+
+ callback_id: Mapped[UUID | None] = mapped_column(
+ Uuid(), ForeignKey("callback.id", ondelete="SET NULL"), nullable=True
+ )
+ callback: Mapped[Callback | None] = relationship("Callback", uselist=False)
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(self, *, connection_id: str, **kwargs):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return (
+ f"<ConnectionTest token={self.token!r} connection_id={self.connection_id!r} state={self.state}>"
Review Comment:
Do we need to show the `id` column here as well?
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -249,6 +256,88 @@ 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_404_NOT_FOUND]),
+ 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 must already be saved. Returns a token that can be used
+ to poll for the test result via GET /connections/test-async/{token}.
+ """
+ if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ "Testing connections is disabled in Airflow configuration. "
+ "Contact your deployment admin to enable it.",
+ )
+
+ try:
+ Connection.get_connection_from_secrets(test_body.connection_id)
+ except AirflowNotFoundException:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ f"The Connection with connection_id: `{test_body.connection_id}` was not found. "
+ "Connection must be saved before testing.",
+ )
+
+ connection_test = ConnectionTest(connection_id=test_body.connection_id)
+ session.add(connection_test)
+ session.flush()
+
+ callback = connection_test.create_callback()
+ session.add(callback)
+ session.flush()
+
+ connection_test.callback_id = callback.id
+ connection_test.state = ConnectionTestState.QUEUED
+ callback.queue()
+
+ return ConnectionTestQueuedResponse(
+ token=connection_test.token,
+ connection_id=connection_test.connection_id,
+ state=connection_test.state,
+ )
+
+
+@connections_router.get(
+ "/test-async/{token}",
Review Comment:
How about renaming `token` to `connection_test_token` here and for the `get_connection_test_status` route, since it might be ambiguous for users on the Iterative OpenAPI site?
It’s fine for `POST test_connection_async` because the `token` field is wrapped in a `ConnectionTestQueuedResponse` body.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2842138493
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,125 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import Boolean, ForeignKey, Index, String, Text, Uuid
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+if TYPE_CHECKING:
+ from airflow.models.callback import Callback
+
+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
+
+
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+# Path used by ExecutorCallback to locate the worker function.
+RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test"
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via ExecutorCallback."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_status: Mapped[bool | None] = mapped_column(Boolean, nullable=True)
Review Comment:
Removed it in the latest push. The state already captures SUCCESS/FAILED, so it was redundant. Thanks for catching that.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2851599684
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -249,6 +256,88 @@ 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_404_NOT_FOUND]),
+ 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 must already be saved. Returns a token that can be used
+ to poll for the test result via GET /connections/test-async/{token}.
+ """
+ if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ "Testing connections is disabled in Airflow configuration. "
+ "Contact your deployment admin to enable it.",
+ )
+
+ try:
+ Connection.get_connection_from_secrets(test_body.connection_id)
+ except AirflowNotFoundException:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ f"The Connection with connection_id: `{test_body.connection_id}` was not found. "
+ "Connection must be saved before testing.",
+ )
+
+ connection_test = ConnectionTest(connection_id=test_body.connection_id)
+ session.add(connection_test)
+ session.flush()
+
+ callback = connection_test.create_callback()
+ session.add(callback)
+ session.flush()
+
+ connection_test.callback_id = callback.id
+ connection_test.state = ConnectionTestState.QUEUED
+ callback.queue()
+
+ return ConnectionTestQueuedResponse(
+ token=connection_test.token,
+ connection_id=connection_test.connection_id,
+ state=connection_test.state,
+ )
+
+
+@connections_router.get(
+ "/test-async/{token}",
Review Comment:
renamed thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3292563283
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -259,6 +283,94 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
+ "/enqueue-test",
+ 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 enqueue_connection_test(
Review Comment:
The POST endpoint authorizes against the `team_name` in `test_body`, but the existing `Connection` row identified by `connection_id` may belong to a different team. Combined with `commit_on_success=True`, a user with POST rights on team A can submit a test for a team B `connection_id` with new credentials, and on a successful test the row gets overwritten with team A's payload (cross-team privilege escalation).
Resolve the team from the row before authz:
```python
actual_team = Connection.get_team_name(test_body.connection_id, session)
requires_access_connection("POST", team_name=actual_team)
```
and reject if `test_body.team_name` is set and doesn't match.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -259,6 +283,94 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
+ "/enqueue-test",
+ 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 enqueue_connection_test(
+ test_body: ConnectionTestRequestBody,
+ session: SessionDep,
+) -> ConnectionTestQueuedResponse:
+ """Enqueue a connection test for deferred execution on a worker; returns a polling 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)
+ try:
+ session.flush()
+ except IntegrityError:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ f"An active connection test already exists for connection_id `{test_body.connection_id}`.",
+ )
+
+ return ConnectionTestQueuedResponse(
+ token=connection_test.token,
+ connection_id=connection_test.connection_id,
+ state=connection_test.state,
+ )
+
+
+@connections_router.get(
+ "/enqueue-test/{connection_test_token}",
Review Comment:
Putting the worker callback token in the URL path means it ends up in API server access logs, reverse-proxy logs, browser history, and any intermediate observability that captures request URIs. A token sitting in a URL is treated as low-sensitivity by every layer between the worker and the API, but it grants the same authority as a JWT in an `Authorization` header.
Two options:
1. Move the token to an `Authorization: Bearer ...` header (or a POST body field) and resolve the request by `request_id` alone.
2. Keep the URL stable but store only a SHA-256 hash of the token, then accept the plaintext via header and compare hashes server-side.
Option 2 is friendlier if the worker has to retry without re-issuing tokens, but option 1 is simpler.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3245,6 +3261,94 @@ def _cleanup_orphaned_asset_state(*, session: Session) -> None:
)
session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+ def _enqueue_connection_tests(self, *, session: Session) -> None:
+ """
+ Enqueue pending connection tests to executors that support them.
+
+ ``max_concurrency`` is per-scheduler, not global: with N HA schedulers
+ the worst-case per-tick dispatch is ``N * max_concurrency``. Connection
+ tests are user-initiated and rare, so the overshoot self-corrects via
+ the reaper. For a true global cap, wrap the budget+claim below in a
+ sentinel-row ``SELECT ... FOR UPDATE``.
+ """
+ max_concurrency = conf.getint("connection_test", "max_concurrency", fallback=4)
+ timeout = conf.getint("connection_test", "timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTestRequest.id)).where(
+ ConnectionTestRequest.state.in_(DISPATCHED_STATES)
+ )
+ )
+ budget = max_concurrency - (active_count or 0)
+ 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)
Review Comment:
`_get_workload_team_name` calls `ct.get_dag_id()`, which returns `None` for `ConnectionTestRequest` because there's no DAG behind a connection test. The `team_name` resolution then collapses to `None`, and every connection-test workload routes through the global executor regardless of which team the connection belongs to. The team-scoped executor pool the rest of the PR sets up is bypassed for this workload type.
Resolve the team directly from the row:
```python
team_name = Connection.get_team_name(ct.connection_id, session)
executor = self._try_to_load_executor(executor_name, team_name=team_name)
```
A regression test that spins up two team-scoped executors and asserts a `ConnectionTest` workload for team B lands on team B's executor, not the global one, would lock this in.
##########
airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_05_02.py:
##########
@@ -0,0 +1,31 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
This PR adds `v2026_05_02.py`, which sits between the already-released `v2026_04_06` and the unreleased `v2026_06_16` / `v2026_06_30`. The execution-API versioning convention is that new `VersionChange` entries go into the most recent **unreleased** version file, not a new file dated between releases. Inserting a backdated version breaks the linear migration chain cadwyn replays.
Delete `v2026_05_02.py` and add `AddConnectionTestEndpoint` to `v2026_06_30.py` (or whichever is the latest unreleased file at merge time).
##########
airflow-core/src/airflow/models/crypto.py:
##########
@@ -93,6 +96,65 @@ 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()
+ self.is_encrypted = fernet.is_encrypted
+ else:
Review Comment:
The pre-PR `Connection.set_password` was a no-op when `value` was falsy. The mixin version here adds an `else` branch that resets `self._password = None` and `self.is_encrypted = False`. Any caller flowing `password=None` (typical for a PATCH with an update mask that omits the password, or a sparse update path that explicitly sends `null`) now wipes the existing encrypted password instead of leaving it alone.
Either restore the no-op semantics (return early when `not value`), or treat this as a behavior change with a newsfragment and confirm every PATCH/PUT path in the connections API filters `None` before calling the setter.
A unit test that round-trips `conn.set_password("secret"); conn.set_password(None); assert conn.password == "secret"` would lock this in.
##########
task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py:
##########
@@ -0,0 +1,87 @@
+# 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 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.exceptions import AirflowTaskTimeout
+from airflow.sdk.execution_time.timeout import TimeoutPosix
+
+__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."""
+ client = Client(base_url=server, token=token)
+
+ try:
+ r = client.connection_tests.get_connection(connection_test_id)
+
+ conn = SDKConnection(
+ conn_id=r.conn_id,
+ conn_type=r.conn_type,
+ host=r.host,
+ login=r.login,
+ password=r.password,
+ schema=r.schema_,
+ port=r.port,
+ extra=r.extra,
+ )
+ with TimeoutPosix(
+ seconds=timeout,
+ error_message=f"Connection test timed out after {timeout}s",
+ ):
+ success, message = conn.test_connection()
Review Comment:
The `_preset_connections` ContextVar short-circuits `_get_connection(conn_id)` from the preset, which works as long as `hook.test_connection()` resolves credentials through `BaseHook.get_connection`. But several real-world hooks resolve their own `conn_id` differently:
- SSH/AWS hooks often read `AIRFLOW_CONN_<ID>` directly in `__init__` and store the resolved fields on the hook instance.
- Anything that goes through `BaseHook.get_connection_from_secrets` hits the configured secrets backend (HashiCorp Vault, AWS Secrets Manager, etc.) before falling back to the metadata DB.
In both cases the preset is silently ignored and `test_connection()` ends up testing whatever the real env/backend says, which can be a different credential (or none at all). The previous in-process implementation set `os.environ[conn_env_var] = conn.get_uri()` exactly to defend against this.
Can you verify with at least one hook in each class (SSH or AWS for the env-var case, any provider that resolves via secrets backend for the second), and if they pick the wrong creds, add an env-var override in the worker supervisor before invoking the hook?
##########
task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py:
##########
@@ -0,0 +1,87 @@
+# 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 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.exceptions import AirflowTaskTimeout
+from airflow.sdk.execution_time.timeout import TimeoutPosix
+
+__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."""
+ client = Client(base_url=server, token=token)
+
+ try:
+ r = client.connection_tests.get_connection(connection_test_id)
+
+ conn = SDKConnection(
+ conn_id=r.conn_id,
+ conn_type=r.conn_type,
+ host=r.host,
+ login=r.login,
+ password=r.password,
+ schema=r.schema_,
+ port=r.port,
+ extra=r.extra,
+ )
+ with TimeoutPosix(
Review Comment:
Following up on the prior SIGALRM/Windows concern. Renaming to `TimeoutPosix` makes the platform constraint explicit, thanks. But the import of `signal.SIGALRM` inside `timeout.py` is still unconditional, so a Windows worker process raises `AttributeError` at import time before it can even decide whether to call the timeout. A `hasattr(signal, "SIGALRM")` gate at the import site, or a `threading.Timer` fallback, closes the loop.
If Windows isn't a supported worker platform for this feature, please document that on `supports_connection_test` (and ideally fail fast at supervisor start rather than at the first connection test) so users understand the boundary.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2932379867
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +331,79 @@ 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_404_NOT_FOUND]),
+ dependencies=[Depends(requires_access_connection(method="POST")), Depends(action_logging())],
Review Comment:
You need 'create' connection Perm to 'test' a connection? Someone that only want to 'update a connection' cannot test it?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2897389432
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Maybe the queue_type should also include something to indicate queue priority and whether that queue is intended to be used FIFO or prioritized?
Something vaguely like
If we add a NamedTuple called `WorkloadQueueDef(scheduling_tier: int, sort_key: int)` into the base workload and they are defined like this in the individual workloads:
```python
class TestConnection:
@property
def queue_def(self) -> WorkloadQueueDef :
return WorkloadQueueDef(
scheduling_tier = -1, # connection tests are super fast so let's always get them out of the way first
sort_key = 0, # FIFO order
)
class ExecuteCallback:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
scheduling_tier = 0, # always finish a callback before starting a new task, but don't block the much faster fast connection tests for a callback
sort_key = 0, # FIFO order
)
class ExecuteTask:
@property
def queue_def(self) -> WorkloadQueueDef:
return WorkloadQueueDef(
# after all connection tests and callbacks are running, take the task with the highest priority
scheduling_tier = 1,
sort_key = self.ti.priority_weight,
)
```
((IN fact, we can set them to both default to 0 and that gets even more simplified.... but you get the idea))
and replace all of the "supports_callbacks", "supports_connection_test", etc in each executor can be replaced by a `supported_workload_types: set`. in LocalExecutr it might look like `supported_workload_types = {"ExecuteTask", "ExecuteCallback", "TestConnection"}`
then the entire chunk of code in _get_workloads_to_schedule that sorts out which jobs to pick up would be something real simple like
```python
def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]:
# create a meta_queue which consists of each sub-queue, sorted by that sub-queue's defined sort_key in the order of the scheduling tiers. Then take the top `open_slots` number of jobs off the list.
all_workloads: list[ExecutorWorkload] = [
workload
for queue in self.executor_queues.values()
for workload in queue.values()
]
# Sort by scheduling tier first, then by sort_key within each tier.
all_workloads.sort(key=lambda workload: workload.queue_def)
return [(workload.key, workload) for workload in all_workloads[:open_slots]]
```
As a bonus, this also assures that future workload types honor slots by default; if some future workload type wanted to break that, it would have to be done intentionally. It also means `slots_available` and `slots_occupied` turn into simple comprehensions along the lines of `sum(len(q) for q in self.executor_queues.values())` that don't have to be updated with every new workload. The code overall gets much cleaner.
I'm not sure, I haven't thought that through fully, but as we add more workload types we may need to get that sorted out before long. I have some rough ideas but they were all overkill for only having two types at the time so i figured this was a good place to get those thoughts "on paper". I'll resolve this comment so nobody sees it as blocking, but feel free to message me on Slack and/or tag me in the PR if you get to it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4111922904
@anishgirianish This PR has a few issues that need to be addressed before it can be reviewed — please see our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :warning: **Unresolved review comments**: This PR has 2 unresolved review threads from maintainers: **@jscheffl** (MEMBER): 2 unresolved threads. Please review and resolve all inline review comments before requesting another review. You can resolve a conversation by clicking 'Resolve conversation' on each thread after addressing the feedback. See [pull request guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst).
**What to do next:**
- The comment informs you what you need to do.
- Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
- There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates.
- Maintainers will then proceed with a normal review.
There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971221961
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
Review Comment:
Hi @jscheffl, for every executor to invoke connection test execution, I tried moving the logic out of local_executor.py into executors/workloads/connection_test.py as a shared, executor-agnostic function. This seems like the ideal location since workloads run on workers. However, this causes the check-sdk-imports hook to fail because the function needs from airflow.sdk.api.client import Client to communicate results via the Execution API. Would appreciate your suggestions on how to handle this. Thank you!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971838499
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
Review Comment:
> LocalExecutor is "one" of the executors means it can serve as template for other executors. But including K8s, Celery, Edge, AWS at least 4 additional exist in the repo. So the logic would need to be added to each.
>
> Technically - w/o looking into the code - I see 2 technical options:
>
> * Implement it in the base class (BaseExecutor) and all individual implementations can inherit
> * Implement via a MixIn, then all Executors "just" need to use it... still all executor classes need to be adjusted
>
> Had only time to skim over the PR and have no answer, else I'd directly propose an alternative. But I worry a bit that complexity grows with this approach in each executor while I see in devlist additional features like https://lists.apache.org/thread/o0z8v01v9qq26r6qmvx8zwbkmho1fnbg might need to take the same route.
Hi @jscheffl thank you very much the feedback. I just realzied, this was discussed with (@ferruzzi @jason810496 ) a bit before in this very pr disscussion as well and was actually planned as fast follow up to be using the mechanism following the pattern in this pr https://github.com/apache/airflow/pull/62645 to make it more generic.
Please let me know your thoughts on this. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3128357139
##########
airflow-core/src/airflow/executors/workloads/connection_test.py:
##########
@@ -0,0 +1,87 @@
+# 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
+
+if TYPE_CHECKING:
+ from airflow.api_fastapi.auth.tokens import JWTGenerator
+ from airflow.models.connection_test import ConnectionTestKey, ConnectionTestState
+
+
+class TestConnection(BaseWorkloadSchema):
+ """Execute a connection test on a worker."""
+
+ connection_test_id: uuid.UUID
+ connection_id: str
+ timeout: int = 60
+ 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 str(self.connection_test_id)
Review Comment:
Done. ConnectionTestKey is now a frozen dataclass as suggested, with the executor dicts, state_class_for_key, and scheduler event handling updated to match. Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3128353460
##########
airflow-core/src/airflow/models/connection.py:
##########
@@ -131,7 +131,7 @@ class Connection(Base, LoggingMixin):
host: Mapped[str | None] = mapped_column(String(500), nullable=True)
schema: Mapped[str | None] = mapped_column(String(500), nullable=True)
login: Mapped[str | None] = mapped_column(Text(), nullable=True)
- _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True)
+ # _password and _extra columns inherited from FernetFieldsMixin
Review Comment:
Checked, still works. conn.password = 'foo' flips is_encrypted to True and round-trips (same for extra); Connection's overrides win over the mixin via MRO. Added pinning tests in the latest push.
Thank you.
##########
airflow-core/src/airflow/models/connection.py:
##########
@@ -131,7 +131,7 @@ class Connection(Base, LoggingMixin):
host: Mapped[str | None] = mapped_column(String(500), nullable=True)
schema: Mapped[str | None] = mapped_column(String(500), nullable=True)
login: Mapped[str | None] = mapped_column(Text(), nullable=True)
- _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True)
+ # _password and _extra columns inherited from FernetFieldsMixin
Review Comment:
Checked, still works. conn.password = 'foo' flips is_encrypted to True and round-trips (same for extra); Connection's overrides win over the mixin via MRO. Added pinning tests in the latest push. Thank you.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3025672542
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +286,83 @@ 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]),
+ 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()
+
+ # Only one active test per connection_id at a time.
+ _check_no_active_test(test_body.connection_id, session)
+
+ 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)
+ session.flush()
+
+ 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"))],
+)
+def get_connection_test(
+ connection_test_token: str,
+ session: SessionDep,
+) -> ConnectionTestStatusResponse:
+ """
+ Poll for the status of an async connection test.
+
+ Knowledge of the token serves as authorization — only the client
+ that initiated the test knows the crypto-random token.
Review Comment:
updated
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3029297767
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -169,6 +178,81 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
+ """
+ Execute a connection test workload.
+
+ Constructs an SDK ``Client``, fetches the connection via the Execution API,
+ enforces a timeout via ``signal.alarm``, and reports all outcomes back
+ through the Execution API.
+
+ :param log: Logger instance
+ :param workload: The TestConnection workload to execute
+ :param team_conf: Team-specific executor configuration
+ """
+ # Lazy import: SDK modules must not be loaded at module level to avoid
+ # coupling core (scheduler-loaded) code to the SDK.
+ from airflow.sdk.api.client import Client
+
+ setproctitle(
+ f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}",
+ log,
+ )
+
+ base_url = team_conf.get("api", "base_url", fallback="/")
+ if base_url.startswith("/"):
+ base_url = f"http://localhost:8080{base_url}"
Review Comment:
Hi @pierrejeambrun , thank you for flagging this. Just wanted to note that this follows the same pattern already used in _execute_work (line 148) for task execution, the localhost:8080 is a fallback when [api] base_url isn't configured and returns a relative path. I completely understand the concern, though. Would you have a preferred approach for handling this? Happy to go whichever direction you suggest. Thank you!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177845967
##########
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:
Done, String(2000) in model + migration, max_length=2000 on ConnectionTestResultBody.result_message. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177829516
##########
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:
Implemented your suggested shape , SDKConnection.test_connection() with _preset_connections encapsulated, Thank you so much
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177837762
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2602,6 +2602,33 @@ scheduler:
type: float
example: ~
default: "120.0"
+ connection_test_timeout:
Review Comment:
Done, moved to a new [connection_test] section with renamed keys (timeout, max_concurrency, reaper_interval). thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177848529
##########
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:
Done, moved ConnectionTestState import to module level, dropped the three runtime imports + dead TYPE_CHECKING entry.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3178763816
##########
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:
Done, added `extra_filters` to `_TableConfig` and applied `state IN ('success','failed')` so in-flight rows are never selectable for deletion regardless of retention. `recency_column` stays at `updated_at` to also respect the retention window from the last state change.
##########
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:
Done, added `extra_filters` to `_TableConfig` and applied `state IN ('success','failed')` so in-flight rows are never selectable for deletion regardless of retention. `recency_column` stays at `updated_at` to also respect the retention window from the last state change. Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3130099366
##########
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
+ commit_on_success: bool = Field(
+ default=False,
+ description="If True, save or update the connection in the connection table when the test succeeds.",
+ )
+ executor: str | None = None
Review Comment:
We probably want to validate the `executor` against the configured executors list.
No need to do a full roundtrip if the executor isn't valid, we can just reject at the API layer level with a 422
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3025380683
##########
airflow-core/src/airflow/migrations/versions/0111_3_2_0_add_connection_test_table.py:
##########
@@ -0,0 +1,77 @@
+#
+# 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: a4c2d171ae18
+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 = "a4c2d171ae18"
+branch_labels = None
+depends_on = None
+airflow_version = "3.2.0"
Review Comment:
This migration is chained after revision `a4c2d171ae18` (which is marked as `airflow_version = "3.3.0"` in `0110_3_3_0_xcom_dag_result.py`), but this file sets `airflow_version = "3.2.0"`. That version label (and the file name `0111_3_2_0_...`) conflicts with the actual alembic chain and results in `3.2.0` appearing *after* `3.3.0` in the revision-head mapping/docs. Please align `airflow_version` (and the migration filename if needed) with the next version after `3.3.0`, or adjust the `down_revision` if this truly belongs in the `3.2.0` chain.
```suggestion
airflow_version = "3.4.0"
```
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +286,83 @@ 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]),
+ 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()
+
+ # Only one active test per connection_id at a time.
+ _check_no_active_test(test_body.connection_id, session)
+
+ 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)
+ session.flush()
Review Comment:
The "only one active test per connection_id" enforcement is a non-atomic check-then-insert (`_check_no_active_test()` followed by `session.add/flush()`). Two concurrent POSTs can both pass the check and create multiple active rows for the same `connection_id`, contradicting the API contract and making later 409 handling unreliable. Consider enforcing this at the DB level (unique constraint/index) or using a transactional locking strategy so the check and insert are serialized for a given `connection_id`.
##########
airflow-core/src/airflow/utils/db.py:
##########
@@ -115,7 +115,7 @@ class MappedClassProtocol(Protocol):
"3.0.3": "fe199e1abd77",
"3.1.0": "cc92b33c6709",
"3.1.8": "509b94a1042d",
- "3.2.0": "1d6611b6ab7c",
+ "3.2.0": "a7e6d4c3b2f1",
"3.3.0": "a4c2d171ae18",
Review Comment:
The `_REVISION_HEADS_MAP` entry for `3.2.0` now points to revision `a7e6d4c3b2f1`, but that revision’s migration file declares `down_revision = "a4c2d171ae18"` (a `3.3.0` migration). This makes the mapping non-monotonic (`3.2.0` > `3.3.0`) and will break logic that relies on version→head ordering. The new migration should be mapped to the *next* Airflow version after `3.3.0` (and `3.2.0` should continue to map to `1d6611b6ab7c`).
```suggestion
"3.2.0": "1d6611b6ab7c",
"3.3.0": "a4c2d171ae18",
"3.4.0": "a7e6d4c3b2f1",
```
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +286,83 @@ 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]),
+ 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()
+
+ # Only one active test per connection_id at a time.
+ _check_no_active_test(test_body.connection_id, session)
+
+ 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)
+ session.flush()
+
+ 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"))],
+)
+def get_connection_test(
+ connection_test_token: str,
+ session: SessionDep,
+) -> ConnectionTestStatusResponse:
+ """
+ Poll for the status of an async connection test.
+
+ Knowledge of the token serves as authorization — only the client
+ that initiated the test knows the crypto-random token.
Review Comment:
The docstring claims that "knowledge of the token serves as authorization", but the route also enforces `requires_access_connection(method="GET")` (and therefore still requires normal authz). This is misleading for API consumers and could encourage unsafe assumptions. Please reword to clarify that the token is an identifier/capability for the test result, while the endpoint still requires authenticated/authorized access.
```suggestion
The `connection_test_token` is a crypto-random identifier used to look up
the specific test result. Access to this endpoint is still governed by
standard authentication and connection-level authorization
(see `requires_access_connection(method="GET")`).
```
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -305,10 +315,24 @@ def heartbeat(self) -> None:
self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads)
self.trigger_tasks(open_slots)
+ self.trigger_connection_tests()
+
# Calling child class sync method
self.log.debug("Calling the %s sync method", self.__class__)
self.sync()
+ def trigger_connection_tests(self) -> None:
+ """Process queued connection tests, respecting available slot capacity."""
+ if not self.supports_connection_test or not self.queued_connection_tests:
+ return
+
+ available = self.slots_available
+ if available <= 0:
+ return
+
+ tests_to_run = list(self.queued_connection_tests.values())[:available]
+ self._process_workloads(tests_to_run)
Review Comment:
`trigger_connection_tests()` uses `self.slots_available` to decide how many queued tests to run. Since `slots_available` subtracts `len(self.queued_connection_tests)`, this can deadlock when the queue reaches `parallelism` (available becomes 0, so nothing is ever processed). It can also under-schedule connection tests whenever tasks/callbacks are queued, because queued items reduce `slots_available` even though they still need to be executed. Compute capacity based on `parallelism - len(self.running)` (like `heartbeat()` does for tasks) and schedule up to the remaining open slots after `trigger_tasks()` has run.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3122,6 +3136,86 @@ 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("core", "max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("core", "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 not None and not executor.supports_connection_test:
+ executor = None
+ if executor is None:
+ reason = (
+ f"No executor matches '{ct.executor}'"
+ if ct.executor
+ else "No executor supports 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("core", "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),
Review Comment:
`_reap_stale_connection_tests()` reaps rows in `ACTIVE_STATES`, which currently includes `PENDING`. This means tests that are merely waiting for capacity (never dispatched) can be marked FAILED after `timeout+grace`, contradicting the doc/config text that targets stuck `QUEUED`/`RUNNING` tests and potentially causing false failures under load. Consider reaping only `DISPATCHED_STATES` (or at least excluding `PENDING`).
```suggestion
ConnectionTestRequest.state.in_(DISPATCHED_STATES),
```
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -867,6 +869,25 @@ def get_detail_response(self, ti_id: uuid.UUID) -> HITLDetailResponse:
return HITLDetailResponse.model_validate_json(resp.read())
+class ConnectionTestOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def get_connection(self, connection_test_id: uuid.UUID) -> ConnectionResponse:
+ """Fetch connection data for a test request from the API server."""
+ resp = self.client.get(f"connection-tests/{connection_test_id}/connection")
+ return ConnectionResponse.model_validate_json(resp.read())
+
Review Comment:
`ConnectionTestOperations.get_connection()` is typed/validated as `ConnectionResponse`, but the execution API endpoint returns `ConnectionTestConnectionResponse` (and the SDK already generated that model). Using the dedicated model here would better reflect the endpoint contract and avoid accidental coupling if `ConnectionResponse` later diverges (e.g., additional runtime-only fields).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3244740574
##########
airflow-core/src/airflow/models/connection.py:
##########
@@ -131,7 +131,7 @@ class Connection(Base, LoggingMixin):
host: Mapped[str | None] = mapped_column(String(500), nullable=True)
schema: Mapped[str | None] = mapped_column(String(500), nullable=True)
login: Mapped[str | None] = mapped_column(Text(), nullable=True)
- _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True)
+ # _password and _extra columns inherited from FernetFieldsMixin
Review Comment:
Sorry about that, thanks for verifying.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3157908409
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -57,24 +60,53 @@
from airflow.configuration import conf
from airflow.exceptions import AirflowNotFoundException
from airflow.models import Connection
+from airflow.models.connection_test import ACTIVE_STATES, ConnectionTestRequest
from airflow.secrets.environment_variables import CONN_ENV_PREFIX
from airflow.utils.db import create_default_connections as db_create_default_connections
from airflow.utils.strings import get_random_string
connections_router = AirflowRouter(tags=["Connection"], prefix="/connections")
+def _ensure_test_connection_enabled() -> None:
+ """Raise 403 if connection testing is not enabled in the Airflow configuration."""
+ if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ "Testing connections is disabled in Airflow configuration. "
+ "Contact your deployment admin to enable it.",
+ )
+
+
+def _check_no_active_test(connection_id: str, session: SessionDep) -> None:
+ """Raise 409 if there is an active connection test request for the given connection_id."""
+ active_test = session.scalar(
+ select(ConnectionTestRequest).filter(
+ ConnectionTestRequest.connection_id == connection_id,
+ ConnectionTestRequest.state.in_(ACTIVE_STATES),
+ )
+ )
+ if active_test is not None:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ f"Cannot modify connection `{connection_id}` while an async test is running. "
+ "This typically takes only a few seconds — please retry shortly.",
+ )
+
Review Comment:
Dropped the app-layer guard. The DB now enforces it: a single column tracks "active test for this connection" and a unique constraint blocks duplicates. Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3157900830
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3141,6 +3160,86 @@ 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 not None and not executor.supports_connection_test:
+ executor = None
+ if executor is None:
+ reason = (
+ f"No executor matches '{ct.executor}'"
+ if ct.executor
+ else "No executor supports connection testing"
+ )
+ ct.state = ConnectionTestState.FAILED
Review Comment:
Reworded to name the executor that was actually tried. Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3169352729
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -872,6 +874,25 @@ def get_detail_response(self, ti_id: uuid.UUID) -> HITLDetailResponse:
return HITLDetailResponse.model_validate_json(resp.read())
+class ConnectionTestOperations:
+ __slots__ = ("client",)
+
+ def __init__(self, client: Client):
+ self.client = client
+
+ def get_connection(self, connection_test_id: uuid.UUID) -> ConnectionResponse:
+ """Fetch connection data for a test request from the API server."""
+ resp = self.client.get(f"connection-tests/{connection_test_id}/connection")
+ return ConnectionResponse.model_validate_json(resp.read())
Review Comment:
Return type is `ConnectionResponse` but the endpoint (`GET /connection-tests/{id}/connection`) returns `ConnectionTestConnectionResponse`. Both models happen to have the same fields today so it works at runtime, but the annotation is wrong and will silently break if the two models ever diverge.
```suggestion
def get_connection(self, connection_test_id: uuid.UUID) -> ConnectionTestConnectionResponse:
"""Fetch connection data for a test request from the API server."""
resp = self.client.get(f"connection-tests/{connection_test_id}/connection")
return ConnectionTestConnectionResponse.model_validate_json(resp.read())
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971256512
##########
airflow-core/newsfragments/62343.feature.rst:
##########
Review Comment:
Thank you. for sure I agree. removed. thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2963468962
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via a TestConnection workload."""
+
+ __tablename__ = "connection_test"
+
+ 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(10), nullable=False, default=ConnectionTestState.PENDING)
+ result_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True)
+ reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0")
+
+ __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),)
+
+ def __init__(
+ self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return f"<ConnectionTest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+
+def run_connection_test(*, conn: Connection) -> tuple[bool, str]:
+ """
+ Worker-side function to execute a connection test.
+
+ Returns a (success, message) tuple. The caller is responsible for
+ reporting the result back via the Execution API.
+ """
+ try:
+ return conn.test_connection()
+ except Exception as e:
+ log.exception("Connection test failed", connection_id=conn.conn_id)
+ return False, str(e)
+
+
+_SNAPSHOT_FIELDS = (
+ "conn_type",
+ "description",
+ "host",
+ "login",
+ "_password",
+ "schema",
+ "port",
+ "_extra",
+ "is_encrypted",
+ "is_extra_encrypted",
+)
+
+
+def snapshot_connection(conn: Connection) -> dict:
Review Comment:
`snapshot_connection()` stores Fernet-encrypted `_password` and `_extra` ciphertext in a JSON column. This expands the attack surface — anyone with read access to the `connection_test` table gets credential ciphertext that was previously confined to the `connection` table.
Two additional problems:
1. If Fernet keys rotate, stored snapshots become undecryptable and reverts will write garbage.
2. DB-level audit/masking rules on `connection._password` won't cover this secondary storage.
If the snapshot approach survives (vs. Pierre's request-buffer redesign), consider storing only non-secret fields and re-fetching credentials at revert time, or encrypting the entire snapshot blob with a short-lived key.
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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.QUEUED, ConnectionTestState.RUNNING))
Review Comment:
`ACTIVE_STATES` includes QUEUED and RUNNING but not PENDING. The reaper (`_reap_stale_connection_tests`) only checks `ACTIVE_STATES`, so tests stuck in PENDING — e.g., scheduler restart between creation and dispatch, or no supporting executor found — accumulate forever.
If those PENDING tests have snapshots (from save-and-test), the connection stays in its post-edit state with no resolution path.
```suggestion
ACTIVE_STATES = frozenset((ConnectionTestState.PENDING, ConnectionTestState.QUEUED, ConnectionTestState.RUNNING))
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2976160287
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
Review Comment:
thank you,
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3025677085
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3122,6 +3136,86 @@ 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("core", "max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("core", "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 not None and not executor.supports_connection_test:
+ executor = None
+ if executor is None:
+ reason = (
+ f"No executor matches '{ct.executor}'"
+ if ct.executor
+ else "No executor supports 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("core", "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),
Review Comment:
This is intentional. The reaper cutoff is timeout + grace_period (default 90 seconds). A PENDING test that has been sitting that long means the scheduler never picked it up something is wrong. Without reaping PENDING rows, they would block the user's polling endpoint indefinitely and prevent future tests for that connection_id (due to the "only one active test" check). Using DISPATCHED_STATES would leave stuck PENDING rows orphaned with no cleanup path.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3103424589
##########
airflow-core/src/airflow/executors/workloads/connection_test.py:
##########
@@ -0,0 +1,87 @@
+# 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
+
+if TYPE_CHECKING:
+ from airflow.api_fastapi.auth.tokens import JWTGenerator
+ from airflow.models.connection_test import ConnectionTestKey, ConnectionTestState
+
+
+class TestConnection(BaseWorkloadSchema):
+ """Execute a connection test on a worker."""
+
+ connection_test_id: uuid.UUID
+ connection_id: str
+ timeout: int = 60
+ 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 str(self.connection_test_id)
Review Comment:
I know it is late for a big change, but one thing I am finding with the recent Workload work is that having the Callback Key be a str is really limiting, and adding a second Workload with a string as a key is potentially going to cause an issue. Can we do this instead:
```
@dataclass(frozen=True, slots=True)
class ConnectionTestKey:
id: str # str(UUID)
def __str__(self) -> str:
return self.id
```
It's something I regret not doing in the Callback because now isinstance(key, CallbackKey) will match on any string that is passed.
##########
airflow-core/src/airflow/models/connection.py:
##########
@@ -131,7 +131,7 @@ class Connection(Base, LoggingMixin):
host: Mapped[str | None] = mapped_column(String(500), nullable=True)
schema: Mapped[str | None] = mapped_column(String(500), nullable=True)
login: Mapped[str | None] = mapped_column(Text(), nullable=True)
- _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True)
+ # _password and _extra columns inherited from FernetFieldsMixin
Review Comment:
Does this break connection.is_encrypted and is_extra_encrypted?? Could you check and make sure that is_encrypted is set correctly after calling conn.password = 'foo'?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3129903590
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3141,6 +3160,86 @@ 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 not None and not executor.supports_connection_test:
+ executor = None
+ if executor is None:
+ reason = (
+ f"No executor matches '{ct.executor}'"
+ if ct.executor
+ else "No executor supports connection testing"
+ )
+ ct.state = ConnectionTestState.FAILED
Review Comment:
Message is confusing. It looks like 'no executor supports testing' while it's more like the executor attached to this workload doesn't support testing. But in a multi executor env, other executor could support testing.
(I specify celeryexecutor for the test, that won't work, but I also have the local executor available)
##########
airflow-core/src/airflow/models/connection.py:
##########
@@ -131,7 +131,7 @@ class Connection(Base, LoggingMixin):
host: Mapped[str | None] = mapped_column(String(500), nullable=True)
schema: Mapped[str | None] = mapped_column(String(500), nullable=True)
login: Mapped[str | None] = mapped_column(Text(), nullable=True)
- _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True)
+ # _password and _extra columns inherited from FernetFieldsMixin
Review Comment:
We need to remove that comment I guess.
##########
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
+ commit_on_success: bool = Field(
+ default=False,
+ description="If True, save or update the connection in the connection table when the test succeeds.",
+ )
+ executor: str | None = None
Review Comment:
We probably want to validate the `executor` against the configured executors list.
No need to to a full roundtrip if the executor isn't valid, we can just reject at the API layer level with a 422
##########
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 = 60
Review Comment:
To remove. This should come from the config and always provided.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -57,24 +60,53 @@
from airflow.configuration import conf
from airflow.exceptions import AirflowNotFoundException
from airflow.models import Connection
+from airflow.models.connection_test import ACTIVE_STATES, ConnectionTestRequest
from airflow.secrets.environment_variables import CONN_ENV_PREFIX
from airflow.utils.db import create_default_connections as db_create_default_connections
from airflow.utils.strings import get_random_string
connections_router = AirflowRouter(tags=["Connection"], prefix="/connections")
+def _ensure_test_connection_enabled() -> None:
+ """Raise 403 if connection testing is not enabled in the Airflow configuration."""
+ if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ "Testing connections is disabled in Airflow configuration. "
+ "Contact your deployment admin to enable it.",
+ )
+
+
+def _check_no_active_test(connection_id: str, session: SessionDep) -> None:
+ """Raise 409 if there is an active connection test request for the given connection_id."""
+ active_test = session.scalar(
+ select(ConnectionTestRequest).filter(
+ ConnectionTestRequest.connection_id == connection_id,
+ ConnectionTestRequest.state.in_(ACTIVE_STATES),
+ )
+ )
+ if active_test is not None:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ f"Cannot modify connection `{connection_id}` while an async test is running. "
+ "This typically takes only a few seconds — please retry shortly.",
+ )
+
Review Comment:
Instead of having this we should rely on a db unique key constraint. This can still produce race conditions and is not atomic. I would remove this and find a db side solution for mysql.
##########
airflow-core/src/airflow/models/crypto.py:
##########
@@ -93,6 +96,66 @@ def rotate(self, msg: bytes | str) -> bytes:
return self._fernet.rotate(msg)
+class FernetFieldsMixin:
+ """
+ Mixin providing Fernet-encrypted ``password`` and ``extra`` fields.
+
+ Subclasses get ``_password`` / ``_extra`` mapped columns and transparent
+ encrypt-on-write / decrypt-on-read via SQLAlchemy synonyms.
+
+ Models that need additional logic (e.g. the ``Connection`` model's
+ ``is_encrypted`` flag) can override ``get_password``, ``set_password``,
+ ``get_extra``, and ``set_extra``.
Review Comment:
Shoudln't we instead use the same pattern for the TestConnection ? Have the `is_encrypted` `is_extra_encrypted` here and the common logic here too?
This way we get a clear `FERNET_KEY configuration is missing` when `is_encrypted=True` and worker doesn't have a key configured.
(To basically mirror the Connection pattern)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4303949661
#protm
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3128342837
##########
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)
Review Comment:
thanks for raising. I see SIGALRM is already used unconditionally in a few other places (task-sdk/.../timeout.py, utils/db.py, utils/helpers.py, dag_processing/importers/python_importer.py), so I followed the existing pattern here. Happy to change if you'd prefer. Please let me know
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3157912925
##########
airflow-core/src/airflow/models/crypto.py:
##########
@@ -93,6 +96,66 @@ def rotate(self, msg: bytes | str) -> bytes:
return self._fernet.rotate(msg)
+class FernetFieldsMixin:
+ """
+ Mixin providing Fernet-encrypted ``password`` and ``extra`` fields.
+
+ Subclasses get ``_password`` / ``_extra`` mapped columns and transparent
+ encrypt-on-write / decrypt-on-read via SQLAlchemy synonyms.
+
+ Models that need additional logic (e.g. the ``Connection`` model's
+ ``is_encrypted`` flag) can override ``get_password``, ``set_password``,
+ ``get_extra``, and ``set_extra``.
Review Comment:
Done, FernetFieldsMixin now owns is_encrypted/is_extra_encrypted and raises ValueError("…FERNET_KEY configuration is missing"). Connection drops its duplicates; ConnectionTestRequest inherits.
thank you
##########
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 = 60
Review Comment:
Removed thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3157921772
##########
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
+ commit_on_success: bool = Field(
+ default=False,
+ description="If True, save or update the connection in the connection table when the test succeeds.",
+ )
+ executor: str | None = None
Review Comment:
Done, added _ensure_executor_is_configured route helper that checks against the configured executors list.
thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3292564464
##########
airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_05_02.py:
##########
@@ -0,0 +1,31 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
**Cadwyn version is backdated; new changes should land in the latest unreleased file.**
This PR adds `v2026_05_02.py`, which sits between the already-released `v2026_04_06` and the unreleased `v2026_06_16`/`v2026_06_30`. The execution-API versioning convention is that new `VersionChange` entries go into the most recent **unreleased** version file, not a new file dated between releases. Inserting a backdated version breaks the linear migration chain cadwyn replays.
Delete `v2026_05_02.py` and add `AddConnectionTestEndpoint` to `v2026_06_30.py` (or whichever is currently the latest unreleased file at merge time).
##########
airflow-core/src/airflow/models/crypto.py:
##########
@@ -93,6 +96,65 @@ 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()
+ self.is_encrypted = fernet.is_encrypted
+ else:
Review Comment:
**Extracting `set_password` to the mixin silently changes falsy-value semantics.**
The pre-PR `Connection.set_password` was a no-op when `value` was falsy. The mixin version here adds an `else` branch that resets `self._password = None` and `self.is_encrypted = False`. Any caller flowing `password=None` (typical for a PATCH with an update mask that omits the password, or a sparse update path that explicitly sends `null`) now wipes the existing encrypted password instead of leaving it alone.
Either restore the no-op semantics (return early when `not value`), or treat this as a behavior change with a newsfragment and confirm every PATCH/PUT path in the connections API filters `None` before calling the setter.
A unit test that round-trips `conn.set_password("secret"); conn.set_password(None); assert conn.password == "secret"` would lock this in.
##########
task-sdk/src/airflow/sdk/execution_time/connection_test_supervisor.py:
##########
@@ -0,0 +1,87 @@
+# 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 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.exceptions import AirflowTaskTimeout
+from airflow.sdk.execution_time.timeout import TimeoutPosix
+
+__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."""
+ client = Client(base_url=server, token=token)
+
+ try:
+ r = client.connection_tests.get_connection(connection_test_id)
+
+ conn = SDKConnection(
+ conn_id=r.conn_id,
+ conn_type=r.conn_type,
+ host=r.host,
+ login=r.login,
+ password=r.password,
+ schema=r.schema_,
+ port=r.port,
+ extra=r.extra,
+ )
+ with TimeoutPosix(
+ seconds=timeout,
+ error_message=f"Connection test timed out after {timeout}s",
+ ):
+ success, message = conn.test_connection()
Review Comment:
**`_preset_connections` ContextVar only intercepts the `_get_connection` path; hooks reading env vars or secrets backends bypass it.**
This short-circuits `_get_connection(conn_id)` from the preset, which works as long as `hook.test_connection()` resolves credentials through `BaseHook.get_connection`. But several real-world hooks resolve their own `conn_id` differently:
- SSH/AWS hooks often read `AIRFLOW_CONN_<ID>` directly in `__init__` and store the resolved fields on the hook instance.
- Anything that goes through `BaseHook.get_connection_from_secrets` hits the configured secrets backend (HashiCorp Vault, AWS Secrets Manager, etc.) before falling back to the metadata DB.
In both cases the preset is silently ignored and `test_connection()` ends up testing whatever the real env/backend says, which can be a different credential (or none at all). The previous in-process implementation set `os.environ[conn_env_var] = conn.get_uri()` exactly to defend against this.
Please verify with at least one hook in each class (SSH or AWS for the env-var case, any provider that resolves via secrets backend for the second), and if they pick the wrong creds, add an env-var override in the worker supervisor before invoking the hook.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-3957676083
@jason810496 Thanks for the review! I have addressed it in the latest push.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2932092188
##########
airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml:
##########
@@ -1631,6 +1631,103 @@ paths:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
+ /api/v2/connections/{connection_id}/test:
Review Comment:
We shouldn't "PATCH Connection and test". We are creating a 'TestConnection' (which the user does not care it's internal implementation), and returning 'true/false' for the test.
This should be a classic 'POST' Action. Not a patch.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2932379867
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +331,79 @@ 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_404_NOT_FOUND]),
+ dependencies=[Depends(requires_access_connection(method="POST")), Depends(action_logging())],
Review Comment:
You need 'create' connection Perm to 'test' a connection? Someone that only want to 'update a connection' cannot test it?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177840823
##########
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:
Done via _only_terminal_states validator on ConnectionTestResultBody; kept the enum import-stable for the generated client (Literal-wrapping made codegen drop the enum). thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177838393
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2602,6 +2602,33 @@ scheduler:
type: float
example: ~
default: "120.0"
+ connection_test_timeout:
+ description: |
+ Maximum number of seconds an async connection test is allowed to run
+ before it is considered timed out. The scheduler reaper uses this value
+ plus a grace period to mark stale tests as failed.
+ version_added: 3.3.0
+ type: integer
+ example: ~
+ default: "60"
+ max_connection_test_concurrency:
Review Comment:
Done thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3029282065
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,234 @@
+# 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 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, Uuid, select
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, synonym
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.models.connection import Connection
+from airflow.models.crypto import get_fernet
+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))
+
+
+class ConnectionTestRequest(Base):
+ """
+ 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)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+
+ # Connection fields — password and extra are Fernet-encrypted.
+ conn_type: Mapped[str] = mapped_column(String(500), nullable=False)
+ host: Mapped[str | None] = mapped_column(String(500), nullable=True)
+ login: Mapped[str | None] = mapped_column(Text, nullable=True)
+ _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True)
+ schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True)
+ port: Mapped[int | None] = mapped_column(Integer, nullable=True)
+ _extra: Mapped[str | None] = mapped_column("extra", Text(), nullable=True)
+ commit_on_success: Mapped[bool] = mapped_column(
+ Boolean, nullable=False, default=False, server_default="0"
+ )
+
+ __table_args__ = (Index("idx_connection_test_request_state_created_at", state, created_at),)
+
+ def __init__(
+ self,
+ *,
+ connection_id: str,
+ conn_type: str,
+ host: str | None = None,
+ login: str | None = None,
+ password: str | None = None,
+ schema: str | None = None,
+ port: int | None = None,
+ extra: str | None = None,
+ commit_on_success: bool = False,
+ executor: str | None = None,
+ queue: str | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.connection_id = connection_id
+ self.conn_type = conn_type
+ self.host = host
+ self.login = login
+ self.password = password
+ self.schema = schema
+ self.port = port
+ self.extra = extra
+ self.commit_on_success = commit_on_success
+ self.executor = executor
+ self.queue = queue
+ self.token = secrets.token_urlsafe(32)
+ self.state = ConnectionTestState.PENDING
+
+ def __repr__(self) -> str:
+ return (
+ f"<ConnectionTestRequest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>"
+ )
+
+ def get_executor_name(self) -> str | None:
+ """Return the executor name for scheduler routing."""
+ return self.executor
+
+ def get_dag_id(self) -> None:
+ """Return None — connection tests are not associated with any DAG."""
+ return None
+
+ def get_password(self) -> str | None:
+ if self._password:
+ fernet = get_fernet()
+ if not fernet.is_encrypted:
+ return self._password
+ return fernet.decrypt(bytes(self._password, "utf-8")).decode()
+ return self._password
+
+ def set_password(self, value: str | None):
+ if value:
+ fernet = get_fernet()
+ self._password = fernet.encrypt(bytes(value, "utf-8")).decode()
+ else:
+ self._password = value
+
+ @declared_attr
+ def password(cls):
+ """Password. The value is decrypted/encrypted when reading/setting the value."""
+ return synonym("_password", descriptor=property(cls.get_password, cls.set_password))
+
+ def get_extra(self) -> str | None:
+ if self._extra:
+ fernet = get_fernet()
+ if not fernet.is_encrypted:
+ return self._extra
+ return fernet.decrypt(bytes(self._extra, "utf-8")).decode()
+ return self._extra
+
+ def set_extra(self, value: str | None):
+ if value:
+ fernet = get_fernet()
+ self._extra = fernet.encrypt(bytes(value, "utf-8")).decode()
+ else:
+ self._extra = value
+
+ @declared_attr
+ def extra(cls):
+ """Extra data. The value is decrypted/encrypted when reading/setting the value."""
+ return synonym("_extra", descriptor=property(cls.get_extra, cls.set_extra))
+
Review Comment:
Done, extracted FernetFieldsMixin into models/crypto.py with shared encrypted _password/_extra columns and synonyms.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2842181156
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py:
##########
@@ -0,0 +1,65 @@
+# 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 fastapi import APIRouter, HTTPException, status
+
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.connection_test import ConnectionTestResultBody
+from airflow.models.connection_test import TERMINAL_STATES, ConnectionTest
+
+router = APIRouter()
+
+
Review Comment:
Agreed, thanks for confirming on #61153. Keeping it for now since it's forward-compatible with callbacks eventually running in TaskSDK, at that point, workers would report results through this endpoint instead of updating the DB directly.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2842239235
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,125 @@
+# 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 datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+import structlog
+import uuid6
+from sqlalchemy import Boolean, ForeignKey, Index, String, Text, Uuid
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+if TYPE_CHECKING:
+ from airflow.models.callback import Callback
+
+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
+
+
+TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED))
+
+# Path used by ExecutorCallback to locate the worker function.
+RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test"
+
+
+class ConnectionTest(Base):
+ """Tracks an async connection test dispatched to a worker via ExecutorCallback."""
+
+ __tablename__ = "connection_test"
+
+ id: Mapped[UUID] = mapped_column(Uuid(), primary_key=True, default=uuid6.uuid7)
+ token: Mapped[str] = mapped_column(String(64), nullable=False, unique=True)
Review Comment:
Thanks for the question! UUIDv7 is sequential, so the crypto-random token prevents authenticated users from enumerating other users' connection test results. This was the approach agreed on in the dev mailing list discussion.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-3946352063
@jason810496 Thanks for the thorough review! Addressed your feedback in the latest push:
- Removed `result_status` column — `state` is sufficient
- Moved `_ImportPathCallbackDef` to `connection_test.py` with a `create_callback()` factory method
Could you please take another look when you get a chance? Thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi merged PR #62343:
URL: https://github.com/apache/airflow/pull/62343
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via ExecutorCallback dispatch to workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-3961720903
> > @pierrejeambrun thank you very much for the review. The dev list discussion ([thread](https://lists.apache.org/thread/xd7zmyp95y77cw36mb5wjp17dyynz100)) received positive feedback , and the implementation follows the suggestions outlined in the thread. Happy to wait for a more formal conclusion on the thread. Would love to know if there's anything you'd like me to revisit or improve in the meantime.
>
> I do plan to review proposal, PR and reply but occupied today. So tomorrow or Friday.
>
> There isn't any hurry either since earliest this would go would be 3.2 anyway and there is some time to go for it.
@kaxil Sure thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971228498
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -257,6 +286,83 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
os.environ.pop(conn_env_var, None)
+@connections_router.post(
Review Comment:
Hi @jscheffl, thank you for raising this. It's a valid concern. I wanted to note that the existing sync POST /connections/test route is already in the public API with the same test_connection=Enabled guard and auth. Moving only the async version to UI-only would create an inconsistency where external clients could test synchronously but not asynchronously. Would you prefer both routes to be public (consistent with current state), both moved to UI-only, or is there something specific about the async route that warrants different treatment? Happy to go whichever direction you think is best. Thank you!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971261156
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -305,10 +315,19 @@ def heartbeat(self) -> None:
self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads)
self.trigger_tasks(open_slots)
+ self.trigger_connection_tests()
+
# Calling child class sync method
self.log.debug("Calling the %s sync method", self.__class__)
self.sync()
+ def trigger_connection_tests(self) -> None:
+ """Process queued connection tests."""
+ if not self.supports_connection_test or not self.queued_connection_tests:
+ return
+
+ self._process_workloads(list(self.queued_connection_tests.values()))
Review Comment:
added the check thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2971266605
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -174,6 +189,70 @@ def bulk_connections(
return BulkConnectionService(session=session, request=request).handle_request()
+@connections_router.patch(
+ "/{connection_id}/test",
+ responses=create_openapi_http_exception_doc(
+ [
+ status.HTTP_400_BAD_REQUEST,
+ status.HTTP_403_FORBIDDEN,
+ status.HTTP_404_NOT_FOUND,
+ ]
+ ),
+ dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())],
+)
+def patch_connection_and_test(
Review Comment:
udpdated thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2876261697
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -212,6 +227,69 @@ def patch_connection(
return connection
+@connections_router.patch(
+ "/{connection_id}/save-and-test",
Review Comment:
Would `PUT` (or remain as `PATCH`) `/test/{connection_id}/` or `/{connection_id}/connection_test` be more RESTful?
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -168,6 +177,84 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
Review Comment:
It seems we need to add dedicated supervisor like #62645 as follow-up.
##########
airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py:
##########
@@ -72,12 +73,46 @@ 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
+ queue: str | None = None
Review Comment:
Maybe we should allow users to specify `executor` here.
For example, specifying `executor` and `queue` would have different meanings at the task level.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -240,10 +242,14 @@ def queue_workload(self, workload: workloads.All, session: Session) -> None:
f"See LocalExecutor or CeleryExecutor for reference implementation."
)
self.queued_callbacks[workload.callback.id] = workload
+ elif isinstance(workload, workloads.TestConnection):
+ if not self.supports_connection_test:
+ raise ValueError(f"Executor {type(self).__name__} does not support connection testing")
Review Comment:
How about raising `NotImplementedError` like `ExecuteCallback` above?
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3179,6 +3189,91 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ @provide_session
+ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None:
+ """Dispatch pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("core", "max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("core", "connection_test_timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTest.id)).where(
+ ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING])
+ )
+ )
+ budget = max_concurrency - (active_count or 0)
+ if budget <= 0:
+ return
+
+ pending_stmt = (
+ select(ConnectionTest)
+ .where(ConnectionTest.state == ConnectionTestState.PENDING)
+ .order_by(ConnectionTest.created_at)
+ .limit(budget)
+ )
+ pending_stmt = with_row_locks(pending_stmt, session, of=ConnectionTest, skip_locked=True)
+ pending_tests = session.scalars(pending_stmt).all()
+
+ if not pending_tests:
+ return
+
+ for ct in pending_tests:
+ executor = self._find_executor_for_connection_test(ct.queue)
+ if executor is None:
+ reason = (
+ f"No executor serves queue '{ct.queue}'"
+ if ct.queue
+ else "No executor supports 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,
+ 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("core", "connection_test_timeout", fallback=60)
+ grace_period = max(30, timeout // 2)
+ cutoff = timezone.utcnow() - timedelta(seconds=timeout + grace_period)
+
+ stale_stmt = select(ConnectionTest).where(
+ ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]),
+ ConnectionTest.updated_at < cutoff,
+ )
+ stale_tests = session.scalars(stale_stmt).all()
Review Comment:
Should we add row-level skip lock like `_dispatch_connection_tests`?
```suggestion
stale_tests = with_row_locks(stale_tests, session, of=ConnectionTest, skip_locked=True)
stale_tests = session.scalars(stale_stmt).all()
```
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3179,6 +3189,91 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ @provide_session
+ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None:
+ """Dispatch pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("core", "max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("core", "connection_test_timeout", fallback=60)
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTest.id)).where(
+ ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING])
+ )
+ )
+ budget = max_concurrency - (active_count or 0)
+ if budget <= 0:
+ return
+
+ pending_stmt = (
+ select(ConnectionTest)
+ .where(ConnectionTest.state == ConnectionTestState.PENDING)
+ .order_by(ConnectionTest.created_at)
+ .limit(budget)
+ )
+ pending_stmt = with_row_locks(pending_stmt, session, of=ConnectionTest, skip_locked=True)
+ pending_tests = session.scalars(pending_stmt).all()
+
+ if not pending_tests:
+ return
+
+ for ct in pending_tests:
+ executor = self._find_executor_for_connection_test(ct.queue)
+ if executor is None:
+ reason = (
+ f"No executor serves queue '{ct.queue}'"
+ if ct.queue
+ else "No executor supports 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,
+ 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("core", "connection_test_timeout", fallback=60)
+ grace_period = max(30, timeout // 2)
+ cutoff = timezone.utcnow() - timedelta(seconds=timeout + grace_period)
+
+ stale_stmt = select(ConnectionTest).where(
+ ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]),
+ ConnectionTest.updated_at < cutoff,
+ )
+ 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)
+ if ct.connection_snapshot:
+ attempt_revert(ct, session=session)
+
+ session.flush()
+
+ def _find_executor_for_connection_test(self, queue: str | None) -> BaseExecutor | None:
+ """Find an executor that supports connection testing, optionally matching a queue."""
+ if queue is not None:
+ for executor in self.executors:
+ if executor.supports_connection_test and executor.team_name == queue:
+ return executor
Review Comment:
I haven’t seen `team_name` mixed with `TI.queue` anywhere so far. I’m concerned this usage might create ambiguity for users.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2876755045
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -168,6 +177,84 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con
raise RuntimeError(error_msg or "Callback execution failed")
+def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None:
Review Comment:
Agreed, for sure will follow up with this change after this
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4038892846
@anishgirianish This PR has been converted to **draft** because it does not yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :x: **Merge conflicts**: This PR has merge conflicts with the `main` branch. Your branch is 31 commits behind `main`. Please rebase your branch (`git fetch origin && git rebase origin/main`), resolve the conflicts, and push again. See [contributing quick start](https://github.com/apache/airflow/blob/main/contributing-docs/03a_contributors_quick_start_beginners.rst).
- :warning: **Unresolved review comments**: This PR has 1 unresolved review thread from maintainers. Please review and resolve all inline review comments before requesting another review. You can resolve a conversation by clicking 'Resolve conversation' on each thread after addressing the feedback. See [pull request guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst).
> **Note:** Your branch is **31 commits behind `main`**. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.
**What to do next:**
- The comment informs you what you need to do.
- Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
- Maintainers will then proceed with a normal review.
Converting a PR to draft is **not** a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3066930789
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2567,6 +2567,33 @@ scheduler:
type: float
example: ~
default: "120.0"
+ connection_test_timeout:
+ description: |
+ Maximum number of seconds an async connection test is allowed to run
+ before it is considered timed out. The scheduler reaper uses this value
+ plus a grace period to mark stale tests as failed.
+ version_added: 3.3.0
+ type: integer
+ example: ~
+ default: "60"
+ max_connection_test_concurrency:
+ description: |
+ Maximum number of connection tests that can be active (QUEUED + RUNNING)
+ at the same time. Excess tests will remain in PENDING state until slots
+ become available.
+ version_added: 3.3.0
+ type: integer
+ example: ~
+ default: "4"
Review Comment:
description was outdated, updated it
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2893498845
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -240,10 +242,18 @@ def queue_workload(self, workload: workloads.All, session: Session) -> None:
f"See LocalExecutor or CeleryExecutor for reference implementation."
)
self.queued_callbacks[workload.callback.id] = workload
+ elif isinstance(workload, workloads.TestConnection):
+ if not self.supports_connection_test:
+ raise NotImplementedError(
+ f"{type(self).__name__} does not support TestConnection workloads. "
+ f"Set supports_connection_test = True and implement connection test handling "
+ f"in _process_workloads(). See LocalExecutor for reference implementation."
+ )
+ self.queued_connection_tests[str(workload.connection_test_id)] = workload
else:
raise ValueError(
f"Un-handled workload type {type(workload).__name__!r} in {type(self).__name__}. "
- f"Workload must be one of: ExecuteTask, ExecuteCallback."
+ f"Workload must be one of: ExecuteTask, ExecuteCallback, TestConnection."
Review Comment:
Great idea, will explore this as a follow-up!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r2901196801
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -573,13 +592,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
@property
def slots_available(self):
- """Number of new workloads (tasks and callbacks) this executor instance can accept."""
- return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
+ """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
+ return (
+ self.parallelism
+ - len(self.running)
+ - len(self.queued_tasks)
+ - len(self.queued_callbacks)
+ - len(self.queued_connection_tests)
+ )
Review Comment:
Hey @ferruzzi I've pulled out the workload queue generalization you suggested into a separate PR #[New Pr point to this branch](https://github.com/anishgirianish/airflow/pull/1) taking out branch from this pr. Would like to request you for your review on this. Thank you.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3169346979
##########
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:
The supervisor imports `airflow.models.connection.Connection` solely to call `test_connection()` on it. That method dispatches to the provider hook via `ProvidersManager` (core), which adds a new ORM dependency to task-sdk and is the wrong direction — task-sdk should not be pulling in core ORM models.
The right fix, which belongs in this PR since connection testing is its stated purpose:
**Add `test_connection()` to `airflow.sdk.definitions.connection.Connection`**, dispatching via `ProvidersManagerTaskRuntime` (already in task-sdk for exactly this kind of hook dispatch):
```python
def test_connection(self) -> tuple[bool, str]:
from airflow.sdk.execution_time.context import _preset_connections
from airflow.sdk.providers_manager_runtime import ProvidersManagerTaskRuntime
from airflow.sdk._shared.module_loading import import_string
hook_info = ProvidersManagerTaskRuntime().hooks.get(self.conn_type)
if not hook_info:
return False, f"Unknown conn_type: {self.conn_type!r}"
outer = _preset_connections.get() or {}
token = _preset_connections.set({**outer, self.conn_id: self})
try:
hook = import_string(hook_info.hook_class_name)(conn_id=self.conn_id)
if not hasattr(hook, "test_connection"):
return False, f"Hook {type(hook).__name__} doesn't implement test_connection"
return hook.test_connection()
finally:
_preset_connections.reset(token)
```
`_preset_connections` is the right injection point (checked before all secrets backends, including user-configured custom ones), but it should be encapsulated inside `test_connection()` rather than wired manually in the supervisor.
That reduces the supervisor to:
```python
conn = SDKConnection(conn_id=r.conn_id, conn_type=r.conn_type, ...)
with TimeoutPosix(seconds=timeout, error_message=f"Connection test timed out after {timeout}s"):
success, message = conn.test_connection()
```
Which also fixes the second issue here: `signal.signal(signal.SIGALRM, ...)` + `signal.alarm()` should use `TimeoutPosix` from `airflow.sdk.execution_time.timeout` — it already exists, handles non-POSIX gracefully, and uses `signal.setitimer` (sub-second capable) rather than integer-only `signal.alarm`.
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -144,7 +145,16 @@ def _convert_variable_result_to_variable(var_result: VariableResult, deserialize
return Variable(**var_result.model_dump(exclude={"type"}))
+_preset_connections: ContextVar[dict[str, Connection]] = ContextVar("_preset_connections", default={})
Review Comment:
`default={}` is a single shared mutable object — every context that hasn't called `.set()` gets back the exact same dict from `.get()`. The current code never mutates it, so it doesn't blow up today, but it's a trap for anyone who writes `_preset_connections.get()[conn_id] = conn` instead of using `.set()`.
Change to `default=None` and guard at call sites:
```python
_preset_connections: ContextVar[dict[str, Connection] | None] = ContextVar("_preset_connections", default=None)
# in _get_connection:
preset = _preset_connections.get()
if preset and conn_id in preset:
...
```
Also: `.set({preset.conn_id: preset})` replaces the entire dict, shadowing any outer preset for the duration of the call. If `test_connection()` is ever called from within a task that has its own preset set (e.g. test tooling), the outer entries disappear until `.reset()`. Merge instead of replace:
```python
outer = _preset_connections.get() or {}
token = _preset_connections.set({**outer, self.conn_id: self})
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177844042
##########
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:
Done, added BaseExecutor.fail_connection_test(key); reaper calls it for every supporting executor after marking FAILED. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177843525
##########
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:
Done, dropped the aggregate parallelism_budget; per-executor slots_available in trigger_connection_tests is the gate.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177844478
##########
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:
Done, num_queued_workloads now includes len(self.queued_connection_tests). thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3177840275
##########
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:
Done, load row first, manual is_authorized_connection(GET, conn_id, team_name), return 404 on unauthorized to avoid the token oracle. thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3179136463
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -216,6 +218,7 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None)
self.team_name: str | None = team_name
self.queued_tasks: dict[TaskInstanceKey, workloads.ExecuteTask] = {}
self.queued_callbacks: dict[str, workloads.ExecuteCallback] = {}
+ self.queued_connection_tests: dict[ConnectionTestKey, workloads.TestConnection] = {}
Review Comment:
@ashb this #63491 exactly has this refactor. Would really appreciate a look when you get a chance, thank you!
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -216,6 +218,7 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None)
self.team_name: str | None = team_name
self.queued_tasks: dict[TaskInstanceKey, workloads.ExecuteTask] = {}
self.queued_callbacks: dict[str, workloads.ExecuteCallback] = {}
+ self.queued_connection_tests: dict[ConnectionTestKey, workloads.TestConnection] = {}
Review Comment:
@ashb this #63491 pr exactly has this refactor. Would really appreciate a look when you get a chance, thank you!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4505144899
Tests are failing in main https://github.com/apache/airflow/pull/67268 opened a sperate pr for fixing those.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4585025031
> I found two blockers in the current head:
>
> 1. The worker success path drops team ownership when `commit_on_success=True` creates a brand-new connection. The public API resolves and authorizes the effective team, then stores it on `ConnectionTestRequest`, but `commit_to_connection_table()` creates the new `Connection(...)` without `team_name=self.team_name`. In multi-team mode that can turn a team-scoped connection credential into a persisted global connection. Please copy the team on creation and add a regression test for a new team-owned connection with `commit_on_success=True`.
> 2. The UI still appears to exercise the old synchronous API-server test path, so the main UI workflow does not get the worker-isolation behavior added by this PR. `useTestConnection.ts` still calls the generated client for `/api/v2/connections/test`, and that route still executes `conn.test_connection()` inside the API server process. If the intent is for UI connection tests to be isolated from the API server, the UI should enqueue the test and poll by token, or the old API-server route should be separately gated/deprecated so enabling connection testing does not keep exposing the synchronous API-server execution path.
>
> Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
Hi @Vamsi-klu thank you so much for your review.
1. Addressed thanks
2. UI changes are not part of this pr and will come along in a follow-up as stated in Pr description.
I think this clears the blockers. Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "Vamsi-klu (via GitHub)" <gi...@apache.org>.
Vamsi-klu commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4584957633
I found two blockers in the current head:
1. The worker success path drops team ownership when `commit_on_success=True` creates a brand-new connection. The public API resolves and authorizes the effective team, then stores it on `ConnectionTestRequest`, but `commit_to_connection_table()` creates the new `Connection(...)` without `team_name=self.team_name`. In multi-team mode that can turn a team-scoped connection credential into a persisted global connection. Please copy the team on creation and add a regression test for a new team-owned connection with `commit_on_success=True`.
2. The UI still appears to exercise the old synchronous API-server test path, so the main UI workflow does not get the worker-isolation behavior added by this PR. `useTestConnection.ts` still calls the generated client for `/api/v2/connections/test`, and that route still executes `conn.test_connection()` inside the API server process. If the intent is for UI connection tests to be isolated from the API server, the UI should enqueue the test and poll by token, or the old API-server route should be separately gated/deprecated so enabling connection testing does not keep exposing the synchronous API-server execution path.
---
Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #62343:
URL: https://github.com/apache/airflow/pull/62343#issuecomment-4426851054
@ashb @kaxil thank you so much for the review. I have addressed the feedback on the latest push. Would like to request you for your re-review on this whenever you get a chance.
Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add async connection testing via workers for security isolation [airflow]
Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3316035444
##########
airflow-core/src/airflow/models/connection_test.py:
##########
@@ -0,0 +1,216 @@
+# 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(String(2000), nullable=True)
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(
+ UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False
+ )
+ executor: Mapped[str | None] = mapped_column(String(256), nullable=True)
+ queue: Mapped[str | None] = mapped_column(String(256), nullable=True)
+
+ conn_type: Mapped[str] = mapped_column(String(500), nullable=False)
+ host: Mapped[str | None] = mapped_column(String(500), nullable=True)
+ login: Mapped[str | None] = mapped_column(Text, nullable=True)
+ schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True)
+ port: Mapped[int | None] = mapped_column(Integer, nullable=True)
+ commit_on_success: Mapped[bool] = mapped_column(
+ Boolean, nullable=False, default=False, server_default="0"
+ )
+ is_encrypted: Mapped[bool] = mapped_column(
Review Comment:
Added rotate_fernet_key() thanks
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org