From eb8e93e8aeb5de25a2c29ef034602dfe449e6694 Mon Sep 17 00:00:00 2001 From: Anish Date: Sun, 22 Feb 2026 23:24:03 -0600 Subject: [PATCH 01/38] Add async connection testing via ExecutorCallback dispatch to workers --- .../core_api/datamodels/connections.py | 28 ++- .../openapi/v2-rest-api-generated.yaml | 163 +++++++++++++++++- .../core_api/routes/public/connections.py | 121 +++++++++++++ .../datamodels/connection_test.py | 29 ++++ .../execution_api/routes/__init__.py | 4 + .../execution_api/routes/connection_tests.py | 65 +++++++ .../0107_3_2_0_add_connection_test_table.py | 73 ++++++++ airflow-core/src/airflow/models/__init__.py | 1 + .../src/airflow/models/connection_test.py | 125 ++++++++++++++ .../airflow/ui/openapi-gen/queries/common.ts | 7 + .../ui/openapi-gen/queries/ensureQueryData.ts | 14 ++ .../ui/openapi-gen/queries/prefetch.ts | 14 ++ .../airflow/ui/openapi-gen/queries/queries.ts | 74 ++++---- .../ui/openapi-gen/queries/suspense.ts | 14 ++ .../ui/openapi-gen/requests/schemas.gen.ts | 85 ++++++++- .../ui/openapi-gen/requests/services.gen.ts | 120 +++++++------ .../ui/openapi-gen/requests/types.gen.ts | 96 ++++++++++- airflow-core/src/airflow/utils/db_cleanup.py | 1 + .../routes/public/test_connections.py | 107 ++++++++++++ .../versions/head/test_connection_tests.py | 80 +++++++++ .../tests/unit/models/test_connection_test.py | 108 ++++++++++++ .../airflowctl/api/datamodels/generated.py | 36 +++- .../airflow/sdk/api/datamodels/_generated.py | 25 +++ 23 files changed, 1285 insertions(+), 105 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py create mode 100644 airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py create mode 100644 airflow-core/src/airflow/models/connection_test.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py create mode 100644 airflow-core/tests/unit/models/test_connection_test.py diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index f7cb944ebbf6b..936b6832fec85 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -19,6 +19,7 @@ import json from collections.abc import Iterable, Mapping +from datetime import datetime from typing import Annotated, Any from pydantic import Field, field_validator @@ -72,12 +73,37 @@ 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 — just the connection_id.""" + + connection_id: str + + +class ConnectionTestQueuedResponse(BaseModel): + """Response returned when an async connection test is queued.""" + + token: str + connection_id: str + state: str + + +class ConnectionTestStatusResponse(BaseModel): + """Response returned when polling for async connection test status.""" + + token: str + connection_id: str + state: str + result_status: bool | None = None + result_message: str | None = None + created_at: datetime + + class ConnectionHookFieldBehavior(BaseModel): """A class to store the behavior of each standard field of a Hook.""" diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 17c4f87a6e44f..57bdb6d00fd65 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -1681,6 +1681,105 @@ 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. + + + The connection must already be saved. Returns a token that can be used + + to poll for the test result via GET /connections/test-async/{token}.' + operationId: test_connection_async + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionTestRequestBody' + required: true + responses: + '202': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionTestQueuedResponse' + '401': + description: Unauthorized + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '403': + description: Forbidden + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '404': + description: Not Found + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' + security: + - OAuth2PasswordBearer: [] + - HTTPBearer: [] + /api/v2/connections/test-async/{token}: + get: + tags: + - Connection + summary: Get Connection Test Status + description: "Poll for the status of an async connection test.\n\nKnowledge\ + \ of the token serves as authorization \u2014 only the client\nthat initiated\ + \ the test knows the crypto-random token." + operationId: get_connection_test_status + parameters: + - name: token + in: path + required: true + schema: + type: string + title: Token + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionTestStatusResponse' + '401': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Unauthorized + '403': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Forbidden + '404': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Not Found + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' /api/v2/connections/defaults: post: tags: @@ -10096,6 +10195,35 @@ components: - team_name title: ConnectionResponse description: Connection serializer for responses. + ConnectionTestQueuedResponse: + properties: + token: + type: string + title: Token + connection_id: + type: string + title: Connection Id + state: + type: string + title: State + type: object + required: + - token + - connection_id + - state + title: ConnectionTestQueuedResponse + description: Response returned when an async connection test is queued. + ConnectionTestRequestBody: + properties: + connection_id: + type: string + title: Connection Id + additionalProperties: false + type: object + required: + - connection_id + title: ConnectionTestRequestBody + description: "Request body for async connection test \u2014 just the connection_id." ConnectionTestResponse: properties: status: @@ -10109,7 +10237,40 @@ components: - status - message title: ConnectionTestResponse - description: Connection Test serializer for responses. + description: Connection Test serializer for synchronous test responses. + ConnectionTestStatusResponse: + properties: + token: + type: string + title: Token + connection_id: + type: string + title: Connection Id + state: + type: string + title: State + result_status: + anyOf: + - type: boolean + - type: 'null' + title: Result Status + result_message: + anyOf: + - type: string + - type: 'null' + title: Result Message + created_at: + type: string + format: date-time + title: Created At + type: object + required: + - token + - connection_id + - state + - created_at + title: ConnectionTestStatusResponse + description: Response returned when polling for async connection test status. CreateAssetEventsBody: properties: asset_id: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 05ec6b642941d..47e6e4b0a207a 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -41,7 +41,10 @@ ConnectionBodyPartial, ConnectionCollectionResponse, ConnectionResponse, + ConnectionTestQueuedResponse, + ConnectionTestRequestBody, ConnectionTestResponse, + ConnectionTestStatusResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc from airflow.api_fastapi.core_api.security import ( @@ -57,6 +60,12 @@ from airflow.configuration import conf from airflow.exceptions import AirflowNotFoundException from airflow.models import Connection +from airflow.models.callback import CallbackFetchMethod, ExecutorCallback +from airflow.models.connection_test import ( + RUN_CONNECTION_TEST_PATH, + ConnectionTest, + ConnectionTestState, +) from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.utils.db import create_default_connections as db_create_default_connections from airflow.utils.strings import get_random_string @@ -64,6 +73,25 @@ connections_router = AirflowRouter(tags=["Connection"], prefix="/connections") +class _ImportPathCallbackDef: + """ + Minimal implementation of ImportPathExecutorCallbackDefProtocol. + + ExecutorCallback.__init__ expects an object satisfying this protocol, but no concrete + implementation exists in airflow-core — the only one (SyncCallback) lives in the task-sdk. + Once #61153 lands and ExecuteCallback.make() is decoupled from DagRun, this adapter can + be replaced with the proper factory method. + """ + + def __init__(self, path: str, kwargs: dict, executor: str | None = None): + self.path = path + self.kwargs = kwargs + self.executor = executor + + def serialize(self) -> dict: + return {"path": self.path, "kwargs": self.kwargs, "executor": self.executor} + + @connections_router.delete( "/{connection_id}", status_code=status.HTTP_204_NO_CONTENT, @@ -257,6 +285,99 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse: os.environ.pop(conn_env_var, None) +@connections_router.post( + "/test-async", + status_code=status.HTTP_202_ACCEPTED, + responses=create_openapi_http_exception_doc([status.HTTP_403_FORBIDDEN, status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_connection(method="POST")), Depends(action_logging())], +) +def test_connection_async( + test_body: ConnectionTestRequestBody, + session: SessionDep, +) -> ConnectionTestQueuedResponse: + """ + Queue an async connection test to be executed on a worker. + + The connection must already be saved. Returns a token that can be used + to poll for the test result via GET /connections/test-async/{token}. + """ + if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled": + raise HTTPException( + status.HTTP_403_FORBIDDEN, + "Testing connections is disabled in Airflow configuration. " + "Contact your deployment admin to enable it.", + ) + + try: + Connection.get_connection_from_secrets(test_body.connection_id) + except AirflowNotFoundException: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"The Connection with connection_id: `{test_body.connection_id}` was not found. " + "Connection must be saved before testing.", + ) + + connection_test = ConnectionTest(connection_id=test_body.connection_id) + session.add(connection_test) + session.flush() + + # ExecutorCallback requires an object satisfying ImportPathExecutorCallbackDefProtocol, + # but the only concrete impl (SyncCallback) lives in task-sdk which core API cannot + # import. This adapter will be replaced by ExecuteCallback.make() once #61153 merges. + callback_def = _ImportPathCallbackDef( + path=RUN_CONNECTION_TEST_PATH, + kwargs={ + "connection_id": test_body.connection_id, + "connection_test_id": str(connection_test.id), + }, + ) + callback = ExecutorCallback(callback_def, fetch_method=CallbackFetchMethod.IMPORT_PATH) + session.add(callback) + session.flush() + + connection_test.callback_id = callback.id + connection_test.state = ConnectionTestState.QUEUED + callback.queue() + + return ConnectionTestQueuedResponse( + token=connection_test.token, + connection_id=connection_test.connection_id, + state=connection_test.state, + ) + + +@connections_router.get( + "/test-async/{token}", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), +) +def get_connection_test_status( + 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. + """ + connection_test = session.scalar(select(ConnectionTest).filter_by(token=token)) + + if connection_test is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"No connection test found for token: `{token}`", + ) + + return ConnectionTestStatusResponse( + token=connection_test.token, + connection_id=connection_test.connection_id, + state=connection_test.state, + result_status=connection_test.result_status, + result_message=connection_test.result_message, + created_at=connection_test.created_at, + ) + + @connections_router.post( "/defaults", status_code=status.HTTP_204_NO_CONTENT, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py new file mode 100644 index 0000000000000..6042906918ab8 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py @@ -0,0 +1,29 @@ +# 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 airflow.api_fastapi.core_api.base import StrictBaseModel +from airflow.models.connection_test import ConnectionTestState + + +class ConnectionTestResultBody(StrictBaseModel): + """Payload sent by workers to report connection test results.""" + + state: ConnectionTestState + result_status: bool | None = None + result_message: str | None = None diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py index a076592d6471a..0d4291bbcc2de 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py @@ -22,6 +22,7 @@ from airflow.api_fastapi.execution_api.routes import ( asset_events, assets, + connection_tests, connections, dag_runs, dags, @@ -42,6 +43,9 @@ authenticated_router.include_router(assets.router, prefix="/assets", tags=["Assets"]) authenticated_router.include_router(asset_events.router, prefix="/asset-events", tags=["Asset Events"]) +authenticated_router.include_router( + connection_tests.router, prefix="/connection-tests", tags=["Connection Tests"] +) authenticated_router.include_router(connections.router, prefix="/connections", tags=["Connections"]) authenticated_router.include_router(dag_runs.router, prefix="/dag-runs", tags=["Dag Runs"]) authenticated_router.include_router(dags.router, prefix="/dags", tags=["Dags"]) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py new file mode 100644 index 0000000000000..d497b1c6e6db3 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -0,0 +1,65 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from uuid import UUID + +from fastapi import APIRouter, HTTPException, status + +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.execution_api.datamodels.connection_test import ConnectionTestResultBody +from airflow.models.connection_test import TERMINAL_STATES, ConnectionTest + +router = APIRouter() + + +@router.patch( + "/{connection_test_id}", + status_code=status.HTTP_204_NO_CONTENT, + responses={ + status.HTTP_404_NOT_FOUND: {"description": "Connection test not found"}, + status.HTTP_409_CONFLICT: {"description": "Connection test already in a terminal state"}, + }, +) +def patch_connection_test( + connection_test_id: UUID, + body: ConnectionTestResultBody, + session: SessionDep, +) -> None: + """Update the result of a connection test (called by workers).""" + ct = session.get(ConnectionTest, connection_test_id, with_for_update=True) + if ct is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={ + "reason": "not_found", + "message": f"Connection test {connection_test_id} not found", + }, + ) + + if ct.state in TERMINAL_STATES: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "reason": "conflict", + "message": f"Connection test {connection_test_id} is already in terminal state: {ct.state}", + }, + ) + + ct.state = body.state + ct.result_status = body.result_status + ct.result_message = body.result_message diff --git a/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py new file mode 100644 index 0000000000000..a3ccef07872e0 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py @@ -0,0 +1,73 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Add connection_test table. + +Revision ID: a7e6d4c3b2f1 +Revises: 509b94a1042d +Create Date: 2026-02-22 00:00:00.000000 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "a7e6d4c3b2f1" +down_revision = "509b94a1042d" +branch_labels = None +depends_on = None +airflow_version = "3.2.0" + + +def upgrade(): + """Create connection_test table.""" + op.create_table( + "connection_test", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("token", sa.String(64), nullable=False), + sa.Column("connection_id", sa.String(250), nullable=False), + sa.Column("state", sa.String(10), nullable=False), + sa.Column("result_status", sa.Boolean(), nullable=True), + sa.Column("result_message", sa.Text(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.Column("updated_at", sa.DateTime(), nullable=False), + sa.Column("callback_id", sa.Uuid(), nullable=True), + sa.PrimaryKeyConstraint("id", name=op.f("connection_test_pkey")), + sa.ForeignKeyConstraint( + ["callback_id"], + ["callback.id"], + name=op.f("connection_test_callback_id_fkey"), + ondelete="SET NULL", + ), + sa.UniqueConstraint("token", name=op.f("connection_test_token_uq")), + ) + op.create_index( + op.f("idx_connection_test_state_created_at"), + "connection_test", + ["state", "created_at"], + ) + + +def downgrade(): + """Drop connection_test table.""" + op.drop_index(op.f("idx_connection_test_state_created_at"), table_name="connection_test") + op.drop_table("connection_test") diff --git a/airflow-core/src/airflow/models/__init__.py b/airflow-core/src/airflow/models/__init__.py index 8e12325f568eb..49a1c1f41129b 100644 --- a/airflow-core/src/airflow/models/__init__.py +++ b/airflow-core/src/airflow/models/__init__.py @@ -62,6 +62,7 @@ def import_all_models(): import airflow.models.asset import airflow.models.backfill + import airflow.models.connection_test import airflow.models.dag_favorite import airflow.models.dag_version import airflow.models.dagbag diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py new file mode 100644 index 0000000000000..898debd7e62c7 --- /dev/null +++ b/airflow-core/src/airflow/models/connection_test.py @@ -0,0 +1,125 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import secrets +from datetime import datetime +from enum import Enum +from typing import TYPE_CHECKING +from uuid import UUID + +import structlog +import uuid6 +from sqlalchemy import Boolean, ForeignKey, Index, String, Text, Uuid +from sqlalchemy.orm import Mapped, mapped_column, relationship + +from airflow._shared.timezones import timezone +from airflow.models.base import Base +from airflow.utils.sqlalchemy import UtcDateTime + +if TYPE_CHECKING: + from airflow.models.callback import Callback + +log = structlog.get_logger(__name__) + + +class ConnectionTestState(str, Enum): + """All possible states of a connection test.""" + + PENDING = "pending" + QUEUED = "queued" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + + def __str__(self) -> str: + return self.value + + +TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED)) + +# Path used by ExecutorCallback to locate the worker function. +RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test" + + +class ConnectionTest(Base): + """Tracks an async connection test dispatched to a worker via ExecutorCallback.""" + + __tablename__ = "connection_test" + + id: Mapped[UUID] = mapped_column(Uuid(), primary_key=True, default=uuid6.uuid7) + token: Mapped[str] = mapped_column(String(64), nullable=False, unique=True) + connection_id: Mapped[str] = mapped_column(String(250), nullable=False) + state: Mapped[str] = mapped_column(String(10), nullable=False, default=ConnectionTestState.PENDING) + result_status: Mapped[bool | None] = mapped_column(Boolean, nullable=True) + result_message: Mapped[str | None] = mapped_column(Text, nullable=True) + created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False) + updated_at: Mapped[datetime] = mapped_column( + UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False + ) + + callback_id: Mapped[UUID | None] = mapped_column( + Uuid(), ForeignKey("callback.id", ondelete="SET NULL"), nullable=True + ) + callback: Mapped[Callback | None] = relationship("Callback", uselist=False) + + __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),) + + def __init__(self, *, connection_id: str, **kwargs): + super().__init__(**kwargs) + self.connection_id = connection_id + self.token = secrets.token_urlsafe(32) + self.state = ConnectionTestState.PENDING + + def __repr__(self) -> str: + return ( + f"" + ) + + +def run_connection_test(*, connection_id: str, connection_test_id: str) -> None: + """ + Worker-side function to execute a connection test. + + This is the function referenced by the ExecutorCallback's import path. + It fetches the connection, runs test_connection(), and reports results + back by updating the ConnectionTest row directly. + """ + from airflow.models.connection import Connection + from airflow.utils.session import create_session + + connection_test_uuid = UUID(connection_test_id) + + with create_session() as session: + ct = session.get(ConnectionTest, connection_test_uuid) + if ct: + ct.state = ConnectionTestState.RUNNING + + try: + conn = Connection.get_connection_from_secrets(connection_id) + test_status, test_message = conn.test_connection() + except Exception as e: + test_status = False + test_message = str(e) + log.exception("Connection test failed", connection_id=connection_id) + + with create_session() as session: + ct = session.get(ConnectionTest, connection_test_uuid) + if ct: + ct.result_status = test_status + ct.result_message = test_message + ct.state = ConnectionTestState.SUCCESS if test_status else ConnectionTestState.FAILED diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts index 612a3d56747e3..9034236d66574 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -122,6 +122,12 @@ export const UseConnectionServiceGetConnectionsKeyFn = ({ connectionIdPattern, l offset?: number; orderBy?: string[]; } = {}, queryKey?: Array) => [useConnectionServiceGetConnectionsKey, ...(queryKey ?? [{ connectionIdPattern, limit, offset, orderBy }])]; +export type ConnectionServiceGetConnectionTestStatusDefaultResponse = Awaited>; +export type ConnectionServiceGetConnectionTestStatusQueryResult = UseQueryResult; +export const useConnectionServiceGetConnectionTestStatusKey = "ConnectionServiceGetConnectionTestStatus"; +export const UseConnectionServiceGetConnectionTestStatusKeyFn = ({ token }: { + token: string; +}, queryKey?: Array) => [useConnectionServiceGetConnectionTestStatusKey, ...(queryKey ?? [{ token }])]; export type ConnectionServiceHookMetaDataDefaultResponse = Awaited>; export type ConnectionServiceHookMetaDataQueryResult = UseQueryResult; export const useConnectionServiceHookMetaDataKey = "ConnectionServiceHookMetaData"; @@ -923,6 +929,7 @@ export type BackfillServiceCreateBackfillMutationResult = Awaited>; export type ConnectionServicePostConnectionMutationResult = Awaited>; export type ConnectionServiceTestConnectionMutationResult = Awaited>; +export type ConnectionServiceTestConnectionAsyncMutationResult = Awaited>; export type ConnectionServiceCreateDefaultConnectionsMutationResult = Awaited>; export type DagRunServiceClearDagRunMutationResult = Awaited>; export type DagRunServiceTriggerDagRunMutationResult = Awaited>; diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts index 8bc9e9df8e6df..2dae237e1ae82 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -225,6 +225,20 @@ export const ensureUseConnectionServiceGetConnectionsData = (queryClient: QueryC orderBy?: string[]; } = {}) => queryClient.ensureQueryData({ queryKey: Common.UseConnectionServiceGetConnectionsKeyFn({ connectionIdPattern, limit, offset, orderBy }), queryFn: () => ConnectionService.getConnections({ connectionIdPattern, limit, offset, orderBy }) }); /** +* Get Connection Test Status +* 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. +* @param data The data for the request. +* @param data.token +* @returns ConnectionTestStatusResponse Successful Response +* @throws ApiError +*/ +export const ensureUseConnectionServiceGetConnectionTestStatusData = (queryClient: QueryClient, { token }: { + token: string; +}) => queryClient.ensureQueryData({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ token }), queryFn: () => ConnectionService.getConnectionTestStatus({ token }) }); +/** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. * @returns ConnectionHookMetaData Successful Response diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts index f4cb6f482bdf6..db0fcf85b25b4 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -225,6 +225,20 @@ export const prefetchUseConnectionServiceGetConnections = (queryClient: QueryCli orderBy?: string[]; } = {}) => queryClient.prefetchQuery({ queryKey: Common.UseConnectionServiceGetConnectionsKeyFn({ connectionIdPattern, limit, offset, orderBy }), queryFn: () => ConnectionService.getConnections({ connectionIdPattern, limit, offset, orderBy }) }); /** +* Get Connection Test Status +* 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. +* @param data The data for the request. +* @param data.token +* @returns ConnectionTestStatusResponse Successful Response +* @throws ApiError +*/ +export const prefetchUseConnectionServiceGetConnectionTestStatus = (queryClient: QueryClient, { token }: { + token: string; +}) => queryClient.prefetchQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ token }), queryFn: () => ConnectionService.getConnectionTestStatus({ token }) }); +/** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. * @returns ConnectionHookMetaData Successful Response diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index 8e9ef5aa29d0c..2b32fe404bc47 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -2,7 +2,7 @@ import { UseMutationOptions, UseQueryOptions, useMutation, useQuery } from "@tanstack/react-query"; import { AssetService, AuthLinksService, BackfillService, CalendarService, ConfigService, ConnectionService, DagParsingService, DagRunService, DagService, DagSourceService, DagStatsService, DagVersionService, DagWarningService, DashboardService, DeadlinesService, DependenciesService, EventLogService, ExperimentalService, ExtraLinksService, GanttService, GridService, ImportErrorService, JobService, LoginService, MonitorService, PartitionedDagRunService, PluginService, PoolService, ProviderService, StructureService, TaskInstanceService, TaskService, TeamsService, VariableService, VersionService, XcomService } from "../requests/services.gen"; -import { BackfillPostBody, BulkBody_BulkTaskInstanceBody_, BulkBody_ConnectionBody_, BulkBody_PoolBody_, BulkBody_VariableBody_, ClearTaskInstancesBody, ConnectionBody, CreateAssetEventsBody, DAGPatchBody, DAGRunClearBody, DAGRunPatchBody, DAGRunsBatchBody, DagRunState, DagWarningType, GenerateTokenBody, PatchTaskInstanceBody, PoolBody, PoolPatchBody, TaskInstancesBatchBody, TriggerDAGRunPostBody, UpdateHITLDetailPayload, VariableBody, XComCreateBody, XComUpdateBody } from "../requests/types.gen"; +import { BackfillPostBody, BulkBody_BulkTaskInstanceBody_, BulkBody_ConnectionBody_, BulkBody_PoolBody_, BulkBody_VariableBody_, ClearTaskInstancesBody, ConnectionBody, ConnectionTestRequestBody, CreateAssetEventsBody, DAGPatchBody, DAGRunClearBody, DAGRunPatchBody, DAGRunsBatchBody, DagRunState, DagWarningType, GenerateTokenBody, PatchTaskInstanceBody, PoolBody, PoolPatchBody, TaskInstancesBatchBody, TriggerDAGRunPostBody, UpdateHITLDetailPayload, VariableBody, XComCreateBody, XComUpdateBody } from "../requests/types.gen"; import * as Common from "./common"; /** * Get Assets @@ -225,6 +225,20 @@ export const useConnectionServiceGetConnections = , "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseConnectionServiceGetConnectionsKeyFn({ connectionIdPattern, limit, offset, orderBy }, queryKey), queryFn: () => ConnectionService.getConnections({ connectionIdPattern, limit, offset, orderBy }) as TData, ...options }); /** +* Get Connection Test Status +* 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. +* @param data The data for the request. +* @param data.token +* @returns ConnectionTestStatusResponse Successful Response +* @throws ApiError +*/ +export const useConnectionServiceGetConnectionTestStatus = = unknown[]>({ token }: { + token: string; +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ token }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ token }) as TData, ...options }); +/** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. * @returns ConnectionHookMetaData Successful Response @@ -1501,32 +1515,6 @@ export const useAuthLinksServiceGetAuthMenus = = unknown[]>(queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseAuthLinksServiceGetCurrentUserInfoKeyFn(queryKey), queryFn: () => AuthLinksService.getCurrentUserInfo() as TData, ...options }); /** -* Get Partitioned Dag Runs -* Return PartitionedDagRuns. Filter by dag_id and/or has_created_dag_run_id. -* @param data The data for the request. -* @param data.dagId -* @param data.hasCreatedDagRunId -* @returns PartitionedDagRunCollectionResponse Successful Response -* @throws ApiError -*/ -export const usePartitionedDagRunServiceGetPartitionedDagRuns = = unknown[]>({ dagId, hasCreatedDagRunId }: { - dagId?: string; - hasCreatedDagRunId?: boolean; -} = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId, hasCreatedDagRunId }, queryKey), queryFn: () => PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId }) as TData, ...options }); -/** -* Get Pending Partitioned Dag Run -* Return full details for pending PartitionedDagRun. -* @param data The data for the request. -* @param data.dagId -* @param data.partitionKey -* @returns PartitionedDagRunDetailResponse Successful Response -* @throws ApiError -*/ -export const usePartitionedDagRunServiceGetPendingPartitionedDagRun = = unknown[]>({ dagId, partitionKey }: { - dagId: string; - partitionKey: string; -}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId, partitionKey }, queryKey), queryFn: () => PartitionedDagRunService.getPendingPartitionedDagRun({ dagId, partitionKey }) as TData, ...options }); -/** * Get Dependencies * Dependencies graph. * @param data The data for the request. @@ -1564,20 +1552,14 @@ export const useDashboardServiceDagStats = = unknown[]>({ dagId, dagRunId, limit, offset, orderBy }: { +export const useDeadlinesServiceGetDagRunDeadlines = = unknown[]>({ dagId, runId }: { dagId: string; - dagRunId: string; - limit?: number; - offset?: number; - orderBy?: string[]; -}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, dagRunId, limit, offset, orderBy }, queryKey), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, dagRunId, limit, offset, orderBy }) as TData, ...options }); + runId: string; +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, runId }, queryKey), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, runId }) as TData, ...options }); /** * Structure Data * Get Structure Data. @@ -1828,6 +1810,22 @@ export const useConnectionServiceTestConnection = ({ mutationFn: ({ requestBody }) => ConnectionService.testConnection({ requestBody }) as unknown as Promise, ...options }); /** +* Test Connection Async +* Queue an async connection test to be executed on a worker. +* +* The connection must already be saved. Returns a token that can be used +* to poll for the test result via GET /connections/test-async/{token}. +* @param data The data for the request. +* @param data.requestBody +* @returns ConnectionTestQueuedResponse Successful Response +* @throws ApiError +*/ +export const useConnectionServiceTestConnectionAsync = (options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ requestBody }) => ConnectionService.testConnectionAsync({ requestBody }) as unknown as Promise, ...options }); +/** * Create Default Connections * Create default connections. * @returns void Successful Response diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts index c4a41691b1a2e..762067e621653 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -225,6 +225,20 @@ export const useConnectionServiceGetConnectionsSuspense = , "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseConnectionServiceGetConnectionsKeyFn({ connectionIdPattern, limit, offset, orderBy }, queryKey), queryFn: () => ConnectionService.getConnections({ connectionIdPattern, limit, offset, orderBy }) as TData, ...options }); /** +* Get Connection Test Status +* 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. +* @param data The data for the request. +* @param data.token +* @returns ConnectionTestStatusResponse Successful Response +* @throws ApiError +*/ +export const useConnectionServiceGetConnectionTestStatusSuspense = = unknown[]>({ token }: { + token: string; +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ token }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ token }) as TData, ...options }); +/** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. * @returns ConnectionHookMetaData Successful Response diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index f47aa107aea24..3da3229f7e82f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1698,6 +1698,41 @@ export const $ConnectionResponse = { description: 'Connection serializer for responses.' } as const; +export const $ConnectionTestQueuedResponse = { + properties: { + token: { + type: 'string', + title: 'Token' + }, + connection_id: { + type: 'string', + title: 'Connection Id' + }, + state: { + type: 'string', + title: 'State' + } + }, + type: 'object', + required: ['token', 'connection_id', 'state'], + title: 'ConnectionTestQueuedResponse', + description: 'Response returned when an async connection test is queued.' +} as const; + +export const $ConnectionTestRequestBody = { + properties: { + connection_id: { + type: 'string', + title: 'Connection Id' + } + }, + additionalProperties: false, + type: 'object', + required: ['connection_id'], + title: 'ConnectionTestRequestBody', + description: 'Request body for async connection test — just the connection_id.' +} as const; + export const $ConnectionTestResponse = { properties: { status: { @@ -1712,7 +1747,55 @@ export const $ConnectionTestResponse = { type: 'object', required: ['status', 'message'], title: 'ConnectionTestResponse', - description: 'Connection Test serializer for responses.' + description: 'Connection Test serializer for synchronous test responses.' +} as const; + +export const $ConnectionTestStatusResponse = { + properties: { + token: { + type: 'string', + title: 'Token' + }, + connection_id: { + type: 'string', + title: 'Connection Id' + }, + state: { + type: 'string', + title: 'State' + }, + result_status: { + anyOf: [ + { + type: 'boolean' + }, + { + type: 'null' + } + ], + title: 'Result Status' + }, + result_message: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Result Message' + }, + created_at: { + type: 'string', + format: 'date-time', + title: 'Created At' + } + }, + type: 'object', + required: ['token', 'connection_id', 'state', 'created_at'], + title: 'ConnectionTestStatusResponse', + description: 'Response returned when polling for async connection test status.' } as const; export const $CreateAssetEventsBody = { diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 6e701bf68dc6e..077a91fd8edb8 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -3,7 +3,7 @@ import type { CancelablePromise } from './core/CancelablePromise'; import { OpenAPI } from './core/OpenAPI'; import { request as __request } from './core/request'; -import type { GetAssetsData, GetAssetsResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetAliasData, GetAssetAliasResponse, GetAssetEventsData, GetAssetEventsResponse, CreateAssetEventData, CreateAssetEventResponse, MaterializeAssetData, MaterializeAssetResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, GetAssetData, GetAssetResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, NextRunAssetsData, NextRunAssetsResponse, ListBackfillsData, ListBackfillsResponse, CreateBackfillData, CreateBackfillResponse, GetBackfillData, GetBackfillResponse, PauseBackfillData, PauseBackfillResponse, UnpauseBackfillData, UnpauseBackfillResponse, CancelBackfillData, CancelBackfillResponse, CreateBackfillDryRunData, CreateBackfillDryRunResponse, ListBackfillsUiData, ListBackfillsUiResponse, DeleteConnectionData, DeleteConnectionResponse, GetConnectionData, GetConnectionResponse, PatchConnectionData, PatchConnectionResponse, GetConnectionsData, GetConnectionsResponse, PostConnectionData, PostConnectionResponse, BulkConnectionsData, BulkConnectionsResponse, TestConnectionData, TestConnectionResponse, CreateDefaultConnectionsResponse, HookMetaDataResponse, GetDagRunData, GetDagRunResponse, DeleteDagRunData, DeleteDagRunResponse, PatchDagRunData, PatchDagRunResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, ClearDagRunData, ClearDagRunResponse, GetDagRunsData, GetDagRunsResponse, TriggerDagRunData, TriggerDagRunResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, ListDagWarningsData, ListDagWarningsResponse, GetDagsData, GetDagsResponse, PatchDagsData, PatchDagsResponse, GetDagData, GetDagResponse, PatchDagData, PatchDagResponse, DeleteDagData, DeleteDagResponse, GetDagDetailsData, GetDagDetailsResponse, FavoriteDagData, FavoriteDagResponse, UnfavoriteDagData, UnfavoriteDagResponse, GetDagTagsData, GetDagTagsResponse, GetDagsUiData, GetDagsUiResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExtraLinksData, GetExtraLinksResponse, GetTaskInstanceData, GetTaskInstanceResponse, PatchTaskInstanceData, PatchTaskInstanceResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, GetTaskInstancesData, GetTaskInstancesResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, GetLogData, GetLogResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetPluginsData, GetPluginsResponse, ImportErrorsResponse, DeletePoolData, DeletePoolResponse, GetPoolData, GetPoolResponse, PatchPoolData, PatchPoolResponse, GetPoolsData, GetPoolsResponse, PostPoolData, PostPoolResponse, BulkPoolsData, BulkPoolsResponse, GetProvidersData, GetProvidersResponse, GetXcomEntryData, GetXcomEntryResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, GetXcomEntriesData, GetXcomEntriesResponse, CreateXcomEntryData, CreateXcomEntryResponse, GetTasksData, GetTasksResponse, GetTaskData, GetTaskResponse, DeleteVariableData, DeleteVariableResponse, GetVariableData, GetVariableResponse, PatchVariableData, PatchVariableResponse, GetVariablesData, GetVariablesResponse, PostVariableData, PostVariableResponse, BulkVariablesData, BulkVariablesResponse, ReparseDagFileData, ReparseDagFileResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetHealthResponse, GetVersionResponse, LoginData, LoginResponse, LogoutResponse, GetAuthMenusResponse, GetCurrentUserInfoResponse, GenerateTokenData, GenerateTokenResponse2, GetPartitionedDagRunsData, GetPartitionedDagRunsResponse, GetPendingPartitionedDagRunData, GetPendingPartitionedDagRunResponse, GetDependenciesData, GetDependenciesResponse, HistoricalMetricsData, HistoricalMetricsResponse, DagStatsResponse2, GetDagRunDeadlinesData, GetDagRunDeadlinesResponse, StructureDataData, StructureDataResponse2, GetDagStructureData, GetDagStructureResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesStreamData, GetGridTiSummariesStreamResponse, GetGanttDataData, GetGanttDataResponse, GetCalendarData, GetCalendarResponse, ListTeamsData, ListTeamsResponse } from './types.gen'; +import type { BulkConnectionsData, BulkConnectionsResponse, BulkPoolsData, BulkPoolsResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, BulkVariablesData, BulkVariablesResponse, CancelBackfillData, CancelBackfillResponse, ClearDagRunData, ClearDagRunResponse, CreateAssetEventData, CreateAssetEventResponse, CreateBackfillData, CreateBackfillDryRunData, CreateBackfillDryRunResponse, CreateBackfillResponse, CreateDefaultConnectionsResponse, CreateXcomEntryData, CreateXcomEntryResponse, DagStatsResponse2, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, DeleteConnectionData, DeleteConnectionResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, DeleteDagData, DeleteDagResponse, DeleteDagRunData, DeleteDagRunResponse, DeletePoolData, DeletePoolResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, DeleteVariableData, DeleteVariableResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, FavoriteDagData, FavoriteDagResponse, GenerateTokenData, GenerateTokenResponse2, GetAssetAliasData, GetAssetAliasResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetData, GetAssetEventsData, GetAssetEventsResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, GetAssetResponse, GetAssetsData, GetAssetsResponse, GetAuthMenusResponse, GetBackfillData, GetBackfillResponse, GetCalendarData, GetCalendarResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, GetConnectionData, GetConnectionResponse, GetConnectionTestStatusData, GetConnectionTestStatusResponse, GetConnectionsData, GetConnectionsResponse, GetCurrentUserInfoResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, GetDagData, GetDagDetailsData, GetDagDetailsResponse, GetDagResponse, GetDagRunData, GetDagRunDeadlinesData, GetDagRunDeadlinesResponse, GetDagRunResponse, GetDagRunsData, GetDagRunsResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetDagStructureData, GetDagStructureResponse, GetDagTagsData, GetDagTagsResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetDagsData, GetDagsResponse, GetDagsUiData, GetDagsUiResponse, GetDependenciesData, GetDependenciesResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, GetExtraLinksData, GetExtraLinksResponse, GetGanttDataData, GetGanttDataResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesData, GetGridTiSummariesResponse, GetGridTiSummariesStreamData, GetGridTiSummariesStreamResponse, GetHealthResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetLogData, GetLogResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetPartitionedDagRunsData, GetPartitionedDagRunsResponse, GetPendingPartitionedDagRunData, GetPendingPartitionedDagRunResponse, GetPluginsData, GetPluginsResponse, GetPoolData, GetPoolResponse, GetPoolsData, GetPoolsResponse, GetProvidersData, GetProvidersResponse, GetTaskData, GetTaskInstanceData, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstancesData, GetTaskInstancesResponse, GetTaskResponse, GetTasksData, GetTasksResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, GetVariableData, GetVariableResponse, GetVariablesData, GetVariablesResponse, GetVersionResponse, GetXcomEntriesData, GetXcomEntriesResponse, GetXcomEntryData, GetXcomEntryResponse, HistoricalMetricsData, HistoricalMetricsResponse, HookMetaDataResponse, ImportErrorsResponse, ListBackfillsData, ListBackfillsResponse, ListBackfillsUiData, ListBackfillsUiResponse, ListDagWarningsData, ListDagWarningsResponse, ListTeamsData, ListTeamsResponse, LoginData, LoginResponse, LogoutResponse, MaterializeAssetData, MaterializeAssetResponse, NextRunAssetsData, NextRunAssetsResponse, PatchConnectionData, PatchConnectionResponse, PatchDagData, PatchDagResponse, PatchDagRunData, PatchDagRunResponse, PatchDagsData, PatchDagsResponse, PatchPoolData, PatchPoolResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, PatchTaskInstanceData, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, PatchTaskInstanceResponse, PatchVariableData, PatchVariableResponse, PauseBackfillData, PauseBackfillResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PostConnectionData, PostConnectionResponse, PostPoolData, PostPoolResponse, PostVariableData, PostVariableResponse, ReparseDagFileData, ReparseDagFileResponse, StructureDataData, StructureDataResponse2, TestConnectionAsyncData, TestConnectionAsyncResponse, TestConnectionData, TestConnectionResponse, TriggerDagRunData, TriggerDagRunResponse, UnfavoriteDagData, UnfavoriteDagResponse, UnpauseBackfillData, UnpauseBackfillResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse } from './types.gen'; export class AssetService { /** @@ -794,6 +794,59 @@ export class ConnectionService { }); } + /** + * Test Connection Async + * Queue an async connection test to be executed on a worker. + * + * The connection must already be saved. Returns a token that can be used + * to poll for the test result via GET /connections/test-async/{token}. + * @param data The data for the request. + * @param data.requestBody + * @returns ConnectionTestQueuedResponse Successful Response + * @throws ApiError + */ + public static testConnectionAsync(data: TestConnectionAsyncData): CancelablePromise { + return __request(OpenAPI, { + method: 'POST', + url: '/api/v2/connections/test-async', + body: data.requestBody, + mediaType: 'application/json', + errors: { + 401: 'Unauthorized', + 403: 'Forbidden', + 404: 'Not Found', + 422: 'Validation Error' + } + }); + } + + /** + * Get Connection Test Status + * 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. + * @param data The data for the request. + * @param data.token + * @returns ConnectionTestStatusResponse Successful Response + * @throws ApiError + */ + public static getConnectionTestStatus(data: GetConnectionTestStatusData): CancelablePromise { + return __request(OpenAPI, { + method: 'GET', + url: '/api/v2/connections/test-async/{token}', + path: { + token: data.token + }, + errors: { + 401: 'Unauthorized', + 403: 'Forbidden', + 404: 'Not Found', + 422: 'Validation Error' + } + }); + } + /** * Create Default Connections * Create default connections. @@ -3871,55 +3924,6 @@ export class AuthLinksService { } -export class PartitionedDagRunService { - /** - * Get Partitioned Dag Runs - * Return PartitionedDagRuns. Filter by dag_id and/or has_created_dag_run_id. - * @param data The data for the request. - * @param data.dagId - * @param data.hasCreatedDagRunId - * @returns PartitionedDagRunCollectionResponse Successful Response - * @throws ApiError - */ - public static getPartitionedDagRuns(data: GetPartitionedDagRunsData = {}): CancelablePromise { - return __request(OpenAPI, { - method: 'GET', - url: '/ui/partitioned_dag_runs', - query: { - dag_id: data.dagId, - has_created_dag_run_id: data.hasCreatedDagRunId - }, - errors: { - 422: 'Validation Error' - } - }); - } - - /** - * Get Pending Partitioned Dag Run - * Return full details for pending PartitionedDagRun. - * @param data The data for the request. - * @param data.dagId - * @param data.partitionKey - * @returns PartitionedDagRunDetailResponse Successful Response - * @throws ApiError - */ - public static getPendingPartitionedDagRun(data: GetPendingPartitionedDagRunData): CancelablePromise { - return __request(OpenAPI, { - method: 'GET', - url: '/ui/pending_partitioned_dag_run/{dag_id}/{partition_key}', - path: { - dag_id: data.dagId, - partition_key: data.partitionKey - }, - errors: { - 422: 'Validation Error' - } - }); - } - -} - export class DependenciesService { /** * Get Dependencies @@ -3993,25 +3997,17 @@ export class DeadlinesService { * Get all deadlines for a specific DAG run. * @param data The data for the request. * @param data.dagId - * @param data.dagRunId - * @param data.limit - * @param data.offset - * @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, deadline_time, created_at, alert_name` - * @returns DeadlineCollectionResponse Successful Response + * @param data.runId + * @returns DeadlineResponse Successful Response * @throws ApiError */ public static getDagRunDeadlines(data: GetDagRunDeadlinesData): CancelablePromise { return __request(OpenAPI, { method: 'GET', - url: '/ui/dags/{dag_id}/dagRuns/{dag_run_id}/deadlines', + url: '/ui/deadlines/{dag_id}/{run_id}', path: { dag_id: data.dagId, - dag_run_id: data.dagRunId - }, - query: { - limit: data.limit, - offset: data.offset, - order_by: data.orderBy + run_id: data.runId }, errors: { 404: 'Not Found', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 7423c08d42c05..ba13d8db91d8e 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -496,13 +496,41 @@ export type ConnectionResponse = { }; /** - * Connection Test serializer for responses. + * Response returned when an async connection test is queued. + */ +export type ConnectionTestQueuedResponse = { + token: string; + connection_id: string; + state: string; +}; + +/** + * Request body for async connection test — just the connection_id. + */ +export type ConnectionTestRequestBody = { + connection_id: string; +}; + +/** + * Connection Test serializer for synchronous test responses. */ export type ConnectionTestResponse = { status: boolean; message: string; }; +/** + * Response returned when polling for async connection test status. + */ +export type ConnectionTestStatusResponse = { + token: string; + connection_id: string; + state: string; + result_status?: boolean | null; + result_message?: string | null; + created_at: string; +}; + /** * Create asset events request. */ @@ -2482,6 +2510,18 @@ export type TestConnectionData = { export type TestConnectionResponse = ConnectionTestResponse; +export type TestConnectionAsyncData = { + requestBody: ConnectionTestRequestBody; +}; + +export type TestConnectionAsyncResponse = ConnectionTestQueuedResponse; + +export type GetConnectionTestStatusData = { + token: string; +}; + +export type GetConnectionTestStatusResponse = ConnectionTestStatusResponse; + export type CreateDefaultConnectionsResponse = void; export type HookMetaDataResponse = Array; @@ -4465,6 +4505,60 @@ export type $OpenApiTs = { }; }; }; + '/api/v2/connections/test-async': { + post: { + req: TestConnectionAsyncData; + res: { + /** + * Successful Response + */ + 202: ConnectionTestQueuedResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; + '/api/v2/connections/test-async/{token}': { + get: { + req: GetConnectionTestStatusData; + res: { + /** + * Successful Response + */ + 200: ConnectionTestStatusResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; '/api/v2/connections/defaults': { post: { res: { diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py index e6b5283669b86..59333cb6b822b 100644 --- a/airflow-core/src/airflow/utils/db_cleanup.py +++ b/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", recency_column_name="created_at"), ] # We need to have `fallback="database"` because this is executed at top level code and provider configuration diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 1f63247cfa9ab..982f03133b7a7 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -29,6 +29,8 @@ from airflow.api_fastapi.core_api.datamodels.connections import ConnectionBody from airflow.api_fastapi.core_api.services.public.connections import BulkConnectionService from airflow.models import Connection +from airflow.models.callback import Callback +from airflow.models.connection_test import RUN_CONNECTION_TEST_PATH, ConnectionTest, ConnectionTestState from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.utils.session import NEW_SESSION, provide_session @@ -1168,6 +1170,111 @@ def test_should_test_new_connection_without_existing(self, test_client): assert response.json()["status"] is True +class TestAsyncConnectionTest(TestConnectionEndpoint): + """Tests for the async connection test endpoints (POST + GET polling).""" + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_post_should_respond_202(self, test_client, session): + """POST /connections/test-async with a saved connection returns 202 + token.""" + self.create_connection() + response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) + assert response.status_code == 202 + body = response.json() + assert "token" in body + assert body["connection_id"] == TEST_CONN_ID + assert body["state"] == "queued" + assert len(body["token"]) > 0 + + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/connections/test-async", json={"connection_id": TEST_CONN_ID} + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/connections/test-async", json={"connection_id": TEST_CONN_ID} + ) + assert response.status_code == 403 + + def test_should_respond_403_by_default(self, test_client): + """Connection testing is disabled by default.""" + response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) + assert response.status_code == 403 + assert response.json() == { + "detail": "Testing connections is disabled in Airflow configuration. " + "Contact your deployment admin to enable it." + } + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_should_respond_404_for_nonexistent_connection(self, test_client): + """Connection must be saved before testing.""" + response = test_client.post("/connections/test-async", json={"connection_id": "nonexistent"}) + assert response.status_code == 404 + assert "was not found" in response.json()["detail"] + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_post_creates_connection_test_and_callback(self, test_client, session): + """POST creates both a ConnectionTest row and an ExecutorCallback.""" + self.create_connection() + response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) + assert response.status_code == 202 + token = response.json()["token"] + + ct = session.scalar(select(ConnectionTest).filter_by(token=token)) + assert ct is not None + assert ct.connection_id == TEST_CONN_ID + assert ct.state == "queued" + assert ct.callback_id is not None + + cb = session.get(Callback, ct.callback_id) + assert cb is not None + assert cb.data["path"] == RUN_CONNECTION_TEST_PATH + assert cb.data["kwargs"]["connection_id"] == TEST_CONN_ID + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_get_status_returns_queued(self, test_client, session): + """GET /connections/test-async/{token} returns current status.""" + self.create_connection() + post_response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) + token = post_response.json()["token"] + + response = test_client.get(f"/connections/test-async/{token}") + assert response.status_code == 200 + body = response.json() + assert body["token"] == token + assert body["connection_id"] == TEST_CONN_ID + assert body["state"] == "queued" + assert body["result_status"] is None + assert body["result_message"] is None + assert "created_at" in body + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_get_status_returns_completed_result(self, test_client, session): + """GET returns result after the worker has updated the ConnectionTest.""" + self.create_connection() + post_response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) + token = post_response.json()["token"] + + ct = session.scalar(select(ConnectionTest).filter_by(token=token)) + ct.state = ConnectionTestState.SUCCESS + ct.result_status = True + ct.result_message = "Connection successfully tested" + session.commit() + + response = test_client.get(f"/connections/test-async/{token}") + assert response.status_code == 200 + body = response.json() + assert body["state"] == "success" + assert body["result_status"] is True + assert body["result_message"] == "Connection successfully tested" + + def test_get_status_returns_404_for_invalid_token(self, test_client): + """GET with an unknown token returns 404.""" + response = test_client.get("/connections/test-async/nonexistent-token") + assert response.status_code == 404 + + class TestCreateDefaultConnections(TestConnectionEndpoint): def test_should_respond_204(self, test_client, session): response = test_client.post("/connections/defaults") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py new file mode 100644 index 0000000000000..b17ddd06ed8c3 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py @@ -0,0 +1,80 @@ +# 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 pytest + +from airflow.models.connection_test import ConnectionTest, ConnectionTestState + +pytestmark = pytest.mark.db_test + + +class TestPatchConnectionTest: + def test_patch_updates_result(self, client, session): + """PATCH sets the state and result fields.""" + ct = ConnectionTest(connection_id="test_conn") + ct.state = ConnectionTestState.RUNNING + session.add(ct) + session.commit() + + response = client.patch( + f"/execution/connection-tests/{ct.id}", + json={ + "state": "success", + "result_status": True, + "result_message": "Connection successfully tested", + }, + ) + assert response.status_code == 204 + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == "success" + assert ct.result_status is True + assert ct.result_message == "Connection successfully tested" + + def test_patch_returns_404_for_nonexistent(self, client): + """PATCH with unknown id returns 404.""" + response = client.patch( + "/execution/connection-tests/00000000-0000-0000-0000-000000000000", + json={"state": "success", "result_status": True, "result_message": "ok"}, + ) + assert response.status_code == 404 + + def test_patch_returns_422_for_invalid_uuid(self, client): + """PATCH with invalid uuid returns 422.""" + response = client.patch( + "/execution/connection-tests/not-a-uuid", + json={"state": "success", "result_status": True, "result_message": "ok"}, + ) + assert response.status_code == 422 + + def test_patch_returns_409_for_terminal_state(self, client, session): + """PATCH on a test already in terminal state returns 409.""" + ct = ConnectionTest(connection_id="test_conn") + ct.state = ConnectionTestState.SUCCESS + ct.result_status = True + ct.result_message = "Already done" + session.add(ct) + session.commit() + + response = client.patch( + f"/execution/connection-tests/{ct.id}", + json={"state": "failed", "result_status": False, "result_message": "retry"}, + ) + assert response.status_code == 409 + assert "terminal state" in response.json()["detail"]["message"] diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py new file mode 100644 index 0000000000000..c4346cbc7b326 --- /dev/null +++ b/airflow-core/tests/unit/models/test_connection_test.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. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.models.connection import Connection +from airflow.models.connection_test import ConnectionTest, ConnectionTestState, run_connection_test + +pytestmark = pytest.mark.db_test + + +class TestConnectionTestModel: + def test_token_is_generated(self): + ct = ConnectionTest(connection_id="test_conn") + assert ct.token is not None + assert len(ct.token) > 0 + + def test_initial_state_is_pending(self): + ct = ConnectionTest(connection_id="test_conn") + assert ct.state == ConnectionTestState.PENDING + + def test_tokens_are_unique(self): + ct1 = ConnectionTest(connection_id="test_conn") + ct2 = ConnectionTest(connection_id="test_conn") + assert ct1.token != ct2.token + + def test_repr(self): + ct = ConnectionTest(connection_id="test_conn") + r = repr(ct) + assert "test_conn" in r + assert "pending" in r + + +class TestRunConnectionTest: + def test_successful_connection_test(self, session): + """Worker function updates state to SUCCESS on successful test.""" + ct = ConnectionTest(connection_id="test_conn") + session.add(ct) + session.commit() + ct_id = str(ct.id) + + with mock.patch.object( + Connection, "get_connection_from_secrets", return_value=mock.MagicMock() + ) as mock_get_conn: + mock_get_conn.return_value.test_connection.return_value = (True, "Connection OK") + run_connection_test(connection_id="test_conn", connection_test_id=ct_id) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.SUCCESS + assert ct.result_status is True + assert ct.result_message == "Connection OK" + + def test_failed_connection_test(self, session): + """Worker function updates state to FAILED when test_connection returns False.""" + ct = ConnectionTest(connection_id="test_conn") + session.add(ct) + session.commit() + ct_id = str(ct.id) + + with mock.patch.object( + Connection, "get_connection_from_secrets", return_value=mock.MagicMock() + ) as mock_get_conn: + mock_get_conn.return_value.test_connection.return_value = (False, "Connection failed") + run_connection_test(connection_id="test_conn", connection_test_id=ct_id) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert ct.result_status is False + assert ct.result_message == "Connection failed" + + def test_exception_during_connection_test(self, session): + """Worker function handles exceptions gracefully.""" + ct = ConnectionTest(connection_id="test_conn") + session.add(ct) + session.commit() + ct_id = str(ct.id) + + with mock.patch.object( + Connection, + "get_connection_from_secrets", + side_effect=Exception("Could not resolve host: db.example.com"), + ): + run_connection_test(connection_id="test_conn", connection_test_id=ct_id) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert ct.result_status is False + assert "Could not resolve host" in ct.result_message diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index bb375f888349c..5e70758038cde 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -244,15 +244,49 @@ 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 — just the connection_id. + """ + + model_config = ConfigDict( + extra="forbid", + ) + connection_id: Annotated[str, Field(title="Connection Id")] + + 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): + """ + Response returned when polling for async connection test status. + """ + + token: Annotated[str, Field(title="Token")] + connection_id: Annotated[str, Field(title="Connection Id")] + state: Annotated[str, Field(title="State")] + result_status: Annotated[bool | None, Field(title="Result Status")] = None + result_message: Annotated[str | None, Field(title="Result Message")] = None + created_at: Annotated[datetime, Field(title="Created At")] + + class CreateAssetEventsBody(BaseModel): """ Create asset events request. diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index b6c08e9d76c82..f9e8a088183f7 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -78,6 +78,18 @@ class ConnectionResponse(BaseModel): extra: Annotated[str | None, Field(title="Extra")] = None +class ConnectionTestState(str, Enum): + """ + All possible states of a connection test. + """ + + PENDING = "pending" + QUEUED = "queued" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + + class DagResponse(BaseModel): """ Schema for DAG response. @@ -535,6 +547,19 @@ class AssetResponse(BaseModel): extra: Annotated[dict[str, JsonValue] | None, Field(title="Extra")] = None +class ConnectionTestResultBody(BaseModel): + """ + Payload sent by workers to report connection test results. + """ + + model_config = ConfigDict( + extra="forbid", + ) + state: ConnectionTestState + result_status: Annotated[bool | None, Field(title="Result Status")] = None + result_message: Annotated[str | None, Field(title="Result Message")] = None + + class HITLDetailRequest(BaseModel): """ Schema for the request part of a Human-in-the-loop detail for a specific task instance. From 928699e729a0b317292373f165f5fedcc1932035 Mon Sep 17 00:00:00 2001 From: Anish Date: Mon, 23 Feb 2026 01:38:02 -0600 Subject: [PATCH 02/38] fix test in ci --- .../versions/0107_3_2_0_add_connection_test_table.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py index a3ccef07872e0..d97e0aa15ead1 100644 --- a/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py @@ -30,6 +30,8 @@ import sqlalchemy as sa from alembic import op +from airflow.utils.sqlalchemy import UtcDateTime + # revision identifiers, used by Alembic. revision = "a7e6d4c3b2f1" down_revision = "509b94a1042d" @@ -48,8 +50,8 @@ def upgrade(): sa.Column("state", sa.String(10), nullable=False), sa.Column("result_status", sa.Boolean(), nullable=True), sa.Column("result_message", sa.Text(), nullable=True), - sa.Column("created_at", sa.DateTime(), nullable=False), - sa.Column("updated_at", sa.DateTime(), nullable=False), + sa.Column("created_at", UtcDateTime(timezone=True), nullable=False), + sa.Column("updated_at", UtcDateTime(timezone=True), nullable=False), sa.Column("callback_id", sa.Uuid(), nullable=True), sa.PrimaryKeyConstraint("id", name=op.f("connection_test_pkey")), sa.ForeignKeyConstraint( From 2324221fdf66daac9339dc356a06d089099e736c Mon Sep 17 00:00:00 2001 From: Anish Date: Mon, 23 Feb 2026 10:36:25 -0600 Subject: [PATCH 03/38] Address review feedback: remove redundant result_status, move callback factory to model --- .../core_api/datamodels/connections.py | 1 - .../openapi/v2-rest-api-generated.yaml | 5 --- .../core_api/routes/public/connections.py | 34 +---------------- .../datamodels/connection_test.py | 1 - .../execution_api/routes/connection_tests.py | 1 - .../0107_3_2_0_add_connection_test_table.py | 1 - .../src/airflow/models/connection_test.py | 38 +++++++++++++++++-- .../ui/openapi-gen/requests/schemas.gen.ts | 11 ------ .../ui/openapi-gen/requests/types.gen.ts | 1 - .../routes/public/test_connections.py | 3 -- .../versions/head/test_connection_tests.py | 9 ++--- .../tests/unit/models/test_connection_test.py | 3 -- .../airflowctl/api/datamodels/generated.py | 1 - .../airflow/sdk/api/datamodels/_generated.py | 1 - 14 files changed, 38 insertions(+), 72 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index 936b6832fec85..9192926a0c77f 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -99,7 +99,6 @@ class ConnectionTestStatusResponse(BaseModel): token: str connection_id: str state: str - result_status: bool | None = None result_message: str | None = None created_at: datetime diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 57bdb6d00fd65..74e5b8e13a59b 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -10249,11 +10249,6 @@ components: state: type: string title: State - result_status: - anyOf: - - type: boolean - - type: 'null' - title: Result Status result_message: anyOf: - type: string diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 47e6e4b0a207a..9d0505f381a7e 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -60,9 +60,7 @@ from airflow.configuration import conf from airflow.exceptions import AirflowNotFoundException from airflow.models import Connection -from airflow.models.callback import CallbackFetchMethod, ExecutorCallback from airflow.models.connection_test import ( - RUN_CONNECTION_TEST_PATH, ConnectionTest, ConnectionTestState, ) @@ -73,25 +71,6 @@ connections_router = AirflowRouter(tags=["Connection"], prefix="/connections") -class _ImportPathCallbackDef: - """ - Minimal implementation of ImportPathExecutorCallbackDefProtocol. - - ExecutorCallback.__init__ expects an object satisfying this protocol, but no concrete - implementation exists in airflow-core — the only one (SyncCallback) lives in the task-sdk. - Once #61153 lands and ExecuteCallback.make() is decoupled from DagRun, this adapter can - be replaced with the proper factory method. - """ - - def __init__(self, path: str, kwargs: dict, executor: str | None = None): - self.path = path - self.kwargs = kwargs - self.executor = executor - - def serialize(self) -> dict: - return {"path": self.path, "kwargs": self.kwargs, "executor": self.executor} - - @connections_router.delete( "/{connection_id}", status_code=status.HTTP_204_NO_CONTENT, @@ -321,17 +300,7 @@ def test_connection_async( session.add(connection_test) session.flush() - # ExecutorCallback requires an object satisfying ImportPathExecutorCallbackDefProtocol, - # but the only concrete impl (SyncCallback) lives in task-sdk which core API cannot - # import. This adapter will be replaced by ExecuteCallback.make() once #61153 merges. - callback_def = _ImportPathCallbackDef( - path=RUN_CONNECTION_TEST_PATH, - kwargs={ - "connection_id": test_body.connection_id, - "connection_test_id": str(connection_test.id), - }, - ) - callback = ExecutorCallback(callback_def, fetch_method=CallbackFetchMethod.IMPORT_PATH) + callback = connection_test.create_callback() session.add(callback) session.flush() @@ -372,7 +341,6 @@ def get_connection_test_status( token=connection_test.token, connection_id=connection_test.connection_id, state=connection_test.state, - result_status=connection_test.result_status, result_message=connection_test.result_message, created_at=connection_test.created_at, ) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py index 6042906918ab8..3fd046e90b6e7 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py @@ -25,5 +25,4 @@ class ConnectionTestResultBody(StrictBaseModel): """Payload sent by workers to report connection test results.""" state: ConnectionTestState - result_status: bool | None = None result_message: str | None = None diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py index d497b1c6e6db3..c61a7b3f4ff3e 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -61,5 +61,4 @@ def patch_connection_test( ) ct.state = body.state - ct.result_status = body.result_status ct.result_message = body.result_message diff --git a/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py index d97e0aa15ead1..02c7ebbf2f4c0 100644 --- a/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py @@ -48,7 +48,6 @@ def upgrade(): sa.Column("token", sa.String(64), nullable=False), sa.Column("connection_id", sa.String(250), nullable=False), sa.Column("state", sa.String(10), nullable=False), - sa.Column("result_status", sa.Boolean(), nullable=True), 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), diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 898debd7e62c7..89c0a234254ee 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -24,7 +24,7 @@ import structlog import uuid6 -from sqlalchemy import Boolean, ForeignKey, Index, String, Text, Uuid +from sqlalchemy import ForeignKey, Index, String, Text, Uuid from sqlalchemy.orm import Mapped, mapped_column, relationship from airflow._shared.timezones import timezone @@ -32,7 +32,7 @@ from airflow.utils.sqlalchemy import UtcDateTime if TYPE_CHECKING: - from airflow.models.callback import Callback + from airflow.models.callback import Callback, ExecutorCallback log = structlog.get_logger(__name__) @@ -56,6 +56,25 @@ def __str__(self) -> str: RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test" +class _ImportPathCallbackDef: + """ + Minimal implementation of ImportPathExecutorCallbackDefProtocol. + + ExecutorCallback.__init__ expects an object satisfying this protocol, but no concrete + implementation exists in airflow-core — the only one (SyncCallback) lives in the task-sdk. + Once #61153 lands and ExecuteCallback.make() is decoupled from DagRun, this adapter can + be replaced with the proper factory method. + """ + + def __init__(self, path: str, kwargs: dict, executor: str | None = None): + self.path = path + self.kwargs = kwargs + self.executor = executor + + def serialize(self) -> dict: + return {"path": self.path, "kwargs": self.kwargs, "executor": self.executor} + + class ConnectionTest(Base): """Tracks an async connection test dispatched to a worker via ExecutorCallback.""" @@ -65,7 +84,6 @@ class ConnectionTest(Base): token: Mapped[str] = mapped_column(String(64), nullable=False, unique=True) connection_id: Mapped[str] = mapped_column(String(250), nullable=False) state: Mapped[str] = mapped_column(String(10), nullable=False, default=ConnectionTestState.PENDING) - result_status: Mapped[bool | None] = mapped_column(Boolean, nullable=True) 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( @@ -90,6 +108,19 @@ def __repr__(self) -> str: f"" ) + def create_callback(self) -> ExecutorCallback: + """Create an ExecutorCallback that will run the connection test on a worker.""" + from airflow.models.callback import CallbackFetchMethod, ExecutorCallback + + callback_def = _ImportPathCallbackDef( + path=RUN_CONNECTION_TEST_PATH, + kwargs={ + "connection_id": self.connection_id, + "connection_test_id": str(self.id), + }, + ) + return ExecutorCallback(callback_def, fetch_method=CallbackFetchMethod.IMPORT_PATH) + def run_connection_test(*, connection_id: str, connection_test_id: str) -> None: """ @@ -120,6 +151,5 @@ def run_connection_test(*, connection_id: str, connection_test_id: str) -> None: with create_session() as session: ct = session.get(ConnectionTest, connection_test_uuid) if ct: - ct.result_status = test_status ct.result_message = test_message ct.state = ConnectionTestState.SUCCESS if test_status else ConnectionTestState.FAILED diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 3da3229f7e82f..5c0987aef17ec 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1764,17 +1764,6 @@ export const $ConnectionTestStatusResponse = { type: 'string', title: 'State' }, - result_status: { - anyOf: [ - { - type: 'boolean' - }, - { - type: 'null' - } - ], - title: 'Result Status' - }, result_message: { anyOf: [ { diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index ba13d8db91d8e..cee0a9fe60962 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -526,7 +526,6 @@ export type ConnectionTestStatusResponse = { token: string; connection_id: string; state: string; - result_status?: boolean | null; result_message?: string | null; created_at: string; }; diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 982f03133b7a7..69a884a55665e 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -1245,7 +1245,6 @@ def test_get_status_returns_queued(self, test_client, session): assert body["token"] == token assert body["connection_id"] == TEST_CONN_ID assert body["state"] == "queued" - assert body["result_status"] is None assert body["result_message"] is None assert "created_at" in body @@ -1258,7 +1257,6 @@ def test_get_status_returns_completed_result(self, test_client, session): ct = session.scalar(select(ConnectionTest).filter_by(token=token)) ct.state = ConnectionTestState.SUCCESS - ct.result_status = True ct.result_message = "Connection successfully tested" session.commit() @@ -1266,7 +1264,6 @@ def test_get_status_returns_completed_result(self, test_client, session): assert response.status_code == 200 body = response.json() assert body["state"] == "success" - assert body["result_status"] is True assert body["result_message"] == "Connection successfully tested" def test_get_status_returns_404_for_invalid_token(self, test_client): diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py index b17ddd06ed8c3..1951d90f6f13a 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py @@ -35,7 +35,6 @@ def test_patch_updates_result(self, client, session): f"/execution/connection-tests/{ct.id}", json={ "state": "success", - "result_status": True, "result_message": "Connection successfully tested", }, ) @@ -44,14 +43,13 @@ def test_patch_updates_result(self, client, session): session.expire_all() ct = session.get(ConnectionTest, ct.id) assert ct.state == "success" - assert ct.result_status is True assert ct.result_message == "Connection successfully tested" def test_patch_returns_404_for_nonexistent(self, client): """PATCH with unknown id returns 404.""" response = client.patch( "/execution/connection-tests/00000000-0000-0000-0000-000000000000", - json={"state": "success", "result_status": True, "result_message": "ok"}, + json={"state": "success", "result_message": "ok"}, ) assert response.status_code == 404 @@ -59,7 +57,7 @@ def test_patch_returns_422_for_invalid_uuid(self, client): """PATCH with invalid uuid returns 422.""" response = client.patch( "/execution/connection-tests/not-a-uuid", - json={"state": "success", "result_status": True, "result_message": "ok"}, + json={"state": "success", "result_message": "ok"}, ) assert response.status_code == 422 @@ -67,14 +65,13 @@ def test_patch_returns_409_for_terminal_state(self, client, session): """PATCH on a test already in terminal state returns 409.""" ct = ConnectionTest(connection_id="test_conn") ct.state = ConnectionTestState.SUCCESS - ct.result_status = True ct.result_message = "Already done" session.add(ct) session.commit() response = client.patch( f"/execution/connection-tests/{ct.id}", - json={"state": "failed", "result_status": False, "result_message": "retry"}, + json={"state": "failed", "result_message": "retry"}, ) assert response.status_code == 409 assert "terminal state" in response.json()["detail"]["message"] diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py index c4346cbc7b326..34d998ef22c98 100644 --- a/airflow-core/tests/unit/models/test_connection_test.py +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -65,7 +65,6 @@ def test_successful_connection_test(self, session): session.expire_all() ct = session.get(ConnectionTest, ct.id) assert ct.state == ConnectionTestState.SUCCESS - assert ct.result_status is True assert ct.result_message == "Connection OK" def test_failed_connection_test(self, session): @@ -84,7 +83,6 @@ def test_failed_connection_test(self, session): session.expire_all() ct = session.get(ConnectionTest, ct.id) assert ct.state == ConnectionTestState.FAILED - assert ct.result_status is False assert ct.result_message == "Connection failed" def test_exception_during_connection_test(self, session): @@ -104,5 +102,4 @@ def test_exception_during_connection_test(self, session): session.expire_all() ct = session.get(ConnectionTest, ct.id) assert ct.state == ConnectionTestState.FAILED - assert ct.result_status is False assert "Could not resolve host" in ct.result_message diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index 5e70758038cde..0734ca8a271b6 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -282,7 +282,6 @@ class ConnectionTestStatusResponse(BaseModel): token: Annotated[str, Field(title="Token")] connection_id: Annotated[str, Field(title="Connection Id")] state: Annotated[str, Field(title="State")] - result_status: Annotated[bool | None, Field(title="Result Status")] = None result_message: Annotated[str | None, Field(title="Result Message")] = None created_at: Annotated[datetime, Field(title="Created At")] diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index f9e8a088183f7..87d7bc73004a1 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -556,7 +556,6 @@ class ConnectionTestResultBody(BaseModel): extra="forbid", ) state: ConnectionTestState - result_status: Annotated[bool | None, Field(title="Result Status")] = None result_message: Annotated[str | None, Field(title="Result Message")] = None From 63b012651bcef4c2f83d3732b3c42b1b9c88753c Mon Sep 17 00:00:00 2001 From: Anish Date: Tue, 24 Feb 2026 14:49:48 -0600 Subject: [PATCH 04/38] ci test fix and rebase issues --- .../airflow/ui/openapi-gen/queries/queries.ts | 42 ++++++++++-- .../ui/openapi-gen/requests/services.gen.ts | 65 +++++++++++++++++-- 2 files changed, 98 insertions(+), 9 deletions(-) diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index 2b32fe404bc47..a6be3842c110e 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -1515,6 +1515,32 @@ export const useAuthLinksServiceGetAuthMenus = = unknown[]>(queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseAuthLinksServiceGetCurrentUserInfoKeyFn(queryKey), queryFn: () => AuthLinksService.getCurrentUserInfo() as TData, ...options }); /** +* Get Partitioned Dag Runs +* Return PartitionedDagRuns. Filter by dag_id and/or has_created_dag_run_id. +* @param data The data for the request. +* @param data.dagId +* @param data.hasCreatedDagRunId +* @returns PartitionedDagRunCollectionResponse Successful Response +* @throws ApiError +*/ +export const usePartitionedDagRunServiceGetPartitionedDagRuns = = unknown[]>({ dagId, hasCreatedDagRunId }: { + dagId?: string; + hasCreatedDagRunId?: boolean; +} = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn({ dagId, hasCreatedDagRunId }, queryKey), queryFn: () => PartitionedDagRunService.getPartitionedDagRuns({ dagId, hasCreatedDagRunId }) as TData, ...options }); +/** +* Get Pending Partitioned Dag Run +* Return full details for pending PartitionedDagRun. +* @param data The data for the request. +* @param data.dagId +* @param data.partitionKey +* @returns PartitionedDagRunDetailResponse Successful Response +* @throws ApiError +*/ +export const usePartitionedDagRunServiceGetPendingPartitionedDagRun = = unknown[]>({ dagId, partitionKey }: { + dagId: string; + partitionKey: string; +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId, partitionKey }, queryKey), queryFn: () => PartitionedDagRunService.getPendingPartitionedDagRun({ dagId, partitionKey }) as TData, ...options }); +/** * Get Dependencies * Dependencies graph. * @param data The data for the request. @@ -1552,14 +1578,20 @@ export const useDashboardServiceDagStats = = unknown[]>({ dagId, runId }: { +export const useDeadlinesServiceGetDagRunDeadlines = = unknown[]>({ dagId, dagRunId, limit, offset, orderBy }: { dagId: string; - runId: string; -}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, runId }, queryKey), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, runId }) as TData, ...options }); + dagRunId: string; + limit?: number; + offset?: number; + orderBy?: string[]; +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseDeadlinesServiceGetDagRunDeadlinesKeyFn({ dagId, dagRunId, limit, offset, orderBy }, queryKey), queryFn: () => DeadlinesService.getDagRunDeadlines({ dagId, dagRunId, limit, offset, orderBy }) as TData, ...options }); /** * Structure Data * Get Structure Data. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 077a91fd8edb8..e6230d3f8e44b 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -3924,6 +3924,55 @@ export class AuthLinksService { } +export class PartitionedDagRunService { + /** + * Get Partitioned Dag Runs + * Return PartitionedDagRuns. Filter by dag_id and/or has_created_dag_run_id. + * @param data The data for the request. + * @param data.dagId + * @param data.hasCreatedDagRunId + * @returns PartitionedDagRunCollectionResponse Successful Response + * @throws ApiError + */ + public static getPartitionedDagRuns(data: GetPartitionedDagRunsData = {}): CancelablePromise { + return __request(OpenAPI, { + method: 'GET', + url: '/ui/partitioned_dag_runs', + query: { + dag_id: data.dagId, + has_created_dag_run_id: data.hasCreatedDagRunId + }, + errors: { + 422: 'Validation Error' + } + }); + } + + /** + * Get Pending Partitioned Dag Run + * Return full details for pending PartitionedDagRun. + * @param data The data for the request. + * @param data.dagId + * @param data.partitionKey + * @returns PartitionedDagRunDetailResponse Successful Response + * @throws ApiError + */ + public static getPendingPartitionedDagRun(data: GetPendingPartitionedDagRunData): CancelablePromise { + return __request(OpenAPI, { + method: 'GET', + url: '/ui/pending_partitioned_dag_run/{dag_id}/{partition_key}', + path: { + dag_id: data.dagId, + partition_key: data.partitionKey + }, + errors: { + 422: 'Validation Error' + } + }); + } + +} + export class DependenciesService { /** * Get Dependencies @@ -3997,17 +4046,25 @@ export class DeadlinesService { * Get all deadlines for a specific DAG run. * @param data The data for the request. * @param data.dagId - * @param data.runId - * @returns DeadlineResponse Successful Response + * @param data.dagRunId + * @param data.limit + * @param data.offset + * @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, deadline_time, created_at, alert_name` + * @returns DealineCollectionResponse Successful Response * @throws ApiError */ public static getDagRunDeadlines(data: GetDagRunDeadlinesData): CancelablePromise { return __request(OpenAPI, { method: 'GET', - url: '/ui/deadlines/{dag_id}/{run_id}', + url: '/ui/dags/{dag_id}/dagRuns/{dag_run_id}/deadlines', path: { dag_id: data.dagId, - run_id: data.runId + dag_run_id: data.dagRunId + }, + query: { + limit: data.limit, + offset: data.offset, + order_by: data.orderBy }, errors: { 404: 'Not Found', From d1e9d318612c0eee0de0132d927db4da24d04b4d Mon Sep 17 00:00:00 2001 From: Anish Date: Wed, 25 Feb 2026 02:16:09 -0600 Subject: [PATCH 05/38] refactor and cleanups --- .../openapi/v2-rest-api-generated.yaml | 6 ++-- .../core_api/routes/public/connections.py | 32 +++++++++---------- .../src/airflow/models/connection_test.py | 3 +- .../airflow/ui/openapi-gen/queries/common.ts | 6 ++-- .../ui/openapi-gen/queries/ensureQueryData.ts | 8 ++--- .../ui/openapi-gen/queries/prefetch.ts | 8 ++--- .../airflow/ui/openapi-gen/queries/queries.ts | 8 ++--- .../ui/openapi-gen/queries/suspense.ts | 8 ++--- .../ui/openapi-gen/requests/services.gen.ts | 6 ++-- .../ui/openapi-gen/requests/types.gen.ts | 4 +-- .../routes/public/test_connections.py | 2 +- 11 files changed, 46 insertions(+), 45 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 74e5b8e13a59b..37e8994cc870f 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -1733,7 +1733,7 @@ paths: security: - OAuth2PasswordBearer: [] - HTTPBearer: [] - /api/v2/connections/test-async/{token}: + /api/v2/connections/test-async/{connection_test_token}: get: tags: - Connection @@ -1743,12 +1743,12 @@ paths: \ the test knows the crypto-random token." operationId: get_connection_test_status parameters: - - name: token + - name: connection_test_token in: path required: true schema: type: string - title: Token + title: Connection Test Token responses: '200': description: Successful Response diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 9d0505f381a7e..4c71a2c962372 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -71,6 +71,16 @@ 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.", + ) + + @connections_router.delete( "/{connection_id}", status_code=status.HTTP_204_NO_CONTENT, @@ -236,12 +246,7 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse: as some hook classes tries to find out the `conn` from their __init__ method & errors out if not found. It also deletes the conn id env connection after the test. """ - 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.", - ) + _ensure_test_connection_enabled() transient_conn_id = get_random_string() conn_env_var = f"{CONN_ENV_PREFIX}{transient_conn_id.upper()}" @@ -280,12 +285,7 @@ def test_connection_async( The connection must already be saved. Returns a token that can be used to poll for the test result via GET /connections/test-async/{token}. """ - if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled": - raise HTTPException( - status.HTTP_403_FORBIDDEN, - "Testing connections is disabled in Airflow configuration. " - "Contact your deployment admin to enable it.", - ) + _ensure_test_connection_enabled() try: Connection.get_connection_from_secrets(test_body.connection_id) @@ -316,11 +316,11 @@ def test_connection_async( @connections_router.get( - "/test-async/{token}", + "/test-async/{connection_test_token}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), ) def get_connection_test_status( - token: str, + connection_test_token: str, session: SessionDep, ) -> ConnectionTestStatusResponse: """ @@ -329,12 +329,12 @@ def get_connection_test_status( Knowledge of the token serves as authorization — only the client that initiated the test knows the crypto-random token. """ - connection_test = session.scalar(select(ConnectionTest).filter_by(token=token)) + connection_test = session.scalar(select(ConnectionTest).filter_by(token=connection_test_token)) if connection_test is None: raise HTTPException( status.HTTP_404_NOT_FOUND, - f"No connection test found for token: `{token}`", + f"No connection test found for token: `{connection_test_token}`", ) return ConnectionTestStatusResponse( diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 89c0a234254ee..3c9c424006d20 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -105,7 +105,8 @@ def __init__(self, *, connection_id: str, **kwargs): def __repr__(self) -> str: return ( - f"" + f"" ) def create_callback(self) -> ExecutorCallback: diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts index 9034236d66574..e0e249462d833 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -125,9 +125,9 @@ export const UseConnectionServiceGetConnectionsKeyFn = ({ connectionIdPattern, l export type ConnectionServiceGetConnectionTestStatusDefaultResponse = Awaited>; export type ConnectionServiceGetConnectionTestStatusQueryResult = UseQueryResult; export const useConnectionServiceGetConnectionTestStatusKey = "ConnectionServiceGetConnectionTestStatus"; -export const UseConnectionServiceGetConnectionTestStatusKeyFn = ({ token }: { - token: string; -}, queryKey?: Array) => [useConnectionServiceGetConnectionTestStatusKey, ...(queryKey ?? [{ token }])]; +export const UseConnectionServiceGetConnectionTestStatusKeyFn = ({ connectionTestToken }: { + connectionTestToken: string; +}, queryKey?: Array) => [useConnectionServiceGetConnectionTestStatusKey, ...(queryKey ?? [{ connectionTestToken }])]; export type ConnectionServiceHookMetaDataDefaultResponse = Awaited>; export type ConnectionServiceHookMetaDataQueryResult = UseQueryResult; export const useConnectionServiceHookMetaDataKey = "ConnectionServiceHookMetaData"; diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts index 2dae237e1ae82..5e551e0ab9375 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -231,13 +231,13 @@ export const ensureUseConnectionServiceGetConnectionsData = (queryClient: QueryC * Knowledge of the token serves as authorization — only the client * that initiated the test knows the crypto-random token. * @param data The data for the request. -* @param data.token +* @param data.connectionTestToken * @returns ConnectionTestStatusResponse Successful Response * @throws ApiError */ -export const ensureUseConnectionServiceGetConnectionTestStatusData = (queryClient: QueryClient, { token }: { - token: string; -}) => queryClient.ensureQueryData({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ token }), queryFn: () => ConnectionService.getConnectionTestStatus({ token }) }); +export const ensureUseConnectionServiceGetConnectionTestStatusData = (queryClient: QueryClient, { connectionTestToken }: { + connectionTestToken: string; +}) => queryClient.ensureQueryData({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ connectionTestToken }), queryFn: () => ConnectionService.getConnectionTestStatus({ connectionTestToken }) }); /** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts index db0fcf85b25b4..5ab8f3ab76a89 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -231,13 +231,13 @@ export const prefetchUseConnectionServiceGetConnections = (queryClient: QueryCli * Knowledge of the token serves as authorization — only the client * that initiated the test knows the crypto-random token. * @param data The data for the request. -* @param data.token +* @param data.connectionTestToken * @returns ConnectionTestStatusResponse Successful Response * @throws ApiError */ -export const prefetchUseConnectionServiceGetConnectionTestStatus = (queryClient: QueryClient, { token }: { - token: string; -}) => queryClient.prefetchQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ token }), queryFn: () => ConnectionService.getConnectionTestStatus({ token }) }); +export const prefetchUseConnectionServiceGetConnectionTestStatus = (queryClient: QueryClient, { connectionTestToken }: { + connectionTestToken: string; +}) => queryClient.prefetchQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ connectionTestToken }), queryFn: () => ConnectionService.getConnectionTestStatus({ connectionTestToken }) }); /** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index a6be3842c110e..11166fa1bd484 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -231,13 +231,13 @@ export const useConnectionServiceGetConnections = = unknown[]>({ token }: { - token: string; -}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ token }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ token }) as TData, ...options }); +export const useConnectionServiceGetConnectionTestStatus = = unknown[]>({ connectionTestToken }: { + connectionTestToken: string; +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ connectionTestToken }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ connectionTestToken }) as TData, ...options }); /** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts index 762067e621653..6c5c042bca865 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -231,13 +231,13 @@ export const useConnectionServiceGetConnectionsSuspense = = unknown[]>({ token }: { - token: string; -}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ token }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ token }) as TData, ...options }); +export const useConnectionServiceGetConnectionTestStatusSuspense = = unknown[]>({ connectionTestToken }: { + connectionTestToken: string; +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ connectionTestToken }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ connectionTestToken }) as TData, ...options }); /** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index e6230d3f8e44b..f4769cc7e269f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -827,16 +827,16 @@ export class ConnectionService { * Knowledge of the token serves as authorization — only the client * that initiated the test knows the crypto-random token. * @param data The data for the request. - * @param data.token + * @param data.connectionTestToken * @returns ConnectionTestStatusResponse Successful Response * @throws ApiError */ public static getConnectionTestStatus(data: GetConnectionTestStatusData): CancelablePromise { return __request(OpenAPI, { method: 'GET', - url: '/api/v2/connections/test-async/{token}', + url: '/api/v2/connections/test-async/{connection_test_token}', path: { - token: data.token + connection_test_token: data.connectionTestToken }, errors: { 401: 'Unauthorized', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index cee0a9fe60962..80981e883a59e 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2516,7 +2516,7 @@ export type TestConnectionAsyncData = { export type TestConnectionAsyncResponse = ConnectionTestQueuedResponse; export type GetConnectionTestStatusData = { - token: string; + connectionTestToken: string; }; export type GetConnectionTestStatusResponse = ConnectionTestStatusResponse; @@ -4531,7 +4531,7 @@ export type $OpenApiTs = { }; }; }; - '/api/v2/connections/test-async/{token}': { + '/api/v2/connections/test-async/{connection_test_token}': { get: { req: GetConnectionTestStatusData; res: { diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 69a884a55665e..62bc77af2652a 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -1234,7 +1234,7 @@ def test_post_creates_connection_test_and_callback(self, test_client, session): @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_get_status_returns_queued(self, test_client, session): - """GET /connections/test-async/{token} returns current status.""" + """GET /connections/test-async/{connection_test_token} returns current status.""" self.create_connection() post_response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) token = post_response.json()["token"] From 2f926f0c135f2d207cfcacfcb8722ad5a4fafea2 Mon Sep 17 00:00:00 2001 From: Anish Date: Fri, 27 Feb 2026 19:50:24 -0600 Subject: [PATCH 06/38] Add async connection testing on workers via dedicated workload --- .../core_api/datamodels/connections.py | 3 +- .../openapi/v2-rest-api-generated.yaml | 7 +- .../core_api/routes/public/connections.py | 15 +- .../execution_api/versions/__init__.py | 2 + .../execution_api/versions/v2026_03_31.py | 12 +- .../src/airflow/config_templates/config.yml | 35 ++++ .../src/airflow/executors/base_executor.py | 29 ++- .../src/airflow/executors/local_executor.py | 146 ++++++++------ .../airflow/executors/workloads/__init__.py | 4 +- .../executors/workloads/connection_test.py | 55 ++++++ .../src/airflow/jobs/scheduler_job_runner.py | 94 +++++++++ .../0107_3_2_0_add_connection_test_table.py | 8 +- .../src/airflow/models/connection_test.py | 82 ++------ .../ui/openapi-gen/requests/schemas.gen.ts | 13 +- .../ui/openapi-gen/requests/types.gen.ts | 3 +- .../routes/public/test_connections.py | 23 +-- .../v2026_03_31/test_connection_tests.py | 60 ++++++ .../unit/executors/test_base_executor.py | 32 +++ .../jobs/test_scheduler_connection_tests.py | 186 ++++++++++++++++++ .../tests/unit/models/test_connection_test.py | 57 +++--- .../airflowctl/api/datamodels/generated.py | 3 +- .../src/tests_common/test_utils/db.py | 10 + 22 files changed, 676 insertions(+), 203 deletions(-) create mode 100644 airflow-core/src/airflow/executors/workloads/connection_test.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py create mode 100644 airflow-core/tests/unit/jobs/test_scheduler_connection_tests.py diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index 9192926a0c77f..de6cadedbe8d2 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -80,9 +80,10 @@ class ConnectionTestResponse(BaseModel): class ConnectionTestRequestBody(StrictBaseModel): - """Request body for async connection test — just the connection_id.""" + """Request body for async connection test.""" connection_id: str + queue: str | None = None class ConnectionTestQueuedResponse(BaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 37e8994cc870f..4dbbafcbffd53 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -10218,12 +10218,17 @@ components: connection_id: type: string title: Connection Id + queue: + anyOf: + - type: string + - type: 'null' + title: Queue additionalProperties: false type: object required: - connection_id title: ConnectionTestRequestBody - description: "Request body for async connection test \u2014 just the connection_id." + description: Request body for async connection test. ConnectionTestResponse: properties: status: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 4c71a2c962372..d8c19adf6c23b 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -60,10 +60,7 @@ from airflow.configuration import conf from airflow.exceptions import AirflowNotFoundException from airflow.models import Connection -from airflow.models.connection_test import ( - ConnectionTest, - ConnectionTestState, -) +from airflow.models.connection_test import ConnectionTest 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 @@ -296,18 +293,10 @@ def test_connection_async( "Connection must be saved before testing.", ) - connection_test = ConnectionTest(connection_id=test_body.connection_id) + connection_test = ConnectionTest(connection_id=test_body.connection_id, queue=test_body.queue) session.add(connection_test) session.flush() - callback = connection_test.create_callback() - session.add(callback) - session.flush() - - connection_test.callback_id = callback.id - connection_test.state = ConnectionTestState.QUEUED - callback.queue() - return ConnectionTestQueuedResponse( token=connection_test.token, connection_id=connection_test.connection_id, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index 2cbe2e3007b3f..2ddbf81fd03d7 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -34,6 +34,7 @@ MovePreviousRunEndpoint, ) from airflow.api_fastapi.execution_api.versions.v2026_03_31 import ( + AddConnectionTestEndpoint, AddNoteField, MakeDagRunStartDateNullable, ModifyDeferredTaskKwargsToJsonValue, @@ -46,6 +47,7 @@ Version("2026-04-13", AddDagEndpoint), Version( "2026-03-31", + AddConnectionTestEndpoint, MakeDagRunStartDateNullable, ModifyDeferredTaskKwargsToJsonValue, RemoveUpstreamMapIndexesField, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py index 2d14493e81fe6..bd5f9c41e0ecf 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py @@ -19,7 +19,7 @@ from typing import Any -from cadwyn import ResponseInfo, VersionChange, convert_response_to_previous_version_for, schema +from cadwyn import ResponseInfo, VersionChange, convert_response_to_previous_version_for, endpoint, schema from airflow.api_fastapi.common.types import UtcDateTime from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( @@ -29,6 +29,16 @@ ) +class AddConnectionTestEndpoint(VersionChange): + """Add connection-tests endpoint for async connection testing.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + endpoint("/connection-tests/{connection_test_id}", ["PATCH"]).didnt_exist, + ) + + class ModifyDeferredTaskKwargsToJsonValue(VersionChange): """Change the types of `trigger_kwargs` and `next_kwargs` in TIDeferredStatePayload to JsonValue.""" diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 0c002f5276cfe..65232cb8288a6 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -481,6 +481,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 + type: integer + example: ~ + default: "4" max_templated_field_length: description: | The maximum length of the rendered template field. If the value to be stored in the @@ -2549,6 +2567,23 @@ scheduler: type: float example: ~ default: "120.0" + connection_test_dispatch_interval: + description: | + How often (in seconds) the scheduler should check for pending + connection tests and dispatch them to an executor. + version_added: 3.2.0 + type: float + example: ~ + default: "2.0" + connection_test_reaper_interval: + description: | + How often (in seconds) the scheduler should check for stale + connection tests (QUEUED or RUNNING past their timeout + grace period) + and mark them as failed. + version_added: 3.2.0 + type: float + example: ~ + default: "30.0" allowed_run_id_pattern: description: | The run_id pattern used to verify the validity of user input to the run_id parameter when diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index d67c25c7bafaa..dd25d0fdfdf3c 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -143,6 +143,7 @@ class BaseExecutor(LoggingMixin): supports_ad_hoc_ti_run: bool = False supports_callbacks: bool = False supports_multi_team: bool = False + supports_connection_test: bool = False sentry_integration: str = "" is_local: bool = False @@ -186,6 +187,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: deque[workloads.TestConnection] = deque() self.running: set[WorkloadKey] = set() self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {} self._task_event_logs: deque[Log] = deque() @@ -220,6 +222,7 @@ def log_task_event(self, *, event: str, extra: str, ti_key: TaskInstanceKey): self._task_event_logs.append(Log(event=event, task_instance=ti_key, extra=extra)) def queue_workload(self, workload: workloads.All, session: Session) -> None: +<<<<<<< HEAD if isinstance(workload, workloads.ExecuteTask): ti = workload.ti self.queued_tasks[ti.key] = workload @@ -231,10 +234,14 @@ def queue_workload(self, workload: workloads.All, session: Session) -> None: f"See LocalExecutor or CeleryExecutor for reference implementation." ) self.queued_callbacks[workload.callback.id] = workload + elif isinstance(workload, workloads.TestConnection): + if not self.supports_connection_test: + raise ValueError(f"Executor {type(self).__name__} does not support connection testing") + self.queued_connection_tests.append(workload) else: raise ValueError( f"Un-handled workload type {type(workload).__name__!r} in {type(self).__name__}. " - f"Workload must be one of: ExecuteTask, ExecuteCallback." + f"Workload must be one of: ExecuteTask, ExecuteCallback, TestConnection." ) def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, workloads.All]]: @@ -305,10 +312,30 @@ def heartbeat(self) -> None: self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads) self.trigger_tasks(open_slots) + if self.supports_connection_test and self.queued_connection_tests: + 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, max_tests: int | None = None) -> None: + """ + Process queued connection tests. + + :param max_tests: Maximum number of tests to trigger. Defaults to all queued. + """ + if not self.queued_connection_tests: + return + + count = max_tests if max_tests is not None else len(self.queued_connection_tests) + test_workloads: list[workloads.TestConnection] = [] + for _ in range(min(count, len(self.queued_connection_tests))): + test_workloads.append(self.queued_connection_tests.popleft()) + + if test_workloads: + self._process_workloads(test_workloads) + def _get_metric_name(self, metric_base_name: str) -> str: return ( f"{metric_base_name}.{self.__class__.__name__}" diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 9b5939a0bd2e7..f0731bccea218 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -35,10 +35,9 @@ import structlog -from airflow.executors import workloads +from airflow.executors import workloads as _workloads from airflow.executors.base_executor import BaseExecutor -from airflow.executors.workloads.callback import execute_callback_workload -from airflow.utils.state import CallbackState, TaskInstanceState +from airflow.utils.state import TaskInstanceState # add logger to parameter of setproctitle to support logging if sys.platform == "darwin": @@ -51,23 +50,13 @@ if TYPE_CHECKING: from structlog.typing import FilteringBoundLogger as Logger - from airflow.executors.workloads.types import WorkloadResultType - - -def _get_executor_process_title_prefix(team_name: str | None) -> str: - """ - Build the process title prefix for LocalExecutor workers. - - :param team_name: Team name from executor configuration - """ - team_suffix = f" [{team_name}]" if team_name else "" - return f"airflow worker -- LocalExecutor{team_suffix}:" + TaskInstanceStateType = tuple[_workloads.TaskInstance, TaskInstanceState, Exception | None] def _run_worker( logger_name: str, - input: SimpleQueue[workloads.All | None], - output: Queue[WorkloadResultType], + input: SimpleQueue[_workloads.All | None], + output: Queue[TaskInstanceStateType], unread_messages: multiprocessing.sharedctypes.Synchronized[int], team_conf, ): @@ -79,8 +68,11 @@ def _run_worker( log = structlog.get_logger(logger_name) log.info("Worker starting up pid=%d", os.getpid()) + # Create team suffix for process title + team_suffix = f" [{team_conf.team_name}]" if team_conf.team_name else "" + while True: - setproctitle(f"{_get_executor_process_title_prefix(team_conf.team_name)} ", log) + setproctitle(f"airflow worker -- LocalExecutor{team_suffix}: ", log) try: workload = input.get() except EOFError: @@ -95,33 +87,34 @@ def _run_worker( # Received poison pill, no more tasks to run return + if isinstance(workload, _workloads.TestConnection): + with unread_messages: + unread_messages.value -= 1 + _execute_connection_test(log, workload, team_conf) + continue + + if not isinstance(workload, _workloads.ExecuteTask): + raise ValueError(f"LocalExecutor does not know how to handle {type(workload)}") + # Decrement this as soon as we pick up a message off the queue with unread_messages: unread_messages.value -= 1 + key = None + if ti := getattr(workload, "ti", None): + key = ti.key + else: + raise TypeError(f"Don't know how to get ti key from {type(workload).__name__}") - # Handle different workload types - if isinstance(workload, workloads.ExecuteTask): - try: - _execute_work(log, workload, team_conf) - output.put((workload.ti.key, TaskInstanceState.SUCCESS, None)) - except Exception as e: - log.exception("Task execution failed.") - output.put((workload.ti.key, TaskInstanceState.FAILED, e)) - - elif isinstance(workload, workloads.ExecuteCallback): - output.put((workload.callback.id, CallbackState.RUNNING, None)) - try: - _execute_callback(log, workload, team_conf) - output.put((workload.callback.id, CallbackState.SUCCESS, None)) - except Exception as e: - log.exception("Callback execution failed") - output.put((workload.callback.id, CallbackState.FAILED, e)) + try: + _execute_work(log, workload, team_conf) - else: - raise ValueError(f"LocalExecutor does not know how to handle {type(workload)}") + output.put((key, TaskInstanceState.SUCCESS, None)) + except Exception as e: + log.exception("uhoh") + output.put((key, TaskInstanceState.FAILED, e)) -def _execute_work(log: Logger, workload: workloads.ExecuteTask, team_conf) -> None: +def _execute_work(log: Logger, workload: _workloads.ExecuteTask, team_conf) -> None: """ Execute command received and stores result state in queue. @@ -131,7 +124,9 @@ def _execute_work(log: Logger, workload: workloads.ExecuteTask, team_conf) -> No """ from airflow.sdk.execution_time.supervisor import supervise - setproctitle(f"{_get_executor_process_title_prefix(team_conf.team_name)} {workload.ti.id}", log) + # Create team suffix for process title + team_suffix = f" [{team_conf.team_name}]" if team_conf.team_name else "" + setproctitle(f"airflow worker -- LocalExecutor{team_suffix}: {workload.ti.id}", log) base_url = team_conf.get("api", "base_url", fallback="/") # If it's a relative URL, use localhost:8080 as the default @@ -152,20 +147,64 @@ def _execute_work(log: Logger, workload: workloads.ExecuteTask, team_conf) -> No ) -def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_conf) -> None: +def _execute_connection_test(log: Logger, workload: _workloads.TestConnection, team_conf) -> None: """ - Execute a callback workload. + Execute a connection test workload with a timeout. + + Results are reported back via the Execution API (workers must not access the metadata DB directly). :param log: Logger instance - :param workload: The ExecuteCallback workload to execute + :param workload: The TestConnection workload to execute :param team_conf: Team-specific executor configuration """ - setproctitle(f"{_get_executor_process_title_prefix(team_conf.team_name)} {workload.callback.id}", log) + from concurrent.futures import ThreadPoolExecutor, TimeoutError + + import httpx - success, error_msg = execute_callback_workload(workload.callback, log) + from airflow.models.connection_test import run_connection_test - if not success: - raise RuntimeError(error_msg or "Callback execution failed") + team_suffix = f" [{team_conf.team_name}]" if team_conf.team_name else "" + setproctitle( + f"airflow worker -- LocalExecutor{team_suffix}: connection-test {workload.connection_id}", log + ) + + # Build execution API URL (same pattern as _execute_work) + base_url = team_conf.get("api", "base_url", fallback="/") + if base_url.startswith("/"): + base_url = f"http://localhost:8080{base_url}" + api_url = f"{base_url.rstrip('/')}/execution/connection-tests/{workload.connection_test_id}" + headers = {"Authorization": f"Bearer {workload.token}"} + + def _patch(state: str, result_message: str | None = None) -> None: + payload: dict[str, str] = {"state": state} + if result_message is not None: + payload["result_message"] = result_message + httpx.patch(api_url, json=payload, headers=headers) + + try: + _patch("running") + + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit( + run_connection_test, + connection_id=workload.connection_id, + ) + success, message = future.result(timeout=workload.timeout) + + _patch("success" if success else "failed", message) + except TimeoutError: + log.error( + "Connection test timed out after %ds", + workload.timeout, + connection_id=workload.connection_id, + ) + _patch("failed", f"Connection test timed out after {workload.timeout}s") + except Exception: + log.exception("Connection test failed unexpectedly", connection_id=workload.connection_id) + try: + _patch("failed", "Connection test failed unexpectedly") + except Exception: + log.exception("Failed to report connection test failure") class LocalExecutor(BaseExecutor): @@ -181,11 +220,11 @@ class LocalExecutor(BaseExecutor): is_mp_using_fork: bool = multiprocessing.get_start_method() == "fork" supports_multi_team: bool = True + supports_connection_test: bool = True serve_logs: bool = True - supports_callbacks: bool = True - activity_queue: SimpleQueue[workloads.All | None] - result_queue: SimpleQueue[WorkloadResultType] + activity_queue: SimpleQueue[_workloads.All | None] + result_queue: SimpleQueue[TaskInstanceStateType] workers: dict[int, multiprocessing.Process] _unread_messages: multiprocessing.sharedctypes.Synchronized[int] @@ -328,14 +367,11 @@ def end(self) -> None: def terminate(self): """Terminate the executor is not doing anything.""" - def _process_workloads(self, workload_list): - for workload in workload_list: + def _process_workloads(self, workloads): + for workload in workloads: self.activity_queue.put(workload) - # Remove from appropriate queue based on workload type - if isinstance(workload, workloads.ExecuteTask): + if isinstance(workload, _workloads.ExecuteTask): del self.queued_tasks[workload.ti.key] - elif isinstance(workload, workloads.ExecuteCallback): - del self.queued_callbacks[workload.callback.id] with self._unread_messages: - self._unread_messages.value += len(workload_list) + self._unread_messages.value += len(workloads) self._check_workers() diff --git a/airflow-core/src/airflow/executors/workloads/__init__.py b/airflow-core/src/airflow/executors/workloads/__init__.py index 462e38ad0aaac..136ab37734cbb 100644 --- a/airflow-core/src/airflow/executors/workloads/__init__.py +++ b/airflow-core/src/airflow/executors/workloads/__init__.py @@ -24,11 +24,12 @@ from airflow.executors.workloads.base import BaseWorkload, BundleInfo from airflow.executors.workloads.callback import CallbackFetchMethod, ExecuteCallback +from airflow.executors.workloads.connection_test import TestConnection from airflow.executors.workloads.task import ExecuteTask, TaskInstanceDTO from airflow.executors.workloads.trigger import RunTrigger All = Annotated[ - ExecuteTask | ExecuteCallback | RunTrigger, + ExecuteTask | ExecuteCallback | RunTrigger | TestConnection, Field(discriminator="type"), ] @@ -43,4 +44,5 @@ "ExecuteTask", "TaskInstance", "TaskInstanceDTO", + "TestConnection", ] diff --git a/airflow-core/src/airflow/executors/workloads/connection_test.py b/airflow-core/src/airflow/executors/workloads/connection_test.py new file mode 100644 index 0000000000000..356d8ac212f6e --- /dev/null +++ b/airflow-core/src/airflow/executors/workloads/connection_test.py @@ -0,0 +1,55 @@ +# 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 + + +class TestConnection(BaseWorkloadSchema): + """Execute a connection test on a worker.""" + + connection_test_id: uuid.UUID + connection_id: str + timeout: int = 60 + + type: Literal["TestConnection"] = Field(init=False, default="TestConnection") + + @classmethod + def make( + cls, + *, + connection_test_id: uuid.UUID, + connection_id: str, + timeout: int = 60, + generator: JWTGenerator | None = None, + ) -> TestConnection: + return cls( + connection_test_id=connection_test_id, + connection_id=connection_id, + timeout=timeout, + token=cls.generate_token(str(connection_test_id), generator), + ) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index c64bd166f1b3a..a5357cebfbaa4 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -73,6 +73,7 @@ ) from airflow.models.backfill import Backfill from airflow.models.callback import Callback, CallbackType, ExecutorCallback +from airflow.models.connection_test import ConnectionTest, ConnectionTestState from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion from airflow.models.dagbag import DBDagBag @@ -1585,6 +1586,16 @@ def _run_scheduler_loop(self) -> None: action=bundle_cleanup_mgr.remove_stale_bundle_versions, ) + if any(x.supports_connection_test for x in self.executors): + timers.call_regular_interval( + delay=conf.getfloat("scheduler", "connection_test_dispatch_interval", fallback=2.0), + action=self._dispatch_connection_tests, + ) + timers.call_regular_interval( + delay=conf.getfloat("scheduler", "connection_test_reaper_interval", fallback=30.0), + action=self._reap_stale_connection_tests, + ) + idle_count = 0 for loop_count in itertools.count(start=1): @@ -3114,6 +3125,89 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]: session.add(warning) existing_warned_dag_ids.add(warning.dag_id) + @provide_session + def _dispatch_connection_tests(self, session: Session = NEW_SESSION) -> None: + """Dispatch pending connection tests to executors that support them.""" + max_concurrency = conf.getint("core", "max_connection_test_concurrency", fallback=4) + timeout = conf.getint("core", "connection_test_timeout", fallback=60) + + active_count = session.scalar( + select(func.count(ConnectionTest.id)).where( + ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]) + ) + ) + budget = max_concurrency - (active_count or 0) + if budget <= 0: + return + + pending_stmt = ( + select(ConnectionTest) + .where(ConnectionTest.state == ConnectionTestState.PENDING) + .order_by(ConnectionTest.created_at) + .limit(budget) + ) + pending_stmt = with_row_locks(pending_stmt, session, of=ConnectionTest, skip_locked=True) + pending_tests = session.scalars(pending_stmt).all() + + if not pending_tests: + return + + for ct in pending_tests: + executor = self._find_executor_for_connection_test(ct.queue) + if executor is None: + reason = ( + f"No executor serves queue '{ct.queue}'" + if ct.queue + else "No executor supports connection testing" + ) + ct.state = ConnectionTestState.FAILED + ct.result_message = reason + self.log.warning("Failing connection test %s: %s", ct.id, reason) + continue + + workload = workloads.TestConnection.make( + connection_test_id=ct.id, + connection_id=ct.connection_id, + timeout=timeout, + generator=executor.jwt_generator, + ) + executor.queue_workload(workload, session=session) + ct.state = ConnectionTestState.QUEUED + + session.flush() + + @provide_session + def _reap_stale_connection_tests(self, session: Session = NEW_SESSION) -> None: + """Mark connection tests that have exceeded their timeout as FAILED.""" + timeout = conf.getint("core", "connection_test_timeout", fallback=60) + grace_period = 30 + cutoff = timezone.utcnow() - timedelta(seconds=timeout + grace_period) + + stale_tests = session.scalars( + select(ConnectionTest).where( + ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]), + ConnectionTest.updated_at < cutoff, + ) + ).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) + + session.flush() + + def _find_executor_for_connection_test(self, queue: str | None) -> BaseExecutor | None: + """Find an executor that supports connection testing, optionally matching a queue.""" + if queue is not None: + for executor in self.executors: + if executor.supports_connection_test and executor.team_name == queue: + return executor + for executor in self.executors: + if executor.supports_connection_test and (queue is None or executor.team_name is None): + return executor + return None + def _executor_to_workloads( self, workloads: Iterable[SchedulerWorkload], diff --git a/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py index 02c7ebbf2f4c0..3ca39f7fc96b4 100644 --- a/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py @@ -51,14 +51,8 @@ def upgrade(): 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("callback_id", sa.Uuid(), nullable=True), + sa.Column("queue", sa.String(256), nullable=True), sa.PrimaryKeyConstraint("id", name=op.f("connection_test_pkey")), - sa.ForeignKeyConstraint( - ["callback_id"], - ["callback.id"], - name=op.f("connection_test_callback_id_fkey"), - ondelete="SET NULL", - ), sa.UniqueConstraint("token", name=op.f("connection_test_token_uq")), ) op.create_index( diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 3c9c424006d20..6092a91511dd0 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -19,21 +19,17 @@ import secrets from datetime import datetime from enum import Enum -from typing import TYPE_CHECKING from uuid import UUID import structlog import uuid6 -from sqlalchemy import ForeignKey, Index, String, Text, Uuid -from sqlalchemy.orm import Mapped, mapped_column, relationship +from sqlalchemy import Index, String, Text, Uuid +from sqlalchemy.orm import Mapped, mapped_column from airflow._shared.timezones import timezone from airflow.models.base import Base from airflow.utils.sqlalchemy import UtcDateTime -if TYPE_CHECKING: - from airflow.models.callback import Callback, ExecutorCallback - log = structlog.get_logger(__name__) @@ -52,31 +48,9 @@ def __str__(self) -> str: TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED)) -# Path used by ExecutorCallback to locate the worker function. -RUN_CONNECTION_TEST_PATH = "airflow.models.connection_test.run_connection_test" - - -class _ImportPathCallbackDef: - """ - Minimal implementation of ImportPathExecutorCallbackDefProtocol. - - ExecutorCallback.__init__ expects an object satisfying this protocol, but no concrete - implementation exists in airflow-core — the only one (SyncCallback) lives in the task-sdk. - Once #61153 lands and ExecuteCallback.make() is decoupled from DagRun, this adapter can - be replaced with the proper factory method. - """ - - def __init__(self, path: str, kwargs: dict, executor: str | None = None): - self.path = path - self.kwargs = kwargs - self.executor = executor - - def serialize(self) -> dict: - return {"path": self.path, "kwargs": self.kwargs, "executor": self.executor} - class ConnectionTest(Base): - """Tracks an async connection test dispatched to a worker via ExecutorCallback.""" + """Tracks an async connection test dispatched to a worker via a TestConnection workload.""" __tablename__ = "connection_test" @@ -89,17 +63,14 @@ class ConnectionTest(Base): updated_at: Mapped[datetime] = mapped_column( UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False ) - - callback_id: Mapped[UUID | None] = mapped_column( - Uuid(), ForeignKey("callback.id", ondelete="SET NULL"), nullable=True - ) - callback: Mapped[Callback | None] = relationship("Callback", uselist=False) + queue: Mapped[str | None] = mapped_column(String(256), nullable=True) __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),) - def __init__(self, *, connection_id: str, **kwargs): + def __init__(self, *, connection_id: str, queue: str | None = None, **kwargs): super().__init__(**kwargs) self.connection_id = connection_id + self.queue = queue self.token = secrets.token_urlsafe(32) self.state = ConnectionTestState.PENDING @@ -109,48 +80,19 @@ def __repr__(self) -> str: f" connection_id={self.connection_id!r} state={self.state}>" ) - def create_callback(self) -> ExecutorCallback: - """Create an ExecutorCallback that will run the connection test on a worker.""" - from airflow.models.callback import CallbackFetchMethod, ExecutorCallback - - callback_def = _ImportPathCallbackDef( - path=RUN_CONNECTION_TEST_PATH, - kwargs={ - "connection_id": self.connection_id, - "connection_test_id": str(self.id), - }, - ) - return ExecutorCallback(callback_def, fetch_method=CallbackFetchMethod.IMPORT_PATH) - -def run_connection_test(*, connection_id: str, connection_test_id: str) -> None: +def run_connection_test(*, connection_id: str) -> tuple[bool, str]: """ - Worker-side function to execute a connection test. + Worker-side pure function to execute a connection test. - This is the function referenced by the ExecutorCallback's import path. - It fetches the connection, runs test_connection(), and reports results - back by updating the ConnectionTest row directly. + Returns a (success, message) tuple. The caller is responsible for + reporting the result back via the Execution API. """ from airflow.models.connection import Connection - from airflow.utils.session import create_session - - connection_test_uuid = UUID(connection_test_id) - - with create_session() as session: - ct = session.get(ConnectionTest, connection_test_uuid) - if ct: - ct.state = ConnectionTestState.RUNNING try: conn = Connection.get_connection_from_secrets(connection_id) - test_status, test_message = conn.test_connection() + return conn.test_connection() except Exception as e: - test_status = False - test_message = str(e) log.exception("Connection test failed", connection_id=connection_id) - - with create_session() as session: - ct = session.get(ConnectionTest, connection_test_uuid) - if ct: - ct.result_message = test_message - ct.state = ConnectionTestState.SUCCESS if test_status else ConnectionTestState.FAILED + return False, str(e) diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 5c0987aef17ec..65da80057b6fe 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1724,13 +1724,24 @@ export const $ConnectionTestRequestBody = { connection_id: { type: 'string', title: 'Connection Id' + }, + queue: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Queue' } }, additionalProperties: false, type: 'object', required: ['connection_id'], title: 'ConnectionTestRequestBody', - description: 'Request body for async connection test — just the connection_id.' + description: 'Request body for async connection test.' } as const; export const $ConnectionTestResponse = { diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 80981e883a59e..388a3a69beae4 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -505,10 +505,11 @@ export type ConnectionTestQueuedResponse = { }; /** - * Request body for async connection test — just the connection_id. + * Request body for async connection test. */ export type ConnectionTestRequestBody = { connection_id: string; + queue?: string | null; }; /** diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 62bc77af2652a..6f5c6472d311f 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -29,8 +29,7 @@ from airflow.api_fastapi.core_api.datamodels.connections import ConnectionBody from airflow.api_fastapi.core_api.services.public.connections import BulkConnectionService from airflow.models import Connection -from airflow.models.callback import Callback -from airflow.models.connection_test import RUN_CONNECTION_TEST_PATH, ConnectionTest, ConnectionTestState +from airflow.models.connection_test import ConnectionTest, ConnectionTestState from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.utils.session import NEW_SESSION, provide_session @@ -1182,7 +1181,7 @@ def test_post_should_respond_202(self, test_client, session): body = response.json() assert "token" in body assert body["connection_id"] == TEST_CONN_ID - assert body["state"] == "queued" + assert body["state"] == "pending" assert len(body["token"]) > 0 def test_should_respond_401(self, unauthenticated_test_client): @@ -1214,8 +1213,8 @@ def test_should_respond_404_for_nonexistent_connection(self, test_client): assert "was not found" in response.json()["detail"] @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_post_creates_connection_test_and_callback(self, test_client, session): - """POST creates both a ConnectionTest row and an ExecutorCallback.""" + def test_post_creates_connection_test_row(self, test_client, session): + """POST creates a ConnectionTest row in PENDING state.""" self.create_connection() response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) assert response.status_code == 202 @@ -1224,17 +1223,11 @@ def test_post_creates_connection_test_and_callback(self, test_client, session): ct = session.scalar(select(ConnectionTest).filter_by(token=token)) assert ct is not None assert ct.connection_id == TEST_CONN_ID - assert ct.state == "queued" - assert ct.callback_id is not None - - cb = session.get(Callback, ct.callback_id) - assert cb is not None - assert cb.data["path"] == RUN_CONNECTION_TEST_PATH - assert cb.data["kwargs"]["connection_id"] == TEST_CONN_ID + assert ct.state == "pending" @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_get_status_returns_queued(self, test_client, session): - """GET /connections/test-async/{connection_test_token} returns current status.""" + def test_get_status_returns_pending(self, test_client, session): + """GET /connections/test-async/{token} returns current status (pending before scheduler dispatch).""" self.create_connection() post_response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) token = post_response.json()["token"] @@ -1244,7 +1237,7 @@ def test_get_status_returns_queued(self, test_client, session): body = response.json() assert body["token"] == token assert body["connection_id"] == TEST_CONN_ID - assert body["state"] == "queued" + assert body["state"] == "pending" assert body["result_message"] is None assert "created_at" in body diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py new file mode 100644 index 0000000000000..d5203253c3081 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py @@ -0,0 +1,60 @@ +# 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 pytest + +from airflow.models.connection_test import ConnectionTest, ConnectionTestState + +pytestmark = pytest.mark.db_test + + +@pytest.fixture +def old_ver_client(client): + """Client configured to use API version before connection-tests endpoint was added.""" + client.headers["Airflow-API-Version"] = "2025-12-08" + return client + + +class TestConnectionTestEndpointVersioning: + """Test that the connection-tests endpoint didn't exist in older API versions.""" + + def test_old_version_returns_404(self, old_ver_client, session): + """PATCH /connection-tests/{id} should not exist in older API versions.""" + ct = ConnectionTest(connection_id="test_conn") + ct.state = ConnectionTestState.RUNNING + session.add(ct) + session.commit() + + response = old_ver_client.patch( + f"/execution/connection-tests/{ct.id}", + json={"state": "success", "result_message": "ok"}, + ) + assert response.status_code == 404 + + def test_head_version_works(self, client, session): + """PATCH /connection-tests/{id} should work in the current API version.""" + ct = ConnectionTest(connection_id="test_conn") + ct.state = ConnectionTestState.RUNNING + session.add(ct) + session.commit() + + response = client.patch( + f"/execution/connection-tests/{ct.id}", + json={"state": "success", "result_message": "ok"}, + ) + assert response.status_code == 204 diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index fa0f311d018fe..76cfc8954c057 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -407,6 +407,38 @@ def test_repr(): assert repr(executor) == "BaseExecutor(parallelism=10, team_name='teamA')" +def test_supports_connection_test_default_value(): + assert not BaseExecutor.supports_connection_test + + +def test_queue_connection_test_workload_rejected_by_default(): + """BaseExecutor (supports_connection_test=False) rejects TestConnection workloads.""" + import uuid + + executor = BaseExecutor() + wl = workloads.TestConnection.make( + connection_test_id=uuid.uuid4(), + connection_id="test_conn", + ) + with pytest.raises(ValueError, match="does not support connection testing"): + executor.queue_workload(wl, session=mock.MagicMock()) + + +def test_queue_connection_test_workload_accepted_when_supported(): + """An executor with supports_connection_test=True accepts TestConnection workloads.""" + import uuid + + executor = LocalExecutor() + executor.queued_connection_tests.clear() + wl = workloads.TestConnection.make( + connection_test_id=uuid.uuid4(), + connection_id="test_conn", + ) + executor.queue_workload(wl, session=mock.MagicMock()) + assert len(executor.queued_connection_tests) == 1 + assert executor.queued_connection_tests[0] is wl + + @mock.patch.dict("os.environ", {}, clear=True) class TestExecutorConf: """Test ExecutorConf shim class that provides team-specific configuration access.""" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_connection_tests.py b/airflow-core/tests/unit/jobs/test_scheduler_connection_tests.py new file mode 100644 index 0000000000000..b82e78d972cb9 --- /dev/null +++ b/airflow-core/tests/unit/jobs/test_scheduler_connection_tests.py @@ -0,0 +1,186 @@ +# 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. +"""Tests for scheduler connection test dispatch and reaper.""" + +from __future__ import annotations + +import logging +import os +from datetime import timedelta +from unittest import mock + +import pytest +import time_machine +from sqlalchemy import delete + +from airflow._shared.timezones import timezone +from airflow.executors.base_executor import BaseExecutor +from airflow.executors.local_executor import LocalExecutor +from airflow.jobs.job import Job +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner +from airflow.models.connection_test import ConnectionTest, ConnectionTestState + +pytestmark = pytest.mark.db_test + + +@pytest.fixture +def scheduler_job_runner(session): + """Create a SchedulerJobRunner with a mock Job and supporting executor.""" + session.execute(delete(ConnectionTest)) + session.commit() + + mock_job = mock.MagicMock(spec=Job) + mock_job.id = 1 + mock_job.max_tis_per_query = 16 + executor = LocalExecutor() + executor.queued_connection_tests.clear() + runner = SchedulerJobRunner.__new__(SchedulerJobRunner) + runner.job = mock_job + runner.executors = [executor] + runner.executor = executor + runner._log = mock.MagicMock(spec=logging.Logger) + return runner + + +class TestDispatchConnectionTests: + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_pending_tests(self, scheduler_job_runner, session): + """Pending connection tests are dispatched to a supporting executor.""" + ct = ConnectionTest(connection_id="test_conn") + session.add(ct) + session.commit() + assert ct.state == ConnectionTestState.PENDING + + scheduler_job_runner._dispatch_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.QUEUED + assert len(scheduler_job_runner.executor.queued_connection_tests) == 1 + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "1", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_respects_concurrency_limit(self, scheduler_job_runner, session): + """Excess pending tests stay PENDING when concurrency is at capacity.""" + # Create one already-queued test to fill the concurrency budget (limit=1) + ct_active = ConnectionTest(connection_id="active_conn") + ct_active.state = ConnectionTestState.QUEUED + session.add(ct_active) + + ct_pending = ConnectionTest(connection_id="pending_conn") + session.add(ct_pending) + session.commit() + + scheduler_job_runner._dispatch_connection_tests(session=session) + + session.expire_all() + ct_pending = session.get(ConnectionTest, ct_pending.id) + assert ct_pending.state == ConnectionTestState.PENDING + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_fails_fast_when_no_executor_supports(self, scheduler_job_runner, session): + """Tests fail immediately when no executor supports connection testing.""" + # Replace executor with one that doesn't support connection tests + unsupporting_executor = BaseExecutor() + unsupporting_executor.supports_connection_test = False + scheduler_job_runner.executors = [unsupporting_executor] + scheduler_job_runner.executor = unsupporting_executor + + ct = ConnectionTest(connection_id="test_conn") + session.add(ct) + session.commit() + + scheduler_job_runner._dispatch_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert "No executor supports connection testing" in ct.result_message + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_with_queue_falls_back_to_global_executor(self, scheduler_job_runner, session): + """Tests with a queue are dispatched to the global executor as fallback.""" + ct = ConnectionTest(connection_id="test_conn", queue="gpu_workers") + session.add(ct) + session.commit() + + # Default LocalExecutor has team_name=None — serves as fallback for any queue + scheduler_job_runner._dispatch_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.QUEUED + assert len(scheduler_job_runner.executor.queued_connection_tests) == 1 + + +class TestReapStaleConnectionTests: + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) + def test_reap_stale_queued_test(self, scheduler_job_runner, session): + """Stale QUEUED tests are marked as FAILED by the reaper.""" + initial_time = timezone.utcnow() + + with time_machine.travel(initial_time, tick=False): + ct = ConnectionTest(connection_id="test_conn") + ct.state = ConnectionTestState.QUEUED + session.add(ct) + session.commit() + + # Jump forward past timeout (60s) + grace period (30s) + with time_machine.travel(initial_time + timedelta(seconds=200), tick=False): + scheduler_job_runner._reap_stale_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert "timed out" in ct.result_message + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) + def test_does_not_reap_fresh_tests(self, scheduler_job_runner, session): + """Fresh QUEUED tests are not reaped.""" + ct = ConnectionTest(connection_id="test_conn") + ct.state = ConnectionTestState.QUEUED + session.add(ct) + session.commit() + + scheduler_job_runner._reap_stale_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.QUEUED diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py index 34d998ef22c98..435c2d1a4c207 100644 --- a/airflow-core/tests/unit/models/test_connection_test.py +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -47,59 +47,46 @@ def test_repr(self): assert "test_conn" in r assert "pending" in r + def test_queue_parameter(self): + ct = ConnectionTest(connection_id="test_conn", queue="my_queue") + assert ct.queue == "my_queue" -class TestRunConnectionTest: - def test_successful_connection_test(self, session): - """Worker function updates state to SUCCESS on successful test.""" + def test_queue_defaults_to_none(self): ct = ConnectionTest(connection_id="test_conn") - session.add(ct) - session.commit() - ct_id = str(ct.id) + assert ct.queue is None + +class TestRunConnectionTest: + def test_successful_connection_test(self): + """Pure function returns (True, message) on successful test.""" with mock.patch.object( Connection, "get_connection_from_secrets", return_value=mock.MagicMock() ) as mock_get_conn: mock_get_conn.return_value.test_connection.return_value = (True, "Connection OK") - run_connection_test(connection_id="test_conn", connection_test_id=ct_id) + success, message = run_connection_test(connection_id="test_conn") - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.SUCCESS - assert ct.result_message == "Connection OK" - - def test_failed_connection_test(self, session): - """Worker function updates state to FAILED when test_connection returns False.""" - ct = ConnectionTest(connection_id="test_conn") - session.add(ct) - session.commit() - ct_id = str(ct.id) + assert success is True + assert message == "Connection OK" + def test_failed_connection_test(self): + """Pure function returns (False, message) when test_connection returns False.""" with mock.patch.object( Connection, "get_connection_from_secrets", return_value=mock.MagicMock() ) as mock_get_conn: mock_get_conn.return_value.test_connection.return_value = (False, "Connection failed") - run_connection_test(connection_id="test_conn", connection_test_id=ct_id) - - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.FAILED - assert ct.result_message == "Connection failed" + success, message = run_connection_test(connection_id="test_conn") - def test_exception_during_connection_test(self, session): - """Worker function handles exceptions gracefully.""" - ct = ConnectionTest(connection_id="test_conn") - session.add(ct) - session.commit() - ct_id = str(ct.id) + assert success is False + assert message == "Connection failed" + def test_exception_during_connection_test(self): + """Pure function returns (False, error_str) on exception.""" with mock.patch.object( Connection, "get_connection_from_secrets", side_effect=Exception("Could not resolve host: db.example.com"), ): - run_connection_test(connection_id="test_conn", connection_test_id=ct_id) + success, message = run_connection_test(connection_id="test_conn") - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.FAILED - assert "Could not resolve host" in ct.result_message + assert success is False + assert "Could not resolve host" in message diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index 0734ca8a271b6..16fb15221f94c 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -256,13 +256,14 @@ class ConnectionTestQueuedResponse(BaseModel): class ConnectionTestRequestBody(BaseModel): """ - Request body for async connection test — just the connection_id. + Request body for async connection test. """ model_config = ConfigDict( extra="forbid", ) connection_id: Annotated[str, Field(title="Connection Id")] + queue: Annotated[str | None, Field(title="Queue")] = None class ConnectionTestResponse(BaseModel): diff --git a/devel-common/src/tests_common/test_utils/db.py b/devel-common/src/tests_common/test_utils/db.py index cbfb0b377ae71..76c453b9952d1 100644 --- a/devel-common/src/tests_common/test_utils/db.py +++ b/devel-common/src/tests_common/test_utils/db.py @@ -470,6 +470,14 @@ def clear_db_teams(): session.execute(delete(Team)) +def clear_db_connection_tests(): + with create_session() as session: + if AIRFLOW_V_3_2_PLUS: + from airflow.models.connection_test import ConnectionTest + + session.execute(delete(ConnectionTest)) + + @_retry_db def clear_db_revoked_tokens(): with create_session() as session: @@ -1001,3 +1009,5 @@ def clear_all(): clear_db_backfills() clear_db_dag_bundles() clear_db_dag_parsing_requests() + if AIRFLOW_V_3_2_PLUS: + clear_db_connection_tests() From 5ce672ee23146b94b1ee7ead2e4f374757e6cdb9 Mon Sep 17 00:00:00 2001 From: Anish Date: Sat, 28 Feb 2026 12:31:50 -0600 Subject: [PATCH 07/38] Add revert-on-failure for save-and-test connection endpoint --- .../core_api/datamodels/connections.py | 9 + .../openapi/v2-rest-api-generated.yaml | 98 ++++++++ .../core_api/routes/public/connections.py | 66 +++++- .../execution_api/routes/connection_tests.py | 10 +- .../src/airflow/jobs/scheduler_job_runner.py | 4 +- .../0108_3_2_0_add_connection_snapshot.py | 52 +++++ .../src/airflow/models/connection_test.py | 117 +++++++++- .../airflow/ui/openapi-gen/queries/common.ts | 1 + .../airflow/ui/openapi-gen/queries/queries.ts | 23 ++ .../ui/openapi-gen/requests/schemas.gen.ts | 25 +++ .../ui/openapi-gen/requests/services.gen.ts | 38 +++- .../ui/openapi-gen/requests/types.gen.ts | 49 ++++ .../routes/public/test_connections.py | 122 ++++++++++ .../versions/head/test_connection_tests.py | 140 +++++++++++- .../jobs/test_scheduler_connection_tests.py | 186 ---------------- .../tests/unit/jobs/test_scheduler_job.py | 210 ++++++++++++++++++ .../tests/unit/models/test_connection_test.py | 189 +++++++++++++++- .../airflowctl/api/datamodels/generated.py | 11 + 18 files changed, 1155 insertions(+), 195 deletions(-) create mode 100644 airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_snapshot.py delete mode 100644 airflow-core/tests/unit/jobs/test_scheduler_connection_tests.py diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index de6cadedbe8d2..0f4ca3fe73673 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -102,6 +102,15 @@ class ConnectionTestStatusResponse(BaseModel): state: str result_message: str | None = None created_at: datetime + reverted: bool = False + + +class ConnectionSaveAndTestResponse(BaseModel): + """Response returned by the combined save-and-test endpoint.""" + + connection: ConnectionResponse + test_token: str + test_state: str class ConnectionHookFieldBehavior(BaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 4dbbafcbffd53..049fc3854bf7d 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -1631,6 +1631,83 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + /api/v2/connections/{connection_id}/save-and-test: + patch: + tags: + - Connection + summary: Patch Connection And Test + description: 'Update a connection and queue an async test with revert-on-failure. + + + Atomically saves the edit and creates a ConnectionTest with snapshots of the + + pre-edit and post-edit state. If the test fails, the connection is automatically + + reverted to its pre-edit values.' + operationId: patch_connection_and_test + security: + - OAuth2PasswordBearer: [] + - HTTPBearer: [] + parameters: + - name: connection_id + in: path + required: true + schema: + type: string + title: Connection Id + - name: update_mask + in: query + required: false + schema: + anyOf: + - type: array + items: + type: string + - type: 'null' + title: Update Mask + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionBody' + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionSaveAndTestResponse' + '401': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Unauthorized + '403': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Forbidden + '400': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Bad Request + '404': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Not Found + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' /api/v2/connections/test: post: tags: @@ -10195,6 +10272,23 @@ components: - team_name title: ConnectionResponse description: Connection serializer for responses. + ConnectionSaveAndTestResponse: + properties: + connection: + $ref: '#/components/schemas/ConnectionResponse' + test_token: + type: string + title: Test Token + test_state: + type: string + title: Test State + type: object + required: + - connection + - test_token + - test_state + title: ConnectionSaveAndTestResponse + description: Response returned by the combined save-and-test endpoint. ConnectionTestQueuedResponse: properties: token: @@ -10263,6 +10357,10 @@ components: type: string format: date-time title: Created At + reverted: + type: boolean + title: Reverted + default: false type: object required: - token diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index d8c19adf6c23b..adad474898f09 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -41,6 +41,7 @@ ConnectionBodyPartial, ConnectionCollectionResponse, ConnectionResponse, + ConnectionSaveAndTestResponse, ConnectionTestQueuedResponse, ConnectionTestRequestBody, ConnectionTestResponse, @@ -60,7 +61,7 @@ from airflow.configuration import conf from airflow.exceptions import AirflowNotFoundException from airflow.models import Connection -from airflow.models.connection_test import ConnectionTest +from airflow.models.connection_test import ConnectionTest, snapshot_connection 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 @@ -234,6 +235,68 @@ def patch_connection( return connection +@connections_router.patch( + "/{connection_id}/save-and-test", + responses=create_openapi_http_exception_doc( + [ + status.HTTP_400_BAD_REQUEST, + status.HTTP_403_FORBIDDEN, + status.HTTP_404_NOT_FOUND, + ] + ), + dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())], +) +def patch_connection_and_test( + connection_id: str, + patch_body: ConnectionBody, + session: SessionDep, + update_mask: list[str] | None = Query(None), +) -> ConnectionSaveAndTestResponse: + """ + Update a connection and queue an async test with revert-on-failure. + + Atomically saves the edit and creates a ConnectionTest with snapshots of the + pre-edit and post-edit state. If the test fails, the connection is automatically + reverted to its pre-edit values. + """ + _ensure_test_connection_enabled() + + if patch_body.connection_id != connection_id: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, + "The connection_id in the request body does not match the URL parameter", + ) + + connection = session.scalar(select(Connection).filter_by(conn_id=connection_id).limit(1)) + if connection is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"The Connection with connection_id: `{connection_id}` was not found", + ) + + try: + ConnectionBody(**patch_body.model_dump()) + except ValidationError as e: + raise RequestValidationError(errors=e.errors()) + + pre_snapshot = snapshot_connection(connection) + + update_orm_from_pydantic(connection, patch_body, update_mask) + + post_snapshot = snapshot_connection(connection) + + connection_test = ConnectionTest(connection_id=connection_id) + connection_test.connection_snapshot = {"pre": pre_snapshot, "post": post_snapshot} + session.add(connection_test) + session.flush() + + return ConnectionSaveAndTestResponse( + connection=connection, + test_token=connection_test.token, + test_state=connection_test.state, + ) + + @connections_router.post("/test", dependencies=[Depends(requires_access_connection(method="POST"))]) def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse: """ @@ -332,6 +395,7 @@ def get_connection_test_status( state=connection_test.state, result_message=connection_test.result_message, created_at=connection_test.created_at, + reverted=connection_test.reverted, ) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py index c61a7b3f4ff3e..1ad2e0192c1bd 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -22,7 +22,12 @@ from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.execution_api.datamodels.connection_test import ConnectionTestResultBody -from airflow.models.connection_test import TERMINAL_STATES, ConnectionTest +from airflow.models.connection_test import ( + TERMINAL_STATES, + ConnectionTest, + ConnectionTestState, + attempt_revert, +) router = APIRouter() @@ -62,3 +67,6 @@ def patch_connection_test( ct.state = body.state ct.result_message = body.result_message + + if body.state == ConnectionTestState.FAILED and ct.connection_snapshot: + attempt_revert(ct, session=session) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index a5357cebfbaa4..06cbbf34ce36f 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -73,7 +73,7 @@ ) from airflow.models.backfill import Backfill from airflow.models.callback import Callback, CallbackType, ExecutorCallback -from airflow.models.connection_test import ConnectionTest, ConnectionTestState +from airflow.models.connection_test import ConnectionTest, ConnectionTestState, attempt_revert from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion from airflow.models.dagbag import DBDagBag @@ -3194,6 +3194,8 @@ def _reap_stale_connection_tests(self, session: Session = NEW_SESSION) -> None: ct.state = ConnectionTestState.FAILED ct.result_message = f"Connection test timed out (exceeded {timeout}s + {grace_period}s grace)" self.log.warning("Reaped stale connection test %s", ct.id) + if ct.connection_snapshot: + attempt_revert(ct, session=session) session.flush() diff --git a/airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_snapshot.py b/airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_snapshot.py new file mode 100644 index 0000000000000..2e4f0f45919c2 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_snapshot.py @@ -0,0 +1,52 @@ +# +# 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_snapshot and reverted columns to connection_test. + +Revision ID: b8f3e5d1a9c2 +Revises: a7e6d4c3b2f1 +Create Date: 2026-02-27 00:00:00.000000 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "b8f3e5d1a9c2" +down_revision = "a7e6d4c3b2f1" +branch_labels = None +depends_on = None +airflow_version = "3.2.0" + + +def upgrade(): + """Add connection_snapshot and reverted columns to connection_test.""" + with op.batch_alter_table("connection_test") as batch_op: + batch_op.add_column(sa.Column("connection_snapshot", sa.JSON(), nullable=True)) + batch_op.add_column(sa.Column("reverted", sa.Boolean(), nullable=False, server_default=sa.false())) + + +def downgrade(): + """Remove connection_snapshot and reverted columns from connection_test.""" + with op.batch_alter_table("connection_test") as batch_op: + batch_op.drop_column("reverted") + batch_op.drop_column("connection_snapshot") diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 6092a91511dd0..679da4fadc7bd 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -19,17 +19,23 @@ 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 Index, String, Text, Uuid +from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select from sqlalchemy.orm import Mapped, mapped_column from airflow._shared.timezones import timezone from airflow.models.base import Base +from airflow.models.connection import Connection +from airflow.models.crypto import get_fernet from airflow.utils.sqlalchemy import UtcDateTime +if TYPE_CHECKING: + from sqlalchemy.orm import Session + log = structlog.get_logger(__name__) @@ -64,6 +70,8 @@ class ConnectionTest(Base): UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False ) queue: Mapped[str | None] = mapped_column(String(256), nullable=True) + connection_snapshot: Mapped[dict | None] = mapped_column(JSON, nullable=True) + reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False) __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),) @@ -88,11 +96,114 @@ def run_connection_test(*, connection_id: str) -> tuple[bool, str]: Returns a (success, message) tuple. The caller is responsible for reporting the result back via the Execution API. """ - from airflow.models.connection import Connection - try: conn = Connection.get_connection_from_secrets(connection_id) return conn.test_connection() except Exception as e: log.exception("Connection test failed", connection_id=connection_id) return False, str(e) + + +_SNAPSHOT_FIELDS = ( + "conn_type", + "description", + "host", + "login", + "_password", + "schema", + "port", + "_extra", + "is_encrypted", + "is_extra_encrypted", + "team_name", +) + + +def snapshot_connection(conn: Connection) -> dict: + """ + Capture raw DB column values from a Connection for later restore. + + Encrypted fields (``_password``, ``_extra``) are stored as ciphertext + so they can be written directly back without re-encryption. + """ + return {field: getattr(conn, field) for field in _SNAPSHOT_FIELDS} + + +def _revert_connection(conn: Connection, snapshot: dict) -> None: + """ + Restore a Connection's columns from a snapshot dict. + + Writes directly to ``_password`` and ``_extra`` (bypassing the + encrypting property setters) so the stored ciphertext is preserved. + """ + for field, value in snapshot.items(): + setattr(conn, field, value) + + +def _decrypt_snapshot_field(snapshot: dict, field: str) -> str | None: + """Decrypt a single encrypted field from a snapshot dict using Fernet.""" + raw = snapshot.get(field) + if raw is None: + return None + encrypted_flag = "is_encrypted" if field == "_password" else "is_extra_encrypted" + if not snapshot.get(encrypted_flag, False): + return raw + fernet = get_fernet() + return fernet.decrypt(bytes(raw, "utf-8")).decode() + + +def _can_safely_revert(conn: Connection, post_snapshot: dict) -> bool: + """ + Check whether the connection's current state matches the post-edit snapshot. + + Compares **decrypted** values for encrypted fields and direct values for + non-encrypted fields. Returns ``False`` if any field differs, indicating + a concurrent edit has occurred and the revert should be skipped. + """ + for field in _SNAPSHOT_FIELDS: + if field in ("is_encrypted", "is_extra_encrypted"): + continue + + if field == "_password": + current_val = conn.password + snapshot_val = _decrypt_snapshot_field(post_snapshot, "_password") + elif field == "_extra": + current_val = conn.extra + snapshot_val = _decrypt_snapshot_field(post_snapshot, "_extra") + else: + current_val = getattr(conn, field) + snapshot_val = post_snapshot.get(field) + + if current_val != snapshot_val: + return False + return True + + +def attempt_revert(ct: ConnectionTest, *, session: Session) -> None: + """Revert a connection to its pre-edit values if no concurrent edit has occurred.""" + if not ct.connection_snapshot: + log.warning("attempt_revert called without snapshot for %s", ct.id) + return + + pre_snapshot = ct.connection_snapshot["pre"] + post_snapshot = ct.connection_snapshot["post"] + + conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id)) + if conn is None: + ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists." + log.warning("Revert skipped: connection %s no longer exists", ct.connection_id) + return + + if not _can_safely_revert(conn, post_snapshot): + ct.result_message = ( + ct.result_message or "" + ) + " | Revert skipped: connection was modified by another user." + log.warning( + "Revert skipped: concurrent edit detected on connection %s", + ct.connection_id, + ) + return + + _revert_connection(conn, pre_snapshot) + ct.reverted = True + log.info("Reverted connection %s to pre-edit state", ct.connection_id) diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts index e0e249462d833..59ee8d16cbf22 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -948,6 +948,7 @@ export type BackfillServiceCancelBackfillMutationResult = Awaited>; export type ConnectionServicePatchConnectionMutationResult = Awaited>; export type ConnectionServiceBulkConnectionsMutationResult = Awaited>; +export type ConnectionServicePatchConnectionAndTestMutationResult = Awaited>; export type DagRunServicePatchDagRunMutationResult = Awaited>; export type DagServicePatchDagsMutationResult = Awaited>; export type DagServicePatchDagMutationResult = Awaited>; diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index 11166fa1bd484..ea58efd99c22d 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -2118,6 +2118,29 @@ export const useConnectionServiceBulkConnections = ({ mutationFn: ({ requestBody }) => ConnectionService.bulkConnections({ requestBody }) as unknown as Promise, ...options }); /** +* Patch Connection And Test +* Update a connection and queue an async test with revert-on-failure. +* +* Atomically saves the edit and creates a ConnectionTest with snapshots of the +* pre-edit and post-edit state. If the test fails, the connection is automatically +* reverted to its pre-edit values. +* @param data The data for the request. +* @param data.connectionId +* @param data.requestBody +* @param data.updateMask +* @returns ConnectionSaveAndTestResponse Successful Response +* @throws ApiError +*/ +export const useConnectionServicePatchConnectionAndTest = (options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ connectionId, requestBody, updateMask }) => ConnectionService.patchConnectionAndTest({ connectionId, requestBody, updateMask }) as unknown as Promise, ...options }); +/** * Patch Dag Run * Modify a DAG Run. * @param data The data for the request. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 65da80057b6fe..ca61dfd80daea 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1698,6 +1698,26 @@ export const $ConnectionResponse = { description: 'Connection serializer for responses.' } as const; +export const $ConnectionSaveAndTestResponse = { + properties: { + connection: { + '$ref': '#/components/schemas/ConnectionResponse' + }, + test_token: { + type: 'string', + title: 'Test Token' + }, + test_state: { + type: 'string', + title: 'Test State' + } + }, + type: 'object', + required: ['connection', 'test_token', 'test_state'], + title: 'ConnectionSaveAndTestResponse', + description: 'Response returned by the combined save-and-test endpoint.' +} as const; + export const $ConnectionTestQueuedResponse = { properties: { token: { @@ -1790,6 +1810,11 @@ export const $ConnectionTestStatusResponse = { type: 'string', format: 'date-time', title: 'Created At' + }, + reverted: { + type: 'boolean', + title: 'Reverted', + default: false } }, type: 'object', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index f4769cc7e269f..f6b0e693de356 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -3,7 +3,7 @@ import type { CancelablePromise } from './core/CancelablePromise'; import { OpenAPI } from './core/OpenAPI'; import { request as __request } from './core/request'; -import type { BulkConnectionsData, BulkConnectionsResponse, BulkPoolsData, BulkPoolsResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, BulkVariablesData, BulkVariablesResponse, CancelBackfillData, CancelBackfillResponse, ClearDagRunData, ClearDagRunResponse, CreateAssetEventData, CreateAssetEventResponse, CreateBackfillData, CreateBackfillDryRunData, CreateBackfillDryRunResponse, CreateBackfillResponse, CreateDefaultConnectionsResponse, CreateXcomEntryData, CreateXcomEntryResponse, DagStatsResponse2, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, DeleteConnectionData, DeleteConnectionResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, DeleteDagData, DeleteDagResponse, DeleteDagRunData, DeleteDagRunResponse, DeletePoolData, DeletePoolResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, DeleteVariableData, DeleteVariableResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, FavoriteDagData, FavoriteDagResponse, GenerateTokenData, GenerateTokenResponse2, GetAssetAliasData, GetAssetAliasResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetData, GetAssetEventsData, GetAssetEventsResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, GetAssetResponse, GetAssetsData, GetAssetsResponse, GetAuthMenusResponse, GetBackfillData, GetBackfillResponse, GetCalendarData, GetCalendarResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, GetConnectionData, GetConnectionResponse, GetConnectionTestStatusData, GetConnectionTestStatusResponse, GetConnectionsData, GetConnectionsResponse, GetCurrentUserInfoResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, GetDagData, GetDagDetailsData, GetDagDetailsResponse, GetDagResponse, GetDagRunData, GetDagRunDeadlinesData, GetDagRunDeadlinesResponse, GetDagRunResponse, GetDagRunsData, GetDagRunsResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetDagStructureData, GetDagStructureResponse, GetDagTagsData, GetDagTagsResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetDagsData, GetDagsResponse, GetDagsUiData, GetDagsUiResponse, GetDependenciesData, GetDependenciesResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, GetExtraLinksData, GetExtraLinksResponse, GetGanttDataData, GetGanttDataResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesData, GetGridTiSummariesResponse, GetGridTiSummariesStreamData, GetGridTiSummariesStreamResponse, GetHealthResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetLogData, GetLogResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetPartitionedDagRunsData, GetPartitionedDagRunsResponse, GetPendingPartitionedDagRunData, GetPendingPartitionedDagRunResponse, GetPluginsData, GetPluginsResponse, GetPoolData, GetPoolResponse, GetPoolsData, GetPoolsResponse, GetProvidersData, GetProvidersResponse, GetTaskData, GetTaskInstanceData, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstancesData, GetTaskInstancesResponse, GetTaskResponse, GetTasksData, GetTasksResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, GetVariableData, GetVariableResponse, GetVariablesData, GetVariablesResponse, GetVersionResponse, GetXcomEntriesData, GetXcomEntriesResponse, GetXcomEntryData, GetXcomEntryResponse, HistoricalMetricsData, HistoricalMetricsResponse, HookMetaDataResponse, ImportErrorsResponse, ListBackfillsData, ListBackfillsResponse, ListBackfillsUiData, ListBackfillsUiResponse, ListDagWarningsData, ListDagWarningsResponse, ListTeamsData, ListTeamsResponse, LoginData, LoginResponse, LogoutResponse, MaterializeAssetData, MaterializeAssetResponse, NextRunAssetsData, NextRunAssetsResponse, PatchConnectionData, PatchConnectionResponse, PatchDagData, PatchDagResponse, PatchDagRunData, PatchDagRunResponse, PatchDagsData, PatchDagsResponse, PatchPoolData, PatchPoolResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, PatchTaskInstanceData, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, PatchTaskInstanceResponse, PatchVariableData, PatchVariableResponse, PauseBackfillData, PauseBackfillResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PostConnectionData, PostConnectionResponse, PostPoolData, PostPoolResponse, PostVariableData, PostVariableResponse, ReparseDagFileData, ReparseDagFileResponse, StructureDataData, StructureDataResponse2, TestConnectionAsyncData, TestConnectionAsyncResponse, TestConnectionData, TestConnectionResponse, TriggerDagRunData, TriggerDagRunResponse, UnfavoriteDagData, UnfavoriteDagResponse, UnpauseBackfillData, UnpauseBackfillResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse } from './types.gen'; +import type { BulkConnectionsData, BulkConnectionsResponse, BulkPoolsData, BulkPoolsResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, BulkVariablesData, BulkVariablesResponse, CancelBackfillData, CancelBackfillResponse, ClearDagRunData, ClearDagRunResponse, CreateAssetEventData, CreateAssetEventResponse, CreateBackfillData, CreateBackfillDryRunData, CreateBackfillDryRunResponse, CreateBackfillResponse, CreateDefaultConnectionsResponse, CreateXcomEntryData, CreateXcomEntryResponse, DagStatsResponse2, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, DeleteConnectionData, DeleteConnectionResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, DeleteDagData, DeleteDagResponse, DeleteDagRunData, DeleteDagRunResponse, DeletePoolData, DeletePoolResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, DeleteVariableData, DeleteVariableResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, FavoriteDagData, FavoriteDagResponse, GenerateTokenData, GenerateTokenResponse2, GetAssetAliasData, GetAssetAliasResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetData, GetAssetEventsData, GetAssetEventsResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, GetAssetResponse, GetAssetsData, GetAssetsResponse, GetAuthMenusResponse, GetBackfillData, GetBackfillResponse, GetCalendarData, GetCalendarResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, GetConnectionData, GetConnectionResponse, GetConnectionTestStatusData, GetConnectionTestStatusResponse, GetConnectionsData, GetConnectionsResponse, GetCurrentUserInfoResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, GetDagData, GetDagDetailsData, GetDagDetailsResponse, GetDagResponse, GetDagRunData, GetDagRunDeadlinesData, GetDagRunDeadlinesResponse, GetDagRunResponse, GetDagRunsData, GetDagRunsResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetDagStructureData, GetDagStructureResponse, GetDagTagsData, GetDagTagsResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetDagsData, GetDagsResponse, GetDagsUiData, GetDagsUiResponse, GetDependenciesData, GetDependenciesResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, GetExtraLinksData, GetExtraLinksResponse, GetGanttDataData, GetGanttDataResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesData, GetGridTiSummariesResponse, GetGridTiSummariesStreamData, GetGridTiSummariesStreamResponse, GetHealthResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetLogData, GetLogResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetPartitionedDagRunsData, GetPartitionedDagRunsResponse, GetPendingPartitionedDagRunData, GetPendingPartitionedDagRunResponse, GetPluginsData, GetPluginsResponse, GetPoolData, GetPoolResponse, GetPoolsData, GetPoolsResponse, GetProvidersData, GetProvidersResponse, GetTaskData, GetTaskInstanceData, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstancesData, GetTaskInstancesResponse, GetTaskResponse, GetTasksData, GetTasksResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, GetVariableData, GetVariableResponse, GetVariablesData, GetVariablesResponse, GetVersionResponse, GetXcomEntriesData, GetXcomEntriesResponse, GetXcomEntryData, GetXcomEntryResponse, HistoricalMetricsData, HistoricalMetricsResponse, HookMetaDataResponse, ImportErrorsResponse, ListBackfillsData, ListBackfillsResponse, ListBackfillsUiData, ListBackfillsUiResponse, ListDagWarningsData, ListDagWarningsResponse, ListTeamsData, ListTeamsResponse, LoginData, LoginResponse, LogoutResponse, MaterializeAssetData, MaterializeAssetResponse, NextRunAssetsData, NextRunAssetsResponse, PatchConnectionAndTestData, PatchConnectionAndTestResponse, PatchConnectionData, PatchConnectionResponse, PatchDagData, PatchDagResponse, PatchDagRunData, PatchDagRunResponse, PatchDagsData, PatchDagsResponse, PatchPoolData, PatchPoolResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, PatchTaskInstanceData, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, PatchTaskInstanceResponse, PatchVariableData, PatchVariableResponse, PauseBackfillData, PauseBackfillResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PostConnectionData, PostConnectionResponse, PostPoolData, PostPoolResponse, PostVariableData, PostVariableResponse, ReparseDagFileData, ReparseDagFileResponse, StructureDataData, StructureDataResponse2, TestConnectionAsyncData, TestConnectionAsyncResponse, TestConnectionData, TestConnectionResponse, TriggerDagRunData, TriggerDagRunResponse, UnfavoriteDagData, UnfavoriteDagResponse, UnpauseBackfillData, UnpauseBackfillResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse } from './types.gen'; export class AssetService { /** @@ -768,6 +768,42 @@ export class ConnectionService { }); } + /** + * Patch Connection And Test + * Update a connection and queue an async test with revert-on-failure. + * + * Atomically saves the edit and creates a ConnectionTest with snapshots of the + * pre-edit and post-edit state. If the test fails, the connection is automatically + * reverted to its pre-edit values. + * @param data The data for the request. + * @param data.connectionId + * @param data.requestBody + * @param data.updateMask + * @returns ConnectionSaveAndTestResponse Successful Response + * @throws ApiError + */ + public static patchConnectionAndTest(data: PatchConnectionAndTestData): CancelablePromise { + return __request(OpenAPI, { + method: 'PATCH', + url: '/api/v2/connections/{connection_id}/save-and-test', + path: { + connection_id: data.connectionId + }, + query: { + update_mask: data.updateMask + }, + body: data.requestBody, + mediaType: 'application/json', + errors: { + 400: 'Bad Request', + 401: 'Unauthorized', + 403: 'Forbidden', + 404: 'Not Found', + 422: 'Validation Error' + } + }); + } + /** * Test Connection * Test an API connection. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 388a3a69beae4..db434f1bf9160 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -495,6 +495,15 @@ export type ConnectionResponse = { team_name: string | null; }; +/** + * Response returned by the combined save-and-test endpoint. + */ +export type ConnectionSaveAndTestResponse = { + connection: ConnectionResponse; + test_token: string; + test_state: string; +}; + /** * Response returned when an async connection test is queued. */ @@ -529,6 +538,7 @@ export type ConnectionTestStatusResponse = { state: string; result_message?: string | null; created_at: string; + reverted?: boolean; }; /** @@ -2504,6 +2514,14 @@ export type BulkConnectionsData = { export type BulkConnectionsResponse = BulkResponse; +export type PatchConnectionAndTestData = { + connectionId: string; + requestBody: ConnectionBody; + updateMask?: Array<(string)> | null; +}; + +export type PatchConnectionAndTestResponse = ConnectionSaveAndTestResponse; + export type TestConnectionData = { requestBody: ConnectionBody; }; @@ -4482,6 +4500,37 @@ export type $OpenApiTs = { }; }; }; + '/api/v2/connections/{connection_id}/save-and-test': { + patch: { + req: PatchConnectionAndTestData; + res: { + /** + * Successful Response + */ + 200: ConnectionSaveAndTestResponse; + /** + * Bad Request + */ + 400: HTTPExceptionResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; '/api/v2/connections/test': { post: { req: TestConnectionData; diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 6f5c6472d311f..316f5542f5493 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -1265,6 +1265,128 @@ def test_get_status_returns_404_for_invalid_token(self, test_client): assert response.status_code == 404 +class TestSaveAndTest(TestConnectionEndpoint): + """Tests for the combined PATCH /{connection_id}/save-and-test endpoint.""" + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_save_and_test_returns_200_with_token(self, test_client, session): + """PATCH save-and-test updates the connection and returns a test token.""" + self.create_connection() + response = test_client.patch( + f"/connections/{TEST_CONN_ID}/save-and-test", + json={ + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "host": "updated-host.example.com", + }, + ) + assert response.status_code == 200 + body = response.json() + assert body["test_token"] + assert body["test_state"] == "pending" + assert body["connection"]["host"] == "updated-host.example.com" + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_save_and_test_creates_snapshot(self, test_client, session): + """PATCH save-and-test creates a ConnectionTest with a connection_snapshot.""" + self.create_connection() + response = test_client.patch( + f"/connections/{TEST_CONN_ID}/save-and-test", + json={ + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "host": "new-host.example.com", + }, + ) + assert response.status_code == 200 + token = response.json()["test_token"] + + ct = session.scalar(select(ConnectionTest).filter_by(token=token)) + assert ct is not None + assert ct.connection_snapshot is not None + snapshot = ct.connection_snapshot + assert "pre" in snapshot + assert "post" in snapshot + assert snapshot["pre"]["host"] == TEST_CONN_HOST + assert snapshot["post"]["host"] == "new-host.example.com" + + def test_save_and_test_403_when_disabled(self, test_client): + """PATCH save-and-test returns 403 when test_connection is disabled.""" + self.create_connection() + response = test_client.patch( + f"/connections/{TEST_CONN_ID}/save-and-test", + json={ + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + }, + ) + assert response.status_code == 403 + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_save_and_test_404_for_nonexistent(self, test_client): + """PATCH save-and-test returns 404 for nonexistent connection.""" + response = test_client.patch( + "/connections/nonexistent/save-and-test", + json={ + "connection_id": "nonexistent", + "conn_type": "http", + }, + ) + assert response.status_code == 404 + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_poll_shows_reverted_true_after_failed_test(self, test_client, session): + """GET status shows reverted=True after a failed test triggers a revert.""" + self.create_connection() + response = test_client.patch( + f"/connections/{TEST_CONN_ID}/save-and-test", + json={ + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "host": "bad-host.example.com", + }, + ) + token = response.json()["test_token"] + + ct = session.scalar(select(ConnectionTest).filter_by(token=token)) + ct.state = ConnectionTestState.FAILED + ct.result_message = "Connection refused" + + from airflow.models.connection_test import attempt_revert + + attempt_revert(ct, session=session) + session.commit() + + poll_response = test_client.get(f"/connections/test-async/{token}") + assert poll_response.status_code == 200 + body = poll_response.json() + assert body["reverted"] is True + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_poll_shows_reverted_false_for_success(self, test_client, session): + """GET status shows reverted=False for a successful test.""" + self.create_connection() + response = test_client.patch( + f"/connections/{TEST_CONN_ID}/save-and-test", + json={ + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "host": "good-host.example.com", + }, + ) + token = response.json()["test_token"] + + ct = session.scalar(select(ConnectionTest).filter_by(token=token)) + ct.state = ConnectionTestState.SUCCESS + ct.result_message = "Connection OK" + session.commit() + + poll_response = test_client.get(f"/connections/test-async/{token}") + assert poll_response.status_code == 200 + body = poll_response.json() + assert body["reverted"] is False + + class TestCreateDefaultConnections(TestConnectionEndpoint): def test_should_respond_204(self, test_client, session): response = test_client.post("/connections/defaults") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py index 1951d90f6f13a..a1c406224f934 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py @@ -17,8 +17,12 @@ from __future__ import annotations import pytest +from sqlalchemy import select -from airflow.models.connection_test import ConnectionTest, ConnectionTestState +from airflow.models.connection import Connection +from airflow.models.connection_test import ConnectionTest, ConnectionTestState, snapshot_connection + +from tests_common.test_utils.db import clear_db_connection_tests, clear_db_connections pytestmark = pytest.mark.db_test @@ -75,3 +79,137 @@ def test_patch_returns_409_for_terminal_state(self, client, session): ) assert response.status_code == 409 assert "terminal state" in response.json()["detail"]["message"] + + +class TestPatchConnectionTestRevert: + """Tests for the revert-on-failure behavior in the execution API.""" + + @pytest.fixture(autouse=True) + def setup_teardown(self): + clear_db_connections(add_default_connections_back=False) + clear_db_connection_tests() + yield + clear_db_connections(add_default_connections_back=False) + clear_db_connection_tests() + + def test_patch_failed_with_snapshot_reverts_connection(self, client, session): + """PATCH with state=failed and snapshot triggers revert.""" + conn = Connection( + conn_id="revert_conn", + conn_type="postgres", + host="old-host.example.com", + login="old_user", + ) + session.add(conn) + session.flush() + + pre_snap = snapshot_connection(conn) + conn.host = "new-host.example.com" + conn.login = "new_user" + post_snap = snapshot_connection(conn) + session.flush() + + ct = ConnectionTest(connection_id="revert_conn") + ct.state = ConnectionTestState.RUNNING + ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} + session.add(ct) + session.commit() + + response = client.patch( + f"/execution/connection-tests/{ct.id}", + json={"state": "failed", "result_message": "Connection refused"}, + ) + assert response.status_code == 204 + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.reverted is True + conn = session.scalar(select(Connection).filter_by(conn_id="revert_conn")) + assert conn.host == "old-host.example.com" + assert conn.login == "old_user" + + def test_patch_success_with_snapshot_no_revert(self, client, session): + """PATCH with state=success does not trigger revert even with snapshot.""" + conn = Connection( + conn_id="no_revert_conn", + conn_type="postgres", + host="old-host.example.com", + ) + session.add(conn) + session.flush() + + pre_snap = snapshot_connection(conn) + conn.host = "new-host.example.com" + post_snap = snapshot_connection(conn) + session.flush() + + ct = ConnectionTest(connection_id="no_revert_conn") + ct.state = ConnectionTestState.RUNNING + ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} + session.add(ct) + session.commit() + + response = client.patch( + f"/execution/connection-tests/{ct.id}", + json={"state": "success", "result_message": "Connection OK"}, + ) + assert response.status_code == 204 + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.reverted is False + conn = session.scalar(select(Connection).filter_by(conn_id="no_revert_conn")) + assert conn.host == "new-host.example.com" + + def test_patch_failed_without_snapshot_no_revert(self, client, session): + """PATCH with state=failed but no snapshot does not trigger revert.""" + ct = ConnectionTest(connection_id="test_conn") + ct.state = ConnectionTestState.RUNNING + session.add(ct) + session.commit() + + response = client.patch( + f"/execution/connection-tests/{ct.id}", + json={"state": "failed", "result_message": "Connection refused"}, + ) + assert response.status_code == 204 + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.reverted is False + + def test_patch_failed_concurrent_edit_skips_revert(self, client, session): + """PATCH with state=failed skips revert when connection was modified concurrently.""" + conn = Connection( + conn_id="concurrent_conn", + conn_type="postgres", + host="old-host.example.com", + ) + session.add(conn) + session.flush() + + pre_snap = snapshot_connection(conn) + conn.host = "new-host.example.com" + post_snap = snapshot_connection(conn) + + conn.host = "third-party-host.example.com" + session.flush() + + ct = ConnectionTest(connection_id="concurrent_conn") + ct.state = ConnectionTestState.RUNNING + ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} + session.add(ct) + session.commit() + + response = client.patch( + f"/execution/connection-tests/{ct.id}", + json={"state": "failed", "result_message": "Connection refused"}, + ) + assert response.status_code == 204 + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.reverted is False + assert "modified by another user" in ct.result_message + conn = session.scalar(select(Connection).filter_by(conn_id="concurrent_conn")) + assert conn.host == "third-party-host.example.com" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_connection_tests.py b/airflow-core/tests/unit/jobs/test_scheduler_connection_tests.py deleted file mode 100644 index b82e78d972cb9..0000000000000 --- a/airflow-core/tests/unit/jobs/test_scheduler_connection_tests.py +++ /dev/null @@ -1,186 +0,0 @@ -# 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. -"""Tests for scheduler connection test dispatch and reaper.""" - -from __future__ import annotations - -import logging -import os -from datetime import timedelta -from unittest import mock - -import pytest -import time_machine -from sqlalchemy import delete - -from airflow._shared.timezones import timezone -from airflow.executors.base_executor import BaseExecutor -from airflow.executors.local_executor import LocalExecutor -from airflow.jobs.job import Job -from airflow.jobs.scheduler_job_runner import SchedulerJobRunner -from airflow.models.connection_test import ConnectionTest, ConnectionTestState - -pytestmark = pytest.mark.db_test - - -@pytest.fixture -def scheduler_job_runner(session): - """Create a SchedulerJobRunner with a mock Job and supporting executor.""" - session.execute(delete(ConnectionTest)) - session.commit() - - mock_job = mock.MagicMock(spec=Job) - mock_job.id = 1 - mock_job.max_tis_per_query = 16 - executor = LocalExecutor() - executor.queued_connection_tests.clear() - runner = SchedulerJobRunner.__new__(SchedulerJobRunner) - runner.job = mock_job - runner.executors = [executor] - runner.executor = executor - runner._log = mock.MagicMock(spec=logging.Logger) - return runner - - -class TestDispatchConnectionTests: - @mock.patch.dict( - os.environ, - { - "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", - "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", - }, - ) - def test_dispatch_pending_tests(self, scheduler_job_runner, session): - """Pending connection tests are dispatched to a supporting executor.""" - ct = ConnectionTest(connection_id="test_conn") - session.add(ct) - session.commit() - assert ct.state == ConnectionTestState.PENDING - - scheduler_job_runner._dispatch_connection_tests(session=session) - - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.QUEUED - assert len(scheduler_job_runner.executor.queued_connection_tests) == 1 - - @mock.patch.dict( - os.environ, - { - "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "1", - "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", - }, - ) - def test_dispatch_respects_concurrency_limit(self, scheduler_job_runner, session): - """Excess pending tests stay PENDING when concurrency is at capacity.""" - # Create one already-queued test to fill the concurrency budget (limit=1) - ct_active = ConnectionTest(connection_id="active_conn") - ct_active.state = ConnectionTestState.QUEUED - session.add(ct_active) - - ct_pending = ConnectionTest(connection_id="pending_conn") - session.add(ct_pending) - session.commit() - - scheduler_job_runner._dispatch_connection_tests(session=session) - - session.expire_all() - ct_pending = session.get(ConnectionTest, ct_pending.id) - assert ct_pending.state == ConnectionTestState.PENDING - - @mock.patch.dict( - os.environ, - { - "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", - "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", - }, - ) - def test_dispatch_fails_fast_when_no_executor_supports(self, scheduler_job_runner, session): - """Tests fail immediately when no executor supports connection testing.""" - # Replace executor with one that doesn't support connection tests - unsupporting_executor = BaseExecutor() - unsupporting_executor.supports_connection_test = False - scheduler_job_runner.executors = [unsupporting_executor] - scheduler_job_runner.executor = unsupporting_executor - - ct = ConnectionTest(connection_id="test_conn") - session.add(ct) - session.commit() - - scheduler_job_runner._dispatch_connection_tests(session=session) - - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.FAILED - assert "No executor supports connection testing" in ct.result_message - - @mock.patch.dict( - os.environ, - { - "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", - "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", - }, - ) - def test_dispatch_with_queue_falls_back_to_global_executor(self, scheduler_job_runner, session): - """Tests with a queue are dispatched to the global executor as fallback.""" - ct = ConnectionTest(connection_id="test_conn", queue="gpu_workers") - session.add(ct) - session.commit() - - # Default LocalExecutor has team_name=None — serves as fallback for any queue - scheduler_job_runner._dispatch_connection_tests(session=session) - - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.QUEUED - assert len(scheduler_job_runner.executor.queued_connection_tests) == 1 - - -class TestReapStaleConnectionTests: - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) - def test_reap_stale_queued_test(self, scheduler_job_runner, session): - """Stale QUEUED tests are marked as FAILED by the reaper.""" - initial_time = timezone.utcnow() - - with time_machine.travel(initial_time, tick=False): - ct = ConnectionTest(connection_id="test_conn") - ct.state = ConnectionTestState.QUEUED - session.add(ct) - session.commit() - - # Jump forward past timeout (60s) + grace period (30s) - with time_machine.travel(initial_time + timedelta(seconds=200), tick=False): - scheduler_job_runner._reap_stale_connection_tests(session=session) - - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.FAILED - assert "timed out" in ct.result_message - - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) - def test_does_not_reap_fresh_tests(self, scheduler_job_runner, session): - """Fresh QUEUED tests are not reaped.""" - ct = ConnectionTest(connection_id="test_conn") - ct.state = ConnectionTestState.QUEUED - session.add(ct) - session.commit() - - scheduler_job_runner._reap_stale_connection_tests(session=session) - - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.QUEUED diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 997ebd24bcf0d..06b90355eae98 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -67,6 +67,8 @@ ) from airflow.models.backfill import Backfill, _create_backfill from airflow.models.callback import ExecutorCallback +from airflow.models.connection import Connection +from airflow.models.connection_test import ConnectionTest, ConnectionTestState, snapshot_connection from airflow.models.dag import DagModel, get_last_dagrun, infer_automated_data_interval from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel @@ -113,6 +115,7 @@ clear_db_assets, clear_db_backfills, clear_db_callbacks, + clear_db_connections, clear_db_dag_bundles, clear_db_dags, clear_db_deadline, @@ -9445,3 +9448,210 @@ def test_fallback_values_used_only_when_dag_version_is_none(self): assert _extract_bundle_name(ti) == "fallback-bundle" assert _extract_bundle_version(ti) == "fallback-v1" + + +@pytest.fixture +def scheduler_job_runner_for_connection_tests(session): + """Create a SchedulerJobRunner with a mock Job and supporting executor.""" + session.execute(delete(ConnectionTest)) + session.commit() + + mock_job = mock.MagicMock(spec=Job) + mock_job.id = 1 + mock_job.max_tis_per_query = 16 + executor = LocalExecutor() + executor.queued_connection_tests.clear() + runner = SchedulerJobRunner.__new__(SchedulerJobRunner) + runner.job = mock_job + runner.executors = [executor] + runner.executor = executor + runner._log = mock.MagicMock(spec=logging.Logger) + return runner + + +class TestDispatchConnectionTests: + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_pending_tests(self, scheduler_job_runner_for_connection_tests, session): + """Pending connection tests are dispatched to a supporting executor.""" + ct = ConnectionTest(connection_id="test_conn") + session.add(ct) + session.commit() + assert ct.state == ConnectionTestState.PENDING + + scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.QUEUED + assert len(scheduler_job_runner_for_connection_tests.executor.queued_connection_tests) == 1 + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "1", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_respects_concurrency_limit(self, scheduler_job_runner_for_connection_tests, session): + """Excess pending tests stay PENDING when concurrency is at capacity.""" + ct_active = ConnectionTest(connection_id="active_conn") + ct_active.state = ConnectionTestState.QUEUED + session.add(ct_active) + + ct_pending = ConnectionTest(connection_id="pending_conn") + session.add(ct_pending) + session.commit() + + scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + + session.expire_all() + ct_pending = session.get(ConnectionTest, ct_pending.id) + assert ct_pending.state == ConnectionTestState.PENDING + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_fails_fast_when_no_executor_supports( + self, scheduler_job_runner_for_connection_tests, session + ): + """Tests fail immediately when no executor supports connection testing.""" + unsupporting_executor = BaseExecutor() + unsupporting_executor.supports_connection_test = False + scheduler_job_runner_for_connection_tests.executors = [unsupporting_executor] + scheduler_job_runner_for_connection_tests.executor = unsupporting_executor + + ct = ConnectionTest(connection_id="test_conn") + session.add(ct) + session.commit() + + scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert "No executor supports connection testing" in ct.result_message + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_with_queue_falls_back_to_global_executor( + self, scheduler_job_runner_for_connection_tests, session + ): + """Tests with a queue are dispatched to the global executor as fallback.""" + ct = ConnectionTest(connection_id="test_conn", queue="gpu_workers") + session.add(ct) + session.commit() + + scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.QUEUED + assert len(scheduler_job_runner_for_connection_tests.executor.queued_connection_tests) == 1 + + +class TestReapStaleConnectionTests: + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) + def test_reap_stale_queued_test(self, scheduler_job_runner_for_connection_tests, session): + """Stale QUEUED tests are marked as FAILED by the reaper.""" + initial_time = timezone.utcnow() + + with time_machine.travel(initial_time, tick=False): + ct = ConnectionTest(connection_id="test_conn") + ct.state = ConnectionTestState.QUEUED + session.add(ct) + session.commit() + + with time_machine.travel(initial_time + timedelta(seconds=200), tick=False): + scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert "timed out" in ct.result_message + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) + def test_does_not_reap_fresh_tests(self, scheduler_job_runner_for_connection_tests, session): + """Fresh QUEUED tests are not reaped.""" + ct = ConnectionTest(connection_id="test_conn") + ct.state = ConnectionTestState.QUEUED + session.add(ct) + session.commit() + + scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.QUEUED + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) + def test_reaper_reverts_connection_on_timeout_with_snapshot( + self, scheduler_job_runner_for_connection_tests, session + ): + """Stale tests with a snapshot trigger revert on timeout.""" + clear_db_connections(add_default_connections_back=False) + + conn = Connection( + conn_id="reaper_conn", + conn_type="postgres", + host="old-host.example.com", + ) + session.add(conn) + session.flush() + + pre_snap = snapshot_connection(conn) + conn.host = "new-host.example.com" + post_snap = snapshot_connection(conn) + session.flush() + + initial_time = timezone.utcnow() + with time_machine.travel(initial_time, tick=False): + ct = ConnectionTest(connection_id="reaper_conn") + ct.state = ConnectionTestState.QUEUED + ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} + session.add(ct) + session.commit() + + with time_machine.travel(initial_time + timedelta(seconds=200), tick=False): + scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert ct.reverted is True + conn = session.scalar(select(Connection).filter_by(conn_id="reaper_conn")) + assert conn.host == "old-host.example.com" + + clear_db_connections(add_default_connections_back=False) + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) + def test_reaper_no_revert_without_snapshot(self, scheduler_job_runner_for_connection_tests, session): + """Stale tests without a snapshot do not trigger revert.""" + initial_time = timezone.utcnow() + with time_machine.travel(initial_time, tick=False): + ct = ConnectionTest(connection_id="no_snap_conn") + ct.state = ConnectionTestState.QUEUED + session.add(ct) + session.commit() + + with time_machine.travel(initial_time + timedelta(seconds=200), tick=False): + scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert ct.reverted is False diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py index 435c2d1a4c207..af1ce9dc2346a 100644 --- a/airflow-core/tests/unit/models/test_connection_test.py +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -21,7 +21,15 @@ import pytest from airflow.models.connection import Connection -from airflow.models.connection_test import ConnectionTest, ConnectionTestState, run_connection_test +from airflow.models.connection_test import ( + ConnectionTest, + ConnectionTestState, + attempt_revert, + run_connection_test, + snapshot_connection, +) + +from tests_common.test_utils.db import clear_db_connection_tests, clear_db_connections pytestmark = pytest.mark.db_test @@ -90,3 +98,182 @@ def test_exception_during_connection_test(self): assert success is False assert "Could not resolve host" in message + + +class TestSnapshotConnection: + def test_snapshot_captures_all_fields(self): + """snapshot_connection captures all expected fields including encrypted ones.""" + conn = Connection( + conn_id="snap_test", + conn_type="postgres", + host="db.example.com", + login="user", + password="secret", + schema="mydb", + port=5432, + extra='{"key": "value"}', + ) + snap = snapshot_connection(conn) + assert snap["conn_type"] == "postgres" + assert snap["host"] == "db.example.com" + assert snap["login"] == "user" + assert snap["schema"] == "mydb" + assert snap["port"] == 5432 + assert snap["_password"] is not None + assert snap["_password"] != "secret" # Should be encrypted + assert snap["_extra"] is not None + assert snap["is_encrypted"] is True + assert snap["is_extra_encrypted"] is True + + def test_snapshot_with_null_password_and_extra(self): + """snapshot_connection handles None password and extra.""" + conn = Connection(conn_id="snap_test", conn_type="http") + snap = snapshot_connection(conn) + assert snap["_password"] is None + assert snap["_extra"] is None + assert not snap["is_encrypted"] + assert not snap["is_extra_encrypted"] + + +class TestAttemptRevert: + @pytest.fixture(autouse=True) + def setup_teardown(self): + clear_db_connections(add_default_connections_back=False) + clear_db_connection_tests() + yield + clear_db_connections(add_default_connections_back=False) + clear_db_connection_tests() + + def test_attempt_revert_success(self, session): + """attempt_revert restores all connection fields and sets reverted=True.""" + conn = Connection( + conn_id="revert_conn", + conn_type="postgres", + host="old-host.example.com", + login="old_user", + password="old_secret", + schema="mydb", + port=5432, + extra='{"key": "old_value"}', + ) + session.add(conn) + session.flush() + + pre_snap = snapshot_connection(conn) + + conn.host = "new-host.example.com" + conn.login = "new_user" + conn.password = "new_secret" + conn.port = 9999 + conn.extra = '{"key": "new_value"}' + + post_snap = snapshot_connection(conn) + + ct = ConnectionTest(connection_id="revert_conn") + ct.state = ConnectionTestState.FAILED + ct.result_message = "Connection refused" + ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} + session.add(ct) + session.flush() + + attempt_revert(ct, session=session) + session.flush() + + assert ct.reverted is True + session.refresh(conn) + assert conn.host == "old-host.example.com" + assert conn.login == "old_user" + assert conn.password == "old_secret" + assert conn.schema == "mydb" + assert conn.port == 5432 + assert conn.extra == '{"key": "old_value"}' + + def test_attempt_revert_skipped_concurrent_edit(self, session): + """attempt_revert skips revert when connection was modified by another user.""" + conn = Connection( + conn_id="concurrent_conn", + conn_type="postgres", + host="old-host.example.com", + ) + session.add(conn) + session.flush() + + pre_snap = snapshot_connection(conn) + conn.host = "new-host.example.com" + post_snap = snapshot_connection(conn) + + conn.host = "third-party-host.example.com" + session.flush() + + ct = ConnectionTest(connection_id="concurrent_conn") + ct.state = ConnectionTestState.FAILED + ct.result_message = "Connection refused" + ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} + session.add(ct) + session.flush() + + attempt_revert(ct, session=session) + + assert ct.reverted is False + assert "modified by another user" in ct.result_message + assert conn.host == "third-party-host.example.com" + + def test_attempt_revert_skipped_concurrent_password_edit(self, session): + """attempt_revert skips revert when password was changed concurrently.""" + conn = Connection( + conn_id="pw_conn", + conn_type="postgres", + host="host.example.com", + password="original_secret", + ) + session.add(conn) + session.flush() + + pre_snap = snapshot_connection(conn) + conn.password = "new_secret" + post_snap = snapshot_connection(conn) + + conn.password = "third_party_secret" + session.flush() + + ct = ConnectionTest(connection_id="pw_conn") + ct.state = ConnectionTestState.FAILED + ct.result_message = "Connection refused" + ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} + session.add(ct) + session.flush() + + attempt_revert(ct, session=session) + + assert ct.reverted is False + assert "modified by another user" in ct.result_message + session.refresh(conn) + assert conn.password == "third_party_secret" + + def test_attempt_revert_skipped_connection_deleted(self, session): + """attempt_revert skips revert when connection no longer exists.""" + conn = Connection( + conn_id="deleted_conn", + conn_type="postgres", + host="old-host.example.com", + ) + session.add(conn) + session.flush() + + pre_snap = snapshot_connection(conn) + conn.host = "new-host.example.com" + post_snap = snapshot_connection(conn) + + ct = ConnectionTest(connection_id="deleted_conn") + ct.state = ConnectionTestState.FAILED + ct.result_message = "Connection refused" + ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} + session.add(ct) + + session.delete(conn) + session.flush() + + attempt_revert(ct, session=session) + + assert ct.reverted is False + assert "no longer exists" in ct.result_message diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index 16fb15221f94c..e41de765672b6 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -244,6 +244,16 @@ class ConnectionResponse(BaseModel): team_name: Annotated[str | None, Field(title="Team Name")] = None +class ConnectionSaveAndTestResponse(BaseModel): + """ + Response returned by the combined save-and-test endpoint. + """ + + connection: ConnectionResponse + test_token: Annotated[str, Field(title="Test Token")] + test_state: Annotated[str, Field(title="Test State")] + + class ConnectionTestQueuedResponse(BaseModel): """ Response returned when an async connection test is queued. @@ -285,6 +295,7 @@ class ConnectionTestStatusResponse(BaseModel): state: Annotated[str, Field(title="State")] result_message: Annotated[str | None, Field(title="Result Message")] = None created_at: Annotated[datetime, Field(title="Created At")] + reverted: Annotated[bool | None, Field(title="Reverted")] = False class CreateAssetEventsBody(BaseModel): From fabebef0b626db696a3f78b0a9ae475a155fe700 Mon Sep 17 00:00:00 2001 From: Anish Date: Sat, 28 Feb 2026 13:26:06 -0600 Subject: [PATCH 08/38] regenrated auto generates after rebase --- airflow-core/docs/migrations-ref.rst | 7 ++++++- airflow-core/src/airflow/executors/base_executor.py | 1 - ...st_table.py => 0109_3_2_0_add_connection_test_table.py} | 4 ++-- ...n_snapshot.py => 0110_3_2_0_add_connection_snapshot.py} | 0 airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts | 2 +- .../src/airflow/ui/openapi-gen/requests/services.gen.ts | 2 +- airflow-core/src/airflow/utils/db.py | 2 +- 7 files changed, 11 insertions(+), 7 deletions(-) rename airflow-core/src/airflow/migrations/versions/{0107_3_2_0_add_connection_test_table.py => 0109_3_2_0_add_connection_test_table.py} (97%) rename airflow-core/src/airflow/migrations/versions/{0108_3_2_0_add_connection_snapshot.py => 0110_3_2_0_add_connection_snapshot.py} (100%) diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index 07d6f5afc16e9..abd2765129d55 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,12 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``1d6611b6ab7c`` (head) | ``888b59e02a5b`` | ``3.2.0`` | Add bundle_name to callback table. | +| ``b8f3e5d1a9c2`` (head) | ``a7e6d4c3b2f1`` | ``3.2.0`` | Add connection_snapshot and reverted columns to | +| | | | connection_test. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``a7e6d4c3b2f1`` | ``1d6611b6ab7c`` | ``3.2.0`` | Add connection_test table. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``1d6611b6ab7c`` | ``888b59e02a5b`` | ``3.2.0`` | Add bundle_name to callback table. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``888b59e02a5b`` | ``6222ce48e289`` | ``3.2.0`` | Fix migration file ORM inconsistencies. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index dd25d0fdfdf3c..00a6e0e151501 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -222,7 +222,6 @@ def log_task_event(self, *, event: str, extra: str, ti_key: TaskInstanceKey): self._task_event_logs.append(Log(event=event, task_instance=ti_key, extra=extra)) def queue_workload(self, workload: workloads.All, session: Session) -> None: -<<<<<<< HEAD if isinstance(workload, workloads.ExecuteTask): ti = workload.ti self.queued_tasks[ti.key] = workload diff --git a/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py similarity index 97% rename from airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py rename to airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py index 3ca39f7fc96b4..82710679aa30b 100644 --- a/airflow-core/src/airflow/migrations/versions/0107_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py @@ -20,7 +20,7 @@ Add connection_test table. Revision ID: a7e6d4c3b2f1 -Revises: 509b94a1042d +Revises: 888b59e02a5b Create Date: 2026-02-22 00:00:00.000000 """ @@ -34,7 +34,7 @@ # revision identifiers, used by Alembic. revision = "a7e6d4c3b2f1" -down_revision = "509b94a1042d" +down_revision = "888b59e02a5b" branch_labels = None depends_on = None airflow_version = "3.2.0" diff --git a/airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_snapshot.py b/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_snapshot.py similarity index 100% rename from airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_snapshot.py rename to airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_snapshot.py diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index ea58efd99c22d..871e937c5eaeb 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -1582,7 +1582,7 @@ export const useDashboardServiceDagStats = = unknown[]>({ dagId, dagRunId, limit, offset, orderBy }: { diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index f6b0e693de356..4540c4b2555b2 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -4086,7 +4086,7 @@ export class DeadlinesService { * @param data.limit * @param data.offset * @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, deadline_time, created_at, alert_name` - * @returns DealineCollectionResponse Successful Response + * @returns DeadlineCollectionResponse Successful Response * @throws ApiError */ public static getDagRunDeadlines(data: GetDagRunDeadlinesData): CancelablePromise { diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 9bc0608611b5a..e4e33efbe4358 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/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": "b8f3e5d1a9c2", } # Prefix used to identify tables holding data moved during migration. From 64a574d7f828d9723455e3bb0fc6afa88b445796 Mon Sep 17 00:00:00 2001 From: Anish Date: Sat, 28 Feb 2026 23:51:10 -0600 Subject: [PATCH 09/38] fix local executer --- .../src/airflow/executors/local_executor.py | 114 +++++++++++------- 1 file changed, 71 insertions(+), 43 deletions(-) diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index f0731bccea218..4d60ce7f834f2 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -35,9 +35,10 @@ import structlog -from airflow.executors import workloads as _workloads +from airflow.executors import workloads from airflow.executors.base_executor import BaseExecutor -from airflow.utils.state import TaskInstanceState +from airflow.executors.workloads.callback import execute_callback_workload +from airflow.utils.state import CallbackState, TaskInstanceState # add logger to parameter of setproctitle to support logging if sys.platform == "darwin": @@ -50,13 +51,23 @@ if TYPE_CHECKING: from structlog.typing import FilteringBoundLogger as Logger - TaskInstanceStateType = tuple[_workloads.TaskInstance, TaskInstanceState, Exception | None] + from airflow.executors.workloads.types import WorkloadResultType + + +def _get_executor_process_title_prefix(team_name: str | None) -> str: + """ + Build the process title prefix for LocalExecutor workers. + + :param team_name: Team name from executor configuration + """ + team_suffix = f" [{team_name}]" if team_name else "" + return f"airflow worker -- LocalExecutor{team_suffix}:" def _run_worker( logger_name: str, - input: SimpleQueue[_workloads.All | None], - output: Queue[TaskInstanceStateType], + input: SimpleQueue[workloads.All | None], + output: Queue[WorkloadResultType], unread_messages: multiprocessing.sharedctypes.Synchronized[int], team_conf, ): @@ -68,11 +79,8 @@ def _run_worker( log = structlog.get_logger(logger_name) log.info("Worker starting up pid=%d", os.getpid()) - # Create team suffix for process title - team_suffix = f" [{team_conf.team_name}]" if team_conf.team_name else "" - while True: - setproctitle(f"airflow worker -- LocalExecutor{team_suffix}: ", log) + setproctitle(f"{_get_executor_process_title_prefix(team_conf.team_name)} ", log) try: workload = input.get() except EOFError: @@ -87,34 +95,36 @@ def _run_worker( # Received poison pill, no more tasks to run return - if isinstance(workload, _workloads.TestConnection): - with unread_messages: - unread_messages.value -= 1 - _execute_connection_test(log, workload, team_conf) - continue - - if not isinstance(workload, _workloads.ExecuteTask): - raise ValueError(f"LocalExecutor does not know how to handle {type(workload)}") - # Decrement this as soon as we pick up a message off the queue with unread_messages: unread_messages.value -= 1 - key = None - if ti := getattr(workload, "ti", None): - key = ti.key - else: - raise TypeError(f"Don't know how to get ti key from {type(workload).__name__}") - try: - _execute_work(log, workload, team_conf) + # Handle different workload types + if isinstance(workload, workloads.ExecuteTask): + try: + _execute_work(log, workload, team_conf) + output.put((workload.ti.key, TaskInstanceState.SUCCESS, None)) + except Exception as e: + log.exception("Task execution failed.") + output.put((workload.ti.key, TaskInstanceState.FAILED, e)) + + elif isinstance(workload, workloads.ExecuteCallback): + output.put((workload.callback.id, CallbackState.RUNNING, None)) + try: + _execute_callback(log, workload, team_conf) + output.put((workload.callback.id, CallbackState.SUCCESS, None)) + except Exception as e: + log.exception("Callback execution failed") + output.put((workload.callback.id, CallbackState.FAILED, e)) + + elif isinstance(workload, workloads.TestConnection): + _execute_connection_test(log, workload, team_conf) - output.put((key, TaskInstanceState.SUCCESS, None)) - except Exception as e: - log.exception("uhoh") - output.put((key, TaskInstanceState.FAILED, e)) + else: + raise ValueError(f"LocalExecutor does not know how to handle {type(workload)}") -def _execute_work(log: Logger, workload: _workloads.ExecuteTask, team_conf) -> None: +def _execute_work(log: Logger, workload: workloads.ExecuteTask, team_conf) -> None: """ Execute command received and stores result state in queue. @@ -124,9 +134,7 @@ def _execute_work(log: Logger, workload: _workloads.ExecuteTask, team_conf) -> N """ from airflow.sdk.execution_time.supervisor import supervise - # Create team suffix for process title - team_suffix = f" [{team_conf.team_name}]" if team_conf.team_name else "" - setproctitle(f"airflow worker -- LocalExecutor{team_suffix}: {workload.ti.id}", log) + setproctitle(f"{_get_executor_process_title_prefix(team_conf.team_name)} {workload.ti.id}", log) base_url = team_conf.get("api", "base_url", fallback="/") # If it's a relative URL, use localhost:8080 as the default @@ -147,7 +155,23 @@ def _execute_work(log: Logger, workload: _workloads.ExecuteTask, team_conf) -> N ) -def _execute_connection_test(log: Logger, workload: _workloads.TestConnection, team_conf) -> None: +def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_conf) -> None: + """ + Execute a callback workload. + + :param log: Logger instance + :param workload: The ExecuteCallback workload to execute + :param team_conf: Team-specific executor configuration + """ + setproctitle(f"{_get_executor_process_title_prefix(team_conf.team_name)} {workload.callback.id}", log) + + success, error_msg = execute_callback_workload(workload.callback, log) + + if not success: + 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 with a timeout. @@ -163,9 +187,9 @@ def _execute_connection_test(log: Logger, workload: _workloads.TestConnection, t from airflow.models.connection_test import run_connection_test - team_suffix = f" [{team_conf.team_name}]" if team_conf.team_name else "" setproctitle( - f"airflow worker -- LocalExecutor{team_suffix}: connection-test {workload.connection_id}", log + f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}", + log, ) # Build execution API URL (same pattern as _execute_work) @@ -220,11 +244,12 @@ class LocalExecutor(BaseExecutor): is_mp_using_fork: bool = multiprocessing.get_start_method() == "fork" supports_multi_team: bool = True - supports_connection_test: bool = True serve_logs: bool = True + supports_callbacks: bool = True + supports_connection_test: bool = True - activity_queue: SimpleQueue[_workloads.All | None] - result_queue: SimpleQueue[TaskInstanceStateType] + activity_queue: SimpleQueue[workloads.All | None] + result_queue: SimpleQueue[WorkloadResultType] workers: dict[int, multiprocessing.Process] _unread_messages: multiprocessing.sharedctypes.Synchronized[int] @@ -367,11 +392,14 @@ def end(self) -> None: def terminate(self): """Terminate the executor is not doing anything.""" - def _process_workloads(self, workloads): - for workload in workloads: + def _process_workloads(self, workload_list): + for workload in workload_list: self.activity_queue.put(workload) - if isinstance(workload, _workloads.ExecuteTask): + # Remove from appropriate queue based on workload type + if isinstance(workload, workloads.ExecuteTask): del self.queued_tasks[workload.ti.key] + elif isinstance(workload, workloads.ExecuteCallback): + del self.queued_callbacks[workload.callback.id] with self._unread_messages: - self._unread_messages.value += len(workloads) + self._unread_messages.value += len(workload_list) self._check_workers() From 7f9029a8247e4cd1631e64c1c8a342d5c14a9b9a Mon Sep 17 00:00:00 2001 From: Anish Date: Sun, 1 Mar 2026 14:11:36 -0600 Subject: [PATCH 10/38] refactor local executor --- .../src/airflow/executors/local_executor.py | 55 ++++++++++--------- .../src/airflow/models/connection_test.py | 5 +- .../unit/executors/test_base_executor.py | 5 +- .../tests/unit/models/test_connection_test.py | 4 +- task-sdk/src/airflow/sdk/api/client.py | 22 ++++++++ 5 files changed, 57 insertions(+), 34 deletions(-) diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 4d60ce7f834f2..32eab13d96f47 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -29,6 +29,7 @@ import multiprocessing import multiprocessing.sharedctypes import os +import signal import sys from multiprocessing import Queue, SimpleQueue from typing import TYPE_CHECKING @@ -38,6 +39,7 @@ from airflow.executors import workloads from airflow.executors.base_executor import BaseExecutor from airflow.executors.workloads.callback import execute_callback_workload +from airflow.models.connection_test import ConnectionTestState, run_connection_test from airflow.utils.state import CallbackState, TaskInstanceState # add logger to parameter of setproctitle to support logging @@ -173,62 +175,63 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None: """ - Execute a connection test workload with a timeout. + Execute a connection test workload. - Results are reported back via the Execution API (workers must not access the metadata DB directly). + Results are reported back via the Execution API. :param log: Logger instance :param workload: The TestConnection workload to execute :param team_conf: Team-specific executor configuration """ - from concurrent.futures import ThreadPoolExecutor, TimeoutError - - import httpx - - from airflow.models.connection_test import run_connection_test + from airflow.sdk.api.client import Client setproctitle( f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}", log, ) - # Build execution API URL (same pattern as _execute_work) base_url = team_conf.get("api", "base_url", fallback="/") if base_url.startswith("/"): base_url = f"http://localhost:8080{base_url}" - api_url = f"{base_url.rstrip('/')}/execution/connection-tests/{workload.connection_test_id}" - headers = {"Authorization": f"Bearer {workload.token}"} + default_execution_api_server = f"{base_url.rstrip('/')}/execution/" + server = team_conf.get("core", "execution_api_server_url", fallback=default_execution_api_server) - def _patch(state: str, result_message: str | None = None) -> None: - payload: dict[str, str] = {"state": state} - if result_message is not None: - payload["result_message"] = result_message - httpx.patch(api_url, json=payload, headers=headers) + client: Client = Client(base_url=server, token=workload.token) - try: - _patch("running") + def _handle_timeout(signum, frame): + raise TimeoutError(f"Connection test timed out after {workload.timeout}s") - with ThreadPoolExecutor(max_workers=1) as pool: - future = pool.submit( - run_connection_test, - connection_id=workload.connection_id, - ) - success, message = future.result(timeout=workload.timeout) + signal.signal(signal.SIGALRM, _handle_timeout) + signal.alarm(workload.timeout) + try: + client.connection_tests.update_state(workload.connection_test_id, ConnectionTestState.RUNNING) + success, message = run_connection_test(connection_id=workload.connection_id) - _patch("success" if success else "failed", message) + state = ConnectionTestState.SUCCESS if success else ConnectionTestState.FAILED + client.connection_tests.update_state(workload.connection_test_id, state, message) except TimeoutError: log.error( "Connection test timed out after %ds", workload.timeout, connection_id=workload.connection_id, ) - _patch("failed", f"Connection test timed out after {workload.timeout}s") + client.connection_tests.update_state( + workload.connection_test_id, + ConnectionTestState.FAILED, + f"Connection test timed out after {workload.timeout}s", + ) except Exception: log.exception("Connection test failed unexpectedly", connection_id=workload.connection_id) try: - _patch("failed", "Connection test failed unexpectedly") + client.connection_tests.update_state( + workload.connection_test_id, + ConnectionTestState.FAILED, + "Connection test failed unexpectedly", + ) except Exception: log.exception("Failed to report connection test failure") + finally: + signal.alarm(0) class LocalExecutor(BaseExecutor): diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 679da4fadc7bd..c0dffd9e45337 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -83,10 +83,7 @@ def __init__(self, *, connection_id: str, queue: str | None = None, **kwargs): self.state = ConnectionTestState.PENDING def __repr__(self) -> str: - return ( - f"" - ) + return f"" def run_connection_test(*, connection_id: str) -> tuple[bool, str]: diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 76cfc8954c057..8263e91492491 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -27,6 +27,7 @@ import pytest import structlog import time_machine +from sqlalchemy.orm import Session from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import CallbackRequest @@ -421,7 +422,7 @@ def test_queue_connection_test_workload_rejected_by_default(): connection_id="test_conn", ) with pytest.raises(ValueError, match="does not support connection testing"): - executor.queue_workload(wl, session=mock.MagicMock()) + executor.queue_workload(wl, session=mock.MagicMock(spec=Session)) def test_queue_connection_test_workload_accepted_when_supported(): @@ -434,7 +435,7 @@ def test_queue_connection_test_workload_accepted_when_supported(): connection_test_id=uuid.uuid4(), connection_id="test_conn", ) - executor.queue_workload(wl, session=mock.MagicMock()) + executor.queue_workload(wl, session=mock.MagicMock(spec=Session)) assert len(executor.queued_connection_tests) == 1 assert executor.queued_connection_tests[0] is wl diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py index af1ce9dc2346a..fb9cb2940846f 100644 --- a/airflow-core/tests/unit/models/test_connection_test.py +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -68,7 +68,7 @@ class TestRunConnectionTest: def test_successful_connection_test(self): """Pure function returns (True, message) on successful test.""" with mock.patch.object( - Connection, "get_connection_from_secrets", return_value=mock.MagicMock() + Connection, "get_connection_from_secrets", return_value=mock.MagicMock(spec=Connection) ) as mock_get_conn: mock_get_conn.return_value.test_connection.return_value = (True, "Connection OK") success, message = run_connection_test(connection_id="test_conn") @@ -79,7 +79,7 @@ def test_successful_connection_test(self): def test_failed_connection_test(self): """Pure function returns (False, message) when test_connection returns False.""" with mock.patch.object( - Connection, "get_connection_from_secrets", return_value=mock.MagicMock() + Connection, "get_connection_from_secrets", return_value=mock.MagicMock(spec=Connection) ) as mock_get_conn: mock_get_conn.return_value.test_connection.return_value = (False, "Connection failed") success, message = run_connection_test(connection_id="test_conn") diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 90374f76be50f..59bf62b0f7729 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -45,6 +45,8 @@ AssetEventsResponse, AssetResponse, ConnectionResponse, + ConnectionTestResultBody, + ConnectionTestState, DagResponse, DagRun, DagRunStateResponse, @@ -851,6 +853,20 @@ 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 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=result_message) + self.client.patch(f"connection-tests/{id}", content=body.model_dump_json()) + + class BearerAuth(httpx.Auth): def __init__(self, token: str): self.token: str = token @@ -1025,6 +1041,12 @@ def hitl(self): """Operations related to HITL Responses.""" return HITLOperations(self) + @lru_cache() # type: ignore[misc] + @property + def connection_tests(self) -> ConnectionTestOperations: + """Operations related to Connection Tests.""" + return ConnectionTestOperations(self) + @lru_cache() # type: ignore[misc] @property def dags(self) -> DagsOperations: From 703ece23fc6a97cb85d4b970fc9cceaa5208790e Mon Sep 17 00:00:00 2001 From: Anish Date: Sun, 1 Mar 2026 18:47:00 -0600 Subject: [PATCH 11/38] Simplify local executor --- .../openapi/v2-rest-api-generated.yaml | 3 + .../core_api/routes/public/connections.py | 1 + .../src/airflow/executors/local_executor.py | 20 +- .../src/airflow/jobs/scheduler_job_runner.py | 16 +- .../0109_3_2_0_add_connection_test_table.py | 2 + .../0110_3_2_0_add_connection_snapshot.py | 52 ----- .../src/airflow/models/connection_test.py | 11 +- airflow-core/src/airflow/utils/db.py | 2 +- .../routes/public/test_connections.py | 13 +- .../versions/head/test_connection_tests.py | 6 + .../v2026_03_31/test_connection_tests.py | 8 + .../unit/executors/test_base_executor.py | 10 +- .../tests/unit/jobs/test_scheduler_job.py | 184 +++++++++++++++++- 13 files changed, 233 insertions(+), 95 deletions(-) delete mode 100644 airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_snapshot.py diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 049fc3854bf7d..67ba27995d551 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -1819,6 +1819,9 @@ paths: \ of the token serves as authorization \u2014 only the client\nthat initiated\ \ the test knows the crypto-random token." operationId: get_connection_test_status + security: + - OAuth2PasswordBearer: [] + - HTTPBearer: [] parameters: - name: connection_test_token in: path diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index adad474898f09..1211bed576b55 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -370,6 +370,7 @@ def test_connection_async( @connections_router.get( "/test-async/{connection_test_token}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_connection(method="GET"))], ) def get_connection_test_status( connection_test_token: str, diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 32eab13d96f47..99631aa86a330 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -120,7 +120,10 @@ def _run_worker( output.put((workload.callback.id, CallbackState.FAILED, e)) elif isinstance(workload, workloads.TestConnection): - _execute_connection_test(log, workload, team_conf) + try: + _execute_connection_test(log, workload, team_conf) + except Exception: + log.exception("Connection test failed") else: raise ValueError(f"LocalExecutor does not know how to handle {type(workload)}") @@ -222,14 +225,11 @@ def _handle_timeout(signum, frame): ) except Exception: log.exception("Connection test failed unexpectedly", connection_id=workload.connection_id) - try: - client.connection_tests.update_state( - workload.connection_test_id, - ConnectionTestState.FAILED, - "Connection test failed unexpectedly", - ) - except Exception: - log.exception("Failed to report connection test failure") + client.connection_tests.update_state( + workload.connection_test_id, + ConnectionTestState.FAILED, + "Connection test failed unexpectedly", + ) finally: signal.alarm(0) @@ -403,6 +403,8 @@ def _process_workloads(self, workload_list): del self.queued_tasks[workload.ti.key] elif isinstance(workload, workloads.ExecuteCallback): del self.queued_callbacks[workload.callback.id] + elif isinstance(workload, workloads.TestConnection): + pass # Already removed from queued_connection_tests by base class with self._unread_messages: self._unread_messages.value += len(workload_list) self._check_workers() diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 06cbbf34ce36f..328e87d229403 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3180,15 +3180,14 @@ def _dispatch_connection_tests(self, session: Session = NEW_SESSION) -> None: 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 = 30 + grace_period = max(30, timeout // 2) cutoff = timezone.utcnow() - timedelta(seconds=timeout + grace_period) - stale_tests = session.scalars( - select(ConnectionTest).where( - ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]), - ConnectionTest.updated_at < cutoff, - ) - ).all() + stale_stmt = select(ConnectionTest).where( + ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]), + ConnectionTest.updated_at < cutoff, + ) + stale_tests = session.scalars(stale_stmt).all() for ct in stale_tests: ct.state = ConnectionTestState.FAILED @@ -3205,8 +3204,9 @@ def _find_executor_for_connection_test(self, queue: str | None) -> BaseExecutor for executor in self.executors: if executor.supports_connection_test and executor.team_name == queue: return executor + return None for executor in self.executors: - if executor.supports_connection_test and (queue is None or executor.team_name is None): + if executor.supports_connection_test: return executor return None diff --git a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py index 82710679aa30b..34e35a14e81f1 100644 --- a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py @@ -52,6 +52,8 @@ def upgrade(): sa.Column("created_at", UtcDateTime(timezone=True), nullable=False), sa.Column("updated_at", UtcDateTime(timezone=True), nullable=False), sa.Column("queue", sa.String(256), nullable=True), + sa.Column("connection_snapshot", sa.JSON(), nullable=True), + sa.Column("reverted", sa.Boolean(), nullable=False, server_default="false"), sa.PrimaryKeyConstraint("id", name=op.f("connection_test_pkey")), sa.UniqueConstraint("token", name=op.f("connection_test_token_uq")), ) diff --git a/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_snapshot.py b/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_snapshot.py deleted file mode 100644 index 2e4f0f45919c2..0000000000000 --- a/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_snapshot.py +++ /dev/null @@ -1,52 +0,0 @@ -# -# 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_snapshot and reverted columns to connection_test. - -Revision ID: b8f3e5d1a9c2 -Revises: a7e6d4c3b2f1 -Create Date: 2026-02-27 00:00:00.000000 - -""" - -from __future__ import annotations - -import sqlalchemy as sa -from alembic import op - -# revision identifiers, used by Alembic. -revision = "b8f3e5d1a9c2" -down_revision = "a7e6d4c3b2f1" -branch_labels = None -depends_on = None -airflow_version = "3.2.0" - - -def upgrade(): - """Add connection_snapshot and reverted columns to connection_test.""" - with op.batch_alter_table("connection_test") as batch_op: - batch_op.add_column(sa.Column("connection_snapshot", sa.JSON(), nullable=True)) - batch_op.add_column(sa.Column("reverted", sa.Boolean(), nullable=False, server_default=sa.false())) - - -def downgrade(): - """Remove connection_snapshot and reverted columns from connection_test.""" - with op.batch_alter_table("connection_test") as batch_op: - batch_op.drop_column("reverted") - batch_op.drop_column("connection_snapshot") diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index c0dffd9e45337..83511df9ceb62 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -179,7 +179,7 @@ def _can_safely_revert(conn: Connection, post_snapshot: dict) -> bool: def attempt_revert(ct: ConnectionTest, *, session: Session) -> None: """Revert a connection to its pre-edit values if no concurrent edit has occurred.""" if not ct.connection_snapshot: - log.warning("attempt_revert called without snapshot for %s", ct.id) + log.warning("attempt_revert called without snapshot", connection_test_id=ct.id) return pre_snapshot = ct.connection_snapshot["pre"] @@ -188,19 +188,16 @@ def attempt_revert(ct: ConnectionTest, *, session: Session) -> None: conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id)) if conn is None: ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists." - log.warning("Revert skipped: connection %s no longer exists", ct.connection_id) + log.warning("Revert skipped: connection no longer exists", connection_id=ct.connection_id) return if not _can_safely_revert(conn, post_snapshot): ct.result_message = ( ct.result_message or "" ) + " | Revert skipped: connection was modified by another user." - log.warning( - "Revert skipped: concurrent edit detected on connection %s", - ct.connection_id, - ) + log.warning("Revert skipped: concurrent edit detected", connection_id=ct.connection_id) return _revert_connection(conn, pre_snapshot) ct.reverted = True - log.info("Reverted connection %s to pre-edit state", ct.connection_id) + log.info("Reverted connection to pre-edit state", connection_id=ct.connection_id) diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index e4e33efbe4358..5479aecf1d20b 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/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": "b8f3e5d1a9c2", + "3.2.0": "a7e6d4c3b2f1", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 316f5542f5493..cca4c3cc78cb0 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -29,13 +29,18 @@ from airflow.api_fastapi.core_api.datamodels.connections import ConnectionBody from airflow.api_fastapi.core_api.services.public.connections import BulkConnectionService from airflow.models import Connection -from airflow.models.connection_test import ConnectionTest, ConnectionTestState +from airflow.models.connection_test import ConnectionTest, ConnectionTestState, attempt_revert from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.utils.session import NEW_SESSION, provide_session from tests_common.test_utils.api_fastapi import _check_last_log from tests_common.test_utils.asserts import assert_queries_count -from tests_common.test_utils.db import clear_db_connections, clear_db_logs, clear_test_connections +from tests_common.test_utils.db import ( + clear_db_connection_tests, + clear_db_connections, + clear_db_logs, + clear_test_connections, +) from tests_common.test_utils.markers import skip_if_force_lowest_dependencies_marker pytestmark = pytest.mark.db_test @@ -95,10 +100,12 @@ class TestConnectionEndpoint: def setup(self) -> None: clear_test_connections(False) clear_db_connections(False) + clear_db_connection_tests() clear_db_logs() def teardown_method(self) -> None: clear_db_connections() + clear_db_connection_tests() def create_connection(self, team_name: str | None = None): _create_connection(team_name=team_name) @@ -1352,8 +1359,6 @@ def test_poll_shows_reverted_true_after_failed_test(self, test_client, session): ct.state = ConnectionTestState.FAILED ct.result_message = "Connection refused" - from airflow.models.connection_test import attempt_revert - attempt_revert(ct, session=session) session.commit() diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py index a1c406224f934..564de6fa1b27f 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py @@ -28,6 +28,12 @@ class TestPatchConnectionTest: + @pytest.fixture(autouse=True) + def setup_teardown(self): + clear_db_connection_tests() + yield + clear_db_connection_tests() + def test_patch_updates_result(self, client, session): """PATCH sets the state and result fields.""" ct = ConnectionTest(connection_id="test_conn") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py index d5203253c3081..f7f8aa44769c8 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py @@ -20,6 +20,8 @@ from airflow.models.connection_test import ConnectionTest, ConnectionTestState +from tests_common.test_utils.db import clear_db_connection_tests + pytestmark = pytest.mark.db_test @@ -33,6 +35,12 @@ def old_ver_client(client): class TestConnectionTestEndpointVersioning: """Test that the connection-tests endpoint didn't exist in older API versions.""" + @pytest.fixture(autouse=True) + def setup_teardown(self): + clear_db_connection_tests() + yield + clear_db_connection_tests() + def test_old_version_returns_404(self, old_ver_client, session): """PATCH /connection-tests/{id} should not exist in older API versions.""" ct = ConnectionTest(connection_id="test_conn") diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 8263e91492491..4d0085a2eee34 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -21,7 +21,7 @@ import textwrap from datetime import timedelta from unittest import mock -from uuid import UUID +from uuid import UUID, uuid4 import pendulum import pytest @@ -414,11 +414,9 @@ def test_supports_connection_test_default_value(): def test_queue_connection_test_workload_rejected_by_default(): """BaseExecutor (supports_connection_test=False) rejects TestConnection workloads.""" - import uuid - executor = BaseExecutor() wl = workloads.TestConnection.make( - connection_test_id=uuid.uuid4(), + connection_test_id=uuid4(), connection_id="test_conn", ) with pytest.raises(ValueError, match="does not support connection testing"): @@ -427,12 +425,10 @@ def test_queue_connection_test_workload_rejected_by_default(): def test_queue_connection_test_workload_accepted_when_supported(): """An executor with supports_connection_test=True accepts TestConnection workloads.""" - import uuid - executor = LocalExecutor() executor.queued_connection_tests.clear() wl = workloads.TestConnection.make( - connection_test_id=uuid.uuid4(), + connection_test_id=uuid4(), connection_id="test_conn", ) executor.queue_workload(wl, session=mock.MagicMock(spec=Session)) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 06b90355eae98..d600f312eac98 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9466,7 +9466,9 @@ def scheduler_job_runner_for_connection_tests(session): runner.executors = [executor] runner.executor = executor runner._log = mock.MagicMock(spec=logging.Logger) - return runner + yield runner + session.execute(delete(ConnectionTest)) + session.commit() class TestDispatchConnectionTests: @@ -9548,10 +9550,10 @@ def test_dispatch_fails_fast_when_no_executor_supports( "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", }, ) - def test_dispatch_with_queue_falls_back_to_global_executor( + def test_dispatch_with_unmatched_queue_fails_fast( self, scheduler_job_runner_for_connection_tests, session ): - """Tests with a queue are dispatched to the global executor as fallback.""" + """Tests requesting a queue with no matching executor are failed immediately.""" ct = ConnectionTest(connection_id="test_conn", queue="gpu_workers") session.add(ct) session.commit() @@ -9560,8 +9562,137 @@ def test_dispatch_with_queue_falls_back_to_global_executor( session.expire_all() ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.QUEUED - assert len(scheduler_job_runner_for_connection_tests.executor.queued_connection_tests) == 1 + assert ct.state == ConnectionTestState.FAILED + assert "gpu_workers" in ct.result_message + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "3", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_budget_dispatches_up_to_remaining_slots( + self, scheduler_job_runner_for_connection_tests, session + ): + """When 1 slot is occupied, only budget (cap - active) pending tests are dispatched.""" + ct_active = ConnectionTest(connection_id="active_conn") + ct_active.state = ConnectionTestState.RUNNING + session.add(ct_active) + + pending_tests = [] + for i in range(3): + ct = ConnectionTest(connection_id=f"pending_{i}") + session.add(ct) + pending_tests.append(ct) + session.commit() + pending_ids = [ct.id for ct in pending_tests] + + scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + + session.expire_all() + states = [session.get(ConnectionTest, pid).state for pid in pending_ids] + assert states.count(ConnectionTestState.QUEUED) == 2 + assert states.count(ConnectionTestState.PENDING) == 1 + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "2", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_order_is_fifo_by_created_at(self, scheduler_job_runner_for_connection_tests, session): + """Pending tests are dispatched in FIFO order based on created_at.""" + initial_time = timezone.utcnow() + + with time_machine.travel(initial_time - timedelta(minutes=5), tick=False): + ct_old = ConnectionTest(connection_id="old_conn") + session.add(ct_old) + session.flush() + + with time_machine.travel(initial_time, tick=False): + ct_new = ConnectionTest(connection_id="new_conn") + session.add(ct_new) + session.flush() + + with time_machine.travel(initial_time + timedelta(minutes=1), tick=False): + ct_newest = ConnectionTest(connection_id="newest_conn") + session.add(ct_newest) + session.flush() + + session.commit() + + scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + + session.expire_all() + assert session.get(ConnectionTest, ct_old.id).state == ConnectionTestState.QUEUED + assert session.get(ConnectionTest, ct_new.id).state == ConnectionTestState.QUEUED + assert session.get(ConnectionTest, ct_newest.id).state == ConnectionTestState.PENDING + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_fails_fast_for_unserved_queue(self, scheduler_job_runner_for_connection_tests, session): + """Tests requesting a queue no executor serves are failed immediately.""" + with mock.patch.object( + scheduler_job_runner_for_connection_tests, + "_find_executor_for_connection_test", + return_value=None, + ): + ct = ConnectionTest(connection_id="test_conn", queue="nonexistent_queue") + session.add(ct) + session.commit() + + scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert "nonexistent_queue" in ct.result_message + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_queue_matched_to_team_name(self, session): + """When queue is specified, executor whose team_name matches is selected.""" + session.execute(delete(ConnectionTest)) + session.commit() + + mock_job = mock.MagicMock(spec=Job) + mock_job.id = 1 + mock_job.max_tis_per_query = 16 + + executor_a = LocalExecutor() + executor_a.team_name = "team_a" + executor_a.queued_connection_tests.clear() + + executor_b = LocalExecutor() + executor_b.team_name = "team_b" + executor_b.queued_connection_tests.clear() + + runner = SchedulerJobRunner.__new__(SchedulerJobRunner) + runner.job = mock_job + runner.executors = [executor_a, executor_b] + runner.executor = executor_a + runner._log = mock.MagicMock(spec=logging.Logger) + + ct = ConnectionTest(connection_id="team_conn", queue="team_b") + session.add(ct) + session.commit() + + runner._dispatch_connection_tests(session=session) + + assert len(executor_b.queued_connection_tests) == 1 + assert len(executor_a.queued_connection_tests) == 0 class TestReapStaleConnectionTests: @@ -9636,8 +9767,6 @@ def test_reaper_reverts_connection_on_timeout_with_snapshot( conn = session.scalar(select(Connection).filter_by(conn_id="reaper_conn")) assert conn.host == "old-host.example.com" - clear_db_connections(add_default_connections_back=False) - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) def test_reaper_no_revert_without_snapshot(self, scheduler_job_runner_for_connection_tests, session): """Stale tests without a snapshot do not trigger revert.""" @@ -9655,3 +9784,44 @@ def test_reaper_no_revert_without_snapshot(self, scheduler_job_runner_for_connec ct = session.get(ConnectionTest, ct.id) assert ct.state == ConnectionTestState.FAILED assert ct.reverted is False + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) + def test_reap_stale_running_test(self, scheduler_job_runner_for_connection_tests, session): + """Stale RUNNING tests are also reaped by the reaper.""" + initial_time = timezone.utcnow() + with time_machine.travel(initial_time, tick=False): + ct = ConnectionTest(connection_id="running_conn") + ct.state = ConnectionTestState.RUNNING + session.add(ct) + session.commit() + + with time_machine.travel(initial_time + timedelta(seconds=200), tick=False): + scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert "timed out" in ct.result_message + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) + def test_reaper_ignores_terminal_states(self, scheduler_job_runner_for_connection_tests, session): + """Tests in terminal states (SUCCESS, FAILED) are not touched by the reaper.""" + initial_time = timezone.utcnow() + with time_machine.travel(initial_time, tick=False): + ct_success = ConnectionTest(connection_id="success_conn") + ct_success.state = ConnectionTestState.SUCCESS + ct_success.result_message = "OK" + session.add(ct_success) + + ct_failed = ConnectionTest(connection_id="failed_conn") + ct_failed.state = ConnectionTestState.FAILED + ct_failed.result_message = "Error" + session.add(ct_failed) + session.commit() + + with time_machine.travel(initial_time + timedelta(seconds=200), tick=False): + scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) + + session.expire_all() + assert session.get(ConnectionTest, ct_success.id).state == ConnectionTestState.SUCCESS + assert session.get(ConnectionTest, ct_failed.id).state == ConnectionTestState.FAILED From ddd02535a599bc253f0122e443faa4ca1ae5b19d Mon Sep 17 00:00:00 2001 From: Anish Date: Sun, 1 Mar 2026 20:26:21 -0600 Subject: [PATCH 12/38] further simplify local execter --- .../src/airflow/executors/base_executor.py | 20 +++++-------------- .../src/airflow/executors/local_executor.py | 4 +--- .../unit/executors/test_base_executor.py | 2 +- 3 files changed, 7 insertions(+), 19 deletions(-) diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 00a6e0e151501..aa2fff0fcd331 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -187,7 +187,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: deque[workloads.TestConnection] = deque() + self.queued_connection_tests: dict[str, workloads.TestConnection] = {} self.running: set[WorkloadKey] = set() self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {} self._task_event_logs: deque[Log] = deque() @@ -236,7 +236,7 @@ def queue_workload(self, workload: workloads.All, session: Session) -> None: elif isinstance(workload, workloads.TestConnection): if not self.supports_connection_test: raise ValueError(f"Executor {type(self).__name__} does not support connection testing") - self.queued_connection_tests.append(workload) + self.queued_connection_tests[str(workload.connection_test_id)] = workload else: raise ValueError( f"Un-handled workload type {type(workload).__name__!r} in {type(self).__name__}. " @@ -318,22 +318,12 @@ def heartbeat(self) -> None: self.log.debug("Calling the %s sync method", self.__class__) self.sync() - def trigger_connection_tests(self, max_tests: int | None = None) -> None: - """ - Process queued connection tests. - - :param max_tests: Maximum number of tests to trigger. Defaults to all queued. - """ + def trigger_connection_tests(self) -> None: + """Process queued connection tests.""" if not self.queued_connection_tests: return - count = max_tests if max_tests is not None else len(self.queued_connection_tests) - test_workloads: list[workloads.TestConnection] = [] - for _ in range(min(count, len(self.queued_connection_tests))): - test_workloads.append(self.queued_connection_tests.popleft()) - - if test_workloads: - self._process_workloads(test_workloads) + self._process_workloads(list(self.queued_connection_tests.values())) def _get_metric_name(self, metric_base_name: str) -> str: return ( diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 99631aa86a330..2640588e75f3f 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -180,8 +180,6 @@ def _execute_connection_test(log: Logger, workload: workloads.TestConnection, te """ Execute a connection test workload. - Results are reported back via the Execution API. - :param log: Logger instance :param workload: The TestConnection workload to execute :param team_conf: Team-specific executor configuration @@ -404,7 +402,7 @@ def _process_workloads(self, workload_list): elif isinstance(workload, workloads.ExecuteCallback): del self.queued_callbacks[workload.callback.id] elif isinstance(workload, workloads.TestConnection): - pass # Already removed from queued_connection_tests by base class + del self.queued_connection_tests[str(workload.connection_test_id)] with self._unread_messages: self._unread_messages.value += len(workload_list) self._check_workers() diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 4d0085a2eee34..9f92d783eaca6 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -433,7 +433,7 @@ def test_queue_connection_test_workload_accepted_when_supported(): ) executor.queue_workload(wl, session=mock.MagicMock(spec=Session)) assert len(executor.queued_connection_tests) == 1 - assert executor.queued_connection_tests[0] is wl + assert executor.queued_connection_tests[str(wl.connection_test_id)] is wl @mock.patch.dict("os.environ", {}, clear=True) From 9fe3946f728eb22c4dbb07e96110a77456c935ce Mon Sep 17 00:00:00 2001 From: Anish Date: Sun, 1 Mar 2026 22:20:27 -0600 Subject: [PATCH 13/38] add queue parmeter --- .../openapi/v2-rest-api-generated.yaml | 10 ++++++++++ .../core_api/routes/public/connections.py | 3 ++- .../airflow/ui/openapi-gen/queries/queries.ts | 5 ++++- .../ui/openapi-gen/requests/services.gen.ts | 4 +++- .../ui/openapi-gen/requests/types.gen.ts | 4 ++++ .../routes/public/test_connections.py | 19 +++++++++++++++++++ 6 files changed, 42 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 67ba27995d551..b6013cc1855b4 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -1665,6 +1665,16 @@ paths: type: string - type: 'null' title: Update Mask + - name: queue + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + description: Executor queue to route the connection test to + title: Queue + description: Executor queue to route the connection test to requestBody: required: true content: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 1211bed576b55..b5006dc752f93 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -251,6 +251,7 @@ def patch_connection_and_test( patch_body: ConnectionBody, session: SessionDep, update_mask: list[str] | None = Query(None), + queue: str | None = Query(None, description="Executor queue to route the connection test to"), ) -> ConnectionSaveAndTestResponse: """ Update a connection and queue an async test with revert-on-failure. @@ -285,7 +286,7 @@ def patch_connection_and_test( post_snapshot = snapshot_connection(connection) - connection_test = ConnectionTest(connection_id=connection_id) + connection_test = ConnectionTest(connection_id=connection_id, queue=queue) connection_test.connection_snapshot = {"pre": pre_snapshot, "post": post_snapshot} session.add(connection_test) session.flush() diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index 871e937c5eaeb..ed4ddc49ebf6f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -2128,18 +2128,21 @@ export const useConnectionServiceBulkConnections = (options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ connectionId, requestBody, updateMask }) => ConnectionService.patchConnectionAndTest({ connectionId, requestBody, updateMask }) as unknown as Promise, ...options }); +}, TContext>({ mutationFn: ({ connectionId, queue, requestBody, updateMask }) => ConnectionService.patchConnectionAndTest({ connectionId, queue, requestBody, updateMask }) as unknown as Promise, ...options }); /** * Patch Dag Run * Modify a DAG Run. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 4540c4b2555b2..7e0b0da932471 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -779,6 +779,7 @@ export class ConnectionService { * @param data.connectionId * @param data.requestBody * @param data.updateMask + * @param data.queue Executor queue to route the connection test to * @returns ConnectionSaveAndTestResponse Successful Response * @throws ApiError */ @@ -790,7 +791,8 @@ export class ConnectionService { connection_id: data.connectionId }, query: { - update_mask: data.updateMask + update_mask: data.updateMask, + queue: data.queue }, body: data.requestBody, mediaType: 'application/json', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index db434f1bf9160..f89d1c9d35d12 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2516,6 +2516,10 @@ export type BulkConnectionsResponse = BulkResponse; export type PatchConnectionAndTestData = { connectionId: string; + /** + * Executor queue to route the connection test to + */ + queue?: string | null; requestBody: ConnectionBody; updateMask?: Array<(string)> | null; }; diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index cca4c3cc78cb0..894d16b739e58 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -1317,6 +1317,25 @@ def test_save_and_test_creates_snapshot(self, test_client, session): assert snapshot["pre"]["host"] == TEST_CONN_HOST assert snapshot["post"]["host"] == "new-host.example.com" + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_save_and_test_passes_queue_parameter(self, test_client, session): + """PATCH save-and-test passes the queue parameter to the ConnectionTest.""" + self.create_connection() + response = test_client.patch( + f"/connections/{TEST_CONN_ID}/save-and-test?queue=team_a", + json={ + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "host": "queued-host.example.com", + }, + ) + assert response.status_code == 200 + token = response.json()["test_token"] + + ct = session.scalar(select(ConnectionTest).filter_by(token=token)) + assert ct is not None + assert ct.queue == "team_a" + def test_save_and_test_403_when_disabled(self, test_client): """PATCH save-and-test returns 403 when test_connection is disabled.""" self.create_connection() From e3c4421e0cdd06ac15d939f8a05aba73d2a87620 Mon Sep 17 00:00:00 2001 From: Anish Date: Sun, 1 Mar 2026 23:31:16 -0600 Subject: [PATCH 14/38] fix ci test failiures --- airflow-core/src/airflow/models/connection_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 83511df9ceb62..b7e23bda3beb3 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -71,7 +71,7 @@ class ConnectionTest(Base): ) queue: Mapped[str | None] = mapped_column(String(256), nullable=True) connection_snapshot: Mapped[dict | None] = mapped_column(JSON, nullable=True) - reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False) + reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="false") __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),) From 82576e75258321c6dfc8da61ddb2fc241e847a1f Mon Sep 17 00:00:00 2001 From: Anish Date: Sun, 1 Mar 2026 23:55:14 -0600 Subject: [PATCH 15/38] changed reverted default server value --- .../migrations/versions/0109_3_2_0_add_connection_test_table.py | 2 +- airflow-core/src/airflow/models/connection_test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py index 34e35a14e81f1..cc870be041af2 100644 --- a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py @@ -53,7 +53,7 @@ def upgrade(): sa.Column("updated_at", UtcDateTime(timezone=True), nullable=False), sa.Column("queue", sa.String(256), nullable=True), sa.Column("connection_snapshot", sa.JSON(), nullable=True), - sa.Column("reverted", sa.Boolean(), nullable=False, server_default="false"), + sa.Column("reverted", sa.Boolean(), nullable=False, server_default="0"), sa.PrimaryKeyConstraint("id", name=op.f("connection_test_pkey")), sa.UniqueConstraint("token", name=op.f("connection_test_token_uq")), ) diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index b7e23bda3beb3..745e31f292b6b 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -71,7 +71,7 @@ class ConnectionTest(Base): ) queue: Mapped[str | None] = mapped_column(String(256), nullable=True) connection_snapshot: Mapped[dict | None] = mapped_column(JSON, nullable=True) - reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="false") + reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0") __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),) From b4f702c8176a4a3a4af3aaa0110a40e11c9853c4 Mon Sep 17 00:00:00 2001 From: Anish Date: Mon, 2 Mar 2026 22:07:33 -0600 Subject: [PATCH 16/38] refactor: move connection test execution logic from local_executor to workload module --- .../src/airflow/executors/local_executor.py | 27 +- .../src/airflow/models/connection_test.py | 7 +- .../unit/executors/test_local_executor.py | 244 +++++++++++++++++- .../tests/unit/models/test_connection_test.py | 37 ++- 4 files changed, 289 insertions(+), 26 deletions(-) diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 2640588e75f3f..8f89e9d32209e 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -39,6 +39,7 @@ from airflow.executors import workloads from airflow.executors.base_executor import BaseExecutor from airflow.executors.workloads.callback import execute_callback_workload +from airflow.models.connection import Connection from airflow.models.connection_test import ConnectionTestState, run_connection_test from airflow.utils.state import CallbackState, TaskInstanceState @@ -180,11 +181,18 @@ def _execute_connection_test(log: Logger, workload: workloads.TestConnection, te """ Execute a connection test workload. + Constructs an SDK ``Client``, fetches the connection via the Execution API, + enforces a timeout via ``signal.alarm``, and reports all outcomes back + through the Execution API. + :param log: Logger instance :param workload: The TestConnection workload to execute :param team_conf: Team-specific executor configuration """ + # Lazy import: SDK modules must not be loaded at module level to avoid + # coupling core (scheduler-loaded) code to the SDK. from airflow.sdk.api.client import Client + from airflow.sdk.execution_time.comms import ErrorResponse setproctitle( f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}", @@ -197,7 +205,7 @@ def _execute_connection_test(log: Logger, workload: workloads.TestConnection, te default_execution_api_server = f"{base_url.rstrip('/')}/execution/" server = team_conf.get("core", "execution_api_server_url", fallback=default_execution_api_server) - client: Client = Client(base_url=server, token=workload.token) + client = Client(base_url=server, token=workload.token) def _handle_timeout(signum, frame): raise TimeoutError(f"Connection test timed out after {workload.timeout}s") @@ -206,7 +214,22 @@ def _handle_timeout(signum, frame): signal.alarm(workload.timeout) try: client.connection_tests.update_state(workload.connection_test_id, ConnectionTestState.RUNNING) - success, message = run_connection_test(connection_id=workload.connection_id) + + conn_response = client.connections.get(workload.connection_id) + if isinstance(conn_response, ErrorResponse): + raise RuntimeError(f"Connection '{workload.connection_id}' not found via Execution API") + + conn = Connection( + conn_id=conn_response.conn_id, + conn_type=conn_response.conn_type, + host=conn_response.host, + login=conn_response.login, + password=conn_response.password, + schema=conn_response.schema_, + port=conn_response.port, + extra=conn_response.extra, + ) + success, message = run_connection_test(conn=conn) state = ConnectionTestState.SUCCESS if success else ConnectionTestState.FAILED client.connection_tests.update_state(workload.connection_test_id, state, message) diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 745e31f292b6b..6c370ba5f5361 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -86,18 +86,17 @@ def __repr__(self) -> str: return f"" -def run_connection_test(*, connection_id: str) -> tuple[bool, str]: +def run_connection_test(*, conn: Connection) -> tuple[bool, str]: """ - Worker-side pure function to execute a connection test. + Worker-side function to execute a connection test. Returns a (success, message) tuple. The caller is responsible for reporting the result back via the Execution API. """ try: - conn = Connection.get_connection_from_secrets(connection_id) return conn.test_connection() except Exception as e: - log.exception("Connection test failed", connection_id=connection_id) + log.exception("Connection test failed", connection_id=conn.conn_id) return False, str(e) diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index 59afffe6833fe..a194ad5e6d3d2 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -23,16 +23,21 @@ from unittest import mock import pytest +import structlog from kgb import spy_on from uuid6 import uuid7 from airflow._shared.timezones import timezone from airflow.executors import workloads -from airflow.executors.local_executor import LocalExecutor, _execute_work +from airflow.executors.base_executor import ExecutorConf +from airflow.executors.local_executor import LocalExecutor, _execute_connection_test, _execute_work from airflow.executors.workloads.base import BundleInfo from airflow.executors.workloads.callback import CallbackDTO from airflow.executors.workloads.task import TaskInstanceDTO from airflow.models.callback import CallbackFetchMethod +from airflow.models.connection_test import ConnectionTestState +from airflow.sdk.api.datamodels._generated import ConnectionResponse +from airflow.sdk.execution_time.comms import ErrorResponse from airflow.settings import Session from airflow.utils.state import State @@ -376,6 +381,243 @@ def test_global_executor_without_team_name(self): executor.end() +class TestLocalExecutorConnectionTestSupport: + def test_supports_connection_test_flag_is_true(self): + executor = LocalExecutor() + assert executor.supports_connection_test is True + + +@mock.patch("airflow.executors.local_executor.signal", autospec=True) +@mock.patch("airflow.sdk.api.client.Client", autospec=True) +class TestLocalExecutorConnectionTestExecution: + def test_successful_connection_test(self, MockClient, _mock_signal): + """Fetches connection via Execution API, runs test, reports SUCCESS.""" + mock_client = MockClient.return_value + mock_client.connections.get.return_value = ConnectionResponse( + conn_id="test_conn", + conn_type="http", + host="httpbin.org", + port=443, + ) + + test_id = uuid7() + workload = workloads.TestConnection( + connection_test_id=test_id, + connection_id="test_conn", + timeout=60, + token="test-token", + ) + + with mock.patch( + "airflow.executors.local_executor.run_connection_test", + return_value=(True, "Connection OK"), + ): + _execute_connection_test( + mock.MagicMock(spec=structlog.typing.FilteringBoundLogger), + workload, + ExecutorConf(team_name=None), + ) + + calls = mock_client.connection_tests.update_state.call_args_list + assert len(calls) == 2 + assert calls[0].args == (test_id, ConnectionTestState.RUNNING) + assert calls[1].args == (test_id, ConnectionTestState.SUCCESS, "Connection OK") + + def test_failed_connection_test(self, MockClient, _mock_signal): + """Fetches connection via Execution API, test fails, reports FAILED.""" + mock_client = MockClient.return_value + mock_client.connections.get.return_value = ConnectionResponse( + conn_id="test_conn", + conn_type="postgres", + host="db.example.com", + ) + + test_id = uuid7() + workload = workloads.TestConnection( + connection_test_id=test_id, + connection_id="test_conn", + timeout=60, + token="test-token", + ) + + with mock.patch( + "airflow.executors.local_executor.run_connection_test", + return_value=(False, "Connection refused"), + ): + _execute_connection_test( + mock.MagicMock(spec=structlog.typing.FilteringBoundLogger), + workload, + ExecutorConf(team_name=None), + ) + + calls = mock_client.connection_tests.update_state.call_args_list + assert len(calls) == 2 + assert calls[0].args == (test_id, ConnectionTestState.RUNNING) + assert calls[1].args == (test_id, ConnectionTestState.FAILED, "Connection refused") + + def test_connection_not_found_via_execution_api(self, MockClient, _mock_signal): + """Reports FAILED when connection is not found via Execution API.""" + mock_client = MockClient.return_value + mock_client.connections.get.return_value = ErrorResponse(detail={"conn_id": "missing_conn"}) + + test_id = uuid7() + workload = workloads.TestConnection( + connection_test_id=test_id, + connection_id="missing_conn", + timeout=60, + token="test-token", + ) + + _execute_connection_test( + mock.MagicMock(spec=structlog.typing.FilteringBoundLogger), + workload, + ExecutorConf(team_name=None), + ) + + calls = mock_client.connection_tests.update_state.call_args_list + assert calls[-1].args[1] == ConnectionTestState.FAILED + assert calls[-1].args[2] == "Connection test failed unexpectedly" + + def test_unexpected_exception_reports_failed(self, MockClient, _mock_signal): + """Reports FAILED when an unexpected exception occurs.""" + mock_client = MockClient.return_value + mock_client.connections.get.return_value = ConnectionResponse( + conn_id="test_conn", + conn_type="http", + ) + + test_id = uuid7() + workload = workloads.TestConnection( + connection_test_id=test_id, + connection_id="test_conn", + timeout=60, + token="test-token", + ) + + with mock.patch( + "airflow.executors.local_executor.run_connection_test", + side_effect=RuntimeError("Something broke"), + ): + _execute_connection_test( + mock.MagicMock(spec=structlog.typing.FilteringBoundLogger), + workload, + ExecutorConf(team_name=None), + ) + + calls = mock_client.connection_tests.update_state.call_args_list + assert calls[-1].args == (test_id, ConnectionTestState.FAILED, "Connection test failed unexpectedly") + + def test_connection_fields_passed_correctly(self, MockClient, _mock_signal): + """Verifies all connection fields from the API response are passed to Connection.""" + mock_client = MockClient.return_value + mock_client.connections.get.return_value = ConnectionResponse( + conn_id="full_conn", + conn_type="postgres", + host="db.example.com", + login="admin", + password="s3cret", + schema="mydb", + port=5432, + extra='{"sslmode": "require"}', + ) + + workload = workloads.TestConnection( + connection_test_id=uuid7(), + connection_id="full_conn", + timeout=60, + token="test-token", + ) + + captured_conn = None + + def capture_conn(*, conn): + nonlocal captured_conn + captured_conn = conn + return True, "OK" + + with mock.patch( + "airflow.executors.local_executor.run_connection_test", + side_effect=capture_conn, + ): + _execute_connection_test( + mock.MagicMock(spec=structlog.typing.FilteringBoundLogger), + workload, + ExecutorConf(team_name=None), + ) + + assert captured_conn is not None + assert captured_conn.conn_id == "full_conn" + assert captured_conn.conn_type == "postgres" + assert captured_conn.host == "db.example.com" + assert captured_conn.login == "admin" + assert captured_conn.password == "s3cret" + assert captured_conn.schema == "mydb" + assert captured_conn.port == 5432 + assert captured_conn.extra == '{"sslmode": "require"}' + + def test_timeout_reports_failed(self, MockClient, _mock_signal): + """Reports FAILED with timeout message when TimeoutError is raised.""" + mock_client = MockClient.return_value + mock_client.connections.get.return_value = ConnectionResponse( + conn_id="test_conn", + conn_type="http", + ) + + test_id = uuid7() + workload = workloads.TestConnection( + connection_test_id=test_id, + connection_id="test_conn", + timeout=30, + token="test-token", + ) + + def raise_timeout(*, conn): + raise TimeoutError("Connection test timed out after 30s") + + with mock.patch( + "airflow.executors.local_executor.run_connection_test", + side_effect=raise_timeout, + ): + _execute_connection_test( + mock.MagicMock(spec=structlog.typing.FilteringBoundLogger), + workload, + ExecutorConf(team_name=None), + ) + + calls = mock_client.connection_tests.update_state.call_args_list + assert calls[-1].args[1] == ConnectionTestState.FAILED + assert "timed out" in calls[-1].args[2] + + def test_alarm_is_cancelled_in_finally(self, MockClient, mock_signal): + """signal.alarm(0) is called to cancel the timer even on success.""" + mock_client = MockClient.return_value + mock_client.connections.get.return_value = ConnectionResponse( + conn_id="test_conn", + conn_type="http", + ) + + workload = workloads.TestConnection( + connection_test_id=uuid7(), + connection_id="test_conn", + timeout=60, + token="test-token", + ) + + with mock.patch( + "airflow.executors.local_executor.run_connection_test", + return_value=(True, "OK"), + ): + _execute_connection_test( + mock.MagicMock(spec=structlog.typing.FilteringBoundLogger), + workload, + ExecutorConf(team_name=None), + ) + + alarm_calls = mock_signal.alarm.call_args_list + assert alarm_calls[0].args == (60,) + assert alarm_calls[-1].args == (0,) + + class TestLocalExecutorCallbackSupport: def test_supports_callbacks_flag_is_true(self): executor = LocalExecutor() diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py index fb9cb2940846f..d12f955055cdb 100644 --- a/airflow-core/tests/unit/models/test_connection_test.py +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -66,35 +66,34 @@ def test_queue_defaults_to_none(self): class TestRunConnectionTest: def test_successful_connection_test(self): - """Pure function returns (True, message) on successful test.""" - with mock.patch.object( - Connection, "get_connection_from_secrets", return_value=mock.MagicMock(spec=Connection) - ) as mock_get_conn: - mock_get_conn.return_value.test_connection.return_value = (True, "Connection OK") - success, message = run_connection_test(connection_id="test_conn") + """Returns (True, message) on successful test.""" + conn = mock.MagicMock(spec=Connection) + conn.conn_id = "test_conn" + conn.test_connection.return_value = (True, "Connection OK") + + success, message = run_connection_test(conn=conn) assert success is True assert message == "Connection OK" def test_failed_connection_test(self): - """Pure function returns (False, message) when test_connection returns False.""" - with mock.patch.object( - Connection, "get_connection_from_secrets", return_value=mock.MagicMock(spec=Connection) - ) as mock_get_conn: - mock_get_conn.return_value.test_connection.return_value = (False, "Connection failed") - success, message = run_connection_test(connection_id="test_conn") + """Returns (False, message) when test_connection returns False.""" + conn = mock.MagicMock(spec=Connection) + conn.conn_id = "test_conn" + conn.test_connection.return_value = (False, "Connection failed") + + success, message = run_connection_test(conn=conn) assert success is False assert message == "Connection failed" def test_exception_during_connection_test(self): - """Pure function returns (False, error_str) on exception.""" - with mock.patch.object( - Connection, - "get_connection_from_secrets", - side_effect=Exception("Could not resolve host: db.example.com"), - ): - success, message = run_connection_test(connection_id="test_conn") + """Returns (False, error_str) on exception.""" + conn = mock.MagicMock(spec=Connection) + conn.conn_id = "test_conn" + conn.test_connection.side_effect = Exception("Could not resolve host: db.example.com") + + success, message = run_connection_test(conn=conn) assert success is False assert "Could not resolve host" in message From a01521044fad14810866f21960fc2439ab20f0fd Mon Sep 17 00:00:00 2001 From: Anish Date: Mon, 2 Mar 2026 22:34:58 -0600 Subject: [PATCH 17/38] update migration description --- airflow-core/docs/migrations-ref.rst | 2 +- .../migrations/versions/0109_3_2_0_add_connection_test_table.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index abd2765129d55..14ca8d6f5eec3 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -42,7 +42,7 @@ Here's the list of all the Database Migrations that are executed via when you ru | ``b8f3e5d1a9c2`` (head) | ``a7e6d4c3b2f1`` | ``3.2.0`` | Add connection_snapshot and reverted columns to | | | | | connection_test. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ -| ``a7e6d4c3b2f1`` | ``1d6611b6ab7c`` | ``3.2.0`` | Add connection_test table. | +| ``a7e6d4c3b2f1`` | ``1d6611b6ab7c`` | ``3.2.0`` | Add connection_test table for async connection testing. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``1d6611b6ab7c`` | ``888b59e02a5b`` | ``3.2.0`` | Add bundle_name to callback table. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py index cc870be041af2..36f94b265f67b 100644 --- a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py @@ -17,7 +17,7 @@ # under the License. """ -Add connection_test table. +Add connection_test table for async connection testing. Revision ID: a7e6d4c3b2f1 Revises: 888b59e02a5b From f82bd1bc4e4ceeb856b3bc86f5bff2b302b60ade Mon Sep 17 00:00:00 2001 From: Anish Date: Mon, 2 Mar 2026 23:18:23 -0600 Subject: [PATCH 18/38] clean ups --- airflow-core/newsfragments/62343.feature.rst | 1 + airflow-core/src/airflow/executors/local_executor.py | 4 ++-- airflow-core/src/airflow/jobs/scheduler_job_runner.py | 4 ++-- airflow-core/tests/unit/executors/test_local_executor.py | 5 +++-- 4 files changed, 8 insertions(+), 6 deletions(-) create mode 100644 airflow-core/newsfragments/62343.feature.rst diff --git a/airflow-core/newsfragments/62343.feature.rst b/airflow-core/newsfragments/62343.feature.rst new file mode 100644 index 0000000000000..52448232e3087 --- /dev/null +++ b/airflow-core/newsfragments/62343.feature.rst @@ -0,0 +1 @@ +Add async connection testing support on workers diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 8f89e9d32209e..43aaa6801b27c 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -244,12 +244,12 @@ def _handle_timeout(signum, frame): ConnectionTestState.FAILED, f"Connection test timed out after {workload.timeout}s", ) - except Exception: + except Exception as e: log.exception("Connection test failed unexpectedly", connection_id=workload.connection_id) client.connection_tests.update_state( workload.connection_test_id, ConnectionTestState.FAILED, - "Connection test failed unexpectedly", + f"Connection test failed unexpectedly: {e}"[:500], ) finally: signal.alarm(0) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 328e87d229403..d1d79519495dd 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3126,7 +3126,7 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]: existing_warned_dag_ids.add(warning.dag_id) @provide_session - def _dispatch_connection_tests(self, session: Session = NEW_SESSION) -> None: + def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None: """Dispatch pending connection tests to executors that support them.""" max_concurrency = conf.getint("core", "max_connection_test_concurrency", fallback=4) timeout = conf.getint("core", "connection_test_timeout", fallback=60) @@ -3177,7 +3177,7 @@ def _dispatch_connection_tests(self, session: Session = NEW_SESSION) -> None: session.flush() @provide_session - def _reap_stale_connection_tests(self, session: Session = NEW_SESSION) -> None: + 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) diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index a194ad5e6d3d2..227b4a3bb21d6 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -476,7 +476,7 @@ def test_connection_not_found_via_execution_api(self, MockClient, _mock_signal): calls = mock_client.connection_tests.update_state.call_args_list assert calls[-1].args[1] == ConnectionTestState.FAILED - assert calls[-1].args[2] == "Connection test failed unexpectedly" + assert "Connection test failed unexpectedly" in calls[-1].args[2] def test_unexpected_exception_reports_failed(self, MockClient, _mock_signal): """Reports FAILED when an unexpected exception occurs.""" @@ -505,7 +505,8 @@ def test_unexpected_exception_reports_failed(self, MockClient, _mock_signal): ) calls = mock_client.connection_tests.update_state.call_args_list - assert calls[-1].args == (test_id, ConnectionTestState.FAILED, "Connection test failed unexpectedly") + assert calls[-1].args[1] == ConnectionTestState.FAILED + assert "Connection test failed unexpectedly: Something broke" in calls[-1].args[2] def test_connection_fields_passed_correctly(self, MockClient, _mock_signal): """Verifies all connection fields from the API response are passed to Connection.""" From d5d812653c7a55a7c94708eb55bb2431c81bfe36 Mon Sep 17 00:00:00 2001 From: Anish Date: Tue, 3 Mar 2026 00:08:47 -0600 Subject: [PATCH 19/38] refactor dispatch and reaper for connection test --- .../src/airflow/executors/base_executor.py | 6 +++++- .../src/airflow/jobs/scheduler_job_runner.py | 17 ++++++++--------- .../tests/unit/executors/test_base_executor.py | 2 +- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index aa2fff0fcd331..eba900e9ac5cc 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -235,7 +235,11 @@ def queue_workload(self, workload: workloads.All, session: Session) -> None: self.queued_callbacks[workload.callback.id] = workload elif isinstance(workload, workloads.TestConnection): if not self.supports_connection_test: - raise ValueError(f"Executor {type(self).__name__} does not support connection testing") + raise NotImplementedError( + f"{type(self).__name__} does not support TestConnection workloads. " + f"Set supports_connection_test = True and implement connection test handling " + f"in _process_workloads(). See LocalExecutor for reference implementation." + ) self.queued_connection_tests[str(workload.connection_test_id)] = workload else: raise ValueError( diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index d1d79519495dd..5514f86530bef 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1586,15 +1586,14 @@ def _run_scheduler_loop(self) -> None: action=bundle_cleanup_mgr.remove_stale_bundle_versions, ) - if any(x.supports_connection_test for x in self.executors): - timers.call_regular_interval( - delay=conf.getfloat("scheduler", "connection_test_dispatch_interval", fallback=2.0), - action=self._dispatch_connection_tests, - ) - timers.call_regular_interval( - delay=conf.getfloat("scheduler", "connection_test_reaper_interval", fallback=30.0), - action=self._reap_stale_connection_tests, - ) + timers.call_regular_interval( + delay=conf.getfloat("scheduler", "connection_test_dispatch_interval", fallback=2.0), + action=self._dispatch_connection_tests, + ) + timers.call_regular_interval( + delay=conf.getfloat("scheduler", "connection_test_reaper_interval", fallback=30.0), + action=self._reap_stale_connection_tests, + ) idle_count = 0 diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 9f92d783eaca6..d126e465e14e9 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -419,7 +419,7 @@ def test_queue_connection_test_workload_rejected_by_default(): connection_test_id=uuid4(), connection_id="test_conn", ) - with pytest.raises(ValueError, match="does not support connection testing"): + with pytest.raises(NotImplementedError, match="does not support TestConnection workloads"): executor.queue_workload(wl, session=mock.MagicMock(spec=Session)) From 918641a54f9cc0afa42b8a18d09dfcd85e8cb5b8 Mon Sep 17 00:00:00 2001 From: Anish Date: Tue, 3 Mar 2026 01:53:33 -0600 Subject: [PATCH 20/38] address review feedback refactor --- .../core_api/datamodels/connections.py | 2 +- .../openapi/v2-rest-api-generated.yaml | 14 +-- .../core_api/routes/public/connections.py | 100 +++++++++--------- .../src/airflow/jobs/scheduler_job_runner.py | 15 +-- .../0109_3_2_0_add_connection_test_table.py | 2 +- .../src/airflow/models/connection_test.py | 6 +- .../airflow/ui/openapi-gen/queries/queries.ts | 8 +- .../ui/openapi-gen/requests/schemas.gen.ts | 4 +- .../ui/openapi-gen/requests/services.gen.ts | 6 +- .../ui/openapi-gen/requests/types.gen.ts | 8 +- .../routes/public/test_connections.py | 30 +++--- .../tests/unit/jobs/test_scheduler_job.py | 20 ++-- .../tests/unit/models/test_connection_test.py | 10 +- .../airflowctl/api/datamodels/generated.py | 2 +- 14 files changed, 115 insertions(+), 112 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index 0f4ca3fe73673..d2f066ac69412 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -83,7 +83,7 @@ class ConnectionTestRequestBody(StrictBaseModel): """Request body for async connection test.""" connection_id: str - queue: str | None = None + executor: str | None = None class ConnectionTestQueuedResponse(BaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index b6013cc1855b4..b0cd6fa32dfa1 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -1631,7 +1631,7 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' - /api/v2/connections/{connection_id}/save-and-test: + /api/v2/connections/{connection_id}/test: patch: tags: - Connection @@ -1665,16 +1665,16 @@ paths: type: string - type: 'null' title: Update Mask - - name: queue + - name: executor in: query required: false schema: anyOf: - type: string - type: 'null' - description: Executor queue to route the connection test to - title: Queue - description: Executor queue to route the connection test to + description: Executor (team) to route the connection test to + title: Executor + description: Executor (team) to route the connection test to requestBody: required: true content: @@ -10325,11 +10325,11 @@ components: connection_id: type: string title: Connection Id - queue: + executor: anyOf: - type: string - type: 'null' - title: Queue + title: Executor additionalProperties: false type: object required: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index b5006dc752f93..a539ae62caf44 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -190,53 +190,7 @@ def bulk_connections( @connections_router.patch( - "/{connection_id}", - responses=create_openapi_http_exception_doc( - [ - status.HTTP_400_BAD_REQUEST, - status.HTTP_404_NOT_FOUND, - ] - ), - dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())], -) -def patch_connection( - connection_id: str, - patch_body: ConnectionBody, - session: SessionDep, - update_mask: list[str] | None = Query(None), -) -> ConnectionResponse: - """Update a connection entry.""" - if patch_body.connection_id != connection_id: - raise HTTPException( - status.HTTP_400_BAD_REQUEST, - "The connection_id in the request body does not match the URL parameter", - ) - - connection = session.scalar(select(Connection).filter_by(conn_id=connection_id).limit(1)) - - if connection is None: - raise HTTPException( - status.HTTP_404_NOT_FOUND, f"The Connection with connection_id: `{connection_id}` was not found" - ) - - if update_mask: - fields_to_update = patch_body.model_fields_set & set(update_mask) - try: - ConnectionBodyPartial(**patch_body.model_dump(include=fields_to_update)) - except ValidationError as e: - raise RequestValidationError(errors=e.errors()) - else: - try: - ConnectionBody(**patch_body.model_dump()) - except ValidationError as e: - raise RequestValidationError(errors=e.errors()) - - update_orm_from_pydantic(connection, patch_body, update_mask) - return connection - - -@connections_router.patch( - "/{connection_id}/save-and-test", + "/{connection_id}/test", responses=create_openapi_http_exception_doc( [ status.HTTP_400_BAD_REQUEST, @@ -251,7 +205,7 @@ def patch_connection_and_test( patch_body: ConnectionBody, session: SessionDep, update_mask: list[str] | None = Query(None), - queue: str | None = Query(None, description="Executor queue to route the connection test to"), + executor: str | None = Query(None, description="Executor (team) to route the connection test to"), ) -> ConnectionSaveAndTestResponse: """ Update a connection and queue an async test with revert-on-failure. @@ -286,7 +240,7 @@ def patch_connection_and_test( post_snapshot = snapshot_connection(connection) - connection_test = ConnectionTest(connection_id=connection_id, queue=queue) + connection_test = ConnectionTest(connection_id=connection_id, executor=executor) connection_test.connection_snapshot = {"pre": pre_snapshot, "post": post_snapshot} session.add(connection_test) session.flush() @@ -298,6 +252,52 @@ def patch_connection_and_test( ) +@connections_router.patch( + "/{connection_id}", + responses=create_openapi_http_exception_doc( + [ + status.HTTP_400_BAD_REQUEST, + status.HTTP_404_NOT_FOUND, + ] + ), + dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())], +) +def patch_connection( + connection_id: str, + patch_body: ConnectionBody, + session: SessionDep, + update_mask: list[str] | None = Query(None), +) -> ConnectionResponse: + """Update a connection entry.""" + if patch_body.connection_id != connection_id: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, + "The connection_id in the request body does not match the URL parameter", + ) + + connection = session.scalar(select(Connection).filter_by(conn_id=connection_id).limit(1)) + + if connection is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, f"The Connection with connection_id: `{connection_id}` was not found" + ) + + if update_mask: + fields_to_update = patch_body.model_fields_set & set(update_mask) + try: + ConnectionBodyPartial(**patch_body.model_dump(include=fields_to_update)) + except ValidationError as e: + raise RequestValidationError(errors=e.errors()) + else: + try: + ConnectionBody(**patch_body.model_dump()) + except ValidationError as e: + raise RequestValidationError(errors=e.errors()) + + update_orm_from_pydantic(connection, patch_body, update_mask) + return connection + + @connections_router.post("/test", dependencies=[Depends(requires_access_connection(method="POST"))]) def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse: """ @@ -357,7 +357,7 @@ def test_connection_async( "Connection must be saved before testing.", ) - connection_test = ConnectionTest(connection_id=test_body.connection_id, queue=test_body.queue) + connection_test = ConnectionTest(connection_id=test_body.connection_id, executor=test_body.executor) session.add(connection_test) session.flush() diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 5514f86530bef..a6c03940aa1a1 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3152,11 +3152,11 @@ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None: return for ct in pending_tests: - executor = self._find_executor_for_connection_test(ct.queue) + executor = self._find_executor_for_connection_test(ct.executor) if executor is None: reason = ( - f"No executor serves queue '{ct.queue}'" - if ct.queue + f"No executor matches '{ct.executor}'" + if ct.executor else "No executor supports connection testing" ) ct.state = ConnectionTestState.FAILED @@ -3186,6 +3186,7 @@ def _reap_stale_connection_tests(self, *, session: Session = NEW_SESSION) -> Non ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]), ConnectionTest.updated_at < cutoff, ) + stale_stmt = with_row_locks(stale_stmt, session, of=ConnectionTest, skip_locked=True) stale_tests = session.scalars(stale_stmt).all() for ct in stale_tests: @@ -3197,11 +3198,11 @@ def _reap_stale_connection_tests(self, *, session: Session = NEW_SESSION) -> Non session.flush() - def _find_executor_for_connection_test(self, queue: str | None) -> BaseExecutor | None: - """Find an executor that supports connection testing, optionally matching a queue.""" - if queue is not None: + def _find_executor_for_connection_test(self, executor_name: str | None) -> BaseExecutor | None: + """Find an executor that supports connection testing, optionally matching by team name.""" + if executor_name is not None: for executor in self.executors: - if executor.supports_connection_test and executor.team_name == queue: + if executor.supports_connection_test and executor.team_name == executor_name: return executor return None for executor in self.executors: diff --git a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py index 36f94b265f67b..a5bbc49b57e93 100644 --- a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py @@ -51,7 +51,7 @@ def upgrade(): 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("queue", sa.String(256), nullable=True), + sa.Column("executor", sa.String(256), nullable=True), sa.Column("connection_snapshot", sa.JSON(), nullable=True), sa.Column("reverted", sa.Boolean(), nullable=False, server_default="0"), sa.PrimaryKeyConstraint("id", name=op.f("connection_test_pkey")), diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 6c370ba5f5361..ff79e359ad6be 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -69,16 +69,16 @@ class ConnectionTest(Base): updated_at: Mapped[datetime] = mapped_column( UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False ) - queue: Mapped[str | None] = mapped_column(String(256), nullable=True) + executor: Mapped[str | None] = mapped_column(String(256), nullable=True) connection_snapshot: Mapped[dict | None] = mapped_column(JSON, nullable=True) reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0") __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),) - def __init__(self, *, connection_id: str, queue: str | None = None, **kwargs): + def __init__(self, *, connection_id: str, executor: str | None = None, **kwargs): super().__init__(**kwargs) self.connection_id = connection_id - self.queue = queue + self.executor = executor self.token = secrets.token_urlsafe(32) self.state = ConnectionTestState.PENDING diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index ed4ddc49ebf6f..305ab01c8fa39 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -2128,21 +2128,21 @@ export const useConnectionServiceBulkConnections = (options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ connectionId, queue, requestBody, updateMask }) => ConnectionService.patchConnectionAndTest({ connectionId, queue, requestBody, updateMask }) as unknown as Promise, ...options }); +}, TContext>({ mutationFn: ({ connectionId, executor, requestBody, updateMask }) => ConnectionService.patchConnectionAndTest({ connectionId, executor, requestBody, updateMask }) as unknown as Promise, ...options }); /** * Patch Dag Run * Modify a DAG Run. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index ca61dfd80daea..752ab534e0bf5 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1745,7 +1745,7 @@ export const $ConnectionTestRequestBody = { type: 'string', title: 'Connection Id' }, - queue: { + executor: { anyOf: [ { type: 'string' @@ -1754,7 +1754,7 @@ export const $ConnectionTestRequestBody = { type: 'null' } ], - title: 'Queue' + title: 'Executor' } }, additionalProperties: false, diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 7e0b0da932471..21a8247b5b5a4 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -779,20 +779,20 @@ export class ConnectionService { * @param data.connectionId * @param data.requestBody * @param data.updateMask - * @param data.queue Executor queue to route the connection test to + * @param data.executor Executor (team) to route the connection test to * @returns ConnectionSaveAndTestResponse Successful Response * @throws ApiError */ public static patchConnectionAndTest(data: PatchConnectionAndTestData): CancelablePromise { return __request(OpenAPI, { method: 'PATCH', - url: '/api/v2/connections/{connection_id}/save-and-test', + url: '/api/v2/connections/{connection_id}/test', path: { connection_id: data.connectionId }, query: { update_mask: data.updateMask, - queue: data.queue + executor: data.executor }, body: data.requestBody, mediaType: 'application/json', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index f89d1c9d35d12..919bb7f53397c 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -518,7 +518,7 @@ export type ConnectionTestQueuedResponse = { */ export type ConnectionTestRequestBody = { connection_id: string; - queue?: string | null; + executor?: string | null; }; /** @@ -2517,9 +2517,9 @@ export type BulkConnectionsResponse = BulkResponse; export type PatchConnectionAndTestData = { connectionId: string; /** - * Executor queue to route the connection test to + * Executor (team) to route the connection test to */ - queue?: string | null; + executor?: string | null; requestBody: ConnectionBody; updateMask?: Array<(string)> | null; }; @@ -4504,7 +4504,7 @@ export type $OpenApiTs = { }; }; }; - '/api/v2/connections/{connection_id}/save-and-test': { + '/api/v2/connections/{connection_id}/test': { patch: { req: PatchConnectionAndTestData; res: { diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 894d16b739e58..bcbf7ca83eae0 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -1273,14 +1273,14 @@ def test_get_status_returns_404_for_invalid_token(self, test_client): class TestSaveAndTest(TestConnectionEndpoint): - """Tests for the combined PATCH /{connection_id}/save-and-test endpoint.""" + """Tests for the combined PATCH /{connection_id}/test endpoint.""" @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_save_and_test_returns_200_with_token(self, test_client, session): - """PATCH save-and-test updates the connection and returns a test token.""" + """PATCH /{connection_id}/test updates the connection and returns a test token.""" self.create_connection() response = test_client.patch( - f"/connections/{TEST_CONN_ID}/save-and-test", + f"/connections/{TEST_CONN_ID}/test", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, @@ -1295,10 +1295,10 @@ def test_save_and_test_returns_200_with_token(self, test_client, session): @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_save_and_test_creates_snapshot(self, test_client, session): - """PATCH save-and-test creates a ConnectionTest with a connection_snapshot.""" + """PATCH /{connection_id}/test creates a ConnectionTest with a connection_snapshot.""" self.create_connection() response = test_client.patch( - f"/connections/{TEST_CONN_ID}/save-and-test", + f"/connections/{TEST_CONN_ID}/test", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, @@ -1318,11 +1318,11 @@ def test_save_and_test_creates_snapshot(self, test_client, session): assert snapshot["post"]["host"] == "new-host.example.com" @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_save_and_test_passes_queue_parameter(self, test_client, session): - """PATCH save-and-test passes the queue parameter to the ConnectionTest.""" + def test_save_and_test_passes_executor_parameter(self, test_client, session): + """PATCH /{connection_id}/test passes the executor parameter to the ConnectionTest.""" self.create_connection() response = test_client.patch( - f"/connections/{TEST_CONN_ID}/save-and-test?queue=team_a", + f"/connections/{TEST_CONN_ID}/test?executor=team_a", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, @@ -1334,13 +1334,13 @@ def test_save_and_test_passes_queue_parameter(self, test_client, session): ct = session.scalar(select(ConnectionTest).filter_by(token=token)) assert ct is not None - assert ct.queue == "team_a" + assert ct.executor == "team_a" def test_save_and_test_403_when_disabled(self, test_client): - """PATCH save-and-test returns 403 when test_connection is disabled.""" + """PATCH /{connection_id}/test returns 403 when test_connection is disabled.""" self.create_connection() response = test_client.patch( - f"/connections/{TEST_CONN_ID}/save-and-test", + f"/connections/{TEST_CONN_ID}/test", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, @@ -1350,9 +1350,9 @@ def test_save_and_test_403_when_disabled(self, test_client): @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_save_and_test_404_for_nonexistent(self, test_client): - """PATCH save-and-test returns 404 for nonexistent connection.""" + """PATCH /{connection_id}/test returns 404 for nonexistent connection.""" response = test_client.patch( - "/connections/nonexistent/save-and-test", + "/connections/nonexistent/test", json={ "connection_id": "nonexistent", "conn_type": "http", @@ -1365,7 +1365,7 @@ def test_poll_shows_reverted_true_after_failed_test(self, test_client, session): """GET status shows reverted=True after a failed test triggers a revert.""" self.create_connection() response = test_client.patch( - f"/connections/{TEST_CONN_ID}/save-and-test", + f"/connections/{TEST_CONN_ID}/test", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, @@ -1391,7 +1391,7 @@ def test_poll_shows_reverted_false_for_success(self, test_client, session): """GET status shows reverted=False for a successful test.""" self.create_connection() response = test_client.patch( - f"/connections/{TEST_CONN_ID}/save-and-test", + f"/connections/{TEST_CONN_ID}/test", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index d600f312eac98..fa9fb6daf897c 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9550,11 +9550,11 @@ def test_dispatch_fails_fast_when_no_executor_supports( "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", }, ) - def test_dispatch_with_unmatched_queue_fails_fast( + def test_dispatch_with_unmatched_executor_fails_fast( self, scheduler_job_runner_for_connection_tests, session ): - """Tests requesting a queue with no matching executor are failed immediately.""" - ct = ConnectionTest(connection_id="test_conn", queue="gpu_workers") + """Tests requesting an executor with no match are failed immediately.""" + ct = ConnectionTest(connection_id="test_conn", executor="gpu_workers") session.add(ct) session.commit() @@ -9637,14 +9637,16 @@ def test_dispatch_order_is_fifo_by_created_at(self, scheduler_job_runner_for_con "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", }, ) - def test_dispatch_fails_fast_for_unserved_queue(self, scheduler_job_runner_for_connection_tests, session): - """Tests requesting a queue no executor serves are failed immediately.""" + def test_dispatch_fails_fast_for_unserved_executor( + self, scheduler_job_runner_for_connection_tests, session + ): + """Tests requesting an executor no team serves are failed immediately.""" with mock.patch.object( scheduler_job_runner_for_connection_tests, "_find_executor_for_connection_test", return_value=None, ): - ct = ConnectionTest(connection_id="test_conn", queue="nonexistent_queue") + ct = ConnectionTest(connection_id="test_conn", executor="nonexistent_queue") session.add(ct) session.commit() @@ -9662,8 +9664,8 @@ def test_dispatch_fails_fast_for_unserved_queue(self, scheduler_job_runner_for_c "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", }, ) - def test_dispatch_queue_matched_to_team_name(self, session): - """When queue is specified, executor whose team_name matches is selected.""" + def test_dispatch_executor_matched_to_team_name(self, session): + """When executor is specified, the executor whose team_name matches is selected.""" session.execute(delete(ConnectionTest)) session.commit() @@ -9685,7 +9687,7 @@ def test_dispatch_queue_matched_to_team_name(self, session): runner.executor = executor_a runner._log = mock.MagicMock(spec=logging.Logger) - ct = ConnectionTest(connection_id="team_conn", queue="team_b") + ct = ConnectionTest(connection_id="team_conn", executor="team_b") session.add(ct) session.commit() diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py index d12f955055cdb..845636c2bad22 100644 --- a/airflow-core/tests/unit/models/test_connection_test.py +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -55,13 +55,13 @@ def test_repr(self): assert "test_conn" in r assert "pending" in r - def test_queue_parameter(self): - ct = ConnectionTest(connection_id="test_conn", queue="my_queue") - assert ct.queue == "my_queue" + def test_executor_parameter(self): + ct = ConnectionTest(connection_id="test_conn", executor="my_executor") + assert ct.executor == "my_executor" - def test_queue_defaults_to_none(self): + def test_executor_defaults_to_none(self): ct = ConnectionTest(connection_id="test_conn") - assert ct.queue is None + assert ct.executor is None class TestRunConnectionTest: diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index e41de765672b6..d960f2c66ffbc 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -273,7 +273,7 @@ class ConnectionTestRequestBody(BaseModel): extra="forbid", ) connection_id: Annotated[str, Field(title="Connection Id")] - queue: Annotated[str | None, Field(title="Queue")] = None + executor: Annotated[str | None, Field(title="Executor")] = None class ConnectionTestResponse(BaseModel): From 1352b73748876858223d1858cf523a0e040defa2 Mon Sep 17 00:00:00 2001 From: Anish Date: Tue, 3 Mar 2026 10:01:55 -0600 Subject: [PATCH 21/38] clean ups --- .../versions/0109_3_2_0_add_connection_test_table.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py index a5bbc49b57e93..1cfd0b38091d5 100644 --- a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py @@ -20,7 +20,11 @@ Add connection_test table for async connection testing. Revision ID: a7e6d4c3b2f1 +<<<<<<<< HEAD:airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py Revises: 888b59e02a5b +======== +Revises: 6222ce48e289 +>>>>>>>> 189776ee4f (clean ups):airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_test_table.py Create Date: 2026-02-22 00:00:00.000000 """ @@ -34,7 +38,11 @@ # revision identifiers, used by Alembic. revision = "a7e6d4c3b2f1" +<<<<<<<< HEAD:airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py down_revision = "888b59e02a5b" +======== +down_revision = "6222ce48e289" +>>>>>>>> 189776ee4f (clean ups):airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_test_table.py branch_labels = None depends_on = None airflow_version = "3.2.0" From 2ec3a61b3dfe49d07f60cb82c24b763ffc570427 Mon Sep 17 00:00:00 2001 From: Anish Date: Wed, 4 Mar 2026 10:08:31 -0600 Subject: [PATCH 22/38] remove queue executer and team name mixed up --- .../core_api/datamodels/connections.py | 1 + .../openapi/v2-rest-api-generated.yaml | 19 ++++++- .../core_api/routes/public/connections.py | 9 ++-- .../src/airflow/executors/base_executor.py | 5 +- .../executors/workloads/connection_test.py | 3 ++ .../src/airflow/jobs/scheduler_job_runner.py | 17 ++++--- .../0109_3_2_0_add_connection_test_table.py | 1 + .../src/airflow/models/connection_test.py | 7 ++- .../airflow/ui/openapi-gen/queries/queries.ts | 7 ++- .../ui/openapi-gen/requests/schemas.gen.ts | 11 +++++ .../ui/openapi-gen/requests/services.gen.ts | 6 ++- .../ui/openapi-gen/requests/types.gen.ts | 7 ++- .../routes/public/test_connections.py | 34 +++++++++++++ .../unit/executors/test_base_executor.py | 9 ++++ .../tests/unit/jobs/test_scheduler_job.py | 49 +++++++++++++++++-- .../tests/unit/models/test_connection_test.py | 8 +++ .../airflowctl/api/datamodels/generated.py | 1 + 17 files changed, 168 insertions(+), 26 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index d2f066ac69412..1e51b8ceb10ff 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -84,6 +84,7 @@ class ConnectionTestRequestBody(StrictBaseModel): connection_id: str executor: str | None = None + queue: str | None = None class ConnectionTestQueuedResponse(BaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index b0cd6fa32dfa1..9de1315908bb4 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -1672,9 +1672,19 @@ paths: anyOf: - type: string - type: 'null' - description: Executor (team) to route the connection test to + description: Executor to route the connection test to title: Executor - description: Executor (team) to route the connection test to + description: Executor to route the connection test to + - name: queue + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + description: Queue to route the connection test to + title: Queue + description: Queue to route the connection test to requestBody: required: true content: @@ -10330,6 +10340,11 @@ components: - type: string - type: 'null' title: Executor + queue: + anyOf: + - type: string + - type: 'null' + title: Queue additionalProperties: false type: object required: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index a539ae62caf44..2a0a18a060ba0 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -205,7 +205,8 @@ def patch_connection_and_test( patch_body: ConnectionBody, session: SessionDep, update_mask: list[str] | None = Query(None), - executor: str | None = Query(None, description="Executor (team) to route the connection test to"), + executor: str | None = Query(None, description="Executor to route the connection test to"), + queue: str | None = Query(None, description="Queue to route the connection test to"), ) -> ConnectionSaveAndTestResponse: """ Update a connection and queue an async test with revert-on-failure. @@ -240,7 +241,7 @@ def patch_connection_and_test( post_snapshot = snapshot_connection(connection) - connection_test = ConnectionTest(connection_id=connection_id, executor=executor) + connection_test = ConnectionTest(connection_id=connection_id, executor=executor, queue=queue) connection_test.connection_snapshot = {"pre": pre_snapshot, "post": post_snapshot} session.add(connection_test) session.flush() @@ -357,7 +358,9 @@ def test_connection_async( "Connection must be saved before testing.", ) - connection_test = ConnectionTest(connection_id=test_body.connection_id, executor=test_body.executor) + connection_test = ConnectionTest( + connection_id=test_body.connection_id, executor=test_body.executor, queue=test_body.queue + ) session.add(connection_test) session.flush() diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index eba900e9ac5cc..99dc9a3ffacbc 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -315,8 +315,7 @@ def heartbeat(self) -> None: self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads) self.trigger_tasks(open_slots) - if self.supports_connection_test and self.queued_connection_tests: - self.trigger_connection_tests() + self.trigger_connection_tests() # Calling child class sync method self.log.debug("Calling the %s sync method", self.__class__) @@ -324,7 +323,7 @@ def heartbeat(self) -> None: def trigger_connection_tests(self) -> None: """Process queued connection tests.""" - if not self.queued_connection_tests: + if not self.supports_connection_test or not self.queued_connection_tests: return self._process_workloads(list(self.queued_connection_tests.values())) diff --git a/airflow-core/src/airflow/executors/workloads/connection_test.py b/airflow-core/src/airflow/executors/workloads/connection_test.py index 356d8ac212f6e..9cbb04229c075 100644 --- a/airflow-core/src/airflow/executors/workloads/connection_test.py +++ b/airflow-core/src/airflow/executors/workloads/connection_test.py @@ -35,6 +35,7 @@ class TestConnection(BaseWorkloadSchema): connection_test_id: uuid.UUID connection_id: str timeout: int = 60 + queue: str | None = None type: Literal["TestConnection"] = Field(init=False, default="TestConnection") @@ -45,11 +46,13 @@ def make( connection_test_id: uuid.UUID, connection_id: str, timeout: int = 60, + queue: str | None = None, generator: JWTGenerator | None = None, ) -> TestConnection: return cls( connection_test_id=connection_test_id, connection_id=connection_id, timeout=timeout, + queue=queue, token=cls.generate_token(str(connection_test_id), generator), ) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index a6c03940aa1a1..9fc501c90b2c3 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -73,7 +73,7 @@ ) from airflow.models.backfill import Backfill from airflow.models.callback import Callback, CallbackType, ExecutorCallback -from airflow.models.connection_test import ConnectionTest, ConnectionTestState, attempt_revert +from airflow.models.connection_test import ACTIVE_STATES, ConnectionTest, ConnectionTestState, attempt_revert from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion from airflow.models.dagbag import DBDagBag @@ -3131,9 +3131,7 @@ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None: timeout = conf.getint("core", "connection_test_timeout", fallback=60) active_count = session.scalar( - select(func.count(ConnectionTest.id)).where( - ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]) - ) + select(func.count(ConnectionTest.id)).where(ConnectionTest.state.in_(ACTIVE_STATES)) ) budget = max_concurrency - (active_count or 0) if budget <= 0: @@ -3168,6 +3166,7 @@ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None: 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) @@ -3183,7 +3182,7 @@ def _reap_stale_connection_tests(self, *, session: Session = NEW_SESSION) -> Non cutoff = timezone.utcnow() - timedelta(seconds=timeout + grace_period) stale_stmt = select(ConnectionTest).where( - ConnectionTest.state.in_([ConnectionTestState.QUEUED, ConnectionTestState.RUNNING]), + ConnectionTest.state.in_(ACTIVE_STATES), ConnectionTest.updated_at < cutoff, ) stale_stmt = with_row_locks(stale_stmt, session, of=ConnectionTest, skip_locked=True) @@ -3199,10 +3198,14 @@ def _reap_stale_connection_tests(self, *, session: Session = NEW_SESSION) -> Non session.flush() def _find_executor_for_connection_test(self, executor_name: str | None) -> BaseExecutor | None: - """Find an executor that supports connection testing, optionally matching by team name.""" + """Find an executor that supports connection testing, optionally matching by name.""" if executor_name is not None: for executor in self.executors: - if executor.supports_connection_test and executor.team_name == executor_name: + if ( + executor.supports_connection_test + and executor.name + and executor_name in (executor.name.alias, executor.name.module_path) + ): return executor return None for executor in self.executors: diff --git a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py index 1cfd0b38091d5..9755c67bd8456 100644 --- a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py @@ -60,6 +60,7 @@ def upgrade(): 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("connection_snapshot", sa.JSON(), nullable=True), sa.Column("reverted", sa.Boolean(), nullable=False, server_default="0"), sa.PrimaryKeyConstraint("id", name=op.f("connection_test_pkey")), diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index ff79e359ad6be..1747257323fae 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -52,6 +52,7 @@ def __str__(self) -> str: return self.value +ACTIVE_STATES = frozenset((ConnectionTestState.QUEUED, ConnectionTestState.RUNNING)) TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED)) @@ -70,15 +71,19 @@ class ConnectionTest(Base): UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False ) executor: Mapped[str | None] = mapped_column(String(256), nullable=True) + queue: Mapped[str | None] = mapped_column(String(256), nullable=True) connection_snapshot: Mapped[dict | None] = mapped_column(JSON, nullable=True) reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0") __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),) - def __init__(self, *, connection_id: str, executor: str | None = None, **kwargs): + def __init__( + self, *, connection_id: str, executor: str | None = None, queue: str | None = None, **kwargs + ): super().__init__(**kwargs) self.connection_id = connection_id self.executor = executor + self.queue = queue self.token = secrets.token_urlsafe(32) self.state = ConnectionTestState.PENDING diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index 305ab01c8fa39..f2f02439a3595 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -2128,21 +2128,24 @@ export const useConnectionServiceBulkConnections = (options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ connectionId, executor, requestBody, updateMask }) => ConnectionService.patchConnectionAndTest({ connectionId, executor, requestBody, updateMask }) as unknown as Promise, ...options }); +}, TContext>({ mutationFn: ({ connectionId, executor, queue, requestBody, updateMask }) => ConnectionService.patchConnectionAndTest({ connectionId, executor, queue, requestBody, updateMask }) as unknown as Promise, ...options }); /** * Patch Dag Run * Modify a DAG Run. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 752ab534e0bf5..4107e3f88e3df 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1755,6 +1755,17 @@ export const $ConnectionTestRequestBody = { } ], title: 'Executor' + }, + queue: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Queue' } }, additionalProperties: false, diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 21a8247b5b5a4..1470ae7a855e7 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -779,7 +779,8 @@ export class ConnectionService { * @param data.connectionId * @param data.requestBody * @param data.updateMask - * @param data.executor Executor (team) to route the connection test to + * @param data.executor Executor to route the connection test to + * @param data.queue Queue to route the connection test to * @returns ConnectionSaveAndTestResponse Successful Response * @throws ApiError */ @@ -792,7 +793,8 @@ export class ConnectionService { }, query: { update_mask: data.updateMask, - executor: data.executor + executor: data.executor, + queue: data.queue }, body: data.requestBody, mediaType: 'application/json', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 919bb7f53397c..4a032f12422d4 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -519,6 +519,7 @@ export type ConnectionTestQueuedResponse = { export type ConnectionTestRequestBody = { connection_id: string; executor?: string | null; + queue?: string | null; }; /** @@ -2517,9 +2518,13 @@ export type BulkConnectionsResponse = BulkResponse; export type PatchConnectionAndTestData = { connectionId: string; /** - * Executor (team) to route the connection test to + * Executor to route the connection test to */ executor?: string | null; + /** + * Queue to route the connection test to + */ + queue?: string | null; requestBody: ConnectionBody; updateMask?: Array<(string)> | null; }; diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index bcbf7ca83eae0..4f15a48dfaa47 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -1232,6 +1232,21 @@ def test_post_creates_connection_test_row(self, test_client, session): assert ct.connection_id == TEST_CONN_ID assert ct.state == "pending" + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_post_passes_queue_parameter(self, test_client, session): + """POST /connections/test-async passes the queue parameter to the ConnectionTest.""" + self.create_connection() + response = test_client.post( + "/connections/test-async", + json={"connection_id": TEST_CONN_ID, "queue": "gpu_workers"}, + ) + assert response.status_code == 202 + token = response.json()["token"] + + ct = session.scalar(select(ConnectionTest).filter_by(token=token)) + assert ct is not None + assert ct.queue == "gpu_workers" + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_get_status_returns_pending(self, test_client, session): """GET /connections/test-async/{token} returns current status (pending before scheduler dispatch).""" @@ -1336,6 +1351,25 @@ def test_save_and_test_passes_executor_parameter(self, test_client, session): assert ct is not None assert ct.executor == "team_a" + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_save_and_test_passes_queue_parameter(self, test_client, session): + """PATCH /{connection_id}/test passes the queue parameter to the ConnectionTest.""" + self.create_connection() + response = test_client.patch( + f"/connections/{TEST_CONN_ID}/test?queue=gpu_workers", + json={ + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "host": "queued-host.example.com", + }, + ) + assert response.status_code == 200 + token = response.json()["test_token"] + + ct = session.scalar(select(ConnectionTest).filter_by(token=token)) + assert ct is not None + assert ct.queue == "gpu_workers" + def test_save_and_test_403_when_disabled(self, test_client): """PATCH /{connection_id}/test returns 403 when test_connection is disabled.""" self.create_connection() diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index d126e465e14e9..8171e223fcffd 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -436,6 +436,15 @@ def test_queue_connection_test_workload_accepted_when_supported(): assert executor.queued_connection_tests[str(wl.connection_test_id)] is wl +def test_trigger_connection_tests_skipped_when_not_supported(): + """trigger_connection_tests is a no-op when supports_connection_test is False.""" + executor = BaseExecutor() + executor.queued_connection_tests["dummy"] = mock.MagicMock(spec=workloads.TestConnection) + with mock.patch.object(executor, "_process_workloads") as mock_process: + executor.trigger_connection_tests() + mock_process.assert_not_called() + + @mock.patch.dict("os.environ", {}, clear=True) class TestExecutorConf: """Test ExecutorConf shim class that provides team-specific configuration access.""" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index fa9fb6daf897c..147d2a215de32 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9664,8 +9664,8 @@ def test_dispatch_fails_fast_for_unserved_executor( "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", }, ) - def test_dispatch_executor_matched_to_team_name(self, session): - """When executor is specified, the executor whose team_name matches is selected.""" + def test_dispatch_executor_matched_by_alias(self, session): + """When executor is specified, the executor whose name.alias matches is selected.""" session.execute(delete(ConnectionTest)) session.commit() @@ -9674,11 +9674,11 @@ def test_dispatch_executor_matched_to_team_name(self, session): mock_job.max_tis_per_query = 16 executor_a = LocalExecutor() - executor_a.team_name = "team_a" + executor_a.name = ExecutorName(module_path="path.to.ExecutorA", alias="executor_a") executor_a.queued_connection_tests.clear() executor_b = LocalExecutor() - executor_b.team_name = "team_b" + executor_b.name = ExecutorName(module_path="path.to.ExecutorB", alias="executor_b") executor_b.queued_connection_tests.clear() runner = SchedulerJobRunner.__new__(SchedulerJobRunner) @@ -9687,7 +9687,46 @@ def test_dispatch_executor_matched_to_team_name(self, session): runner.executor = executor_a runner._log = mock.MagicMock(spec=logging.Logger) - ct = ConnectionTest(connection_id="team_conn", executor="team_b") + ct = ConnectionTest(connection_id="team_conn", executor="executor_b") + session.add(ct) + session.commit() + + runner._dispatch_connection_tests(session=session) + + assert len(executor_b.queued_connection_tests) == 1 + assert len(executor_a.queued_connection_tests) == 0 + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_executor_matched_by_module_path(self, session): + """When executor is specified by module_path, the matching executor is selected.""" + session.execute(delete(ConnectionTest)) + session.commit() + + mock_job = mock.MagicMock(spec=Job) + mock_job.id = 1 + mock_job.max_tis_per_query = 16 + + executor_a = LocalExecutor() + executor_a.name = ExecutorName(module_path="path.to.ExecutorA", alias="executor_a") + executor_a.queued_connection_tests.clear() + + executor_b = LocalExecutor() + executor_b.name = ExecutorName(module_path="path.to.ExecutorB", alias="executor_b") + executor_b.queued_connection_tests.clear() + + runner = SchedulerJobRunner.__new__(SchedulerJobRunner) + runner.job = mock_job + runner.executors = [executor_a, executor_b] + runner.executor = executor_a + runner._log = mock.MagicMock(spec=logging.Logger) + + ct = ConnectionTest(connection_id="team_conn", executor="path.to.ExecutorB") session.add(ct) session.commit() diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py index 845636c2bad22..3138fdc82640f 100644 --- a/airflow-core/tests/unit/models/test_connection_test.py +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -63,6 +63,14 @@ def test_executor_defaults_to_none(self): ct = ConnectionTest(connection_id="test_conn") assert ct.executor is None + def test_queue_parameter(self): + ct = ConnectionTest(connection_id="test_conn", queue="my_queue") + assert ct.queue == "my_queue" + + def test_queue_defaults_to_none(self): + ct = ConnectionTest(connection_id="test_conn") + assert ct.queue is None + class TestRunConnectionTest: def test_successful_connection_test(self): diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index d960f2c66ffbc..405f16843d67e 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -274,6 +274,7 @@ class ConnectionTestRequestBody(BaseModel): ) connection_id: Annotated[str, Field(title="Connection Id")] executor: Annotated[str | None, Field(title="Executor")] = None + queue: Annotated[str | None, Field(title="Queue")] = None class ConnectionTestResponse(BaseModel): From 15ebe5df4e471f2123940f8caa3c59ff6f5e0cc7 Mon Sep 17 00:00:00 2001 From: Anish Date: Wed, 4 Mar 2026 18:07:15 -0600 Subject: [PATCH 23/38] used version api --- .../api_fastapi/execution_api/routes/connection_tests.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py index 1ad2e0192c1bd..95c20380223c2 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -18,7 +18,8 @@ from uuid import UUID -from fastapi import APIRouter, HTTPException, status +from cadwyn import VersionedAPIRouter +from fastapi import HTTPException, status from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.execution_api.datamodels.connection_test import ConnectionTestResultBody @@ -29,7 +30,7 @@ attempt_revert, ) -router = APIRouter() +router = VersionedAPIRouter() @router.patch( From 60861aae55250117194acc1d57884bfecb1ba613 Mon Sep 17 00:00:00 2001 From: Anish Date: Wed, 4 Mar 2026 19:15:49 -0600 Subject: [PATCH 24/38] further refined executer finder --- .../src/airflow/jobs/scheduler_job_runner.py | 7 +++- .../tests/unit/jobs/test_scheduler_job.py | 32 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 9fc501c90b2c3..555c87445758c 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3204,7 +3204,12 @@ def _find_executor_for_connection_test(self, executor_name: str | None) -> BaseE if ( executor.supports_connection_test and executor.name - and executor_name in (executor.name.alias, executor.name.module_path) + and executor_name + in ( + executor.name.alias, + executor.name.module_path, + executor.name.module_path.split(".")[-1], + ) ): return executor return None diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 147d2a215de32..66e4bcf7d90bb 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9735,6 +9735,38 @@ def test_dispatch_executor_matched_by_module_path(self, session): assert len(executor_b.queued_connection_tests) == 1 assert len(executor_a.queued_connection_tests) == 0 + def test_dispatch_executor_matched_by_class_name(self, session): + """When executor is specified by class name only, the matching executor is selected.""" + session.execute(delete(ConnectionTest)) + session.commit() + + mock_job = mock.MagicMock(spec=Job) + mock_job.id = 1 + mock_job.max_tis_per_query = 16 + + executor_a = LocalExecutor() + executor_a.name = ExecutorName(module_path="path.to.ExecutorA", alias="executor_a") + executor_a.queued_connection_tests.clear() + + executor_b = LocalExecutor() + executor_b.name = ExecutorName(module_path="path.to.ExecutorB", alias="executor_b") + executor_b.queued_connection_tests.clear() + + runner = SchedulerJobRunner.__new__(SchedulerJobRunner) + runner.job = mock_job + runner.executors = [executor_a, executor_b] + runner.executor = executor_a + runner._log = mock.MagicMock(spec=logging.Logger) + + ct = ConnectionTest(connection_id="team_conn", executor="ExecutorB") + session.add(ct) + session.commit() + + runner._dispatch_connection_tests(session=session) + + assert len(executor_b.queued_connection_tests) == 1 + assert len(executor_a.queued_connection_tests) == 0 + class TestReapStaleConnectionTests: @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) From b17f75bfc6fa0be638e4fe622783ef0d81d32e99 Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 5 Mar 2026 13:53:49 -0600 Subject: [PATCH 25/38] Clear connection_snapshot from DB after terminal state is reached --- .../api_fastapi/execution_api/routes/connection_tests.py | 2 ++ airflow-core/src/airflow/jobs/scheduler_job_runner.py | 1 + airflow-core/src/airflow/models/connection_test.py | 2 ++ .../execution_api/versions/head/test_connection_tests.py | 4 ++++ airflow-core/tests/unit/jobs/test_scheduler_job.py | 2 ++ airflow-core/tests/unit/models/test_connection_test.py | 4 ++++ 6 files changed, 15 insertions(+) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py index 95c20380223c2..13e5a45fe7c74 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -71,3 +71,5 @@ def patch_connection_test( if body.state == ConnectionTestState.FAILED and ct.connection_snapshot: attempt_revert(ct, session=session) + else: + ct.connection_snapshot = None diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 555c87445758c..460e445ed67c8 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3159,6 +3159,7 @@ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None: ) ct.state = ConnectionTestState.FAILED ct.result_message = reason + ct.connection_snapshot = None self.log.warning("Failing connection test %s: %s", ct.id, reason) continue diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 1747257323fae..0ba8a6ff93bde 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -189,6 +189,8 @@ def attempt_revert(ct: ConnectionTest, *, session: Session) -> None: pre_snapshot = ct.connection_snapshot["pre"] post_snapshot = ct.connection_snapshot["post"] + ct.connection_snapshot = None + conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id)) if conn is None: ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists." diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py index 564de6fa1b27f..23db82cd76786 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py @@ -54,6 +54,7 @@ def test_patch_updates_result(self, client, session): ct = session.get(ConnectionTest, ct.id) assert ct.state == "success" assert ct.result_message == "Connection successfully tested" + assert ct.connection_snapshot is None def test_patch_returns_404_for_nonexistent(self, client): """PATCH with unknown id returns 404.""" @@ -130,6 +131,7 @@ def test_patch_failed_with_snapshot_reverts_connection(self, client, session): session.expire_all() ct = session.get(ConnectionTest, ct.id) assert ct.reverted is True + assert ct.connection_snapshot is None conn = session.scalar(select(Connection).filter_by(conn_id="revert_conn")) assert conn.host == "old-host.example.com" assert conn.login == "old_user" @@ -164,6 +166,7 @@ def test_patch_success_with_snapshot_no_revert(self, client, session): session.expire_all() ct = session.get(ConnectionTest, ct.id) assert ct.reverted is False + assert ct.connection_snapshot is None conn = session.scalar(select(Connection).filter_by(conn_id="no_revert_conn")) assert conn.host == "new-host.example.com" @@ -216,6 +219,7 @@ def test_patch_failed_concurrent_edit_skips_revert(self, client, session): session.expire_all() ct = session.get(ConnectionTest, ct.id) assert ct.reverted is False + assert ct.connection_snapshot is None assert "modified by another user" in ct.result_message conn = session.scalar(select(Connection).filter_by(conn_id="concurrent_conn")) assert conn.host == "third-party-host.example.com" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 66e4bcf7d90bb..5a9a7584d7dc9 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9541,6 +9541,7 @@ def test_dispatch_fails_fast_when_no_executor_supports( session.expire_all() ct = session.get(ConnectionTest, ct.id) assert ct.state == ConnectionTestState.FAILED + assert ct.connection_snapshot is None assert "No executor supports connection testing" in ct.result_message @mock.patch.dict( @@ -9837,6 +9838,7 @@ def test_reaper_reverts_connection_on_timeout_with_snapshot( ct = session.get(ConnectionTest, ct.id) assert ct.state == ConnectionTestState.FAILED assert ct.reverted is True + assert ct.connection_snapshot is None conn = session.scalar(select(Connection).filter_by(conn_id="reaper_conn")) assert conn.host == "old-host.example.com" diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py index 3138fdc82640f..c67242f760aad 100644 --- a/airflow-core/tests/unit/models/test_connection_test.py +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -187,6 +187,7 @@ def test_attempt_revert_success(self, session): session.flush() assert ct.reverted is True + assert ct.connection_snapshot is None session.refresh(conn) assert conn.host == "old-host.example.com" assert conn.login == "old_user" @@ -222,6 +223,7 @@ def test_attempt_revert_skipped_concurrent_edit(self, session): attempt_revert(ct, session=session) assert ct.reverted is False + assert ct.connection_snapshot is None assert "modified by another user" in ct.result_message assert conn.host == "third-party-host.example.com" @@ -253,6 +255,7 @@ def test_attempt_revert_skipped_concurrent_password_edit(self, session): attempt_revert(ct, session=session) assert ct.reverted is False + assert ct.connection_snapshot is None assert "modified by another user" in ct.result_message session.refresh(conn) assert conn.password == "third_party_secret" @@ -283,4 +286,5 @@ def test_attempt_revert_skipped_connection_deleted(self, session): attempt_revert(ct, session=session) assert ct.reverted is False + assert ct.connection_snapshot is None assert "no longer exists" in ct.result_message From 9c708f35e727e59a648c1809f38dfb7db15838a3 Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 5 Mar 2026 14:50:24 -0600 Subject: [PATCH 26/38] Address executor routing review feedback for async connection tests --- .../src/airflow/config_templates/config.yml | 8 --- .../src/airflow/executors/base_executor.py | 19 +++++-- .../src/airflow/executors/workloads/types.py | 3 +- .../src/airflow/jobs/scheduler_job_runner.py | 43 +++++--------- .../src/airflow/models/connection_test.py | 9 ++- .../routes/public/test_connections.py | 4 +- .../tests/unit/jobs/test_scheduler_job.py | 56 ++++++++++++++++++- 7 files changed, 94 insertions(+), 48 deletions(-) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 65232cb8288a6..cc86cc92dec1d 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2567,14 +2567,6 @@ scheduler: type: float example: ~ default: "120.0" - connection_test_dispatch_interval: - description: | - How often (in seconds) the scheduler should check for pending - connection tests and dispatch them to an executor. - version_added: 3.2.0 - type: float - example: ~ - default: "2.0" connection_test_reaper_interval: description: | How often (in seconds) the scheduler should check for stale diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 99dc9a3ffacbc..ddaf8c199c22c 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -548,13 +548,24 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task @property def slots_available(self): - """Number of new workloads (tasks and callbacks) this executor instance can accept.""" - return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks) + """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept.""" + return ( + self.parallelism + - len(self.running) + - len(self.queued_tasks) + - len(self.queued_callbacks) + - len(self.queued_connection_tests) + ) @property def slots_occupied(self): - """Number of workloads (tasks and callbacks) this executor instance is currently managing.""" - return len(self.running) + len(self.queued_tasks) + len(self.queued_callbacks) + """Number of workloads (tasks, callbacks, and connection tests) this executor instance is currently managing.""" + return ( + len(self.running) + + len(self.queued_tasks) + + len(self.queued_callbacks) + + len(self.queued_connection_tests) + ) def debug_dump(self): """Get called in response to SIGUSR2 by the scheduler.""" diff --git a/airflow-core/src/airflow/executors/workloads/types.py b/airflow-core/src/airflow/executors/workloads/types.py index 31cda7028466f..de7804f169383 100644 --- a/airflow-core/src/airflow/executors/workloads/types.py +++ b/airflow-core/src/airflow/executors/workloads/types.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, TypeAlias from airflow.models.callback import ExecutorCallback +from airflow.models.connection_test import ConnectionTest from airflow.models.taskinstance import TaskInstance if TYPE_CHECKING: @@ -37,4 +38,4 @@ # Type alias for scheduler workloads (ORM models that can be routed to executors) # Must be outside TYPE_CHECKING for use in function signatures -SchedulerWorkload: TypeAlias = TaskInstance | ExecutorCallback +SchedulerWorkload: TypeAlias = TaskInstance | ExecutorCallback | ConnectionTest diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 460e445ed67c8..bb64e0b1ed902 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1586,10 +1586,6 @@ def _run_scheduler_loop(self) -> None: action=bundle_cleanup_mgr.remove_stale_bundle_versions, ) - timers.call_regular_interval( - delay=conf.getfloat("scheduler", "connection_test_dispatch_interval", fallback=2.0), - action=self._dispatch_connection_tests, - ) timers.call_regular_interval( delay=conf.getfloat("scheduler", "connection_test_reaper_interval", fallback=30.0), action=self._reap_stale_connection_tests, @@ -1639,6 +1635,9 @@ def _run_scheduler_loop(self) -> None: # Route ExecutorCallback workloads to executors (similar to task routing) self._enqueue_executor_callbacks(session) + # Dispatch pending connection tests to executors + self._dispatch_connection_tests(session=session) + # Heartbeat the scheduler periodically perform_heartbeat( job=self.job, heartbeat_callback=self.heartbeat_callback, only_if_necessary=True @@ -3124,16 +3123,21 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]: session.add(warning) existing_warned_dag_ids.add(warning.dag_id) - @provide_session - def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None: + def _dispatch_connection_tests(self, *, session: Session) -> None: """Dispatch pending connection tests to executors that support them.""" max_concurrency = conf.getint("core", "max_connection_test_concurrency", fallback=4) timeout = conf.getint("core", "connection_test_timeout", fallback=60) + 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(ConnectionTest.id)).where(ConnectionTest.state.in_(ACTIVE_STATES)) ) - budget = max_concurrency - (active_count or 0) + concurrency_budget = max_concurrency - (active_count or 0) + budget = min(concurrency_budget, parallelism_budget) if budget <= 0: return @@ -3150,7 +3154,9 @@ def _dispatch_connection_tests(self, *, session: Session = NEW_SESSION) -> None: return for ct in pending_tests: - executor = self._find_executor_for_connection_test(ct.executor) + 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}'" @@ -3198,27 +3204,6 @@ def _reap_stale_connection_tests(self, *, session: Session = NEW_SESSION) -> Non session.flush() - def _find_executor_for_connection_test(self, executor_name: str | None) -> BaseExecutor | None: - """Find an executor that supports connection testing, optionally matching by name.""" - if executor_name is not None: - for executor in self.executors: - if ( - executor.supports_connection_test - and executor.name - and executor_name - in ( - executor.name.alias, - executor.name.module_path, - executor.name.module_path.split(".")[-1], - ) - ): - return executor - return None - for executor in self.executors: - if executor.supports_connection_test: - return executor - return None - def _executor_to_workloads( self, workloads: Iterable[SchedulerWorkload], diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 0ba8a6ff93bde..5fc20b5a6a153 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -90,6 +90,14 @@ def __init__( def __repr__(self) -> str: return f"" + def get_executor_name(self) -> str | None: + """Return the executor name for scheduler routing.""" + return self.executor + + def get_dag_id(self) -> None: + """Return None — connection tests are not associated with any DAG.""" + return None + def run_connection_test(*, conn: Connection) -> tuple[bool, str]: """ @@ -116,7 +124,6 @@ def run_connection_test(*, conn: Connection) -> tuple[bool, str]: "_extra", "is_encrypted", "is_extra_encrypted", - "team_name", ) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 4f15a48dfaa47..a14b70fd4e3ea 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -1337,7 +1337,7 @@ def test_save_and_test_passes_executor_parameter(self, test_client, session): """PATCH /{connection_id}/test passes the executor parameter to the ConnectionTest.""" self.create_connection() response = test_client.patch( - f"/connections/{TEST_CONN_ID}/test?executor=team_a", + f"/connections/{TEST_CONN_ID}/test?executor=my_executor", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, @@ -1349,7 +1349,7 @@ def test_save_and_test_passes_executor_parameter(self, test_client, session): ct = session.scalar(select(ConnectionTest).filter_by(token=token)) assert ct is not None - assert ct.executor == "team_a" + assert ct.executor == "my_executor" @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_save_and_test_passes_queue_parameter(self, test_client, session): diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 5a9a7584d7dc9..75756adc4145c 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9644,10 +9644,10 @@ def test_dispatch_fails_fast_for_unserved_executor( """Tests requesting an executor no team serves are failed immediately.""" with mock.patch.object( scheduler_job_runner_for_connection_tests, - "_find_executor_for_connection_test", + "_try_to_load_executor", return_value=None, ): - ct = ConnectionTest(connection_id="test_conn", executor="nonexistent_queue") + ct = ConnectionTest(connection_id="test_conn", executor="nonexistent_executor") session.add(ct) session.commit() @@ -9656,7 +9656,7 @@ def test_dispatch_fails_fast_for_unserved_executor( session.expire_all() ct = session.get(ConnectionTest, ct.id) assert ct.state == ConnectionTestState.FAILED - assert "nonexistent_queue" in ct.result_message + assert "nonexistent_executor" in ct.result_message @mock.patch.dict( os.environ, @@ -9768,6 +9768,56 @@ def test_dispatch_executor_matched_by_class_name(self, session): assert len(executor_b.queued_connection_tests) == 1 assert len(executor_a.queued_connection_tests) == 0 + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + "AIRFLOW__CORE__PARALLELISM": "1", + }, + ) + def test_dispatch_respects_parallelism_budget(self, scheduler_job_runner_for_connection_tests, session): + """Connection tests are not dispatched when core.parallelism is exhausted.""" + executor = scheduler_job_runner_for_connection_tests.executor + # Simulate 1 running task so all parallelism slots are occupied + executor.running = {"fake_task_key"} + + ct = ConnectionTest(connection_id="test_conn") + session.add(ct) + session.commit() + + scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.PENDING + + @mock.patch.dict( + os.environ, + { + "AIRFLOW__CORE__MAX_CONNECTION_TEST_CONCURRENCY": "4", + "AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60", + }, + ) + def test_dispatch_fails_when_executor_does_not_support_connection_test( + self, scheduler_job_runner_for_connection_tests, session + ): + """When the resolved executor does not support connection tests, the test is failed gracefully.""" + executor = scheduler_job_runner_for_connection_tests.executor + executor.supports_connection_test = False + + ct = ConnectionTest(connection_id="test_conn") + session.add(ct) + session.commit() + + scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + + session.expire_all() + ct = session.get(ConnectionTest, ct.id) + assert ct.state == ConnectionTestState.FAILED + assert ct.connection_snapshot is None + assert "No executor supports connection testing" in ct.result_message + class TestReapStaleConnectionTests: @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) From 9b26a15d844d3fffa09f7a1e4d8a95dd97e6042a Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 5 Mar 2026 16:58:02 -0600 Subject: [PATCH 27/38] fix reverted flag not updating issue --- .../api_fastapi/execution_api/routes/connection_tests.py | 2 +- airflow-core/src/airflow/models/connection_test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py index 13e5a45fe7c74..d40026b7ef1e1 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -71,5 +71,5 @@ def patch_connection_test( if body.state == ConnectionTestState.FAILED and ct.connection_snapshot: attempt_revert(ct, session=session) - else: + elif body.state == ConnectionTestState.SUCCESS: ct.connection_snapshot = None diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 5fc20b5a6a153..64c306f9b53cf 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -72,7 +72,7 @@ class ConnectionTest(Base): ) executor: Mapped[str | None] = mapped_column(String(256), nullable=True) queue: Mapped[str | None] = mapped_column(String(256), nullable=True) - connection_snapshot: Mapped[dict | None] = mapped_column(JSON, nullable=True) + connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True) reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0") __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),) From 9931da78b374273b286db4c8ed2c79de7896f7d5 Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 5 Mar 2026 17:14:07 -0600 Subject: [PATCH 28/38] refactor namings --- .../src/airflow/jobs/scheduler_job_runner.py | 8 +++---- .../tests/unit/jobs/test_scheduler_job.py | 24 +++++++++---------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index bb64e0b1ed902..563b83cde658f 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1635,8 +1635,8 @@ def _run_scheduler_loop(self) -> None: # Route ExecutorCallback workloads to executors (similar to task routing) self._enqueue_executor_callbacks(session) - # Dispatch pending connection tests to executors - self._dispatch_connection_tests(session=session) + # Enqueue pending connection tests to executors + self._enqueue_connection_tests(session=session) # Heartbeat the scheduler periodically perform_heartbeat( @@ -3123,8 +3123,8 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]: session.add(warning) existing_warned_dag_ids.add(warning.dag_id) - def _dispatch_connection_tests(self, *, session: Session) -> None: - """Dispatch pending connection tests to executors that support them.""" + 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) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 75756adc4145c..05c1b33f220fb 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -9486,7 +9486,7 @@ def test_dispatch_pending_tests(self, scheduler_job_runner_for_connection_tests, session.commit() assert ct.state == ConnectionTestState.PENDING - scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() ct = session.get(ConnectionTest, ct.id) @@ -9510,7 +9510,7 @@ def test_dispatch_respects_concurrency_limit(self, scheduler_job_runner_for_conn session.add(ct_pending) session.commit() - scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() ct_pending = session.get(ConnectionTest, ct_pending.id) @@ -9536,7 +9536,7 @@ def test_dispatch_fails_fast_when_no_executor_supports( session.add(ct) session.commit() - scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() ct = session.get(ConnectionTest, ct.id) @@ -9559,7 +9559,7 @@ def test_dispatch_with_unmatched_executor_fails_fast( session.add(ct) session.commit() - scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() ct = session.get(ConnectionTest, ct.id) @@ -9589,7 +9589,7 @@ def test_dispatch_budget_dispatches_up_to_remaining_slots( session.commit() pending_ids = [ct.id for ct in pending_tests] - scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() states = [session.get(ConnectionTest, pid).state for pid in pending_ids] @@ -9624,7 +9624,7 @@ def test_dispatch_order_is_fifo_by_created_at(self, scheduler_job_runner_for_con session.commit() - scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() assert session.get(ConnectionTest, ct_old.id).state == ConnectionTestState.QUEUED @@ -9651,7 +9651,7 @@ def test_dispatch_fails_fast_for_unserved_executor( session.add(ct) session.commit() - scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() ct = session.get(ConnectionTest, ct.id) @@ -9692,7 +9692,7 @@ def test_dispatch_executor_matched_by_alias(self, session): session.add(ct) session.commit() - runner._dispatch_connection_tests(session=session) + runner._enqueue_connection_tests(session=session) assert len(executor_b.queued_connection_tests) == 1 assert len(executor_a.queued_connection_tests) == 0 @@ -9731,7 +9731,7 @@ def test_dispatch_executor_matched_by_module_path(self, session): session.add(ct) session.commit() - runner._dispatch_connection_tests(session=session) + runner._enqueue_connection_tests(session=session) assert len(executor_b.queued_connection_tests) == 1 assert len(executor_a.queued_connection_tests) == 0 @@ -9763,7 +9763,7 @@ def test_dispatch_executor_matched_by_class_name(self, session): session.add(ct) session.commit() - runner._dispatch_connection_tests(session=session) + runner._enqueue_connection_tests(session=session) assert len(executor_b.queued_connection_tests) == 1 assert len(executor_a.queued_connection_tests) == 0 @@ -9786,7 +9786,7 @@ def test_dispatch_respects_parallelism_budget(self, scheduler_job_runner_for_con session.add(ct) session.commit() - scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() ct = session.get(ConnectionTest, ct.id) @@ -9810,7 +9810,7 @@ def test_dispatch_fails_when_executor_does_not_support_connection_test( session.add(ct) session.commit() - scheduler_job_runner_for_connection_tests._dispatch_connection_tests(session=session) + scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() ct = session.get(ConnectionTest, ct.id) From 874b489d6ab873c101e60fa9e22cf826151740a5 Mon Sep 17 00:00:00 2001 From: Anish Date: Tue, 10 Mar 2026 18:29:46 -0500 Subject: [PATCH 29/38] regenerated erd --- airflow-core/docs/img/airflow_erd.sha256 | 1 + airflow-core/docs/img/airflow_erd.svg | 3411 ++++++++++++++++++++++ 2 files changed, 3412 insertions(+) create mode 100644 airflow-core/docs/img/airflow_erd.sha256 create mode 100644 airflow-core/docs/img/airflow_erd.svg diff --git a/airflow-core/docs/img/airflow_erd.sha256 b/airflow-core/docs/img/airflow_erd.sha256 new file mode 100644 index 0000000000000..605175c74b83f --- /dev/null +++ b/airflow-core/docs/img/airflow_erd.sha256 @@ -0,0 +1 @@ +ecba28581229c5a0d1107d162a9caac0254e0bf60845f19757eb3349b43b78e5 \ No newline at end of file diff --git a/airflow-core/docs/img/airflow_erd.svg b/airflow-core/docs/img/airflow_erd.svg new file mode 100644 index 0000000000000..7bb08eac9306c --- /dev/null +++ b/airflow-core/docs/img/airflow_erd.svg @@ -0,0 +1,3411 @@ + + + + + + +%3 + + + +dag_bundle_team + +dag_bundle_team + +dag_bundle_name + + [VARCHAR(250)] + NOT NULL + +team_name + + [VARCHAR(50)] + NOT NULL + + + +dag_bundle + +dag_bundle + +name + + [VARCHAR(250)] + NOT NULL + +active + + [BOOLEAN] + +last_refreshed + + [TIMESTAMP] + +signed_url_template + + [TEXT] + +template_params + + [JSON] + +version + + [VARCHAR(200)] + + + +dag_bundle:name--dag_bundle_team:dag_bundle_name + +0..N +1 + + + +dag + +dag + +dag_id + + [VARCHAR(250)] + NOT NULL + +allowed_run_types + + [JSON] + +asset_expression + + [JSON] + +bundle_name + + [VARCHAR(250)] + NOT NULL + +bundle_version + + [VARCHAR(200)] + +dag_display_name + + [VARCHAR(2000)] + +deadline + + [JSON] + +description + + [TEXT] + +exceeds_max_non_backfill + + [BOOLEAN] + NOT NULL + +fail_fast + + [BOOLEAN] + NOT NULL + +fileloc + + [VARCHAR(2000)] + +has_import_errors + + [BOOLEAN] + NOT NULL + +has_task_concurrency_limits + + [BOOLEAN] + NOT NULL + +is_paused + + [BOOLEAN] + NOT NULL + +is_stale + + [BOOLEAN] + NOT NULL + +last_expired + + [TIMESTAMP] + +last_parse_duration + + [DOUBLE PRECISION] + +last_parsed_time + + [TIMESTAMP] + +max_active_runs + + [INTEGER] + +max_active_tasks + + [INTEGER] + NOT NULL + +max_consecutive_failed_dag_runs + + [INTEGER] + NOT NULL + +next_dagrun + + [TIMESTAMP] + +next_dagrun_create_after + + [TIMESTAMP] + +next_dagrun_data_interval_end + + [TIMESTAMP] + +next_dagrun_data_interval_start + + [TIMESTAMP] + +next_dagrun_partition_date + + [TIMESTAMP] + +next_dagrun_partition_key + + [VARCHAR(255)] + +owners + + [VARCHAR(2000)] + +relative_fileloc + + [VARCHAR(2000)] + +timetable_description + + [VARCHAR(1000)] + +timetable_partitioned + + [BOOLEAN] + NOT NULL + +timetable_summary + + [TEXT] + +timetable_type + + [VARCHAR(255)] + NOT NULL + + + +dag_bundle:name--dag:bundle_name + +0..N +1 + + + +team + +team + +name + + [VARCHAR(50)] + NOT NULL + + + +team:name--dag_bundle_team:team_name + +0..N +1 + + + +connection + +connection + +id + + [INTEGER] + NOT NULL + +conn_id + + [VARCHAR(250)] + NOT NULL + +conn_type + + [VARCHAR(500)] + NOT NULL + +description + + [TEXT] + +extra + + [TEXT] + +host + + [VARCHAR(500)] + +is_encrypted + + [BOOLEAN] + NOT NULL + +is_extra_encrypted + + [BOOLEAN] + NOT NULL + +login + + [TEXT] + +password + + [TEXT] + +port + + [INTEGER] + +schema + + [VARCHAR(500)] + +team_name + + [VARCHAR(50)] + + + +team:name--connection:team_name + +0..N +{0,1} + + + +slot_pool + +slot_pool + +id + + [INTEGER] + NOT NULL + +description + + [TEXT] + +include_deferred + + [BOOLEAN] + NOT NULL + +pool + + [VARCHAR(256)] + NOT NULL + +slots + + [INTEGER] + NOT NULL + +team_name + + [VARCHAR(50)] + + + +team:name--slot_pool:team_name + +0..N +{0,1} + + + +variable + +variable + +id + + [INTEGER] + NOT NULL + +description + + [TEXT] + +is_encrypted + + [BOOLEAN] + NOT NULL + +key + + [VARCHAR(250)] + NOT NULL + +team_name + + [VARCHAR(50)] + +val + + [TEXT] + NOT NULL + + + +team:name--variable:team_name + +0..N +{0,1} + + + +trigger + +trigger + +id + + [INTEGER] + NOT NULL + +classpath + + [VARCHAR(1000)] + NOT NULL + +created_date + + [TIMESTAMP] + NOT NULL + +kwargs + + [TEXT] + NOT NULL + +queue + + [VARCHAR(256)] + +triggerer_id + + [INTEGER] + + + +callback + +callback + +id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +data + + [JSONB] + NOT NULL + +fetch_method + + [VARCHAR(20)] + NOT NULL + +output + + [TEXT] + +priority_weight + + [INTEGER] + NOT NULL + +state + + [VARCHAR(10)] + +trigger_id + + [INTEGER] + +type + + [VARCHAR(20)] + NOT NULL + + + +trigger:id--callback:trigger_id + +0..N +{0,1} + + + +asset_watcher + +asset_watcher + +asset_id + + [INTEGER] + NOT NULL + +trigger_id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + + + +trigger:id--asset_watcher:trigger_id + +0..N +1 + + + +task_instance + +task_instance + +id + + [UUID] + NOT NULL + +context_carrier + + [JSONB] + +custom_operator_name + + [VARCHAR(1000)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + +duration + + [DOUBLE PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + NOT NULL + +external_executor_id + + [TEXT] + +hostname + + [VARCHAR(1000)] + NOT NULL + +last_heartbeat_at + + [TIMESTAMP] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + NOT NULL + +next_kwargs + + [JSONB] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + NOT NULL + +queue + + [VARCHAR(256)] + NOT NULL + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + NOT NULL + +updated_at + + [TIMESTAMP] + + + +trigger:id--task_instance:trigger_id + +0..N +{0,1} + + + +deadline + +deadline + +id + + [UUID] + NOT NULL + +callback_id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +dagrun_id + + [INTEGER] + +deadline_alert_id + + [UUID] + +deadline_time + + [TIMESTAMP] + NOT NULL + +last_updated_at + + [TIMESTAMP] + NOT NULL + +missed + + [BOOLEAN] + NOT NULL + + + +callback:id--deadline:callback_id + +0..N +1 + + + +asset_alias + +asset_alias + +id + + [INTEGER] + NOT NULL + +group + + [VARCHAR(1500)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + + + +asset_alias_asset + +asset_alias_asset + +alias_id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL + + + +asset_alias:id--asset_alias_asset:alias_id + +0..N +1 + + + +asset_alias_asset_event + +asset_alias_asset_event + +alias_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL + + + +asset_alias:id--asset_alias_asset_event:alias_id + +0..N +1 + + + +dag_schedule_asset_alias_reference + +dag_schedule_asset_alias_reference + +alias_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + + + +asset_alias:id--dag_schedule_asset_alias_reference:alias_id + +0..N +1 + + + +asset + +asset + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +extra + + [JSON] + NOT NULL + +group + + [VARCHAR(1500)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL + + + +asset:id--asset_alias_asset:asset_id + +0..N +1 + + + +asset:id--asset_watcher:asset_id + +0..N +1 + + + +asset_active + +asset_active + +name + + [VARCHAR(1500)] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL + + + +asset:uri--asset_active:uri + +1 +1 + + + +asset:name--asset_active:name + +1 +1 + + + +dag_schedule_asset_reference + +dag_schedule_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + + + +asset:id--dag_schedule_asset_reference:asset_id + +0..N +1 + + + +task_outlet_asset_reference + +task_outlet_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + + + +asset:id--task_outlet_asset_reference:asset_id + +0..N +1 + + + +task_inlet_asset_reference + +task_inlet_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + + + +asset:id--task_inlet_asset_reference:asset_id + +0..N +1 + + + +asset_dag_run_queue + +asset_dag_run_queue + +asset_id + + [INTEGER] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + + + +asset:id--asset_dag_run_queue:asset_id + +0..N +1 + + + +asset_event + +asset_event + +id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL + +extra + + [JSON] + NOT NULL + +partition_key + + [VARCHAR(250)] + +source_dag_id + + [VARCHAR(250)] + +source_map_index + + [INTEGER] + +source_run_id + + [VARCHAR(250)] + +source_task_id + + [VARCHAR(250)] + +timestamp + + [TIMESTAMP] + NOT NULL + + + +asset_event:id--asset_alias_asset_event:event_id + +0..N +1 + + + +dagrun_asset_event + +dagrun_asset_event + +dag_run_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL + + + +asset_event:id--dagrun_asset_event:event_id + +0..N +1 + + + +job + +job + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + +end_date + + [TIMESTAMP] + +executor_class + + [VARCHAR(500)] + +hostname + + [VARCHAR(500)] + +job_type + + [VARCHAR(30)] + +latest_heartbeat + + [TIMESTAMP] + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +unixname + + [VARCHAR(1000)] + + + +partitioned_asset_key_log + +partitioned_asset_key_log + +id + + [INTEGER] + NOT NULL + +asset_event_id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL + +asset_partition_dag_run_id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +source_partition_key + + [VARCHAR(250)] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +target_partition_key + + [VARCHAR(250)] + NOT NULL + + + +log + +log + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + +dttm + + [TIMESTAMP] + NOT NULL + +event + + [VARCHAR(60)] + NOT NULL + +extra + + [TEXT] + +logical_date + + [TIMESTAMP] + +map_index + + [INTEGER] + +owner + + [VARCHAR(500)] + +owner_display_name + + [VARCHAR(500)] + +run_id + + [VARCHAR(250)] + +task_id + + [VARCHAR(250)] + +try_number + + [INTEGER] + + + +connection_test + +connection_test + +id + + [UUID] + NOT NULL + +connection_id + + [VARCHAR(250)] + NOT NULL + +connection_snapshot + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +executor + + [VARCHAR(256)] + +queue + + [VARCHAR(256)] + +result_message + + [TEXT] + +reverted + + [BOOLEAN] + NOT NULL + +state + + [VARCHAR(10)] + NOT NULL + +token + + [VARCHAR(64)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + + + +dag_priority_parsing_request + +dag_priority_parsing_request + +id + + [VARCHAR(32)] + NOT NULL + +bundle_name + + [VARCHAR(250)] + NOT NULL + +relative_fileloc + + [VARCHAR(2000)] + NOT NULL + + + +import_error + +import_error + +id + + [INTEGER] + NOT NULL + +bundle_name + + [VARCHAR(250)] + +filename + + [VARCHAR(1024)] + +stacktrace + + [TEXT] + +timestamp + + [TIMESTAMP] + + + +revoked_token + +revoked_token + +jti + + [VARCHAR(32)] + NOT NULL + +exp + + [TIMESTAMP] + NOT NULL + + + +dag_schedule_asset_name_reference + +dag_schedule_asset_name_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + + + +dag:dag_id--dag_schedule_asset_name_reference:dag_id + +0..N +1 + + + +dag_schedule_asset_uri_reference + +dag_schedule_asset_uri_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + + + +dag:dag_id--dag_schedule_asset_uri_reference:dag_id + +0..N +1 + + + +dag:dag_id--dag_schedule_asset_alias_reference:dag_id + +0..N +1 + + + +dag:dag_id--dag_schedule_asset_reference:dag_id + +0..N +1 + + + +dag:dag_id--task_outlet_asset_reference:dag_id + +0..N +1 + + + +dag:dag_id--task_inlet_asset_reference:dag_id + +0..N +1 + + + +dag:dag_id--asset_dag_run_queue:target_dag_id + +0..N +1 + + + +dag_version + +dag_version + +id + + [UUID] + NOT NULL + +bundle_name + + [VARCHAR(250)] + +bundle_version + + [VARCHAR(250)] + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +last_updated + + [TIMESTAMP] + NOT NULL + +version_number + + [INTEGER] + NOT NULL + + + +dag:dag_id--dag_version:dag_id + +0..N +1 + + + +dag_tag + +dag_tag + +dag_id + + [VARCHAR(250)] + NOT NULL + +name + + [VARCHAR(100)] + NOT NULL + + + +dag:dag_id--dag_tag:dag_id + +0..N +1 + + + +dag_owner_attributes + +dag_owner_attributes + +dag_id + + [VARCHAR(250)] + NOT NULL + +owner + + [VARCHAR(500)] + NOT NULL + +link + + [VARCHAR(500)] + NOT NULL + + + +dag:dag_id--dag_owner_attributes:dag_id + +0..N +1 + + + +dag_warning + +dag_warning + +dag_id + + [VARCHAR(250)] + NOT NULL + +warning_type + + [VARCHAR(50)] + NOT NULL + +message + + [TEXT] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL + + + +dag:dag_id--dag_warning:dag_id + +0..N +1 + + + +dag_favorite + +dag_favorite + +dag_id + + [VARCHAR(250)] + NOT NULL + +user_id + + [VARCHAR(250)] + NOT NULL + + + +dag:dag_id--dag_favorite:dag_id + +0..N +1 + + + +dag_run + +dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + +bundle_version + + [VARCHAR(250)] + +clear_number + + [INTEGER] + NOT NULL + +conf + + [JSONB] + +context_carrier + + [JSONB] + +created_at + + [TIMESTAMP] + NOT NULL + +created_dag_version_id + + [UUID] + +creating_job_id + + [INTEGER] + +dag_id + + [VARCHAR(250)] + NOT NULL + +data_interval_end + + [TIMESTAMP] + +data_interval_start + + [TIMESTAMP] + +end_date + + [TIMESTAMP] + +last_scheduling_decision + + [TIMESTAMP] + +log_template_id + + [INTEGER] + NOT NULL + +logical_date + + [TIMESTAMP] + +partition_date + + [TIMESTAMP] + +partition_key + + [VARCHAR(250)] + +queued_at + + [TIMESTAMP] + +run_after + + [TIMESTAMP] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +run_type + + [VARCHAR(50)] + NOT NULL + +scheduled_by_job_id + + [INTEGER] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(50)] + NOT NULL + +triggered_by + + [VARCHAR(50)] + +triggering_user_name + + [VARCHAR(512)] + +updated_at + + [TIMESTAMP] + NOT NULL + + + +dag_version:id--dag_run:created_dag_version_id + +0..N +{0,1} + + + +dag_code + +dag_code + +id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + NOT NULL + +fileloc + + [VARCHAR(2000)] + NOT NULL + +last_updated + + [TIMESTAMP] + NOT NULL + +source_code + + [TEXT] + NOT NULL + +source_code_hash + + [VARCHAR(32)] + NOT NULL + + + +dag_version:id--dag_code:dag_version_id + +0..N +1 + + + +serialized_dag + +serialized_dag + +id + + [UUID] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +dag_hash + + [VARCHAR(32)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + NOT NULL + +data + + [JSONB] + +data_compressed + + [BYTEA] + +last_updated + + [TIMESTAMP] + NOT NULL + + + +dag_version:id--serialized_dag:dag_version_id + +0..N +1 + + + +dag_version:id--task_instance:dag_version_id + +0..N +{0,1} + + + +log_template + +log_template + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +elasticsearch_id + + [TEXT] + NOT NULL + +filename + + [TEXT] + NOT NULL + + + +log_template:id--dag_run:log_template_id + +0..N +1 + + + +dag_run:id--dagrun_asset_event:dag_run_id + +0..N +1 + + + +asset_partition_dag_run + +asset_partition_dag_run + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +created_dag_run_id + + [INTEGER] + +partition_key + + [VARCHAR(250)] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + + + +dag_run:id--asset_partition_dag_run:created_dag_run_id + +0..N +{0,1} + + + +dag_run:run_id--task_instance:run_id + +0..N +1 + + + +dag_run:dag_id--task_instance:dag_id + +0..N +1 + + + +backfill_dag_run + +backfill_dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + NOT NULL + +dag_run_id + + [INTEGER] + +exception_reason + + [VARCHAR(250)] + +logical_date + + [TIMESTAMP] + +partition_key + + [VARCHAR(250)] + +sort_ordinal + + [INTEGER] + NOT NULL + + + +dag_run:id--backfill_dag_run:dag_run_id + +0..N +{0,1} + + + +dag_run_note + +dag_run_note + +dag_run_id + + [INTEGER] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] + + + +dag_run:id--dag_run_note:dag_run_id + +1 +1 + + + +dag_run:id--deadline:dagrun_id + +0..N +{0,1} + + + +backfill + +backfill + +id + + [INTEGER] + NOT NULL + +completed_at + + [TIMESTAMP] + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_run_conf + + [JSON] + NOT NULL + +from_date + + [TIMESTAMP] + NOT NULL + +is_paused + + [BOOLEAN] + +max_active_runs + + [INTEGER] + NOT NULL + +reprocess_behavior + + [VARCHAR(250)] + NOT NULL + +to_date + + [TIMESTAMP] + NOT NULL + +triggering_user_name + + [VARCHAR(512)] + +updated_at + + [TIMESTAMP] + NOT NULL + + + +backfill:id--dag_run:backfill_id + +0..N +{0,1} + + + +backfill:id--backfill_dag_run:backfill_id + +0..N +1 + + + +deadline_alert + +deadline_alert + +id + + [UUID] + NOT NULL + +callback_def + + [JSON] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +description + + [TEXT] + +interval + + [DOUBLE PRECISION] + NOT NULL + +name + + [VARCHAR(250)] + +reference + + [JSON] + NOT NULL + +serialized_dag_id + + [UUID] + NOT NULL + + + +serialized_dag:id--deadline_alert:serialized_dag_id + +0..N +1 + + + +deadline_alert:id--deadline:deadline_alert_id + +0..N +{0,1} + + + +hitl_detail + +hitl_detail + +ti_id + + [UUID] + NOT NULL + +assignees + + [JSON] + +body + + [TEXT] + +chosen_options + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +defaults + + [JSON] + +multiple + + [BOOLEAN] + +options + + [JSON] + NOT NULL + +params + + [JSON] + NOT NULL + +params_input + + [JSON] + NOT NULL + +responded_at + + [TIMESTAMP] + +responded_by + + [JSON] + +subject + + [TEXT] + NOT NULL + + + +task_instance:id--hitl_detail:ti_id + +1 +1 + + + +task_map + +task_map + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +keys + + [JSONB] + +length + + [INTEGER] + NOT NULL + + + +task_instance:dag_id--task_map:dag_id + +0..N +1 + + + +task_instance:task_id--task_map:task_id + +0..N +1 + + + +task_instance:run_id--task_map:run_id + +0..N +1 + + + +task_instance:map_index--task_map:map_index + +0..N +1 + + + +task_reschedule + +task_reschedule + +id + + [INTEGER] + NOT NULL + +duration + + [INTEGER] + NOT NULL + +end_date + + [TIMESTAMP] + NOT NULL + +reschedule_date + + [TIMESTAMP] + NOT NULL + +start_date + + [TIMESTAMP] + NOT NULL + +ti_id + + [UUID] + NOT NULL + + + +task_instance:id--task_reschedule:ti_id + +0..N +1 + + + +xcom + +xcom + +dag_run_id + + [INTEGER] + NOT NULL + +key + + [VARCHAR(512)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL + +value + + [JSONB] + + + +task_instance:dag_id--xcom:dag_id + +0..N +1 + + + +task_instance:map_index--xcom:map_index + +0..N +1 + + + +task_instance:task_id--xcom:task_id + +0..N +1 + + + +task_instance:run_id--xcom:run_id + +0..N +1 + + + +task_instance_note + +task_instance_note + +ti_id + + [UUID] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] + + + +task_instance:id--task_instance_note:ti_id + +1 +1 + + + +task_instance_history + +task_instance_history + +task_instance_id + + [UUID] + NOT NULL + +context_carrier + + [JSONB] + +custom_operator_name + + [VARCHAR(1000)] + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + +duration + + [DOUBLE PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + +external_executor_id + + [TEXT] + +hostname + + [VARCHAR(1000)] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + +next_kwargs + + [JSONB] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + +queue + + [VARCHAR(256)] + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + +updated_at + + [TIMESTAMP] + + + +task_instance:dag_id--task_instance_history:dag_id + +0..N +1 + + + +task_instance:run_id--task_instance_history:run_id + +0..N +1 + + + +task_instance:map_index--task_instance_history:map_index + +0..N +1 + + + +task_instance:task_id--task_instance_history:task_id + +0..N +1 + + + +rendered_task_instance_fields + +rendered_task_instance_fields + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +k8s_pod_yaml + + [JSON] + +rendered_fields + + [JSON] + NOT NULL + + + +task_instance:task_id--rendered_task_instance_fields:task_id + +0..N +1 + + + +task_instance:dag_id--rendered_task_instance_fields:dag_id + +0..N +1 + + + +task_instance:map_index--rendered_task_instance_fields:map_index + +0..N +1 + + + +task_instance:run_id--rendered_task_instance_fields:run_id + +0..N +1 + + + +edge_job + +edge_job + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +try_number + + [INTEGER] + NOT NULL + +command + + [VARCHAR(2048)] + NOT NULL + +concurrency_slots + + [INTEGER] + NOT NULL + +edge_worker + + [VARCHAR(64)] + +last_update + + [TIMESTAMP] + +queue + + [VARCHAR(256)] + NOT NULL + +queued_dttm + + [TIMESTAMP] + +state + + [VARCHAR(20)] + NOT NULL + + + +hitl_detail_history + +hitl_detail_history + +ti_history_id + + [UUID] + NOT NULL + +assignees + + [JSON] + +body + + [TEXT] + +chosen_options + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +defaults + + [JSON] + +multiple + + [BOOLEAN] + +options + + [JSON] + NOT NULL + +params + + [JSON] + NOT NULL + +params_input + + [JSON] + NOT NULL + +responded_at + + [TIMESTAMP] + +responded_by + + [JSON] + +subject + + [TEXT] + NOT NULL + + + +task_instance_history:task_instance_id--hitl_detail_history:ti_history_id + +1 +1 + + + +alembic_version + +alembic_version + +version_num + + [VARCHAR(32)] + NOT NULL + + + +edge_logs + +edge_logs + +dag_id + + [VARCHAR(250)] + NOT NULL + +log_chunk_time + + [TIMESTAMP] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +try_number + + [INTEGER] + NOT NULL + +log_chunk_data + + [TEXT] + NOT NULL + + + +edge_worker + +edge_worker + +worker_name + + [VARCHAR(64)] + NOT NULL + +concurrency + + [INTEGER] + +first_online + + [TIMESTAMP] + +jobs_active + + [INTEGER] + NOT NULL + +jobs_failed + + [INTEGER] + NOT NULL + +jobs_success + + [INTEGER] + NOT NULL + +jobs_taken + + [INTEGER] + NOT NULL + +last_update + + [TIMESTAMP] + +maintenance_comment + + [VARCHAR(1024)] + +queues + + [VARCHAR(256)] + +state + + [VARCHAR(20)] + NOT NULL + +sysinfo + + [VARCHAR(256)] + + + +alembic_version_edge3 + +alembic_version_edge3 + +version_num + + [VARCHAR(32)] + NOT NULL + + + +alembic_version_fab + +alembic_version_fab + +version_num + + [VARCHAR(32)] + NOT NULL + + + +ab_user + +ab_user + +id + + [INTEGER] + NOT NULL + +active + + [BOOLEAN] + +changed_by_fk + + [INTEGER] + +changed_on + + [TIMESTAMP] + +created_by_fk + + [INTEGER] + +created_on + + [TIMESTAMP] + +email + + [VARCHAR(512)] + NOT NULL + +fail_login_count + + [INTEGER] + +first_name + + [VARCHAR(256)] + NOT NULL + +last_login + + [TIMESTAMP] + +last_name + + [VARCHAR(256)] + NOT NULL + +login_count + + [INTEGER] + +password + + [VARCHAR(256)] + +username + + [VARCHAR(512)] + NOT NULL + + + +ab_user:id--ab_user:changed_by_fk + +0..N +{0,1} + + + +ab_user:id--ab_user:created_by_fk + +0..N +{0,1} + + + +ab_user_role + +ab_user_role + +id + + [INTEGER] + NOT NULL + +role_id + + [INTEGER] + +user_id + + [INTEGER] + + + +ab_user:id--ab_user_role:user_id + +0..N +{0,1} + + + +ab_user_group + +ab_user_group + +id + + [INTEGER] + NOT NULL + +group_id + + [INTEGER] + +user_id + + [INTEGER] + + + +ab_user:id--ab_user_group:user_id + +0..N +{0,1} + + + +ab_register_user + +ab_register_user + +id + + [INTEGER] + NOT NULL + +email + + [VARCHAR(512)] + NOT NULL + +first_name + + [VARCHAR(256)] + NOT NULL + +last_name + + [VARCHAR(256)] + NOT NULL + +password + + [VARCHAR(256)] + +registration_date + + [TIMESTAMP] + +registration_hash + + [VARCHAR(256)] + +username + + [VARCHAR(512)] + NOT NULL + + + +ab_group + +ab_group + +id + + [INTEGER] + NOT NULL + +description + + [VARCHAR(512)] + +label + + [VARCHAR(150)] + +name + + [VARCHAR(100)] + NOT NULL + + + +ab_group_role + +ab_group_role + +id + + [INTEGER] + NOT NULL + +group_id + + [INTEGER] + +role_id + + [INTEGER] + + + +ab_group:id--ab_group_role:group_id + +0..N +{0,1} + + + +ab_group:id--ab_user_group:group_id + +0..N +{0,1} + + + +ab_role + +ab_role + +id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(64)] + NOT NULL + + + +ab_role:id--ab_group_role:role_id + +0..N +{0,1} + + + +ab_role:id--ab_user_role:role_id + +0..N +{0,1} + + + +ab_permission_view_role + +ab_permission_view_role + +id + + [INTEGER] + NOT NULL + +permission_view_id + + [INTEGER] + +role_id + + [INTEGER] + + + +ab_role:id--ab_permission_view_role:role_id + +0..N +{0,1} + + + +ab_permission + +ab_permission + +id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(100)] + NOT NULL + + + +ab_permission_view + +ab_permission_view + +id + + [INTEGER] + NOT NULL + +permission_id + + [INTEGER] + NOT NULL + +view_menu_id + + [INTEGER] + NOT NULL + + + +ab_permission:id--ab_permission_view:permission_id + +0..N +1 + + + +ab_permission_view:id--ab_permission_view_role:permission_view_id + +0..N +{0,1} + + + +ab_view_menu + +ab_view_menu + +id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(250)] + NOT NULL + + + +ab_view_menu:id--ab_permission_view:view_menu_id + +0..N +1 + + + +session + +session + +id + + [INTEGER] + NOT NULL + +data + + [BYTEA] + +expiry + + [TIMESTAMP] + +session_id + + [VARCHAR(255)] + + + From 7b73a81460ac11d1d51d0b3fcdcfbded692f4155 Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 12 Mar 2026 09:51:33 -0500 Subject: [PATCH 30/38] regenrated erd --- airflow-core/docs/img/airflow_erd.sha256 | 2 +- airflow-core/docs/img/airflow_erd.svg | 5753 ++++++++++------------ 2 files changed, 2586 insertions(+), 3169 deletions(-) diff --git a/airflow-core/docs/img/airflow_erd.sha256 b/airflow-core/docs/img/airflow_erd.sha256 index 605175c74b83f..a296665a3fd06 100644 --- a/airflow-core/docs/img/airflow_erd.sha256 +++ b/airflow-core/docs/img/airflow_erd.sha256 @@ -1 +1 @@ -ecba28581229c5a0d1107d162a9caac0254e0bf60845f19757eb3349b43b78e5 \ No newline at end of file +f3c5aa57ed9f3bf4fb4599bc3b84e20830d54c08f59bd9096436ecc2793e8b37 diff --git a/airflow-core/docs/img/airflow_erd.svg b/airflow-core/docs/img/airflow_erd.svg index 7bb08eac9306c..15c4775cacabe 100644 --- a/airflow-core/docs/img/airflow_erd.svg +++ b/airflow-core/docs/img/airflow_erd.svg @@ -4,701 +4,502 @@ - - + + %3 - - + + -dag_bundle_team - -dag_bundle_team - -dag_bundle_name - - [VARCHAR(250)] - NOT NULL - -team_name - - [VARCHAR(50)] - NOT NULL +connection + +connection + +id + + [INTEGER] + NOT NULL + +conn_id + + [VARCHAR(250)] + NOT NULL + +conn_type + + [VARCHAR(500)] + NOT NULL + +description + + [TEXT] + +extra + + [TEXT] + +host + + [VARCHAR(500)] + +is_encrypted + + [BOOLEAN] + NOT NULL + +is_extra_encrypted + + [BOOLEAN] + NOT NULL + +login + + [TEXT] + +password + + [TEXT] + +port + + [INTEGER] + +schema + + [VARCHAR(500)] + +team_name + + [VARCHAR(50)] - + -dag_bundle - -dag_bundle - -name - - [VARCHAR(250)] - NOT NULL - -active - - [BOOLEAN] - -last_refreshed - - [TIMESTAMP] - -signed_url_template - - [TEXT] - -template_params - - [JSON] - -version - - [VARCHAR(200)] - - - -dag_bundle:name--dag_bundle_team:dag_bundle_name - -0..N -1 - - - -dag - -dag - -dag_id - - [VARCHAR(250)] - NOT NULL - -allowed_run_types - - [JSON] - -asset_expression - - [JSON] - -bundle_name - - [VARCHAR(250)] - NOT NULL - -bundle_version - - [VARCHAR(200)] - -dag_display_name - - [VARCHAR(2000)] - -deadline - - [JSON] - -description - - [TEXT] - -exceeds_max_non_backfill - - [BOOLEAN] - NOT NULL - -fail_fast - - [BOOLEAN] - NOT NULL - -fileloc - - [VARCHAR(2000)] - -has_import_errors - - [BOOLEAN] - NOT NULL - -has_task_concurrency_limits - - [BOOLEAN] - NOT NULL - -is_paused - - [BOOLEAN] - NOT NULL - -is_stale - - [BOOLEAN] - NOT NULL - -last_expired - - [TIMESTAMP] - -last_parse_duration - - [DOUBLE PRECISION] - -last_parsed_time - - [TIMESTAMP] - -max_active_runs - - [INTEGER] - -max_active_tasks - - [INTEGER] - NOT NULL - -max_consecutive_failed_dag_runs - - [INTEGER] - NOT NULL - -next_dagrun - - [TIMESTAMP] - -next_dagrun_create_after - - [TIMESTAMP] - -next_dagrun_data_interval_end - - [TIMESTAMP] - -next_dagrun_data_interval_start - - [TIMESTAMP] - -next_dagrun_partition_date - - [TIMESTAMP] - -next_dagrun_partition_key - - [VARCHAR(255)] - -owners - - [VARCHAR(2000)] - -relative_fileloc - - [VARCHAR(2000)] - -timetable_description - - [VARCHAR(1000)] - -timetable_partitioned - - [BOOLEAN] - NOT NULL - -timetable_summary - - [TEXT] - -timetable_type - - [VARCHAR(255)] - NOT NULL - - - -dag_bundle:name--dag:bundle_name - -0..N -1 +dag_bundle_team + +dag_bundle_team + +dag_bundle_name + + [VARCHAR(250)] + NOT NULL + +team_name + + [VARCHAR(50)] + NOT NULL team - -team - -name - - [VARCHAR(50)] - NOT NULL + +team + +name + + [VARCHAR(50)] + NOT NULL + + + +team:name--connection:team_name + +0..N +{0,1} team:name--dag_bundle_team:team_name - -0..N -1 - - - -connection - -connection - -id - - [INTEGER] - NOT NULL - -conn_id - - [VARCHAR(250)] - NOT NULL - -conn_type - - [VARCHAR(500)] - NOT NULL - -description - - [TEXT] - -extra - - [TEXT] - -host - - [VARCHAR(500)] - -is_encrypted - - [BOOLEAN] - NOT NULL - -is_extra_encrypted - - [BOOLEAN] - NOT NULL - -login - - [TEXT] - -password - - [TEXT] - -port - - [INTEGER] - -schema - - [VARCHAR(500)] - -team_name - - [VARCHAR(50)] - - - -team:name--connection:team_name - -0..N -{0,1} + +0..N +1 - + slot_pool - -slot_pool - -id - - [INTEGER] - NOT NULL - -description - - [TEXT] - -include_deferred - - [BOOLEAN] - NOT NULL - -pool - - [VARCHAR(256)] - NOT NULL - -slots - - [INTEGER] - NOT NULL - -team_name - - [VARCHAR(50)] + +slot_pool + +id + + [INTEGER] + NOT NULL + +description + + [TEXT] + +include_deferred + + [BOOLEAN] + NOT NULL + +pool + + [VARCHAR(256)] + NOT NULL + +slots + + [INTEGER] + NOT NULL + +team_name + + [VARCHAR(50)] - + team:name--slot_pool:team_name - -0..N -{0,1} + +0..N +{0,1} - + variable - -variable - -id - - [INTEGER] - NOT NULL - -description - - [TEXT] - -is_encrypted - - [BOOLEAN] - NOT NULL - -key - - [VARCHAR(250)] - NOT NULL - -team_name - - [VARCHAR(50)] - -val - - [TEXT] - NOT NULL + +variable + +id + + [INTEGER] + NOT NULL + +description + + [TEXT] + +is_encrypted + + [BOOLEAN] + NOT NULL + +key + + [VARCHAR(250)] + NOT NULL + +team_name + + [VARCHAR(50)] + +val + + [TEXT] + NOT NULL - + team:name--variable:team_name - -0..N -{0,1} + +0..N +{0,1} - + -trigger - -trigger - -id - - [INTEGER] - NOT NULL - -classpath - - [VARCHAR(1000)] - NOT NULL - -created_date - - [TIMESTAMP] - NOT NULL - -kwargs - - [TEXT] - NOT NULL - -queue - - [VARCHAR(256)] - -triggerer_id - - [INTEGER] - - - -callback - -callback - -id - - [UUID] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -data - - [JSONB] - NOT NULL - -fetch_method - - [VARCHAR(20)] - NOT NULL - -output - - [TEXT] - -priority_weight - - [INTEGER] - NOT NULL - -state - - [VARCHAR(10)] - -trigger_id - - [INTEGER] - -type - - [VARCHAR(20)] - NOT NULL +dag_bundle + +dag_bundle + +name + + [VARCHAR(250)] + NOT NULL + +active + + [BOOLEAN] + +last_refreshed + + [TIMESTAMP] + +signed_url_template + + [TEXT] + +template_params + + [JSON] + +version + + [VARCHAR(200)] - + -trigger:id--callback:trigger_id - -0..N -{0,1} +dag_bundle:name--dag_bundle_team:dag_bundle_name + +0..N +1 - - -asset_watcher - -asset_watcher - -asset_id - - [INTEGER] - NOT NULL - -trigger_id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL + + +dag + +dag + +dag_id + + [VARCHAR(250)] + NOT NULL + +allowed_run_types + + [JSON] + +asset_expression + + [JSON] + +bundle_name + + [VARCHAR(250)] + NOT NULL + +bundle_version + + [VARCHAR(200)] + +dag_display_name + + [VARCHAR(2000)] + +deadline + + [JSON] + +description + + [TEXT] + +exceeds_max_non_backfill + + [BOOLEAN] + NOT NULL + +fail_fast + + [BOOLEAN] + NOT NULL + +fileloc + + [VARCHAR(2000)] + +has_import_errors + + [BOOLEAN] + NOT NULL + +has_task_concurrency_limits + + [BOOLEAN] + NOT NULL + +is_paused + + [BOOLEAN] + NOT NULL + +is_stale + + [BOOLEAN] + NOT NULL + +last_expired + + [TIMESTAMP] + +last_parse_duration + + [FLOAT] + +last_parsed_time + + [TIMESTAMP] + +max_active_runs + + [INTEGER] + +max_active_tasks + + [INTEGER] + NOT NULL + +max_consecutive_failed_dag_runs + + [INTEGER] + NOT NULL + +next_dagrun + + [TIMESTAMP] + +next_dagrun_create_after + + [TIMESTAMP] + +next_dagrun_data_interval_end + + [TIMESTAMP] + +next_dagrun_data_interval_start + + [TIMESTAMP] + +next_dagrun_partition_date + + [TIMESTAMP] + +next_dagrun_partition_key + + [VARCHAR(255)] + +owners + + [VARCHAR(2000)] + +relative_fileloc + + [VARCHAR(2000)] + +timetable_description + + [VARCHAR(1000)] + +timetable_partitioned + + [BOOLEAN] + NOT NULL + +timetable_summary + + [TEXT] + +timetable_type + + [VARCHAR(255)] + NOT NULL - - -trigger:id--asset_watcher:trigger_id - -0..N -1 + + +dag_bundle:name--dag:bundle_name + +0..N +1 - - -task_instance - -task_instance - -id - - [UUID] - NOT NULL - -context_carrier - - [JSONB] - -custom_operator_name - - [VARCHAR(1000)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - -duration - - [DOUBLE PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - NOT NULL - -external_executor_id - - [TEXT] - -hostname - - [VARCHAR(1000)] - NOT NULL - -last_heartbeat_at - - [TIMESTAMP] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - NOT NULL - -next_kwargs - - [JSONB] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - NOT NULL - -queue - - [VARCHAR(256)] - NOT NULL - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - NOT NULL - -updated_at - - [TIMESTAMP] + + +job + +job + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + +end_date + + [TIMESTAMP] + +executor_class + + [VARCHAR(500)] + +hostname + + [VARCHAR(500)] + +job_type + + [VARCHAR(30)] + +latest_heartbeat + + [TIMESTAMP] + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +unixname + + [VARCHAR(1000)] - - -trigger:id--task_instance:trigger_id - -0..N -{0,1} + + +callback + +callback + +id + + [CHAR(32)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +data + + [JSON] + NOT NULL + +fetch_method + + [VARCHAR(20)] + NOT NULL + +output + + [TEXT] + +priority_weight + + [INTEGER] + NOT NULL + +state + + [VARCHAR(10)] + +trigger_id + + [INTEGER] + +type + + [VARCHAR(20)] + NOT NULL - + deadline deadline id - [UUID] - NOT NULL + [CHAR(32)] + NOT NULL callback_id - [UUID] - NOT NULL + [CHAR(32)] + NOT NULL created_at @@ -711,7 +512,7 @@ deadline_alert_id - [UUID] + [CHAR(32)] deadline_time @@ -729,2683 +530,2299 @@ NOT NULL - + callback:id--deadline:callback_id - + 0..N -1 - - - -asset_alias - -asset_alias - -id - - [INTEGER] - NOT NULL - -group - - [VARCHAR(1500)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL +1 - + asset_alias_asset - -asset_alias_asset - -alias_id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - - - -asset_alias:id--asset_alias_asset:alias_id - -0..N -1 + +asset_alias_asset + +alias_id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL asset_alias_asset_event - -asset_alias_asset_event - -alias_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL + +asset_alias_asset_event + +alias_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL + + + +asset_watcher + +asset_watcher + +asset_id + + [INTEGER] + NOT NULL + +trigger_id + + [INTEGER] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + + + +asset_alias + +asset_alias + +id + + [INTEGER] + NOT NULL + +group + + [VARCHAR(1500)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + + + +asset_alias:id--asset_alias_asset:alias_id + +0..N +1 - + asset_alias:id--asset_alias_asset_event:alias_id - -0..N -1 + +0..N +1 - + dag_schedule_asset_alias_reference - -dag_schedule_asset_alias_reference - -alias_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_alias_reference + +alias_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + asset_alias:id--dag_schedule_asset_alias_reference:alias_id - -0..N -1 + +0..N +1 - + asset - -asset - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -extra - - [JSON] - NOT NULL - -group - - [VARCHAR(1500)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL + +asset + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +extra + + [JSON] + NOT NULL + +group + + [VARCHAR(1500)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL - + asset:id--asset_alias_asset:asset_id - -0..N -1 + +0..N +1 - + asset:id--asset_watcher:asset_id - -0..N -1 + +0..N +1 - + asset_active - -asset_active - -name - - [VARCHAR(1500)] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL + +asset_active + +name + + [VARCHAR(1500)] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL - -asset:uri--asset_active:uri - -1 -1 + +asset:name--asset_active:name + +1 +1 - -asset:name--asset_active:name - -1 -1 + +asset:uri--asset_active:uri + +1 +1 - + dag_schedule_asset_reference - -dag_schedule_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +dag_schedule_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL asset:id--dag_schedule_asset_reference:asset_id - -0..N -1 + +0..N +1 - + task_outlet_asset_reference - -task_outlet_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +task_outlet_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL asset:id--task_outlet_asset_reference:asset_id - -0..N -1 + +0..N +1 - + task_inlet_asset_reference - -task_inlet_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +task_inlet_asset_reference + +asset_id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + asset:id--task_inlet_asset_reference:asset_id - -0..N -1 + +0..N +1 - + asset_dag_run_queue - -asset_dag_run_queue - -asset_id - - [INTEGER] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + +asset_dag_run_queue + +asset_id + + [INTEGER] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL asset:id--asset_dag_run_queue:asset_id - -0..N -1 + +0..N +1 + + + +dag_schedule_asset_name_reference + +dag_schedule_asset_name_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +name + + [VARCHAR(1500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + + + +dag_schedule_asset_uri_reference + +dag_schedule_asset_uri_reference + +dag_id + + [VARCHAR(250)] + NOT NULL + +uri + + [VARCHAR(1500)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + + + +dagrun_asset_event + +dagrun_asset_event + +dag_run_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL - + asset_event - -asset_event - -id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - -extra - - [JSON] - NOT NULL - -partition_key - - [VARCHAR(250)] - -source_dag_id - - [VARCHAR(250)] - -source_map_index - - [INTEGER] - -source_run_id - - [VARCHAR(250)] - -source_task_id - - [VARCHAR(250)] - -timestamp - - [TIMESTAMP] - NOT NULL + +asset_event + +id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL + +extra + + [JSON] + NOT NULL + +partition_key + + [VARCHAR(250)] + +source_dag_id + + [VARCHAR(250)] + +source_map_index + + [INTEGER] + +source_run_id + + [VARCHAR(250)] + +source_task_id + + [VARCHAR(250)] + +timestamp + + [TIMESTAMP] + NOT NULL - + asset_event:id--asset_alias_asset_event:event_id - -0..N -1 - - - -dagrun_asset_event - -dagrun_asset_event - -dag_run_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL + +0..N +1 - + asset_event:id--dagrun_asset_event:event_id - -0..N -1 + +0..N +1 - - -job - -job - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - -end_date - - [TIMESTAMP] - -executor_class - - [VARCHAR(500)] - -hostname - - [VARCHAR(500)] - -job_type - - [VARCHAR(30)] - -latest_heartbeat - - [TIMESTAMP] - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -unixname - - [VARCHAR(1000)] + + +asset_partition_dag_run + +asset_partition_dag_run + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +created_dag_run_id + + [INTEGER] + +partition_key + + [VARCHAR(250)] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - + partitioned_asset_key_log - -partitioned_asset_key_log - -id - - [INTEGER] - NOT NULL - -asset_event_id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - -asset_partition_dag_run_id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -source_partition_key - - [VARCHAR(250)] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -target_partition_key - - [VARCHAR(250)] - NOT NULL + +partitioned_asset_key_log + +id + + [INTEGER] + NOT NULL + +asset_event_id + + [INTEGER] + NOT NULL + +asset_id + + [INTEGER] + NOT NULL + +asset_partition_dag_run_id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +source_partition_key + + [VARCHAR(250)] + NOT NULL + +target_dag_id + + [VARCHAR(250)] + NOT NULL + +target_partition_key + + [VARCHAR(250)] + NOT NULL - - -log - -log - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - -dttm - - [TIMESTAMP] - NOT NULL - -event - - [VARCHAR(60)] - NOT NULL - -extra - - [TEXT] - -logical_date - - [TIMESTAMP] - -map_index - - [INTEGER] - -owner - - [VARCHAR(500)] - -owner_display_name - - [VARCHAR(500)] - -run_id - - [VARCHAR(250)] - -task_id - - [VARCHAR(250)] - -try_number - - [INTEGER] + + +dag_version + +dag_version + +id + + [CHAR(32)] + NOT NULL + +bundle_name + + [VARCHAR(250)] + +bundle_version + + [VARCHAR(250)] + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +last_updated + + [TIMESTAMP] + NOT NULL + +version_number + + [INTEGER] + NOT NULL - - -connection_test - -connection_test - -id - - [UUID] - NOT NULL - -connection_id - - [VARCHAR(250)] - NOT NULL - -connection_snapshot - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -executor - - [VARCHAR(256)] - -queue - - [VARCHAR(256)] - -result_message - - [TEXT] - -reverted - - [BOOLEAN] - NOT NULL - -state - - [VARCHAR(10)] - NOT NULL - -token - - [VARCHAR(64)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + + +task_instance + +task_instance + +id + + [CHAR(32)] + NOT NULL + +context_carrier + + [JSON] + +custom_operator_name + + [VARCHAR(1000)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [CHAR(32)] + +duration + + [FLOAT] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BLOB] + NOT NULL + +external_executor_id + + [TEXT] + +hostname + + [VARCHAR(1000)] + NOT NULL + +last_heartbeat_at + + [TIMESTAMP] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + NOT NULL + +next_kwargs + + [JSON] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + NOT NULL + +queue + + [VARCHAR(256)] + NOT NULL + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + NOT NULL + +updated_at + + [TIMESTAMP] - - -dag_priority_parsing_request - -dag_priority_parsing_request - -id - - [VARCHAR(32)] - NOT NULL - -bundle_name - - [VARCHAR(250)] - NOT NULL - -relative_fileloc - - [VARCHAR(2000)] - NOT NULL + + +dag_version:id--task_instance:dag_version_id + +0..N +{0,1} - - -import_error - -import_error - -id - - [INTEGER] - NOT NULL - -bundle_name - - [VARCHAR(250)] - -filename - - [VARCHAR(1024)] - -stacktrace - - [TEXT] - -timestamp - - [TIMESTAMP] + + +dag_run + +dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + +bundle_version + + [VARCHAR(250)] + +clear_number + + [INTEGER] + NOT NULL + +conf + + [JSON] + +context_carrier + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +created_dag_version_id + + [CHAR(32)] + +creating_job_id + + [INTEGER] + +dag_id + + [VARCHAR(250)] + NOT NULL + +data_interval_end + + [TIMESTAMP] + +data_interval_start + + [TIMESTAMP] + +end_date + + [TIMESTAMP] + +last_scheduling_decision + + [TIMESTAMP] + +log_template_id + + [INTEGER] + NOT NULL + +logical_date + + [TIMESTAMP] + +partition_date + + [TIMESTAMP] + +partition_key + + [VARCHAR(250)] + +queued_at + + [TIMESTAMP] + +run_after + + [TIMESTAMP] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +run_type + + [VARCHAR(50)] + NOT NULL + +scheduled_by_job_id + + [INTEGER] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(50)] + NOT NULL + +triggered_by + + [VARCHAR(50)] + +triggering_user_name + + [VARCHAR(512)] + +updated_at + + [TIMESTAMP] + NOT NULL - - -revoked_token - -revoked_token - -jti - - [VARCHAR(32)] - NOT NULL - -exp - - [TIMESTAMP] - NOT NULL + + +dag_version:id--dag_run:created_dag_version_id + +0..N +{0,1} - - -dag_schedule_asset_name_reference - -dag_schedule_asset_name_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + + +dag_code + +dag_code + +id + + [CHAR(32)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [CHAR(32)] + NOT NULL + +fileloc + + [VARCHAR(2000)] + NOT NULL + +last_updated + + [TIMESTAMP] + NOT NULL + +source_code + + [TEXT] + NOT NULL + +source_code_hash + + [VARCHAR(32)] + NOT NULL - - -dag:dag_id--dag_schedule_asset_name_reference:dag_id - -0..N -1 + + +dag_version:id--dag_code:dag_version_id + +1 +1 - - -dag_schedule_asset_uri_reference - -dag_schedule_asset_uri_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL + + +serialized_dag + +serialized_dag + +id + + [CHAR(32)] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +dag_hash + + [VARCHAR(32)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [CHAR(32)] + NOT NULL + +data + + [JSON] + +data_compressed + + [BLOB] + +last_updated + + [TIMESTAMP] + NOT NULL - - -dag:dag_id--dag_schedule_asset_uri_reference:dag_id - -0..N -1 + + +dag_version:id--serialized_dag:dag_version_id + +1 +1 - - -dag:dag_id--dag_schedule_asset_alias_reference:dag_id - -0..N -1 + + +deadline_alert + +deadline_alert + +id + + [CHAR(32)] + NOT NULL + +callback_def + + [JSON] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +description + + [TEXT] + +interval + + [FLOAT] + NOT NULL + +name + + [VARCHAR(250)] + +reference + + [JSON] + NOT NULL + +serialized_dag_id + + [CHAR(32)] + NOT NULL - - -dag:dag_id--dag_schedule_asset_reference:dag_id - -0..N -1 + + +deadline_alert:id--deadline:deadline_alert_id + +0..N +{0,1} - - -dag:dag_id--task_outlet_asset_reference:dag_id - -0..N -1 + + +hitl_detail + +hitl_detail + +ti_id + + [CHAR(32)] + NOT NULL + +assignees + + [JSON] + +body + + [TEXT] + +chosen_options + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +defaults + + [JSON] + +multiple + + [BOOLEAN] + +options + + [JSON] + NOT NULL + +params + + [JSON] + NOT NULL + +params_input + + [JSON] + NOT NULL + +responded_at + + [TIMESTAMP] + +responded_by + + [JSON] + +subject + + [TEXT] + NOT NULL - - -dag:dag_id--task_inlet_asset_reference:dag_id - -0..N -1 + + +log + +log + +id + + [INTEGER] + NOT NULL + +dag_id + + [VARCHAR(250)] + +dttm + + [TIMESTAMP] + NOT NULL + +event + + [VARCHAR(60)] + NOT NULL + +extra + + [TEXT] + +logical_date + + [TIMESTAMP] + +map_index + + [INTEGER] + +owner + + [VARCHAR(500)] + +owner_display_name + + [VARCHAR(500)] + +run_id + + [VARCHAR(250)] + +task_id + + [VARCHAR(250)] + +try_number + + [INTEGER] - - -dag:dag_id--asset_dag_run_queue:target_dag_id - -0..N -1 + + +task_map + +task_map + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +keys + + [JSON] + +length + + [INTEGER] + NOT NULL - - -dag_version - -dag_version - -id - - [UUID] - NOT NULL - -bundle_name - - [VARCHAR(250)] - -bundle_version - - [VARCHAR(250)] - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -last_updated - - [TIMESTAMP] - NOT NULL - -version_number - - [INTEGER] - NOT NULL + + +task_reschedule + +task_reschedule + +id + + [INTEGER] + NOT NULL + +duration + + [INTEGER] + NOT NULL + +end_date + + [TIMESTAMP] + NOT NULL + +reschedule_date + + [TIMESTAMP] + NOT NULL + +start_date + + [TIMESTAMP] + NOT NULL + +ti_id + + [CHAR(32)] + NOT NULL - - -dag:dag_id--dag_version:dag_id - -0..N -1 + + +xcom + +xcom + +dag_run_id + + [INTEGER] + NOT NULL + +key + + [VARCHAR(512)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL + +value + + [JSON] - - -dag_tag - -dag_tag - -dag_id - - [VARCHAR(250)] - NOT NULL - -name - - [VARCHAR(100)] - NOT NULL + + +task_instance:id--hitl_detail:ti_id + +1 +1 - - -dag:dag_id--dag_tag:dag_id - -0..N -1 + + +task_instance:run_id--task_map:run_id + +0..N +1 - - -dag_owner_attributes - -dag_owner_attributes - -dag_id - - [VARCHAR(250)] - NOT NULL - -owner - - [VARCHAR(500)] - NOT NULL - -link - - [VARCHAR(500)] - NOT NULL + + +task_instance:task_id--task_map:task_id + +0..N +1 - - -dag:dag_id--dag_owner_attributes:dag_id - -0..N -1 + + +task_instance:map_index--task_map:map_index + +0..N +1 - - -dag_warning - -dag_warning - -dag_id - - [VARCHAR(250)] - NOT NULL - -warning_type - - [VARCHAR(50)] - NOT NULL - -message - - [TEXT] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL + + +task_instance:dag_id--task_map:dag_id + +0..N +1 - - -dag:dag_id--dag_warning:dag_id - -0..N -1 + + +task_instance:id--task_reschedule:ti_id + +0..N +1 - - -dag_favorite - -dag_favorite - -dag_id - - [VARCHAR(250)] - NOT NULL - -user_id - - [VARCHAR(250)] - NOT NULL + + +task_instance:dag_id--xcom:dag_id + +0..N +1 - - -dag:dag_id--dag_favorite:dag_id - -0..N -1 + + +task_instance:map_index--xcom:map_index + +0..N +1 - + + +task_instance:run_id--xcom:run_id + +0..N +1 + + + +task_instance:task_id--xcom:task_id + +0..N +1 + + + +task_instance_note + +task_instance_note + +ti_id + + [CHAR(32)] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] + + + +task_instance:id--task_instance_note:ti_id + +1 +1 + + -dag_run - -dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - -bundle_version - - [VARCHAR(250)] - -clear_number - - [INTEGER] - NOT NULL - -conf - - [JSONB] - -context_carrier - - [JSONB] - -created_at - - [TIMESTAMP] - NOT NULL - -created_dag_version_id - - [UUID] - -creating_job_id - - [INTEGER] - -dag_id - - [VARCHAR(250)] - NOT NULL - -data_interval_end - - [TIMESTAMP] - -data_interval_start - - [TIMESTAMP] - -end_date - - [TIMESTAMP] - -last_scheduling_decision - - [TIMESTAMP] - -log_template_id - - [INTEGER] - NOT NULL - -logical_date - - [TIMESTAMP] - -partition_date - - [TIMESTAMP] - -partition_key - - [VARCHAR(250)] - -queued_at - - [TIMESTAMP] - -run_after - - [TIMESTAMP] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -run_type - - [VARCHAR(50)] - NOT NULL - -scheduled_by_job_id - - [INTEGER] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(50)] - NOT NULL - -triggered_by - - [VARCHAR(50)] - -triggering_user_name - - [VARCHAR(512)] - -updated_at - - [TIMESTAMP] - NOT NULL +task_instance_history + +task_instance_history + +task_instance_id + + [CHAR(32)] + NOT NULL + +context_carrier + + [JSON] + +custom_operator_name + + [VARCHAR(1000)] + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [CHAR(32)] + +duration + + [FLOAT] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BLOB] + +external_executor_id + + [TEXT] + +hostname + + [VARCHAR(1000)] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + +next_kwargs + + [JSON] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + +queue + + [VARCHAR(256)] + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [DATETIME] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + +updated_at + + [TIMESTAMP] - - -dag_version:id--dag_run:created_dag_version_id - -0..N -{0,1} + + +task_instance:dag_id--task_instance_history:dag_id + +0..N +1 - - -dag_code - -dag_code - -id - - [UUID] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - NOT NULL - -fileloc - - [VARCHAR(2000)] - NOT NULL - -last_updated - - [TIMESTAMP] - NOT NULL - -source_code - - [TEXT] - NOT NULL - -source_code_hash - - [VARCHAR(32)] - NOT NULL + + +task_instance:map_index--task_instance_history:map_index + +0..N +1 - - -dag_version:id--dag_code:dag_version_id - -0..N -1 + + +task_instance:run_id--task_instance_history:run_id + +0..N +1 - - -serialized_dag - -serialized_dag - -id - - [UUID] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -dag_hash - - [VARCHAR(32)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - NOT NULL - -data - - [JSONB] - -data_compressed - - [BYTEA] - -last_updated - - [TIMESTAMP] - NOT NULL + + +task_instance:task_id--task_instance_history:task_id + +0..N +1 - - -dag_version:id--serialized_dag:dag_version_id - -0..N -1 + + +rendered_task_instance_fields + +rendered_task_instance_fields + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +k8s_pod_yaml + + [JSON] + +rendered_fields + + [JSON] + NOT NULL - - -dag_version:id--task_instance:dag_version_id - -0..N -{0,1} + + +task_instance:run_id--rendered_task_instance_fields:run_id + +0..N +1 - + + +task_instance:map_index--rendered_task_instance_fields:map_index + +0..N +1 + + + +task_instance:task_id--rendered_task_instance_fields:task_id + +0..N +1 + + + +task_instance:dag_id--rendered_task_instance_fields:dag_id + +0..N +1 + + + +backfill + +backfill + +id + + [INTEGER] + NOT NULL + +completed_at + + [TIMESTAMP] + +created_at + + [TIMESTAMP] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_run_conf + + [JSON] + NOT NULL + +from_date + + [TIMESTAMP] + NOT NULL + +is_paused + + [BOOLEAN] + +max_active_runs + + [INTEGER] + NOT NULL + +reprocess_behavior + + [VARCHAR(250)] + NOT NULL + +to_date + + [TIMESTAMP] + NOT NULL + +triggering_user_name + + [VARCHAR(512)] + +updated_at + + [TIMESTAMP] + NOT NULL + + + +backfill_dag_run + +backfill_dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + NOT NULL + +dag_run_id + + [INTEGER] + +exception_reason + + [VARCHAR(250)] + +logical_date + + [TIMESTAMP] + +partition_key + + [VARCHAR(250)] + +sort_ordinal + + [INTEGER] + NOT NULL + + + +backfill:id--backfill_dag_run:backfill_id + +0..N +1 + + + +backfill:id--dag_run:backfill_id + +0..N +{0,1} + + +hitl_detail_history + +hitl_detail_history + +ti_history_id + + [CHAR(32)] + NOT NULL + +assignees + + [JSON] + +body + + [TEXT] + +chosen_options + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +defaults + + [JSON] + +multiple + + [BOOLEAN] + +options + + [JSON] + NOT NULL + +params + + [JSON] + NOT NULL + +params_input + + [JSON] + NOT NULL + +responded_at + + [TIMESTAMP] + +responded_by + + [JSON] + +subject + + [TEXT] + NOT NULL + + + +task_instance_history:task_instance_id--hitl_detail_history:ti_history_id + +1 +1 + + + log_template - -log_template - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -elasticsearch_id - - [TEXT] - NOT NULL - -filename - - [TEXT] - NOT NULL + +log_template + +id + + [INTEGER] + NOT NULL + +created_at + + [TIMESTAMP] + NOT NULL + +elasticsearch_id + + [TEXT] + NOT NULL + +filename + + [TEXT] + NOT NULL - + log_template:id--dag_run:log_template_id - -0..N -1 + +0..N +1 + + + +dag_run:id--deadline:dagrun_id + +0..N +{0,1} - + dag_run:id--dagrun_asset_event:dag_run_id - -0..N -1 - - - -asset_partition_dag_run - -asset_partition_dag_run - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -created_dag_run_id - - [INTEGER] - -partition_key - - [VARCHAR(250)] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL + +0..N +1 - + dag_run:id--asset_partition_dag_run:created_dag_run_id - -0..N -{0,1} + +0..N +{0,1} - + dag_run:run_id--task_instance:run_id - -0..N -1 + +0..N +1 - + dag_run:dag_id--task_instance:dag_id - -0..N -1 - - - -backfill_dag_run - -backfill_dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - NOT NULL - -dag_run_id - - [INTEGER] - -exception_reason - - [VARCHAR(250)] - -logical_date - - [TIMESTAMP] - -partition_key - - [VARCHAR(250)] - -sort_ordinal - - [INTEGER] - NOT NULL + +0..N +1 - + dag_run:id--backfill_dag_run:dag_run_id - -0..N -{0,1} + +0..N +{0,1} - + dag_run_note - -dag_run_note - -dag_run_id - - [INTEGER] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] + +dag_run_note + +dag_run_id + + [INTEGER] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] - + dag_run:id--dag_run_note:dag_run_id - -1 -1 - - - -dag_run:id--deadline:dagrun_id - -0..N -{0,1} + +1 +1 - - -backfill - -backfill - -id - - [INTEGER] - NOT NULL - -completed_at - - [TIMESTAMP] - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_run_conf - - [JSON] - NOT NULL - -from_date - - [TIMESTAMP] - NOT NULL - -is_paused - - [BOOLEAN] - -max_active_runs - - [INTEGER] - NOT NULL - -reprocess_behavior - - [VARCHAR(250)] - NOT NULL - -to_date - - [TIMESTAMP] - NOT NULL - -triggering_user_name - - [VARCHAR(512)] - -updated_at - - [TIMESTAMP] - NOT NULL - - - -backfill:id--dag_run:backfill_id - -0..N -{0,1} - - - -backfill:id--backfill_dag_run:backfill_id - -0..N -1 - - - -deadline_alert - -deadline_alert - -id - - [UUID] - NOT NULL - -callback_def - - [JSON] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -description - - [TEXT] - -interval - - [DOUBLE PRECISION] - NOT NULL - -name - - [VARCHAR(250)] - -reference - - [JSON] - NOT NULL - -serialized_dag_id - - [UUID] - NOT NULL - - - -serialized_dag:id--deadline_alert:serialized_dag_id - -0..N -1 - - - -deadline_alert:id--deadline:deadline_alert_id - -0..N -{0,1} - - - -hitl_detail - -hitl_detail - -ti_id - - [UUID] - NOT NULL - -assignees - - [JSON] - -body - - [TEXT] - -chosen_options - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -defaults - - [JSON] - -multiple - - [BOOLEAN] - -options - - [JSON] - NOT NULL - -params - - [JSON] - NOT NULL - -params_input - - [JSON] - NOT NULL - -responded_at - - [TIMESTAMP] - -responded_by - - [JSON] - -subject - - [TEXT] - NOT NULL - - - -task_instance:id--hitl_detail:ti_id - -1 -1 + + +dag_tag + +dag_tag + +dag_id + + [VARCHAR(250)] + NOT NULL + +name + + [VARCHAR(100)] + NOT NULL - - -task_map - -task_map - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -keys - - [JSONB] - -length - - [INTEGER] - NOT NULL + + +dag_owner_attributes + +dag_owner_attributes + +dag_id + + [VARCHAR(250)] + NOT NULL + +owner + + [VARCHAR(500)] + NOT NULL + +link + + [VARCHAR(500)] + NOT NULL - - -task_instance:dag_id--task_map:dag_id - -0..N -1 + + +dag:dag_id--dag_schedule_asset_name_reference:dag_id + +0..N +1 - - -task_instance:task_id--task_map:task_id - -0..N -1 + + +dag:dag_id--dag_schedule_asset_uri_reference:dag_id + +0..N +1 - - -task_instance:run_id--task_map:run_id - -0..N -1 + + +dag:dag_id--dag_schedule_asset_alias_reference:dag_id + +0..N +1 - - -task_instance:map_index--task_map:map_index - -0..N -1 + + +dag:dag_id--dag_schedule_asset_reference:dag_id + +0..N +1 - - -task_reschedule - -task_reschedule - -id - - [INTEGER] - NOT NULL - -duration - - [INTEGER] - NOT NULL - -end_date - - [TIMESTAMP] - NOT NULL - -reschedule_date - - [TIMESTAMP] - NOT NULL - -start_date - - [TIMESTAMP] - NOT NULL - -ti_id - - [UUID] - NOT NULL + + +dag:dag_id--task_outlet_asset_reference:dag_id + +0..N +1 - - -task_instance:id--task_reschedule:ti_id - -0..N -1 + + +dag:dag_id--task_inlet_asset_reference:dag_id + +0..N +1 - - -xcom - -xcom - -dag_run_id - - [INTEGER] - NOT NULL - -key - - [VARCHAR(512)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL - -value - - [JSONB] + + +dag:dag_id--asset_dag_run_queue:target_dag_id + +0..N +1 - - -task_instance:dag_id--xcom:dag_id - -0..N -1 + + +dag:dag_id--dag_version:dag_id + +0..N +1 - + -task_instance:map_index--xcom:map_index - -0..N -1 +dag:dag_id--dag_tag:dag_id + +0..N +1 - + -task_instance:task_id--xcom:task_id - -0..N -1 - - - -task_instance:run_id--xcom:run_id - -0..N -1 +dag:dag_id--dag_owner_attributes:dag_id + +0..N +1 - - -task_instance_note - -task_instance_note - -ti_id - - [UUID] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] + + +dag_warning + +dag_warning + +dag_id + + [VARCHAR(250)] + NOT NULL + +warning_type + + [VARCHAR(50)] + NOT NULL + +message + + [TEXT] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL - + -task_instance:id--task_instance_note:ti_id - -1 -1 - - - -task_instance_history - -task_instance_history - -task_instance_id - - [UUID] - NOT NULL - -context_carrier - - [JSONB] - -custom_operator_name - - [VARCHAR(1000)] - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - -duration - - [DOUBLE PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - -external_executor_id - - [TEXT] - -hostname - - [VARCHAR(1000)] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - -next_kwargs - - [JSONB] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - -queue - - [VARCHAR(256)] - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - -updated_at - - [TIMESTAMP] +dag:dag_id--dag_warning:dag_id + +0..N +1 - - -task_instance:dag_id--task_instance_history:dag_id - -0..N -1 + + +dag_favorite + +dag_favorite + +dag_id + + [VARCHAR(250)] + NOT NULL + +user_id + + [VARCHAR(250)] + NOT NULL - - -task_instance:run_id--task_instance_history:run_id - -0..N -1 + + +dag:dag_id--dag_favorite:dag_id + +0..N +1 - - -task_instance:map_index--task_instance_history:map_index - -0..N -1 + + +trigger + +trigger + +id + + [INTEGER] + NOT NULL + +classpath + + [VARCHAR(1000)] + NOT NULL + +created_date + + [TIMESTAMP] + NOT NULL + +kwargs + + [TEXT] + NOT NULL + +queue + + [VARCHAR(256)] + +triggerer_id + + [INTEGER] - - -task_instance:task_id--task_instance_history:task_id - -0..N -1 + + +trigger:id--callback:trigger_id + +0..N +{0,1} - - -rendered_task_instance_fields - -rendered_task_instance_fields - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -k8s_pod_yaml - - [JSON] - -rendered_fields - - [JSON] - NOT NULL + + +trigger:id--asset_watcher:trigger_id + +0..N +1 - - -task_instance:task_id--rendered_task_instance_fields:task_id - -0..N -1 + + +trigger:id--task_instance:trigger_id + +0..N +{0,1} - - -task_instance:dag_id--rendered_task_instance_fields:dag_id - -0..N -1 + + +connection_test + +connection_test + +id + + [CHAR(32)] + NOT NULL + +connection_id + + [VARCHAR(250)] + NOT NULL + +connection_snapshot + + [JSON] + +created_at + + [TIMESTAMP] + NOT NULL + +executor + + [VARCHAR(256)] + +queue + + [VARCHAR(256)] + +result_message + + [TEXT] + +reverted + + [BOOLEAN] + NOT NULL + +state + + [VARCHAR(10)] + NOT NULL + +token + + [VARCHAR(64)] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL - - -task_instance:map_index--rendered_task_instance_fields:map_index - -0..N -1 + + +dag_priority_parsing_request + +dag_priority_parsing_request + +id + + [VARCHAR(32)] + NOT NULL + +bundle_name + + [VARCHAR(250)] + NOT NULL + +relative_fileloc + + [VARCHAR(2000)] + NOT NULL - - -task_instance:run_id--rendered_task_instance_fields:run_id - -0..N -1 + + +import_error + +import_error + +id + + [INTEGER] + NOT NULL + +bundle_name + + [VARCHAR(250)] + +filename + + [VARCHAR(1024)] + +stacktrace + + [TEXT] + +timestamp + + [TIMESTAMP] - + -edge_job - -edge_job - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -try_number - - [INTEGER] - NOT NULL - -command - - [VARCHAR(2048)] - NOT NULL - -concurrency_slots - - [INTEGER] - NOT NULL - -edge_worker - - [VARCHAR(64)] - -last_update - - [TIMESTAMP] - -queue - - [VARCHAR(256)] - NOT NULL - -queued_dttm - - [TIMESTAMP] - -state - - [VARCHAR(20)] - NOT NULL - - - -hitl_detail_history - -hitl_detail_history - -ti_history_id - - [UUID] - NOT NULL - -assignees - - [JSON] - -body - - [TEXT] - -chosen_options - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -defaults - - [JSON] - -multiple - - [BOOLEAN] - -options - - [JSON] - NOT NULL - -params - - [JSON] - NOT NULL - -params_input - - [JSON] - NOT NULL - -responded_at - - [TIMESTAMP] - -responded_by - - [JSON] - -subject - - [TEXT] - NOT NULL +revoked_token + +revoked_token + +jti + + [VARCHAR(32)] + NOT NULL + +exp + + [TIMESTAMP] + NOT NULL - - -task_instance_history:task_instance_id--hitl_detail_history:ti_history_id - -1 -1 - - - -alembic_version - -alembic_version - -version_num - - [VARCHAR(32)] - NOT NULL - - - -edge_logs - -edge_logs - -dag_id - - [VARCHAR(250)] - NOT NULL - -log_chunk_time - - [TIMESTAMP] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -try_number - - [INTEGER] - NOT NULL - -log_chunk_data - - [TEXT] - NOT NULL - - - -edge_worker - -edge_worker - -worker_name - - [VARCHAR(64)] - NOT NULL - -concurrency - - [INTEGER] - -first_online - - [TIMESTAMP] - -jobs_active - - [INTEGER] - NOT NULL - -jobs_failed - - [INTEGER] - NOT NULL - -jobs_success - - [INTEGER] - NOT NULL - -jobs_taken - - [INTEGER] - NOT NULL - -last_update - - [TIMESTAMP] - -maintenance_comment - - [VARCHAR(1024)] - -queues - - [VARCHAR(256)] - -state - - [VARCHAR(20)] - NOT NULL - -sysinfo - - [VARCHAR(256)] - - - -alembic_version_edge3 - -alembic_version_edge3 - -version_num - - [VARCHAR(32)] - NOT NULL - - - -alembic_version_fab - -alembic_version_fab - -version_num - - [VARCHAR(32)] - NOT NULL - - - -ab_user - -ab_user - -id - - [INTEGER] - NOT NULL - -active - - [BOOLEAN] - -changed_by_fk - - [INTEGER] - -changed_on - - [TIMESTAMP] - -created_by_fk - - [INTEGER] - -created_on - - [TIMESTAMP] - -email - - [VARCHAR(512)] - NOT NULL - -fail_login_count - - [INTEGER] - -first_name - - [VARCHAR(256)] - NOT NULL - -last_login - - [TIMESTAMP] - -last_name - - [VARCHAR(256)] - NOT NULL - -login_count - - [INTEGER] - -password - - [VARCHAR(256)] - -username - - [VARCHAR(512)] - NOT NULL - - - -ab_user:id--ab_user:changed_by_fk - -0..N -{0,1} - - - -ab_user:id--ab_user:created_by_fk - -0..N -{0,1} - - - -ab_user_role - -ab_user_role - -id - - [INTEGER] - NOT NULL - -role_id - - [INTEGER] - -user_id - - [INTEGER] - - - -ab_user:id--ab_user_role:user_id - -0..N -{0,1} - - - -ab_user_group - -ab_user_group - -id - - [INTEGER] - NOT NULL - -group_id - - [INTEGER] - -user_id - - [INTEGER] - - - -ab_user:id--ab_user_group:user_id - -0..N -{0,1} - - - -ab_register_user - -ab_register_user - -id - - [INTEGER] - NOT NULL - -email - - [VARCHAR(512)] - NOT NULL - -first_name - - [VARCHAR(256)] - NOT NULL - -last_name - - [VARCHAR(256)] - NOT NULL - -password - - [VARCHAR(256)] - -registration_date - - [TIMESTAMP] - -registration_hash - - [VARCHAR(256)] - -username - - [VARCHAR(512)] - NOT NULL - - - -ab_group - -ab_group - -id - - [INTEGER] - NOT NULL - -description - - [VARCHAR(512)] - -label - - [VARCHAR(150)] - -name - - [VARCHAR(100)] - NOT NULL - - - -ab_group_role - -ab_group_role - -id - - [INTEGER] - NOT NULL - -group_id - - [INTEGER] - -role_id - - [INTEGER] - - - -ab_group:id--ab_group_role:group_id - -0..N -{0,1} - - - -ab_group:id--ab_user_group:group_id - -0..N -{0,1} - - - -ab_role - -ab_role - -id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(64)] - NOT NULL - - - -ab_role:id--ab_group_role:role_id - -0..N -{0,1} - - - -ab_role:id--ab_user_role:role_id - -0..N -{0,1} - - - -ab_permission_view_role - -ab_permission_view_role - -id - - [INTEGER] - NOT NULL - -permission_view_id - - [INTEGER] - -role_id - - [INTEGER] - - - -ab_role:id--ab_permission_view_role:role_id - -0..N -{0,1} - - - -ab_permission - -ab_permission - -id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(100)] - NOT NULL - - - -ab_permission_view - -ab_permission_view - -id - - [INTEGER] - NOT NULL - -permission_id - - [INTEGER] - NOT NULL - -view_menu_id - - [INTEGER] - NOT NULL - - - -ab_permission:id--ab_permission_view:permission_id - -0..N -1 - - - -ab_permission_view:id--ab_permission_view_role:permission_view_id - -0..N -{0,1} - - - -ab_view_menu - -ab_view_menu - -id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(250)] - NOT NULL - - - -ab_view_menu:id--ab_permission_view:view_menu_id - -0..N -1 - - - -session - -session - -id - - [INTEGER] - NOT NULL - -data - - [BYTEA] - -expiry - - [TIMESTAMP] - -session_id - - [VARCHAR(255)] + + +serialized_dag:id--deadline_alert:serialized_dag_id + +0..N +1 From 324cc47d82011e18880aec1ed7508b0f1eb3b349 Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 12 Mar 2026 10:02:40 -0500 Subject: [PATCH 31/38] Remove committed ERD files --- airflow-core/docs/img/airflow_erd.sha256 | 1 - airflow-core/docs/img/airflow_erd.svg | 2828 ----------------- .../0109_3_2_0_add_connection_test_table.py | 8 - 3 files changed, 2837 deletions(-) delete mode 100644 airflow-core/docs/img/airflow_erd.sha256 delete mode 100644 airflow-core/docs/img/airflow_erd.svg diff --git a/airflow-core/docs/img/airflow_erd.sha256 b/airflow-core/docs/img/airflow_erd.sha256 deleted file mode 100644 index a296665a3fd06..0000000000000 --- a/airflow-core/docs/img/airflow_erd.sha256 +++ /dev/null @@ -1 +0,0 @@ -f3c5aa57ed9f3bf4fb4599bc3b84e20830d54c08f59bd9096436ecc2793e8b37 diff --git a/airflow-core/docs/img/airflow_erd.svg b/airflow-core/docs/img/airflow_erd.svg deleted file mode 100644 index 15c4775cacabe..0000000000000 --- a/airflow-core/docs/img/airflow_erd.svg +++ /dev/null @@ -1,2828 +0,0 @@ - - - - - - -%3 - - - -connection - -connection - -id - - [INTEGER] - NOT NULL - -conn_id - - [VARCHAR(250)] - NOT NULL - -conn_type - - [VARCHAR(500)] - NOT NULL - -description - - [TEXT] - -extra - - [TEXT] - -host - - [VARCHAR(500)] - -is_encrypted - - [BOOLEAN] - NOT NULL - -is_extra_encrypted - - [BOOLEAN] - NOT NULL - -login - - [TEXT] - -password - - [TEXT] - -port - - [INTEGER] - -schema - - [VARCHAR(500)] - -team_name - - [VARCHAR(50)] - - - -dag_bundle_team - -dag_bundle_team - -dag_bundle_name - - [VARCHAR(250)] - NOT NULL - -team_name - - [VARCHAR(50)] - NOT NULL - - - -team - -team - -name - - [VARCHAR(50)] - NOT NULL - - - -team:name--connection:team_name - -0..N -{0,1} - - - -team:name--dag_bundle_team:team_name - -0..N -1 - - - -slot_pool - -slot_pool - -id - - [INTEGER] - NOT NULL - -description - - [TEXT] - -include_deferred - - [BOOLEAN] - NOT NULL - -pool - - [VARCHAR(256)] - NOT NULL - -slots - - [INTEGER] - NOT NULL - -team_name - - [VARCHAR(50)] - - - -team:name--slot_pool:team_name - -0..N -{0,1} - - - -variable - -variable - -id - - [INTEGER] - NOT NULL - -description - - [TEXT] - -is_encrypted - - [BOOLEAN] - NOT NULL - -key - - [VARCHAR(250)] - NOT NULL - -team_name - - [VARCHAR(50)] - -val - - [TEXT] - NOT NULL - - - -team:name--variable:team_name - -0..N -{0,1} - - - -dag_bundle - -dag_bundle - -name - - [VARCHAR(250)] - NOT NULL - -active - - [BOOLEAN] - -last_refreshed - - [TIMESTAMP] - -signed_url_template - - [TEXT] - -template_params - - [JSON] - -version - - [VARCHAR(200)] - - - -dag_bundle:name--dag_bundle_team:dag_bundle_name - -0..N -1 - - - -dag - -dag - -dag_id - - [VARCHAR(250)] - NOT NULL - -allowed_run_types - - [JSON] - -asset_expression - - [JSON] - -bundle_name - - [VARCHAR(250)] - NOT NULL - -bundle_version - - [VARCHAR(200)] - -dag_display_name - - [VARCHAR(2000)] - -deadline - - [JSON] - -description - - [TEXT] - -exceeds_max_non_backfill - - [BOOLEAN] - NOT NULL - -fail_fast - - [BOOLEAN] - NOT NULL - -fileloc - - [VARCHAR(2000)] - -has_import_errors - - [BOOLEAN] - NOT NULL - -has_task_concurrency_limits - - [BOOLEAN] - NOT NULL - -is_paused - - [BOOLEAN] - NOT NULL - -is_stale - - [BOOLEAN] - NOT NULL - -last_expired - - [TIMESTAMP] - -last_parse_duration - - [FLOAT] - -last_parsed_time - - [TIMESTAMP] - -max_active_runs - - [INTEGER] - -max_active_tasks - - [INTEGER] - NOT NULL - -max_consecutive_failed_dag_runs - - [INTEGER] - NOT NULL - -next_dagrun - - [TIMESTAMP] - -next_dagrun_create_after - - [TIMESTAMP] - -next_dagrun_data_interval_end - - [TIMESTAMP] - -next_dagrun_data_interval_start - - [TIMESTAMP] - -next_dagrun_partition_date - - [TIMESTAMP] - -next_dagrun_partition_key - - [VARCHAR(255)] - -owners - - [VARCHAR(2000)] - -relative_fileloc - - [VARCHAR(2000)] - -timetable_description - - [VARCHAR(1000)] - -timetable_partitioned - - [BOOLEAN] - NOT NULL - -timetable_summary - - [TEXT] - -timetable_type - - [VARCHAR(255)] - NOT NULL - - - -dag_bundle:name--dag:bundle_name - -0..N -1 - - - -job - -job - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - -end_date - - [TIMESTAMP] - -executor_class - - [VARCHAR(500)] - -hostname - - [VARCHAR(500)] - -job_type - - [VARCHAR(30)] - -latest_heartbeat - - [TIMESTAMP] - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -unixname - - [VARCHAR(1000)] - - - -callback - -callback - -id - - [CHAR(32)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -data - - [JSON] - NOT NULL - -fetch_method - - [VARCHAR(20)] - NOT NULL - -output - - [TEXT] - -priority_weight - - [INTEGER] - NOT NULL - -state - - [VARCHAR(10)] - -trigger_id - - [INTEGER] - -type - - [VARCHAR(20)] - NOT NULL - - - -deadline - -deadline - -id - - [CHAR(32)] - NOT NULL - -callback_id - - [CHAR(32)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -dagrun_id - - [INTEGER] - -deadline_alert_id - - [CHAR(32)] - -deadline_time - - [TIMESTAMP] - NOT NULL - -last_updated_at - - [TIMESTAMP] - NOT NULL - -missed - - [BOOLEAN] - NOT NULL - - - -callback:id--deadline:callback_id - -0..N -1 - - - -asset_alias_asset - -asset_alias_asset - -alias_id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - - - -asset_alias_asset_event - -asset_alias_asset_event - -alias_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL - - - -asset_watcher - -asset_watcher - -asset_id - - [INTEGER] - NOT NULL - -trigger_id - - [INTEGER] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - - - -asset_alias - -asset_alias - -id - - [INTEGER] - NOT NULL - -group - - [VARCHAR(1500)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - - - -asset_alias:id--asset_alias_asset:alias_id - -0..N -1 - - - -asset_alias:id--asset_alias_asset_event:alias_id - -0..N -1 - - - -dag_schedule_asset_alias_reference - -dag_schedule_asset_alias_reference - -alias_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - - - -asset_alias:id--dag_schedule_asset_alias_reference:alias_id - -0..N -1 - - - -asset - -asset - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -extra - - [JSON] - NOT NULL - -group - - [VARCHAR(1500)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL - - - -asset:id--asset_alias_asset:asset_id - -0..N -1 - - - -asset:id--asset_watcher:asset_id - -0..N -1 - - - -asset_active - -asset_active - -name - - [VARCHAR(1500)] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL - - - -asset:name--asset_active:name - -1 -1 - - - -asset:uri--asset_active:uri - -1 -1 - - - -dag_schedule_asset_reference - -dag_schedule_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - - - -asset:id--dag_schedule_asset_reference:asset_id - -0..N -1 - - - -task_outlet_asset_reference - -task_outlet_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - - - -asset:id--task_outlet_asset_reference:asset_id - -0..N -1 - - - -task_inlet_asset_reference - -task_inlet_asset_reference - -asset_id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - - - -asset:id--task_inlet_asset_reference:asset_id - -0..N -1 - - - -asset_dag_run_queue - -asset_dag_run_queue - -asset_id - - [INTEGER] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - - - -asset:id--asset_dag_run_queue:asset_id - -0..N -1 - - - -dag_schedule_asset_name_reference - -dag_schedule_asset_name_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -name - - [VARCHAR(1500)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - - - -dag_schedule_asset_uri_reference - -dag_schedule_asset_uri_reference - -dag_id - - [VARCHAR(250)] - NOT NULL - -uri - - [VARCHAR(1500)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - - - -dagrun_asset_event - -dagrun_asset_event - -dag_run_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL - - - -asset_event - -asset_event - -id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - -extra - - [JSON] - NOT NULL - -partition_key - - [VARCHAR(250)] - -source_dag_id - - [VARCHAR(250)] - -source_map_index - - [INTEGER] - -source_run_id - - [VARCHAR(250)] - -source_task_id - - [VARCHAR(250)] - -timestamp - - [TIMESTAMP] - NOT NULL - - - -asset_event:id--asset_alias_asset_event:event_id - -0..N -1 - - - -asset_event:id--dagrun_asset_event:event_id - -0..N -1 - - - -asset_partition_dag_run - -asset_partition_dag_run - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -created_dag_run_id - - [INTEGER] - -partition_key - - [VARCHAR(250)] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - - - -partitioned_asset_key_log - -partitioned_asset_key_log - -id - - [INTEGER] - NOT NULL - -asset_event_id - - [INTEGER] - NOT NULL - -asset_id - - [INTEGER] - NOT NULL - -asset_partition_dag_run_id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -source_partition_key - - [VARCHAR(250)] - NOT NULL - -target_dag_id - - [VARCHAR(250)] - NOT NULL - -target_partition_key - - [VARCHAR(250)] - NOT NULL - - - -dag_version - -dag_version - -id - - [CHAR(32)] - NOT NULL - -bundle_name - - [VARCHAR(250)] - -bundle_version - - [VARCHAR(250)] - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -last_updated - - [TIMESTAMP] - NOT NULL - -version_number - - [INTEGER] - NOT NULL - - - -task_instance - -task_instance - -id - - [CHAR(32)] - NOT NULL - -context_carrier - - [JSON] - -custom_operator_name - - [VARCHAR(1000)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [CHAR(32)] - -duration - - [FLOAT] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BLOB] - NOT NULL - -external_executor_id - - [TEXT] - -hostname - - [VARCHAR(1000)] - NOT NULL - -last_heartbeat_at - - [TIMESTAMP] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - NOT NULL - -next_kwargs - - [JSON] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - NOT NULL - -queue - - [VARCHAR(256)] - NOT NULL - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - NOT NULL - -updated_at - - [TIMESTAMP] - - - -dag_version:id--task_instance:dag_version_id - -0..N -{0,1} - - - -dag_run - -dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - -bundle_version - - [VARCHAR(250)] - -clear_number - - [INTEGER] - NOT NULL - -conf - - [JSON] - -context_carrier - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -created_dag_version_id - - [CHAR(32)] - -creating_job_id - - [INTEGER] - -dag_id - - [VARCHAR(250)] - NOT NULL - -data_interval_end - - [TIMESTAMP] - -data_interval_start - - [TIMESTAMP] - -end_date - - [TIMESTAMP] - -last_scheduling_decision - - [TIMESTAMP] - -log_template_id - - [INTEGER] - NOT NULL - -logical_date - - [TIMESTAMP] - -partition_date - - [TIMESTAMP] - -partition_key - - [VARCHAR(250)] - -queued_at - - [TIMESTAMP] - -run_after - - [TIMESTAMP] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -run_type - - [VARCHAR(50)] - NOT NULL - -scheduled_by_job_id - - [INTEGER] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(50)] - NOT NULL - -triggered_by - - [VARCHAR(50)] - -triggering_user_name - - [VARCHAR(512)] - -updated_at - - [TIMESTAMP] - NOT NULL - - - -dag_version:id--dag_run:created_dag_version_id - -0..N -{0,1} - - - -dag_code - -dag_code - -id - - [CHAR(32)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [CHAR(32)] - NOT NULL - -fileloc - - [VARCHAR(2000)] - NOT NULL - -last_updated - - [TIMESTAMP] - NOT NULL - -source_code - - [TEXT] - NOT NULL - -source_code_hash - - [VARCHAR(32)] - NOT NULL - - - -dag_version:id--dag_code:dag_version_id - -1 -1 - - - -serialized_dag - -serialized_dag - -id - - [CHAR(32)] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -dag_hash - - [VARCHAR(32)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [CHAR(32)] - NOT NULL - -data - - [JSON] - -data_compressed - - [BLOB] - -last_updated - - [TIMESTAMP] - NOT NULL - - - -dag_version:id--serialized_dag:dag_version_id - -1 -1 - - - -deadline_alert - -deadline_alert - -id - - [CHAR(32)] - NOT NULL - -callback_def - - [JSON] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -description - - [TEXT] - -interval - - [FLOAT] - NOT NULL - -name - - [VARCHAR(250)] - -reference - - [JSON] - NOT NULL - -serialized_dag_id - - [CHAR(32)] - NOT NULL - - - -deadline_alert:id--deadline:deadline_alert_id - -0..N -{0,1} - - - -hitl_detail - -hitl_detail - -ti_id - - [CHAR(32)] - NOT NULL - -assignees - - [JSON] - -body - - [TEXT] - -chosen_options - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -defaults - - [JSON] - -multiple - - [BOOLEAN] - -options - - [JSON] - NOT NULL - -params - - [JSON] - NOT NULL - -params_input - - [JSON] - NOT NULL - -responded_at - - [TIMESTAMP] - -responded_by - - [JSON] - -subject - - [TEXT] - NOT NULL - - - -log - -log - -id - - [INTEGER] - NOT NULL - -dag_id - - [VARCHAR(250)] - -dttm - - [TIMESTAMP] - NOT NULL - -event - - [VARCHAR(60)] - NOT NULL - -extra - - [TEXT] - -logical_date - - [TIMESTAMP] - -map_index - - [INTEGER] - -owner - - [VARCHAR(500)] - -owner_display_name - - [VARCHAR(500)] - -run_id - - [VARCHAR(250)] - -task_id - - [VARCHAR(250)] - -try_number - - [INTEGER] - - - -task_map - -task_map - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -keys - - [JSON] - -length - - [INTEGER] - NOT NULL - - - -task_reschedule - -task_reschedule - -id - - [INTEGER] - NOT NULL - -duration - - [INTEGER] - NOT NULL - -end_date - - [TIMESTAMP] - NOT NULL - -reschedule_date - - [TIMESTAMP] - NOT NULL - -start_date - - [TIMESTAMP] - NOT NULL - -ti_id - - [CHAR(32)] - NOT NULL - - - -xcom - -xcom - -dag_run_id - - [INTEGER] - NOT NULL - -key - - [VARCHAR(512)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL - -value - - [JSON] - - - -task_instance:id--hitl_detail:ti_id - -1 -1 - - - -task_instance:run_id--task_map:run_id - -0..N -1 - - - -task_instance:task_id--task_map:task_id - -0..N -1 - - - -task_instance:map_index--task_map:map_index - -0..N -1 - - - -task_instance:dag_id--task_map:dag_id - -0..N -1 - - - -task_instance:id--task_reschedule:ti_id - -0..N -1 - - - -task_instance:dag_id--xcom:dag_id - -0..N -1 - - - -task_instance:map_index--xcom:map_index - -0..N -1 - - - -task_instance:run_id--xcom:run_id - -0..N -1 - - - -task_instance:task_id--xcom:task_id - -0..N -1 - - - -task_instance_note - -task_instance_note - -ti_id - - [CHAR(32)] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] - - - -task_instance:id--task_instance_note:ti_id - -1 -1 - - - -task_instance_history - -task_instance_history - -task_instance_id - - [CHAR(32)] - NOT NULL - -context_carrier - - [JSON] - -custom_operator_name - - [VARCHAR(1000)] - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [CHAR(32)] - -duration - - [FLOAT] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BLOB] - -external_executor_id - - [TEXT] - -hostname - - [VARCHAR(1000)] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - -next_kwargs - - [JSON] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - -queue - - [VARCHAR(256)] - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [DATETIME] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - -updated_at - - [TIMESTAMP] - - - -task_instance:dag_id--task_instance_history:dag_id - -0..N -1 - - - -task_instance:map_index--task_instance_history:map_index - -0..N -1 - - - -task_instance:run_id--task_instance_history:run_id - -0..N -1 - - - -task_instance:task_id--task_instance_history:task_id - -0..N -1 - - - -rendered_task_instance_fields - -rendered_task_instance_fields - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -k8s_pod_yaml - - [JSON] - -rendered_fields - - [JSON] - NOT NULL - - - -task_instance:run_id--rendered_task_instance_fields:run_id - -0..N -1 - - - -task_instance:map_index--rendered_task_instance_fields:map_index - -0..N -1 - - - -task_instance:task_id--rendered_task_instance_fields:task_id - -0..N -1 - - - -task_instance:dag_id--rendered_task_instance_fields:dag_id - -0..N -1 - - - -backfill - -backfill - -id - - [INTEGER] - NOT NULL - -completed_at - - [TIMESTAMP] - -created_at - - [TIMESTAMP] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_run_conf - - [JSON] - NOT NULL - -from_date - - [TIMESTAMP] - NOT NULL - -is_paused - - [BOOLEAN] - -max_active_runs - - [INTEGER] - NOT NULL - -reprocess_behavior - - [VARCHAR(250)] - NOT NULL - -to_date - - [TIMESTAMP] - NOT NULL - -triggering_user_name - - [VARCHAR(512)] - -updated_at - - [TIMESTAMP] - NOT NULL - - - -backfill_dag_run - -backfill_dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - NOT NULL - -dag_run_id - - [INTEGER] - -exception_reason - - [VARCHAR(250)] - -logical_date - - [TIMESTAMP] - -partition_key - - [VARCHAR(250)] - -sort_ordinal - - [INTEGER] - NOT NULL - - - -backfill:id--backfill_dag_run:backfill_id - -0..N -1 - - - -backfill:id--dag_run:backfill_id - -0..N -{0,1} - - - -hitl_detail_history - -hitl_detail_history - -ti_history_id - - [CHAR(32)] - NOT NULL - -assignees - - [JSON] - -body - - [TEXT] - -chosen_options - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -defaults - - [JSON] - -multiple - - [BOOLEAN] - -options - - [JSON] - NOT NULL - -params - - [JSON] - NOT NULL - -params_input - - [JSON] - NOT NULL - -responded_at - - [TIMESTAMP] - -responded_by - - [JSON] - -subject - - [TEXT] - NOT NULL - - - -task_instance_history:task_instance_id--hitl_detail_history:ti_history_id - -1 -1 - - - -log_template - -log_template - -id - - [INTEGER] - NOT NULL - -created_at - - [TIMESTAMP] - NOT NULL - -elasticsearch_id - - [TEXT] - NOT NULL - -filename - - [TEXT] - NOT NULL - - - -log_template:id--dag_run:log_template_id - -0..N -1 - - - -dag_run:id--deadline:dagrun_id - -0..N -{0,1} - - - -dag_run:id--dagrun_asset_event:dag_run_id - -0..N -1 - - - -dag_run:id--asset_partition_dag_run:created_dag_run_id - -0..N -{0,1} - - - -dag_run:run_id--task_instance:run_id - -0..N -1 - - - -dag_run:dag_id--task_instance:dag_id - -0..N -1 - - - -dag_run:id--backfill_dag_run:dag_run_id - -0..N -{0,1} - - - -dag_run_note - -dag_run_note - -dag_run_id - - [INTEGER] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] - - - -dag_run:id--dag_run_note:dag_run_id - -1 -1 - - - -dag_tag - -dag_tag - -dag_id - - [VARCHAR(250)] - NOT NULL - -name - - [VARCHAR(100)] - NOT NULL - - - -dag_owner_attributes - -dag_owner_attributes - -dag_id - - [VARCHAR(250)] - NOT NULL - -owner - - [VARCHAR(500)] - NOT NULL - -link - - [VARCHAR(500)] - NOT NULL - - - -dag:dag_id--dag_schedule_asset_name_reference:dag_id - -0..N -1 - - - -dag:dag_id--dag_schedule_asset_uri_reference:dag_id - -0..N -1 - - - -dag:dag_id--dag_schedule_asset_alias_reference:dag_id - -0..N -1 - - - -dag:dag_id--dag_schedule_asset_reference:dag_id - -0..N -1 - - - -dag:dag_id--task_outlet_asset_reference:dag_id - -0..N -1 - - - -dag:dag_id--task_inlet_asset_reference:dag_id - -0..N -1 - - - -dag:dag_id--asset_dag_run_queue:target_dag_id - -0..N -1 - - - -dag:dag_id--dag_version:dag_id - -0..N -1 - - - -dag:dag_id--dag_tag:dag_id - -0..N -1 - - - -dag:dag_id--dag_owner_attributes:dag_id - -0..N -1 - - - -dag_warning - -dag_warning - -dag_id - - [VARCHAR(250)] - NOT NULL - -warning_type - - [VARCHAR(50)] - NOT NULL - -message - - [TEXT] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL - - - -dag:dag_id--dag_warning:dag_id - -0..N -1 - - - -dag_favorite - -dag_favorite - -dag_id - - [VARCHAR(250)] - NOT NULL - -user_id - - [VARCHAR(250)] - NOT NULL - - - -dag:dag_id--dag_favorite:dag_id - -0..N -1 - - - -trigger - -trigger - -id - - [INTEGER] - NOT NULL - -classpath - - [VARCHAR(1000)] - NOT NULL - -created_date - - [TIMESTAMP] - NOT NULL - -kwargs - - [TEXT] - NOT NULL - -queue - - [VARCHAR(256)] - -triggerer_id - - [INTEGER] - - - -trigger:id--callback:trigger_id - -0..N -{0,1} - - - -trigger:id--asset_watcher:trigger_id - -0..N -1 - - - -trigger:id--task_instance:trigger_id - -0..N -{0,1} - - - -connection_test - -connection_test - -id - - [CHAR(32)] - NOT NULL - -connection_id - - [VARCHAR(250)] - NOT NULL - -connection_snapshot - - [JSON] - -created_at - - [TIMESTAMP] - NOT NULL - -executor - - [VARCHAR(256)] - -queue - - [VARCHAR(256)] - -result_message - - [TEXT] - -reverted - - [BOOLEAN] - NOT NULL - -state - - [VARCHAR(10)] - NOT NULL - -token - - [VARCHAR(64)] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - - - -dag_priority_parsing_request - -dag_priority_parsing_request - -id - - [VARCHAR(32)] - NOT NULL - -bundle_name - - [VARCHAR(250)] - NOT NULL - -relative_fileloc - - [VARCHAR(2000)] - NOT NULL - - - -import_error - -import_error - -id - - [INTEGER] - NOT NULL - -bundle_name - - [VARCHAR(250)] - -filename - - [VARCHAR(1024)] - -stacktrace - - [TEXT] - -timestamp - - [TIMESTAMP] - - - -revoked_token - -revoked_token - -jti - - [VARCHAR(32)] - NOT NULL - -exp - - [TIMESTAMP] - NOT NULL - - - -serialized_dag:id--deadline_alert:serialized_dag_id - -0..N -1 - - - diff --git a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py index 9755c67bd8456..ed0c607da9bd4 100644 --- a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py @@ -20,11 +20,7 @@ Add connection_test table for async connection testing. Revision ID: a7e6d4c3b2f1 -<<<<<<<< HEAD:airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py Revises: 888b59e02a5b -======== -Revises: 6222ce48e289 ->>>>>>>> 189776ee4f (clean ups):airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_test_table.py Create Date: 2026-02-22 00:00:00.000000 """ @@ -38,11 +34,7 @@ # revision identifiers, used by Alembic. revision = "a7e6d4c3b2f1" -<<<<<<<< HEAD:airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py down_revision = "888b59e02a5b" -======== -down_revision = "6222ce48e289" ->>>>>>>> 189776ee4f (clean ups):airflow-core/src/airflow/migrations/versions/0108_3_2_0_add_connection_test_table.py branch_labels = None depends_on = None airflow_version = "3.2.0" From 3a1182303f830e514df8d4cdebd1bc96b2ef1b06 Mon Sep 17 00:00:00 2001 From: Anish Date: Thu, 19 Mar 2026 16:33:53 -0500 Subject: [PATCH 32/38] auto gen files --- airflow-core/docs/migrations-ref.rst | 5 +---- ...test_table.py => 0110_3_2_0_add_connection_test_table.py} | 4 ++-- .../src/airflow/ui/openapi-gen/requests/services.gen.ts | 2 +- 3 files changed, 4 insertions(+), 7 deletions(-) rename airflow-core/src/airflow/migrations/versions/{0109_3_2_0_add_connection_test_table.py => 0110_3_2_0_add_connection_test_table.py} (97%) diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index 14ca8d6f5eec3..d1e02fbf22c70 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,10 +39,7 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``b8f3e5d1a9c2`` (head) | ``a7e6d4c3b2f1`` | ``3.2.0`` | Add connection_snapshot and reverted columns to | -| | | | connection_test. | -+-------------------------+------------------+-------------------+--------------------------------------------------------------+ -| ``a7e6d4c3b2f1`` | ``1d6611b6ab7c`` | ``3.2.0`` | Add connection_test table for async connection testing. | +| ``a7e6d4c3b2f1`` (head) | ``1d6611b6ab7c`` | ``3.2.0`` | Add connection_test table for async connection testing. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``1d6611b6ab7c`` | ``888b59e02a5b`` | ``3.2.0`` | Add bundle_name to callback table. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py similarity index 97% rename from airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py rename to airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py index ed0c607da9bd4..33a3625b95025 100644 --- a/airflow-core/src/airflow/migrations/versions/0109_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py @@ -20,7 +20,7 @@ Add connection_test table for async connection testing. Revision ID: a7e6d4c3b2f1 -Revises: 888b59e02a5b +Revises: 1d6611b6ab7c Create Date: 2026-02-22 00:00:00.000000 """ @@ -34,7 +34,7 @@ # revision identifiers, used by Alembic. revision = "a7e6d4c3b2f1" -down_revision = "888b59e02a5b" +down_revision = "1d6611b6ab7c" branch_labels = None depends_on = None airflow_version = "3.2.0" diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 1470ae7a855e7..58e97cd997c40 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -3,7 +3,7 @@ import type { CancelablePromise } from './core/CancelablePromise'; import { OpenAPI } from './core/OpenAPI'; import { request as __request } from './core/request'; -import type { BulkConnectionsData, BulkConnectionsResponse, BulkPoolsData, BulkPoolsResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, BulkVariablesData, BulkVariablesResponse, CancelBackfillData, CancelBackfillResponse, ClearDagRunData, ClearDagRunResponse, CreateAssetEventData, CreateAssetEventResponse, CreateBackfillData, CreateBackfillDryRunData, CreateBackfillDryRunResponse, CreateBackfillResponse, CreateDefaultConnectionsResponse, CreateXcomEntryData, CreateXcomEntryResponse, DagStatsResponse2, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, DeleteConnectionData, DeleteConnectionResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, DeleteDagData, DeleteDagResponse, DeleteDagRunData, DeleteDagRunResponse, DeletePoolData, DeletePoolResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, DeleteVariableData, DeleteVariableResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, FavoriteDagData, FavoriteDagResponse, GenerateTokenData, GenerateTokenResponse2, GetAssetAliasData, GetAssetAliasResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetData, GetAssetEventsData, GetAssetEventsResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, GetAssetResponse, GetAssetsData, GetAssetsResponse, GetAuthMenusResponse, GetBackfillData, GetBackfillResponse, GetCalendarData, GetCalendarResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, GetConnectionData, GetConnectionResponse, GetConnectionTestStatusData, GetConnectionTestStatusResponse, GetConnectionsData, GetConnectionsResponse, GetCurrentUserInfoResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, GetDagData, GetDagDetailsData, GetDagDetailsResponse, GetDagResponse, GetDagRunData, GetDagRunDeadlinesData, GetDagRunDeadlinesResponse, GetDagRunResponse, GetDagRunsData, GetDagRunsResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetDagStructureData, GetDagStructureResponse, GetDagTagsData, GetDagTagsResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetDagsData, GetDagsResponse, GetDagsUiData, GetDagsUiResponse, GetDependenciesData, GetDependenciesResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, GetExtraLinksData, GetExtraLinksResponse, GetGanttDataData, GetGanttDataResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesData, GetGridTiSummariesResponse, GetGridTiSummariesStreamData, GetGridTiSummariesStreamResponse, GetHealthResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetLogData, GetLogResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetPartitionedDagRunsData, GetPartitionedDagRunsResponse, GetPendingPartitionedDagRunData, GetPendingPartitionedDagRunResponse, GetPluginsData, GetPluginsResponse, GetPoolData, GetPoolResponse, GetPoolsData, GetPoolsResponse, GetProvidersData, GetProvidersResponse, GetTaskData, GetTaskInstanceData, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstancesData, GetTaskInstancesResponse, GetTaskResponse, GetTasksData, GetTasksResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, GetVariableData, GetVariableResponse, GetVariablesData, GetVariablesResponse, GetVersionResponse, GetXcomEntriesData, GetXcomEntriesResponse, GetXcomEntryData, GetXcomEntryResponse, HistoricalMetricsData, HistoricalMetricsResponse, HookMetaDataResponse, ImportErrorsResponse, ListBackfillsData, ListBackfillsResponse, ListBackfillsUiData, ListBackfillsUiResponse, ListDagWarningsData, ListDagWarningsResponse, ListTeamsData, ListTeamsResponse, LoginData, LoginResponse, LogoutResponse, MaterializeAssetData, MaterializeAssetResponse, NextRunAssetsData, NextRunAssetsResponse, PatchConnectionAndTestData, PatchConnectionAndTestResponse, PatchConnectionData, PatchConnectionResponse, PatchDagData, PatchDagResponse, PatchDagRunData, PatchDagRunResponse, PatchDagsData, PatchDagsResponse, PatchPoolData, PatchPoolResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, PatchTaskInstanceData, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, PatchTaskInstanceResponse, PatchVariableData, PatchVariableResponse, PauseBackfillData, PauseBackfillResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PostConnectionData, PostConnectionResponse, PostPoolData, PostPoolResponse, PostVariableData, PostVariableResponse, ReparseDagFileData, ReparseDagFileResponse, StructureDataData, StructureDataResponse2, TestConnectionAsyncData, TestConnectionAsyncResponse, TestConnectionData, TestConnectionResponse, TriggerDagRunData, TriggerDagRunResponse, UnfavoriteDagData, UnfavoriteDagResponse, UnpauseBackfillData, UnpauseBackfillResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse } from './types.gen'; +import type { GetAssetsData, GetAssetsResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetAliasData, GetAssetAliasResponse, GetAssetEventsData, GetAssetEventsResponse, CreateAssetEventData, CreateAssetEventResponse, MaterializeAssetData, MaterializeAssetResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, GetAssetData, GetAssetResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, NextRunAssetsData, NextRunAssetsResponse, ListBackfillsData, ListBackfillsResponse, CreateBackfillData, CreateBackfillResponse, GetBackfillData, GetBackfillResponse, PauseBackfillData, PauseBackfillResponse, UnpauseBackfillData, UnpauseBackfillResponse, CancelBackfillData, CancelBackfillResponse, CreateBackfillDryRunData, CreateBackfillDryRunResponse, ListBackfillsUiData, ListBackfillsUiResponse, DeleteConnectionData, DeleteConnectionResponse, GetConnectionData, GetConnectionResponse, PatchConnectionData, PatchConnectionResponse, GetConnectionsData, GetConnectionsResponse, PostConnectionData, PostConnectionResponse, BulkConnectionsData, BulkConnectionsResponse, PatchConnectionAndTestData, PatchConnectionAndTestResponse, TestConnectionData, TestConnectionResponse, TestConnectionAsyncData, TestConnectionAsyncResponse, GetConnectionTestStatusData, GetConnectionTestStatusResponse, CreateDefaultConnectionsResponse, HookMetaDataResponse, GetDagRunData, GetDagRunResponse, DeleteDagRunData, DeleteDagRunResponse, PatchDagRunData, PatchDagRunResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, ClearDagRunData, ClearDagRunResponse, GetDagRunsData, GetDagRunsResponse, TriggerDagRunData, TriggerDagRunResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, ListDagWarningsData, ListDagWarningsResponse, GetDagsData, GetDagsResponse, PatchDagsData, PatchDagsResponse, GetDagData, GetDagResponse, PatchDagData, PatchDagResponse, DeleteDagData, DeleteDagResponse, GetDagDetailsData, GetDagDetailsResponse, FavoriteDagData, FavoriteDagResponse, UnfavoriteDagData, UnfavoriteDagResponse, GetDagTagsData, GetDagTagsResponse, GetDagsUiData, GetDagsUiResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExtraLinksData, GetExtraLinksResponse, GetTaskInstanceData, GetTaskInstanceResponse, PatchTaskInstanceData, PatchTaskInstanceResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, GetTaskInstancesData, GetTaskInstancesResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, GetLogData, GetLogResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetPluginsData, GetPluginsResponse, ImportErrorsResponse, DeletePoolData, DeletePoolResponse, GetPoolData, GetPoolResponse, PatchPoolData, PatchPoolResponse, GetPoolsData, GetPoolsResponse, PostPoolData, PostPoolResponse, BulkPoolsData, BulkPoolsResponse, GetProvidersData, GetProvidersResponse, GetXcomEntryData, GetXcomEntryResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, GetXcomEntriesData, GetXcomEntriesResponse, CreateXcomEntryData, CreateXcomEntryResponse, GetTasksData, GetTasksResponse, GetTaskData, GetTaskResponse, DeleteVariableData, DeleteVariableResponse, GetVariableData, GetVariableResponse, PatchVariableData, PatchVariableResponse, GetVariablesData, GetVariablesResponse, PostVariableData, PostVariableResponse, BulkVariablesData, BulkVariablesResponse, ReparseDagFileData, ReparseDagFileResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetHealthResponse, GetVersionResponse, LoginData, LoginResponse, LogoutResponse, GetAuthMenusResponse, GetCurrentUserInfoResponse, GenerateTokenData, GenerateTokenResponse2, GetPartitionedDagRunsData, GetPartitionedDagRunsResponse, GetPendingPartitionedDagRunData, GetPendingPartitionedDagRunResponse, GetDependenciesData, GetDependenciesResponse, HistoricalMetricsData, HistoricalMetricsResponse, DagStatsResponse2, GetDagRunDeadlinesData, GetDagRunDeadlinesResponse, StructureDataData, StructureDataResponse2, GetDagStructureData, GetDagStructureResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesStreamData, GetGridTiSummariesStreamResponse, GetGanttDataData, GetGanttDataResponse, GetCalendarData, GetCalendarResponse, ListTeamsData, ListTeamsResponse } from './types.gen'; export class AssetService { /** From b31d1c333bdda8a6604f8d7df373864f4ead32d6 Mon Sep 17 00:00:00 2001 From: Anish Date: Sat, 21 Mar 2026 14:27:46 -0500 Subject: [PATCH 33/38] Redesign async connection testing to use request-buffer pattern --- airflow-core/docs/migrations-ref.rst | 3 +- .../core_api/datamodels/connections.py | 17 +- .../openapi/v2-rest-api-generated.yaml | 176 ++----- .../core_api/routes/public/connections.py | 127 ++--- .../datamodels/connection_test.py | 17 +- .../execution_api/routes/connection_tests.py | 49 +- .../execution_api/versions/v2026_03_31.py | 3 +- .../src/airflow/executors/base_executor.py | 9 +- .../src/airflow/executors/local_executor.py | 7 +- .../src/airflow/jobs/scheduler_job_runner.py | 30 +- .../0110_3_2_0_add_connection_test_table.py | 32 +- .../src/airflow/models/connection_test.py | 248 ++++----- .../airflow/ui/openapi-gen/queries/common.ts | 11 +- .../ui/openapi-gen/queries/ensureQueryData.ts | 6 +- .../ui/openapi-gen/queries/prefetch.ts | 6 +- .../airflow/ui/openapi-gen/queries/queries.ts | 40 +- .../ui/openapi-gen/queries/suspense.ts | 6 +- .../ui/openapi-gen/requests/schemas.gen.ts | 102 +++- .../ui/openapi-gen/requests/services.gen.ts | 54 +- .../ui/openapi-gen/requests/types.gen.ts | 79 +-- airflow-core/src/airflow/utils/db_cleanup.py | 2 +- .../routes/public/test_connections.py | 242 +++------ .../versions/head/test_connection_tests.py | 176 ++++--- .../tests/unit/jobs/test_scheduler_job.py | 151 ++---- .../tests/unit/models/test_connection_test.py | 310 +++++------ .../airflowctl/api/datamodels/generated.py | 19 +- dev/test_async_connection.py | 483 ++++++++++++++++++ .../src/tests_common/test_utils/db.py | 4 +- task-sdk/src/airflow/sdk/api/client.py | 5 + .../airflow/sdk/api/datamodels/_generated.py | 15 + 30 files changed, 1311 insertions(+), 1118 deletions(-) create mode 100644 dev/test_async_connection.py diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index d1e02fbf22c70..27d26c86d281a 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,8 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``a7e6d4c3b2f1`` (head) | ``1d6611b6ab7c`` | ``3.2.0`` | Add connection_test table for async connection testing. | +| ``a7e6d4c3b2f1`` (head) | ``1d6611b6ab7c`` | ``3.2.0`` | Add connection_test_request table for async connection | +| | | | testing. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``1d6611b6ab7c`` | ``888b59e02a5b`` | ``3.2.0`` | Add bundle_name to callback table. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index 1e51b8ceb10ff..f13c40c0c194b 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -83,6 +83,14 @@ 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 executor: str | None = None queue: str | None = None @@ -103,15 +111,6 @@ class ConnectionTestStatusResponse(BaseModel): state: str result_message: str | None = None created_at: datetime - reverted: bool = False - - -class ConnectionSaveAndTestResponse(BaseModel): - """Response returned by the combined save-and-test endpoint.""" - - connection: ConnectionResponse - test_token: str - test_state: str class ConnectionHookFieldBehavior(BaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 9de1315908bb4..7bde23abb32e2 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -1335,6 +1335,12 @@ paths: schema: $ref: '#/components/schemas/HTTPExceptionResponse' description: Not Found + '409': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Conflict '422': description: Validation Error content: @@ -1631,103 +1637,6 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' - /api/v2/connections/{connection_id}/test: - patch: - tags: - - Connection - summary: Patch Connection And Test - description: 'Update a connection and queue an async test with revert-on-failure. - - - Atomically saves the edit and creates a ConnectionTest with snapshots of the - - pre-edit and post-edit state. If the test fails, the connection is automatically - - reverted to its pre-edit values.' - operationId: patch_connection_and_test - security: - - OAuth2PasswordBearer: [] - - HTTPBearer: [] - parameters: - - name: connection_id - in: path - required: true - schema: - type: string - title: Connection Id - - name: update_mask - in: query - required: false - schema: - anyOf: - - type: array - items: - type: string - - type: 'null' - title: Update Mask - - name: executor - in: query - required: false - schema: - anyOf: - - type: string - - type: 'null' - description: Executor to route the connection test to - title: Executor - description: Executor to route the connection test to - - name: queue - in: query - required: false - schema: - anyOf: - - type: string - - type: 'null' - description: Queue to route the connection test to - title: Queue - description: Queue to route the connection test to - requestBody: - required: true - content: - application/json: - schema: - $ref: '#/components/schemas/ConnectionBody' - responses: - '200': - description: Successful Response - content: - application/json: - schema: - $ref: '#/components/schemas/ConnectionSaveAndTestResponse' - '401': - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPExceptionResponse' - description: Unauthorized - '403': - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPExceptionResponse' - description: Forbidden - '400': - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPExceptionResponse' - description: Bad Request - '404': - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPExceptionResponse' - description: Not Found - '422': - description: Validation Error - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPValidationError' /api/v2/connections/test: post: tags: @@ -1786,9 +1695,11 @@ paths: description: 'Queue an async connection test to be executed on a worker. - The connection must already be saved. Returns a token that can be used + The connection data is stored in the test request table and the worker - to poll for the test result via GET /connections/test-async/{token}.' + reads from there. Returns a token to poll for the result via + + GET /connections/test-async/{token}.' operationId: test_connection_async requestBody: content: @@ -1815,8 +1726,8 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPExceptionResponse' - '404': - description: Not Found + '409': + description: Conflict content: application/json: schema: @@ -1834,11 +1745,11 @@ paths: get: tags: - Connection - summary: Get Connection Test Status + summary: Get Connection Test description: "Poll for the status of an async connection test.\n\nKnowledge\ \ of the token serves as authorization \u2014 only the client\nthat initiated\ \ the test knows the crypto-random token." - operationId: get_connection_test_status + operationId: get_connection_test security: - OAuth2PasswordBearer: [] - HTTPBearer: [] @@ -10295,23 +10206,6 @@ components: - team_name title: ConnectionResponse description: Connection serializer for responses. - ConnectionSaveAndTestResponse: - properties: - connection: - $ref: '#/components/schemas/ConnectionResponse' - test_token: - type: string - title: Test Token - test_state: - type: string - title: Test State - type: object - required: - - connection - - test_token - - test_state - title: ConnectionSaveAndTestResponse - description: Response returned by the combined save-and-test endpoint. ConnectionTestQueuedResponse: properties: token: @@ -10335,6 +10229,43 @@ components: connection_id: type: string title: Connection Id + conn_type: + type: string + title: Conn Type + host: + anyOf: + - type: string + - type: 'null' + title: Host + login: + anyOf: + - type: string + - type: 'null' + title: Login + schema: + anyOf: + - type: string + - type: 'null' + title: Schema + port: + anyOf: + - type: integer + - type: 'null' + title: Port + password: + anyOf: + - type: string + - type: 'null' + title: Password + extra: + anyOf: + - type: string + - type: 'null' + title: Extra + commit_on_success: + type: boolean + title: Commit On Success + default: false executor: anyOf: - type: string @@ -10349,6 +10280,7 @@ components: type: object required: - connection_id + - conn_type title: ConnectionTestRequestBody description: Request body for async connection test. ConnectionTestResponse: @@ -10385,10 +10317,6 @@ components: type: string format: date-time title: Created At - reverted: - type: boolean - title: Reverted - default: false type: object required: - token diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 2a0a18a060ba0..ac086caa2e434 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -41,7 +41,6 @@ ConnectionBodyPartial, ConnectionCollectionResponse, ConnectionResponse, - ConnectionSaveAndTestResponse, ConnectionTestQueuedResponse, ConnectionTestRequestBody, ConnectionTestResponse, @@ -61,7 +60,7 @@ from airflow.configuration import conf from airflow.exceptions import AirflowNotFoundException from airflow.models import Connection -from airflow.models.connection_test import ConnectionTest, snapshot_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 @@ -79,10 +78,26 @@ def _ensure_test_connection_enabled() -> None: ) +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.", + ) + + @connections_router.delete( "/{connection_id}", status_code=status.HTTP_204_NO_CONTENT, - responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_409_CONFLICT]), dependencies=[Depends(requires_access_connection(method="DELETE")), Depends(action_logging())], ) def delete_connection( @@ -90,6 +105,8 @@ def delete_connection( session: SessionDep, ): """Delete a connection entry.""" + _check_no_active_test(connection_id, session) + connection = session.scalar(select(Connection).filter_by(conn_id=connection_id)) if connection is None: @@ -189,70 +206,6 @@ def bulk_connections( return BulkConnectionService(session=session, request=request).handle_request() -@connections_router.patch( - "/{connection_id}/test", - responses=create_openapi_http_exception_doc( - [ - status.HTTP_400_BAD_REQUEST, - status.HTTP_403_FORBIDDEN, - status.HTTP_404_NOT_FOUND, - ] - ), - dependencies=[Depends(requires_access_connection(method="PUT")), Depends(action_logging())], -) -def patch_connection_and_test( - connection_id: str, - patch_body: ConnectionBody, - session: SessionDep, - update_mask: list[str] | None = Query(None), - executor: str | None = Query(None, description="Executor to route the connection test to"), - queue: str | None = Query(None, description="Queue to route the connection test to"), -) -> ConnectionSaveAndTestResponse: - """ - Update a connection and queue an async test with revert-on-failure. - - Atomically saves the edit and creates a ConnectionTest with snapshots of the - pre-edit and post-edit state. If the test fails, the connection is automatically - reverted to its pre-edit values. - """ - _ensure_test_connection_enabled() - - if patch_body.connection_id != connection_id: - raise HTTPException( - status.HTTP_400_BAD_REQUEST, - "The connection_id in the request body does not match the URL parameter", - ) - - connection = session.scalar(select(Connection).filter_by(conn_id=connection_id).limit(1)) - if connection is None: - raise HTTPException( - status.HTTP_404_NOT_FOUND, - f"The Connection with connection_id: `{connection_id}` was not found", - ) - - try: - ConnectionBody(**patch_body.model_dump()) - except ValidationError as e: - raise RequestValidationError(errors=e.errors()) - - pre_snapshot = snapshot_connection(connection) - - update_orm_from_pydantic(connection, patch_body, update_mask) - - post_snapshot = snapshot_connection(connection) - - connection_test = ConnectionTest(connection_id=connection_id, executor=executor, queue=queue) - connection_test.connection_snapshot = {"pre": pre_snapshot, "post": post_snapshot} - session.add(connection_test) - session.flush() - - return ConnectionSaveAndTestResponse( - connection=connection, - test_token=connection_test.token, - test_state=connection_test.state, - ) - - @connections_router.patch( "/{connection_id}", responses=create_openapi_http_exception_doc( @@ -270,6 +223,8 @@ def patch_connection( update_mask: list[str] | None = Query(None), ) -> ConnectionResponse: """Update a connection entry.""" + _check_no_active_test(connection_id, session) + if patch_body.connection_id != connection_id: raise HTTPException( status.HTTP_400_BAD_REQUEST, @@ -334,7 +289,7 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse: @connections_router.post( "/test-async", status_code=status.HTTP_202_ACCEPTED, - responses=create_openapi_http_exception_doc([status.HTTP_403_FORBIDDEN, status.HTTP_404_NOT_FOUND]), + 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( @@ -344,22 +299,27 @@ def test_connection_async( """ Queue an async connection test to be executed on a worker. - The connection must already be saved. Returns a token that can be used - to poll for the test result via GET /connections/test-async/{token}. + 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() - try: - Connection.get_connection_from_secrets(test_body.connection_id) - except AirflowNotFoundException: - raise HTTPException( - status.HTTP_404_NOT_FOUND, - f"The Connection with connection_id: `{test_body.connection_id}` was not found. " - "Connection must be saved before testing.", - ) - - connection_test = ConnectionTest( - connection_id=test_body.connection_id, executor=test_body.executor, queue=test_body.queue + # 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() @@ -376,7 +336,7 @@ def test_connection_async( responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), dependencies=[Depends(requires_access_connection(method="GET"))], ) -def get_connection_test_status( +def get_connection_test( connection_test_token: str, session: SessionDep, ) -> ConnectionTestStatusResponse: @@ -386,7 +346,7 @@ def get_connection_test_status( Knowledge of the token serves as authorization — only the client that initiated the test knows the crypto-random token. """ - connection_test = session.scalar(select(ConnectionTest).filter_by(token=connection_test_token)) + connection_test = session.scalar(select(ConnectionTestRequest).filter_by(token=connection_test_token)) if connection_test is None: raise HTTPException( @@ -400,7 +360,6 @@ def get_connection_test_status( state=connection_test.state, result_message=connection_test.result_message, created_at=connection_test.created_at, - reverted=connection_test.reverted, ) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py index 3fd046e90b6e7..e0b63e8dae55c 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/connection_test.py @@ -17,7 +17,9 @@ from __future__ import annotations -from airflow.api_fastapi.core_api.base import StrictBaseModel +from pydantic import Field + +from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel from airflow.models.connection_test import ConnectionTestState @@ -26,3 +28,16 @@ class ConnectionTestResultBody(StrictBaseModel): state: ConnectionTestState result_message: str | None = None + + +class ConnectionTestConnectionResponse(BaseModel): + """Connection data returned to workers from a test request.""" + + conn_id: str + conn_type: str + host: str | None = None + login: str | None = None + password: str | None = None + schema_: str | None = Field(None, alias="schema") + port: int | None = None + extra: str | None = None diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py index d40026b7ef1e1..de6f643ea2de3 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -22,17 +22,52 @@ from fastapi import HTTPException, status from airflow.api_fastapi.common.db.common import SessionDep -from airflow.api_fastapi.execution_api.datamodels.connection_test import ConnectionTestResultBody +from airflow.api_fastapi.execution_api.datamodels.connection_test import ( + ConnectionTestConnectionResponse, + ConnectionTestResultBody, +) from airflow.models.connection_test import ( TERMINAL_STATES, - ConnectionTest, + ConnectionTestRequest, ConnectionTestState, - attempt_revert, ) router = VersionedAPIRouter() +@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, + ) + + @router.patch( "/{connection_test_id}", status_code=status.HTTP_204_NO_CONTENT, @@ -47,7 +82,7 @@ def patch_connection_test( session: SessionDep, ) -> None: """Update the result of a connection test (called by workers).""" - ct = session.get(ConnectionTest, connection_test_id, with_for_update=True) + ct = session.get(ConnectionTestRequest, connection_test_id, with_for_update=True) if ct is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -69,7 +104,5 @@ def patch_connection_test( ct.state = body.state ct.result_message = body.result_message - if body.state == ConnectionTestState.FAILED and ct.connection_snapshot: - attempt_revert(ct, session=session) - elif body.state == ConnectionTestState.SUCCESS: - ct.connection_snapshot = None + if body.state == ConnectionTestState.SUCCESS and ct.commit_on_success: + ct.commit_to_connection_table(session=session) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py index bd5f9c41e0ecf..76dab90340abe 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py @@ -30,12 +30,13 @@ class AddConnectionTestEndpoint(VersionChange): - """Add connection-tests endpoint for async connection testing.""" + """Add connection-tests endpoints for async connection testing.""" description = __doc__ instructions_to_migrate_to_previous_version = ( endpoint("/connection-tests/{connection_test_id}", ["PATCH"]).didnt_exist, + endpoint("/connection-tests/{connection_test_id}/connection", ["GET"]).didnt_exist, ) diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index ddaf8c199c22c..24f77add4bec3 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -322,11 +322,16 @@ def heartbeat(self) -> None: self.sync() def trigger_connection_tests(self) -> None: - """Process queued connection tests.""" + """Process queued connection tests, respecting available slot capacity.""" if not self.supports_connection_test or not self.queued_connection_tests: return - self._process_workloads(list(self.queued_connection_tests.values())) + available = self.slots_available + if available <= 0: + return + + tests_to_run = list(self.queued_connection_tests.values())[:available] + self._process_workloads(tests_to_run) def _get_metric_name(self, metric_base_name: str) -> str: return ( diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 43aaa6801b27c..1bb705593f22b 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -192,7 +192,6 @@ def _execute_connection_test(log: Logger, workload: workloads.TestConnection, te # Lazy import: SDK modules must not be loaded at module level to avoid # coupling core (scheduler-loaded) code to the SDK. from airflow.sdk.api.client import Client - from airflow.sdk.execution_time.comms import ErrorResponse setproctitle( f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}", @@ -215,9 +214,7 @@ def _handle_timeout(signum, frame): try: client.connection_tests.update_state(workload.connection_test_id, ConnectionTestState.RUNNING) - conn_response = client.connections.get(workload.connection_id) - if isinstance(conn_response, ErrorResponse): - raise RuntimeError(f"Connection '{workload.connection_id}' not found via Execution API") + conn_response = client.connection_tests.get_connection(workload.connection_test_id) conn = Connection( conn_id=conn_response.conn_id, @@ -249,7 +246,7 @@ def _handle_timeout(signum, frame): client.connection_tests.update_state( workload.connection_test_id, ConnectionTestState.FAILED, - f"Connection test failed unexpectedly: {e}"[:500], + f"Connection test failed unexpectedly: {type(e).__name__}", ) finally: signal.alarm(0) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 563b83cde658f..5783d2ddef874 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -73,7 +73,12 @@ ) from airflow.models.backfill import Backfill from airflow.models.callback import Callback, CallbackType, ExecutorCallback -from airflow.models.connection_test import ACTIVE_STATES, ConnectionTest, ConnectionTestState, attempt_revert +from airflow.models.connection_test import ( + ACTIVE_STATES, + DISPATCHED_STATES, + ConnectionTestRequest, + ConnectionTestState, +) from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion from airflow.models.dagbag import DBDagBag @@ -3134,7 +3139,9 @@ def _enqueue_connection_tests(self, *, session: Session) -> None: return active_count = session.scalar( - select(func.count(ConnectionTest.id)).where(ConnectionTest.state.in_(ACTIVE_STATES)) + 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) @@ -3142,12 +3149,12 @@ def _enqueue_connection_tests(self, *, session: Session) -> None: return pending_stmt = ( - select(ConnectionTest) - .where(ConnectionTest.state == ConnectionTestState.PENDING) - .order_by(ConnectionTest.created_at) + select(ConnectionTestRequest) + .where(ConnectionTestRequest.state == ConnectionTestState.PENDING) + .order_by(ConnectionTestRequest.created_at) .limit(budget) ) - pending_stmt = with_row_locks(pending_stmt, session, of=ConnectionTest, skip_locked=True) + pending_stmt = with_row_locks(pending_stmt, session, of=ConnectionTestRequest, skip_locked=True) pending_tests = session.scalars(pending_stmt).all() if not pending_tests: @@ -3165,7 +3172,6 @@ def _enqueue_connection_tests(self, *, session: Session) -> None: ) ct.state = ConnectionTestState.FAILED ct.result_message = reason - ct.connection_snapshot = None self.log.warning("Failing connection test %s: %s", ct.id, reason) continue @@ -3188,19 +3194,17 @@ def _reap_stale_connection_tests(self, *, session: Session = NEW_SESSION) -> Non grace_period = max(30, timeout // 2) cutoff = timezone.utcnow() - timedelta(seconds=timeout + grace_period) - stale_stmt = select(ConnectionTest).where( - ConnectionTest.state.in_(ACTIVE_STATES), - ConnectionTest.updated_at < cutoff, + stale_stmt = select(ConnectionTestRequest).where( + ConnectionTestRequest.state.in_(ACTIVE_STATES), + ConnectionTestRequest.updated_at < cutoff, ) - stale_stmt = with_row_locks(stale_stmt, session, of=ConnectionTest, skip_locked=True) + 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) - if ct.connection_snapshot: - attempt_revert(ct, session=session) session.flush() diff --git a/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py b/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py index 33a3625b95025..d8d3c360eb0c9 100644 --- a/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py +++ b/airflow-core/src/airflow/migrations/versions/0110_3_2_0_add_connection_test_table.py @@ -17,7 +17,7 @@ # under the License. """ -Add connection_test table for async connection testing. +Add connection_test_request table for async connection testing. Revision ID: a7e6d4c3b2f1 Revises: 1d6611b6ab7c @@ -41,31 +41,37 @@ def upgrade(): - """Create connection_test table.""" + """Create connection_test_request table.""" op.create_table( - "connection_test", + "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(10), 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("connection_snapshot", sa.JSON(), nullable=True), - sa.Column("reverted", sa.Boolean(), nullable=False, server_default="0"), - sa.PrimaryKeyConstraint("id", name=op.f("connection_test_pkey")), - sa.UniqueConstraint("token", name=op.f("connection_test_token_uq")), + sa.Column("conn_type", sa.String(500), nullable=False), + sa.Column("host", sa.String(500), nullable=True), + sa.Column("login", sa.Text(), nullable=True), + sa.Column("password", sa.Text(), nullable=True), + sa.Column("schema", sa.String(500), nullable=True), + sa.Column("port", sa.Integer(), nullable=True), + sa.Column("extra", sa.Text(), nullable=True), + sa.Column("commit_on_success", sa.Boolean(), nullable=False, server_default="0"), + sa.PrimaryKeyConstraint("id", name=op.f("connection_test_request_pkey")), + sa.UniqueConstraint("token", name=op.f("connection_test_request_token_uq")), ) op.create_index( - op.f("idx_connection_test_state_created_at"), - "connection_test", + op.f("idx_connection_test_request_state_created_at"), + "connection_test_request", ["state", "created_at"], ) def downgrade(): - """Drop connection_test table.""" - op.drop_index(op.f("idx_connection_test_state_created_at"), table_name="connection_test") - op.drop_table("connection_test") + """Drop connection_test_request table.""" + op.drop_index(op.f("idx_connection_test_request_state_created_at"), table_name="connection_test_request") + op.drop_table("connection_test_request") diff --git a/airflow-core/src/airflow/models/connection_test.py b/airflow-core/src/airflow/models/connection_test.py index 64c306f9b53cf..8f8ca4abdbbb2 100644 --- a/airflow-core/src/airflow/models/connection_test.py +++ b/airflow-core/src/airflow/models/connection_test.py @@ -24,8 +24,8 @@ import structlog import uuid6 -from sqlalchemy import JSON, Boolean, Index, String, Text, Uuid, select -from sqlalchemy.orm import Mapped, mapped_column +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 @@ -52,19 +52,28 @@ def __str__(self) -> str: return self.value -ACTIVE_STATES = frozenset((ConnectionTestState.QUEUED, ConnectionTestState.RUNNING)) +ACTIVE_STATES = frozenset( + (ConnectionTestState.PENDING, ConnectionTestState.QUEUED, ConnectionTestState.RUNNING) +) +DISPATCHED_STATES = frozenset((ConnectionTestState.QUEUED, ConnectionTestState.RUNNING)) TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED)) -class ConnectionTest(Base): - """Tracks an async connection test dispatched to a worker via a TestConnection workload.""" +class ConnectionTestRequest(Base): + """ + Tracks an async connection test request dispatched to a worker. - __tablename__ = "connection_test" + 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(10), nullable=False, default=ConnectionTestState.PENDING) + 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( @@ -72,23 +81,56 @@ class ConnectionTest(Base): ) executor: Mapped[str | None] = mapped_column(String(256), nullable=True) queue: Mapped[str | None] = mapped_column(String(256), nullable=True) - connection_snapshot: Mapped[dict | None] = mapped_column(JSON(none_as_null=True), nullable=True) - reverted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default="0") - __table_args__ = (Index("idx_connection_test_state_created_at", state, created_at),) + # 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, executor: str | None = None, queue: str | None = None, **kwargs + self, + *, + connection_id: str, + conn_type: str, + host: str | None = None, + login: str | None = None, + password: str | None = None, + schema: str | None = None, + port: int | None = None, + extra: str | None = None, + commit_on_success: bool = False, + executor: str | None = None, + queue: str | None = None, + **kwargs, ): super().__init__(**kwargs) self.connection_id = connection_id + self.conn_type = conn_type + self.host = host + self.login = login + self.password = password + self.schema = schema + self.port = port + self.extra = extra + self.commit_on_success = commit_on_success self.executor = executor self.queue = queue self.token = secrets.token_urlsafe(32) self.state = ConnectionTestState.PENDING def __repr__(self) -> str: - return f"" + return ( + f"" + ) def get_executor_name(self) -> str | None: """Return the executor name for scheduler routing.""" @@ -98,6 +140,85 @@ def get_dag_id(self) -> None: """Return None — connection tests are not associated with any DAG.""" return None + def get_password(self) -> str | None: + if self._password: + fernet = get_fernet() + if not fernet.is_encrypted: + return self._password + return fernet.decrypt(bytes(self._password, "utf-8")).decode() + return self._password + + def set_password(self, value: str | None): + if value: + fernet = get_fernet() + self._password = fernet.encrypt(bytes(value, "utf-8")).decode() + else: + self._password = value + + @declared_attr + def password(cls): + """Password. The value is decrypted/encrypted when reading/setting the value.""" + return synonym("_password", descriptor=property(cls.get_password, cls.set_password)) + + def get_extra(self) -> str | None: + if self._extra: + fernet = get_fernet() + if not fernet.is_encrypted: + return self._extra + return fernet.decrypt(bytes(self._extra, "utf-8")).decode() + return self._extra + + def set_extra(self, value: str | None): + if value: + fernet = get_fernet() + self._extra = fernet.encrypt(bytes(value, "utf-8")).decode() + else: + self._extra = value + + @declared_attr + def extra(cls): + """Extra data. The value is decrypted/encrypted when reading/setting the value.""" + return synonym("_extra", descriptor=property(cls.get_extra, cls.set_extra)) + + def to_connection(self) -> Connection: + """Build a transient Connection object from the stored fields for testing.""" + return Connection( + conn_id=self.connection_id, + conn_type=self.conn_type, + host=self.host, + login=self.login, + password=self.password, + schema=self.schema, + port=self.port, + extra=self.extra, + ) + + def commit_to_connection_table(self, *, session: Session) -> None: + """Upsert the tested connection into the real ``connection`` table.""" + conn = session.scalar(select(Connection).filter_by(conn_id=self.connection_id)) + if conn is None: + conn = Connection( + conn_id=self.connection_id, + conn_type=self.conn_type, + host=self.host, + login=self.login, + password=self.password, + schema=self.schema, + port=self.port, + extra=self.extra, + ) + session.add(conn) + log.info("Created new connection from successful test", connection_id=self.connection_id) + else: + conn.conn_type = self.conn_type + conn.host = self.host + conn.login = self.login + conn.password = self.password + conn.schema = self.schema + conn.port = self.port + conn.extra = self.extra + log.info("Updated existing connection from successful test", connection_id=self.connection_id) + def run_connection_test(*, conn: Connection) -> tuple[bool, str]: """ @@ -111,106 +232,3 @@ def run_connection_test(*, conn: Connection) -> tuple[bool, str]: except Exception as e: log.exception("Connection test failed", connection_id=conn.conn_id) return False, str(e) - - -_SNAPSHOT_FIELDS = ( - "conn_type", - "description", - "host", - "login", - "_password", - "schema", - "port", - "_extra", - "is_encrypted", - "is_extra_encrypted", -) - - -def snapshot_connection(conn: Connection) -> dict: - """ - Capture raw DB column values from a Connection for later restore. - - Encrypted fields (``_password``, ``_extra``) are stored as ciphertext - so they can be written directly back without re-encryption. - """ - return {field: getattr(conn, field) for field in _SNAPSHOT_FIELDS} - - -def _revert_connection(conn: Connection, snapshot: dict) -> None: - """ - Restore a Connection's columns from a snapshot dict. - - Writes directly to ``_password`` and ``_extra`` (bypassing the - encrypting property setters) so the stored ciphertext is preserved. - """ - for field, value in snapshot.items(): - setattr(conn, field, value) - - -def _decrypt_snapshot_field(snapshot: dict, field: str) -> str | None: - """Decrypt a single encrypted field from a snapshot dict using Fernet.""" - raw = snapshot.get(field) - if raw is None: - return None - encrypted_flag = "is_encrypted" if field == "_password" else "is_extra_encrypted" - if not snapshot.get(encrypted_flag, False): - return raw - fernet = get_fernet() - return fernet.decrypt(bytes(raw, "utf-8")).decode() - - -def _can_safely_revert(conn: Connection, post_snapshot: dict) -> bool: - """ - Check whether the connection's current state matches the post-edit snapshot. - - Compares **decrypted** values for encrypted fields and direct values for - non-encrypted fields. Returns ``False`` if any field differs, indicating - a concurrent edit has occurred and the revert should be skipped. - """ - for field in _SNAPSHOT_FIELDS: - if field in ("is_encrypted", "is_extra_encrypted"): - continue - - if field == "_password": - current_val = conn.password - snapshot_val = _decrypt_snapshot_field(post_snapshot, "_password") - elif field == "_extra": - current_val = conn.extra - snapshot_val = _decrypt_snapshot_field(post_snapshot, "_extra") - else: - current_val = getattr(conn, field) - snapshot_val = post_snapshot.get(field) - - if current_val != snapshot_val: - return False - return True - - -def attempt_revert(ct: ConnectionTest, *, session: Session) -> None: - """Revert a connection to its pre-edit values if no concurrent edit has occurred.""" - if not ct.connection_snapshot: - log.warning("attempt_revert called without snapshot", connection_test_id=ct.id) - return - - pre_snapshot = ct.connection_snapshot["pre"] - post_snapshot = ct.connection_snapshot["post"] - - ct.connection_snapshot = None - - conn = session.scalar(select(Connection).filter_by(conn_id=ct.connection_id)) - if conn is None: - ct.result_message = (ct.result_message or "") + " | Revert skipped: connection no longer exists." - log.warning("Revert skipped: connection no longer exists", connection_id=ct.connection_id) - return - - if not _can_safely_revert(conn, post_snapshot): - ct.result_message = ( - ct.result_message or "" - ) + " | Revert skipped: connection was modified by another user." - log.warning("Revert skipped: concurrent edit detected", connection_id=ct.connection_id) - return - - _revert_connection(conn, pre_snapshot) - ct.reverted = True - log.info("Reverted connection to pre-edit state", connection_id=ct.connection_id) diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts index 59ee8d16cbf22..c68fe4bd49e49 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -122,12 +122,12 @@ export const UseConnectionServiceGetConnectionsKeyFn = ({ connectionIdPattern, l offset?: number; orderBy?: string[]; } = {}, queryKey?: Array) => [useConnectionServiceGetConnectionsKey, ...(queryKey ?? [{ connectionIdPattern, limit, offset, orderBy }])]; -export type ConnectionServiceGetConnectionTestStatusDefaultResponse = Awaited>; -export type ConnectionServiceGetConnectionTestStatusQueryResult = UseQueryResult; -export const useConnectionServiceGetConnectionTestStatusKey = "ConnectionServiceGetConnectionTestStatus"; -export const UseConnectionServiceGetConnectionTestStatusKeyFn = ({ connectionTestToken }: { +export type ConnectionServiceGetConnectionTestDefaultResponse = Awaited>; +export type ConnectionServiceGetConnectionTestQueryResult = UseQueryResult; +export const useConnectionServiceGetConnectionTestKey = "ConnectionServiceGetConnectionTest"; +export const UseConnectionServiceGetConnectionTestKeyFn = ({ connectionTestToken }: { connectionTestToken: string; -}, queryKey?: Array) => [useConnectionServiceGetConnectionTestStatusKey, ...(queryKey ?? [{ connectionTestToken }])]; +}, queryKey?: Array) => [useConnectionServiceGetConnectionTestKey, ...(queryKey ?? [{ connectionTestToken }])]; export type ConnectionServiceHookMetaDataDefaultResponse = Awaited>; export type ConnectionServiceHookMetaDataQueryResult = UseQueryResult; export const useConnectionServiceHookMetaDataKey = "ConnectionServiceHookMetaData"; @@ -948,7 +948,6 @@ export type BackfillServiceCancelBackfillMutationResult = Awaited>; export type ConnectionServicePatchConnectionMutationResult = Awaited>; export type ConnectionServiceBulkConnectionsMutationResult = Awaited>; -export type ConnectionServicePatchConnectionAndTestMutationResult = Awaited>; export type DagRunServicePatchDagRunMutationResult = Awaited>; export type DagServicePatchDagsMutationResult = Awaited>; export type DagServicePatchDagMutationResult = Awaited>; diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts index 5e551e0ab9375..e1751b9116ee1 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -225,7 +225,7 @@ export const ensureUseConnectionServiceGetConnectionsData = (queryClient: QueryC orderBy?: string[]; } = {}) => queryClient.ensureQueryData({ queryKey: Common.UseConnectionServiceGetConnectionsKeyFn({ connectionIdPattern, limit, offset, orderBy }), queryFn: () => ConnectionService.getConnections({ connectionIdPattern, limit, offset, orderBy }) }); /** -* Get Connection Test Status +* Get Connection Test * Poll for the status of an async connection test. * * Knowledge of the token serves as authorization — only the client @@ -235,9 +235,9 @@ export const ensureUseConnectionServiceGetConnectionsData = (queryClient: QueryC * @returns ConnectionTestStatusResponse Successful Response * @throws ApiError */ -export const ensureUseConnectionServiceGetConnectionTestStatusData = (queryClient: QueryClient, { connectionTestToken }: { +export const ensureUseConnectionServiceGetConnectionTestData = (queryClient: QueryClient, { connectionTestToken }: { connectionTestToken: string; -}) => queryClient.ensureQueryData({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ connectionTestToken }), queryFn: () => ConnectionService.getConnectionTestStatus({ connectionTestToken }) }); +}) => queryClient.ensureQueryData({ queryKey: Common.UseConnectionServiceGetConnectionTestKeyFn({ connectionTestToken }), queryFn: () => ConnectionService.getConnectionTest({ connectionTestToken }) }); /** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts index 5ab8f3ab76a89..65b9ef1a0dc1c 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -225,7 +225,7 @@ export const prefetchUseConnectionServiceGetConnections = (queryClient: QueryCli orderBy?: string[]; } = {}) => queryClient.prefetchQuery({ queryKey: Common.UseConnectionServiceGetConnectionsKeyFn({ connectionIdPattern, limit, offset, orderBy }), queryFn: () => ConnectionService.getConnections({ connectionIdPattern, limit, offset, orderBy }) }); /** -* Get Connection Test Status +* Get Connection Test * Poll for the status of an async connection test. * * Knowledge of the token serves as authorization — only the client @@ -235,9 +235,9 @@ export const prefetchUseConnectionServiceGetConnections = (queryClient: QueryCli * @returns ConnectionTestStatusResponse Successful Response * @throws ApiError */ -export const prefetchUseConnectionServiceGetConnectionTestStatus = (queryClient: QueryClient, { connectionTestToken }: { +export const prefetchUseConnectionServiceGetConnectionTest = (queryClient: QueryClient, { connectionTestToken }: { connectionTestToken: string; -}) => queryClient.prefetchQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ connectionTestToken }), queryFn: () => ConnectionService.getConnectionTestStatus({ connectionTestToken }) }); +}) => queryClient.prefetchQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestKeyFn({ connectionTestToken }), queryFn: () => ConnectionService.getConnectionTest({ connectionTestToken }) }); /** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index f2f02439a3595..dd341e301993c 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -225,7 +225,7 @@ export const useConnectionServiceGetConnections = , "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseConnectionServiceGetConnectionsKeyFn({ connectionIdPattern, limit, offset, orderBy }, queryKey), queryFn: () => ConnectionService.getConnections({ connectionIdPattern, limit, offset, orderBy }) as TData, ...options }); /** -* Get Connection Test Status +* Get Connection Test * Poll for the status of an async connection test. * * Knowledge of the token serves as authorization — only the client @@ -235,9 +235,9 @@ export const useConnectionServiceGetConnections = = unknown[]>({ connectionTestToken }: { +export const useConnectionServiceGetConnectionTest = = unknown[]>({ connectionTestToken }: { connectionTestToken: string; -}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ connectionTestToken }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ connectionTestToken }) as TData, ...options }); +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestKeyFn({ connectionTestToken }, queryKey), queryFn: () => ConnectionService.getConnectionTest({ connectionTestToken }) as TData, ...options }); /** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. @@ -1845,8 +1845,9 @@ export const useConnectionServiceTestConnection = ({ mutationFn: ({ requestBody }) => ConnectionService.bulkConnections({ requestBody }) as unknown as Promise, ...options }); /** -* Patch Connection And Test -* Update a connection and queue an async test with revert-on-failure. -* -* Atomically saves the edit and creates a ConnectionTest with snapshots of the -* pre-edit and post-edit state. If the test fails, the connection is automatically -* reverted to its pre-edit values. -* @param data The data for the request. -* @param data.connectionId -* @param data.requestBody -* @param data.updateMask -* @param data.executor Executor to route the connection test to -* @param data.queue Queue to route the connection test to -* @returns ConnectionSaveAndTestResponse Successful Response -* @throws ApiError -*/ -export const useConnectionServicePatchConnectionAndTest = (options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ connectionId, executor, queue, requestBody, updateMask }) => ConnectionService.patchConnectionAndTest({ connectionId, executor, queue, requestBody, updateMask }) as unknown as Promise, ...options }); -/** * Patch Dag Run * Modify a DAG Run. * @param data The data for the request. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts index 6c5c042bca865..e2f532998f5eb 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -225,7 +225,7 @@ export const useConnectionServiceGetConnectionsSuspense = , "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseConnectionServiceGetConnectionsKeyFn({ connectionIdPattern, limit, offset, orderBy }, queryKey), queryFn: () => ConnectionService.getConnections({ connectionIdPattern, limit, offset, orderBy }) as TData, ...options }); /** -* Get Connection Test Status +* Get Connection Test * Poll for the status of an async connection test. * * Knowledge of the token serves as authorization — only the client @@ -235,9 +235,9 @@ export const useConnectionServiceGetConnectionsSuspense = = unknown[]>({ connectionTestToken }: { +export const useConnectionServiceGetConnectionTestSuspense = = unknown[]>({ connectionTestToken }: { connectionTestToken: string; -}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestStatusKeyFn({ connectionTestToken }, queryKey), queryFn: () => ConnectionService.getConnectionTestStatus({ connectionTestToken }) as TData, ...options }); +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseConnectionServiceGetConnectionTestKeyFn({ connectionTestToken }, queryKey), queryFn: () => ConnectionService.getConnectionTest({ connectionTestToken }) as TData, ...options }); /** * Hook Meta Data * Retrieve information about available connection types (hook classes) and their parameters. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 4107e3f88e3df..189e703f162ac 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1698,26 +1698,6 @@ export const $ConnectionResponse = { description: 'Connection serializer for responses.' } as const; -export const $ConnectionSaveAndTestResponse = { - properties: { - connection: { - '$ref': '#/components/schemas/ConnectionResponse' - }, - test_token: { - type: 'string', - title: 'Test Token' - }, - test_state: { - type: 'string', - title: 'Test State' - } - }, - type: 'object', - required: ['connection', 'test_token', 'test_state'], - title: 'ConnectionSaveAndTestResponse', - description: 'Response returned by the combined save-and-test endpoint.' -} as const; - export const $ConnectionTestQueuedResponse = { properties: { token: { @@ -1745,6 +1725,81 @@ export const $ConnectionTestRequestBody = { type: 'string', title: 'Connection Id' }, + conn_type: { + type: 'string', + title: 'Conn Type' + }, + host: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Host' + }, + login: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Login' + }, + schema: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Schema' + }, + port: { + anyOf: [ + { + type: 'integer' + }, + { + type: 'null' + } + ], + title: 'Port' + }, + password: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Password' + }, + extra: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Extra' + }, + commit_on_success: { + type: 'boolean', + title: 'Commit On Success', + default: false + }, executor: { anyOf: [ { @@ -1770,7 +1825,7 @@ export const $ConnectionTestRequestBody = { }, additionalProperties: false, type: 'object', - required: ['connection_id'], + required: ['connection_id', 'conn_type'], title: 'ConnectionTestRequestBody', description: 'Request body for async connection test.' } as const; @@ -1821,11 +1876,6 @@ export const $ConnectionTestStatusResponse = { type: 'string', format: 'date-time', title: 'Created At' - }, - reverted: { - type: 'boolean', - title: 'Reverted', - default: false } }, type: 'object', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 58e97cd997c40..7a633a252face 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -3,7 +3,7 @@ import type { CancelablePromise } from './core/CancelablePromise'; import { OpenAPI } from './core/OpenAPI'; import { request as __request } from './core/request'; -import type { GetAssetsData, GetAssetsResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetAliasData, GetAssetAliasResponse, GetAssetEventsData, GetAssetEventsResponse, CreateAssetEventData, CreateAssetEventResponse, MaterializeAssetData, MaterializeAssetResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, GetAssetData, GetAssetResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, NextRunAssetsData, NextRunAssetsResponse, ListBackfillsData, ListBackfillsResponse, CreateBackfillData, CreateBackfillResponse, GetBackfillData, GetBackfillResponse, PauseBackfillData, PauseBackfillResponse, UnpauseBackfillData, UnpauseBackfillResponse, CancelBackfillData, CancelBackfillResponse, CreateBackfillDryRunData, CreateBackfillDryRunResponse, ListBackfillsUiData, ListBackfillsUiResponse, DeleteConnectionData, DeleteConnectionResponse, GetConnectionData, GetConnectionResponse, PatchConnectionData, PatchConnectionResponse, GetConnectionsData, GetConnectionsResponse, PostConnectionData, PostConnectionResponse, BulkConnectionsData, BulkConnectionsResponse, PatchConnectionAndTestData, PatchConnectionAndTestResponse, TestConnectionData, TestConnectionResponse, TestConnectionAsyncData, TestConnectionAsyncResponse, GetConnectionTestStatusData, GetConnectionTestStatusResponse, CreateDefaultConnectionsResponse, HookMetaDataResponse, GetDagRunData, GetDagRunResponse, DeleteDagRunData, DeleteDagRunResponse, PatchDagRunData, PatchDagRunResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, ClearDagRunData, ClearDagRunResponse, GetDagRunsData, GetDagRunsResponse, TriggerDagRunData, TriggerDagRunResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, ListDagWarningsData, ListDagWarningsResponse, GetDagsData, GetDagsResponse, PatchDagsData, PatchDagsResponse, GetDagData, GetDagResponse, PatchDagData, PatchDagResponse, DeleteDagData, DeleteDagResponse, GetDagDetailsData, GetDagDetailsResponse, FavoriteDagData, FavoriteDagResponse, UnfavoriteDagData, UnfavoriteDagResponse, GetDagTagsData, GetDagTagsResponse, GetDagsUiData, GetDagsUiResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExtraLinksData, GetExtraLinksResponse, GetTaskInstanceData, GetTaskInstanceResponse, PatchTaskInstanceData, PatchTaskInstanceResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, GetTaskInstancesData, GetTaskInstancesResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, GetLogData, GetLogResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetPluginsData, GetPluginsResponse, ImportErrorsResponse, DeletePoolData, DeletePoolResponse, GetPoolData, GetPoolResponse, PatchPoolData, PatchPoolResponse, GetPoolsData, GetPoolsResponse, PostPoolData, PostPoolResponse, BulkPoolsData, BulkPoolsResponse, GetProvidersData, GetProvidersResponse, GetXcomEntryData, GetXcomEntryResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, GetXcomEntriesData, GetXcomEntriesResponse, CreateXcomEntryData, CreateXcomEntryResponse, GetTasksData, GetTasksResponse, GetTaskData, GetTaskResponse, DeleteVariableData, DeleteVariableResponse, GetVariableData, GetVariableResponse, PatchVariableData, PatchVariableResponse, GetVariablesData, GetVariablesResponse, PostVariableData, PostVariableResponse, BulkVariablesData, BulkVariablesResponse, ReparseDagFileData, ReparseDagFileResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetHealthResponse, GetVersionResponse, LoginData, LoginResponse, LogoutResponse, GetAuthMenusResponse, GetCurrentUserInfoResponse, GenerateTokenData, GenerateTokenResponse2, GetPartitionedDagRunsData, GetPartitionedDagRunsResponse, GetPendingPartitionedDagRunData, GetPendingPartitionedDagRunResponse, GetDependenciesData, GetDependenciesResponse, HistoricalMetricsData, HistoricalMetricsResponse, DagStatsResponse2, GetDagRunDeadlinesData, GetDagRunDeadlinesResponse, StructureDataData, StructureDataResponse2, GetDagStructureData, GetDagStructureResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesStreamData, GetGridTiSummariesStreamResponse, GetGanttDataData, GetGanttDataResponse, GetCalendarData, GetCalendarResponse, ListTeamsData, ListTeamsResponse } from './types.gen'; +import type { GetAssetsData, GetAssetsResponse, GetAssetAliasesData, GetAssetAliasesResponse, GetAssetAliasData, GetAssetAliasResponse, GetAssetEventsData, GetAssetEventsResponse, CreateAssetEventData, CreateAssetEventResponse, MaterializeAssetData, MaterializeAssetResponse, GetAssetQueuedEventsData, GetAssetQueuedEventsResponse, DeleteAssetQueuedEventsData, DeleteAssetQueuedEventsResponse, GetAssetData, GetAssetResponse, GetDagAssetQueuedEventsData, GetDagAssetQueuedEventsResponse, DeleteDagAssetQueuedEventsData, DeleteDagAssetQueuedEventsResponse, GetDagAssetQueuedEventData, GetDagAssetQueuedEventResponse, DeleteDagAssetQueuedEventData, DeleteDagAssetQueuedEventResponse, NextRunAssetsData, NextRunAssetsResponse, ListBackfillsData, ListBackfillsResponse, CreateBackfillData, CreateBackfillResponse, GetBackfillData, GetBackfillResponse, PauseBackfillData, PauseBackfillResponse, UnpauseBackfillData, UnpauseBackfillResponse, CancelBackfillData, CancelBackfillResponse, CreateBackfillDryRunData, CreateBackfillDryRunResponse, ListBackfillsUiData, ListBackfillsUiResponse, DeleteConnectionData, DeleteConnectionResponse, GetConnectionData, GetConnectionResponse, PatchConnectionData, PatchConnectionResponse, GetConnectionsData, GetConnectionsResponse, PostConnectionData, PostConnectionResponse, BulkConnectionsData, BulkConnectionsResponse, TestConnectionData, TestConnectionResponse, TestConnectionAsyncData, TestConnectionAsyncResponse, GetConnectionTestData, GetConnectionTestResponse, CreateDefaultConnectionsResponse, HookMetaDataResponse, GetDagRunData, GetDagRunResponse, DeleteDagRunData, DeleteDagRunResponse, PatchDagRunData, PatchDagRunResponse, GetUpstreamAssetEventsData, GetUpstreamAssetEventsResponse, ClearDagRunData, ClearDagRunResponse, GetDagRunsData, GetDagRunsResponse, TriggerDagRunData, TriggerDagRunResponse, WaitDagRunUntilFinishedData, WaitDagRunUntilFinishedResponse, GetListDagRunsBatchData, GetListDagRunsBatchResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, GetDagStatsResponse, GetConfigData, GetConfigResponse, GetConfigValueData, GetConfigValueResponse, GetConfigsResponse, ListDagWarningsData, ListDagWarningsResponse, GetDagsData, GetDagsResponse, PatchDagsData, PatchDagsResponse, GetDagData, GetDagResponse, PatchDagData, PatchDagResponse, DeleteDagData, DeleteDagResponse, GetDagDetailsData, GetDagDetailsResponse, FavoriteDagData, FavoriteDagResponse, UnfavoriteDagData, UnfavoriteDagResponse, GetDagTagsData, GetDagTagsResponse, GetDagsUiData, GetDagsUiResponse, GetLatestRunInfoData, GetLatestRunInfoResponse, GetEventLogData, GetEventLogResponse, GetEventLogsData, GetEventLogsResponse, GetExtraLinksData, GetExtraLinksResponse, GetTaskInstanceData, GetTaskInstanceResponse, PatchTaskInstanceData, PatchTaskInstanceResponse, DeleteTaskInstanceData, DeleteTaskInstanceResponse, GetMappedTaskInstancesData, GetMappedTaskInstancesResponse, GetTaskInstanceDependenciesByMapIndexData, GetTaskInstanceDependenciesByMapIndexResponse, GetTaskInstanceDependenciesData, GetTaskInstanceDependenciesResponse, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, GetMappedTaskInstanceTriesData, GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, PatchTaskInstanceByMapIndexData, PatchTaskInstanceByMapIndexResponse, GetTaskInstancesData, GetTaskInstancesResponse, BulkTaskInstancesData, BulkTaskInstancesResponse, GetTaskInstancesBatchData, GetTaskInstancesBatchResponse, GetTaskInstanceTryDetailsData, GetTaskInstanceTryDetailsResponse, GetMappedTaskInstanceTryDetailsData, GetMappedTaskInstanceTryDetailsResponse, PostClearTaskInstancesData, PostClearTaskInstancesResponse, PatchTaskInstanceDryRunByMapIndexData, PatchTaskInstanceDryRunByMapIndexResponse, PatchTaskInstanceDryRunData, PatchTaskInstanceDryRunResponse, GetLogData, GetLogResponse, GetExternalLogUrlData, GetExternalLogUrlResponse, UpdateHitlDetailData, UpdateHitlDetailResponse, GetHitlDetailData, GetHitlDetailResponse, GetHitlDetailTryDetailData, GetHitlDetailTryDetailResponse, GetHitlDetailsData, GetHitlDetailsResponse, GetImportErrorData, GetImportErrorResponse, GetImportErrorsData, GetImportErrorsResponse, GetJobsData, GetJobsResponse, GetPluginsData, GetPluginsResponse, ImportErrorsResponse, DeletePoolData, DeletePoolResponse, GetPoolData, GetPoolResponse, PatchPoolData, PatchPoolResponse, GetPoolsData, GetPoolsResponse, PostPoolData, PostPoolResponse, BulkPoolsData, BulkPoolsResponse, GetProvidersData, GetProvidersResponse, GetXcomEntryData, GetXcomEntryResponse, UpdateXcomEntryData, UpdateXcomEntryResponse, DeleteXcomEntryData, DeleteXcomEntryResponse, GetXcomEntriesData, GetXcomEntriesResponse, CreateXcomEntryData, CreateXcomEntryResponse, GetTasksData, GetTasksResponse, GetTaskData, GetTaskResponse, DeleteVariableData, DeleteVariableResponse, GetVariableData, GetVariableResponse, PatchVariableData, PatchVariableResponse, GetVariablesData, GetVariablesResponse, PostVariableData, PostVariableResponse, BulkVariablesData, BulkVariablesResponse, ReparseDagFileData, ReparseDagFileResponse, GetDagVersionData, GetDagVersionResponse, GetDagVersionsData, GetDagVersionsResponse, GetHealthResponse, GetVersionResponse, LoginData, LoginResponse, LogoutResponse, GetAuthMenusResponse, GetCurrentUserInfoResponse, GenerateTokenData, GenerateTokenResponse2, GetPartitionedDagRunsData, GetPartitionedDagRunsResponse, GetPendingPartitionedDagRunData, GetPendingPartitionedDagRunResponse, GetDependenciesData, GetDependenciesResponse, HistoricalMetricsData, HistoricalMetricsResponse, DagStatsResponse2, GetDagRunDeadlinesData, GetDagRunDeadlinesResponse, StructureDataData, StructureDataResponse2, GetDagStructureData, GetDagStructureResponse, GetGridRunsData, GetGridRunsResponse, GetGridTiSummariesStreamData, GetGridTiSummariesStreamResponse, GetGanttDataData, GetGanttDataResponse, GetCalendarData, GetCalendarResponse, ListTeamsData, ListTeamsResponse } from './types.gen'; export class AssetService { /** @@ -632,6 +632,7 @@ export class ConnectionService { 401: 'Unauthorized', 403: 'Forbidden', 404: 'Not Found', + 409: 'Conflict', 422: 'Validation Error' } }); @@ -768,46 +769,6 @@ export class ConnectionService { }); } - /** - * Patch Connection And Test - * Update a connection and queue an async test with revert-on-failure. - * - * Atomically saves the edit and creates a ConnectionTest with snapshots of the - * pre-edit and post-edit state. If the test fails, the connection is automatically - * reverted to its pre-edit values. - * @param data The data for the request. - * @param data.connectionId - * @param data.requestBody - * @param data.updateMask - * @param data.executor Executor to route the connection test to - * @param data.queue Queue to route the connection test to - * @returns ConnectionSaveAndTestResponse Successful Response - * @throws ApiError - */ - public static patchConnectionAndTest(data: PatchConnectionAndTestData): CancelablePromise { - return __request(OpenAPI, { - method: 'PATCH', - url: '/api/v2/connections/{connection_id}/test', - path: { - connection_id: data.connectionId - }, - query: { - update_mask: data.updateMask, - executor: data.executor, - queue: data.queue - }, - body: data.requestBody, - mediaType: 'application/json', - errors: { - 400: 'Bad Request', - 401: 'Unauthorized', - 403: 'Forbidden', - 404: 'Not Found', - 422: 'Validation Error' - } - }); - } - /** * Test Connection * Test an API connection. @@ -838,8 +799,9 @@ export class ConnectionService { * Test Connection Async * Queue an async connection test to be executed on a worker. * - * The connection must already be saved. Returns a token that can be used - * to poll for the test result via GET /connections/test-async/{token}. + * 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}. * @param data The data for the request. * @param data.requestBody * @returns ConnectionTestQueuedResponse Successful Response @@ -854,14 +816,14 @@ export class ConnectionService { errors: { 401: 'Unauthorized', 403: 'Forbidden', - 404: 'Not Found', + 409: 'Conflict', 422: 'Validation Error' } }); } /** - * Get Connection Test Status + * Get Connection Test * Poll for the status of an async connection test. * * Knowledge of the token serves as authorization — only the client @@ -871,7 +833,7 @@ export class ConnectionService { * @returns ConnectionTestStatusResponse Successful Response * @throws ApiError */ - public static getConnectionTestStatus(data: GetConnectionTestStatusData): CancelablePromise { + public static getConnectionTest(data: GetConnectionTestData): CancelablePromise { return __request(OpenAPI, { method: 'GET', url: '/api/v2/connections/test-async/{connection_test_token}', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 4a032f12422d4..c56c0b59cb55e 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -495,15 +495,6 @@ export type ConnectionResponse = { team_name: string | null; }; -/** - * Response returned by the combined save-and-test endpoint. - */ -export type ConnectionSaveAndTestResponse = { - connection: ConnectionResponse; - test_token: string; - test_state: string; -}; - /** * Response returned when an async connection test is queued. */ @@ -518,6 +509,14 @@ export type ConnectionTestQueuedResponse = { */ export type ConnectionTestRequestBody = { connection_id: string; + conn_type: string; + host?: string | null; + login?: string | null; + schema?: string | null; + port?: number | null; + password?: string | null; + extra?: string | null; + commit_on_success?: boolean; executor?: string | null; queue?: string | null; }; @@ -539,7 +538,6 @@ export type ConnectionTestStatusResponse = { state: string; result_message?: string | null; created_at: string; - reverted?: boolean; }; /** @@ -2515,22 +2513,6 @@ export type BulkConnectionsData = { export type BulkConnectionsResponse = BulkResponse; -export type PatchConnectionAndTestData = { - connectionId: string; - /** - * Executor to route the connection test to - */ - executor?: string | null; - /** - * Queue to route the connection test to - */ - queue?: string | null; - requestBody: ConnectionBody; - updateMask?: Array<(string)> | null; -}; - -export type PatchConnectionAndTestResponse = ConnectionSaveAndTestResponse; - export type TestConnectionData = { requestBody: ConnectionBody; }; @@ -2543,11 +2525,11 @@ export type TestConnectionAsyncData = { export type TestConnectionAsyncResponse = ConnectionTestQueuedResponse; -export type GetConnectionTestStatusData = { +export type GetConnectionTestData = { connectionTestToken: string; }; -export type GetConnectionTestStatusResponse = ConnectionTestStatusResponse; +export type GetConnectionTestResponse = ConnectionTestStatusResponse; export type CreateDefaultConnectionsResponse = void; @@ -4375,6 +4357,10 @@ export type $OpenApiTs = { * Not Found */ 404: HTTPExceptionResponse; + /** + * Conflict + */ + 409: HTTPExceptionResponse; /** * Validation Error */ @@ -4509,37 +4495,6 @@ export type $OpenApiTs = { }; }; }; - '/api/v2/connections/{connection_id}/test': { - patch: { - req: PatchConnectionAndTestData; - res: { - /** - * Successful Response - */ - 200: ConnectionSaveAndTestResponse; - /** - * Bad Request - */ - 400: HTTPExceptionResponse; - /** - * Unauthorized - */ - 401: HTTPExceptionResponse; - /** - * Forbidden - */ - 403: HTTPExceptionResponse; - /** - * Not Found - */ - 404: HTTPExceptionResponse; - /** - * Validation Error - */ - 422: HTTPValidationError; - }; - }; - }; '/api/v2/connections/test': { post: { req: TestConnectionData; @@ -4580,9 +4535,9 @@ export type $OpenApiTs = { */ 403: HTTPExceptionResponse; /** - * Not Found + * Conflict */ - 404: HTTPExceptionResponse; + 409: HTTPExceptionResponse; /** * Validation Error */ @@ -4592,7 +4547,7 @@ export type $OpenApiTs = { }; '/api/v2/connections/test-async/{connection_test_token}': { get: { - req: GetConnectionTestStatusData; + req: GetConnectionTestData; res: { /** * Successful Response diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py index 59333cb6b822b..46fa983e7c1bd 100644 --- a/airflow-core/src/airflow/utils/db_cleanup.py +++ b/airflow-core/src/airflow/utils/db_cleanup.py @@ -172,7 +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", recency_column_name="created_at"), + _TableConfig(table_name="connection_test_request", recency_column_name="created_at"), ] # We need to have `fallback="database"` because this is executed at top level code and provider configuration diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index a14b70fd4e3ea..fead99e94ca0b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -29,7 +29,7 @@ from airflow.api_fastapi.core_api.datamodels.connections import ConnectionBody from airflow.api_fastapi.core_api.services.public.connections import BulkConnectionService from airflow.models import Connection -from airflow.models.connection_test import ConnectionTest, ConnectionTestState, attempt_revert +from airflow.models.connection_test import ConnectionTestRequest, ConnectionTestState from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.utils.session import NEW_SESSION, provide_session @@ -1179,11 +1179,16 @@ def test_should_test_new_connection_without_existing(self, test_client): class TestAsyncConnectionTest(TestConnectionEndpoint): """Tests for the async connection test endpoints (POST + GET polling).""" + TEST_REQUEST_BODY = { + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "host": TEST_CONN_HOST, + } + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_post_should_respond_202(self, test_client, session): - """POST /connections/test-async with a saved connection returns 202 + token.""" - self.create_connection() - response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) + """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 @@ -1192,66 +1197,70 @@ def test_post_should_respond_202(self, test_client, session): assert len(body["token"]) > 0 def test_should_respond_401(self, unauthenticated_test_client): - response = unauthenticated_test_client.post( - "/connections/test-async", json={"connection_id": TEST_CONN_ID} - ) + 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={"connection_id": TEST_CONN_ID} - ) + 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={"connection_id": TEST_CONN_ID}) + response = test_client.post("/connections/test-async", json=self.TEST_REQUEST_BODY) assert response.status_code == 403 - assert response.json() == { - "detail": "Testing connections is disabled in Airflow configuration. " - "Contact your deployment admin to enable it." - } @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_should_respond_404_for_nonexistent_connection(self, test_client): - """Connection must be saved before testing.""" - response = test_client.post("/connections/test-async", json={"connection_id": "nonexistent"}) - assert response.status_code == 404 - assert "was not found" in response.json()["detail"] - - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_post_creates_connection_test_row(self, test_client, session): - """POST creates a ConnectionTest row in PENDING state.""" - self.create_connection() - response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) + 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(ConnectionTest).filter_by(token=token)) + ct = session.scalar(select(ConnectionTestRequest).filter_by(token=token)) assert ct is not None assert ct.connection_id == TEST_CONN_ID + assert ct.conn_type == TEST_CONN_TYPE + assert ct.host == TEST_CONN_HOST assert ct.state == "pending" @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_post_passes_queue_parameter(self, test_client, session): - """POST /connections/test-async passes the queue parameter to the ConnectionTest.""" - self.create_connection() - response = test_client.post( - "/connections/test-async", - json={"connection_id": TEST_CONN_ID, "queue": "gpu_workers"}, - ) + """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(ConnectionTest).filter_by(token=token)) + ct = session.scalar(select(ConnectionTestRequest).filter_by(token=token)) assert ct is not None assert ct.queue == "gpu_workers" + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_post_stores_commit_on_success(self, test_client, session): + """POST /connections/test-async stores the commit_on_success flag.""" + body = {**self.TEST_REQUEST_BODY, "commit_on_success": True} + response = test_client.post("/connections/test-async", json=body) + assert response.status_code == 202 + token = response.json()["token"] + + ct = session.scalar(select(ConnectionTestRequest).filter_by(token=token)) + assert ct is not None + assert ct.commit_on_success is True + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_post_returns_409_for_duplicate_active_test(self, test_client, session): + """POST returns 409 when there's already an active test for the same connection_id.""" + response = test_client.post("/connections/test-async", json=self.TEST_REQUEST_BODY) + assert response.status_code == 202 + + response = test_client.post("/connections/test-async", json=self.TEST_REQUEST_BODY) + assert response.status_code == 409 + assert "async test is running" in response.json()["detail"] + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_get_status_returns_pending(self, test_client, session): - """GET /connections/test-async/{token} returns current status (pending before scheduler dispatch).""" - self.create_connection() - post_response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) + """GET /connections/test-async/{token} returns current status.""" + post_response = test_client.post("/connections/test-async", json=self.TEST_REQUEST_BODY) token = post_response.json()["token"] response = test_client.get(f"/connections/test-async/{token}") @@ -1262,15 +1271,15 @@ def test_get_status_returns_pending(self, test_client, session): assert body["state"] == "pending" assert body["result_message"] is None assert "created_at" in body + assert "reverted" not in body @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_get_status_returns_completed_result(self, test_client, session): - """GET returns result after the worker has updated the ConnectionTest.""" - self.create_connection() - post_response = test_client.post("/connections/test-async", json={"connection_id": TEST_CONN_ID}) + """GET returns result after the worker has updated the test.""" + post_response = test_client.post("/connections/test-async", json=self.TEST_REQUEST_BODY) token = post_response.json()["token"] - ct = session.scalar(select(ConnectionTest).filter_by(token=token)) + ct = session.scalar(select(ConnectionTestRequest).filter_by(token=token)) ct.state = ConnectionTestState.SUCCESS ct.result_message = "Connection successfully tested" session.commit() @@ -1287,162 +1296,49 @@ def test_get_status_returns_404_for_invalid_token(self, test_client): assert response.status_code == 404 -class TestSaveAndTest(TestConnectionEndpoint): - """Tests for the combined PATCH /{connection_id}/test endpoint.""" - - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_save_and_test_returns_200_with_token(self, test_client, session): - """PATCH /{connection_id}/test updates the connection and returns a test token.""" - self.create_connection() - response = test_client.patch( - f"/connections/{TEST_CONN_ID}/test", - json={ - "connection_id": TEST_CONN_ID, - "conn_type": TEST_CONN_TYPE, - "host": "updated-host.example.com", - }, - ) - assert response.status_code == 200 - body = response.json() - assert body["test_token"] - assert body["test_state"] == "pending" - assert body["connection"]["host"] == "updated-host.example.com" +class TestBlockEditDeleteDuringActiveTest(TestConnectionEndpoint): + """Tests that edit/delete is blocked while an async test is running.""" @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_save_and_test_creates_snapshot(self, test_client, session): - """PATCH /{connection_id}/test creates a ConnectionTest with a connection_snapshot.""" + def test_patch_blocked_during_active_test(self, test_client, session): + """PATCH /{connection_id} returns 409 when an active test exists.""" self.create_connection() - response = test_client.patch( - f"/connections/{TEST_CONN_ID}/test", - json={ - "connection_id": TEST_CONN_ID, - "conn_type": TEST_CONN_TYPE, - "host": "new-host.example.com", - }, - ) - assert response.status_code == 200 - token = response.json()["test_token"] - - ct = session.scalar(select(ConnectionTest).filter_by(token=token)) - assert ct is not None - assert ct.connection_snapshot is not None - snapshot = ct.connection_snapshot - assert "pre" in snapshot - assert "post" in snapshot - assert snapshot["pre"]["host"] == TEST_CONN_HOST - assert snapshot["post"]["host"] == "new-host.example.com" - - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_save_and_test_passes_executor_parameter(self, test_client, session): - """PATCH /{connection_id}/test passes the executor parameter to the ConnectionTest.""" - self.create_connection() - response = test_client.patch( - f"/connections/{TEST_CONN_ID}/test?executor=my_executor", - json={ - "connection_id": TEST_CONN_ID, - "conn_type": TEST_CONN_TYPE, - "host": "queued-host.example.com", - }, - ) - assert response.status_code == 200 - token = response.json()["test_token"] - - ct = session.scalar(select(ConnectionTest).filter_by(token=token)) - assert ct is not None - assert ct.executor == "my_executor" - - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_save_and_test_passes_queue_parameter(self, test_client, session): - """PATCH /{connection_id}/test passes the queue parameter to the ConnectionTest.""" - self.create_connection() - response = test_client.patch( - f"/connections/{TEST_CONN_ID}/test?queue=gpu_workers", - json={ - "connection_id": TEST_CONN_ID, - "conn_type": TEST_CONN_TYPE, - "host": "queued-host.example.com", - }, - ) - assert response.status_code == 200 - token = response.json()["test_token"] - - ct = session.scalar(select(ConnectionTest).filter_by(token=token)) - assert ct is not None - assert ct.queue == "gpu_workers" - - def test_save_and_test_403_when_disabled(self, test_client): - """PATCH /{connection_id}/test returns 403 when test_connection is disabled.""" - self.create_connection() - response = test_client.patch( - f"/connections/{TEST_CONN_ID}/test", + test_client.post( + "/connections/test-async", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, + "host": TEST_CONN_HOST, }, ) - assert response.status_code == 403 - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_save_and_test_404_for_nonexistent(self, test_client): - """PATCH /{connection_id}/test returns 404 for nonexistent connection.""" response = test_client.patch( - "/connections/nonexistent/test", - json={ - "connection_id": "nonexistent", - "conn_type": "http", - }, - ) - assert response.status_code == 404 - - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_poll_shows_reverted_true_after_failed_test(self, test_client, session): - """GET status shows reverted=True after a failed test triggers a revert.""" - self.create_connection() - response = test_client.patch( - f"/connections/{TEST_CONN_ID}/test", + f"/connections/{TEST_CONN_ID}", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, - "host": "bad-host.example.com", + "host": "updated-host.example.com", }, ) - token = response.json()["test_token"] - - ct = session.scalar(select(ConnectionTest).filter_by(token=token)) - ct.state = ConnectionTestState.FAILED - ct.result_message = "Connection refused" - - attempt_revert(ct, session=session) - session.commit() - - poll_response = test_client.get(f"/connections/test-async/{token}") - assert poll_response.status_code == 200 - body = poll_response.json() - assert body["reverted"] is True + assert response.status_code == 409 + assert "async test is running" in response.json()["detail"] @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) - def test_poll_shows_reverted_false_for_success(self, test_client, session): - """GET status shows reverted=False for a successful test.""" + def test_delete_blocked_during_active_test(self, test_client, session): + """DELETE /{connection_id} returns 409 when an active test exists.""" self.create_connection() - response = test_client.patch( - f"/connections/{TEST_CONN_ID}/test", + test_client.post( + "/connections/test-async", json={ "connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, - "host": "good-host.example.com", + "host": TEST_CONN_HOST, }, ) - token = response.json()["test_token"] - - ct = session.scalar(select(ConnectionTest).filter_by(token=token)) - ct.state = ConnectionTestState.SUCCESS - ct.result_message = "Connection OK" - session.commit() - poll_response = test_client.get(f"/connections/test-async/{token}") - assert poll_response.status_code == 200 - body = poll_response.json() - assert body["reverted"] is False + response = test_client.delete(f"/connections/{TEST_CONN_ID}") + assert response.status_code == 409 + assert "async test is running" in response.json()["detail"] class TestCreateDefaultConnections(TestConnectionEndpoint): diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py index 23db82cd76786..ac31cfd487159 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_connection_tests.py @@ -20,7 +20,7 @@ from sqlalchemy import select from airflow.models.connection import Connection -from airflow.models.connection_test import ConnectionTest, ConnectionTestState, snapshot_connection +from airflow.models.connection_test import ConnectionTestRequest, ConnectionTestState from tests_common.test_utils.db import clear_db_connection_tests, clear_db_connections @@ -36,7 +36,7 @@ def setup_teardown(self): def test_patch_updates_result(self, client, session): """PATCH sets the state and result fields.""" - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres") ct.state = ConnectionTestState.RUNNING session.add(ct) session.commit() @@ -51,10 +51,9 @@ def test_patch_updates_result(self, client, session): assert response.status_code == 204 session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == "success" assert ct.result_message == "Connection successfully tested" - assert ct.connection_snapshot is None def test_patch_returns_404_for_nonexistent(self, client): """PATCH with unknown id returns 404.""" @@ -74,7 +73,7 @@ def test_patch_returns_422_for_invalid_uuid(self, client): def test_patch_returns_409_for_terminal_state(self, client, session): """PATCH on a test already in terminal state returns 409.""" - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres") ct.state = ConnectionTestState.SUCCESS ct.result_message = "Already done" session.add(ct) @@ -88,8 +87,8 @@ def test_patch_returns_409_for_terminal_state(self, client, session): assert "terminal state" in response.json()["detail"]["message"] -class TestPatchConnectionTestRevert: - """Tests for the revert-on-failure behavior in the execution API.""" +class TestPatchConnectionTestCommitOnSuccess: + """Tests for the commit_on_success behavior in the execution API.""" @pytest.fixture(autouse=True) def setup_teardown(self): @@ -99,61 +98,45 @@ def setup_teardown(self): clear_db_connections(add_default_connections_back=False) clear_db_connection_tests() - def test_patch_failed_with_snapshot_reverts_connection(self, client, session): - """PATCH with state=failed and snapshot triggers revert.""" - conn = Connection( - conn_id="revert_conn", + def test_success_with_commit_creates_connection(self, client, session): + """PATCH with state=success and commit_on_success creates a new connection.""" + ct = ConnectionTestRequest( + connection_id="new_conn", conn_type="postgres", - host="old-host.example.com", - login="old_user", + host="db.example.com", + login="user", + password="secret", + commit_on_success=True, ) - session.add(conn) - session.flush() - - pre_snap = snapshot_connection(conn) - conn.host = "new-host.example.com" - conn.login = "new_user" - post_snap = snapshot_connection(conn) - session.flush() - - ct = ConnectionTest(connection_id="revert_conn") ct.state = ConnectionTestState.RUNNING - ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} session.add(ct) session.commit() response = client.patch( f"/execution/connection-tests/{ct.id}", - json={"state": "failed", "result_message": "Connection refused"}, + json={"state": "success", "result_message": "Connection OK"}, ) assert response.status_code == 204 - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.reverted is True - assert ct.connection_snapshot is None - conn = session.scalar(select(Connection).filter_by(conn_id="revert_conn")) - assert conn.host == "old-host.example.com" - assert conn.login == "old_user" - - def test_patch_success_with_snapshot_no_revert(self, client, session): - """PATCH with state=success does not trigger revert even with snapshot.""" - conn = Connection( - conn_id="no_revert_conn", - conn_type="postgres", - host="old-host.example.com", - ) - session.add(conn) - session.flush() + conn = session.scalar(select(Connection).filter_by(conn_id="new_conn")) + assert conn is not None + assert conn.conn_type == "postgres" + assert conn.host == "db.example.com" - pre_snap = snapshot_connection(conn) - conn.host = "new-host.example.com" - post_snap = snapshot_connection(conn) + def test_success_with_commit_updates_existing(self, client, session): + """PATCH with state=success and commit_on_success updates an existing connection.""" + conn = Connection(conn_id="existing_conn", conn_type="http", host="old-host.example.com") + session.add(conn) session.flush() - ct = ConnectionTest(connection_id="no_revert_conn") + ct = ConnectionTestRequest( + connection_id="existing_conn", + conn_type="postgres", + host="new-host.example.com", + login="new_user", + commit_on_success=True, + ) ct.state = ConnectionTestState.RUNNING - ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} session.add(ct) session.commit() @@ -164,49 +147,40 @@ def test_patch_success_with_snapshot_no_revert(self, client, session): assert response.status_code == 204 session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.reverted is False - assert ct.connection_snapshot is None - conn = session.scalar(select(Connection).filter_by(conn_id="no_revert_conn")) + conn = session.scalar(select(Connection).filter_by(conn_id="existing_conn")) + assert conn.conn_type == "postgres" assert conn.host == "new-host.example.com" - def test_patch_failed_without_snapshot_no_revert(self, client, session): - """PATCH with state=failed but no snapshot does not trigger revert.""" - ct = ConnectionTest(connection_id="test_conn") + def test_success_without_commit_does_not_create(self, client, session): + """PATCH with state=success but commit_on_success=False does not create a connection.""" + ct = ConnectionTestRequest( + connection_id="no_commit_conn", + conn_type="postgres", + host="db.example.com", + commit_on_success=False, + ) ct.state = ConnectionTestState.RUNNING session.add(ct) session.commit() response = client.patch( f"/execution/connection-tests/{ct.id}", - json={"state": "failed", "result_message": "Connection refused"}, + json={"state": "success", "result_message": "Connection OK"}, ) assert response.status_code == 204 - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.reverted is False + conn = session.scalar(select(Connection).filter_by(conn_id="no_commit_conn")) + assert conn is None - def test_patch_failed_concurrent_edit_skips_revert(self, client, session): - """PATCH with state=failed skips revert when connection was modified concurrently.""" - conn = Connection( - conn_id="concurrent_conn", + def test_failed_with_commit_does_not_create(self, client, session): + """PATCH with state=failed and commit_on_success=True does NOT create a connection.""" + ct = ConnectionTestRequest( + connection_id="fail_conn", conn_type="postgres", - host="old-host.example.com", + host="db.example.com", + commit_on_success=True, ) - session.add(conn) - session.flush() - - pre_snap = snapshot_connection(conn) - conn.host = "new-host.example.com" - post_snap = snapshot_connection(conn) - - conn.host = "third-party-host.example.com" - session.flush() - - ct = ConnectionTest(connection_id="concurrent_conn") ct.state = ConnectionTestState.RUNNING - ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} session.add(ct) session.commit() @@ -216,10 +190,48 @@ def test_patch_failed_concurrent_edit_skips_revert(self, client, session): ) assert response.status_code == 204 - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.reverted is False - assert ct.connection_snapshot is None - assert "modified by another user" in ct.result_message - conn = session.scalar(select(Connection).filter_by(conn_id="concurrent_conn")) - assert conn.host == "third-party-host.example.com" + conn = session.scalar(select(Connection).filter_by(conn_id="fail_conn")) + assert conn is None + + +class TestGetConnectionTestConnection: + """Tests for the GET /{connection_test_id}/connection endpoint.""" + + @pytest.fixture(autouse=True) + def setup_teardown(self): + clear_db_connection_tests() + yield + clear_db_connection_tests() + + def test_get_connection_returns_data(self, client, session): + """GET returns decrypted connection data from the test request.""" + ct = ConnectionTestRequest( + connection_id="test_conn", + conn_type="postgres", + host="db.example.com", + login="user", + password="secret", + schema="mydb", + port=5432, + extra='{"key": "value"}', + ) + session.add(ct) + session.commit() + + response = client.get(f"/execution/connection-tests/{ct.id}/connection") + assert response.status_code == 200 + + data = response.json() + assert data["conn_id"] == "test_conn" + assert data["conn_type"] == "postgres" + assert data["host"] == "db.example.com" + assert data["login"] == "user" + assert data["password"] == "secret" + assert data["schema"] == "mydb" + assert data["port"] == 5432 + assert data["extra"] == '{"key": "value"}' + + def test_get_connection_returns_404_for_nonexistent(self, client): + """GET with unknown id returns 404.""" + response = client.get("/execution/connection-tests/00000000-0000-0000-0000-000000000000/connection") + assert response.status_code == 404 diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 05c1b33f220fb..10e0e2c0db6ea 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -67,8 +67,7 @@ ) from airflow.models.backfill import Backfill, _create_backfill from airflow.models.callback import ExecutorCallback -from airflow.models.connection import Connection -from airflow.models.connection_test import ConnectionTest, ConnectionTestState, snapshot_connection +from airflow.models.connection_test import ConnectionTestRequest, ConnectionTestState from airflow.models.dag import DagModel, get_last_dagrun, infer_automated_data_interval from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel @@ -115,7 +114,6 @@ clear_db_assets, clear_db_backfills, clear_db_callbacks, - clear_db_connections, clear_db_dag_bundles, clear_db_dags, clear_db_deadline, @@ -9453,7 +9451,7 @@ def test_fallback_values_used_only_when_dag_version_is_none(self): @pytest.fixture def scheduler_job_runner_for_connection_tests(session): """Create a SchedulerJobRunner with a mock Job and supporting executor.""" - session.execute(delete(ConnectionTest)) + session.execute(delete(ConnectionTestRequest)) session.commit() mock_job = mock.MagicMock(spec=Job) @@ -9467,7 +9465,7 @@ def scheduler_job_runner_for_connection_tests(session): runner.executor = executor runner._log = mock.MagicMock(spec=logging.Logger) yield runner - session.execute(delete(ConnectionTest)) + session.execute(delete(ConnectionTestRequest)) session.commit() @@ -9481,7 +9479,7 @@ class TestDispatchConnectionTests: ) def test_dispatch_pending_tests(self, scheduler_job_runner_for_connection_tests, session): """Pending connection tests are dispatched to a supporting executor.""" - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="test_conn") session.add(ct) session.commit() assert ct.state == ConnectionTestState.PENDING @@ -9489,7 +9487,7 @@ def test_dispatch_pending_tests(self, scheduler_job_runner_for_connection_tests, scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == ConnectionTestState.QUEUED assert len(scheduler_job_runner_for_connection_tests.executor.queued_connection_tests) == 1 @@ -9502,18 +9500,18 @@ def test_dispatch_pending_tests(self, scheduler_job_runner_for_connection_tests, ) def test_dispatch_respects_concurrency_limit(self, scheduler_job_runner_for_connection_tests, session): """Excess pending tests stay PENDING when concurrency is at capacity.""" - ct_active = ConnectionTest(connection_id="active_conn") + ct_active = ConnectionTestRequest(conn_type="test_type", connection_id="active_conn") ct_active.state = ConnectionTestState.QUEUED session.add(ct_active) - ct_pending = ConnectionTest(connection_id="pending_conn") + ct_pending = ConnectionTestRequest(conn_type="test_type", connection_id="pending_conn") session.add(ct_pending) session.commit() scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() - ct_pending = session.get(ConnectionTest, ct_pending.id) + ct_pending = session.get(ConnectionTestRequest, ct_pending.id) assert ct_pending.state == ConnectionTestState.PENDING @mock.patch.dict( @@ -9532,16 +9530,15 @@ def test_dispatch_fails_fast_when_no_executor_supports( scheduler_job_runner_for_connection_tests.executors = [unsupporting_executor] scheduler_job_runner_for_connection_tests.executor = unsupporting_executor - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="test_conn") session.add(ct) session.commit() scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == ConnectionTestState.FAILED - assert ct.connection_snapshot is None assert "No executor supports connection testing" in ct.result_message @mock.patch.dict( @@ -9555,14 +9552,14 @@ def test_dispatch_with_unmatched_executor_fails_fast( self, scheduler_job_runner_for_connection_tests, session ): """Tests requesting an executor with no match are failed immediately.""" - ct = ConnectionTest(connection_id="test_conn", executor="gpu_workers") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="test_conn", executor="gpu_workers") session.add(ct) session.commit() scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == ConnectionTestState.FAILED assert "gpu_workers" in ct.result_message @@ -9577,13 +9574,13 @@ def test_dispatch_budget_dispatches_up_to_remaining_slots( self, scheduler_job_runner_for_connection_tests, session ): """When 1 slot is occupied, only budget (cap - active) pending tests are dispatched.""" - ct_active = ConnectionTest(connection_id="active_conn") + ct_active = ConnectionTestRequest(conn_type="test_type", connection_id="active_conn") ct_active.state = ConnectionTestState.RUNNING session.add(ct_active) pending_tests = [] for i in range(3): - ct = ConnectionTest(connection_id=f"pending_{i}") + ct = ConnectionTestRequest(conn_type="test_type", connection_id=f"pending_{i}") session.add(ct) pending_tests.append(ct) session.commit() @@ -9592,7 +9589,7 @@ def test_dispatch_budget_dispatches_up_to_remaining_slots( scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() - states = [session.get(ConnectionTest, pid).state for pid in pending_ids] + states = [session.get(ConnectionTestRequest, pid).state for pid in pending_ids] assert states.count(ConnectionTestState.QUEUED) == 2 assert states.count(ConnectionTestState.PENDING) == 1 @@ -9608,17 +9605,17 @@ def test_dispatch_order_is_fifo_by_created_at(self, scheduler_job_runner_for_con initial_time = timezone.utcnow() with time_machine.travel(initial_time - timedelta(minutes=5), tick=False): - ct_old = ConnectionTest(connection_id="old_conn") + ct_old = ConnectionTestRequest(conn_type="test_type", connection_id="old_conn") session.add(ct_old) session.flush() with time_machine.travel(initial_time, tick=False): - ct_new = ConnectionTest(connection_id="new_conn") + ct_new = ConnectionTestRequest(conn_type="test_type", connection_id="new_conn") session.add(ct_new) session.flush() with time_machine.travel(initial_time + timedelta(minutes=1), tick=False): - ct_newest = ConnectionTest(connection_id="newest_conn") + ct_newest = ConnectionTestRequest(conn_type="test_type", connection_id="newest_conn") session.add(ct_newest) session.flush() @@ -9627,9 +9624,9 @@ def test_dispatch_order_is_fifo_by_created_at(self, scheduler_job_runner_for_con scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() - assert session.get(ConnectionTest, ct_old.id).state == ConnectionTestState.QUEUED - assert session.get(ConnectionTest, ct_new.id).state == ConnectionTestState.QUEUED - assert session.get(ConnectionTest, ct_newest.id).state == ConnectionTestState.PENDING + assert session.get(ConnectionTestRequest, ct_old.id).state == ConnectionTestState.QUEUED + assert session.get(ConnectionTestRequest, ct_new.id).state == ConnectionTestState.QUEUED + assert session.get(ConnectionTestRequest, ct_newest.id).state == ConnectionTestState.PENDING @mock.patch.dict( os.environ, @@ -9647,14 +9644,16 @@ def test_dispatch_fails_fast_for_unserved_executor( "_try_to_load_executor", return_value=None, ): - ct = ConnectionTest(connection_id="test_conn", executor="nonexistent_executor") + ct = ConnectionTestRequest( + conn_type="test_type", connection_id="test_conn", executor="nonexistent_executor" + ) session.add(ct) session.commit() scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == ConnectionTestState.FAILED assert "nonexistent_executor" in ct.result_message @@ -9667,7 +9666,7 @@ def test_dispatch_fails_fast_for_unserved_executor( ) def test_dispatch_executor_matched_by_alias(self, session): """When executor is specified, the executor whose name.alias matches is selected.""" - session.execute(delete(ConnectionTest)) + session.execute(delete(ConnectionTestRequest)) session.commit() mock_job = mock.MagicMock(spec=Job) @@ -9688,7 +9687,7 @@ def test_dispatch_executor_matched_by_alias(self, session): runner.executor = executor_a runner._log = mock.MagicMock(spec=logging.Logger) - ct = ConnectionTest(connection_id="team_conn", executor="executor_b") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="team_conn", executor="executor_b") session.add(ct) session.commit() @@ -9706,7 +9705,7 @@ def test_dispatch_executor_matched_by_alias(self, session): ) def test_dispatch_executor_matched_by_module_path(self, session): """When executor is specified by module_path, the matching executor is selected.""" - session.execute(delete(ConnectionTest)) + session.execute(delete(ConnectionTestRequest)) session.commit() mock_job = mock.MagicMock(spec=Job) @@ -9727,7 +9726,9 @@ def test_dispatch_executor_matched_by_module_path(self, session): runner.executor = executor_a runner._log = mock.MagicMock(spec=logging.Logger) - ct = ConnectionTest(connection_id="team_conn", executor="path.to.ExecutorB") + ct = ConnectionTestRequest( + conn_type="test_type", connection_id="team_conn", executor="path.to.ExecutorB" + ) session.add(ct) session.commit() @@ -9738,7 +9739,7 @@ def test_dispatch_executor_matched_by_module_path(self, session): def test_dispatch_executor_matched_by_class_name(self, session): """When executor is specified by class name only, the matching executor is selected.""" - session.execute(delete(ConnectionTest)) + session.execute(delete(ConnectionTestRequest)) session.commit() mock_job = mock.MagicMock(spec=Job) @@ -9759,7 +9760,7 @@ def test_dispatch_executor_matched_by_class_name(self, session): runner.executor = executor_a runner._log = mock.MagicMock(spec=logging.Logger) - ct = ConnectionTest(connection_id="team_conn", executor="ExecutorB") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="team_conn", executor="ExecutorB") session.add(ct) session.commit() @@ -9782,14 +9783,14 @@ def test_dispatch_respects_parallelism_budget(self, scheduler_job_runner_for_con # Simulate 1 running task so all parallelism slots are occupied executor.running = {"fake_task_key"} - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="test_conn") session.add(ct) session.commit() scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == ConnectionTestState.PENDING @mock.patch.dict( @@ -9806,16 +9807,15 @@ def test_dispatch_fails_when_executor_does_not_support_connection_test( executor = scheduler_job_runner_for_connection_tests.executor executor.supports_connection_test = False - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="test_conn") session.add(ct) session.commit() scheduler_job_runner_for_connection_tests._enqueue_connection_tests(session=session) session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == ConnectionTestState.FAILED - assert ct.connection_snapshot is None assert "No executor supports connection testing" in ct.result_message @@ -9826,7 +9826,7 @@ def test_reap_stale_queued_test(self, scheduler_job_runner_for_connection_tests, initial_time = timezone.utcnow() with time_machine.travel(initial_time, tick=False): - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="test_conn") ct.state = ConnectionTestState.QUEUED session.add(ct) session.commit() @@ -9835,14 +9835,14 @@ def test_reap_stale_queued_test(self, scheduler_job_runner_for_connection_tests, scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == ConnectionTestState.FAILED assert "timed out" in ct.result_message @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) def test_does_not_reap_fresh_tests(self, scheduler_job_runner_for_connection_tests, session): """Fresh QUEUED tests are not reaped.""" - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="test_conn") ct.state = ConnectionTestState.QUEUED session.add(ct) session.commit() @@ -9850,72 +9850,15 @@ def test_does_not_reap_fresh_tests(self, scheduler_job_runner_for_connection_tes scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == ConnectionTestState.QUEUED - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) - def test_reaper_reverts_connection_on_timeout_with_snapshot( - self, scheduler_job_runner_for_connection_tests, session - ): - """Stale tests with a snapshot trigger revert on timeout.""" - clear_db_connections(add_default_connections_back=False) - - conn = Connection( - conn_id="reaper_conn", - conn_type="postgres", - host="old-host.example.com", - ) - session.add(conn) - session.flush() - - pre_snap = snapshot_connection(conn) - conn.host = "new-host.example.com" - post_snap = snapshot_connection(conn) - session.flush() - - initial_time = timezone.utcnow() - with time_machine.travel(initial_time, tick=False): - ct = ConnectionTest(connection_id="reaper_conn") - ct.state = ConnectionTestState.QUEUED - ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} - session.add(ct) - session.commit() - - with time_machine.travel(initial_time + timedelta(seconds=200), tick=False): - scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) - - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.FAILED - assert ct.reverted is True - assert ct.connection_snapshot is None - conn = session.scalar(select(Connection).filter_by(conn_id="reaper_conn")) - assert conn.host == "old-host.example.com" - - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) - def test_reaper_no_revert_without_snapshot(self, scheduler_job_runner_for_connection_tests, session): - """Stale tests without a snapshot do not trigger revert.""" - initial_time = timezone.utcnow() - with time_machine.travel(initial_time, tick=False): - ct = ConnectionTest(connection_id="no_snap_conn") - ct.state = ConnectionTestState.QUEUED - session.add(ct) - session.commit() - - with time_machine.travel(initial_time + timedelta(seconds=200), tick=False): - scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) - - session.expire_all() - ct = session.get(ConnectionTest, ct.id) - assert ct.state == ConnectionTestState.FAILED - assert ct.reverted is False - @mock.patch.dict(os.environ, {"AIRFLOW__CORE__CONNECTION_TEST_TIMEOUT": "60"}) def test_reap_stale_running_test(self, scheduler_job_runner_for_connection_tests, session): """Stale RUNNING tests are also reaped by the reaper.""" initial_time = timezone.utcnow() with time_machine.travel(initial_time, tick=False): - ct = ConnectionTest(connection_id="running_conn") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="running_conn") ct.state = ConnectionTestState.RUNNING session.add(ct) session.commit() @@ -9924,7 +9867,7 @@ def test_reap_stale_running_test(self, scheduler_job_runner_for_connection_tests scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) session.expire_all() - ct = session.get(ConnectionTest, ct.id) + ct = session.get(ConnectionTestRequest, ct.id) assert ct.state == ConnectionTestState.FAILED assert "timed out" in ct.result_message @@ -9933,12 +9876,12 @@ def test_reaper_ignores_terminal_states(self, scheduler_job_runner_for_connectio """Tests in terminal states (SUCCESS, FAILED) are not touched by the reaper.""" initial_time = timezone.utcnow() with time_machine.travel(initial_time, tick=False): - ct_success = ConnectionTest(connection_id="success_conn") + ct_success = ConnectionTestRequest(conn_type="test_type", connection_id="success_conn") ct_success.state = ConnectionTestState.SUCCESS ct_success.result_message = "OK" session.add(ct_success) - ct_failed = ConnectionTest(connection_id="failed_conn") + ct_failed = ConnectionTestRequest(conn_type="test_type", connection_id="failed_conn") ct_failed.state = ConnectionTestState.FAILED ct_failed.result_message = "Error" session.add(ct_failed) @@ -9948,5 +9891,5 @@ def test_reaper_ignores_terminal_states(self, scheduler_job_runner_for_connectio scheduler_job_runner_for_connection_tests._reap_stale_connection_tests(session=session) session.expire_all() - assert session.get(ConnectionTest, ct_success.id).state == ConnectionTestState.SUCCESS - assert session.get(ConnectionTest, ct_failed.id).state == ConnectionTestState.FAILED + assert session.get(ConnectionTestRequest, ct_success.id).state == ConnectionTestState.SUCCESS + assert session.get(ConnectionTestRequest, ct_failed.id).state == ConnectionTestState.FAILED diff --git a/airflow-core/tests/unit/models/test_connection_test.py b/airflow-core/tests/unit/models/test_connection_test.py index c67242f760aad..0d8da8cef695a 100644 --- a/airflow-core/tests/unit/models/test_connection_test.py +++ b/airflow-core/tests/unit/models/test_connection_test.py @@ -22,11 +22,9 @@ from airflow.models.connection import Connection from airflow.models.connection_test import ( - ConnectionTest, + ConnectionTestRequest, ConnectionTestState, - attempt_revert, run_connection_test, - snapshot_connection, ) from tests_common.test_utils.db import clear_db_connection_tests, clear_db_connections @@ -34,84 +32,92 @@ pytestmark = pytest.mark.db_test -class TestConnectionTestModel: +class TestConnectionTestRequestModel: def test_token_is_generated(self): - ct = ConnectionTest(connection_id="test_conn") + 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 = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres") assert ct.state == ConnectionTestState.PENDING def test_tokens_are_unique(self): - ct1 = ConnectionTest(connection_id="test_conn") - ct2 = ConnectionTest(connection_id="test_conn") + 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 = ConnectionTest(connection_id="test_conn") + 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 = ConnectionTest(connection_id="test_conn", executor="my_executor") + 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 = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres") assert ct.executor is None def test_queue_parameter(self): - ct = ConnectionTest(connection_id="test_conn", queue="my_queue") + 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 = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres") assert ct.queue is None - -class TestRunConnectionTest: - def test_successful_connection_test(self): - """Returns (True, message) on successful test.""" - conn = mock.MagicMock(spec=Connection) - conn.conn_id = "test_conn" - conn.test_connection.return_value = (True, "Connection OK") - - success, message = run_connection_test(conn=conn) - - assert success is True - assert message == "Connection OK" - - def test_failed_connection_test(self): - """Returns (False, message) when test_connection returns False.""" - conn = mock.MagicMock(spec=Connection) - conn.conn_id = "test_conn" - conn.test_connection.return_value = (False, "Connection failed") - - success, message = run_connection_test(conn=conn) - - assert success is False - assert message == "Connection failed" - - def test_exception_during_connection_test(self): - """Returns (False, error_str) on exception.""" - conn = mock.MagicMock(spec=Connection) - conn.conn_id = "test_conn" - conn.test_connection.side_effect = Exception("Could not resolve host: db.example.com") - - success, message = run_connection_test(conn=conn) - - assert success is False - assert "Could not resolve host" in message - - -class TestSnapshotConnection: - def test_snapshot_captures_all_fields(self): - """snapshot_connection captures all expected fields including encrypted ones.""" - conn = Connection( - conn_id="snap_test", + def test_connection_fields_stored(self): + ct = ConnectionTestRequest( + connection_id="test_conn", + conn_type="postgres", + host="db.example.com", + login="user", + password="secret", + schema="mydb", + port=5432, + extra='{"key": "value"}', + ) + assert ct.conn_type == "postgres" + assert ct.host == "db.example.com" + assert ct.login == "user" + assert ct.password == "secret" + assert ct.schema == "mydb" + assert ct.port == 5432 + assert ct.extra == '{"key": "value"}' + + def test_password_is_encrypted(self): + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", password="secret") + assert ct._password is not None + assert ct._password != "secret" + assert ct.password == "secret" + + def test_extra_is_encrypted(self): + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", extra='{"key": "val"}') + assert ct._extra is not None + assert ct._extra != '{"key": "val"}' + assert ct.extra == '{"key": "val"}' + + def test_null_password_and_extra(self): + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="http") + assert ct._password is None + assert ct._extra is None + + def test_commit_on_success_default(self): + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres") + assert ct.commit_on_success is False + + def test_commit_on_success_true(self): + ct = ConnectionTestRequest(connection_id="test_conn", conn_type="postgres", commit_on_success=True) + assert ct.commit_on_success is True + + +class TestToConnection: + def test_to_connection_returns_transient_connection(self): + ct = ConnectionTestRequest( + connection_id="test_conn", conn_type="postgres", host="db.example.com", login="user", @@ -120,29 +126,19 @@ def test_snapshot_captures_all_fields(self): port=5432, extra='{"key": "value"}', ) - snap = snapshot_connection(conn) - assert snap["conn_type"] == "postgres" - assert snap["host"] == "db.example.com" - assert snap["login"] == "user" - assert snap["schema"] == "mydb" - assert snap["port"] == 5432 - assert snap["_password"] is not None - assert snap["_password"] != "secret" # Should be encrypted - assert snap["_extra"] is not None - assert snap["is_encrypted"] is True - assert snap["is_extra_encrypted"] is True - - def test_snapshot_with_null_password_and_extra(self): - """snapshot_connection handles None password and extra.""" - conn = Connection(conn_id="snap_test", conn_type="http") - snap = snapshot_connection(conn) - assert snap["_password"] is None - assert snap["_extra"] is None - assert not snap["is_encrypted"] - assert not snap["is_extra_encrypted"] - - -class TestAttemptRevert: + conn = ct.to_connection() + assert isinstance(conn, Connection) + assert conn.conn_id == "test_conn" + assert conn.conn_type == "postgres" + assert conn.host == "db.example.com" + assert conn.login == "user" + assert conn.password == "secret" + assert conn.schema == "mydb" + assert conn.port == 5432 + assert conn.extra == '{"key": "value"}' + + +class TestCommitToConnectionTable: @pytest.fixture(autouse=True) def setup_teardown(self): clear_db_connections(add_default_connections_back=False) @@ -151,140 +147,82 @@ def setup_teardown(self): clear_db_connections(add_default_connections_back=False) clear_db_connection_tests() - def test_attempt_revert_success(self, session): - """attempt_revert restores all connection fields and sets reverted=True.""" - conn = Connection( - conn_id="revert_conn", + def test_creates_new_connection(self, session): + ct = ConnectionTestRequest( + connection_id="new_conn", conn_type="postgres", - host="old-host.example.com", - login="old_user", - password="old_secret", + host="db.example.com", + login="user", + password="secret", schema="mydb", port=5432, - extra='{"key": "old_value"}', ) - session.add(conn) - session.flush() - - pre_snap = snapshot_connection(conn) - - conn.host = "new-host.example.com" - conn.login = "new_user" - conn.password = "new_secret" - conn.port = 9999 - conn.extra = '{"key": "new_value"}' - - post_snap = snapshot_connection(conn) - - ct = ConnectionTest(connection_id="revert_conn") - ct.state = ConnectionTestState.FAILED - ct.result_message = "Connection refused" - ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} session.add(ct) session.flush() - attempt_revert(ct, session=session) + ct.commit_to_connection_table(session=session) session.flush() - assert ct.reverted is True - assert ct.connection_snapshot is None - session.refresh(conn) - assert conn.host == "old-host.example.com" - assert conn.login == "old_user" - assert conn.password == "old_secret" - assert conn.schema == "mydb" - assert conn.port == 5432 - assert conn.extra == '{"key": "old_value"}' + from sqlalchemy import select - def test_attempt_revert_skipped_concurrent_edit(self, session): - """attempt_revert skips revert when connection was modified by another user.""" - conn = Connection( - conn_id="concurrent_conn", - conn_type="postgres", - host="old-host.example.com", - ) - session.add(conn) - session.flush() - - pre_snap = snapshot_connection(conn) - conn.host = "new-host.example.com" - post_snap = snapshot_connection(conn) + conn = session.scalar(select(Connection).filter_by(conn_id="new_conn")) + assert conn is not None + assert conn.conn_type == "postgres" + assert conn.host == "db.example.com" + assert conn.password == "secret" - conn.host = "third-party-host.example.com" + def test_updates_existing_connection(self, session): + conn = Connection(conn_id="existing_conn", conn_type="http", host="old-host.example.com") + session.add(conn) session.flush() - ct = ConnectionTest(connection_id="concurrent_conn") - ct.state = ConnectionTestState.FAILED - ct.result_message = "Connection refused" - ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} + ct = ConnectionTestRequest( + connection_id="existing_conn", + conn_type="postgres", + host="new-host.example.com", + login="new_user", + password="new_secret", + ) session.add(ct) session.flush() - attempt_revert(ct, session=session) - - assert ct.reverted is False - assert ct.connection_snapshot is None - assert "modified by another user" in ct.result_message - assert conn.host == "third-party-host.example.com" - - def test_attempt_revert_skipped_concurrent_password_edit(self, session): - """attempt_revert skips revert when password was changed concurrently.""" - conn = Connection( - conn_id="pw_conn", - conn_type="postgres", - host="host.example.com", - password="original_secret", - ) - session.add(conn) + ct.commit_to_connection_table(session=session) session.flush() + session.refresh(conn) - pre_snap = snapshot_connection(conn) - conn.password = "new_secret" - post_snap = snapshot_connection(conn) + assert conn.conn_type == "postgres" + assert conn.host == "new-host.example.com" + assert conn.login == "new_user" + assert conn.password == "new_secret" - conn.password = "third_party_secret" - session.flush() - ct = ConnectionTest(connection_id="pw_conn") - ct.state = ConnectionTestState.FAILED - ct.result_message = "Connection refused" - ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} - session.add(ct) - session.flush() +class TestRunConnectionTest: + def test_successful_connection_test(self): + conn = mock.MagicMock(spec=Connection) + conn.conn_id = "test_conn" + conn.test_connection.return_value = (True, "Connection OK") - attempt_revert(ct, session=session) + success, message = run_connection_test(conn=conn) - assert ct.reverted is False - assert ct.connection_snapshot is None - assert "modified by another user" in ct.result_message - session.refresh(conn) - assert conn.password == "third_party_secret" + assert success is True + assert message == "Connection OK" - def test_attempt_revert_skipped_connection_deleted(self, session): - """attempt_revert skips revert when connection no longer exists.""" - conn = Connection( - conn_id="deleted_conn", - conn_type="postgres", - host="old-host.example.com", - ) - session.add(conn) - session.flush() + def test_failed_connection_test(self): + conn = mock.MagicMock(spec=Connection) + conn.conn_id = "test_conn" + conn.test_connection.return_value = (False, "Connection failed") - pre_snap = snapshot_connection(conn) - conn.host = "new-host.example.com" - post_snap = snapshot_connection(conn) + success, message = run_connection_test(conn=conn) - ct = ConnectionTest(connection_id="deleted_conn") - ct.state = ConnectionTestState.FAILED - ct.result_message = "Connection refused" - ct.connection_snapshot = {"pre": pre_snap, "post": post_snap} - session.add(ct) + assert success is False + assert message == "Connection failed" - session.delete(conn) - session.flush() + def test_exception_during_connection_test(self): + conn = mock.MagicMock(spec=Connection) + conn.conn_id = "test_conn" + conn.test_connection.side_effect = Exception("Could not resolve host: db.example.com") - attempt_revert(ct, session=session) + success, message = run_connection_test(conn=conn) - assert ct.reverted is False - assert ct.connection_snapshot is None - assert "no longer exists" in ct.result_message + assert success is False + assert "Could not resolve host" in message diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index 405f16843d67e..9758167d72e6d 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -244,16 +244,6 @@ class ConnectionResponse(BaseModel): team_name: Annotated[str | None, Field(title="Team Name")] = None -class ConnectionSaveAndTestResponse(BaseModel): - """ - Response returned by the combined save-and-test endpoint. - """ - - connection: ConnectionResponse - test_token: Annotated[str, Field(title="Test Token")] - test_state: Annotated[str, Field(title="Test State")] - - class ConnectionTestQueuedResponse(BaseModel): """ Response returned when an async connection test is queued. @@ -273,6 +263,14 @@ class ConnectionTestRequestBody(BaseModel): 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 @@ -296,7 +294,6 @@ class ConnectionTestStatusResponse(BaseModel): state: Annotated[str, Field(title="State")] result_message: Annotated[str | None, Field(title="Result Message")] = None created_at: Annotated[datetime, Field(title="Created At")] - reverted: Annotated[bool | None, Field(title="Reverted")] = False class CreateAssetEventsBody(BaseModel): diff --git a/dev/test_async_connection.py b/dev/test_async_connection.py new file mode 100644 index 0000000000000..4d5f20a5d4ec5 --- /dev/null +++ b/dev/test_async_connection.py @@ -0,0 +1,483 @@ +#!/usr/bin/env python +# 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. +""" +Manual end-to-end tests for async connection testing via the REST API. + +Run against a live Breeze server: + uv run --project dev python dev/test_async_connection.py + +Expects: + - Breeze API server at http://localhost:8080 + - AIRFLOW__CORE__TEST_CONNECTION=Enabled + - Default admin:admin credentials +""" + +from __future__ import annotations + +import sys +import time + +import requests + +BASE = "http://localhost:28080/api/v2" +AUTH_URL = "http://localhost:28080/auth/token" +AUTH = None # Set by get_token() +PASS_COUNT = 0 +FAIL_COUNT = 0 + + +def get_token() -> dict: + """Get a JWT token and return headers dict for requests.""" + r = requests.post(AUTH_URL, json={"username": "admin", "password": "admin"}, timeout=5) + r.raise_for_status() + token = r.json()["access_token"] + return {"Authorization": f"Bearer {token}"} + + +def header(title: str) -> None: + print(f"\n{'=' * 60}") + print(f" {title}") + print(f"{'=' * 60}") + + +def check(label: str, condition: bool, detail: str = "") -> None: + global PASS_COUNT, FAIL_COUNT + status = "PASS" if condition else "FAIL" + if condition: + PASS_COUNT += 1 + else: + FAIL_COUNT += 1 + msg = f" [{status}] {label}" + if detail and not condition: + msg += f" -- {detail}" + print(msg) + + +def cleanup_connection(conn_id: str) -> None: + """Delete a connection if it exists, ignoring 404/409.""" + r = requests.delete(f"{BASE}/connections/{conn_id}", headers=AUTH) + # 204=deleted, 404=didn't exist, 409=active test blocking + if r.status_code == 409: + # Wait for active test to finish then retry + time.sleep(3) + requests.delete(f"{BASE}/connections/{conn_id}", headers=AUTH) + + +def poll_until_terminal(token: str, timeout: int = 30) -> dict: + """Poll test status until terminal or timeout.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + r = requests.get(f"{BASE}/connections/test-async/{token}", headers=AUTH) + body = r.json() + if body.get("state") in ("success", "failed"): + return body + time.sleep(1) + return body + + +# ============================================================ +# TEST 1: Basic async test with connection fields +# ============================================================ +def test_basic_async_test(): + header("TEST 1: Basic async test — POST with connection fields") + + r = requests.post( + f"{BASE}/connections/test-async", + headers=AUTH, + json={ + "connection_id": "manual_test_http", + "conn_type": "http", + "host": "https://httpbin.org", + }, + ) + check("POST returns 202", r.status_code == 202, f"got {r.status_code}: {r.text}") + if r.status_code != 202: + return + + body = r.json() + check("Response has token", "token" in body and len(body["token"]) > 0) + check("Response has connection_id", body.get("connection_id") == "manual_test_http") + check("State is pending", body.get("state") == "pending") + + token = body["token"] + print(f" Token: {token[:20]}...") + + # Poll for status + print(" Polling for result...") + result = poll_until_terminal(token) + check( + "Test reached terminal state", + result.get("state") in ("success", "failed"), + f"state={result.get('state')}", + ) + check("Response has no 'reverted' field", "reverted" not in result) + print(f" Final state: {result.get('state')}") + print(f" Message: {result.get('result_message', '(none)')}") + + +# ============================================================ +# TEST 2: 409 duplicate active test +# ============================================================ +def test_duplicate_409(): + header("TEST 2: 409 for duplicate active test on same connection_id") + + body = { + "connection_id": "dup_test_conn", + "conn_type": "http", + "host": "https://httpbin.org", + } + + r1 = requests.post(f"{BASE}/connections/test-async", headers=AUTH, json=body) + check("First POST returns 202", r1.status_code == 202, f"got {r1.status_code}") + + r2 = requests.post(f"{BASE}/connections/test-async", headers=AUTH, json=body) + check("Second POST returns 409", r2.status_code == 409, f"got {r2.status_code}: {r2.text}") + if r2.status_code == 409: + check("409 message mentions async test", "async test" in r2.json().get("detail", "")) + + # Wait for first test to finish so cleanup works + if r1.status_code == 202: + poll_until_terminal(r1.json()["token"]) + + +# ============================================================ +# TEST 3: 403 when test_connection is disabled +# ============================================================ +def test_disabled_403(): + header("TEST 3: 403 when test_connection is disabled (skip if enabled)") + + r = requests.post( + f"{BASE}/connections/test-async", + headers=AUTH, + json={"connection_id": "x", "conn_type": "http"}, + ) + if r.status_code == 403: + check("Returns 403 when disabled", True) + print(" NOTE: test_connection is disabled. Remaining tests will skip.") + return True # signal disabled + if r.status_code == 202: + check("test_connection is enabled (403 test skipped)", True) + poll_until_terminal(r.json()["token"]) + return False + check("Unexpected status", False, f"got {r.status_code}: {r.text}") + return True + + +# ============================================================ +# TEST 4: Block PATCH while test is active +# ============================================================ +def test_block_patch(): + header("TEST 4: PATCH blocked while async test is running") + + # Create a real connection first + requests.post( + f"{BASE}/connections", + headers=AUTH, + json={"connection_id": "block_patch_conn", "conn_type": "http", "host": "example.com"}, + ) + + # Start an async test + r = requests.post( + f"{BASE}/connections/test-async", + headers=AUTH, + json={"connection_id": "block_patch_conn", "conn_type": "http", "host": "example.com"}, + ) + check("Async test started", r.status_code == 202, f"got {r.status_code}") + token = r.json().get("token") if r.status_code == 202 else None + + # Try to PATCH the same connection + r2 = requests.patch( + f"{BASE}/connections/block_patch_conn", + headers=AUTH, + json={"connection_id": "block_patch_conn", "conn_type": "http", "host": "new-host.com"}, + ) + check("PATCH returns 409", r2.status_code == 409, f"got {r2.status_code}: {r2.text}") + + # Wait for test to finish + if token: + poll_until_terminal(token) + + # PATCH should work now + r3 = requests.patch( + f"{BASE}/connections/block_patch_conn", + headers=AUTH, + json={"connection_id": "block_patch_conn", "conn_type": "http", "host": "new-host.com"}, + ) + check("PATCH succeeds after test completes", r3.status_code == 200, f"got {r3.status_code}: {r3.text}") + + cleanup_connection("block_patch_conn") + + +# ============================================================ +# TEST 5: Block DELETE while test is active +# ============================================================ +def test_block_delete(): + header("TEST 5: DELETE blocked while async test is running") + + requests.post( + f"{BASE}/connections", + headers=AUTH, + json={"connection_id": "block_del_conn", "conn_type": "http", "host": "example.com"}, + ) + + r = requests.post( + f"{BASE}/connections/test-async", + headers=AUTH, + json={"connection_id": "block_del_conn", "conn_type": "http", "host": "example.com"}, + ) + check("Async test started", r.status_code == 202, f"got {r.status_code}") + token = r.json().get("token") if r.status_code == 202 else None + + r2 = requests.delete(f"{BASE}/connections/block_del_conn", headers=AUTH) + check("DELETE returns 409", r2.status_code == 409, f"got {r2.status_code}: {r2.text}") + + if token: + poll_until_terminal(token) + + r3 = requests.delete(f"{BASE}/connections/block_del_conn", headers=AUTH) + check("DELETE succeeds after test completes", r3.status_code == 204, f"got {r3.status_code}: {r3.text}") + + +# ============================================================ +# TEST 6: commit_on_success=True creates a new connection +# ============================================================ +def test_commit_on_success_create(): + header("TEST 6: commit_on_success=True creates connection on success") + + cleanup_connection("commit_new_conn") + + # Verify connection doesn't exist + r = requests.get(f"{BASE}/connections/commit_new_conn", headers=AUTH) + check("Connection does not exist yet", r.status_code == 404) + + r = requests.post( + f"{BASE}/connections/test-async", + headers=AUTH, + json={ + "connection_id": "commit_new_conn", + "conn_type": "http", + "host": "https://httpbin.org", + "commit_on_success": True, + }, + ) + check("POST returns 202", r.status_code == 202, f"got {r.status_code}") + if r.status_code != 202: + return + + token = r.json()["token"] + result = poll_until_terminal(token) + print(f" Test result: {result.get('state')} - {result.get('result_message', '')}") + + if result.get("state") == "success": + r2 = requests.get(f"{BASE}/connections/commit_new_conn", headers=AUTH) + check("Connection was created", r2.status_code == 200, f"got {r2.status_code}") + if r2.status_code == 200: + conn = r2.json() + check("conn_type matches", conn.get("conn_type") == "http") + check("host matches", conn.get("host") == "https://httpbin.org") + else: + print(" Test failed — cannot verify commit (this is OK if no HTTP hook)") + check("Test did not succeed (commit not applicable)", True) + + cleanup_connection("commit_new_conn") + + +# ============================================================ +# TEST 7: commit_on_success=False does NOT create connection +# ============================================================ +def test_commit_on_success_false(): + header("TEST 7: commit_on_success=False does NOT create connection") + + cleanup_connection("no_commit_conn") + + r = requests.post( + f"{BASE}/connections/test-async", + headers=AUTH, + json={ + "connection_id": "no_commit_conn", + "conn_type": "http", + "host": "https://httpbin.org", + "commit_on_success": False, + }, + ) + check("POST returns 202", r.status_code == 202, f"got {r.status_code}") + if r.status_code != 202: + return + + token = r.json()["token"] + result = poll_until_terminal(token) + print(f" Test result: {result.get('state')}") + + r2 = requests.get(f"{BASE}/connections/no_commit_conn", headers=AUTH) + check("Connection was NOT created", r2.status_code == 404, f"got {r2.status_code}") + + +# ============================================================ +# TEST 8: commit_on_success=True updates existing connection +# ============================================================ +def test_commit_on_success_update(): + header("TEST 8: commit_on_success=True updates existing connection") + + cleanup_connection("commit_update_conn") + + # Create a connection first + requests.post( + f"{BASE}/connections", + headers=AUTH, + json={"connection_id": "commit_update_conn", "conn_type": "http", "host": "old-host.example.com"}, + ) + + r = requests.post( + f"{BASE}/connections/test-async", + headers=AUTH, + json={ + "connection_id": "commit_update_conn", + "conn_type": "http", + "host": "https://httpbin.org", + "login": "new_user", + "commit_on_success": True, + }, + ) + check("POST returns 202", r.status_code == 202, f"got {r.status_code}") + if r.status_code != 202: + return + + token = r.json()["token"] + result = poll_until_terminal(token) + print(f" Test result: {result.get('state')}") + + if result.get("state") == "success": + r2 = requests.get(f"{BASE}/connections/commit_update_conn", headers=AUTH) + check("Connection still exists", r2.status_code == 200) + if r2.status_code == 200: + conn = r2.json() + check("Host was updated", conn.get("host") == "https://httpbin.org", f"got {conn.get('host')}") + check("Login was updated", conn.get("login") == "new_user", f"got {conn.get('login')}") + else: + print(" Test failed — cannot verify update (this is OK if no HTTP hook)") + + cleanup_connection("commit_update_conn") + + +# ============================================================ +# TEST 9: Polling with invalid token returns 404 +# ============================================================ +def test_poll_invalid_token(): + header("TEST 9: Polling with invalid token returns 404") + + r = requests.get(f"{BASE}/connections/test-async/this-token-does-not-exist", headers=AUTH) + check("GET returns 404", r.status_code == 404, f"got {r.status_code}") + + +# ============================================================ +# TEST 10: Sync test endpoint still works (unchanged) +# ============================================================ +def test_sync_test_still_works(): + header("TEST 10: Sync test endpoint still works") + + r = requests.post( + f"{BASE}/connections/test", + headers=AUTH, + json={"connection_id": "sync_test", "conn_type": "http", "host": "https://httpbin.org"}, + ) + # 200 if test ran, 403 if disabled + check("Sync test returns 200 or 403", r.status_code in (200, 403), f"got {r.status_code}: {r.text}") + if r.status_code == 200: + body = r.json() + check("Response has status field", "status" in body) + check("Response has message field", "message" in body) + + +# ============================================================ +# TEST 11: Old PATCH /{connection_id}/test endpoint is removed +# ============================================================ +def test_old_patch_endpoint_removed(): + header("TEST 11: Old PATCH /{connection_id}/test endpoint is removed") + + r = requests.patch( + f"{BASE}/connections/some_conn/test", + headers=AUTH, + json={"connection_id": "some_conn", "conn_type": "http"}, + ) + # Should be 404 (route doesn't exist) or 405 (method not allowed) + # FastAPI with no matching route returns 404 + check( + "PATCH /connections/{id}/test returns 404 or 405", + r.status_code in (404, 405), + f"got {r.status_code}: {r.text}", + ) + + +# ============================================================ +# MAIN +# ============================================================ +def main(): + + print("Async Connection Test — Manual E2E Test Suite") + print(f"Target: {BASE}") + print() + + global AUTH + + # Check server is reachable and get token + try: + r = requests.get(f"{BASE.rsplit('/api', 1)[0]}/api/v2/version", timeout=5) + if r.status_code != 200: + print(f"Server returned unexpected status {r.status_code}. Is Breeze running?") + sys.exit(1) + AUTH = get_token() + print("Authenticated with JWT token") + except requests.ConnectionError: + print("Cannot connect to server. Start Breeze first:") + print(" breeze start-airflow") + sys.exit(1) + + # Check if test_connection is enabled + disabled = test_disabled_403() + if disabled: + print("\n*** test_connection is disabled. Enable it with:") + print(" AIRFLOW__CORE__TEST_CONNECTION=Enabled") + print(" Then restart Breeze and re-run this script.") + sys.exit(1) + + # Run all tests + test_basic_async_test() + test_duplicate_409() + test_block_patch() + test_block_delete() + test_commit_on_success_create() + test_commit_on_success_false() + test_commit_on_success_update() + test_poll_invalid_token() + test_sync_test_still_works() + test_old_patch_endpoint_removed() + + # Summary + header("SUMMARY") + total = PASS_COUNT + FAIL_COUNT + print(f" {PASS_COUNT}/{total} checks passed") + if FAIL_COUNT > 0: + print(f" {FAIL_COUNT} checks FAILED") + sys.exit(1) + else: + print(" All checks passed!") + sys.exit(0) + + +if __name__ == "__main__": + main() diff --git a/devel-common/src/tests_common/test_utils/db.py b/devel-common/src/tests_common/test_utils/db.py index 76c453b9952d1..dfe181284888e 100644 --- a/devel-common/src/tests_common/test_utils/db.py +++ b/devel-common/src/tests_common/test_utils/db.py @@ -473,9 +473,9 @@ def clear_db_teams(): def clear_db_connection_tests(): with create_session() as session: if AIRFLOW_V_3_2_PLUS: - from airflow.models.connection_test import ConnectionTest + from airflow.models.connection_test import ConnectionTestRequest - session.execute(delete(ConnectionTest)) + session.execute(delete(ConnectionTestRequest)) @_retry_db diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 59bf62b0f7729..1f882b97c20a9 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -859,6 +859,11 @@ class ConnectionTestOperations: 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()) + def update_state( self, id: uuid.UUID, state: ConnectionTestState, result_message: str | None = None ) -> None: diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 87d7bc73004a1..c5f442b2b7741 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -78,6 +78,21 @@ class ConnectionResponse(BaseModel): extra: Annotated[str | None, Field(title="Extra")] = None +class ConnectionTestConnectionResponse(BaseModel): + """ + Connection data returned to workers from a test request. + """ + + conn_id: Annotated[str, Field(title="Conn 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 + password: Annotated[str | None, Field(title="Password")] = None + schema_: Annotated[str | None, Field(alias="schema", title="Schema")] = None + port: Annotated[int | None, Field(title="Port")] = None + extra: Annotated[str | None, Field(title="Extra")] = None + + class ConnectionTestState(str, Enum): """ All possible states of a connection test. From fd3fbf78d8ca90c1cb221ed0cf7134c397016b4a Mon Sep 17 00:00:00 2001 From: Anish Date: Sat, 21 Mar 2026 14:56:34 -0500 Subject: [PATCH 34/38] fix ci faliure --- .../api_fastapi/execution_api/routes/connection_tests.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py index de6f643ea2de3..47986523271cf 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/connection_tests.py @@ -62,7 +62,7 @@ def get_connection_test_connection( host=ct.host, login=ct.login, password=ct.password, - schema_=ct.schema, + schema=ct.schema, port=ct.port, extra=ct.extra, ) From 42c9d0786d924fe17e5259355a11643f700228a7 Mon Sep 17 00:00:00 2001 From: Anish Date: Sat, 21 Mar 2026 18:54:07 -0500 Subject: [PATCH 35/38] fix failing tests --- .../v2026_03_31/test_connection_tests.py | 6 +++--- .../unit/executors/test_local_executor.py | 19 +++++++++---------- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py index f7f8aa44769c8..f907611f38ae9 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_03_31/test_connection_tests.py @@ -18,7 +18,7 @@ import pytest -from airflow.models.connection_test import ConnectionTest, ConnectionTestState +from airflow.models.connection_test import ConnectionTestRequest, ConnectionTestState from tests_common.test_utils.db import clear_db_connection_tests @@ -43,7 +43,7 @@ def setup_teardown(self): def test_old_version_returns_404(self, old_ver_client, session): """PATCH /connection-tests/{id} should not exist in older API versions.""" - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="test_conn") ct.state = ConnectionTestState.RUNNING session.add(ct) session.commit() @@ -56,7 +56,7 @@ def test_old_version_returns_404(self, old_ver_client, session): def test_head_version_works(self, client, session): """PATCH /connection-tests/{id} should work in the current API version.""" - ct = ConnectionTest(connection_id="test_conn") + ct = ConnectionTestRequest(conn_type="test_type", connection_id="test_conn") ct.state = ConnectionTestState.RUNNING session.add(ct) session.commit() diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index 227b4a3bb21d6..4bca67542ddf3 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -37,7 +37,6 @@ from airflow.models.callback import CallbackFetchMethod from airflow.models.connection_test import ConnectionTestState from airflow.sdk.api.datamodels._generated import ConnectionResponse -from airflow.sdk.execution_time.comms import ErrorResponse from airflow.settings import Session from airflow.utils.state import State @@ -393,7 +392,7 @@ class TestLocalExecutorConnectionTestExecution: def test_successful_connection_test(self, MockClient, _mock_signal): """Fetches connection via Execution API, runs test, reports SUCCESS.""" mock_client = MockClient.return_value - mock_client.connections.get.return_value = ConnectionResponse( + mock_client.connection_tests.get_connection.return_value = ConnectionResponse( conn_id="test_conn", conn_type="http", host="httpbin.org", @@ -426,7 +425,7 @@ def test_successful_connection_test(self, MockClient, _mock_signal): def test_failed_connection_test(self, MockClient, _mock_signal): """Fetches connection via Execution API, test fails, reports FAILED.""" mock_client = MockClient.return_value - mock_client.connections.get.return_value = ConnectionResponse( + mock_client.connection_tests.get_connection.return_value = ConnectionResponse( conn_id="test_conn", conn_type="postgres", host="db.example.com", @@ -456,9 +455,9 @@ def test_failed_connection_test(self, MockClient, _mock_signal): assert calls[1].args == (test_id, ConnectionTestState.FAILED, "Connection refused") def test_connection_not_found_via_execution_api(self, MockClient, _mock_signal): - """Reports FAILED when connection is not found via Execution API.""" + """Reports FAILED when connection test is not found via Execution API.""" mock_client = MockClient.return_value - mock_client.connections.get.return_value = ErrorResponse(detail={"conn_id": "missing_conn"}) + mock_client.connection_tests.get_connection.side_effect = RuntimeError("Connection test not found") test_id = uuid7() workload = workloads.TestConnection( @@ -481,7 +480,7 @@ def test_connection_not_found_via_execution_api(self, MockClient, _mock_signal): def test_unexpected_exception_reports_failed(self, MockClient, _mock_signal): """Reports FAILED when an unexpected exception occurs.""" mock_client = MockClient.return_value - mock_client.connections.get.return_value = ConnectionResponse( + mock_client.connection_tests.get_connection.return_value = ConnectionResponse( conn_id="test_conn", conn_type="http", ) @@ -506,12 +505,12 @@ def test_unexpected_exception_reports_failed(self, MockClient, _mock_signal): calls = mock_client.connection_tests.update_state.call_args_list assert calls[-1].args[1] == ConnectionTestState.FAILED - assert "Connection test failed unexpectedly: Something broke" in calls[-1].args[2] + assert "Connection test failed unexpectedly: RuntimeError" in calls[-1].args[2] def test_connection_fields_passed_correctly(self, MockClient, _mock_signal): """Verifies all connection fields from the API response are passed to Connection.""" mock_client = MockClient.return_value - mock_client.connections.get.return_value = ConnectionResponse( + mock_client.connection_tests.get_connection.return_value = ConnectionResponse( conn_id="full_conn", conn_type="postgres", host="db.example.com", @@ -559,7 +558,7 @@ def capture_conn(*, conn): def test_timeout_reports_failed(self, MockClient, _mock_signal): """Reports FAILED with timeout message when TimeoutError is raised.""" mock_client = MockClient.return_value - mock_client.connections.get.return_value = ConnectionResponse( + mock_client.connection_tests.get_connection.return_value = ConnectionResponse( conn_id="test_conn", conn_type="http", ) @@ -592,7 +591,7 @@ def raise_timeout(*, conn): def test_alarm_is_cancelled_in_finally(self, MockClient, mock_signal): """signal.alarm(0) is called to cancel the timer even on success.""" mock_client = MockClient.return_value - mock_client.connections.get.return_value = ConnectionResponse( + mock_client.connection_tests.get_connection.return_value = ConnectionResponse( conn_id="test_conn", conn_type="http", ) From c6dd945b955e16504ed80c515c215de6b6377fad Mon Sep 17 00:00:00 2001 From: Anish Date: Sat, 21 Mar 2026 20:08:57 -0500 Subject: [PATCH 36/38] fix type issue --- airflow-core/src/airflow/executors/workloads/types.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/executors/workloads/types.py b/airflow-core/src/airflow/executors/workloads/types.py index de7804f169383..fa9e5f7ff774f 100644 --- a/airflow-core/src/airflow/executors/workloads/types.py +++ b/airflow-core/src/airflow/executors/workloads/types.py @@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, TypeAlias from airflow.models.callback import ExecutorCallback -from airflow.models.connection_test import ConnectionTest +from airflow.models.connection_test import ConnectionTestRequest from airflow.models.taskinstance import TaskInstance if TYPE_CHECKING: @@ -38,4 +38,4 @@ # Type alias for scheduler workloads (ORM models that can be routed to executors) # Must be outside TYPE_CHECKING for use in function signatures -SchedulerWorkload: TypeAlias = TaskInstance | ExecutorCallback | ConnectionTest +SchedulerWorkload: TypeAlias = TaskInstance | ExecutorCallback | ConnectionTestRequest From 54b3dec209aec92fca6c8dcbfab880f804643fbe Mon Sep 17 00:00:00 2001 From: Anish Date: Sun, 22 Mar 2026 03:38:03 -0500 Subject: [PATCH 37/38] removed fragment --- airflow-core/newsfragments/62343.feature.rst | 1 - 1 file changed, 1 deletion(-) delete mode 100644 airflow-core/newsfragments/62343.feature.rst diff --git a/airflow-core/newsfragments/62343.feature.rst b/airflow-core/newsfragments/62343.feature.rst deleted file mode 100644 index 52448232e3087..0000000000000 --- a/airflow-core/newsfragments/62343.feature.rst +++ /dev/null @@ -1 +0,0 @@ -Add async connection testing support on workers From 01e353c2ea7e8d4928239d6d2d5a8839948002b9 Mon Sep 17 00:00:00 2001 From: Anish Date: Sun, 22 Mar 2026 03:55:27 -0500 Subject: [PATCH 38/38] remove unwanted file --- dev/test_async_connection.py | 483 ----------------------------------- 1 file changed, 483 deletions(-) delete mode 100644 dev/test_async_connection.py diff --git a/dev/test_async_connection.py b/dev/test_async_connection.py deleted file mode 100644 index 4d5f20a5d4ec5..0000000000000 --- a/dev/test_async_connection.py +++ /dev/null @@ -1,483 +0,0 @@ -#!/usr/bin/env python -# 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. -""" -Manual end-to-end tests for async connection testing via the REST API. - -Run against a live Breeze server: - uv run --project dev python dev/test_async_connection.py - -Expects: - - Breeze API server at http://localhost:8080 - - AIRFLOW__CORE__TEST_CONNECTION=Enabled - - Default admin:admin credentials -""" - -from __future__ import annotations - -import sys -import time - -import requests - -BASE = "http://localhost:28080/api/v2" -AUTH_URL = "http://localhost:28080/auth/token" -AUTH = None # Set by get_token() -PASS_COUNT = 0 -FAIL_COUNT = 0 - - -def get_token() -> dict: - """Get a JWT token and return headers dict for requests.""" - r = requests.post(AUTH_URL, json={"username": "admin", "password": "admin"}, timeout=5) - r.raise_for_status() - token = r.json()["access_token"] - return {"Authorization": f"Bearer {token}"} - - -def header(title: str) -> None: - print(f"\n{'=' * 60}") - print(f" {title}") - print(f"{'=' * 60}") - - -def check(label: str, condition: bool, detail: str = "") -> None: - global PASS_COUNT, FAIL_COUNT - status = "PASS" if condition else "FAIL" - if condition: - PASS_COUNT += 1 - else: - FAIL_COUNT += 1 - msg = f" [{status}] {label}" - if detail and not condition: - msg += f" -- {detail}" - print(msg) - - -def cleanup_connection(conn_id: str) -> None: - """Delete a connection if it exists, ignoring 404/409.""" - r = requests.delete(f"{BASE}/connections/{conn_id}", headers=AUTH) - # 204=deleted, 404=didn't exist, 409=active test blocking - if r.status_code == 409: - # Wait for active test to finish then retry - time.sleep(3) - requests.delete(f"{BASE}/connections/{conn_id}", headers=AUTH) - - -def poll_until_terminal(token: str, timeout: int = 30) -> dict: - """Poll test status until terminal or timeout.""" - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - r = requests.get(f"{BASE}/connections/test-async/{token}", headers=AUTH) - body = r.json() - if body.get("state") in ("success", "failed"): - return body - time.sleep(1) - return body - - -# ============================================================ -# TEST 1: Basic async test with connection fields -# ============================================================ -def test_basic_async_test(): - header("TEST 1: Basic async test — POST with connection fields") - - r = requests.post( - f"{BASE}/connections/test-async", - headers=AUTH, - json={ - "connection_id": "manual_test_http", - "conn_type": "http", - "host": "https://httpbin.org", - }, - ) - check("POST returns 202", r.status_code == 202, f"got {r.status_code}: {r.text}") - if r.status_code != 202: - return - - body = r.json() - check("Response has token", "token" in body and len(body["token"]) > 0) - check("Response has connection_id", body.get("connection_id") == "manual_test_http") - check("State is pending", body.get("state") == "pending") - - token = body["token"] - print(f" Token: {token[:20]}...") - - # Poll for status - print(" Polling for result...") - result = poll_until_terminal(token) - check( - "Test reached terminal state", - result.get("state") in ("success", "failed"), - f"state={result.get('state')}", - ) - check("Response has no 'reverted' field", "reverted" not in result) - print(f" Final state: {result.get('state')}") - print(f" Message: {result.get('result_message', '(none)')}") - - -# ============================================================ -# TEST 2: 409 duplicate active test -# ============================================================ -def test_duplicate_409(): - header("TEST 2: 409 for duplicate active test on same connection_id") - - body = { - "connection_id": "dup_test_conn", - "conn_type": "http", - "host": "https://httpbin.org", - } - - r1 = requests.post(f"{BASE}/connections/test-async", headers=AUTH, json=body) - check("First POST returns 202", r1.status_code == 202, f"got {r1.status_code}") - - r2 = requests.post(f"{BASE}/connections/test-async", headers=AUTH, json=body) - check("Second POST returns 409", r2.status_code == 409, f"got {r2.status_code}: {r2.text}") - if r2.status_code == 409: - check("409 message mentions async test", "async test" in r2.json().get("detail", "")) - - # Wait for first test to finish so cleanup works - if r1.status_code == 202: - poll_until_terminal(r1.json()["token"]) - - -# ============================================================ -# TEST 3: 403 when test_connection is disabled -# ============================================================ -def test_disabled_403(): - header("TEST 3: 403 when test_connection is disabled (skip if enabled)") - - r = requests.post( - f"{BASE}/connections/test-async", - headers=AUTH, - json={"connection_id": "x", "conn_type": "http"}, - ) - if r.status_code == 403: - check("Returns 403 when disabled", True) - print(" NOTE: test_connection is disabled. Remaining tests will skip.") - return True # signal disabled - if r.status_code == 202: - check("test_connection is enabled (403 test skipped)", True) - poll_until_terminal(r.json()["token"]) - return False - check("Unexpected status", False, f"got {r.status_code}: {r.text}") - return True - - -# ============================================================ -# TEST 4: Block PATCH while test is active -# ============================================================ -def test_block_patch(): - header("TEST 4: PATCH blocked while async test is running") - - # Create a real connection first - requests.post( - f"{BASE}/connections", - headers=AUTH, - json={"connection_id": "block_patch_conn", "conn_type": "http", "host": "example.com"}, - ) - - # Start an async test - r = requests.post( - f"{BASE}/connections/test-async", - headers=AUTH, - json={"connection_id": "block_patch_conn", "conn_type": "http", "host": "example.com"}, - ) - check("Async test started", r.status_code == 202, f"got {r.status_code}") - token = r.json().get("token") if r.status_code == 202 else None - - # Try to PATCH the same connection - r2 = requests.patch( - f"{BASE}/connections/block_patch_conn", - headers=AUTH, - json={"connection_id": "block_patch_conn", "conn_type": "http", "host": "new-host.com"}, - ) - check("PATCH returns 409", r2.status_code == 409, f"got {r2.status_code}: {r2.text}") - - # Wait for test to finish - if token: - poll_until_terminal(token) - - # PATCH should work now - r3 = requests.patch( - f"{BASE}/connections/block_patch_conn", - headers=AUTH, - json={"connection_id": "block_patch_conn", "conn_type": "http", "host": "new-host.com"}, - ) - check("PATCH succeeds after test completes", r3.status_code == 200, f"got {r3.status_code}: {r3.text}") - - cleanup_connection("block_patch_conn") - - -# ============================================================ -# TEST 5: Block DELETE while test is active -# ============================================================ -def test_block_delete(): - header("TEST 5: DELETE blocked while async test is running") - - requests.post( - f"{BASE}/connections", - headers=AUTH, - json={"connection_id": "block_del_conn", "conn_type": "http", "host": "example.com"}, - ) - - r = requests.post( - f"{BASE}/connections/test-async", - headers=AUTH, - json={"connection_id": "block_del_conn", "conn_type": "http", "host": "example.com"}, - ) - check("Async test started", r.status_code == 202, f"got {r.status_code}") - token = r.json().get("token") if r.status_code == 202 else None - - r2 = requests.delete(f"{BASE}/connections/block_del_conn", headers=AUTH) - check("DELETE returns 409", r2.status_code == 409, f"got {r2.status_code}: {r2.text}") - - if token: - poll_until_terminal(token) - - r3 = requests.delete(f"{BASE}/connections/block_del_conn", headers=AUTH) - check("DELETE succeeds after test completes", r3.status_code == 204, f"got {r3.status_code}: {r3.text}") - - -# ============================================================ -# TEST 6: commit_on_success=True creates a new connection -# ============================================================ -def test_commit_on_success_create(): - header("TEST 6: commit_on_success=True creates connection on success") - - cleanup_connection("commit_new_conn") - - # Verify connection doesn't exist - r = requests.get(f"{BASE}/connections/commit_new_conn", headers=AUTH) - check("Connection does not exist yet", r.status_code == 404) - - r = requests.post( - f"{BASE}/connections/test-async", - headers=AUTH, - json={ - "connection_id": "commit_new_conn", - "conn_type": "http", - "host": "https://httpbin.org", - "commit_on_success": True, - }, - ) - check("POST returns 202", r.status_code == 202, f"got {r.status_code}") - if r.status_code != 202: - return - - token = r.json()["token"] - result = poll_until_terminal(token) - print(f" Test result: {result.get('state')} - {result.get('result_message', '')}") - - if result.get("state") == "success": - r2 = requests.get(f"{BASE}/connections/commit_new_conn", headers=AUTH) - check("Connection was created", r2.status_code == 200, f"got {r2.status_code}") - if r2.status_code == 200: - conn = r2.json() - check("conn_type matches", conn.get("conn_type") == "http") - check("host matches", conn.get("host") == "https://httpbin.org") - else: - print(" Test failed — cannot verify commit (this is OK if no HTTP hook)") - check("Test did not succeed (commit not applicable)", True) - - cleanup_connection("commit_new_conn") - - -# ============================================================ -# TEST 7: commit_on_success=False does NOT create connection -# ============================================================ -def test_commit_on_success_false(): - header("TEST 7: commit_on_success=False does NOT create connection") - - cleanup_connection("no_commit_conn") - - r = requests.post( - f"{BASE}/connections/test-async", - headers=AUTH, - json={ - "connection_id": "no_commit_conn", - "conn_type": "http", - "host": "https://httpbin.org", - "commit_on_success": False, - }, - ) - check("POST returns 202", r.status_code == 202, f"got {r.status_code}") - if r.status_code != 202: - return - - token = r.json()["token"] - result = poll_until_terminal(token) - print(f" Test result: {result.get('state')}") - - r2 = requests.get(f"{BASE}/connections/no_commit_conn", headers=AUTH) - check("Connection was NOT created", r2.status_code == 404, f"got {r2.status_code}") - - -# ============================================================ -# TEST 8: commit_on_success=True updates existing connection -# ============================================================ -def test_commit_on_success_update(): - header("TEST 8: commit_on_success=True updates existing connection") - - cleanup_connection("commit_update_conn") - - # Create a connection first - requests.post( - f"{BASE}/connections", - headers=AUTH, - json={"connection_id": "commit_update_conn", "conn_type": "http", "host": "old-host.example.com"}, - ) - - r = requests.post( - f"{BASE}/connections/test-async", - headers=AUTH, - json={ - "connection_id": "commit_update_conn", - "conn_type": "http", - "host": "https://httpbin.org", - "login": "new_user", - "commit_on_success": True, - }, - ) - check("POST returns 202", r.status_code == 202, f"got {r.status_code}") - if r.status_code != 202: - return - - token = r.json()["token"] - result = poll_until_terminal(token) - print(f" Test result: {result.get('state')}") - - if result.get("state") == "success": - r2 = requests.get(f"{BASE}/connections/commit_update_conn", headers=AUTH) - check("Connection still exists", r2.status_code == 200) - if r2.status_code == 200: - conn = r2.json() - check("Host was updated", conn.get("host") == "https://httpbin.org", f"got {conn.get('host')}") - check("Login was updated", conn.get("login") == "new_user", f"got {conn.get('login')}") - else: - print(" Test failed — cannot verify update (this is OK if no HTTP hook)") - - cleanup_connection("commit_update_conn") - - -# ============================================================ -# TEST 9: Polling with invalid token returns 404 -# ============================================================ -def test_poll_invalid_token(): - header("TEST 9: Polling with invalid token returns 404") - - r = requests.get(f"{BASE}/connections/test-async/this-token-does-not-exist", headers=AUTH) - check("GET returns 404", r.status_code == 404, f"got {r.status_code}") - - -# ============================================================ -# TEST 10: Sync test endpoint still works (unchanged) -# ============================================================ -def test_sync_test_still_works(): - header("TEST 10: Sync test endpoint still works") - - r = requests.post( - f"{BASE}/connections/test", - headers=AUTH, - json={"connection_id": "sync_test", "conn_type": "http", "host": "https://httpbin.org"}, - ) - # 200 if test ran, 403 if disabled - check("Sync test returns 200 or 403", r.status_code in (200, 403), f"got {r.status_code}: {r.text}") - if r.status_code == 200: - body = r.json() - check("Response has status field", "status" in body) - check("Response has message field", "message" in body) - - -# ============================================================ -# TEST 11: Old PATCH /{connection_id}/test endpoint is removed -# ============================================================ -def test_old_patch_endpoint_removed(): - header("TEST 11: Old PATCH /{connection_id}/test endpoint is removed") - - r = requests.patch( - f"{BASE}/connections/some_conn/test", - headers=AUTH, - json={"connection_id": "some_conn", "conn_type": "http"}, - ) - # Should be 404 (route doesn't exist) or 405 (method not allowed) - # FastAPI with no matching route returns 404 - check( - "PATCH /connections/{id}/test returns 404 or 405", - r.status_code in (404, 405), - f"got {r.status_code}: {r.text}", - ) - - -# ============================================================ -# MAIN -# ============================================================ -def main(): - - print("Async Connection Test — Manual E2E Test Suite") - print(f"Target: {BASE}") - print() - - global AUTH - - # Check server is reachable and get token - try: - r = requests.get(f"{BASE.rsplit('/api', 1)[0]}/api/v2/version", timeout=5) - if r.status_code != 200: - print(f"Server returned unexpected status {r.status_code}. Is Breeze running?") - sys.exit(1) - AUTH = get_token() - print("Authenticated with JWT token") - except requests.ConnectionError: - print("Cannot connect to server. Start Breeze first:") - print(" breeze start-airflow") - sys.exit(1) - - # Check if test_connection is enabled - disabled = test_disabled_403() - if disabled: - print("\n*** test_connection is disabled. Enable it with:") - print(" AIRFLOW__CORE__TEST_CONNECTION=Enabled") - print(" Then restart Breeze and re-run this script.") - sys.exit(1) - - # Run all tests - test_basic_async_test() - test_duplicate_409() - test_block_patch() - test_block_delete() - test_commit_on_success_create() - test_commit_on_success_false() - test_commit_on_success_update() - test_poll_invalid_token() - test_sync_test_still_works() - test_old_patch_endpoint_removed() - - # Summary - header("SUMMARY") - total = PASS_COUNT + FAIL_COUNT - print(f" {PASS_COUNT}/{total} checks passed") - if FAIL_COUNT > 0: - print(f" {FAIL_COUNT} checks FAILED") - sys.exit(1) - else: - print(" All checks passed!") - sys.exit(0) - - -if __name__ == "__main__": - main()