diff --git a/airflow-core/newsfragments/70927.significant.rst b/airflow-core/newsfragments/70927.significant.rst new file mode 100644 index 0000000000000..59a2359935494 --- /dev/null +++ b/airflow-core/newsfragments/70927.significant.rst @@ -0,0 +1,28 @@ +Callbacks are claimed through the Execution API before they run + +A CeleryExecutor worker now claims a callback through the Execution API +(``POST /execution/callbacks/{callback_id}/run``) before importing and running +it, transitioning the callback ``QUEUED -> RUNNING``. Previously a callback was +executed without the worker making any authenticated call first, so the workload +token minted for it was never validated by the API server. + +**Behaviour changes:** + +- A ``3.4`` worker requires an API server that serves the new endpoint. Mixed + deployments negotiate the API version, so a worker talking to an older API + server continues to behave as though the endpoint does not exist. +- Callbacks now emit an intermediate ``RUNNING`` state, matching the documented + ``QUEUED -> RUNNING -> SUCCESS/FAILED`` lifecycle. +- A redelivered or replayed callback message that reaches a callback already + ``RUNNING`` or in a terminal state is refused rather than executed again. + +* Types of change + + * [ ] Dag changes + * [ ] Config changes + * [x] API changes + * [ ] CLI changes + * [x] Behaviour changes + * [ ] Plugin changes + * [ ] Dependency changes + * [ ] Code interface changes diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/callback.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/callback.py new file mode 100644 index 0000000000000..cc7caeb5c5c80 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/callback.py @@ -0,0 +1,30 @@ +# 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 pydantic import BaseModel + +from airflow.utils.state import CallbackState + + +class CallbackRunResponse(BaseModel): + """Returned when a worker claims a callback for execution.""" + + id: UUID + state: CallbackState 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 7b19f3ddd3055..61fa022f4ced4 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 @@ -23,6 +23,7 @@ asset_events, asset_state_store, assets, + callbacks, connection_tests, connections, dag_runs, @@ -49,6 +50,7 @@ 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(callbacks.router, prefix="/callbacks", tags=["Callbacks"]) authenticated_router.include_router( connection_tests.router, prefix="/connection-tests", tags=["Connection Tests"] ) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py new file mode 100644 index 0000000000000..4779f806d554e --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py @@ -0,0 +1,82 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from uuid import UUID + +from cadwyn import VersionedAPIRouter +from fastapi import HTTPException, Security, status + +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.execution_api.datamodels.callback import CallbackRunResponse +from airflow.api_fastapi.execution_api.security import ExecutionAPIRoute, require_auth +from airflow.models.callback import Callback +from airflow.utils.state import CallbackState + +# ``cb:self`` makes the server enforce that the presented workload token was minted for *this* +# callback id (see ``require_auth``). Combined with ``token:workload``, a worker cannot reach this +# route without a validly-signed token whose subject is the callback it is claiming. +router = VersionedAPIRouter( + route_class=ExecutionAPIRoute, + dependencies=[ + Security(require_auth, scopes=["cb:self", "token:workload"]), + ], +) + + +@router.post( + "/{callback_id}/run", + responses={ + status.HTTP_404_NOT_FOUND: {"description": "Callback not found"}, + status.HTTP_409_CONFLICT: {"description": "Callback is not in a claimable state"}, + }, +) +def run_callback( + callback_id: UUID, + session: SessionDep, +) -> CallbackRunResponse: + """ + Atomically claim a callback for execution and transition it to RUNNING. + + The worker calls this *before* importing and invoking the callback. Its purpose is twofold: + the ``Security`` dependency forces the API server to validate the worker's token before any + callback code runs, and the ``QUEUED -> RUNNING`` transition is single-shot, so a redelivered + or replayed message that reaches a callback already RUNNING or terminal is refused rather than + executed a second time. + """ + callback = session.get(Callback, callback_id, with_for_update=True) + if callback is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"reason": "not_found", "message": f"Callback {callback_id} not found"}, + ) + + if callback.state != CallbackState.QUEUED: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "reason": "conflict", + "message": ( + f"Callback {callback_id} is in state {callback.state}; it can only be claimed " + "while QUEUED." + ), + }, + ) + + callback.state = CallbackState.RUNNING + + return CallbackRunResponse(id=callback.id, state=CallbackState.RUNNING) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/security.py b/airflow-core/src/airflow/api_fastapi/execution_api/security.py index 9de3493061f32..aae99783461c1 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -196,6 +196,13 @@ async def require_auth( status_code=status.HTTP_403_FORBIDDEN, detail="Token subject does not match connection test ID", ) + elif "cb:self" in security_scopes.scopes: + cb_self_id = str(request.path_params["callback_id"]) + if str(token.id) != cb_self_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Token subject does not match callback ID", + ) return token 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 dc7035d31e3c9..e6ece949edf19 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 @@ -51,9 +51,14 @@ AddTeamNameField, AddVariableKeysEndpoint, ) +from airflow.api_fastapi.execution_api.versions.v2026_08_01 import AddCallbackRunEndpoint bundle = VersionBundle( HeadVersion(), + Version( + "2026-08-01", + AddCallbackRunEndpoint, + ), Version( "2026-06-30", AddVariableKeysEndpoint, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_08_01.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_08_01.py new file mode 100644 index 0000000000000..35d85e7c6c3b4 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_08_01.py @@ -0,0 +1,30 @@ +# 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 cadwyn import VersionChange, endpoint + + +class AddCallbackRunEndpoint(VersionChange): + """Add the callbacks/{callback_id}/run endpoint that a worker uses to claim a callback.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + endpoint("/callbacks/{callback_id}/run", ["POST"]).didnt_exist, + ) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/conftest.py b/airflow-core/tests/unit/api_fastapi/execution_api/conftest.py index fda87d21e21a9..be6e7f259c54d 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/conftest.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/conftest.py @@ -58,7 +58,10 @@ async def mock_require_auth(request: Request) -> TIToken: raw_id = request.path_params.get( "task_instance_id", - request.path_params.get("connection_test_id", "00000000-0000-0000-0000-000000000000"), + request.path_params.get( + "connection_test_id", + request.path_params.get("callback_id", "00000000-0000-0000-0000-000000000000"), + ), ) try: ti_id = UUID(raw_id) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py index e98d918385933..af4fd28a27a64 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py @@ -257,6 +257,62 @@ def test_mismatched_subject_is_rejected(self, app): assert "does not match" in resp.json()["detail"] +class TestCbSelfScopeEnforcement: + """Routes with the ``cb:self`` scope reject a token minted for a different callback. + + This is the check that makes a workload token unforgeable at the callback boundary: a worker + presenting a token whose subject is not this callback id is refused before the route body runs. + """ + + PATH_CB_ID = "00000000-0000-0000-0000-0000000000c1" + OTHER_CB_ID = "00000000-0000-0000-0000-0000000000c2" + + @pytest.fixture + def app(self): + app = FastAPI() + + authenticated_router = APIRouter(dependencies=[Security(require_auth)]) + cb_self_router = APIRouter(dependencies=[Security(require_auth, scopes=["cb:self"])]) + + @cb_self_router.post("/{callback_id}/run") + def run_endpoint(callback_id: str): + return {"ok": True} + + authenticated_router.include_router(cb_self_router, prefix="/callbacks") + app.include_router(authenticated_router) + return app + + def _override_jwt(self, app: FastAPI, token_cb_id: UUID): + async def mock_jwt(request: Request): + return TIToken(id=token_cb_id, claims=TIClaims(scope="execution")) + + app.dependency_overrides[_jwt_bearer] = mock_jwt + + def test_matching_subject_is_accepted(self, app): + self._override_jwt(app, self.PATH_CB_ID) + client = TestClient(app) + + resp = client.post( + f"/callbacks/{self.PATH_CB_ID}/run", + headers={"Authorization": "Bearer fake"}, + ) + + assert resp.status_code == 200 + + def test_mismatched_subject_is_rejected(self, app): + """A token minted for one callback cannot claim another.""" + self._override_jwt(app, self.OTHER_CB_ID) + client = TestClient(app) + + resp = client.post( + f"/callbacks/{self.PATH_CB_ID}/run", + headers={"Authorization": "Bearer fake"}, + ) + + assert resp.status_code == 403 + assert "does not match" in resp.json()["detail"] + + class TestGetTeamNameDep: """Tests for get_team_name_dep avoiding unnecessary async sessions.""" diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py index bbf2be8704f46..6176a4977d346 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py @@ -49,6 +49,8 @@ # Connection test routes run from a queued worker context (workload-only). "PATCH /connection-tests/{connection_test_id}": {"workload"}, "GET /connection-tests/{connection_test_id}/connection": {"workload"}, + # A queued worker claims its callback (workload-only) before running it. + "POST /callbacks/{callback_id}/run": {"workload"}, } diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py new file mode 100644 index 0000000000000..201a26224e1f3 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py @@ -0,0 +1,64 @@ +# 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.callback import Callback, CallbackFetchMethod, ExecutorCallback +from airflow.sdk.definitions.callback import SyncCallback +from airflow.utils.state import CallbackState + +pytestmark = pytest.mark.db_test + + +def _make_callback(session, state: CallbackState) -> Callback: + cb = ExecutorCallback(SyncCallback("os.getcwd"), fetch_method=CallbackFetchMethod.IMPORT_PATH) + cb.state = state + session.add(cb) + session.commit() + return cb + + +class TestRunCallback: + def test_claim_transitions_queued_to_running(self, client, session): + """A QUEUED callback is claimed and moved to RUNNING; the response reports the new state.""" + cb = _make_callback(session, CallbackState.QUEUED) + + response = client.post(f"/execution/callbacks/{cb.id}/run") + + assert response.status_code == 200 + assert response.json() == {"id": str(cb.id), "state": "running"} + + session.expire_all() + assert session.get(Callback, cb.id).state == CallbackState.RUNNING + + def test_returns_404_for_unknown_callback(self, client): + response = client.post("/execution/callbacks/00000000-0000-0000-0000-000000000000/run") + assert response.status_code == 404 + + def test_returns_422_for_invalid_uuid(self, client): + response = client.post("/execution/callbacks/not-a-uuid/run") + assert response.status_code == 422 + + @pytest.mark.parametrize("state", [CallbackState.RUNNING, CallbackState.SUCCESS, CallbackState.FAILED]) + def test_a_callback_can_only_be_claimed_once(self, client, session, state): + """A callback already RUNNING or terminal cannot be claimed again — replay/redelivery is refused.""" + cb = _make_callback(session, state) + + response = client.post(f"/execution/callbacks/{cb.id}/run") + + assert response.status_code == 409 diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 7e681c1114956..95f6e9dbfb154 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -50,6 +50,7 @@ AssetResponse, AssetStateStorePutBody, AssetStateStoreResponse, + CallbackRunResponse, ConnectionResponse, ConnectionTestConnectionResponse, ConnectionTestResultBody, @@ -1067,6 +1068,25 @@ def get_detail_response(self, ti_id: uuid.UUID) -> HITLDetailResponse: return HITLDetailResponse.model_validate_json(resp.read()) +class CallbackOperations: + __slots__ = ("client",) + + def __init__(self, client: Client): + self.client = client + + def run(self, callback_id: uuid.UUID) -> CallbackRunResponse: + """ + Claim a callback for execution, transitioning it to RUNNING on the server. + + Called by the callback supervisor *before* the callback is imported and invoked. The + request carries the workload token minted for this callback; the server rejects it if the + token is missing, forged, or minted for a different callback, so this call is what stops + an unauthenticated message on the broker from reaching callback execution. + """ + resp = self.client.post(f"callbacks/{callback_id}/run") + return CallbackRunResponse.model_validate_json(resp.read()) + + class ConnectionTestOperations: __slots__ = ("client",) @@ -1298,6 +1318,12 @@ def connection_tests(self) -> ConnectionTestOperations: """Operations related to Connection Tests.""" return ConnectionTestOperations(self) + @lru_cache() # type: ignore[misc] + @property + def callbacks(self) -> CallbackOperations: + """Operations related to Callbacks.""" + return CallbackOperations(self) + @lru_cache() # type: ignore[misc] @property def dags(self) -> DagsOperations: diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index cc3c7eb0a8f20..5a8c4d8dcff04 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -27,7 +27,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, RootModel -API_VERSION: Final[str] = "2026-06-30" +API_VERSION: Final[str] = "2026-08-01" class AssetAliasReferenceAssetEventDagRun(BaseModel): @@ -63,6 +63,19 @@ class AssetProfile(BaseModel): type: Annotated[str, Field(title="Type")] +class CallbackState(str, Enum): + """ + All possible states of callbacks. + """ + + SCHEDULED = "scheduled" + PENDING = "pending" + QUEUED = "queued" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + + class ConnectionResponse(BaseModel): """ Connection schema for responses with fields that are needed for Runtime. @@ -654,6 +667,15 @@ class AssetStateStoreResponse(BaseModel): value: JsonValue | None +class CallbackRunResponse(BaseModel): + """ + Returned when a worker claims a callback for execution. + """ + + id: Annotated[UUID, Field(title="Id")] + state: CallbackState + + class ConnectionTestResultBody(BaseModel): """ Result a worker reports back for a connection test. diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py index 9830771701293..fda32af546c4a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -22,6 +22,7 @@ import signal import sys import time +import uuid from importlib import import_module from importlib.util import module_from_spec, spec_from_file_location from pathlib import Path @@ -433,6 +434,14 @@ def supervise_callback( else: logger = structlog.get_logger(logger_name="callback").bind() + # Claim the callback through the Execution API *before* importing and running it. This is + # the only point at which the workload token is redeemed: unlike a task, a callback makes + # no other authenticated API call, so without this the token minted for it is never + # checked. The server validates the token (workload scope, subject == this callback id) + # and transitions the callback QUEUED -> RUNNING; a missing, forged, or mismatched token + # is rejected here and the callback body never runs. + client.callbacks.run(uuid.UUID(str(id))) + try: process = CallbackSubprocess.start( id=id, diff --git a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py index a6858abcfcdd1..ea583626228a3 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py @@ -547,3 +547,57 @@ def test_callback_supervisor_should_exit_on_error(self, base_start_kwargs): self.mock_super_start.call_args.kwargs["target"]() assert exc_info.value.code == 1 + + +class TestSuperviseCallbackClaimsBeforeExecuting: + """The callback is claimed through the Execution API *before* it is imported and run. + + This is the boundary that stops an unauthenticated message on the Celery broker from reaching + callback execution: the claim call redeems the workload token, and nothing runs until it does. + """ + + CB_ID = "00000000-0000-0000-0000-0000000000cb" + + def _supervise(self, client, start_mock): + from airflow.sdk.execution_time.callback_supervisor import supervise_callback + + with patch( + "airflow.sdk.execution_time.callback_supervisor.CallbackSubprocess.start", + start_mock, + ): + return supervise_callback( + id=self.CB_ID, + callback_path="tests.does.not.matter", + callback_kwargs={}, + dag_rel_path=Path("dag.py"), + token="a-token", + server="http://localhost:9999", + client=client, + ) + + def test_claim_runs_before_the_subprocess_starts(self): + import uuid + + order = [] + client = Mock() + client.callbacks.run.side_effect = lambda *a, **k: order.append("claim") + + start_mock = Mock() + start_mock.side_effect = lambda *a, **k: (order.append("start"), Mock(wait=Mock(return_value=0)))[1] + + self._supervise(client, start_mock) + + client.callbacks.run.assert_called_once_with(uuid.UUID(self.CB_ID)) + assert order == ["claim", "start"], "the callback must be claimed before it is started" + + def test_a_rejected_claim_prevents_the_subprocess_from_starting(self): + """If the server refuses the claim (forged/mismatched/replayed token), nothing executes.""" + client = Mock() + client.callbacks.run.side_effect = RuntimeError("403 Forbidden: token does not match callback") + + start_mock = Mock() + + with pytest.raises(RuntimeError, match="token does not match callback"): + self._supervise(client, start_mock) + + start_mock.assert_not_called()