diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py index fe657354e02fb..8fd8d93e5ed64 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py @@ -35,6 +35,8 @@ logger = structlog.get_logger(logger_name=__name__) +TRIGGER_TI_ID_HEADER = "Airflow-Trigger-Ti-Id" + # Finalizer arguments must not reference the transport, or it cannot be garbage collected. def _shutdown_loop( @@ -51,12 +53,16 @@ def _shutdown_loop( thread.join(timeout=5) -def _configure_in_process_auth(app: FastAPI) -> None: +def _configure_in_process_auth(app: FastAPI, *, triggerer_id: int | None = None) -> None: from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken from airflow.api_fastapi.execution_api.routes.connections import has_connection_access from airflow.api_fastapi.execution_api.routes.variables import has_variable_access from airflow.api_fastapi.execution_api.routes.xcoms import has_xcom_access - from airflow.api_fastapi.execution_api.security import _IN_PROCESS_NON_TI_CALLER, _jwt_bearer + from airflow.api_fastapi.execution_api.security import ( + _IN_PROCESS_NON_TI_CALLER, + _IN_PROCESS_TRIGGERER_ID, + _jwt_bearer, + ) async def always_allow(request: Request): from uuid import UUID @@ -71,6 +77,8 @@ async def always_allow(request: Request): "task_instance_id" not in request.path_params and "X-Airflow-In-Process-Attempt-Id" not in request.headers ) + if triggerer_id is not None: + request.scope[_IN_PROCESS_TRIGGERER_ID] = triggerer_id claims = TIClaims(scope="execution") return TIToken(id=ti_id, claims=claims) @@ -90,6 +98,7 @@ class InProcessExecutionAPI: """ _app: FastAPI | None = None + triggerer_id: int | None = None @cached_property def app(self): @@ -105,7 +114,7 @@ def app(self): self._app = create_task_execution_api_app(lifespan=attrs.evolve(lifespan, registry=registry)) # Auth dependency overrides bypass validation, so the stub skips validator setup in lifespan(). registry.register_value(JWTValidator, None) - _configure_in_process_auth(self._app) + _configure_in_process_auth(self._app, triggerer_id=self.triggerer_id) return self._app diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py index c6ecaa15c766b..3cf2c536514e4 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py @@ -24,6 +24,7 @@ from uuid import UUID import attrs +from cadwyn import current_dependency_solver from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status from pydantic import JsonValue from sqlalchemy import select @@ -33,18 +34,20 @@ from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.core_api.base import BaseModel from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.execution_api.datamodels.token import DagParseToken, ExecutionToken +from airflow.api_fastapi.execution_api.datamodels.token import DagParseToken, ExecutionToken, TIToken from airflow.api_fastapi.execution_api.datamodels.xcom import ( XComResponse, XComSequenceIndexResponse, XComSequenceSliceResponse, ) +from airflow.api_fastapi.execution_api.in_process import TRIGGER_TI_ID_HEADER from airflow.api_fastapi.execution_api.security import ( + _IN_PROCESS_TRIGGERER_ID, CurrentExecutionToken, - CurrentTIToken, DagInGrantedBundle, ExecutionAPIRoute, ExecutionOrDagParseToken, + require_task_token, ) from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates from airflow.exceptions import TaskNotFound @@ -62,6 +65,7 @@ enclosing_loop, ) from airflow.models.taskinstance import TaskInstance +from airflow.models.trigger import Trigger from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, xcom_entity from airflow.utils.db import get_query_count from airflow.utils.state import TaskInstanceState @@ -252,29 +256,86 @@ def _build_xcom_read( raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error -def xcom_query( - dag_id: str, - run_id: str, - task_id: str, - key: str, +def _resolve_trigger_token( + request: Request, session: SessionDep, - dag_bag: DagBagDep, - map_index: Annotated[int | None, Query()] = None, - previous_iteration: bool = False, token: ExecutionToken = CurrentExecutionToken, -) -> Select: - return _build_xcom_read( - dag_id=dag_id, - run_id=run_id, - task_id=task_id, - key=key, - session=session, - dag_bag=dag_bag, - token=token, - map_index=map_index, - previous_iteration=previous_iteration, - sequence=True, - ).statement + dependency_solver=Depends(current_dependency_solver), +) -> ExecutionToken: + """ + Return the identity an XCom request acts as. + + A request made by a trigger's code carries its task instance's id in ``TRIGGER_TI_ID_HEADER``. The + triggerer subprocess can forge the header, so it is only honoured on the triggerer's in-process API. + + A write acts as the claimed task instance only while it is deferred on a trigger assigned to the calling + triggerer. A read acts as it while that try is live, in any state, because a trigger's ``on_kill`` runs + after the task instance leaves ``deferred`` and the trigger row is gone; reads are not a security + boundary on this API, which already serves every XCom. Anything else gets a 403. + """ + header = request.headers.get(TRIGGER_TI_ID_HEADER) + # The endpoint runs with Cadwyn's re-solved dependencies, so checking there reuses the endpoint's session. + if header is None or dependency_solver != "cadwyn": + return token + + triggerer_id = request.scope.get(_IN_PROCESS_TRIGGERER_ID) + if triggerer_id is None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "reason": "access_denied", + "message": f"{TRIGGER_TI_ID_HEADER} is only accepted from the triggerer", + }, + ) + + try: + claimed = UUID(header) + except ValueError: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"reason": "invalid_request", "message": f"{TRIGGER_TI_ID_HEADER} must be a UUID"}, + ) + + if request.method in ("GET", "HEAD"): + permitted = ( + session.scalar( + select(TaskInstance.working_set) + .where(TaskInstance.id == claimed) + .execution_options(include_all_attempts=True) + ) + is True + ) + denial = "Task instance is not a live attempt" + else: + permitted = ( + session.scalar( + select(TaskInstance.id) + .join(Trigger, Trigger.id == TaskInstance.trigger_id) + .where( + TaskInstance.id == claimed, + Trigger.triggerer_id == triggerer_id, + TaskInstance.state == TaskInstanceState.DEFERRED, + ) + ) + is not None + ) + denial = "Task instance is not deferred on a trigger assigned to this triggerer" + if not permitted: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"reason": "access_denied", "message": denial}, + ) + return token.model_copy(update={"id": claimed}) + + +_TriggerAwareToken = Depends(_resolve_trigger_token) + + +def _resolve_trigger_ti_token(token: ExecutionToken = _TriggerAwareToken) -> TIToken: + return require_task_token(token) + + +_TriggerAwareTIToken = Depends(_resolve_trigger_ti_token) @router.get( @@ -297,7 +358,7 @@ def get_mapped_xcom_by_index( session: SessionDep, dag_bag: DagBagDep, previous_iteration: bool = False, - token: ExecutionToken = CurrentExecutionToken, + token: ExecutionToken = _TriggerAwareToken, ) -> XComSequenceIndexResponse: xcom_read = _build_xcom_read( dag_id=dag_id, @@ -364,7 +425,7 @@ def get_mapped_xcom_by_slice( params: Annotated[GetXComSliceFilterParams, Query()], session: SessionDep, dag_bag: DagBagDep, - token: ExecutionToken = CurrentExecutionToken, + token: ExecutionToken = _TriggerAwareToken, ) -> XComSequenceSliceResponse: xcom_read = _build_xcom_read( dag_id=dag_id, @@ -448,7 +509,10 @@ def get_mapped_xcom_by_slice( dependencies=[ExecutionOrDagParseToken, DagInGrantedBundle], responses={ **create_openapi_http_exception_doc( - [(status.HTTP_400_BAD_REQUEST, "map_index cannot be specified in a HEAD request")] + [ + (status.HTTP_400_BAD_REQUEST, "map_index cannot be specified in a HEAD request"), + (status.HTTP_409_CONFLICT, "Several live task instances produce the requested value"), + ] ), status.HTTP_200_OK: { "description": "Metadata about the number of matching XCom values", @@ -463,10 +527,16 @@ def get_mapped_xcom_by_slice( description="Returns the count of mapped XCom values found in the `Content-Range` response header", ) def head_xcom( + dag_id: str, + run_id: str, + task_id: str, + key: str, response: Response, session: SessionDep, - xcom_query: Annotated[Select, Depends(xcom_query)], + dag_bag: DagBagDep, map_index: Annotated[int | None, Query()] = None, + previous_iteration: bool = False, + token: ExecutionToken = _TriggerAwareToken, ) -> None: """Get the count of XComs from database - not other XCom Backends.""" if map_index is not None: @@ -475,6 +545,17 @@ def head_xcom( detail={"reason": "invalid_request", "message": "Cannot specify map_index in a HEAD request"}, ) + xcom_query = _build_xcom_read( + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + key=key, + session=session, + dag_bag=dag_bag, + token=token, + previous_iteration=previous_iteration, + sequence=True, + ).statement count = get_query_count(xcom_query, session=session) # Tell the caller how many items in this query. We define a custom range unit (HTTP spec only defines # "bytes" but we can add our own) @@ -509,7 +590,7 @@ def get_xcom( session: SessionDep, params: Annotated[GetXcomFilterParams, Query()], dag_bag: DagBagDep, - token: ExecutionToken = CurrentExecutionToken, + token: ExecutionToken = _TriggerAwareToken, ) -> XComResponse: """Get an Airflow XCom from database - not other XCom Backends.""" owner = None @@ -643,7 +724,7 @@ def set_xcom( mapped_length: Annotated[ int | None, Query(ge=0, description="Number of mapped tasks this value expands into") ] = None, - token=CurrentTIToken, + token=_TriggerAwareTIToken, ): """Set an Airflow XCom.""" from airflow.configuration import conf @@ -773,7 +854,7 @@ def delete_xcom( task_id: str, key: Annotated[str, Path(min_length=1)], map_index: Annotated[int, Query()] = -1, - token=CurrentTIToken, + token=_TriggerAwareTIToken, ): """Delete a single XCom Value.""" owner = _find_writer_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 d354b5bc5e9a3..0d02409d578d5 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -141,6 +141,7 @@ _REQUEST_SCOPE_LIVE_ATTEMPT_KEY = "live_attempt_checked" _REQUEST_SCOPE_JOB_KEY = "dag_processor_job_id" _IN_PROCESS_NON_TI_CALLER = "airflow_in_process_non_ti_caller" +_IN_PROCESS_TRIGGERER_ID = "airflow_in_process_triggerer_id" _SKIP_AUTO_TI_ATTEMPT_LIVE = "skip_auto_ti_attempt_live" _P = ParamSpec("_P") _R = TypeVar("_R") diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 9552f90d63b07..24a1432a2d8fd 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -17,6 +17,7 @@ from __future__ import annotations import asyncio +import contextvars import functools import logging import math @@ -50,7 +51,7 @@ from airflow._shared.observability.metrics import stats from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.datamodels.taskinstance import task_instance_to_runtime -from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI +from airflow.api_fastapi.execution_api.in_process import TRIGGER_TI_ID_HEADER, InProcessExecutionAPI from airflow.configuration import conf from airflow.exceptions import TaskNotFound from airflow.executors import workloads @@ -140,6 +141,9 @@ from airflow.utils.session import create_session, provide_session if TYPE_CHECKING: + from uuid import UUID + + import httpx2 from opentelemetry.util._decorator import _AgnosticContextManager from sqlalchemy.orm import Session from structlog.typing import FilteringBoundLogger, WrappedLogger @@ -160,6 +164,23 @@ _ON_CANCEL_TIMEOUT: int = conf.getint("triggerer", "on_kill_timeout", fallback=30) +# Task instance of the trigger whose code is running, so its SDK requests act as that task instance; tasks +# serving several triggers must not inherit it. +_current_trigger_ti_id: contextvars.ContextVar[UUID | None] = contextvars.ContextVar( + "airflow_current_trigger_ti_id", default=None +) + +# Task instance of the trigger whose request the supervisor is handling, sent to the API as +# ``TRIGGER_TI_ID_HEADER``. +_handling_trigger_ti_id: contextvars.ContextVar[UUID | None] = contextvars.ContextVar( + "airflow_handling_trigger_ti_id", default=None +) + + +def _add_trigger_ti_id_header(request: httpx2.Request) -> None: + if (ti_id := _handling_trigger_ti_id.get()) is not None: + request.headers[TRIGGER_TI_ID_HEADER] = str(ti_id) + def _make_trigger_span( ti: TaskInstanceDTO | None, trigger_id: int, name: str @@ -581,11 +602,23 @@ def make_client(self) -> Client: """ from airflow.sdk.api.client import Client - client = Client(base_url=None, token="", dry_run=True, transport=InProcessExecutionAPI().transport) + triggerer_id = self.job.id if self.job is not None else None + in_process_api = InProcessExecutionAPI(triggerer_id=triggerer_id) + client = Client(base_url=None, token="", dry_run=True, transport=in_process_api.transport) # Mypy is wrong -- the setter accepts a string on the property setter! `URLType = URL | str` client.base_url = "http://in-process.invalid./" + if triggerer_id is not None: + client.event_hooks["request"].append(_add_trigger_ti_id_header) return client + @contextmanager + def _request_scope(self, request: _RequestFrame) -> Iterator[None]: + token = _handling_trigger_ti_id.set(request.ti_id) + try: + yield + finally: + _handling_trigger_ti_id.reset(token) + def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger, req_id: int) -> None: resp: BaseModel | None = None @@ -807,7 +840,9 @@ def handle_failed_triggers(self): def on_trigger_failure(self, trigger_id: int, exc: list[str] | None) -> None: """Record that a trigger failed.""" - Trigger.submit_failure(trigger_id=trigger_id, exc=exc) + Trigger.submit_failure( + trigger_id=trigger_id, exc=exc, triggerer_id=self.job.id if self.job is not None else None + ) def metric_tags(self) -> dict[str, str]: """ @@ -1212,7 +1247,9 @@ async def asend(self, msg: ToTriggerSupervisor) -> ToTriggerRunner | None: # non-blocking for the foreign loop. cf = asyncio.run_coroutine_threadsafe(self.asend(msg), self._loop) return await asyncio.wrap_future(cf) - frame = _RequestFrame(id=next(self.id_counter), body=msg.model_dump()) + frame = _RequestFrame( + id=next(self.id_counter), body=msg.model_dump(), ti_id=_current_trigger_ti_id.get() + ) future: asyncio.Future = current_loop.create_future() self._pending[frame.id] = future try: @@ -1686,6 +1723,8 @@ async def run_trigger( await greenback.ensure_portal() ti = trigger.task_instance + if ti is not None: + _current_trigger_ti_id.set(ti.id) bind_log_contextvars( trigger_id=trigger_id, ti_id=str(ti.id) if ti else None, @@ -1727,8 +1766,15 @@ async def run_trigger( shared_key = None if shared_key is not None and event_trigger is not None: - shared_stream = self._shared_streams.subscribe( - trigger_id=trigger_id, trigger=event_trigger, key=shared_key + # subscribe() may start the group's poll task, which copies the current context but + # serves several triggers. + subscribe_ctx = contextvars.copy_context() + subscribe_ctx.run(_current_trigger_ti_id.set, None) + shared_stream = subscribe_ctx.run( + self._shared_streams.subscribe, + trigger_id=trigger_id, + trigger=event_trigger, + key=shared_key, ) event_stream = event_trigger.filter_shared_stream(shared_stream) else: diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index 624b09d276dc9..1c0e1461604f0 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -305,10 +305,15 @@ def submit_event(cls, trigger_id, event: TriggerEvent, *, session: Session = NEW @classmethod @provide_session - def submit_failure(cls, trigger_id, exc=None, *, session: Session = NEW_SESSION) -> None: + def submit_failure( + cls, trigger_id, exc=None, *, triggerer_id: int | None = None, session: Session = NEW_SESSION + ) -> None: """ When a trigger has failed unexpectedly, mark everything that depended on it as failed. + With ``triggerer_id``, only act while that triggerer still owns the trigger: a reassigned + trigger is being run elsewhere, so a failure from its previous owner must not fail the task. + Notably, we have to actually run the failure code from a worker as it may have linked callbacks, so hilariously we have to re-schedule the task instances to a worker just so they can then fail. @@ -317,11 +322,12 @@ def submit_failure(cls, trigger_id, exc=None, *, session: Session = NEW_SESSION) the runtime code understands as immediate-fail, and pack the error into next_kwargs. """ - for task_instance in session.scalars( - select(TaskInstance).where( - TaskInstance.trigger_id == trigger_id, TaskInstance.state == TaskInstanceState.DEFERRED - ) - ): + query = select(TaskInstance).where( + TaskInstance.trigger_id == trigger_id, TaskInstance.state == TaskInstanceState.DEFERRED + ) + if triggerer_id is not None: + query = query.join(cls, cls.id == TaskInstance.trigger_id).where(cls.triggerer_id == triggerer_id) + for task_instance in session.scalars(query): # Add the error and set the next_method to the fail state if isinstance(exc, BaseException): traceback = format_exception(type(exc), exc, exc.__traceback__) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py index a846cdd9368a6..fce0ea5285c07 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py @@ -25,6 +25,7 @@ import pytest from fastapi import FastAPI, HTTPException, Path, Request, status +from fastapi.testclient import TestClient from sqlalchemy import delete, event, select, update from sqlalchemy.orm import Session @@ -32,6 +33,7 @@ from airflow.api_fastapi.auth.tokens import JWTValidator from airflow.api_fastapi.execution_api.datamodels.token import ExecutionToken, TIClaims, TIToken from airflow.api_fastapi.execution_api.datamodels.xcom import XComResponse +from airflow.api_fastapi.execution_api.in_process import TRIGGER_TI_ID_HEADER, InProcessExecutionAPI from airflow.api_fastapi.execution_api.security import _jwt_bearer, require_auth from airflow.exceptions import TaskNotFound from airflow.models.dag import DagModel @@ -39,6 +41,7 @@ from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion from airflow.models.task_coordinates import TaskCoordinateResolver from airflow.models.taskinstance import LegacyTaskDataOwner, TaskInstance +from airflow.models.trigger import Trigger from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, XComModelV1, XComModelV2, xcom_entity from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import PythonOperator @@ -50,9 +53,12 @@ from tests_common.test_utils.asserts import capture_orm_selects from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.db import clear_db_triggers pytestmark = pytest.mark.db_test +TRIGGERER_ID = 7 + @pytest.mark.parametrize("suffix", ["", "/item/0", "/item/-1", "/slice"]) def test_regional_mapped_xcom_reads_resolve_live_producers_before_slicing(client, dag_maker, session, suffix): @@ -522,6 +528,156 @@ def test_loop_task_reads_from_inside_the_loop_stay_on_one_pass(client, outside_l assert client.get(f"{fanout_url}/item/-1").json() == "1-1" +@pytest.fixture +def deferred_loop_pass(session, outside_loop_run): + dr, _, _, steps, _ = outside_loop_run + trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + trigger.triggerer_id = TRIGGERER_ID + session.add(trigger) + session.flush() + steps[1].state = TaskInstanceState.DEFERRED + steps[1].trigger_id = trigger.id + session.commit() + yield dr, steps, trigger + clear_db_triggers() + + +def test_trigger_in_a_loop_reads_and_writes_as_its_deferred_pass(session, deferred_loop_pass): + dr, steps, _ = deferred_loop_pass + headers = {TRIGGER_TI_ID_HEADER: str(steps[1].id)} + step_url = f"/xcoms/{dr.dag_id}/{dr.run_id}/body.step/key" + + sessions_begun = [] + + def record(api_session, *_): + sessions_begun.append(api_session) + + with TestClient(InProcessExecutionAPI(triggerer_id=TRIGGERER_ID).app) as api: + event.listen(Session, "after_begin", record) + try: + with capture_orm_selects("task_instance") as task_instance_selects: + assert api.get(step_url, headers=headers).json() == {"key": "key", "value": "step-1"} + assert api.head(step_url, headers=headers).headers["Content-Range"] == "map_indexes 1" + assert api.get(f"{step_url}/slice", headers=headers).json() == ["step-1"] + assert api.get(f"{step_url}/item/0", headers=headers).json() == "step-1" + assert api.post(f"{step_url}-written", json="from-trigger", headers=headers).status_code == 201 + assert api.delete(step_url, headers=headers).status_code == 200 + finally: + event.remove(Session, "after_begin", record) + + liveness_checks = [ + statement + for statement in task_instance_selects + if statement.startswith("SELECT task_instance.working_set FROM") + ] + assert len(liveness_checks) == 1 + assert len(set(map(id, sessions_begun))) == 6 + session.expire_all() + assert XComModelV2.get_for_attempt(steps[1].id, "key-written", session=session).value == "from-trigger" + assert XComModelV2.get_for_attempt(steps[1].id, "key", session=session) is None + assert XComModelV2.get_for_attempt(steps[0].id, "key", session=session).value == "step-0" + + +@pytest.mark.parametrize( + ("api_triggerer_id", "header", "expected_status"), + [ + pytest.param(None, None, 403, id="api-without-triggerer-id"), + pytest.param(TRIGGERER_ID, "not-a-uuid", 400, id="malformed-header"), + ], +) +def test_trigger_ti_id_header_rejected(deferred_loop_pass, api_triggerer_id, header, expected_status): + dr, steps, _ = deferred_loop_pass + + with TestClient(InProcessExecutionAPI(triggerer_id=api_triggerer_id).app) as api: + response = api.get( + f"/xcoms/{dr.dag_id}/{dr.run_id}/body.step/key", + headers={TRIGGER_TI_ID_HEADER: header or str(steps[1].id)}, + ) + + assert response.status_code == expected_status + + +NOT_DEFERRED_ON_THIS_TRIGGERER = "Task instance is not deferred on a trigger assigned to this triggerer" +NOT_A_LIVE_ATTEMPT = "Task instance is not a live attempt" + + +@pytest.mark.parametrize( + ("method", "change", "expected_status", "expected_value", "expected_denial"), + [ + pytest.param( + "post", + {"triggerer_id": TRIGGERER_ID + 1}, + 403, + None, + NOT_DEFERRED_ON_THIS_TRIGGERER, + id="write-trigger-assigned-to-another-triggerer", + ), + pytest.param( + "post", + {"state": TaskInstanceState.RUNNING}, + 403, + None, + NOT_DEFERRED_ON_THIS_TRIGGERER, + id="write-task-not-deferred", + ), + pytest.param( + "delete", + {"state": TaskInstanceState.RUNNING}, + 403, + None, + NOT_DEFERRED_ON_THIS_TRIGGERER, + id="delete-task-not-deferred", + ), + pytest.param( + "get", + {"state": TaskInstanceState.FAILED, "trigger_removed": True}, + 200, + "step-1", + None, + id="read-live-task-no-longer-deferred", + ), + pytest.param( + "head", + {"state": TaskInstanceState.FAILED, "trigger_removed": True}, + 200, + None, + None, + id="head-live-task-no-longer-deferred", + ), + pytest.param("get", {"archived": True}, 403, None, NOT_A_LIVE_ATTEMPT, id="read-archived-try"), + ], +) +def test_trigger_ti_id_claim_check( + session, deferred_loop_pass, method, change, expected_status, expected_value, expected_denial +): + dr, steps, trigger = deferred_loop_pass + trigger.triggerer_id = change.get("triggerer_id", TRIGGERER_ID) + steps[1].state = change.get("state", TaskInstanceState.DEFERRED) + if change.get("trigger_removed"): + steps[1].trigger_id = None + session.delete(trigger) + session.commit() + if change.get("archived"): + steps[1].archive(reason="cleared", session=session) + session.commit() + + with TestClient(InProcessExecutionAPI(triggerer_id=TRIGGERER_ID).app) as api: + response = api.request( + method, + f"/xcoms/{dr.dag_id}/{dr.run_id}/body.step/key", + headers={TRIGGER_TI_ID_HEADER: str(steps[1].id)}, + json="from-trigger" if method == "post" else None, + ) + + body = response.json() if response.content else {} + denial = body.get("detail", {}).get("message") if expected_status == 403 else None + assert (response.status_code, body.get("value"), denial) == ( + expected_status, + expected_value, + expected_denial, + ) + + @pytest.mark.parametrize("key", ["_airflow_loop_decision", "_airflow_loop_other"]) def test_normal_xcom_push_rejects_loop_reserved_prefix(client, create_task_instance, key): ti = create_task_instance() diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index 674c3737f08a6..212a40aa529f6 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -35,6 +35,7 @@ from unittest.mock import ANY, AsyncMock, MagicMock, patch import greenback +import httpx2 import msgspec import pendulum import pytest @@ -50,6 +51,7 @@ from structlog.typing import FilteringBoundLogger from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.in_process import TRIGGER_TI_ID_HEADER from airflow.executors import workloads from airflow.executors.workloads.task import TaskInstanceDTO from airflow.executors.workloads.trigger import RunTrigger @@ -65,6 +67,7 @@ TriggerLoggingFactory, TriggerRunner, TriggerRunnerSupervisor, + _current_trigger_ti_id, _make_trigger_span, messages, ) @@ -95,6 +98,7 @@ ErrorResponse, GetAssetStateStoreByName, GetAssetStateStoreByUri, + GetXCom, OKResponse, SetAssetStateStoreByName, SetAssetStateStoreByUri, @@ -518,6 +522,40 @@ def test_client_delegates_to_make_client_and_caches_result(supervisor_builder, m make_client.assert_called_once_with(supervisor) +@pytest.mark.parametrize("has_job", [True, False]) +def test_handle_requests_sends_the_trigger_ti_id_header_only_for_the_frames_own_requests( + supervisor_builder, jobless_supervisor, mocker, has_job +): + ti_id = uuid.uuid4() + expected = str(ti_id) if has_job else None + seen: dict[str, str | None] = {} + supervisor = supervisor_builder() if has_job else jobless_supervisor + + def handler(request: httpx2.Request) -> httpx2.Response: + key = request.url.path.rsplit("/", 1)[-1] + seen[key] = request.headers.get(TRIGGER_TI_ID_HEADER) + if key == "with-trigger": + other_thread = threading.Thread( + target=lambda: supervisor.client.xcoms.get("d", "r", "t", "from-thread") + ) + other_thread.start() + other_thread.join() + return httpx2.Response(200, json={"key": key, "value": 1}) + + in_process_api = mocker.patch("airflow.jobs.triggerer_job_runner.InProcessExecutionAPI", autospec=True) + in_process_api.return_value.transport = httpx2.MockTransport(handler) + requests = supervisor.handle_requests(log=mocker.Mock(spec=FilteringBoundLogger)) + next(requests) + + def get_xcom_body(key: str) -> dict: + return GetXCom(key=key, dag_id="d", run_id="r", task_id="t").model_dump() + + requests.send(_RequestFrame(id=1, body=get_xcom_body("with-trigger"), ti_id=ti_id)) + requests.send(_RequestFrame(id=2, body=get_xcom_body("without-trigger"))) + + assert seen == {"with-trigger": expected, "from-thread": None, "without-trigger": None} + + def test_run_context_exits_when_subprocess_dies(supervisor_builder, mocker): """Breaking out of the loop on a dead subprocess still unwinds run_context.""" from contextlib import contextmanager @@ -2331,6 +2369,36 @@ def test_failed_trigger(session, dag_maker, supervisor_builder): assert task_instance.next_kwargs["traceback"][-1] == "ModuleNotFoundError: No module named 'fake'\n" +def test_failure_from_triggerer_that_lost_the_trigger_leaves_task_deferred( + session, dag_maker, supervisor_builder +): + trigger_orm = Trigger(classpath="fake.classpath", kwargs={}) + session.add(trigger_orm) + session.flush() + with dag_maker(dag_id="test_reassigned_trigger", session=session): + EmptyOperator(task_id="dummy1") + task_instance = dag_maker.create_dagrun().task_instances[0] + task_instance.state = TaskInstanceState.DEFERRED + task_instance.trigger_id = trigger_orm.id + session.commit() + supervisor = supervisor_builder() + supervisor.load_triggers() + other_triggerer = Job() + session.add(other_triggerer) + session.flush() + session.execute( + update(Trigger).where(Trigger.id == trigger_orm.id).values(triggerer_id=other_triggerer.id) + ) + session.commit() + supervisor.failed_triggers.append((trigger_orm.id, ["ModuleNotFoundError: No module named 'fake'\n"])) + + supervisor.handle_failed_triggers() + + task_instance.refresh_from_db() + assert task_instance.state == TaskInstanceState.DEFERRED + assert task_instance.trigger_id == trigger_orm.id + + class CustomTrigger(BaseTrigger): """Custom Trigger that will access one Variable and one Connection.""" @@ -3159,6 +3227,22 @@ def get_type_names(union_type): ) +def _make_ti_dto(task_id="my_task", map_index=-1, context_carrier=None): + return TaskInstanceDTO( + id=uuid.uuid4(), + dag_version_id=uuid.uuid4(), + task_id=task_id, + dag_id="test_dag", + run_id="test_run", + try_number=1, + map_index=map_index, + pool_slots=1, + queue="default", + priority_weight=1, + context_carrier=context_carrier, + ) + + class TestMakeTriggerSpan: """Tests for the _make_trigger_span helper in the triggerer job runner.""" @@ -3171,29 +3255,14 @@ def sdk_tracer_provider(self): with mock.patch("airflow.jobs.triggerer_job_runner.tracer", test_tracer): yield - def _make_ti_dto(self, task_id="my_task", map_index=-1, context_carrier=None): - return TaskInstanceDTO( - id=uuid.uuid4(), - dag_version_id=uuid.uuid4(), - task_id=task_id, - dag_id="test_dag", - run_id="test_run", - try_number=1, - map_index=map_index, - pool_slots=1, - queue="default", - priority_weight=1, - context_carrier=context_carrier, - ) - def test_make_trigger_span_name_with_task_instance(self): - ti = self._make_ti_dto(task_id="sensor_task", map_index=-1) + ti = _make_ti_dto(task_id="sensor_task", map_index=-1) with _make_trigger_span(ti=ti, trigger_id=1, name="MySensor"): pass assert self.exporter.get_finished_spans()[0].name == "trigger.sensor_task" def test_make_trigger_span_name_with_mapped_task(self): - ti = self._make_ti_dto(task_id="sensor_task", map_index=2) + ti = _make_ti_dto(task_id="sensor_task", map_index=2) with _make_trigger_span(ti=ti, trigger_id=1, name="MySensor"): pass assert self.exporter.get_finished_spans()[0].name == "trigger.sensor_task_2" @@ -3213,7 +3282,7 @@ def test_make_trigger_span_uses_task_context_carrier(self): TraceContextTextMapPropagator().inject(ti_carrier, context=parent_ctx) expected_parent_span_id = parent_span.get_span_context().span_id - ti = self._make_ti_dto(context_carrier=ti_carrier) + ti = _make_ti_dto(context_carrier=ti_carrier) with _make_trigger_span(ti=ti, trigger_id=1, name="MySensor"): pass @@ -3223,7 +3292,7 @@ def test_make_trigger_span_uses_task_context_carrier(self): assert spans[0].parent.span_id == expected_parent_span_id def test_make_trigger_span_sets_attributes_with_ti(self): - ti = self._make_ti_dto(task_id="my_task", map_index=1) + ti = _make_ti_dto(task_id="my_task", map_index=1) with _make_trigger_span(ti=ti, trigger_id=5, name="MyTrigger"): pass @@ -3417,6 +3486,34 @@ async def test_unknown_frame_id_doesnt_crash_reader(decoder_pair): assert not decoder._reader_task.done(), "reader loop crashed unexpectedly" +@pytest.mark.asyncio +@pytest.mark.execution_timeout(15) +async def test_asend_frame_carries_the_current_trigger_ti_id(decoder_pair): + decoder, server_sock = decoder_pair + frames: list[_RequestFrame] = [] + + def serve(): + frame = _read_frame_sync(server_sock) + frames.append(frame) + server_sock.sendall( + _ResponseFrame( + id=frame.id, body={"type": "TriggerStateSync", "to_create": [], "to_cancel": []} + ).as_bytes() + ) + + thread = threading.Thread(target=serve, daemon=True) + thread.start() + ti_id = uuid.uuid4() + token = _current_trigger_ti_id.set(ti_id) + try: + await decoder.asend(messages.TriggerStateChanges(events=None, finished=[1], failures=None)) + finally: + _current_trigger_ti_id.reset(token) + thread.join(timeout=5) + + assert [frame.ti_id for frame in frames] == [ti_id] + + def test_trigger_state_messages_round_trip_with_seqs(): """Event triples and persist confirmations survive the wire encode/decode.""" changes_adapter = TypeAdapter(ToTriggerSupervisor) @@ -3624,6 +3721,78 @@ async def _drive(): assert seq is None +@pytest.mark.asyncio +@pytest.mark.parametrize("has_task_instance", [True, False]) +async def test_run_trigger_sets_current_trigger_ti_id_only_for_task_instance_triggers(has_task_instance): + ti = _make_ti_dto() + seen: list[uuid.UUID | None] = [] + + class _RecordingTrigger(BaseTrigger): + def serialize(self): + return (f"{type(self).__module__}.{type(self).__qualname__}", {}) + + async def run(self): + seen.append(_current_trigger_ti_id.get()) + yield TriggerEvent(True) + + trigger_runner = TriggerRunner() + trigger_runner.triggers = { + 5: {"task": MagicMock(spec=asyncio.Task), "is_watcher": False, "name": "t", "events": 0} + } + trigger = _RecordingTrigger() + if has_task_instance: + trigger.task_instance = ti + + await asyncio.create_task(trigger_runner.run_trigger(5, trigger)) + + assert seen == [ti.id if has_task_instance else None] + + +@pytest.mark.asyncio +async def test_shared_stream_group_tasks_do_not_inherit_the_subscribing_trigger_ti_id(): + ti = _make_ti_dto() + seen: dict[str, uuid.UUID | None] = {} + + class _SharedTrigger(BaseEventTrigger): + def serialize(self): + return (f"{type(self).__module__}.{type(self).__qualname__}", {}) + + def shared_stream_key(self): + return "key" + + @classmethod + async def open_shared_stream(cls, kwargs): + seen["group"] = _current_trigger_ti_id.get() + yield {} + await asyncio.Event().wait() + + async def filter_shared_stream(self, shared_stream): + seen["trigger"] = _current_trigger_ti_id.get() + async for raw in shared_stream: + yield TriggerEvent(raw) + + async def run(self): # pragma: no cover - replaced by filter_shared_stream + yield TriggerEvent({}) + + trigger_runner = TriggerRunner() + trigger_runner.triggers = { + 5: {"task": MagicMock(spec=asyncio.Task), "is_watcher": True, "name": "t", "events": 0} + } + trigger = _SharedTrigger() + trigger.task_instance = ti + + run_task = asyncio.create_task(trigger_runner.run_trigger(5, trigger)) + for _ in range(100): + await asyncio.sleep(0.01) + if trigger_runner.events: + break + run_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await run_task + + assert seen == {"group": None, "trigger": ti.id} + + @pytest.mark.asyncio async def test_trigger_event_payload_not_logged_at_info(cap_structlog): """Ensure the full event payload is not logged at INFO level.""" diff --git a/airflow-core/tests/unit/models/test_trigger.py b/airflow-core/tests/unit/models/test_trigger.py index ed7e414e0571d..716c0f53bfb35 100644 --- a/airflow-core/tests/unit/models/test_trigger.py +++ b/airflow-core/tests/unit/models/test_trigger.py @@ -341,24 +341,33 @@ def test_handle_event_submit_fails_task_when_the_event_payload_cannot_be_seriali assert "event payload could not be serialized" in task_instance.next_kwargs["error"] -def test_submit_failure(session, create_task_instance): +@pytest.mark.parametrize( + ("caller_triggerer_id", "expected_state", "expected_next_method"), + [ + pytest.param(None, State.SCHEDULED, "__fail__", id="unguarded"), + pytest.param(1, State.SCHEDULED, "__fail__", id="owning-triggerer"), + pytest.param(2, State.DEFERRED, None, id="previous-owner"), + ], +) +def test_submit_failure( + session, create_task_instance, caller_triggerer_id, expected_state, expected_next_method +): """ Tests that failures submitted to a trigger fail their dependent - task instances if not using a TaskEndEvent. + task instances if not using a TaskEndEvent, unless the caller no longer owns the trigger. """ - # Make a trigger trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}) + trigger.triggerer_id = 1 session.add(trigger) - # Make a TaskInstance that's deferred and waiting on it task_instance = create_task_instance(task_id="fake", logical_date=timezone.utcnow(), state=State.DEFERRED) task_instance.trigger_id = trigger.id session.commit() - # Call submit_event - Trigger.submit_failure(trigger.id, session=session) - # Check that the task instance is now scheduled to fail + + Trigger.submit_failure(trigger.id, triggerer_id=caller_triggerer_id, session=session) + updated_task_instance = session.scalar(select(TaskInstance)) - assert updated_task_instance.state == State.SCHEDULED - assert updated_task_instance.next_method == "__fail__" + assert updated_task_instance.state == expected_state + assert updated_task_instance.next_method == expected_next_method @pytest.mark.parametrize( diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 7c957027ba792..9c6993f2bdb4e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -208,6 +208,11 @@ class _RequestFrame(_FrameMixin, msgspec.Struct, array_like=True, frozen=True, o calls, so that server-side spans (e.g. POST /xcoms/…) appear as children of the correct task span rather than under the supervisor's own span. """ + ti_id: UUID | None = None + """Id of the task instance whose trigger made this request, set only by the triggerer's trigger runner. + + The supervisor forwards it so the API can act as that task instance. + """ class _ResponseFrame(_FrameMixin, msgspec.Struct, array_like=True, frozen=True, omit_defaults=True): # type: ignore[call-arg] diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 92295bd4f8ddb..960b2a01d1394 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -35,7 +35,7 @@ import weakref from collections import deque from collections.abc import Callable, Generator -from contextlib import contextmanager, suppress +from contextlib import contextmanager, nullcontext, suppress from datetime import UTC, datetime from enum import Enum, auto from http import HTTPStatus @@ -185,6 +185,7 @@ _trace_propagator = TraceContextTextMapPropagator() if TYPE_CHECKING: + from contextlib import AbstractContextManager from typing import Self from structlog.typing import FilteringBoundLogger, WrappedLogger @@ -1154,7 +1155,8 @@ def handle_requests(self, log: FilteringBoundLogger) -> Generator[None, _Request if request.context_carrier: ctx = _trace_propagator.extract(request.context_carrier) token = otel_context.attach(ctx) - self._handle_request(msg, log, request.id) + with self._request_scope(request): + self._handle_request(msg, log, request.id) except ServerResponseError as e: error_details = e.response.json() if e.response else None log.error( @@ -1203,6 +1205,10 @@ def handle_requests(self, log: FilteringBoundLogger) -> Generator[None, _Request if token is not None: otel_context.detach(token) + def _request_scope(self, request: _RequestFrame) -> AbstractContextManager[None]: + """Return a context manager that wraps the handling of ``request``; subclasses override to scope API calls.""" + return nullcontext() + def _handle_request(self, msg, log: FilteringBoundLogger, req_id: int) -> None: if self._request_handlers is None: raise NotImplementedError(f"{type(self).__name__} must declare its request handlers")