You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2025/03/19 17:37:51 UTC
(airflow) branch main updated: Fix bug in XComModel deserialization for value serialized in Task SDK (#47961)
This is an automated email from the ASF dual-hosted git repository.
kaxilnaik 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 59137f980fd Fix bug in XComModel deserialization for value serialized in Task SDK (#47961)
59137f980fd is described below
commit 59137f980fd9baad6a2acc50c0f0eea28f042a8c
Author: Kaxil Naik <ka...@apache.org>
AuthorDate: Wed Mar 19 23:07:44 2025 +0530
Fix bug in XComModel deserialization for value serialized in Task SDK (#47961)
Currently, the Task SDK uses `airflow.serialization` module to serialize XCom value during execution. This is because
we have agreed with the contract of the execution client (Python, Golang and other SDK clients) being responsible for serialization and de-serialization.
To do this, as part of https://github.com/apache/airflow/issues/45481, we skip XCom serialization from XComModel and instead directly store in the DB.
This works well when deserialization is done on Execution side itself. However, for some cases we deserialize it on the Scheduler or API server only. Some of those:
- Extra Operator Links
- Branching & Skipping logic in `NotPreviouslySkippedDep` where we pull XCom.
closes https://github.com/apache/airflow/issues/47907
closes https://github.com/apache/airflow/pull/47918
---
airflow/models/xcom.py | 37 +++++++++++-
.../api_fastapi/execution_api/routes/test_xcoms.py | 65 +++++++++++++++++++++-
tests/models/test_xcom.py | 51 +++++++++++++++++
3 files changed, 150 insertions(+), 3 deletions(-)
diff --git a/airflow/models/xcom.py b/airflow/models/xcom.py
index a8f1d3ee310..e2fdf49ad6f 100644
--- a/airflow/models/xcom.py
+++ b/airflow/models/xcom.py
@@ -331,11 +331,44 @@ class XComModel(TaskInstanceDependencies):
@staticmethod
def deserialize_value(result) -> Any:
- """Deserialize XCom value from str objects."""
+ """
+ Deserialize XCom value from a database result.
+
+ If deserialization fails, the raw value is returned, which must still be a valid Python JSON-compatible
+ type (e.g., ``dict``, ``list``, ``str``, ``int``, ``float``, or ``bool``).
+
+ XCom values are stored as JSON in the database, and SQLAlchemy automatically handles
+ serialization (``json.dumps``) and deserialization (``json.loads``). However, we
+ use a custom encoder for serialization (``serialize_value``) and deserialization to handle special
+ cases, such as encoding tuples via the Airflow Serialization module. These must be decoded
+ using ``XComDecoder`` to restore original types.
+
+ Some XCom values, such as those set via the Task Execution API, bypass ``serialize_value``
+ and are stored directly in JSON format. Since these values are already deserialized
+ by SQLAlchemy, they are returned as-is.
+
+ **Example: Handling a tuple**:
+
+ .. code-block:: python
+
+ original_value = (1, 2, 3)
+ serialized_value = XComModel.serialize_value(original_value)
+ print(serialized_value)
+ # '{"__classname__": "builtins.tuple", "__version__": 1, "__data__": [1, 2, 3]}'
+
+ This serialized value is stored in the database. When deserialized, the value is restored to the original tuple.
+
+ :param result: The XCom database row or object containing a ``value`` attribute.
+ :return: The deserialized Python object.
+ """
if result.value is None:
return None
- return json.loads(result.value, cls=XComDecoder)
+ try:
+ return json.loads(result.value, cls=XComDecoder)
+ except (ValueError, TypeError):
+ # Already deserialized (e.g., set via Task Execution API)
+ return result.value
class LazyXComSelectSequence(LazySelectSequence[Any]):
diff --git a/tests/api_fastapi/execution_api/routes/test_xcoms.py b/tests/api_fastapi/execution_api/routes/test_xcoms.py
index ae21fefdbf4..c2b49841b3a 100644
--- a/tests/api_fastapi/execution_api/routes/test_xcoms.py
+++ b/tests/api_fastapi/execution_api/routes/test_xcoms.py
@@ -29,7 +29,7 @@ from airflow.api_fastapi.execution_api.datamodels.xcom import XComResponse
from airflow.models.dagrun import DagRun
from airflow.models.taskmap import TaskMap
from airflow.models.xcom import XComModel
-from airflow.serialization.serde import serialize
+from airflow.serialization.serde import deserialize, serialize
from airflow.utils.session import create_session
pytestmark = pytest.mark.db_test
@@ -163,6 +163,69 @@ class TestXComsSetEndpoint:
task_map = session.query(TaskMap).filter_by(task_id=ti.task_id, dag_id=ti.dag_id).one_or_none()
assert task_map is None, "Should not be mapped"
+ @pytest.mark.parametrize(
+ "orig_value, ser_value, deser_value",
+ [
+ pytest.param(1, 1, 1, id="int"),
+ pytest.param(1.0, 1.0, 1.0, id="float"),
+ pytest.param("string", "string", "string", id="str"),
+ pytest.param(True, True, True, id="bool"),
+ pytest.param({"key": "value"}, {"key": "value"}, {"key": "value"}, id="dict"),
+ pytest.param([1, 2], [1, 2], [1, 2], id="list"),
+ pytest.param(
+ (1, 2),
+ # Client serializes tuple as encoded list, send the encoded list to the API
+ {"__classname__": "builtins.tuple", "__data__": [1, 2], "__version__": 1},
+ # The API will send the encoded list to the DB and sends the same encoded list back
+ # during the response to the client as it is the clients responsibility to
+ # serialize it into a JSON object & deserialize value into a native object.
+ {"__classname__": "builtins.tuple", "__data__": [1, 2], "__version__": 1},
+ id="tuple",
+ ),
+ ],
+ )
+ def test_xcom_round_trip(self, client, create_task_instance, session, orig_value, ser_value, deser_value):
+ """
+ Test that deserialization works when XCom values are stored directly in the DB with API Server.
+
+ This tests the case where the XCom value is stored from the Task API where the value is serialized
+ via Client SDK into JSON object and passed via the API Server to the DB. It by-passes
+ the XComModel.serialize_value and stores valid Python JSON compatible objects to DB.
+
+ This test is to ensure that the deserialization works correctly in this case as well as
+ checks that the value is stored correctly before it hits the API.
+ """
+
+ ti = create_task_instance()
+ session.commit()
+
+ # Serialize the value to simulate the client SDK
+ value = serialize(orig_value)
+
+ # Test that the value is serialized correctly
+ assert value == ser_value
+
+ response = client.post(
+ f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/xcom_1",
+ json=value,
+ )
+
+ assert response.status_code == 201
+
+ stored_value = XComModel.get_many(
+ key="xcom_1",
+ dag_ids=ti.dag_id,
+ task_ids=ti.task_id,
+ run_id=ti.run_id,
+ session=session,
+ ).first()
+ deserialized_value = XComModel.deserialize_value(stored_value)
+
+ assert deserialized_value == deser_value
+
+ # Ensure that the deserialized value on the client side is the same as the original value
+ assert deserialize(deserialized_value) == orig_value
+
def test_xcom_set_mapped(self, client, create_task_instance, session):
ti = create_task_instance()
session.commit()
diff --git a/tests/models/test_xcom.py b/tests/models/test_xcom.py
index c01c6a88c05..ce407d1fc69 100644
--- a/tests/models/test_xcom.py
+++ b/tests/models/test_xcom.py
@@ -370,3 +370,54 @@ class TestXComClear:
session=session,
)
assert session.query(XComModel).count() == 1
+
+
+class TestXComRoundTrip:
+ @pytest.mark.parametrize(
+ "value, expected_value",
+ [
+ pytest.param(1, 1, id="int"),
+ pytest.param(1.0, 1.0, id="float"),
+ pytest.param("string", "string", id="str"),
+ pytest.param(True, True, id="bool"),
+ pytest.param({"key": "value"}, {"key": "value"}, id="dict"),
+ pytest.param([1, 2, 3], [1, 2, 3], id="list"),
+ pytest.param((1, 2, 3), (1, 2, 3), id="tuple"), # tuple is preserved
+ pytest.param(None, None, id="none"),
+ ],
+ )
+ def test_xcom_round_trip(self, value, expected_value, push_simple_json_xcom, task_instance, session):
+ """Test that XComModel serialization and deserialization work as expected."""
+ push_simple_json_xcom(ti=task_instance, key="xcom_1", value=value)
+
+ stored_value = XComModel.get_many(
+ key="xcom_1",
+ dag_ids=task_instance.dag_id,
+ task_ids=task_instance.task_id,
+ run_id=task_instance.run_id,
+ session=session,
+ ).first()
+ deserialized_value = XComModel.deserialize_value(stored_value)
+
+ assert deserialized_value == expected_value
+
+ @pytest.mark.parametrize(
+ "value, expected_value",
+ [
+ pytest.param(1, 1, id="int"),
+ pytest.param(1.0, 1.0, id="float"),
+ pytest.param("string", "string", id="str"),
+ pytest.param(True, True, id="bool"),
+ pytest.param({"key": "value"}, {"key": "value"}, id="dict"),
+ pytest.param([1, 2, 3], [1, 2, 3], id="list"),
+ pytest.param((1, 2, 3), (1, 2, 3), id="tuple"), # tuple is preserved
+ pytest.param(None, None, id="none"),
+ ],
+ )
+ def test_xcom_deser_fallback(self, value, expected_value):
+ """Test fallback in deserialization."""
+
+ mock_xcom = MagicMock(value=value)
+ deserialized_value = XComModel.deserialize_value(mock_xcom)
+
+ assert deserialized_value == expected_value