This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 83b7552da0a Resolve a trigger's XCom requests as its deferred task
instance (#74772) (#74868)
83b7552da0a is described below
commit 83b7552da0ad8416b115d24184a83f1b76d072d2
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Sat Oct 17 19:43:41 2026 +0100
Resolve a trigger's XCom requests as its deferred task instance (#74772)
(#74868)
* Resolve a trigger's XCom requests as its deferred task instance
Note: _None_ of this is a security boundary, and should not be treated as
such.
The Execution API resolves XCom reads and writes for a task inside a loop
to the caller's own iteration, using the task identity in the caller's JWT.
Triggers have no such identity: the triggerer talks to an in-process API
that treats every request as the all-zero task instance. A trigger's XCom
requests therefore resolve as an anonymous caller rather than as the
deferred task it runs for, so it can neither read (a 400) nor write the
XComs of a task inside a loop.
Triggers can also share one runner coroutine and one comms channel (the
"shared streams" idea), so only the asyncio task running a trigger's code
knows which trigger it is; a context variable carries that trigger's task
instance id to each request frame. Shared-stream groups are started from
whichever trigger subscribes first but serve several triggers, so they start
with it cleared rather than inheriting that trigger's id. This isn't ideal,
but will do "for now".
The trigger runner process (the one running async code) sets the task
instance id, and since that process runs user code a trigger can claim to be
any task instance. The (in-process) API therefore checks the claim. A write
acts as the task instance only when it is deferred and its trigger is
already assigned to the calling triggerer. A read acts as it whenever that
try is still live, because after a user marks a deferred task failed or
successful the trigger row is gone and the triggerer still runs the
trigger's on_kill(). The most a misbehaving trigger can do is write as
another task that is deferred on the same triggerer, or read as any live
task, and it can already read and write every task's XComs by coordinates
through the in-process API. As stated, this is _not_ a security boundary we
can rely on
Requests that fail the check are rejected rather than falling back to the
anonymous identity, deliberately: a trigger whose task instance is no
longer deferred, or whose trigger was reassigned, should not still be
writing its task's XComs, and falling back would only work for tasks
outside loops. Clearing a deferred task archives its try before on_kill()
runs, so on_kill() after a clear cannot read its task's XComs; a trigger
that needs a job id there should keep it on the trigger, as the in-tree
triggers already do.
Only the in-process API that the triggerer creates knows which triggerer
is calling, so every other API rejects the header rather than trusting it.
A triggerer that eventually talks to a real API server will need its own
credential before this can work there.
Only the XCom routes honour the trigger's task instance id for now.
get_previous_task_instance also resolves relative to the caller, so a
trigger asking for its previous task instance still gets the latest pass
rather than the one at its own loop position. Every other request from the
triggerer keeps its current identity, so access rules and the live-attempt
check on those routes are unchanged.
* fixup! Resolve a trigger's XCom requests as its deferred task instance
Pin the HEAD and DELETE rules in the trigger claim check and assert which
check denied each 403, so loosening the read-method test no longer passes.
* fixup! Resolve a trigger's XCom requests as its deferred task instance
Ignore a trigger failure reported by a triggerer that no longer owns the
trigger. The XCom write check now answers 403 to a stale copy on its
previous owner, which raises before yielding and would otherwise have
Trigger.submit_failure send the task to __fail__ while the new owner is
still running it.
(cherry picked from commit 61f1d4b85e)
Adapted to v3-4-test's CurrentTIToken: the typed execution tokens from the
Dag processor scoped access work (#74370) are not on this branch.
---
.../api_fastapi/execution_api/in_process.py | 15 +-
.../api_fastapi/execution_api/routes/xcoms.py | 135 +++++++++++---
.../airflow/api_fastapi/execution_api/security.py | 1 +
.../src/airflow/jobs/triggerer_job_runner.py | 58 +++++-
airflow-core/src/airflow/models/trigger.py | 18 +-
.../execution_api/versions/head/test_xcoms.py | 156 ++++++++++++++++
airflow-core/tests/unit/jobs/test_triggerer_job.py | 207 +++++++++++++++++++--
airflow-core/tests/unit/models/test_trigger.py | 27 ++-
task-sdk/src/airflow/sdk/execution_time/comms.py | 5 +
.../src/airflow/sdk/execution_time/supervisor.py | 10 +-
10 files changed, 556 insertions(+), 76 deletions(-)
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 fe657354e02..8fd8d93e5ed 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 @@ if TYPE_CHECKING:
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 @@ def _configure_in_process_auth(app: FastAPI) -> None:
"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 @@ class InProcessExecutionAPI:
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 02fb6eceb6b..ad9d1d42939 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 typing import TYPE_CHECKING, Annotated
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,12 +34,14 @@ from airflow.api_fastapi.common.dagbag import DagBagDep
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 TIToken
from airflow.api_fastapi.execution_api.datamodels.xcom import (
XComResponse,
XComSequenceIndexResponse,
XComSequenceSliceResponse,
)
-from airflow.api_fastapi.execution_api.security import CurrentTIToken
+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, CurrentTIToken
from airflow.api_fastapi.execution_api.versions.v2026_10_30 import
IdentifyArchivedTaskStateUpdates
from airflow.exceptions import TaskNotFound
from airflow.models.dagrun import DagRun, InvalidLoopDecision,
validate_loop_decision
@@ -55,6 +58,7 @@ from airflow.models.task_coordinates import (
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
@@ -64,7 +68,6 @@ if TYPE_CHECKING:
from sqlalchemy.orm.util import AliasedClass
from sqlalchemy.sql.elements import ColumnElement
- from airflow.api_fastapi.execution_api.datamodels.token import TIToken
from airflow.models.dagbag import DBDagBag
@@ -242,29 +245,79 @@ 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=CurrentTIToken,
-) -> 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
+ token: TIToken = CurrentTIToken,
+ dependency_solver=Depends(current_dependency_solver),
+) -> TIToken:
+ """
+ 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)
@router.get(
@@ -286,7 +339,7 @@ def get_mapped_xcom_by_index(
session: SessionDep,
dag_bag: DagBagDep,
previous_iteration: bool = False,
- token=CurrentTIToken,
+ token=_TriggerAwareToken,
) -> XComSequenceIndexResponse:
xcom_read = _build_xcom_read(
dag_id=dag_id,
@@ -352,7 +405,7 @@ def get_mapped_xcom_by_slice(
params: Annotated[GetXComSliceFilterParams, Query()],
session: SessionDep,
dag_bag: DagBagDep,
- token=CurrentTIToken,
+ token=_TriggerAwareToken,
) -> XComSequenceSliceResponse:
xcom_read = _build_xcom_read(
dag_id=dag_id,
@@ -435,7 +488,10 @@ def get_mapped_xcom_by_slice(
"/{dag_id}/{run_id}/{task_id}/{key:path}",
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",
@@ -450,10 +506,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=_TriggerAwareToken,
) -> None:
"""Get the count of XComs from database - not other XCom Backends."""
if map_index is not None:
@@ -462,6 +524,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)
@@ -495,7 +568,7 @@ def get_xcom(
session: SessionDep,
params: Annotated[GetXcomFilterParams, Query()],
dag_bag: DagBagDep,
- token=CurrentTIToken,
+ token=_TriggerAwareToken,
) -> XComResponse:
"""Get an Airflow XCom from database - not other XCom Backends."""
owner = session.execute(
@@ -627,7 +700,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=_TriggerAwareToken,
):
"""Set an Airflow XCom."""
from airflow.configuration import conf
@@ -757,7 +830,7 @@ def delete_xcom(
task_id: str,
key: Annotated[str, Path(min_length=1)],
map_index: Annotated[int, Query()] = -1,
- token=CurrentTIToken,
+ token=_TriggerAwareToken,
):
"""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 6efd453c4f8..2c121b6e4ca 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py
@@ -95,6 +95,7 @@ VALID_TOKEN_TYPES: frozenset[str] =
frozenset(get_args(TokenScope))
_REQUEST_SCOPE_TOKEN_KEY = "ti_token"
_REQUEST_SCOPE_LIVE_ATTEMPT_KEY = "live_attempt_checked"
_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 9552f90d63b..24a1432a2d8 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.module_loading import import_string
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.log.task_log_address import
prepare_task_log_contexts, render
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 @@ _REASSIGNED_CANCEL_MSG = "__airflow_reassigned__"
_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 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
"""
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 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
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 @@ class TriggerCommsDecoder(CommsDecoder[ToTriggerRunner,
ToTriggerSupervisor]):
# 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 @@ class TriggerRunner:
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 @@ class TriggerRunner:
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 624b09d276d..1c0e1461604 100644
--- a/airflow-core/src/airflow/models/trigger.py
+++ b/airflow-core/src/airflow/models/trigger.py
@@ -305,10 +305,15 @@ class Trigger(Base):
@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 @@ class Trigger(Base):
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 5c8d49c3f1e..80f844eb3cf 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 @@ from uuid import uuid4
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,12 +33,14 @@ from airflow._shared.timezones import timezone
from airflow.api_fastapi.auth.tokens import JWTValidator
from airflow.api_fastapi.execution_api.datamodels.token import 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.dagrun import DagRun
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
@@ -49,9 +52,12 @@ from airflow.utils.state import DagRunState,
TaskInstanceState
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):
@@ -486,6 +492,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"
[email protected]
+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"
+
+
[email protected](
+ ("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"
+
+
[email protected](
+ ("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 674c3737f08..212a40aa529 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 import mock
from unittest.mock import ANY, AsyncMock, MagicMock, patch
import greenback
+import httpx2
import msgspec
import pendulum
import pytest
@@ -50,6 +51,7 @@ from sqlalchemy import delete, update
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 @@ from airflow.jobs.triggerer_job_runner import (
TriggerLoggingFactory,
TriggerRunner,
TriggerRunnerSupervisor,
+ _current_trigger_ti_id,
_make_trigger_span,
messages,
)
@@ -95,6 +98,7 @@ from airflow.sdk.execution_time.comms import (
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)
[email protected]("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 @@ class TestTriggererMessageTypes:
)
+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 @@ class TestMakeTriggerSpan:
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 @@ class TestMakeTriggerSpan:
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 @@ class TestMakeTriggerSpan:
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"
[email protected]
[email protected]_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 @@ def
test_run_trigger_appends_none_seq_for_non_shared_trigger():
assert seq is None
[email protected]
[email protected]("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]
+
+
[email protected]
+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 ed7e414e057..716c0f53bfb 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):
[email protected](
+ ("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 7c957027ba7..9c6993f2bdb 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 7af3b98616d..d6ff0b03f40 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 time
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
@@ -184,6 +184,7 @@ from opentelemetry.trace.propagation.tracecontext import
TraceContextTextMapProp
_trace_propagator = TraceContextTextMapPropagator()
if TYPE_CHECKING:
+ from contextlib import AbstractContextManager
from typing import Self
from structlog.typing import FilteringBoundLogger, WrappedLogger
@@ -1140,7 +1141,8 @@ class WatchedSubprocess:
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(
@@ -1189,6 +1191,10 @@ class WatchedSubprocess:
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")
