pierrejeambrun commented on code in PR #62343: URL: https://github.com/apache/airflow/pull/62343#discussion_r3028331645
########## airflow-core/src/airflow/models/connection_test.py: ########## @@ -0,0 +1,234 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import secrets +from datetime import datetime +from enum import Enum +from typing import TYPE_CHECKING +from uuid import UUID + +import structlog +import uuid6 +from sqlalchemy import Boolean, Index, Integer, String, Text, Uuid, select +from sqlalchemy.orm import Mapped, declared_attr, mapped_column, synonym + +from airflow._shared.timezones import timezone +from airflow.models.base import Base +from airflow.models.connection import Connection +from airflow.models.crypto import get_fernet +from airflow.utils.sqlalchemy import UtcDateTime + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + +log = structlog.get_logger(__name__) + + +class ConnectionTestState(str, Enum): + """All possible states of a connection test.""" + + PENDING = "pending" + QUEUED = "queued" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + + def __str__(self) -> str: + return self.value + + +ACTIVE_STATES = frozenset( + (ConnectionTestState.PENDING, ConnectionTestState.QUEUED, ConnectionTestState.RUNNING) +) +DISPATCHED_STATES = frozenset((ConnectionTestState.QUEUED, ConnectionTestState.RUNNING)) +TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED)) + + +class ConnectionTestRequest(Base): + """ + Tracks an async connection test request dispatched to a worker. + + Stores the full connection details so the worker reads from this table + instead of the real ``connection`` table. The real ``connection`` table + is only modified if the test succeeds and ``commit_on_success`` is True. + """ + Review Comment: Can you add a db composite unique constraint on `connection_id` where state is in `('pending', 'queued', 'running')`. To make sure that there cannot be two ConnectionTestRequest for the same ID. ########## airflow-core/src/airflow/config_templates/config.yml: ########## @@ -491,6 +491,24 @@ core: type: string example: ~ default: "Disabled" + connection_test_timeout: + description: | + Maximum number of seconds an async connection test is allowed to run + before it is considered timed out. The scheduler reaper uses this value + plus a grace period to mark stale tests as failed. + version_added: 3.2.0 + type: integer + example: ~ + default: "60" + max_connection_test_concurrency: + description: | + Maximum number of connection tests that can be active (QUEUED + RUNNING) + at the same time. Excess tests will remain in PENDING state until slots + become available. + version_added: 3.2.0 Review Comment: version_added isn't right. ########## airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py: ########## @@ -77,12 +78,46 @@ class ConnectionCollectionResponse(BaseModel): class ConnectionTestResponse(BaseModel): - """Connection Test serializer for responses.""" + """Connection Test serializer for synchronous test responses.""" status: bool message: str +class ConnectionTestRequestBody(StrictBaseModel): + """Request body for async connection test.""" + + connection_id: str + conn_type: str + host: str | None = None + login: str | None = None + schema_: str | None = Field(None, alias="schema") + port: int | None = None + password: str | None = None + extra: str | None = None + commit_on_success: bool = False Review Comment: We should probably add a documentation on what `commit_on_success` is doing. So the generated openapi spec mention that piece of documentation. ########## airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py: ########## @@ -257,6 +286,84 @@ def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse: os.environ.pop(conn_env_var, None) +@connections_router.post( + "/test-async", + status_code=status.HTTP_202_ACCEPTED, + responses=create_openapi_http_exception_doc([status.HTTP_403_FORBIDDEN, status.HTTP_409_CONFLICT]), + dependencies=[Depends(requires_access_connection(method="POST")), Depends(action_logging())], +) +def test_connection_async( + test_body: ConnectionTestRequestBody, + session: SessionDep, +) -> ConnectionTestQueuedResponse: + """ + Queue an async connection test to be executed on a worker. + + The connection data is stored in the test request table and the worker + reads from there. Returns a token to poll for the result via + GET /connections/test-async/{token}. + """ + _ensure_test_connection_enabled() + + # Only one active test per connection_id at a time. + _check_no_active_test(test_body.connection_id, session) + + connection_test = ConnectionTestRequest( + connection_id=test_body.connection_id, + conn_type=test_body.conn_type, + host=test_body.host, + login=test_body.login, + password=test_body.password, + schema=test_body.schema_, + port=test_body.port, + extra=test_body.extra, + commit_on_success=test_body.commit_on_success, + executor=test_body.executor, + queue=test_body.queue, + ) + session.add(connection_test) + session.flush() Review Comment: I don't think we need that manual flush. ########## airflow-core/src/airflow/executors/local_executor.py: ########## @@ -169,6 +178,81 @@ def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_con raise RuntimeError(error_msg or "Callback execution failed") +def _execute_connection_test(log: Logger, workload: workloads.TestConnection, team_conf) -> None: + """ + Execute a connection test workload. + + Constructs an SDK ``Client``, fetches the connection via the Execution API, + enforces a timeout via ``signal.alarm``, and reports all outcomes back + through the Execution API. + + :param log: Logger instance + :param workload: The TestConnection workload to execute + :param team_conf: Team-specific executor configuration + """ + # Lazy import: SDK modules must not be loaded at module level to avoid + # coupling core (scheduler-loaded) code to the SDK. + from airflow.sdk.api.client import Client + + setproctitle( + f"{_get_executor_process_title_prefix(team_conf.team_name)} connection-test {workload.connection_id}", + log, + ) + + base_url = team_conf.get("api", "base_url", fallback="/") + if base_url.startswith("/"): + base_url = f"http://localhost:8080{base_url}"; Review Comment: not sure that hardcoding `http://localhost:8080` makes sense. ########## airflow-ctl/src/airflowctl/api/datamodels/generated.py: ########## @@ -245,15 +245,58 @@ class ConnectionResponse(BaseModel): team_name: Annotated[str | None, Field(title="Team Name")] = None +class ConnectionTestQueuedResponse(BaseModel): + """ + Response returned when an async connection test is queued. + """ + + token: Annotated[str, Field(title="Token")] + connection_id: Annotated[str, Field(title="Connection Id")] + state: Annotated[str, Field(title="State")] + + +class ConnectionTestRequestBody(BaseModel): + """ + Request body for async connection test. + """ + + model_config = ConfigDict( + extra="forbid", + ) + connection_id: Annotated[str, Field(title="Connection Id")] + conn_type: Annotated[str, Field(title="Conn Type")] + host: Annotated[str | None, Field(title="Host")] = None + login: Annotated[str | None, Field(title="Login")] = None + schema_: Annotated[str | None, Field(alias="schema", title="Schema")] = None + port: Annotated[int | None, Field(title="Port")] = None + password: Annotated[str | None, Field(title="Password")] = None + extra: Annotated[str | None, Field(title="Extra")] = None + commit_on_success: Annotated[bool | None, Field(title="Commit On Success")] = False + executor: Annotated[str | None, Field(title="Executor")] = None + queue: Annotated[str | None, Field(title="Queue")] = None + + class ConnectionTestResponse(BaseModel): """ - Connection Test serializer for responses. + Connection Test serializer for synchronous test responses. """ status: Annotated[bool, Field(title="Status")] message: Annotated[str, Field(title="Message")] +class ConnectionTestStatusResponse(BaseModel): Review Comment: ```suggestion class AsyncConnectionTestResponse(BaseModel): ``` ########## airflow-core/src/airflow/config_templates/config.yml: ########## @@ -491,6 +491,24 @@ core: type: string example: ~ default: "Disabled" + connection_test_timeout: + description: | + Maximum number of seconds an async connection test is allowed to run + before it is considered timed out. The scheduler reaper uses this value + plus a grace period to mark stale tests as failed. + version_added: 3.2.0 + type: integer + example: ~ + default: "60" + max_connection_test_concurrency: + description: | + Maximum number of connection tests that can be active (QUEUED + RUNNING) + at the same time. Excess tests will remain in PENDING state until slots + become available. + version_added: 3.2.0 Review Comment: Config keys should be consistently in [scheduler] section `connection_test_timeout` and `max_connection_test_concurrency` are in [core] but `connection_test_reaper_interval` is in [scheduler]. All three seems only consumed by the scheduler. Moving them all to [scheduler] would be consistent with how other scheduler-only settings are organized. ########## airflow-core/src/airflow/models/connection_test.py: ########## @@ -0,0 +1,234 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import secrets +from datetime import datetime +from enum import Enum +from typing import TYPE_CHECKING +from uuid import UUID + +import structlog +import uuid6 +from sqlalchemy import Boolean, Index, Integer, String, Text, Uuid, select +from sqlalchemy.orm import Mapped, declared_attr, mapped_column, synonym + +from airflow._shared.timezones import timezone +from airflow.models.base import Base +from airflow.models.connection import Connection +from airflow.models.crypto import get_fernet +from airflow.utils.sqlalchemy import UtcDateTime + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + +log = structlog.get_logger(__name__) + + +class ConnectionTestState(str, Enum): + """All possible states of a connection test.""" + + PENDING = "pending" + QUEUED = "queued" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + + def __str__(self) -> str: + return self.value + + +ACTIVE_STATES = frozenset( + (ConnectionTestState.PENDING, ConnectionTestState.QUEUED, ConnectionTestState.RUNNING) +) +DISPATCHED_STATES = frozenset((ConnectionTestState.QUEUED, ConnectionTestState.RUNNING)) +TERMINAL_STATES = frozenset((ConnectionTestState.SUCCESS, ConnectionTestState.FAILED)) + + +class ConnectionTestRequest(Base): + """ + Tracks an async connection test request dispatched to a worker. + + Stores the full connection details so the worker reads from this table + instead of the real ``connection`` table. The real ``connection`` table + is only modified if the test succeeds and ``commit_on_success`` is True. + """ + + __tablename__ = "connection_test_request" + + id: Mapped[UUID] = mapped_column(Uuid(), primary_key=True, default=uuid6.uuid7) + token: Mapped[str] = mapped_column(String(64), nullable=False, unique=True) + connection_id: Mapped[str] = mapped_column(String(250), nullable=False) + state: Mapped[str] = mapped_column(String(20), nullable=False, default=ConnectionTestState.PENDING) + result_message: Mapped[str | None] = mapped_column(Text, nullable=True) + created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False) + updated_at: Mapped[datetime] = mapped_column( + UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow, nullable=False + ) + executor: Mapped[str | None] = mapped_column(String(256), nullable=True) + queue: Mapped[str | None] = mapped_column(String(256), nullable=True) + + # Connection fields — password and extra are Fernet-encrypted. + conn_type: Mapped[str] = mapped_column(String(500), nullable=False) + host: Mapped[str | None] = mapped_column(String(500), nullable=True) + login: Mapped[str | None] = mapped_column(Text, nullable=True) + _password: Mapped[str | None] = mapped_column("password", Text(), nullable=True) + schema: Mapped[str | None] = mapped_column("schema", String(500), nullable=True) + port: Mapped[int | None] = mapped_column(Integer, nullable=True) + _extra: Mapped[str | None] = mapped_column("extra", Text(), nullable=True) + commit_on_success: Mapped[bool] = mapped_column( + Boolean, nullable=False, default=False, server_default="0" + ) + + __table_args__ = (Index("idx_connection_test_request_state_created_at", state, created_at),) + + def __init__( + self, + *, + connection_id: str, + conn_type: str, + host: str | None = None, + login: str | None = None, + password: str | None = None, + schema: str | None = None, + port: int | None = None, + extra: str | None = None, + commit_on_success: bool = False, + executor: str | None = None, + queue: str | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self.connection_id = connection_id + self.conn_type = conn_type + self.host = host + self.login = login + self.password = password + self.schema = schema + self.port = port + self.extra = extra + self.commit_on_success = commit_on_success + self.executor = executor + self.queue = queue + self.token = secrets.token_urlsafe(32) + self.state = ConnectionTestState.PENDING + + def __repr__(self) -> str: + return ( + f"<ConnectionTestRequest id={self.id!r} connection_id={self.connection_id!r} state={self.state}>" + ) + + def get_executor_name(self) -> str | None: + """Return the executor name for scheduler routing.""" + return self.executor + + def get_dag_id(self) -> None: + """Return None — connection tests are not associated with any DAG.""" + return None + + def get_password(self) -> str | None: + if self._password: + fernet = get_fernet() + if not fernet.is_encrypted: + return self._password + return fernet.decrypt(bytes(self._password, "utf-8")).decode() + return self._password + + def set_password(self, value: str | None): + if value: + fernet = get_fernet() + self._password = fernet.encrypt(bytes(value, "utf-8")).decode() + else: + self._password = value + + @declared_attr + def password(cls): + """Password. The value is decrypted/encrypted when reading/setting the value.""" + return synonym("_password", descriptor=property(cls.get_password, cls.set_password)) + + def get_extra(self) -> str | None: + if self._extra: + fernet = get_fernet() + if not fernet.is_encrypted: + return self._extra + return fernet.decrypt(bytes(self._extra, "utf-8")).decode() + return self._extra + + def set_extra(self, value: str | None): + if value: + fernet = get_fernet() + self._extra = fernet.encrypt(bytes(value, "utf-8")).decode() + else: + self._extra = value + + @declared_attr + def extra(cls): + """Extra data. The value is decrypted/encrypted when reading/setting the value.""" + return synonym("_extra", descriptor=property(cls.get_extra, cls.set_extra)) + Review Comment: I think all that is duplicated from `Connection` model. Can we consider factorizing this into a common mixin (ConnectionMixin, FernetFieldMixin) or shared utility. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
