ephraimbuddy commented on code in PR #72297:
URL: https://github.com/apache/airflow/pull/72297#discussion_r4455806078
##########
task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py:
##########
@@ -133,3 +134,13 @@ class
AddExpandedTICountToSupervisorTIRunContext(VersionChange):
instructions_to_migrate_to_previous_version = (
schema(TIRunContext).field("expanded_ti_count").didnt_exist,
)
+
+
+class AddTaskInstanceIdToSetXCom(VersionChange):
Review Comment:
**Wrong version file.** `v2027_02_28.py` is now the in-progress supervisor
schema version (added in #75387), and `schema/AGENTS.md` says PRs append to the
most recent file. Adding to `2026-10-30` changes a frozen contract after the
fact. Please move `AddTaskInstanceIdToSetXCom` there and register it under
`Version("2027-02-28", ...)`. `check-supervisor-schemas-versions` only checks
that some file under `versions/` changed, so CI won't catch this.
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -1001,6 +1001,9 @@ class SetXCom(BaseModel):
map_index: int | None = None
dag_result: bool = False
mapped_length: int | None = None
+ # When set (own-TI / callback writes), the API pins the row to this
attempt instead of
+ # resolving by coordinates alone — so a stale callback cannot write the
successor TI.
+ task_instance_id: UUID | None = None
Review Comment:
With this, an `on_retry_callback` push lands on the archived try. Reads
resolve to the live try, so neither the retry nor downstream tasks will see the
value, and that's the Slack-thread case from the description. It's better than
overwriting the next try, but users won't expect it. Either document it in the
callback docs, or reject the write (see the 410 question on `_find_writer_id`)
so it fails loudly.
##########
task-sdk/src/airflow/sdk/execution_time/schema/schema.json:
##########
@@ -3947,6 +3947,19 @@
"default": null,
"title": "Mapped Length"
},
+ "task_instance_id": {
Review Comment:
**Vendored schema copies are stale.** `go-sdk/schema/supervisor-schema.json`
and `ts-sdk/schema/supervisor-schema.json` must match this file exactly, and on
this head they don't, so `sync-go-sdk-schemas` and `sync-ts-sdk-schemas` will
fail static checks. Please run them and commit the refreshed copies, along with
whatever the Go and TS generated-code checks regenerate from them.
##########
airflow-core/tests/unit/dag_processing/test_processor.py:
##########
@@ -2655,6 +2681,128 @@ def
test_handle_request_get_variable_masks_value_with_key(self, proc):
"type": "VariableResult",
}
+ def test_handle_request_set_xcom(self, proc):
+ with patch.object(DagFileProcessorProcess, "send_msg", autospec=True)
as mock_send_msg:
+ proc._handle_request(
+ SetXCom(key="slack_thread_ts", value="123.456", dag_id="d",
run_id="r", task_id="t"),
+ structlog.get_logger(),
+ req_id=789,
+ )
+
+ proc.client.xcoms.set.assert_called_once_with(
+ "d",
+ "r",
+ "t",
+ "slack_thread_ts",
+ "123.456",
+ None,
+ dag_result=False,
+ mapped_length=None,
+ task_instance_id=None,
+ )
+ mock_send_msg.assert_called_once()
+ _, args, kwargs = mock_send_msg.mock_calls[0]
+ assert args[0] is proc
+ assert args[1] is None
+ assert kwargs["request_id"] == 789
+ assert kwargs["error"] is None
+
+ @staticmethod
+ def _make_dag_run(dag_maker, dag_id: str, team_name: str | None) -> DagRun:
+ """Create a dag run, optionally owning the dag by ``team_name`` via
its bundle."""
+ from airflow.models import DagModel
+ from airflow.models.dagbundle import DagBundleModel
+ from airflow.models.team import Team
+ from airflow.providers.standard.operators.empty import EmptyOperator
+ from airflow.utils.session import create_session
+
+ with dag_maker(dag_id=dag_id):
+ EmptyOperator(task_id="task")
+ dag_run = dag_maker.create_dagrun(run_id="run1",
state=DagRunState.RUNNING)
+
+ if team_name is not None:
+ with create_session() as session:
+ bundle = DagBundleModel(name=f"bundle-{dag_id}")
+ team = session.scalar(select(Team).where(Team.name ==
team_name)) or Team(name=team_name)
+ bundle.teams.append(team)
+ session.add(bundle)
+ session.flush()
+ session.execute(
+ update(DagModel).where(DagModel.dag_id ==
dag_id).values(bundle_name=bundle.name)
+ )
+ return dag_run
+
+ @pytest.mark.parametrize(
+ ("multi_team", "team_owned"),
+ [
+ pytest.param(False, True, id="multi_team_disabled"),
+ pytest.param(True, False, id="multi_team_enabled_global_dag"),
+ pytest.param(True, True, id="multi_team_enabled_team_owned_dag"),
+ ],
+ )
+ def test_handle_request_set_xcom_via_execution_api(
+ self, dag_maker, inprocess_client, multi_team, team_owned
+ ):
+ """
+ SetXCom through the real Execution API, not a mocked client.
+
+ ``test_handle_request_set_xcom`` above only proves the message reaches
+ ``client.xcoms.set`` with the right arguments; the client itself is a
mock. This
+ exercises the real ``InProcessExecutionAPI`` round trip the Dag
processor uses in
+ production. The Dag processor's in-process client bypasses
``has_xcom_access``
+ entirely (see that dependency's docstring), so the write succeeds
regardless of
+ multi-team mode or which team owns the target Dag.
+ """
+ dag_id = f"set_xcom_dag_{uuid.uuid4().hex}"
+ dag_run = self._make_dag_run(dag_maker, dag_id, "team_a" if team_owned
else None)
+
+ proc_mock = MagicMock()
+ proc_mock.create_time.return_value = 0.0
+ r, w = socketpair()
+ proc = DagFileProcessorProcess(
+ process_log=structlog.get_logger().bind(),
+ id=uuid.uuid4(),
+ pid=1234,
+ process=proc_mock,
+ stdin=w,
+ logger_filehandle=MagicMock(spec=BinaryIO),
+ client=inprocess_client,
+ bundle_name="mybundle",
+ dag_file_rel_path="dags/my_dag.py",
+ )
+ proc._open_sockets.clear()
+ r.close()
+ w.close()
+
+ with create_session() as session:
+ ti_id = dag_run.get_task_instance("task", session=session).id
+
+ msg = SetXCom(
+ key="slack_thread_ts",
+ value="123.456",
+ dag_id=dag_id,
+ run_id="run1",
+ task_id="task",
+ task_instance_id=ti_id,
+ )
+
+ with conf_vars({("core", "multi_team"): str(multi_team)}):
+ with patch.object(DagFileProcessorProcess, "send_msg",
autospec=True):
+ proc._handle_request(msg, structlog.get_logger(), req_id=1)
+
+ with create_session() as session:
+ xcom = session.scalar(
+ select(XComModel).where(
+ XComModel.dag_id == dag_id,
+ XComModel.run_id == "run1",
+ XComModel.task_id == "task",
+ XComModel.key == "slack_thread_ts",
+ )
+ )
+ assert xcom is not None
+ assert XComModel.deserialize_value(xcom) == 123.456
Review Comment:
`"123.456"` was pushed as a string. The float here comes from
`XComModel.deserialize_value` running `json.loads` on it; the SDK's read path
returns `"123.456"`. Please assert `xcom.value == "123.456"`, or read it back
through `inprocess_client.xcoms.get(...)`.
##########
airflow-core/tests/unit/dag_processing/test_processor.py:
##########
@@ -2655,6 +2681,128 @@ def
test_handle_request_get_variable_masks_value_with_key(self, proc):
"type": "VariableResult",
}
+ def test_handle_request_set_xcom(self, proc):
+ with patch.object(DagFileProcessorProcess, "send_msg", autospec=True)
as mock_send_msg:
+ proc._handle_request(
+ SetXCom(key="slack_thread_ts", value="123.456", dag_id="d",
run_id="r", task_id="t"),
+ structlog.get_logger(),
+ req_id=789,
+ )
+
+ proc.client.xcoms.set.assert_called_once_with(
+ "d",
+ "r",
+ "t",
+ "slack_thread_ts",
+ "123.456",
+ None,
+ dag_result=False,
+ mapped_length=None,
+ task_instance_id=None,
+ )
+ mock_send_msg.assert_called_once()
+ _, args, kwargs = mock_send_msg.mock_calls[0]
+ assert args[0] is proc
+ assert args[1] is None
+ assert kwargs["request_id"] == 789
+ assert kwargs["error"] is None
+
+ @staticmethod
+ def _make_dag_run(dag_maker, dag_id: str, team_name: str | None) -> DagRun:
+ """Create a dag run, optionally owning the dag by ``team_name`` via
its bundle."""
+ from airflow.models import DagModel
+ from airflow.models.dagbundle import DagBundleModel
+ from airflow.models.team import Team
+ from airflow.providers.standard.operators.empty import EmptyOperator
+ from airflow.utils.session import create_session
+
+ with dag_maker(dag_id=dag_id):
+ EmptyOperator(task_id="task")
+ dag_run = dag_maker.create_dagrun(run_id="run1",
state=DagRunState.RUNNING)
+
+ if team_name is not None:
+ with create_session() as session:
+ bundle = DagBundleModel(name=f"bundle-{dag_id}")
+ team = session.scalar(select(Team).where(Team.name ==
team_name)) or Team(name=team_name)
+ bundle.teams.append(team)
+ session.add(bundle)
+ session.flush()
+ session.execute(
+ update(DagModel).where(DagModel.dag_id ==
dag_id).values(bundle_name=bundle.name)
+ )
+ return dag_run
+
+ @pytest.mark.parametrize(
+ ("multi_team", "team_owned"),
+ [
+ pytest.param(False, True, id="multi_team_disabled"),
Review Comment:
These cases only exercise the in-process `has_xcom_access` override, which
this PR doesn't change, so they pass or fail together. Please drop the
parametrization and `_make_dag_run` (its inline imports also break the
imports-at-top rule). The HTTP tests in `test_xcoms.py` cover the real scope.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -850,17 +870,48 @@ def delete_xcom(
def _find_writer_id(
- token: TIToken,
+ token: ExecutionToken,
*,
dag_id: str,
run_id: str,
task_id: str,
map_index: int,
session: Session,
dag_bag: DBDagBag,
+ task_instance_id: UUID | None = None,
) -> UUID | None:
"""Resolve the task instance a write targets with the rule reads use,
working from the caller."""
- caller = session.get(TaskInstance, token.id,
execution_options={"include_all_attempts": True})
+ # Explicit attempt id wins: pin the write (including archived) and never
fall back to a successor.
+ if task_instance_id is not None:
+ target = session.get(TaskInstance, task_instance_id,
execution_options={"include_all_attempts": True})
+ if target is None:
+ return None
+ if (target.dag_id, target.run_id, target.task_id) != (dag_id, run_id,
task_id):
Review Comment:
The new 400s (coordinate mismatch, `map_index` mismatch) and the 404 for an
unknown `task_instance_id` aren't covered. Please add cases to
`TestXComsSetEndpoint`.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -850,17 +870,48 @@ def delete_xcom(
def _find_writer_id(
- token: TIToken,
+ token: ExecutionToken,
*,
dag_id: str,
run_id: str,
task_id: str,
map_index: int,
session: Session,
dag_bag: DBDagBag,
+ task_instance_id: UUID | None = None,
) -> UUID | None:
"""Resolve the task instance a write targets with the rule reads use,
working from the caller."""
- caller = session.get(TaskInstance, token.id,
execution_options={"include_all_attempts": True})
+ # Explicit attempt id wins: pin the write (including archived) and never
fall back to a successor.
+ if task_instance_id is not None:
Review Comment:
**This lets any caller write to any try.** The branch writes to whichever
try `task_instance_id` names, including archived ones, and only checks that its
coordinates match the URL.
- A task token is never required to name its own try, so any task can write
into another task's archived history. Main's `_require_live_attempt` takes the
opposite position: "Historical attempts never grant access".
- It skips the producer resolver, so the 400 for "target is in a loop the
caller is not in" no longer applies here.
The team and bundle boundaries still hold, since the target must match the
path `dag_id`, so this doesn't break the security model. It's cheap to close,
though: for task tokens, accept only `task_instance_id == token.id` (the
in-process Dag processor arrives as a non-TI caller and needs its own case).
Please also decide whether a write to an archived try should be stored or
rejected with 410, the way `_require_live_attempt` treats archived tries.
##########
airflow-core/tests/unit/api_fastapi/execution_api/test_security.py:
##########
@@ -844,7 +844,8 @@ def
test_exchange_rejects_invalid_grants_and_file_identity(self, client, body, s
("POST", f"/jobs/{JOB_ID}/complete", {}, 403),
("POST", f"/jobs/{JOB_ID}/parse-token", {}, 403),
("POST", "/jobs", {}, 403),
- ("POST", "/xcoms/granted_dag/run/task/key", {}, 403),
+ # Auth succeeds for dag_parse; missing TI yields 404 (set was
previously execution-only → 403).
Review Comment:
nit: this describes history rather than the code. Either drop it or reword
it to something like `# Admitted for dag_parse; 404 because the TI doesn't
exist.`
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -725,9 +726,27 @@ def set_xcom(
mapped_length: Annotated[
int | None, Query(ge=0, description="Number of mapped tasks this value
expands into")
] = None,
- token=_TriggerAwareTIToken,
+ task_instance_id: Annotated[
+ UUID | None,
+ Query(
+ description=(
+ "When set, write to this task-instance attempt (including
archived). "
+ "Coordinates must match; the write does not fall back to a
successor attempt."
+ ),
+ ),
+ ] = None,
+ # Dag-parse tokens are admitted so callbacks during parsing can persist
XComs; scope is the
Review Comment:
This comment isn't accurate:
- `PUT /variables` has no bundle bound, only `has_variable_access`, so this
isn't "the same as parse-time Variable writes".
- The manager still uses the in-process client, whose `always_allow` token
is a `TIToken`. `DagInGrantedBundle` does nothing for that token, so today's
Dag processor can write any Dag's XComs. The bundle scope only applies to
`dag_parse` tokens, which the manager doesn't use yet.
Please say that in the docstring and drop this comment, since the two say
the same thing.
##########
airflow-core/tests/unit/dag_processing/test_processor.py:
##########
@@ -2436,6 +2437,31 @@ def get_type_names(union_type):
+ "\n\nEither handle these types in ToDagProcessor or update
in_task_runner_but_not_in_dag_processing_process list."
)
+ def test_to_manager_accepts_set_xcom(self):
+ msg = DagFileProcessorProcess.decoder.validate_python(
+ {
+ "type": "SetXCom",
+ "key": "slack_thread_ts",
+ "value": "123.456",
+ "dag_id": "d",
+ "run_id": "r",
+ "task_id": "t",
+ }
+ )
+ assert isinstance(msg, SetXCom)
+
+ def test_to_manager_rejects_unsupported_delete_xcom(self):
Review Comment:
This passes without the PR. The
`in_task_runner_but_not_in_dag_processing_process` check already enforces the
union, which also makes `test_to_manager_accepts_set_xcom` redundant. The same
goes for the mocked `test_handle_request_set_xcom`, which the round-trip test
covers. Please drop all three.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
