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/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index f9e7cb1725000..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,10 +148,10 @@ 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 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 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..b152e98ce9541 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py @@ -0,0 +1,85 @@ +# 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, 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.deps import DepContainer +from airflow.api_fastapi.execution_api.security import ( + ExecutionAPIRoute, + issue_execution_token, + require_auth, +) +from airflow.models.callback import Callback +from airflow.utils.state import CallbackState + +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, +) -> None: + """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( + 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": callback.state, + }, + ) + + callback.state = CallbackState.RUNNING + + 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..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,14 +70,15 @@ from typing import Any, get_args import structlog -from fastapi import Depends, HTTPException, Request, status +import svcs +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 @@ -196,6 +197,13 @@ async def require_auth( 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 @@ -203,6 +211,12 @@ async def require_auth( CurrentTIToken: TIToken = Depends(require_auth) +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"}) + + class ExecutionAPIRoute(APIRoute): """ Custom route class that precomputes allowed token types from Security scopes. 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/src/airflow/executors/workloads/base.py b/airflow-core/src/airflow/executors/workloads/base.py index 41334d68f3038..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) -> 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": "workload"}, + 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 c1842a5900685..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.""" 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/test_token_scope_boundaries.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py index bbf2be8704f46..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,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"}, + # 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 new file mode 100644 index 0000000000000..da24c5c5f9f0c --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py @@ -0,0 +1,131 @@ +# 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""" + pass + + +@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"] = "test_dag" + callback.state = CallbackState.QUEUED + session.add(callback) + session.commit() + return callback + + +class TestRunCallback: + 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 + payload = jwt.decode(response.headers["Refreshed-API-Token"], options={"verify_signature": False}) + assert payload["scope"] == "execution" + assert payload["sub"] == str(queued_callback.id) + + session.expire_all() + session.refresh(queued_callback) + assert queued_callback.state == CallbackState.RUNNING + + 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.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.""" + queued_callback.state = state + session.commit() + + response = client.patch(f"/execution/callbacks/{queued_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 JWT validation runs end-to-end.""" + exec_app.dependency_overrides.pop(require_auth, None) + + +@pytest.mark.usefixtures("_use_real_jwt_bearer") +@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": sub, + "scope": token_scope, + "exp": 9999999999, + "iat": 1000000000, + "nbf": 1000000000, + } + lifespan.registry.register_value(JWTValidator, validator) + + 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/__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..c652cede7b0f8 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_08_01/test_callbacks.py @@ -0,0 +1,43 @@ +# 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 + +pytestmark = pytest.mark.db_test + +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.""" + + 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 = client.patch(MISSING_CALLBACK_URL) + + assert response.status_code == 404 + assert response.json() == {"detail": "Not Found"} + + 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) + + assert response.status_code == 404 + assert response.json()["detail"]["reason"] == "not_found" 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" diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index a0c0a30de088f..ea220c68fde71 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -1097,6 +1097,17 @@ 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 an execution token.""" + self.client.patch(f"callbacks/{callback_id}/run") + + class BearerAuth(httpx.Auth): def __init__(self, token: str): self.token: str = token @@ -1304,6 +1315,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..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,6 +433,9 @@ 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. + 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..93e71943dbe47 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -2254,3 +2254,36 @@ 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): + 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) + + assert client.auth is not None + assert client.auth.token == "execution-token" + + def test_run_raises_on_conflict(self): + 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