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 34c3dc35406f8..5083bc9036fc8 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.types import UtcDateTime @@ -73,6 +72,7 @@ CurrentTIToken, ExecutionAPIRoute, get_team_name_for_ti, + issue_execution_token, require_auth, ) from airflow.configuration import conf @@ -324,9 +324,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..dbc3d955c2242 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 @@ -196,6 +196,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 +210,12 @@ async def require_auth( CurrentTIToken: TIToken = Depends(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.""" + 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/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..6646b165104bb --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py @@ -0,0 +1,176 @@ +# 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 + + +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 + + +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 + + 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 JWT validation runs end-to-end.""" + 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": "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 7e681c1114956..233e636b0ae39 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -1091,6 +1091,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 @@ -1298,6 +1309,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 0f9b8130e2654..dc70c9e0268da 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -2182,3 +2182,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