From 6355f6780a51f0b80ae55c18cb3123d69706adc0 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Mon, 13 Jul 2026 22:34:57 +0000 Subject: [PATCH 01/11] Add single-use callback token for deadline callback context fetch Deadline callbacks run in a subprocess that needs to read the DagRun context (and connections/variables/xcoms) from the Execution API. PR also accept the long-lived ``workload`` token via ``token:workload`` opt-ins. That over-broadened the workload token's reach (scope creep): a long-lived token could read arbitrary DagRun/connection/variable/xcom data for the whole queue-wait lifetime, and Ash asked for a single-use credential instead. This re-lands the security core of #66608 with a tighter design: * New ``callback`` token scope. The executor mints a ``callback``-scoped token whose ``sub`` is the callback id. That token is accepted ONLY on the new ``PATCH /callbacks/{id}/run`` exchange endpoint (pinned by a ``callback:self`` scope) and nowhere else. * Single-use exchange. ``run_callback`` performs an atomic ``QUEUED -> RUNNING`` transition under ``SELECT ... FOR UPDATE`` and returns a fresh short-lived ``execution`` token via the ``Refreshed-API-Token`` header. A second presentation of the same callback token finds the row already RUNNING and gets 409, so the token is exchangeable exactly once. This mirrors the Task Instance ``/run`` pattern. The subprocess calls the exchange once, before any context read, then uses the execution token for the actual reads. * Reverted the #66608 ``token:workload`` opt-ins on the DagRun, connections, variables, and xcoms GET endpoints back to execution-only. * ``JWTReissueMiddleware`` skips ``callback`` tokens (as it already does for ``workload``), since they are long-lived and exchanged exactly once. Two shared-helper refactors keep the new code DRY and also clean up pre-existing duplication: * ``require_auth`` now drives the three ``*:self`` scope checks (``ti:self`` / ``ct:self`` / ``callback:self``) from a single ``SELF_SCOPE_PATH_PARAMS`` mapping + loop instead of three near-identical branches. Behavior for ``ti:self`` and ``ct:self`` is unchanged. * New ``issue_execution_token`` helper mints an execution-scope token and sets the ``Refreshed-API-Token`` header. It replaces the duplicated mint block in the Task Instance ``/run`` handler (still guarded by the ``workload`` scope check) and the new callback exchange endpoint. --- .../airflow/api_fastapi/execution_api/app.py | 6 +- .../execution_api/datamodels/token.py | 2 +- .../execution_api/routes/__init__.py | 2 + .../execution_api/routes/callbacks.py | 104 ++++++++++ .../execution_api/routes/task_instances.py | 6 +- .../api_fastapi/execution_api/security.py | 49 +++-- .../src/airflow/executors/workloads/base.py | 4 +- .../airflow/executors/workloads/callback.py | 2 +- .../api_fastapi/execution_api/conftest.py | 5 +- .../execution_api/test_security.py | 10 + .../versions/head/test_callbacks.py | 178 ++++++++++++++++++ task-sdk/src/airflow/sdk/api/client.py | 26 +++ .../sdk/execution_time/callback_supervisor.py | 4 + task-sdk/tests/task_sdk/api/test_client.py | 36 ++++ 14 files changed, 405 insertions(+), 29 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index f9e7cb1725000..73afc9a4f9af8 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -148,10 +148,8 @@ async def dispatch(self, request: Request, call_next): validator: JWTValidator = await services.aget(JWTValidator) claims = await validator.avalidated_claims(token, {}) - # Workload tokens are long-lived and meant to survive queue - # wait times so avoid refreshing them. If avalidated_claims - # raises for a workload token, the outer except handles it. - if claims.get("scope") == "workload": + # Workload and callback tokens are long-lived (exchanged once, not refreshed here). + if claims.get("scope") in ("workload", "callback"): return response now = int(time.time()) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/token.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/token.py index 4c3b935f5aa62..eb858fac1dfc4 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/token.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/token.py @@ -24,7 +24,7 @@ from airflow.api_fastapi.core_api.base import BaseModel -TokenScope = Literal["execution", "workload"] +TokenScope = Literal["execution", "workload", "callback"] class TIClaims(BaseModel): 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..5662378c87087 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, @@ -53,6 +54,7 @@ connection_tests.router, prefix="/connection-tests", tags=["Connection Tests"] ) authenticated_router.include_router(connections.router, prefix="/connections", tags=["Connections"]) +authenticated_router.include_router(callbacks.router, prefix="/callbacks", tags=["Callbacks"]) authenticated_router.include_router(dag_runs.router, prefix="/dag-runs", tags=["Dag Runs"]) authenticated_router.include_router(dags.router, prefix="/dags", tags=["Dags"]) authenticated_router.include_router(task_instances.router, prefix="/task-instances", tags=["Task Instances"]) 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..ce5fcde289087 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py @@ -0,0 +1,104 @@ +# 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 + +import structlog +from cadwyn import VersionedAPIRouter +from fastapi import HTTPException, Response, Security, status + +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +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, + issue_execution_token, + require_auth, +) +from airflow.models.callback import Callback +from airflow.utils.state import CallbackState + +log = structlog.get_logger(logger_name=__name__) + +# The ``callback``-scoped token is accepted only on this router; ``callback:self`` pins it +# to its own callback id and it can reach only the exchange endpoint below. +router = VersionedAPIRouter( + route_class=ExecutionAPIRoute, + dependencies=[ + Security(require_auth, scopes=["callback:self", "token:callback"]), + ], +) + + +@router.patch( + "/{callback_id}/run", + status_code=status.HTTP_204_NO_CONTENT, + responses=create_openapi_http_exception_doc( + [ + (status.HTTP_404_NOT_FOUND, "Callback not found"), + (status.HTTP_409_CONFLICT, "The callback token was already exchanged"), + ] + ), +) +def run_callback( + callback_id: UUID, + response: Response, + session: SessionDep, + services=DepContainer, + token: TIToken = CurrentTIToken, +) -> None: + """ + Exchange a single-use callback token for a short-lived execution token. + + The deadline-callback subprocess calls this once, before reading any + context. The exchange is gated on an atomic ``QUEUED -> RUNNING`` transition + under ``SELECT ... FOR UPDATE``: the first call transitions the row and + returns a fresh ``execution`` token via the ``Refreshed-API-Token`` header; + any second presentation of the same callback token finds the row already in + ``RUNNING`` (or terminal) and gets ``409``. This makes the callback token + exchangeable exactly once, mirroring the Task Instance ``/run`` pattern. + """ + 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": "invalid_state", + "message": ( + f"Callback {callback_id} is in state {callback.state}; its token can only be " + "exchanged once while QUEUED." + ), + "previous_state": str(callback.state) if callback.state is not None else None, + }, + ) + + callback.state = CallbackState.RUNNING + log.info("Callback token exchanged; marked running", callback_id=str(callback_id)) + + issue_execution_token(services, response, sub=str(callback_id)) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 02d723a82185e..44450e385641d 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -43,7 +43,6 @@ from airflow._shared.observability.traces import override_ids from airflow._shared.state import TaskScope from airflow._shared.timezones import timezone -from airflow.api_fastapi.auth.tokens import JWTGenerator from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.db.dags import eager_load_teams @@ -75,6 +74,7 @@ CurrentTIToken, ExecutionAPIRoute, get_team_name_for_ti, + issue_execution_token, require_auth, ) from airflow.api_fastapi.execution_api.services.task_instances import ( @@ -351,9 +351,7 @@ def ti_run( # JWTReissueMiddleware also writes Refreshed-API-Token but skips workload tokens, so we set it here for the workload→execution swap. if token.claims.scope == "workload": - generator: JWTGenerator = services.get(JWTGenerator) - execution_token = generator.generate(extras={"sub": str(task_instance_id), "scope": "execution"}) - response.headers["Refreshed-API-Token"] = execution_token + issue_execution_token(services, response, sub=str(task_instance_id)) return context 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..ae372e5678b50 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -70,14 +70,14 @@ from typing import Any, get_args import structlog -from fastapi import Depends, HTTPException, Request, status +from fastapi import Depends, HTTPException, Request, Response, status from fastapi.params import Security as SecurityParam from fastapi.routing import APIRoute from fastapi.security import HTTPBearer, SecurityScopes from pydantic import ValidationError from sqlalchemy import select -from airflow.api_fastapi.auth.tokens import JWTValidator +from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken, TokenScope from airflow.api_fastapi.execution_api.deps import DepContainer @@ -85,6 +85,14 @@ VALID_TOKEN_TYPES: frozenset[str] = frozenset(get_args(TokenScope)) +# ``*:self`` Security scopes pin a token to a single resource. Each maps the scope name to +# the route path parameter carrying the resource id that the JWT ``sub`` must match. +SELF_SCOPE_PATH_PARAMS: dict[str, str] = { + "ti:self": "task_instance_id", + "ct:self": "connection_test_id", + "callback:self": "callback_id", +} + _REQUEST_SCOPE_TOKEN_KEY = "ti_token" @@ -182,20 +190,15 @@ async def require_auth( f"Allowed types: {', '.join(sorted(allowed_token_types))}", ) - if "ti:self" in security_scopes.scopes: - ti_self_id = str(request.path_params["task_instance_id"]) - if str(token.id) != ti_self_id: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="Token subject does not match task instance ID", - ) - elif "ct:self" in security_scopes.scopes: - ct_self_id = str(request.path_params["connection_test_id"]) - if str(token.id) != ct_self_id: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="Token subject does not match connection test ID", - ) + # Enforce the first ``*:self`` scope present: the JWT ``sub`` must match the path param. + for scope, path_param in SELF_SCOPE_PATH_PARAMS.items(): + if scope in security_scopes.scopes: + if str(token.id) != str(request.path_params[path_param]): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Token subject does not match {path_param}", + ) + break return token @@ -203,6 +206,20 @@ async def require_auth( CurrentTIToken: TIToken = Depends(require_auth) +def issue_execution_token(services: Any, response: Response, sub: str) -> None: + """ + Mint a short-lived ``execution``-scoped token and set it on the response. + + Used by endpoints that swap a longer-lived token (``workload`` on the Task Instance + ``/run`` endpoint, ``callback`` on the callback exchange endpoint) for an ``execution`` + token: it mints via the request-scoped ``JWTGenerator`` and sets the ``Refreshed-API-Token`` + header the SDK client adopts automatically. ``JWTReissueMiddleware`` skips those longer-lived + scopes, so this is the single place the swap happens. + """ + generator: JWTGenerator = services.get(JWTGenerator) + response.headers["Refreshed-API-Token"] = generator.generate(extras={"sub": sub, "scope": "execution"}) + + class ExecutionAPIRoute(APIRoute): """ Custom route class that precomputes allowed token types from Security scopes. diff --git a/airflow-core/src/airflow/executors/workloads/base.py b/airflow-core/src/airflow/executors/workloads/base.py index 41334d68f3038..98788bc49e59b 100644 --- a/airflow-core/src/airflow/executors/workloads/base.py +++ b/airflow-core/src/airflow/executors/workloads/base.py @@ -84,12 +84,12 @@ class BaseWorkloadSchema(BaseModel): """The identity token for this workload""" @staticmethod - def generate_token(sub_id: str, generator: JWTGenerator | None = None) -> str: + def generate_token(sub_id: str, generator: JWTGenerator | None = None, scope: str = "workload") -> str: if not generator: return "" valid_for = conf.getfloat("scheduler", "task_queued_timeout") return generator.generate( - extras={"sub": sub_id, "scope": "workload"}, + extras={"sub": sub_id, "scope": scope}, valid_for=valid_for, ) diff --git a/airflow-core/src/airflow/executors/workloads/callback.py b/airflow-core/src/airflow/executors/workloads/callback.py index c1842a5900685..702e849fa407b 100644 --- a/airflow-core/src/airflow/executors/workloads/callback.py +++ b/airflow-core/src/airflow/executors/workloads/callback.py @@ -127,7 +127,7 @@ def make( return cls( callback=CallbackDTO.model_validate(callback, from_attributes=True), dag_rel_path=dag_rel_path or Path(dag_run.dag_model.relative_fileloc or ""), - token=cls.generate_token(str(callback.id), generator), + token=cls.generate_token(str(callback.id), generator, scope="callback"), log_path=fname, bundle_info=bundle_info, ) 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..7fd2fcdb4f064 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 @@ -80,6 +80,16 @@ def test_ignores_non_token_scopes(self): ) assert route.allowed_token_types == frozenset({"execution"}) + def test_extracts_callback_token_scope(self): + route = ExecutionAPIRoute( + path="/test", + endpoint=lambda: None, + dependencies=[ + Security(require_auth, scopes=["callback:self", "token:callback"]), + ], + ) + assert route.allowed_token_types == frozenset({"callback"}) + def test_rejects_invalid_token_types(self): with pytest.raises(ValueError, match="Invalid token types"): ExecutionAPIRoute( 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..d39de34649a51 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py @@ -0,0 +1,178 @@ +# 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 jwt +import pytest + +from airflow.api_fastapi.auth.tokens import JWTValidator +from airflow.api_fastapi.execution_api.app import lifespan +from airflow.api_fastapi.execution_api.security import require_auth +from airflow.models.callback import CallbackFetchMethod, CallbackState, ExecutorCallback +from airflow.sdk.definitions.callback import SyncCallback + +from tests_common.test_utils.db import clear_db_callbacks + +pytestmark = pytest.mark.db_test + + +def sync_callback(): + """Empty (sync) callable used for unit tests.""" + + +def _make_queued_callback(session, dag_id="test_dag"): + """Create and persist an ExecutorCallback in QUEUED state.""" + callback = ExecutorCallback( + callback_def=SyncCallback(sync_callback, kwargs={}), + fetch_method=CallbackFetchMethod.IMPORT_PATH, + ) + callback.data["dag_id"] = dag_id + callback.state = CallbackState.QUEUED + session.add(callback) + session.commit() + return callback + + +class TestRunCallback: + @pytest.fixture(autouse=True) + def setup_teardown(self): + clear_db_callbacks() + yield + clear_db_callbacks() + + def test_run_marks_running_and_swaps_token(self, client, session): + """First call transitions QUEUED -> RUNNING and returns a fresh execution token.""" + callback = _make_queued_callback(session) + + response = client.patch(f"/execution/callbacks/{callback.id}/run") + + assert response.status_code == 204 + + # A real execution-scope JWT is minted with the callback id as its subject. + refreshed = response.headers["Refreshed-API-Token"] + payload = jwt.decode(refreshed, options={"verify_signature": False}) + assert payload["scope"] == "execution" + assert payload["sub"] == str(callback.id) + + session.expire_all() + session.refresh(callback) + assert callback.state == CallbackState.RUNNING + + def test_run_is_single_use(self, client, session): + """A second exchange of the same callback token is rejected with 409.""" + callback = _make_queued_callback(session) + + first = client.patch(f"/execution/callbacks/{callback.id}/run") + assert first.status_code == 204 + + second = client.patch(f"/execution/callbacks/{callback.id}/run") + assert second.status_code == 409 + detail = second.json()["detail"] + assert detail["reason"] == "invalid_state" + assert detail["previous_state"] == CallbackState.RUNNING + + @pytest.mark.parametrize( + "state", + [CallbackState.RUNNING, CallbackState.SUCCESS, CallbackState.FAILED, CallbackState.PENDING], + ) + def test_run_rejects_non_queued_state(self, client, session, state): + """The token can only be exchanged while the callback is QUEUED.""" + callback = _make_queued_callback(session) + callback.state = state + session.commit() + + response = client.patch(f"/execution/callbacks/{callback.id}/run") + assert response.status_code == 409 + assert response.json()["detail"]["reason"] == "invalid_state" + + def test_run_returns_404_for_nonexistent(self, client): + """Exchanging a token for an unknown callback returns 404.""" + response = client.patch("/execution/callbacks/00000000-0000-0000-0000-000000000000/run") + assert response.status_code == 404 + assert response.json()["detail"]["reason"] == "not_found" + + +@pytest.fixture +def _use_real_jwt_bearer(exec_app): + """Remove the mock require_auth override so the real scope validation runs.""" + exec_app.dependency_overrides.pop(require_auth, None) + + +@pytest.mark.usefixtures("_use_real_jwt_bearer") +def test_run_accepts_callback_token_matching_sub(client, session): + """The exchange endpoint accepts a callback-scope JWT whose sub is the callback id.""" + clear_db_callbacks() + callback = _make_queued_callback(session) + + validator = mock.AsyncMock(spec=JWTValidator) + validator.avalidated_claims.return_value = { + "sub": str(callback.id), + "scope": "callback", + "exp": 9999999999, + "iat": 1000000000, + "nbf": 1000000000, + } + lifespan.registry.register_value(JWTValidator, validator) + + resp = client.patch(f"/execution/callbacks/{callback.id}/run") + assert resp.status_code == 204, resp.json() + clear_db_callbacks() + + +@pytest.mark.usefixtures("_use_real_jwt_bearer") +def test_run_rejects_callback_token_for_other_callback(client, session): + """callback:self blocks a callback token whose sub is a different callback id.""" + clear_db_callbacks() + callback = _make_queued_callback(session) + + validator = mock.AsyncMock(spec=JWTValidator) + validator.avalidated_claims.return_value = { + # sub belongs to a different callback than the one in the URL. + "sub": "11111111-1111-1111-1111-111111111111", + "scope": "callback", + "exp": 9999999999, + "iat": 1000000000, + "nbf": 1000000000, + } + lifespan.registry.register_value(JWTValidator, validator) + + resp = client.patch(f"/execution/callbacks/{callback.id}/run") + assert resp.status_code == 403, resp.json() + clear_db_callbacks() + + +@pytest.mark.usefixtures("_use_real_jwt_bearer") +def test_run_rejects_execution_scope_token(client, session): + """Only callback-scope tokens are accepted on the exchange endpoint.""" + clear_db_callbacks() + callback = _make_queued_callback(session) + + validator = mock.AsyncMock(spec=JWTValidator) + validator.avalidated_claims.return_value = { + "sub": str(callback.id), + "scope": "execution", + "exp": 9999999999, + "iat": 1000000000, + "nbf": 1000000000, + } + lifespan.registry.register_value(JWTValidator, validator) + + resp = client.patch(f"/execution/callbacks/{callback.id}/run") + assert resp.status_code == 403, resp.json() + clear_db_callbacks() diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index a0c0a30de088f..a74f5fc78a09b 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -1097,6 +1097,26 @@ def update_state( self.client.patch(f"connection-tests/{id}", content=body.model_dump_json()) +class CallbackOperations: + __slots__ = ("client",) + + def __init__(self, client: Client): + self.client = client + + def run(self, callback_id: uuid.UUID) -> None: + """ + Exchange the single-use callback token for a short-lived execution token. + + Marks the callback RUNNING on the server and swaps this client's + ``callback``-scoped token for an ``execution`` token via the + ``Refreshed-API-Token`` response header (adopted automatically by + ``_update_auth``). Must be called once, before any context read; a + second call raises ``ServerResponseError`` (409) because the callback + is no longer QUEUED. + """ + self.client.patch(f"callbacks/{callback_id}/run") + + class BearerAuth(httpx.Auth): def __init__(self, token: str): self.token: str = token @@ -1304,6 +1324,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/execution_time/callback_supervisor.py b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py index 9830771701293..da9802d086861 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -433,6 +433,10 @@ def supervise_callback( else: logger = structlog.get_logger(logger_name="callback").bind() + # Swap the single-use callback token for an execution token before any context read. + # A duplicate delivery finds the callback already RUNNING and raises, so it runs at most once. + client.callbacks.run(UUID(id)) + try: process = CallbackSubprocess.start( id=id, diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 8527349f56b26..5d84240eb9b53 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -2254,3 +2254,39 @@ def handle_request(request: httpx.Request) -> httpx.Response: client = make_client(transport=httpx.MockTransport(handle_request)) result = client.asset_state_store.clear(uri="s3://bucket/key") assert result == OKResponse(ok=True) + + +class TestCallbackOperations: + def test_run_exchanges_token(self): + """run() PATCHes the exchange endpoint and adopts the refreshed execution token.""" + callback_id = uuid7() + + def handle_request(request: httpx.Request) -> httpx.Response: + assert request.method == "PATCH" + assert request.url.path == f"/callbacks/{callback_id}/run" + return httpx.Response( + status_code=204, + headers={"Refreshed-API-Token": "execution-token"}, + ) + + client = make_client(transport=httpx.MockTransport(handle_request)) + client.callbacks.run(callback_id) + + # The single-use callback token has been swapped for the execution token. + assert client.auth is not None + assert client.auth.token == "execution-token" + + def test_run_raises_on_conflict(self): + """A second exchange (callback no longer QUEUED) surfaces as ServerResponseError.""" + callback_id = uuid7() + + def handle_request(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=409, + json={"detail": {"reason": "invalid_state", "previous_state": "running"}}, + ) + + client = make_client(transport=httpx.MockTransport(handle_request)) + with pytest.raises(ServerResponseError) as err: + client.callbacks.run(callback_id) + assert err.value.response.status_code == 409 From ebbeeb82ea649f166a89dfda101d5a0780f6269d Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Wed, 15 Jul 2026 00:09:37 +0000 Subject: [PATCH 02/11] Declare callback /run route in NON_DEFAULT_TOKEN_POLICY The new PATCH /callbacks/{callback_id}/run route accepts the single-use 'callback' token, a non-default (non-execution) scope. The token scope boundary regression guard requires such routes to be declared explicitly in NON_DEFAULT_TOKEN_POLICY; add it there per the test's documented maintenance step. --- .../api_fastapi/execution_api/test_token_scope_boundaries.py | 3 +++ 1 file changed, 3 insertions(+) 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..e15e53ff1cab0 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,9 @@ # 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"}, + # The callback /run endpoint exchanges a single-use callback token for a + # short-lived execution token (callback-only). + "PATCH /callbacks/{callback_id}/run": {"callback"}, } From c783947975dcf76f4c04f61ce734c2600e195e79 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Fri, 24 Jul 2026 23:40:26 +0000 Subject: [PATCH 03/11] Trim over-explained comments --- .../api_fastapi/execution_api/routes/callbacks.py | 11 ++--------- .../airflow/api_fastapi/execution_api/security.py | 13 ++++--------- .../execution_api/test_token_scope_boundaries.py | 3 +-- .../execution_api/versions/head/test_callbacks.py | 2 -- task-sdk/src/airflow/sdk/api/client.py | 12 ++++-------- .../sdk/execution_time/callback_supervisor.py | 3 +-- task-sdk/tests/task_sdk/api/test_client.py | 1 - 7 files changed, 12 insertions(+), 33 deletions(-) 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 index ce5fcde289087..031ef87af75cc 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py @@ -37,8 +37,6 @@ log = structlog.get_logger(logger_name=__name__) -# The ``callback``-scoped token is accepted only on this router; ``callback:self`` pins it -# to its own callback id and it can reach only the exchange endpoint below. router = VersionedAPIRouter( route_class=ExecutionAPIRoute, dependencies=[ @@ -67,13 +65,8 @@ def run_callback( """ Exchange a single-use callback token for a short-lived execution token. - The deadline-callback subprocess calls this once, before reading any - context. The exchange is gated on an atomic ``QUEUED -> RUNNING`` transition - under ``SELECT ... FOR UPDATE``: the first call transitions the row and - returns a fresh ``execution`` token via the ``Refreshed-API-Token`` header; - any second presentation of the same callback token finds the row already in - ``RUNNING`` (or terminal) and gets ``409``. This makes the callback token - exchangeable exactly once, mirroring the Task Instance ``/run`` pattern. + Gated on an atomic ``QUEUED -> RUNNING`` transition, so the exchange + succeeds at most once; the new token is returned via ``Refreshed-API-Token``. """ callback = session.get(Callback, callback_id, with_for_update=True) if callback is None: 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 ae372e5678b50..b7bb6d9e153ab 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -85,8 +85,7 @@ VALID_TOKEN_TYPES: frozenset[str] = frozenset(get_args(TokenScope)) -# ``*:self`` Security scopes pin a token to a single resource. Each maps the scope name to -# the route path parameter carrying the resource id that the JWT ``sub`` must match. +# Maps each ``*:self`` scope to the path param whose value the JWT ``sub`` must match. SELF_SCOPE_PATH_PARAMS: dict[str, str] = { "ti:self": "task_instance_id", "ct:self": "connection_test_id", @@ -190,7 +189,6 @@ async def require_auth( f"Allowed types: {', '.join(sorted(allowed_token_types))}", ) - # Enforce the first ``*:self`` scope present: the JWT ``sub`` must match the path param. for scope, path_param in SELF_SCOPE_PATH_PARAMS.items(): if scope in security_scopes.scopes: if str(token.id) != str(request.path_params[path_param]): @@ -208,13 +206,10 @@ async def require_auth( def issue_execution_token(services: Any, response: Response, sub: str) -> None: """ - Mint a short-lived ``execution``-scoped token and set it on the response. + Mint an ``execution``-scoped token and set it on the ``Refreshed-API-Token`` header. - Used by endpoints that swap a longer-lived token (``workload`` on the Task Instance - ``/run`` endpoint, ``callback`` on the callback exchange endpoint) for an ``execution`` - token: it mints via the request-scoped ``JWTGenerator`` and sets the ``Refreshed-API-Token`` - header the SDK client adopts automatically. ``JWTReissueMiddleware`` skips those longer-lived - scopes, so this is the single place the swap happens. + ``JWTReissueMiddleware`` skips ``workload``/``callback`` tokens, so endpoints that + swap them for an execution token must set the header here themselves. """ generator: JWTGenerator = services.get(JWTGenerator) response.headers["Refreshed-API-Token"] = generator.generate(extras={"sub": sub, "scope": "execution"}) 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 e15e53ff1cab0..195503f2bf080 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,8 +49,7 @@ # 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"}, - # The callback /run endpoint exchanges a single-use callback token for a - # short-lived execution token (callback-only). + # Callback /run exchanges a single-use callback token for an execution token. "PATCH /callbacks/{callback_id}/run": {"callback"}, } 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 index d39de34649a51..7c95df4641a19 100644 --- 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 @@ -64,7 +64,6 @@ def test_run_marks_running_and_swaps_token(self, client, session): assert response.status_code == 204 - # A real execution-scope JWT is minted with the callback id as its subject. refreshed = response.headers["Refreshed-API-Token"] payload = jwt.decode(refreshed, options={"verify_signature": False}) assert payload["scope"] == "execution" @@ -143,7 +142,6 @@ def test_run_rejects_callback_token_for_other_callback(client, session): validator = mock.AsyncMock(spec=JWTValidator) validator.avalidated_claims.return_value = { - # sub belongs to a different callback than the one in the URL. "sub": "11111111-1111-1111-1111-111111111111", "scope": "callback", "exp": 9999999999, diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index a74f5fc78a09b..1660f9a5a716c 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -1105,14 +1105,10 @@ def __init__(self, client: Client): def run(self, callback_id: uuid.UUID) -> None: """ - Exchange the single-use callback token for a short-lived execution token. - - Marks the callback RUNNING on the server and swaps this client's - ``callback``-scoped token for an ``execution`` token via the - ``Refreshed-API-Token`` response header (adopted automatically by - ``_update_auth``). Must be called once, before any context read; a - second call raises ``ServerResponseError`` (409) because the callback - is no longer QUEUED. + Exchange the single-use callback token for an execution token. + + Must be called before any context read; a second call raises + ``ServerResponseError`` (409). """ self.client.patch(f"callbacks/{callback_id}/run") 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 da9802d086861..0b51a5a31580a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -433,8 +433,7 @@ def supervise_callback( else: logger = structlog.get_logger(logger_name="callback").bind() - # Swap the single-use callback token for an execution token before any context read. - # A duplicate delivery finds the callback already RUNNING and raises, so it runs at most once. + # Exchange the single-use callback token before any context read; raises on duplicate delivery. client.callbacks.run(UUID(id)) try: diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 5d84240eb9b53..6176a353ae083 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -2272,7 +2272,6 @@ def handle_request(request: httpx.Request) -> httpx.Response: client = make_client(transport=httpx.MockTransport(handle_request)) client.callbacks.run(callback_id) - # The single-use callback token has been swapped for the execution token. assert client.auth is not None assert client.auth.token == "execution-token" From 50e9d8ffd910c5f3876d558534ece3e5d5a80a50 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Sat, 25 Jul 2026 00:32:45 +0000 Subject: [PATCH 04/11] Match surrounding code style --- .../airflow/api_fastapi/execution_api/app.py | 4 +- .../execution_api/routes/callbacks.py | 16 +------ .../api_fastapi/execution_api/security.py | 43 ++++++++++--------- .../versions/head/test_callbacks.py | 6 +-- task-sdk/src/airflow/sdk/api/client.py | 7 +-- .../sdk/execution_time/callback_supervisor.py | 2 +- task-sdk/tests/task_sdk/api/test_client.py | 2 - 7 files changed, 32 insertions(+), 48 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index 73afc9a4f9af8..53be7e5af6769 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -148,7 +148,9 @@ async def dispatch(self, request: Request, call_next): validator: JWTValidator = await services.aget(JWTValidator) claims = await validator.avalidated_claims(token, {}) - # Workload and callback tokens are long-lived (exchanged once, not refreshed here). + # Workload and callback tokens are long-lived and meant to survive + # queue wait times so avoid refreshing them. If avalidated_claims + # raises for such a token, the outer except handles it. if claims.get("scope") in ("workload", "callback"): return response 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 index 031ef87af75cc..b152e98ce9541 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py @@ -18,16 +18,13 @@ from uuid import UUID -import structlog from cadwyn import VersionedAPIRouter from fastapi import HTTPException, Response, Security, status from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -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, issue_execution_token, require_auth, @@ -35,8 +32,6 @@ from airflow.models.callback import Callback from airflow.utils.state import CallbackState -log = structlog.get_logger(logger_name=__name__) - router = VersionedAPIRouter( route_class=ExecutionAPIRoute, dependencies=[ @@ -60,14 +55,8 @@ def run_callback( response: Response, session: SessionDep, services=DepContainer, - token: TIToken = CurrentTIToken, ) -> None: - """ - Exchange a single-use callback token for a short-lived execution token. - - Gated on an atomic ``QUEUED -> RUNNING`` transition, so the exchange - succeeds at most once; the new token is returned via ``Refreshed-API-Token``. - """ + """Exchange a single-use callback token for a short-lived execution token.""" callback = session.get(Callback, callback_id, with_for_update=True) if callback is None: raise HTTPException( @@ -87,11 +76,10 @@ def run_callback( f"Callback {callback_id} is in state {callback.state}; its token can only be " "exchanged once while QUEUED." ), - "previous_state": str(callback.state) if callback.state is not None else None, + "previous_state": callback.state, }, ) callback.state = CallbackState.RUNNING - log.info("Callback token exchanged; marked running", callback_id=str(callback_id)) issue_execution_token(services, response, sub=str(callback_id)) 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 b7bb6d9e153ab..dbc3d955c2242 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -85,13 +85,6 @@ VALID_TOKEN_TYPES: frozenset[str] = frozenset(get_args(TokenScope)) -# Maps each ``*:self`` scope to the path param whose value the JWT ``sub`` must match. -SELF_SCOPE_PATH_PARAMS: dict[str, str] = { - "ti:self": "task_instance_id", - "ct:self": "connection_test_id", - "callback:self": "callback_id", -} - _REQUEST_SCOPE_TOKEN_KEY = "ti_token" @@ -189,14 +182,27 @@ async def require_auth( f"Allowed types: {', '.join(sorted(allowed_token_types))}", ) - for scope, path_param in SELF_SCOPE_PATH_PARAMS.items(): - if scope in security_scopes.scopes: - if str(token.id) != str(request.path_params[path_param]): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"Token subject does not match {path_param}", - ) - break + if "ti:self" in security_scopes.scopes: + ti_self_id = str(request.path_params["task_instance_id"]) + if str(token.id) != ti_self_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Token subject does not match task instance ID", + ) + elif "ct:self" in security_scopes.scopes: + ct_self_id = str(request.path_params["connection_test_id"]) + if str(token.id) != ct_self_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Token subject does not match connection test ID", + ) + elif "callback: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 @@ -205,12 +211,7 @@ async def require_auth( def issue_execution_token(services: Any, response: Response, sub: str) -> None: - """ - Mint an ``execution``-scoped token and set it on the ``Refreshed-API-Token`` header. - - ``JWTReissueMiddleware`` skips ``workload``/``callback`` tokens, so endpoints that - swap them for an execution token must set the header here themselves. - """ + """Mint an ``execution``-scoped token and set it on the ``Refreshed-API-Token`` header.""" generator: JWTGenerator = services.get(JWTGenerator) response.headers["Refreshed-API-Token"] = generator.generate(extras={"sub": sub, "scope": "execution"}) 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 index 7c95df4641a19..6646b165104bb 100644 --- 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 @@ -33,11 +33,11 @@ def sync_callback(): - """Empty (sync) callable used for unit tests.""" + """Empty (sync) callable used for unit tests""" + pass def _make_queued_callback(session, dag_id="test_dag"): - """Create and persist an ExecutorCallback in QUEUED state.""" callback = ExecutorCallback( callback_def=SyncCallback(sync_callback, kwargs={}), fetch_method=CallbackFetchMethod.IMPORT_PATH, @@ -109,7 +109,7 @@ def test_run_returns_404_for_nonexistent(self, client): @pytest.fixture def _use_real_jwt_bearer(exec_app): - """Remove the mock require_auth override so the real scope validation runs.""" + """Remove the mock require_auth override so the real JWT validation runs end-to-end.""" exec_app.dependency_overrides.pop(require_auth, None) diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 1660f9a5a716c..ea220c68fde71 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -1104,12 +1104,7 @@ def __init__(self, client: Client): self.client = client def run(self, callback_id: uuid.UUID) -> None: - """ - Exchange the single-use callback token for an execution token. - - Must be called before any context read; a second call raises - ``ServerResponseError`` (409). - """ + """Exchange the single-use callback token for an execution token.""" self.client.patch(f"callbacks/{callback_id}/run") 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 0b51a5a31580a..6b7650b9e4c63 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -433,7 +433,7 @@ def supervise_callback( else: logger = structlog.get_logger(logger_name="callback").bind() - # Exchange the single-use callback token before any context read; raises on duplicate delivery. + # Swap the single-use callback token for an execution token before any context read. client.callbacks.run(UUID(id)) try: diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 6176a353ae083..93e71943dbe47 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -2258,7 +2258,6 @@ def handle_request(request: httpx.Request) -> httpx.Response: class TestCallbackOperations: def test_run_exchanges_token(self): - """run() PATCHes the exchange endpoint and adopts the refreshed execution token.""" callback_id = uuid7() def handle_request(request: httpx.Request) -> httpx.Response: @@ -2276,7 +2275,6 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert client.auth.token == "execution-token" def test_run_raises_on_conflict(self): - """A second exchange (callback no longer QUEUED) surfaces as ServerResponseError.""" callback_id = uuid7() def handle_request(request: httpx.Request) -> httpx.Response: From 352502afebe0e0c64a970fd69eb4fe6b046f1783 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Wed, 29 Jul 2026 05:45:58 +0000 Subject: [PATCH 05/11] Re-trigger CI From 9f18d72320eeb85d7e21022517481e616c33cf08 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli <58916776+seanghaeli@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:08:35 -0700 Subject: [PATCH 06/11] Update airflow-core/src/airflow/api_fastapi/execution_api/security.py Co-authored-by: Ash Berlin-Taylor --- airflow-core/src/airflow/api_fastapi/execution_api/security.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 dbc3d955c2242..c11f4b58e8ee6 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -210,7 +210,7 @@ async def require_auth( CurrentTIToken: TIToken = Depends(require_auth) -def issue_execution_token(services: Any, response: Response, sub: str) -> None: +def issue_execution_token(services: svcs.Container, response: Response, sub: str) -> None: """Mint an ``execution``-scoped token and set it on the ``Refreshed-API-Token`` header.""" generator: JWTGenerator = services.get(JWTGenerator) response.headers["Refreshed-API-Token"] = generator.generate(extras={"sub": sub, "scope": "execution"}) From ce5f1780e954a4847bf59a6d3f2b6e444374f531 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Tue, 4 Aug 2026 21:54:08 +0000 Subject: [PATCH 07/11] Add missing svcs import to execution_api security The applied suggestion referenced svcs.Container without importing it; this file disables deferred annotations (I002), so import at runtime. --- airflow-core/src/airflow/api_fastapi/execution_api/security.py | 1 + 1 file changed, 1 insertion(+) 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 c11f4b58e8ee6..8ad4d2299ab8b 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -70,6 +70,7 @@ from typing import Any, get_args import structlog +import svcs from fastapi import Depends, HTTPException, Request, Response, status from fastapi.params import Security as SecurityParam from fastapi.routing import APIRoute From 81556254baf054e5d2cbc6667b587cfbfb272ce3 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Tue, 4 Aug 2026 22:05:38 +0000 Subject: [PATCH 08/11] Make workload token scope a class-level invariant Review feedback from ashb: the scope claim is a property of the workload type, not something callers should supply per call. --- .../src/airflow/executors/workloads/base.py | 11 +++++++---- .../airflow/executors/workloads/callback.py | 6 ++++-- .../tests/unit/executors/test_workloads.py | 19 +++++++++++++++++++ 3 files changed, 30 insertions(+), 6 deletions(-) diff --git a/airflow-core/src/airflow/executors/workloads/base.py b/airflow-core/src/airflow/executors/workloads/base.py index 98788bc49e59b..43e47065b664a 100644 --- a/airflow-core/src/airflow/executors/workloads/base.py +++ b/airflow-core/src/airflow/executors/workloads/base.py @@ -21,7 +21,7 @@ import os from abc import ABC, abstractmethod from collections.abc import Hashable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar from pydantic import BaseModel, ConfigDict, Field @@ -83,13 +83,16 @@ class BaseWorkloadSchema(BaseModel): token: str = Field(repr=False) """The identity token for this workload""" - @staticmethod - def generate_token(sub_id: str, generator: JWTGenerator | None = None, scope: str = "workload") -> str: + token_scope: ClassVar[str] = "workload" + """Scope claim stamped into tokens minted for this workload type.""" + + @classmethod + def generate_token(cls, sub_id: str, generator: JWTGenerator | None = None) -> str: if not generator: return "" valid_for = conf.getfloat("scheduler", "task_queued_timeout") return generator.generate( - extras={"sub": sub_id, "scope": scope}, + extras={"sub": sub_id, "scope": cls.token_scope}, valid_for=valid_for, ) diff --git a/airflow-core/src/airflow/executors/workloads/callback.py b/airflow-core/src/airflow/executors/workloads/callback.py index 702e849fa407b..29d6234123ab7 100644 --- a/airflow-core/src/airflow/executors/workloads/callback.py +++ b/airflow-core/src/airflow/executors/workloads/callback.py @@ -20,7 +20,7 @@ from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, ClassVar, Literal from uuid import UUID import structlog @@ -77,6 +77,8 @@ class ExecuteCallback(BaseDagBundleWorkload): type: Literal["ExecuteCallback"] = Field(init=False, default="ExecuteCallback") + token_scope: ClassVar[str] = "callback" + @property def key(self) -> CallbackKey: """Return the callback key for this workload.""" @@ -127,7 +129,7 @@ def make( return cls( callback=CallbackDTO.model_validate(callback, from_attributes=True), dag_rel_path=dag_rel_path or Path(dag_run.dag_model.relative_fileloc or ""), - token=cls.generate_token(str(callback.id), generator, scope="callback"), + token=cls.generate_token(str(callback.id), generator), log_path=fname, bundle_info=bundle_info, ) diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py index 37fbcd96ce950..7c639c43594cd 100644 --- a/airflow-core/tests/unit/executors/test_workloads.py +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -90,6 +90,25 @@ def test_generate_token_without_generator(): assert BaseWorkloadSchema.generate_token("ti-123", None) == "" +def test_token_scope_is_a_class_level_invariant(): + """Token scope is fixed per workload type, not caller-suppliable.""" + assert BaseWorkloadSchema.token_scope == "workload" + assert ExecuteTask.token_scope == "workload" + assert ExecuteCallback.token_scope == "callback" + + +def test_generate_token_uses_subclass_token_scope(monkeypatch): + """ExecuteCallback.generate_token should stamp its own 'callback' scope.""" + monkeypatch.setattr(workloads_base.conf, "getfloat", lambda section, key: 86400.0) + + generator = JWTGenerator(secret_key="test-secret", audience="test", valid_for=60) + token = ExecuteCallback.generate_token("cb-123", generator) + + claims = jwt.decode(token, "test-secret", algorithms=["HS512"], audience="test") + assert claims["sub"] == "cb-123" + assert claims["scope"] == "callback" + + def test_callback_key_is_frozen_and_hashable(): """CallbackKey must be usable as a dict key (hashable) and immutable (frozen).""" cid = "some-uuid-value" From 5981ac06278b42e603f7e7c9a99baf16b8a44dce Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Wed, 5 Aug 2026 17:09:29 +0000 Subject: [PATCH 09/11] Re-trigger CI / re-associate PR head From 55b0e20d86e6d20d7e3f09e2dc7c6c6b609bc150 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Wed, 12 Aug 2026 19:41:28 +0000 Subject: [PATCH 10/11] Version the callback claim endpoint in the Execution API --- .../newsfragments/71192.significant.rst | 30 ++++++++ .../execution_api/versions/__init__.py | 5 ++ .../execution_api/versions/v2026_08_01.py | 30 ++++++++ .../versions/v2026_08_01/__init__.py | 16 ++++ .../versions/v2026_08_01/test_callbacks.py | 76 +++++++++++++++++++ 5 files changed, 157 insertions(+) create mode 100644 airflow-core/newsfragments/71192.significant.rst create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_08_01.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/__init__.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/test_callbacks.py diff --git a/airflow-core/newsfragments/71192.significant.rst b/airflow-core/newsfragments/71192.significant.rst new file mode 100644 index 0000000000000..bf375a4dc03fb --- /dev/null +++ b/airflow-core/newsfragments/71192.significant.rst @@ -0,0 +1,30 @@ +Callbacks exchange a single-use token through the Execution API before they run + +A worker now exchanges a callback's single-use token through the Execution API +(``PATCH /execution/callbacks/{callback_id}/run``) before importing and running +it, transitioning the callback ``QUEUED -> RUNNING`` and swapping the callback +token for a short-lived ``execution``-scoped token. Previously a callback was +executed without the worker making any authenticated call first, so the 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. +- The callback token is single-use: a redelivered or replayed 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/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index d56ec735c8f13..01621782a70aa 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,11 +51,16 @@ AddTeamNameField, AddVariableKeysEndpoint, ) +from airflow.api_fastapi.execution_api.versions.v2026_08_01 import AddCallbackRunEndpoint from airflow.api_fastapi.execution_api.versions.v2026_10_30 import AddArgBindingsToTIRunContext bundle = VersionBundle( HeadVersion(), Version("2026-10-30", AddArgBindingsToTIRunContext), + 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..c27693cc196f0 --- /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 a worker uses to exchange its single-use callback token.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + endpoint("/callbacks/{callback_id}/run", ["PATCH"]).didnt_exist, + ) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/__init__.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/test_callbacks.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/test_callbacks.py new file mode 100644 index 0000000000000..4f0325477dc4f --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/test_callbacks.py @@ -0,0 +1,76 @@ +# 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 CallbackFetchMethod, CallbackState, ExecutorCallback +from airflow.sdk.definitions.callback import SyncCallback + +from tests_common.test_utils.db import clear_db_callbacks + +pytestmark = pytest.mark.db_test + + +def sync_callback(): + """Empty (sync) callable used for the version-gating tests.""" + pass + + +def _make_queued_callback(session, dag_id="test_dag"): + callback = ExecutorCallback( + callback_def=SyncCallback(sync_callback, kwargs={}), + fetch_method=CallbackFetchMethod.IMPORT_PATH, + ) + callback.data["dag_id"] = dag_id + callback.state = CallbackState.QUEUED + session.add(callback) + session.commit() + return callback + + +@pytest.fixture +def old_ver_client(client): + """Client configured to use the API version before the callback run endpoint was added.""" + client.headers["Airflow-API-Version"] = "2026-06-30" + return client + + +class TestRunCallbackEndpointVersioning: + """The callbacks/{callback_id}/run endpoint didn't exist before the 2026-08-01 API version.""" + + @pytest.fixture(autouse=True) + def setup_teardown(self): + clear_db_callbacks() + yield + clear_db_callbacks() + + def test_old_version_returns_404(self, old_ver_client, session): + """PATCH /callbacks/{callback_id}/run should not exist in older API versions.""" + callback = _make_queued_callback(session) + + response = old_ver_client.patch(f"/execution/callbacks/{callback.id}/run") + + assert response.status_code == 404 + + def test_head_version_works(self, client, session): + """PATCH /callbacks/{callback_id}/run should work in the current API version.""" + callback = _make_queued_callback(session) + + response = client.patch(f"/execution/callbacks/{callback.id}/run") + + assert response.status_code == 204 From d1063ae2bd1cb3eec95ed261695ec8a299de58f0 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Thu, 13 Aug 2026 22:09:19 +0000 Subject: [PATCH 11/11] Tighten callback token tests --- .../versions/head/test_callbacks.py | 129 ++++++------------ .../versions/v2026_08_01/test_callbacks.py | 55 ++------ 2 files changed, 53 insertions(+), 131 deletions(-) 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 index 6646b165104bb..da24c5c5f9f0c 100644 --- 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 @@ -37,12 +37,20 @@ def sync_callback(): pass -def _make_queued_callback(session, dag_id="test_dag"): +@pytest.fixture(autouse=True) +def clean_callbacks(): + clear_db_callbacks() + yield + clear_db_callbacks() + + +@pytest.fixture +def queued_callback(session): callback = ExecutorCallback( callback_def=SyncCallback(sync_callback, kwargs={}), fetch_method=CallbackFetchMethod.IMPORT_PATH, ) - callback.data["dag_id"] = dag_id + callback.data["dag_id"] = "test_dag" callback.state = CallbackState.QUEUED session.add(callback) session.commit() @@ -50,53 +58,32 @@ def _make_queued_callback(session, dag_id="test_dag"): class TestRunCallback: - @pytest.fixture(autouse=True) - def setup_teardown(self): - clear_db_callbacks() - yield - clear_db_callbacks() - - def test_run_marks_running_and_swaps_token(self, client, session): - """First call transitions QUEUED -> RUNNING and returns a fresh execution token.""" - callback = _make_queued_callback(session) - - response = client.patch(f"/execution/callbacks/{callback.id}/run") + def test_run_is_single_use(self, client, session, queued_callback): + """First call moves QUEUED -> RUNNING and swaps the token; a replay is rejected with 409.""" + response = client.patch(f"/execution/callbacks/{queued_callback.id}/run") assert response.status_code == 204 - - refreshed = response.headers["Refreshed-API-Token"] - payload = jwt.decode(refreshed, options={"verify_signature": False}) + payload = jwt.decode(response.headers["Refreshed-API-Token"], options={"verify_signature": False}) assert payload["scope"] == "execution" - assert payload["sub"] == str(callback.id) + assert payload["sub"] == str(queued_callback.id) session.expire_all() - session.refresh(callback) - assert callback.state == CallbackState.RUNNING - - def test_run_is_single_use(self, client, session): - """A second exchange of the same callback token is rejected with 409.""" - callback = _make_queued_callback(session) + session.refresh(queued_callback) + assert queued_callback.state == CallbackState.RUNNING - first = client.patch(f"/execution/callbacks/{callback.id}/run") - assert first.status_code == 204 - - second = client.patch(f"/execution/callbacks/{callback.id}/run") + second = client.patch(f"/execution/callbacks/{queued_callback.id}/run") assert second.status_code == 409 detail = second.json()["detail"] assert detail["reason"] == "invalid_state" assert detail["previous_state"] == CallbackState.RUNNING - @pytest.mark.parametrize( - "state", - [CallbackState.RUNNING, CallbackState.SUCCESS, CallbackState.FAILED, CallbackState.PENDING], - ) - def test_run_rejects_non_queued_state(self, client, session, state): + @pytest.mark.parametrize("state", [CallbackState.SUCCESS, CallbackState.FAILED, CallbackState.PENDING]) + def test_run_rejects_non_queued_state(self, client, session, queued_callback, state): """The token can only be exchanged while the callback is QUEUED.""" - callback = _make_queued_callback(session) - callback.state = state + queued_callback.state = state session.commit() - response = client.patch(f"/execution/callbacks/{callback.id}/run") + response = client.patch(f"/execution/callbacks/{queued_callback.id}/run") assert response.status_code == 409 assert response.json()["detail"]["reason"] == "invalid_state" @@ -114,63 +101,31 @@ def _use_real_jwt_bearer(exec_app): @pytest.mark.usefixtures("_use_real_jwt_bearer") -def test_run_accepts_callback_token_matching_sub(client, session): - """The exchange endpoint accepts a callback-scope JWT whose sub is the callback id.""" - clear_db_callbacks() - callback = _make_queued_callback(session) - - validator = mock.AsyncMock(spec=JWTValidator) - validator.avalidated_claims.return_value = { - "sub": str(callback.id), - "scope": "callback", - "exp": 9999999999, - "iat": 1000000000, - "nbf": 1000000000, - } - lifespan.registry.register_value(JWTValidator, validator) - - resp = client.patch(f"/execution/callbacks/{callback.id}/run") - assert resp.status_code == 204, resp.json() - clear_db_callbacks() - - -@pytest.mark.usefixtures("_use_real_jwt_bearer") -def test_run_rejects_callback_token_for_other_callback(client, session): - """callback:self blocks a callback token whose sub is a different callback id.""" - clear_db_callbacks() - callback = _make_queued_callback(session) - +@pytest.mark.parametrize( + ("token_sub", "token_scope", "expected_status"), + [ + pytest.param("self", "callback", 204, id="callback-token-matching-sub-accepted"), + pytest.param( + "11111111-1111-1111-1111-111111111111", + "callback", + 403, + id="callback-token-for-other-callback-rejected", + ), + pytest.param("self", "execution", 403, id="execution-scope-token-rejected"), + ], +) +def test_run_validates_token_scope_and_sub(client, queued_callback, token_sub, token_scope, expected_status): + """The exchange endpoint only accepts a callback-scope JWT whose sub is this callback's id.""" + sub = str(queued_callback.id) if token_sub == "self" else token_sub validator = mock.AsyncMock(spec=JWTValidator) validator.avalidated_claims.return_value = { - "sub": "11111111-1111-1111-1111-111111111111", - "scope": "callback", + "sub": sub, + "scope": token_scope, "exp": 9999999999, "iat": 1000000000, "nbf": 1000000000, } lifespan.registry.register_value(JWTValidator, validator) - resp = client.patch(f"/execution/callbacks/{callback.id}/run") - assert resp.status_code == 403, resp.json() - clear_db_callbacks() - - -@pytest.mark.usefixtures("_use_real_jwt_bearer") -def test_run_rejects_execution_scope_token(client, session): - """Only callback-scope tokens are accepted on the exchange endpoint.""" - clear_db_callbacks() - callback = _make_queued_callback(session) - - validator = mock.AsyncMock(spec=JWTValidator) - validator.avalidated_claims.return_value = { - "sub": str(callback.id), - "scope": "execution", - "exp": 9999999999, - "iat": 1000000000, - "nbf": 1000000000, - } - lifespan.registry.register_value(JWTValidator, validator) - - resp = client.patch(f"/execution/callbacks/{callback.id}/run") - assert resp.status_code == 403, resp.json() - clear_db_callbacks() + resp = client.patch(f"/execution/callbacks/{queued_callback.id}/run") + assert resp.status_code == expected_status, resp.json() diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/test_callbacks.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/test_callbacks.py index 4f0325477dc4f..c652cede7b0f8 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/test_callbacks.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/test_callbacks.py @@ -18,59 +18,26 @@ import pytest -from airflow.models.callback import CallbackFetchMethod, CallbackState, ExecutorCallback -from airflow.sdk.definitions.callback import SyncCallback - -from tests_common.test_utils.db import clear_db_callbacks - pytestmark = pytest.mark.db_test - -def sync_callback(): - """Empty (sync) callable used for the version-gating tests.""" - pass - - -def _make_queued_callback(session, dag_id="test_dag"): - callback = ExecutorCallback( - callback_def=SyncCallback(sync_callback, kwargs={}), - fetch_method=CallbackFetchMethod.IMPORT_PATH, - ) - callback.data["dag_id"] = dag_id - callback.state = CallbackState.QUEUED - session.add(callback) - session.commit() - return callback - - -@pytest.fixture -def old_ver_client(client): - """Client configured to use the API version before the callback run endpoint was added.""" - client.headers["Airflow-API-Version"] = "2026-06-30" - return client +MISSING_CALLBACK_URL = "/execution/callbacks/00000000-0000-0000-0000-000000000000/run" class TestRunCallbackEndpointVersioning: """The callbacks/{callback_id}/run endpoint didn't exist before the 2026-08-01 API version.""" - @pytest.fixture(autouse=True) - def setup_teardown(self): - clear_db_callbacks() - yield - clear_db_callbacks() - - def test_old_version_returns_404(self, old_ver_client, session): - """PATCH /callbacks/{callback_id}/run should not exist in older API versions.""" - callback = _make_queued_callback(session) + def test_old_version_returns_404(self, client): + """Before 2026-08-01 the route is absent, so routing itself 404s (no endpoint-shaped detail).""" + client.headers["Airflow-API-Version"] = "2026-06-30" - response = old_ver_client.patch(f"/execution/callbacks/{callback.id}/run") + response = client.patch(MISSING_CALLBACK_URL) assert response.status_code == 404 + assert response.json() == {"detail": "Not Found"} - def test_head_version_works(self, client, session): - """PATCH /callbacks/{callback_id}/run should work in the current API version.""" - callback = _make_queued_callback(session) + def test_head_version_routes_to_endpoint(self, client): + """At head the route exists: the same request reaches the endpoint's own 404 handling.""" + response = client.patch(MISSING_CALLBACK_URL) - response = client.patch(f"/execution/callbacks/{callback.id}/run") - - assert response.status_code == 204 + assert response.status_code == 404 + assert response.json()["detail"]["reason"] == "not_found"