Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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",
+),
+)
+
+
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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:
+ret
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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(ba
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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`
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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,
+}
+
[email protected](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
+
[email protected](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"
+
[email protected](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"
+
[email protected](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
+
[email protected](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"
+
[email protected](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_execu
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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.00
+
+"""
+
+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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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"})
+respons
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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 cal
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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())],
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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
+#
+# Unles
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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",
+hos
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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.e
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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
+
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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
+
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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,
+
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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.00 + +""" + +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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add async connection testing via workers for security isolation [airflow]
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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
