diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py index 792988a4187c7..0d9dd086953a0 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import json from dataclasses import dataclass from typing import Annotated from uuid import UUID @@ -66,6 +67,7 @@ from airflow.models.task_coordinates import TaskCoordinateResolver, public_map_index_expression from airflow.models.taskinstance import TaskInstance from airflow.models.xcom import XComModel, build_xcom_read_query, select_producers, xcom_entity +from airflow.serialization.stringify import stringify as stringify_xcom xcom_router = AirflowRouter( tags=["XCom"], prefix="/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/xcomEntries" @@ -140,28 +142,15 @@ def get_xcom_entry( value = result.value if deserialize: - # Custom XCom backends may store references (eg: object storage paths) in the database. - # The custom XCom backend's deserialize_value() resolves these to actual values, but that is only - # used on workers during task execution. The API reads directly from the database and uses - # stringify() to convert DB values (references or serialized data) to human readable - # format for UI display or for API users. - import json - - from airflow.serialization.stringify import ( - StringifyNotSupportedError, - stringify as stringify_xcom, - ) - + # The value is rendered from what is stored in the database: no XCom backend is involved and no + # object is reconstructed from it. A custom backend may store a reference (eg: an object + # storage path) there; only workers resolve it into the actual value. try: parsed_value = json.loads(result.value) except (ValueError, TypeError): # Already deserialized (e.g., set via Task Execution API) parsed_value = result.value - - try: - value = stringify_xcom(parsed_value) - except StringifyNotSupportedError: - value = XComModel.deserialize_value(result) + value = stringify_xcom(parsed_value) else: # For native format, return the raw serialized value from the database # This preserves the JSON string format that the API expects diff --git a/airflow-core/src/airflow/serialization/definitions/operatorlink.py b/airflow-core/src/airflow/serialization/definitions/operatorlink.py index 51d8d04f39a8f..d1c6d2f286726 100644 --- a/airflow-core/src/airflow/serialization/definitions/operatorlink.py +++ b/airflow-core/src/airflow/serialization/definitions/operatorlink.py @@ -26,6 +26,7 @@ from airflow.models.dynamic_region import SENTINEL_REGION_ID from airflow.models.xcom import XComModel, xcom_entity +from airflow.serialization.stringify import stringify as stringify_xcom from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import create_session @@ -79,21 +80,10 @@ def get_link( ) return "" - from airflow.serialization.stringify import ( - StringifyNotSupportedError, - stringify as stringify_xcom, - ) - try: parsed_value = json.loads(result.value) except (ValueError, TypeError): # Handling for cases when types do not need to be deserialized (e.g. when value is a simple string link) parsed_value = result.value - try: - return str(stringify_xcom(parsed_value)) - except StringifyNotSupportedError: - # If stringify doesn't support the type, return the raw value as a string. - # This avoids the XComModel.deserialize_value() call that could - # instantiate arbitrary classes from untrusted XCom data. - return str(parsed_value) + return str(stringify_xcom(parsed_value)) diff --git a/airflow-core/src/airflow/serialization/stringify.py b/airflow-core/src/airflow/serialization/stringify.py index 74b654aa4511b..4be7087ecf3c4 100644 --- a/airflow-core/src/airflow/serialization/stringify.py +++ b/airflow-core/src/airflow/serialization/stringify.py @@ -28,14 +28,7 @@ # but it has no place in the human-readable XCom value rendering. _DAGBAG_PREFIX_RE = re.compile(r"unusual_prefix_[a-f0-9]{40}_") - -class StringifyNotSupportedError(ValueError): - """ - Raised when stringify() cannot handle a serialized object and full deserialization is required. - - This typically occurs when trying to stringify Airflow classes that need to be fully deserialized - using XComModel.deserialize_value() instead. - """ +_FIELD_RENDERED_MODULE = "airflow.sdk.definitions.asset." CLASSNAME = "__classname__" @@ -82,13 +75,74 @@ def _stringify_builtin_collection(classname: str, value: Any): return None +def _render_builtin_collection_field(o: dict) -> str | None: + """Render a serialized tuple, set or frozenset field the way ``repr`` renders it, else ``None``.""" + classname, data = o[CLASSNAME], o.get(DATA) + if classname not in ("builtins.tuple", "builtins.set", "builtins.frozenset") or not isinstance( + data, (list, tuple, set, frozenset) + ): + return None + items = ", ".join(_render_field(v) for v in data) + if classname == "builtins.tuple": + return f"({items},)" if len(data) == 1 else f"({items})" + if classname == "builtins.set": + return f"{{{items}}}" if data else "set()" + return f"frozenset({{{items}}})" if data else "frozenset()" + + +def _render_field(o: Any) -> str: + """Render a field value of an Airflow object the way ``repr`` would render the deserialized value.""" + if isinstance(o, str): + return repr(o) + if isinstance(o, list): + return f"[{', '.join(_render_field(v) for v in o)}]" + if isinstance(o, dict): + o = _convert(o) + if CLASSNAME in o and VERSION in o: + if (rendered := _render_builtin_collection_field(o)) is not None: + return rendered + return str(stringify(o)) + return "{" + ", ".join(f"{k!r}: {_render_field(v)}" for k, v in o.items()) + "}" + return str(stringify(o)) + + +def _stringify_airflow_object(classname: str, value: Any) -> str: + """ + Render an Airflow class from its serialized fields, e.g. ``AssetAlias(name='x', group='asset')``. + + The class is neither imported nor instantiated, so whether it is an attrs type, whose serialized + fields mirror its ``repr``, cannot be checked here. Fields are rendered only for the asset + definitions, which are attrs types. Other Airflow classes may define ``serialize()`` returning + internal state, such as the storage options of ``ObjectStoragePath``; they are shown by their + ``path`` or ``uri`` field only. + """ + short_name = classname.rsplit(".", 1)[-1] + if not classname.startswith(_FIELD_RENDERED_MODULE): + locator = next( + (value[k] for k in ("path", "uri") if isinstance(value, dict) and isinstance(value.get(k), str)), + None, + ) + return f"{short_name}({locator!r})" if locator is not None else f"{short_name}(...)" + if value is None: + args = "" + elif isinstance(value, dict): + args = ", ".join(f"{k}={_render_field(v)}" for k, v in value.items()) + elif isinstance(value, _builtin_collections): + args = ", ".join(_render_field(v) for v in value) + else: + args = _render_field(value) + return f"{short_name}({args})" + + def stringify(o: T | None) -> object: """ - Convert a serialized object to a human readable representation. + Convert a serialized object to a human readable representation without importing any class. - Matches the behavior of deserialize(full=False) exactly: + Matches the behavior of deserialize(full=False), except for Airflow classes: - Returns objects as-is for primitives, collections, plain dicts - Only returns string when encountering a serialized object (__classname__) + - Asset definitions are rendered like their ``repr``, e.g. ``AssetAlias(name='x', group='asset')`` + - Other Airflow classes are rendered by their path or URI only, e.g. ``ObjectStoragePath('s3://b/k')`` """ if o is None: return o @@ -123,12 +177,8 @@ def stringify(o: T | None) -> object: if not classname: raise TypeError("classname cannot be empty") - # Early detection for `airflow.` classes. These classes will need full deserialization, not just stringification - if isinstance(classname, str) and classname.startswith("airflow."): - raise StringifyNotSupportedError( - f"Cannot stringify Airflow class '{classname}'. " - f"Use XComModel.deserialize_value() to deserialize Airflow classes." - ) + if classname.startswith("airflow."): + return _stringify_airflow_object(classname, value) result = _stringify_builtin_collection(classname, value) if result is not None: diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py index bee30deaea1af..bbe638fbbf1fd 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py @@ -37,7 +37,7 @@ from airflow.models.team import Team from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, xcom_entity from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.sdk import DAG, AssetAlias, task_group +from airflow.sdk import DAG, Asset, AssetAlias, task_group from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.definitions._internal.loop import create_loop from airflow.sdk.execution_time.xcom import resolve_xcom_backend @@ -417,8 +417,16 @@ def test_custom_xcom_deserialize(self, params: str, expected_value: int | str, t ("xcom_value", "expected_return"), [ pytest.param(42, 42, id="jsonable"), + pytest.param("text", "text", id="str"), + pytest.param({"a": [1, 2], "b": None}, {"a": [1, 2], "b": None}, id="dict"), + pytest.param([1, "two", 3.0], [1, "two", 3.0], id="list"), pytest.param(AssetAlias("x"), "AssetAlias(name='x', group='asset')", id="nonjsonable"), pytest.param([42, AssetAlias("x")], [42, "AssetAlias(name='x', group='asset')"], id="nested"), + pytest.param( + Asset("s3://bucket/key", extra={"a": 1}), + repr(Asset("s3://bucket/key", extra={"a": 1})), + id="asset", + ), ], ) def test_stringify_false(self, test_client, xcom_value, expected_return): @@ -429,6 +437,61 @@ def test_stringify_false(self, test_client, xcom_value, expected_return): assert response.status_code == 200 assert response.json()["value"] == expected_return + @pytest.mark.parametrize( + ("xcom_value", "expected_return"), + [ + pytest.param(42, "42", id="int"), + pytest.param({"a": [1, 2]}, "{'a': [1, 2]}", id="dict"), + pytest.param(AssetAlias("x"), "AssetAlias(name='x', group='asset')", id="airflow-class"), + ], + ) + def test_stringify_true(self, test_client, xcom_value, expected_return): + XComModel.set(TEST_XCOM_KEY, xcom_value, dag_id=TEST_DAG_ID, task_id=TEST_TASK_ID, run_id=run_id) + + url = f"/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY}" + response = test_client.get(url, params={"deserialize": True, "stringify": True}) + assert response.status_code == 200 + assert response.json()["value"] == expected_return + + @mock.patch("airflow.sdk.serde.serializers.deltalake.deserialize", autospec=True) + @mock.patch("airflow.sdk.serde.deserialize", autospec=True) + def test_deserialize_renders_stored_payload_without_building_objects( + self, serde_deserialize, deltalake_deserialize, test_client + ): + """A ``deserialize=true`` read renders the stored payload; it never imports or builds the classes it names.""" + payload = { + "__classname__": "airflow.sdk.definitions.asset.AssetAlias", + "__version__": 0, + "__data__": { + "name": "x", + "group": { + "__classname__": "deltalake.table.DeltaTable", + "__version__": 1, + "__data__": {"table_uri": "s3://bucket/table", "version": 0}, + }, + }, + } + # Stored as-is, the way the Execution API stores values a task pushes. + XComModel.set( + TEST_XCOM_KEY, + json.dumps(payload), + dag_id=TEST_DAG_ID, + task_id=TEST_TASK_ID, + run_id=run_id, + serialize=False, + ) + + url = f"/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY}" + response = test_client.get(url, params={"deserialize": True}) + + assert response.status_code == 200 + assert response.json()["value"] == ( + "AssetAlias(name='x', group=deltalake.table.DeltaTable@version=1" + "(table_uri='s3://bucket/table', version=0))" + ) + serde_deserialize.assert_not_called() + deltalake_deserialize.assert_not_called() + class TestGetXComEntries(TestXComEndpoint): @pytest.fixture(autouse=True) diff --git a/airflow-core/tests/unit/serialization/test_operatorlink.py b/airflow-core/tests/unit/serialization/test_operatorlink.py index ef3a399279cfc..d5497766593ba 100644 --- a/airflow-core/tests/unit/serialization/test_operatorlink.py +++ b/airflow-core/tests/unit/serialization/test_operatorlink.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import json from types import SimpleNamespace from unittest import mock @@ -41,3 +42,23 @@ def test_operator_link_lookup_is_bounded(create_session): session.execute.assert_called_once() statement = session.execute.call_args.args[0] assert "LIMIT 1" in str(statement.compile(compile_kwargs={"literal_binds": True})) + + +@mock.patch("airflow.sdk.serde.deserialize", autospec=True) +@mock.patch("airflow.serialization.definitions.operatorlink.create_session", autospec=True) +def test_operator_link_renders_airflow_class_payload(create_session, serde_deserialize): + session = mock.create_autospec(Session, instance=True) + create_session.return_value.__enter__.return_value = session + result = mock.create_autospec(Result, instance=True) + payload = { + "__classname__": "airflow.sdk.definitions.asset.AssetAlias", + "__version__": 0, + "__data__": {"name": "x", "group": "asset"}, + } + result.first.return_value = SimpleNamespace(value=json.dumps(payload)) + session.execute.return_value = result + link = XComOperatorLink(name="Example", xcom_key="_link_example") + ti_key = TaskInstanceKey(dag_id="example", task_id="task", run_id="run", try_number=1, map_index=-1) + + assert link.get_link(None, ti_key=ti_key) == "AssetAlias(name='x', group='asset')" + serde_deserialize.assert_not_called() diff --git a/airflow-core/tests/unit/serialization/test_stringify.py b/airflow-core/tests/unit/serialization/test_stringify.py index 6f9c817db774b..c88d77c8f2eba 100644 --- a/airflow-core/tests/unit/serialization/test_stringify.py +++ b/airflow-core/tests/unit/serialization/test_stringify.py @@ -19,10 +19,12 @@ from dataclasses import dataclass from typing import ClassVar +from unittest import mock import pytest from airflow._shared.module_loading import qualname +from airflow.sdk import Asset, AssetAlias from airflow.sdk.serde import serialize from airflow.serialization.stringify import CLASSNAME, VERSION, stringify @@ -42,6 +44,11 @@ class V: c: int +@dataclass +class Thing: + f: object + + class TestStringify: # Existing ported over tests def test_stringify(self): @@ -253,9 +260,83 @@ def test_stringify_dict_with_nested_serialized(self): result = stringify(e) assert result == {"key1": "(1,2)", "key2": "value"} - def test_error_thrown_for_airflow_classes(self): - from airflow.sdk import AssetAlias + @pytest.mark.parametrize( + "value", + [ + pytest.param(AssetAlias("x"), id="asset-alias"), + pytest.param(Asset("s3://bucket/key", extra={"a": [1, "b"]}), id="asset-with-nested-class"), + ], + ) + def test_stringify_airflow_class_reads_like_repr(self, value): + assert stringify(serialize(value)) == repr(value) - e = serialize(AssetAlias("x")) - with pytest.raises(ValueError, match="Cannot stringify Airflow class"): - stringify(e) + @pytest.mark.parametrize( + ("data", "expected"), + [ + pytest.param(None, "Thing()", id="no-data"), + pytest.param("x", "Thing('x')", id="str"), + pytest.param(3, "Thing(3)", id="int"), + pytest.param([1, "a", {"k": "v"}], "Thing(1, 'a', {'k': 'v'})", id="list"), + pytest.param( + {"inner": {"__type": "dict", "__var": {"k": None}}}, "Thing(inner={'k': None})", id="old-dict" + ), + ], + ) + def test_stringify_airflow_class_data_shapes(self, data, expected): + assert ( + stringify({CLASSNAME: "airflow.sdk.definitions.asset.Thing", VERSION: 0, "__data__": data}) + == expected + ) + + @pytest.mark.parametrize( + "field", + [ + pytest.param(("a", "b"), id="tuple"), + pytest.param(("a",), id="single-tuple"), + pytest.param((), id="empty-tuple"), + pytest.param((1, ("x", (2,))), id="nested-tuple"), + pytest.param({"a"}, id="set"), + pytest.param(set(), id="empty-set"), + pytest.param(frozenset({"a"}), id="frozenset"), + pytest.param(frozenset(), id="empty-frozenset"), + pytest.param([("a",), {"k": ("v", 1)}], id="tuples-in-list-and-dict"), + ], + ) + def test_stringify_airflow_class_collection_fields_read_like_repr(self, field): + e = { + CLASSNAME: "airflow.sdk.definitions.asset.Thing", + VERSION: 0, + "__data__": {"f": serialize(field)}, + } + assert stringify(e) == repr(Thing(field)) + + @pytest.mark.parametrize( + ("data", "expected"), + [ + pytest.param( + {"path": "s3://conn@bucket/key", "conn_id": "conn", "kwargs": {"key": "k", "token": "t"}}, + "ObjectStoragePath('s3://conn@bucket/key')", + id="path", + ), + pytest.param( + {"uri": "s3://b", "options": {"token": "t"}}, "ObjectStoragePath('s3://b')", id="uri" + ), + pytest.param( + {"conn_id": "conn", "kwargs": {"anon": False}}, "ObjectStoragePath(...)", id="no-locator" + ), + pytest.param("s3://b", "ObjectStoragePath(...)", id="non-dict"), + ], + ) + def test_stringify_airflow_class_with_custom_serialization_is_compact(self, data, expected): + e = {CLASSNAME: "airflow.sdk.io.path.ObjectStoragePath", VERSION: 1, "__data__": data} + assert stringify(e) == expected + + @mock.patch("importlib.import_module", autospec=True) + def test_stringify_airflow_class_does_not_import_nested_class(self, import_module): + e = { + CLASSNAME: "airflow.sdk.definitions.asset.Outer", + VERSION: 1, + "__data__": {"inner": {CLASSNAME: "other.mod.Inner", VERSION: 2, "__data__": {"uri": "s3://b"}}}, + } + assert stringify(e) == "Outer(inner=other.mod.Inner@version=2(uri='s3://b'))" + import_module.assert_not_called()