pierrejeambrun commented on code in PR #62343:
URL: https://github.com/apache/airflow/pull/62343#discussion_r3129903590
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3141,6 +3160,86 @@ def _activate_assets_generate_warnings() ->
Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ def _enqueue_connection_tests(self, *, session: Session) -> None:
+ """Enqueue pending connection tests to executors that support them."""
+ max_concurrency = conf.getint("scheduler",
"max_connection_test_concurrency", fallback=4)
+ timeout = conf.getint("scheduler", "connection_test_timeout",
fallback=60)
+
+ num_occupied_slots = sum(executor.slots_occupied for executor in
self.executors)
+ parallelism_budget = conf.getint("core", "parallelism") -
num_occupied_slots
+ if parallelism_budget <= 0:
+ return
+
+ active_count = session.scalar(
+ select(func.count(ConnectionTestRequest.id)).where(
+ ConnectionTestRequest.state.in_(DISPATCHED_STATES)
+ )
+ )
+ concurrency_budget = max_concurrency - (active_count or 0)
+ budget = min(concurrency_budget, parallelism_budget)
+ if budget <= 0:
+ return
+
+ pending_stmt = (
+ select(ConnectionTestRequest)
+ .where(ConnectionTestRequest.state == ConnectionTestState.PENDING)
+ .order_by(ConnectionTestRequest.created_at)
+ .limit(budget)
+ )
+ pending_stmt = with_row_locks(pending_stmt, session,
of=ConnectionTestRequest, skip_locked=True)
+ pending_tests = session.scalars(pending_stmt).all()
+
+ if not pending_tests:
+ return
+
+ for ct in pending_tests:
+ executor = self._try_to_load_executor(ct, session)
+ if executor is not None and not executor.supports_connection_test:
+ executor = None
+ if executor is None:
+ reason = (
+ f"No executor matches '{ct.executor}'"
+ if ct.executor
+ else "No executor supports connection testing"
+ )
+ ct.state = ConnectionTestState.FAILED
Review Comment:
Message is confusing. It looks like 'no executor supports testing' while
it's more like the executor attached to this workload doesn't support testing.
But in a multi executor env, other executor could support testing.
(I specify celeryexecutor for the test, that won't work, but I also have the
local executor available)
##########
airflow-core/src/airflow/models/connection.py:
##########
@@ -131,7 +131,7 @@ class Connection(Base, LoggingMixin):
host: Mapped[str | None] = mapped_column(String(500), nullable=True)
schema: Mapped[str | None] = mapped_column(String(500), nullable=True)
login: Mapped[str | None] = mapped_column(Text(), nullable=True)
- _password: Mapped[str | None] = mapped_column("password", Text(),
nullable=True)
+ # _password and _extra columns inherited from FernetFieldsMixin
Review Comment:
We need to remove that comment I guess.
##########
airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py:
##########
@@ -78,12 +79,49 @@ class ConnectionCollectionResponse(BaseModel):
class ConnectionTestResponse(BaseModel):
- """Connection Test serializer for responses."""
+ """Connection Test serializer for synchronous test responses."""
status: bool
message: str
+class ConnectionTestRequestBody(StrictBaseModel):
+ """Request body for async connection test."""
+
+ connection_id: str
+ conn_type: str
+ host: str | None = None
+ login: str | None = None
+ schema_: str | None = Field(None, alias="schema")
+ port: int | None = None
+ password: str | None = None
+ extra: str | None = None
+ commit_on_success: bool = Field(
+ default=False,
+ description="If True, save or update the connection in the connection
table when the test succeeds.",
+ )
+ executor: str | None = None
Review Comment:
We probably want to validate the `executor` against the configured executors
list.
No need to to a full roundtrip if the executor isn't valid, we can just
reject at the API layer level with a 422
##########
airflow-core/src/airflow/executors/workloads/connection_test.py:
##########
@@ -0,0 +1,88 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Connection test workload schema for executor communication."""
+
+from __future__ import annotations
+
+import uuid
+from typing import TYPE_CHECKING, Literal
+
+from pydantic import Field
+
+from airflow.executors.workloads.base import BaseWorkloadSchema
+from airflow.models.connection_test import ConnectionTestKey
+
+if TYPE_CHECKING:
+ from airflow.api_fastapi.auth.tokens import JWTGenerator
+ from airflow.models.connection_test import ConnectionTestState
+
+
+class TestConnection(BaseWorkloadSchema):
+ """Execute a connection test on a worker."""
+
+ connection_test_id: uuid.UUID
+ connection_id: str
+ timeout: int = 60
Review Comment:
To remove. This should come from the config and always provided.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py:
##########
@@ -57,24 +60,53 @@
from airflow.configuration import conf
from airflow.exceptions import AirflowNotFoundException
from airflow.models import Connection
+from airflow.models.connection_test import ACTIVE_STATES, ConnectionTestRequest
from airflow.secrets.environment_variables import CONN_ENV_PREFIX
from airflow.utils.db import create_default_connections as
db_create_default_connections
from airflow.utils.strings import get_random_string
connections_router = AirflowRouter(tags=["Connection"], prefix="/connections")
+def _ensure_test_connection_enabled() -> None:
+ """Raise 403 if connection testing is not enabled in the Airflow
configuration."""
+ if conf.get("core", "test_connection",
fallback="Disabled").lower().strip() != "enabled":
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ "Testing connections is disabled in Airflow configuration. "
+ "Contact your deployment admin to enable it.",
+ )
+
+
+def _check_no_active_test(connection_id: str, session: SessionDep) -> None:
+ """Raise 409 if there is an active connection test request for the given
connection_id."""
+ active_test = session.scalar(
+ select(ConnectionTestRequest).filter(
+ ConnectionTestRequest.connection_id == connection_id,
+ ConnectionTestRequest.state.in_(ACTIVE_STATES),
+ )
+ )
+ if active_test is not None:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ f"Cannot modify connection `{connection_id}` while an async test
is running. "
+ "This typically takes only a few seconds — please retry shortly.",
+ )
+
Review Comment:
Instead of having this we should rely on a db unique key constraint. This
can still produce race conditions and is not atomic. I would remove this and
find a db side solution for mysql.
##########
airflow-core/src/airflow/models/crypto.py:
##########
@@ -93,6 +96,66 @@ def rotate(self, msg: bytes | str) -> bytes:
return self._fernet.rotate(msg)
+class FernetFieldsMixin:
+ """
+ Mixin providing Fernet-encrypted ``password`` and ``extra`` fields.
+
+ Subclasses get ``_password`` / ``_extra`` mapped columns and transparent
+ encrypt-on-write / decrypt-on-read via SQLAlchemy synonyms.
+
+ Models that need additional logic (e.g. the ``Connection`` model's
+ ``is_encrypted`` flag) can override ``get_password``, ``set_password``,
+ ``get_extra``, and ``set_extra``.
Review Comment:
Shoudln't we instead use the same pattern for the TestConnection ? Have the
`is_encrypted` `is_extra_encrypted` here and the common logic here too?
This way we get a clear `FERNET_KEY configuration is missing` when
`is_encrypted=True` and worker doesn't have a key configured.
(To basically mirror the Connection pattern)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
