Skip to content

Commit cee3476

Browse files
sdhilip200sdhilip200
committed
Add server-side XComModel value helper
1 parent fe3ad5f commit cee3476

5 files changed

Lines changed: 98 additions & 1 deletion

File tree

‎airflow-core/src/airflow/models/xcom.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,8 @@
5454
from sqlalchemy.orm import Session
5555
from sqlalchemy.sql.expression import Select, TextClause
5656

57+
from airflow.models.taskinstancekey import TaskInstanceKey
58+
5759

5860
XCOM_RETURN_KEY = "return_value"
5961

@@ -332,6 +334,51 @@ def get_many(
332334
return query.limit(limit)
333335
return query
334336

337+
@classmethod
338+
@provide_session
339+
def get_one(
340+
cls,
341+
*,
342+
key: str,
343+
dag_id: str,
344+
task_id: str,
345+
run_id: str,
346+
map_index: int | None = None,
347+
include_prior_dates: bool = False,
348+
session: Session = NEW_SESSION,
349+
) -> Any | None:
350+
"""Retrieve and deserialize a single XCom value from the metadata database."""
351+
result = session.execute(
352+
cls.get_many(
353+
key=key,
354+
dag_ids=dag_id,
355+
task_ids=task_id,
356+
run_id=run_id,
357+
map_indexes=map_index,
358+
include_prior_dates=include_prior_dates,
359+
limit=1,
360+
).with_only_columns(cls.value)
361+
).first()
362+
if result is None:
363+
return None
364+
return cls.deserialize_value(result)
365+
366+
@classmethod
367+
def get_value(
368+
cls,
369+
*,
370+
ti_key: TaskInstanceKey,
371+
key: str,
372+
) -> Any | None:
373+
"""Retrieve and deserialize an XCom value for a task instance key."""
374+
return cls.get_one(
375+
key=key,
376+
dag_id=ti_key.dag_id,
377+
task_id=ti_key.task_id,
378+
run_id=ti_key.run_id,
379+
map_index=ti_key.map_index,
380+
)
381+
335382
@staticmethod
336383
def serialize_value(
337384
value: Any,

‎airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,21 @@ class TryNumberPlugin(AirflowPlugin):
7272
operator_extra_links = [TryNumberLink()]
7373

7474

75+
class XComModelLink(BaseOperatorLink):
76+
name = "XCom Model"
77+
operators = [CustomOperator]
78+
79+
def get_link(self, operator, *, ti_key):
80+
from airflow.providers.common.compat.sdk import XComModel
81+
82+
return XComModel.get_value(key="plugin_link", ti_key=ti_key) or ""
83+
84+
85+
class XComModelPlugin(AirflowPlugin):
86+
name = "xcom_model_plugin"
87+
operator_extra_links = [XComModelLink()]
88+
89+
7590
@pytest.mark.mock_plugin_manager(plugins=[])
7691
class TestGetExtraLinks:
7792
dag_id = "TEST_DAG_ID"
@@ -265,6 +280,23 @@ def test_should_respond_200_support_plugins(self, test_client):
265280
).model_dump()
266281
)
267282

283+
@pytest.mark.mock_plugin_manager(plugins=[XComModelPlugin])
284+
def test_plugin_link_can_read_xcom_with_xcommodel(self, test_client):
285+
XCom.set(
286+
key="plugin_link",
287+
value="https://example.com/plugin-link",
288+
task_id=self.task_single_link,
289+
dag_id=self.dag_id,
290+
run_id=self.dag_run_id,
291+
)
292+
293+
response = test_client.get(
294+
f"/dags/{self.dag_id}/dagRuns/{self.dag_run_id}/taskInstances/{self.task_single_link}/links",
295+
)
296+
297+
assert response.status_code == 200
298+
assert response.json()["extra_links"]["XCom Model"] == "https://example.com/plugin-link"
299+
268300
def test_should_respond_200_mapped_task_instance(self, test_client, session):
269301
for map_index, value in enumerate(["TEST_LINK_VALUE_1", "TEST_LINK_VALUE_2"]):
270302
XCom.set(

‎airflow-core/tests/unit/models/test_xcom.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,16 @@ def test_xcom_get_one(self, session, task_instance):
240240
).first()
241241
assert XComModel.deserialize_value(stored_value) == {"key": "value"}
242242

243+
def test_xcommodel_get_value(self, task_instance, push_simple_json_xcom):
244+
push_simple_json_xcom(ti=task_instance, key="xcom_1", value={"key": "value"})
245+
246+
value = XComModel.get_value(
247+
ti_key=task_instance.key,
248+
key="xcom_1",
249+
)
250+
251+
assert value == {"key": "value"}
252+
243253
@pytest.fixture
244254
def tis_for_xcom_get_one_from_prior_date(self, task_instance_factory, push_simple_json_xcom):
245255
date1 = timezone.datetime(2021, 12, 3, 4, 56)

‎providers/common/compat/src/airflow/providers/common/compat/sdk.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
if TYPE_CHECKING:
3131
import airflow.sdk.io as io # noqa: F401
3232
import airflow.sdk.timezone as timezone # noqa: F401
33-
from airflow.models.xcom import XCOM_RETURN_KEY as XCOM_RETURN_KEY
33+
from airflow.models.xcom import XCOM_RETURN_KEY as XCOM_RETURN_KEY, XComModel as XComModel
3434
from airflow.sdk import (
3535
DAG as DAG,
3636
Asset as Asset,
@@ -135,6 +135,7 @@
135135
"AssetAlias": ("airflow.sdk", "airflow.datasets", "DatasetAlias"),
136136
"AssetAll": ("airflow.sdk", "airflow.datasets", "DatasetAll"),
137137
"AssetAny": ("airflow.sdk", "airflow.datasets", "DatasetAny"),
138+
"XComModel": ("airflow.models.xcom", "airflow.models.xcom", "XCom"),
138139
}
139140

140141
# Airflow 3-only renames (not available in Airflow 2)

‎providers/common/compat/tests/unit/common/compat/test_sdk.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,3 +52,10 @@ def test_invalid_import_raises_attribute_error():
5252

5353
with pytest.raises(AttributeError, match="has no attribute 'NonExistentClass'"):
5454
_ = sdk.NonExistentClass
55+
56+
57+
def test_xcom_model_imports_server_side_model():
58+
from airflow.models.xcom import XComModel
59+
from airflow.providers.common.compat import sdk
60+
61+
assert sdk.XComModel is XComModel

0 commit comments

Comments
 (0)