ashb commented on code in PR #66141: URL: https://github.com/apache/airflow/pull/66141#discussion_r3519967224
########## airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_07_01.py: ########## Review Comment: Since this has missed 3.3.0 we will want to future-date this version to ~2026-07-31 (and the release manager will fix the version when it's included in a release) ########## airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py: ########## @@ -0,0 +1,153 @@ +# 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 typing import Annotated +from uuid import UUID + +import structlog +from cadwyn import VersionedAPIRouter +from fastapi import Body, HTTPException, Response, Security, status +from structlog.contextvars import bind_contextvars + +from airflow.api_fastapi.auth.tokens import JWTGenerator +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.execution_api.datamodels.callback import CallbackTerminalStatePayload +from airflow.api_fastapi.execution_api.datamodels.token import TIToken +from airflow.api_fastapi.execution_api.deps import DepContainer +from airflow.api_fastapi.execution_api.security import CurrentTIToken, ExecutionAPIRoute, require_auth +from airflow.models.callback import Callback +from airflow.utils.state import CallbackState + +log = structlog.get_logger(__name__) + +router = VersionedAPIRouter(route_class=ExecutionAPIRoute) + + +def _require_self(token: TIToken, callback_id: UUID) -> None: + """Mirror the ``ti:self`` enforcement from security.py for callback routes.""" + if str(token.id) != str(callback_id): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Token subject does not match callback id", + ) + + [email protected]( + "/{callback_id}/run", + status_code=status.HTTP_204_NO_CONTENT, + dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], + responses={ + status.HTTP_403_FORBIDDEN: {"description": "Token subject does not match callback id"}, + status.HTTP_404_NOT_FOUND: {"description": "Callback not found"}, + status.HTTP_409_CONFLICT: {"description": "Callback is not in a state that can be marked running"}, + }, +) +def callback_run( + callback_id: UUID, + response: Response, + session: SessionDep, + services=DepContainer, + token: TIToken = CurrentTIToken, +) -> Response: + """ + Mark a callback as RUNNING. + + Mirrors ``PATCH /task-instances/{id}/run``: this is the single endpoint that + accepts a workload-scoped token and atomically (a) transitions the callback + from QUEUED to RUNNING and (b) issues a fresh execution-scoped token via the + ``Refreshed-API-Token`` response header. All subsequent supervisor calls hit + execution-only routes. Review Comment: I believe this doc string is what shows up in the OpenAPI generated spec, and thus as inline docs for any clients generated off this. As a result we should re-word this to make it be more user/client focused. (It's pretty close, I just don't think the "Mirrors PATCH" part is relevant to users. That can be a comment inside the fn itself if you think that is relevant. ########## airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py: ########## @@ -0,0 +1,300 @@ +# 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 typing import Literal +from unittest import mock +from uuid import UUID, uuid4 + +import pytest +from fastapi import Request + +from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator +from airflow.api_fastapi.execution_api.app import lifespan +from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken +from airflow.api_fastapi.execution_api.security import require_auth +from airflow.executors.workloads.callback import CallbackFetchMethod +from airflow.models.callback import ExecutorCallback +from airflow.utils.state import CallbackState + +from tests_common.test_utils.db import clear_db_callbacks + +pytestmark = pytest.mark.db_test + + +class _FakeCallbackDef: + """Minimal CallbackDefinitionProtocol stand-in for tests.""" + + path: str = "tests.fake.callback" + kwargs: dict = {} + executor: str | None = None + + def serialize(self) -> dict: + return {"path": self.path, "kwargs": self.kwargs, "executor": self.executor} + + +def _make_callback(state: CallbackState, session) -> ExecutorCallback: + cb = ExecutorCallback(callback_def=_FakeCallbackDef(), fetch_method=CallbackFetchMethod.IMPORT_PATH) + cb.state = state + session.add(cb) + session.commit() + return cb + + +def _override_require_auth(exec_app, scope: Literal["execution", "workload"] = "execution") -> None: + """Override require_auth to return a token whose sub matches the path callback_id.""" + + async def _token(request: Request) -> TIToken: + path_id = request.path_params.get("callback_id") + cb_id = UUID(path_id) if path_id else uuid4() + return TIToken(id=cb_id, claims=TIClaims(scope=scope)) + + exec_app.dependency_overrides[require_auth] = _token + + [email protected] +def _use_real_jwt_bearer(exec_app): + """Remove the mock require_auth override so the real JWT validation runs end-to-end.""" + exec_app.dependency_overrides.pop(require_auth, None) + + +class TestCallbackRun: + def setup_method(self): + clear_db_callbacks() + + def teardown_method(self): + clear_db_callbacks() + + @pytest.mark.parametrize("scope", ["workload", "execution"]) + def test_run_marks_callback_running_and_swaps_workload_token(self, client, exec_app, session, scope): + cb = _make_callback(CallbackState.QUEUED, session) + + mock_gen = mock.MagicMock(spec=JWTGenerator) + mock_gen.generate.return_value = "mock-execution-token" + lifespan.registry.register_value(JWTGenerator, mock_gen) + + _override_require_auth(exec_app, scope=scope) + + response = client.post(f"/execution/callbacks/{cb.id}/run") + + exec_app.dependency_overrides.pop(require_auth, None) + + assert response.status_code == 204 + + session.expire_all() + cb_after = session.get(ExecutorCallback, cb.id) + assert cb_after.state == CallbackState.RUNNING + + if scope == "workload": + assert response.headers["Refreshed-API-Token"] == "mock-execution-token" + mock_gen.generate.assert_called_once() + extras = mock_gen.generate.call_args.kwargs["extras"] + assert extras == {"sub": str(cb.id), "scope": "execution"} + else: + # Execution-scoped tokens skip the swap; the middleware handles refresh elsewhere. + assert "Refreshed-API-Token" not in response.headers + mock_gen.generate.assert_not_called() + + @pytest.mark.parametrize("scope", ["workload", "execution"]) + def test_run_is_idempotent_when_already_running(self, client, exec_app, session, scope): + """ + A retried ``POST /callbacks/{id}/run`` against a callback already in RUNNING + state must succeed (204) without regressing the row, and a workload-scoped + retry must still receive a fresh execution token via ``Refreshed-API-Token`` + so the supervisor can complete the workload→execution swap that may have + been lost when the original response was dropped. + """ + cb = _make_callback(CallbackState.RUNNING, session) + + mock_gen = mock.MagicMock(spec=JWTGenerator) + mock_gen.generate.return_value = "mock-execution-token" + lifespan.registry.register_value(JWTGenerator, mock_gen) + + _override_require_auth(exec_app, scope=scope) + + response = client.post(f"/execution/callbacks/{cb.id}/run") + + exec_app.dependency_overrides.pop(require_auth, None) + + assert response.status_code == 204 + + session.expire_all() + cb_after = session.get(ExecutorCallback, cb.id) + assert cb_after.state == CallbackState.RUNNING + + if scope == "workload": + assert response.headers["Refreshed-API-Token"] == "mock-execution-token" + mock_gen.generate.assert_called_once() + extras = mock_gen.generate.call_args.kwargs["extras"] + assert extras == {"sub": str(cb.id), "scope": "execution"} + else: + assert "Refreshed-API-Token" not in response.headers + mock_gen.generate.assert_not_called() + + @pytest.mark.parametrize( + "state", + [ + CallbackState.PENDING, + CallbackState.SCHEDULED, + CallbackState.SUCCESS, + CallbackState.FAILED, + ], + ) + def test_run_returns_409_for_non_runnable_state(self, client, exec_app, session, state): + cb = _make_callback(state, session) + _override_require_auth(exec_app, scope="workload") + + response = client.post(f"/execution/callbacks/{cb.id}/run") + exec_app.dependency_overrides.pop(require_auth, None) + + assert response.status_code == 409 + assert response.json()["detail"]["reason"] == "invalid_state" + assert response.json()["detail"]["current_state"] == state.value + + def test_run_returns_404_when_callback_missing(self, client, exec_app): + missing_id = uuid4() + + async def _token(request: Request) -> TIToken: + return TIToken(id=missing_id, claims=TIClaims(scope="workload")) + + exec_app.dependency_overrides[require_auth] = _token + + response = client.post(f"/execution/callbacks/{missing_id}/run") + exec_app.dependency_overrides.pop(require_auth, None) + + assert response.status_code == 404 + assert response.json()["detail"]["reason"] == "not_found" + + def test_run_rejects_mismatched_sub(self, client, exec_app, session): + cb = _make_callback(CallbackState.QUEUED, session) + + async def _token(request: Request) -> TIToken: + return TIToken(id=uuid4(), claims=TIClaims(scope="workload")) + + exec_app.dependency_overrides[require_auth] = _token + + response = client.post(f"/execution/callbacks/{cb.id}/run") + exec_app.dependency_overrides.pop(require_auth, None) + + assert response.status_code == 403 + assert response.json()["detail"] == "Token subject does not match callback id" + + +class TestCallbackUpdateState: + def setup_method(self): + clear_db_callbacks() + + def teardown_method(self): + clear_db_callbacks() + + @pytest.mark.parametrize( + ("payload_state", "expected_state"), + [ + ("success", CallbackState.SUCCESS), + ("failed", CallbackState.FAILED), + ], + ) + def test_update_state_writes_terminal_state( + self, client, exec_app, session, payload_state, expected_state + ): + cb = _make_callback(CallbackState.RUNNING, session) + _override_require_auth(exec_app, scope="execution") + + response = client.patch( + f"/execution/callbacks/{cb.id}/state", + json={"state": payload_state, "output": "an output"}, + ) + exec_app.dependency_overrides.pop(require_auth, None) + + assert response.status_code == 204 + session.expire_all() + cb_after = session.get(ExecutorCallback, cb.id) + assert cb_after.state == expected_state + assert cb_after.output == "an output" + + @pytest.mark.parametrize( + "state", + [ + CallbackState.QUEUED, + CallbackState.PENDING, + CallbackState.SCHEDULED, + CallbackState.SUCCESS, + CallbackState.FAILED, + ], + ) + def test_update_state_returns_409_when_not_running(self, client, exec_app, session, state): + cb = _make_callback(state, session) + _override_require_auth(exec_app, scope="execution") + + response = client.patch( + f"/execution/callbacks/{cb.id}/state", + json={"state": "success"}, + ) + exec_app.dependency_overrides.pop(require_auth, None) + + assert response.status_code == 409 + assert response.json()["detail"]["reason"] == "invalid_state" + + def test_update_state_returns_404_when_callback_missing(self, client, exec_app): + missing_id = uuid4() + _override_require_auth(exec_app, scope="execution") + + response = client.patch( + f"/execution/callbacks/{missing_id}/state", + json={"state": "success"}, + ) + exec_app.dependency_overrides.pop(require_auth, None) + + assert response.status_code == 404 + + def test_update_state_rejects_mismatched_sub(self, client, exec_app, session): + cb = _make_callback(CallbackState.RUNNING, session) + + async def _token(request: Request) -> TIToken: + return TIToken(id=uuid4(), claims=TIClaims(scope="execution")) Review Comment: I think we also need to rename `TIToken` -- given it's. now used for Callbacks (which is the right thing to do) calling this `TI` isn't correct. That or we need to make a new type, `CallbackToken` and then have the security decorator return `TIToken | CallbackToken` (and then something inside TI routes ensure it's of the right type) -- 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]
