You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ur...@apache.org on 2025/04/11 09:51:23 UTC
(airflow) branch main updated: Allow pulling XCom from inlet events (#49054)
This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 3ed428fcf76 Allow pulling XCom from inlet events (#49054)
3ed428fcf76 is described below
commit 3ed428fcf76c16f9b07dc5b647b6902c8dbb35d3
Author: Tzu-ping Chung <ur...@gmail.com>
AuthorDate: Fri Apr 11 17:51:16 2025 +0800
Allow pulling XCom from inlet events (#49054)
---
task-sdk/src/airflow/sdk/bases/xcom.py | 11 ++-
task-sdk/src/airflow/sdk/execution_time/comms.py | 75 +++++++++++++++
task-sdk/src/airflow/sdk/execution_time/context.py | 23 ++---
.../src/airflow/sdk/execution_time/task_runner.py | 7 +-
.../tests/task_sdk/execution_time/test_context.py | 107 ++++++++++++++++++---
.../task_sdk/execution_time/test_task_runner.py | 4 +-
6 files changed, 199 insertions(+), 28 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/bases/xcom.py b/task-sdk/src/airflow/sdk/bases/xcom.py
index caff009a3d8..4129f265462 100644
--- a/task-sdk/src/airflow/sdk/bases/xcom.py
+++ b/task-sdk/src/airflow/sdk/bases/xcom.py
@@ -17,7 +17,7 @@
from __future__ import annotations
-from typing import Any
+from typing import Any, Protocol
import structlog
@@ -26,6 +26,13 @@ from airflow.sdk.execution_time.comms import DeleteXCom, GetXCom, SetXCom, XComR
log = structlog.get_logger(logger_name="task")
+class TIKeyProtocol(Protocol):
+ dag_id: str
+ task_id: str
+ run_id: str
+ map_index: int
+
+
class BaseXCom:
"""BaseXcom is an interface now to interact with XCom backends."""
@@ -116,7 +123,7 @@ class BaseXCom:
def get_value(
cls,
*,
- ti_key: Any,
+ ti_key: TIKeyProtocol,
key: str,
) -> Any:
"""
diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py
index 76c858f3ac1..a9c27934577 100644
--- a/task-sdk/src/airflow/sdk/execution_time/comms.py
+++ b/task-sdk/src/airflow/sdk/execution_time/comms.py
@@ -43,14 +43,19 @@ Execution API server is because:
from __future__ import annotations
+from collections.abc import Iterator
from datetime import datetime
+from functools import cached_property
from typing import Annotated, Any, Literal, Union
from uuid import UUID
+import attrs
from fastapi import Body
from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, field_serializer
from airflow.sdk.api.datamodels._generated import (
+ AssetEventDagRunReference,
+ AssetEventResponse,
AssetEventsResponse,
AssetResponse,
BundleInfo,
@@ -108,6 +113,50 @@ class AssetResult(AssetResponse):
return cls(**asset_response.model_dump(exclude_defaults=True), type="AssetResult")
+@attrs.define(kw_only=True)
+class AssetEventSourceTaskInstance:
+ """Used in AssetEventResult."""
+
+ dag_id: str
+ task_id: str
+ run_id: str
+ map_index: int
+
+ def xcom_pull(
+ self,
+ *,
+ key: str = "return_value", # TODO: Make this a constant; see RuntimeTaskInstance.
+ default: Any = None,
+ ) -> Any:
+ from airflow.sdk.execution_time.xcom import XCom
+
+ if (value := XCom.get_value(ti_key=self, key=key)) is None:
+ return default
+ return value
+
+
+class AssetEventResult(AssetEventResponse):
+ """Used in AssetEventsResult."""
+
+ @classmethod
+ def from_asset_event_response(cls, asset_event_response: AssetEventResponse) -> AssetEventResult:
+ return cls(**asset_event_response.model_dump(exclude_defaults=True))
+
+ @cached_property
+ def source_task_instance(self) -> AssetEventSourceTaskInstance | None:
+ if not (self.source_task_id and self.source_dag_id and self.source_run_id):
+ return None
+ if self.source_map_index is None:
+ return None
+
+ return AssetEventSourceTaskInstance(
+ dag_id=self.source_dag_id,
+ task_id=self.source_task_id,
+ run_id=self.source_run_id,
+ map_index=self.source_map_index,
+ )
+
+
class AssetEventsResult(AssetEventsResponse):
"""Response to GetAssetEvent request."""
@@ -129,6 +178,32 @@ class AssetEventsResult(AssetEventsResponse):
type="AssetEventsResult",
)
+ def iter_asset_event_results(self) -> Iterator[AssetEventResult]:
+ return (AssetEventResult.from_asset_event_response(event) for event in self.asset_events)
+
+
+class AssetEventDagRunReferenceResult(AssetEventDagRunReference):
+ @classmethod
+ def from_asset_event_dag_run_reference(
+ cls,
+ asset_event_dag_run_reference: AssetEventDagRunReference,
+ ) -> AssetEventDagRunReferenceResult:
+ return cls(**asset_event_dag_run_reference.model_dump(exclude_defaults=True))
+
+ @cached_property
+ def source_task_instance(self) -> AssetEventSourceTaskInstance | None:
+ if not (self.source_task_id and self.source_dag_id and self.source_run_id):
+ return None
+ if self.source_map_index is None:
+ return None
+
+ return AssetEventSourceTaskInstance(
+ dag_id=self.source_dag_id,
+ task_id=self.source_task_id,
+ run_id=self.source_run_id,
+ map_index=self.source_map_index,
+ )
+
class XComResult(XComResponse):
"""Response to ReadXCom request."""
diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py
index fee91f3efaf..e79c4a6efde 100644
--- a/task-sdk/src/airflow/sdk/execution_time/context.py
+++ b/task-sdk/src/airflow/sdk/execution_time/context.py
@@ -44,11 +44,12 @@ if TYPE_CHECKING:
from uuid import UUID
from airflow.sdk import Variable
- from airflow.sdk.api.datamodels._generated import AssetEventDagRunReference, AssetEventResponse
from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions.connection import Connection
from airflow.sdk.definitions.context import Context
from airflow.sdk.execution_time.comms import (
+ AssetEventDagRunReferenceResult,
+ AssetEventResult,
AssetEventsResult,
AssetResult,
ConnectionResult,
@@ -377,20 +378,20 @@ class _AssetRefResolutionMixin:
@attrs.define
class TriggeringAssetEventsAccessor(
_AssetRefResolutionMixin,
- Mapping[Union[Asset, AssetAlias, AssetRef], Sequence["AssetEventDagRunReference"]],
+ Mapping[Union[Asset, AssetAlias, AssetRef], Sequence["AssetEventDagRunReferenceResult"]],
):
"""Lazy mapping of triggering asset events."""
- _events: Mapping[BaseAssetUniqueKey, Sequence[AssetEventDagRunReference]]
+ _events: Mapping[BaseAssetUniqueKey, Sequence[AssetEventDagRunReferenceResult]]
@classmethod
- def build(cls, events: Iterable[AssetEventDagRunReference]) -> TriggeringAssetEventsAccessor:
- collected: dict[BaseAssetUniqueKey, list[AssetEventDagRunReference]] = collections.defaultdict(list)
+ def build(cls, events: Iterable[AssetEventDagRunReferenceResult]) -> TriggeringAssetEventsAccessor:
+ coll: dict[BaseAssetUniqueKey, list[AssetEventDagRunReferenceResult]] = collections.defaultdict(list)
for event in events:
- collected[AssetUniqueKey(name=event.asset.name, uri=event.asset.uri)].append(event)
+ coll[AssetUniqueKey(name=event.asset.name, uri=event.asset.uri)].append(event)
for alias in event.source_aliases:
- collected[AssetAliasUniqueKey(name=alias.name)].append(event)
- return cls(collected)
+ coll[AssetAliasUniqueKey(name=alias.name)].append(event)
+ return cls(coll)
def __str__(self) -> str:
return f"TriggeringAssetEventAccessor(_events={self._events})"
@@ -404,7 +405,7 @@ class TriggeringAssetEventsAccessor(
def __len__(self) -> int:
return len(self._events)
- def __getitem__(self, key: Asset | AssetAlias | AssetRef) -> Sequence[AssetEventDagRunReference]:
+ def __getitem__(self, key: Asset | AssetAlias | AssetRef) -> Sequence[AssetEventDagRunReferenceResult]:
hashable_key: BaseAssetUniqueKey
if isinstance(key, Asset):
hashable_key = AssetUniqueKey.from_asset(key)
@@ -531,7 +532,7 @@ class InletEventsAccessors(Mapping[Union[int, Asset, AssetAlias, AssetRef], Any]
def __len__(self) -> int:
return len(self._inlets)
- def __getitem__(self, key: int | Asset | AssetAlias | AssetRef) -> list[AssetEventResponse]:
+ def __getitem__(self, key: int | Asset | AssetAlias | AssetRef) -> list[AssetEventResult]:
from airflow.sdk.definitions.asset import Asset
from airflow.sdk.execution_time.comms import (
ErrorResponse,
@@ -573,7 +574,7 @@ class InletEventsAccessors(Mapping[Union[int, Asset, AssetAlias, AssetRef], Any]
if TYPE_CHECKING:
assert isinstance(msg, AssetEventsResult)
- return msg.asset_events
+ return list(msg.iter_asset_event_results())
@cache # Prevent multiple API access.
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 59da7cefbb9..102ca78b524 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -56,6 +56,7 @@ from airflow.sdk.definitions.param import process_params
from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType
from airflow.sdk.execution_time.callback_runner import create_executable_runner
from airflow.sdk.execution_time.comms import (
+ AssetEventDagRunReferenceResult,
DagRunStateResult,
DeferTask,
DRCount,
@@ -170,7 +171,6 @@ class RuntimeTaskInstance(TaskInstance):
"params": validated_params,
# TODO: Make this go through Public API longer term.
# "test_mode": task_instance.test_mode,
- # "triggering_asset_events": lazy_object_proxy.Proxy(get_triggering_events),
"var": {
"json": VariableAccessor(deserialize_json=True),
"value": VariableAccessor(deserialize_json=False),
@@ -182,7 +182,10 @@ class RuntimeTaskInstance(TaskInstance):
context_from_server: Context = {
# TODO: Assess if we need to pass these through timezone.coerce_datetime
"dag_run": dag_run, # type: ignore[typeddict-item] # Removable after #46522
- "triggering_asset_events": TriggeringAssetEventsAccessor.build(dag_run.consumed_asset_events),
+ "triggering_asset_events": TriggeringAssetEventsAccessor.build(
+ AssetEventDagRunReferenceResult.from_asset_event_dag_run_reference(event)
+ for event in dag_run.consumed_asset_events
+ ),
"task_instance_key_str": f"{self.task.dag_id}__{self.task.task_id}__{dag_run.run_id}",
"task_reschedule_count": from_server.task_reschedule_count or 0,
"prev_start_date_success": lazy_object_proxy.Proxy(
diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py
index 2cff0703a13..eeca6595338 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_context.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_context.py
@@ -24,11 +24,7 @@ from unittest.mock import MagicMock, patch
import pytest
from airflow.sdk import get_current_context
-from airflow.sdk.api.datamodels._generated import (
- AssetEventDagRunReference,
- AssetEventResponse,
- AssetResponse,
-)
+from airflow.sdk.api.datamodels._generated import AssetEventResponse, AssetResponse
from airflow.sdk.definitions.asset import (
Asset,
AssetAlias,
@@ -40,13 +36,19 @@ from airflow.sdk.definitions.connection import Connection
from airflow.sdk.definitions.variable import Variable
from airflow.sdk.exceptions import ErrorType
from airflow.sdk.execution_time.comms import (
+ AssetEventDagRunReferenceResult,
+ AssetEventResult,
+ AssetEventSourceTaskInstance,
AssetEventsResult,
AssetResult,
ConnectionResult,
ErrorResponse,
GetAssetByName,
GetAssetByUri,
+ GetAssetEventByAsset,
+ GetXCom,
VariableResult,
+ XComResult,
)
from airflow.sdk.execution_time.context import (
ConnectionAccessor,
@@ -474,7 +476,7 @@ class TestTriggeringAssetEventsAccessor:
@pytest.fixture
def accessor(self, event_data):
return TriggeringAssetEventsAccessor.build(
- [AssetEventDagRunReference.model_validate(d) for d in event_data],
+ AssetEventDagRunReferenceResult.model_validate(d) for d in event_data
)
@pytest.mark.parametrize(
@@ -487,7 +489,7 @@ class TestTriggeringAssetEventsAccessor:
],
)
def test_getitem(self, event_data, accessor, key, result_indexes):
- expected = [AssetEventDagRunReference.model_validate(event_data[index]) for index in result_indexes]
+ expected = [AssetEventDagRunReferenceResult.model_validate(event_data[i]) for i in result_indexes]
assert accessor[key] == expected
@pytest.mark.parametrize(
@@ -507,7 +509,7 @@ class TestTriggeringAssetEventsAccessor:
result_indexes,
):
mock_supervisor_comms.get_message.return_value = resolved_asset
- expected = [AssetEventDagRunReference.model_validate(event_data[index]) for index in result_indexes]
+ expected = [AssetEventDagRunReferenceResult.model_validate(event_data[i]) for i in result_indexes]
assert accessor[Asset.ref(name=name)] == expected
assert len(mock_supervisor_comms.send_request.mock_calls) == 1
assert mock_supervisor_comms.send_request.mock_calls[0].kwargs["msg"] == GetAssetByName(name=name)
@@ -530,12 +532,36 @@ class TestTriggeringAssetEventsAccessor:
result_indexes,
):
mock_supervisor_comms.get_message.return_value = resolved_asset
- expected = [AssetEventDagRunReference.model_validate(event_data[index]) for index in result_indexes]
+ expected = [AssetEventDagRunReferenceResult.model_validate(event_data[i]) for i in result_indexes]
assert accessor[Asset.ref(uri=uri)] == expected
assert len(mock_supervisor_comms.send_request.mock_calls) == 1
assert mock_supervisor_comms.send_request.mock_calls[0].kwargs["msg"] == GetAssetByUri(uri=uri)
assert _AssetRefResolutionMixin._asset_ref_cache
+ def test_source_task_instance_xcom_pull(self, mock_supervisor_comms, accessor):
+ events = accessor[Asset("2")]
+ assert len(events) == 1
+ source = events[0].source_task_instance
+ assert source == AssetEventSourceTaskInstance(dag_id="d1", task_id="t2", run_id="r1", map_index=-1)
+
+ mock_supervisor_comms.reset_mock()
+ mock_supervisor_comms.get_message.side_effect = [
+ XComResult(key="return_value", value="__example_xcom_value__"),
+ ]
+ assert source.xcom_pull() == "__example_xcom_value__"
+ assert mock_supervisor_comms.send_request.mock_calls == [
+ mock.call(
+ log=mock.ANY,
+ msg=GetXCom(
+ key="return_value",
+ dag_id="d1",
+ run_id="r1",
+ task_id="t2",
+ map_index=-1,
+ ),
+ ),
+ ]
+
TEST_ASSET = Asset(name="test_uri", uri="test://test")
TEST_ASSET_ALIAS = AssetAlias(name="name")
@@ -617,7 +643,9 @@ class TestInletEventAccessor:
AssetResult(name="test_uri", uri="test://test", group="asset"),
AssetResult(name="test_uri", uri="test://test", group="asset"),
]
- return InletEventsAccessors(inlets=TEST_INLETS)
+ obj = InletEventsAccessors(inlets=TEST_INLETS)
+ mock_supervisor_comms.reset_mock()
+ return obj
@pytest.mark.usefixtures("mock_supervisor_comms")
def test__iter__(self, sample_inlet_evnets_accessor):
@@ -632,7 +660,7 @@ class TestInletEventAccessor:
def test__get_item__(self, key, sample_inlet_evnets_accessor, mock_supervisor_comms):
# This test only verifies a valid key can be used to access inlet events,
# but not access asset events are fetched. That is verified in test_asset_events in execution_api
- asset_event_resp = AssetEventResponse(
+ asset_event_resp = AssetEventResult(
id=1,
created_dagruns=[],
timestamp=timezone.utcnow(),
@@ -665,3 +693,60 @@ class TestInletEventAccessor:
def test_for_asset_alias(self, mocked__getitem__, sample_inlet_evnets_accessor):
sample_inlet_evnets_accessor.for_asset_alias(name="name")
assert mocked__getitem__.call_args[0][0] == TEST_ASSET_ALIAS
+
+ def test_source_task_instance_xcom_pull(self, sample_inlet_evnets_accessor, mock_supervisor_comms):
+ mock_supervisor_comms.get_message.side_effect = [
+ AssetEventsResult(
+ asset_events=[
+ AssetEventResponse(
+ id=1,
+ timestamp=timezone.utcnow(),
+ asset=AssetResponse(name="test_uri", uri="test://test", group="asset"),
+ created_dagruns=[],
+ source_dag_id="__dag__",
+ source_run_id="__run__",
+ source_task_id="__task__",
+ source_map_index=0,
+ ),
+ AssetEventResponse(
+ id=1,
+ timestamp=timezone.utcnow(),
+ asset=AssetResponse(name="test_uri", uri="test://test", group="asset"),
+ created_dagruns=[],
+ ),
+ ],
+ )
+ ]
+ events = sample_inlet_evnets_accessor[Asset.ref(name="test_uri")]
+ assert mock_supervisor_comms.send_request.mock_calls == [
+ mock.call(log=mock.ANY, msg=GetAssetEventByAsset(name="test_uri", uri=None)),
+ ]
+
+ assert len(events) == 2
+ assert events[1].source_task_instance is None
+
+ source = events[0].source_task_instance
+ assert source == AssetEventSourceTaskInstance(
+ dag_id="__dag__",
+ run_id="__run__",
+ task_id="__task__",
+ map_index=0,
+ )
+
+ mock_supervisor_comms.reset_mock()
+ mock_supervisor_comms.get_message.side_effect = [
+ XComResult(key="return_value", value="__example_xcom_value__"),
+ ]
+ assert source.xcom_pull() == "__example_xcom_value__"
+ assert mock_supervisor_comms.send_request.mock_calls == [
+ mock.call(
+ log=mock.ANY,
+ msg=GetXCom(
+ key="return_value",
+ dag_id="__dag__",
+ run_id="__run__",
+ task_id="__task__",
+ map_index=0,
+ ),
+ ),
+ ]
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index 1db6efa1b14..acdace83443 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -49,7 +49,6 @@ from airflow.listeners.listener import get_listener_manager
from airflow.providers.standard.operators.python import PythonOperator
from airflow.sdk import DAG, BaseOperator, Connection, dag as dag_decorator, get_current_context
from airflow.sdk.api.datamodels._generated import (
- AssetEventResponse,
AssetProfile,
AssetResponse,
DagRunState,
@@ -62,6 +61,7 @@ from airflow.sdk.definitions.asset import Asset, AssetAlias, Dataset, Model
from airflow.sdk.definitions.param import DagParam
from airflow.sdk.exceptions import ErrorType
from airflow.sdk.execution_time.comms import (
+ AssetEventResult,
AssetEventsResult,
BundleInfo,
ConnectionResult,
@@ -821,7 +821,7 @@ def test_run_with_asset_outlets(
def test_run_with_asset_inlets(create_runtime_ti, mock_supervisor_comms):
"""Test running a basic task that contains asset inlets."""
- asset_event_resp = AssetEventResponse(
+ asset_event_resp = AssetEventResult(
id=1,
created_dagruns=[],
timestamp=timezone.utcnow(),