From 4600076e678643ae99bd13f00b80d6e8003ba2e2 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Mon, 19 Oct 2026 16:56:31 +0200 Subject: [PATCH] [v3-4-test] Only rebuild built-in value types when converting next_kwargs for older workers (#73577) * Only rebuild built-in value types when converting next_kwargs for older workers When serving a task instance to a worker on an older Execution API version, the version converter decodes the stored serde next_kwargs and re-encodes them in the legacy format. Decoding reconstructs every value in the API server, including instances of arbitrary classes, although only the worker ever needs those objects. Only decode payloads whose values are built-in value types - the types the legacy format can carry anyway - and pass anything else through as stored. Generated-by: Claude Opus 5 * Keep converting Asset values in next_kwargs for older workers Task SDK 1.1.x always decodes next_kwargs with BaseSerialization.deserialize, so passing a serde payload through unchanged makes the deferred task crash on resume with KeyError('__var'). That broke resuming with Asset or AssetAlias values, which the legacy format supports natively. Assets and asset aliases are now translated structurally from their encoded fields into the legacy envelope, so nothing beyond built-in value types is instantiated in the API server. Any other class makes the task resume through the trigger-failure path, so older workers fail it with a TaskDeferralError naming the type instead of an opaque KeyError. Payloads already in the legacy format are still passed through. Generated-by: Claude Opus 5 * Translate tuples and sets of assets in next_kwargs for older workers A tuple or set holding an Asset or AssetAlias failed the built-in value check because of the nested asset, so its serde envelope was treated as unconvertible and the task resumed through the trigger-failure path, even though the legacy format carries tuples and sets of assets. Tuple, set and frozenset envelopes are now translated structurally into the legacy tuple/set encoding, with their items converted recursively and set items sorted when comparable, as BaseSerialization.serialize does. The legacy format has no frozenset, so it is sent as a set. No class beyond the built-in value types is constructed. Generated-by: Claude Opus 5 * Only treat exact legacy envelopes as legacy and fail clearly on frozensets of assets A dict was taken for a legacy BaseSerialization envelope as soon as it had a "__type" key, so user data such as {"__type": "record", "asset": ...} resumed the task through the failure path, and a top-level dict with both "__type" and "__var" plus other keys was passed through unconverted. Only a dict whose keys are exactly "__type" and "__var" is now treated as an envelope; any other dict is converted recursively. Frozensets holding assets were sent as legacy sets, which older workers cannot decode when the frozenset sits inside a set. The legacy format has no frozenset, so these now resume through the trigger-failure path with a clear message. Frozensets of built-in values keep the built-in path. Generated-by: Claude Opus 5 * Never deserialize non-builtin classes behind legacy-looking dicts in next_kwargs The built-in value check followed only "__var" of any dict holding "__type" and "__var", while serde's decoder unwraps such a dict and then decodes the inner dict as a class envelope without re-checking it. A class envelope placed inside a legacy dict envelope was therefore passed as built-in and constructed in the API server, and other keys of such a dict were silently dropped. The check now mirrors the decoder and fails safe: only exact envelopes count as envelopes, a dict envelope's content must be a plain dict, class envelopes must have exactly the serde keys, and any other dict holding a reserved key is left to the structural translation. A top-level payload already in the legacy format is passed through only when every dict in it is an envelope the legacy decoder can read; otherwise the task resumes through the trigger-failure path. Legacy payloads are no longer decoded in the API server at all. Generated-by: Claude Opus 5 * Keep asset access control when translating next_kwargs for older workers The structural Asset translation copied only name, uri, group and extra, so a non-default access_control was dropped, while encode_asset_like() emits it whenever producer_teams is set, consumer_teams is not None or allow_global is false. The translation now builds the access_control dict exactly as encode_asset_like() does, from the serde data of AssetAccessControl, with nothing constructed. An Asset or AssetAlias whose serde data holds a field or shape that cannot be mapped with confidence (unknown keys, watchers, class envelopes in extra, malformed access control) resumes through the trigger-failure path rather than losing data silently. Generated-by: Claude Opus 5 * Add newsfragment for older-worker next_kwargs conversion Generated-by: Claude Opus 5 * Tell users how to avoid next_kwargs failures on older workers Generated-by: Claude Opus 5 * Name Task SDK 1.1.x as the affected worker versions (cherry picked from commit b0239d215e1f7c7f244a8e3fa7c3bc36fc729fe0) Co-authored-by: Jarek Potiuk Generated-by: Claude Opus 5 --- .../newsfragments/73577.significant.rst | 5 + .../execution_api/versions/v2026_04_06.py | 224 +++++++++++- .../v2026_04_06/test_task_instances.py | 339 ++++++++++++++++++ 3 files changed, 562 insertions(+), 6 deletions(-) create mode 100644 airflow-core/newsfragments/73577.significant.rst diff --git a/airflow-core/newsfragments/73577.significant.rst b/airflow-core/newsfragments/73577.significant.rst new file mode 100644 index 0000000000000..fd83f53a8adc4 --- /dev/null +++ b/airflow-core/newsfragments/73577.significant.rst @@ -0,0 +1,5 @@ +Older workers can no longer resume deferred tasks whose ``next_kwargs`` hold custom classes + +When a deferred task resumes on a worker with an older Task SDK, the API server no longer reconstructs arbitrary classes from the stored ``next_kwargs`` to convert them to the legacy format. Built-in value types, assets and asset aliases (including asset access control) are still converted, without constructing them. A deferred task whose ``next_kwargs`` hold any other class now fails on such a worker with an error naming the type and asking to upgrade the Task SDK, instead of crashing on a payload it cannot decode. + +**What to do:** if you run workers with ``apache-airflow-task-sdk`` 1.1.x or older and your deferrable operators or triggers pass instances of custom classes in ``next_kwargs`` (for example via ``self.defer(..., kwargs=...)`` or in the trigger event payload), upgrade the Task SDK on those workers to 1.2.0 or later. Newer workers receive ``next_kwargs`` unchanged and are not affected. Until all workers are upgraded, pass only built-in values, assets and asset aliases in ``next_kwargs`` of tasks that may resume on older workers. diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py index 4710de4f66549..73566d44c769d 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py @@ -17,6 +17,7 @@ from __future__ import annotations +import contextlib from typing import Any from cadwyn import ResponseInfo, VersionChange, convert_response_to_previous_version_for, endpoint, schema @@ -109,6 +110,207 @@ def ensure_start_date_in_dag_run(response: ResponseInfo) -> None: # type: ignor response.body["start_date"] = response.body.get("run_after") +# serde class names whose deserializers rebuild plain value types. Everything else (attrs, +# dataclasses, pydantic models, classes with their own ``deserialize``) is rebuilt by calling into +# the named class. +_BUILTIN_VALUE_CLASSNAMES = frozenset( + { + "builtins.frozenset", + "builtins.set", + "builtins.tuple", + "datetime.date", + "datetime.datetime", + "datetime.timedelta", + "decimal.Decimal", + "pendulum.date.Date", + "pendulum.datetime.DateTime", + "pendulum.tz.timezone.FixedTimezone", + "pendulum.tz.timezone.Timezone", + "uuid.UUID", + "zoneinfo.ZoneInfo", + } +) + + +def _holds_only_builtin_values(value: Any) -> bool: + """ + Whether serde ``deserialize`` would rebuild nothing but built-in value types from an encoded value. + + Mirrors how the decoder reads each dict, and fails safe: a dict it could misread, such as one + holding envelope keys alongside other keys, counts as not built-in. + """ + from airflow._shared.serialization import ( + CLASSNAME, + DATA, + OLD_DATA, + OLD_DICT, + OLD_TYPE, + OLD_TYPE_TO_FULL_QUALNAME, + VERSION, + ) + + reserved_keys = {CLASSNAME, VERSION, DATA, OLD_TYPE, OLD_DATA} + + def is_plain_builtin_dict(d: Any) -> bool: + return ( + isinstance(d, dict) + and not d.keys() & reserved_keys + and all(_holds_only_builtin_values(item) for item in d.values()) + ) + + if isinstance(value, list): + return all(_holds_only_builtin_values(item) for item in value) + if not isinstance(value, dict): + return True + if _is_legacy_envelope(value): + if value[OLD_TYPE] == OLD_DICT: + # The decoder unwraps this envelope without checking the inner dict for another one. + return is_plain_builtin_dict(value[OLD_DATA]) + return OLD_TYPE_TO_FULL_QUALNAME.get( + value[OLD_TYPE] + ) in _BUILTIN_VALUE_CLASSNAMES and _holds_only_builtin_values(value[OLD_DATA]) + if CLASSNAME in value: + return ( + value.keys() == {CLASSNAME, VERSION, DATA} + and value[CLASSNAME] in _BUILTIN_VALUE_CLASSNAMES + and _holds_only_builtin_values(value[DATA]) + ) + return is_plain_builtin_dict(value) + + +_ASSET_CLASSNAME = "airflow.sdk.definitions.asset.Asset" +_ASSET_ALIAS_CLASSNAME = "airflow.sdk.definitions.asset.AssetAlias" +_ASSET_ACCESS_CONTROL_CLASSNAME = "airflow.sdk.definitions.asset.access_control.AssetAccessControl" + + +class _UnconvertibleNextKwargsError(ValueError): + """next_kwargs hold a value that cannot be translated to the legacy BaseSerialization format.""" + + def __init__(self, what: str) -> None: + super().__init__( + f"The deferred task cannot be resumed: its next_kwargs hold {what}, " + "which cannot be sent to this worker's Task SDK version. Upgrade the Task SDK on the worker." + ) + + +def _is_legacy_envelope(value: Any) -> bool: + """Whether a value is a BaseSerialization ``{"__type": ..., "__var": ...}`` envelope.""" + from airflow._shared.serialization import OLD_DATA, OLD_TYPE + + return isinstance(value, dict) and value.keys() == {OLD_TYPE, OLD_DATA} + + +def _is_well_formed_legacy(value: Any) -> bool: + """Whether every dict in a BaseSerialization-encoded value is an envelope its decoder can read.""" + from airflow._shared.serialization import OLD_DATA, OLD_DICT, OLD_TYPE + from airflow.serialization.enums import DagAttributeTypes as DAT + + if isinstance(value, list): + return all(_is_well_formed_legacy(item) for item in value) + if not isinstance(value, dict): + return True + if not _is_legacy_envelope(value): + return False + data = value[OLD_DATA] + if value[OLD_TYPE] == OLD_DICT: + return isinstance(data, dict) and all(_is_well_formed_legacy(item) for item in data.values()) + if value[OLD_TYPE] in (DAT.TUPLE, DAT.SET): + return isinstance(data, list) and _is_well_formed_legacy(data) + return True + + +def _convert_to_legacy(value: Any) -> Any: + """ + Translate a serde-encoded value to the legacy BaseSerialization encoding. + + Assets and asset aliases are translated structurally from their encoded fields, so no class + beyond the built-in value types is ever instantiated here. + """ + from airflow._shared.serialization import CLASSNAME, DATA, OLD_TYPE + from airflow.sdk.serde import deserialize + from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding + from airflow.serialization.serialized_objects import BaseSerialization + + if _holds_only_builtin_values(value): + return BaseSerialization.serialize(deserialize(value)) + if isinstance(value, list): + return [_convert_to_legacy(item) for item in value] + if CLASSNAME not in value: + if _is_legacy_envelope(value): + raise _UnconvertibleNextKwargsError(f"a value of type {str(value[OLD_TYPE])!r}") + return { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: {key: _convert_to_legacy(item) for key, item in value.items()}, + } + classname = value[CLASSNAME] + data = value.get(DATA) + # No frozenset: the legacy format has none, and sending a set instead breaks nesting it in a set. + # Frozensets of built-in values never get here; they are handled with the other built-ins. + collection_types = {"builtins.tuple": DAT.TUPLE, "builtins.set": DAT.SET} + if classname in collection_types and isinstance(data, list): + items = [_convert_to_legacy(item) for item in data] + if collection_types[classname] == DAT.SET: + # Mirrors BaseSerialization.serialize, which sorts set items when they are comparable. + with contextlib.suppress(TypeError): + items = sorted(items) + return {Encoding.TYPE: collection_types[classname], Encoding.VAR: items} + if classname == _ASSET_CLASSNAME: + return {Encoding.TYPE: DAT.ASSET, Encoding.VAR: _translate_asset(data)} + if classname == _ASSET_ALIAS_CLASSNAME and isinstance(data, dict) and data.keys() == {"name", "group"}: + return {Encoding.TYPE: DAT.ASSET_ALIAS, Encoding.VAR: {"name": data["name"], "group": data["group"]}} + raise _UnconvertibleNextKwargsError(f"a value of type {str(classname)!r}") + + +def _is_plain_json(value: Any) -> bool: + """Whether a serde-encoded value holds only JSON values, with no class envelope anywhere.""" + from airflow._shared.serialization import CLASSNAME + + if isinstance(value, list): + return all(_is_plain_json(item) for item in value) + if isinstance(value, dict): + return CLASSNAME not in value and all(_is_plain_json(item) for item in value.values()) + return value is None or isinstance(value, (str, int, float, bool)) + + +def _is_team_list(value: Any) -> bool: + return isinstance(value, list) and all(isinstance(team, str) for team in value) + + +def _translate_asset(data: Any) -> dict[str, Any]: + """Build the legacy encoding of an Asset from its serde data, field for field as ``encode_asset_like``.""" + from airflow._shared.serialization import CLASSNAME, DATA + + unconvertible = _UnconvertibleNextKwargsError(f"a value of type {_ASSET_CLASSNAME!r}") + required = {"name", "uri", "group", "extra"} + if ( + not isinstance(data, dict) + or not required <= data.keys() <= required | {"watchers", "access_control"} + or data.get("watchers") + or not _is_plain_json(data["extra"]) + ): + raise unconvertible + var = {key: data[key] for key in ("name", "uri", "group", "extra")} + + if (envelope := data.get("access_control")) is None: + return var + if not isinstance(envelope, dict) or envelope.get(CLASSNAME) != _ASSET_ACCESS_CONTROL_CLASSNAME: + raise unconvertible + ac = envelope.get(DATA) + if ( + not isinstance(ac, dict) + or ac.keys() != {"producer_teams", "consumer_teams", "allow_global"} + or not _is_team_list(ac["producer_teams"]) + or not (ac["consumer_teams"] is None or _is_team_list(ac["consumer_teams"])) + or not isinstance(ac["allow_global"], bool) + ): + raise unconvertible + if ac["producer_teams"] or ac["consumer_teams"] is not None or not ac["allow_global"]: + var["access_control"] = {"producer_teams": ac["producer_teams"], "allow_global": ac["allow_global"]} + if ac["consumer_teams"] is not None: + var["access_control"]["consumer_teams"] = ac["consumer_teams"] + return var + + class ModifyDeferredTaskKwargsToJsonValue(VersionChange): """Change the types of `trigger_kwargs` and `next_kwargs` in TIDeferredStatePayload to JsonValue.""" @@ -132,21 +334,31 @@ def convert_next_kwargs_to_base_serialization(response: ResponseInfo) -> None: timedelta, etc.), then re-serialize with BaseSerialization so old workers get proper typed values instead of raw {"__classname__": ...} dicts. """ + from airflow.models.trigger import TRIGGER_FAIL_REPR + from airflow.serialization.serialized_objects import BaseSerialization + next_kwargs = response.body.get("next_kwargs") if next_kwargs is None: return - from airflow.sdk.serde import deserialize - from airflow.serialization.serialized_objects import BaseSerialization - + # Only built-in value types are rebuilt here; assets and asset aliases are translated + # structurally. Anything else would be reconstructed by calling into its class in the API + # server, so the task is resumed through the trigger-failure path instead, which old + # workers report as a TaskDeferralError carrying this message. try: - plain = deserialize(next_kwargs) + if _is_legacy_envelope(next_kwargs): + if not _is_well_formed_legacy(next_kwargs): + raise _UnconvertibleNextKwargsError("a malformed BaseSerialization-encoded value") + # Already in BaseSerialization format (rolling upgrade, old data in DB) + return + response.body["next_kwargs"] = _convert_to_legacy(next_kwargs) + except _UnconvertibleNextKwargsError as e: + response.body["next_method"] = TRIGGER_FAIL_REPR + response.body["next_kwargs"] = BaseSerialization.serialize({"error": str(e)}) except (ImportError, KeyError, AttributeError, TypeError): # Already in BaseSerialization format (rolling upgrade, old data in DB) return - response.body["next_kwargs"] = BaseSerialization.serialize(plain) - class RemoveUpstreamMapIndexesField(VersionChange): """Remove upstream_map_indexes field from TIRunContext - now computed by Task SDK.""" diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_06/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_06/test_task_instances.py index d855291727756..8eac551a1a564 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_06/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_06/test_task_instances.py @@ -17,9 +17,13 @@ from __future__ import annotations +from dataclasses import dataclass + import pytest from airflow._shared.timezones import timezone +from airflow.sdk import Asset, AssetAlias +from airflow.sdk.definitions.asset.access_control import AssetAccessControl from airflow.serialization.serialized_objects import BaseSerialization from airflow.utils.state import DagRunState, State @@ -40,6 +44,11 @@ } +@dataclass +class _ResumePayload: + value: int + + @pytest.fixture def old_ver_client(client): """Last released execution API before nullable DagRun.start_date (2026-04-06 bundle).""" @@ -191,6 +200,336 @@ def test_old_version_deserializes_complex_types(self, old_ver_client, session, c # datetime must come back as a datetime, not a {"__classname__": ...} dict assert result["event"] == TIMESTAMP + def test_old_version_translates_assets_without_constructing_them( + self, old_ver_client, session, create_task_instance, monkeypatch + ): + """Assets and asset aliases reach old workers in the legacy format without being rebuilt here.""" + from airflow.sdk.serde import serialize as serde_serialize + + original = { + "assets": [Asset("resume-data", uri="s3://bucket/key", group="g", extra={"k": 1})], + "alias": AssetAlias("resume-alias"), + "event": TIMESTAMP, + } + expected = BaseSerialization.serialize(original) + + ti = create_task_instance( + task_id="test_next_kwargs_assets", + state=State.QUEUED, + session=session, + start_date=TIMESTAMP, + ) + ti.next_method = "execute_complete" + ti.next_kwargs = serde_serialize(original) + session.commit() + + constructed = [] + for cls in (Asset, AssetAlias): + original_init = cls.__init__ + + def spy(self, *args, _original_init=original_init, **kwargs): + constructed.append(type(self)) + _original_init(self, *args, **kwargs) + + monkeypatch.setattr(cls, "__init__", spy) + + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + + assert response.status_code == 200 + assert constructed == [] + assert response.json()["next_method"] == "execute_complete" + next_kwargs = response.json()["next_kwargs"] + assert next_kwargs == expected + result = BaseSerialization.deserialize(next_kwargs) + assert [(a.name, a.uri, a.group, a.extra) for a in result["assets"]] == [ + ("resume-data", "s3://bucket/key", "g", {"k": 1}) + ] + assert (result["alias"].name, result["alias"].group) == ("resume-alias", "asset") + assert result["event"] == TIMESTAMP + + @pytest.mark.parametrize( + ("access_control", "expected_access_control"), + [ + pytest.param( + AssetAccessControl(producer_teams=["team_a"], allow_global=False), + {"producer_teams": ["team_a"], "allow_global": False}, + id="producers-only", + ), + pytest.param( + AssetAccessControl(consumer_teams=["team_b"]), + {"producer_teams": [], "allow_global": True, "consumer_teams": ["team_b"]}, + id="consumers", + ), + ], + ) + def test_old_version_keeps_asset_access_control( + self, + old_ver_client, + session, + create_task_instance, + monkeypatch, + access_control, + expected_access_control, + ): + """Asset access restrictions survive the translation without the classes being rebuilt here.""" + from airflow.sdk.serde import serialize as serde_serialize + + original = {"asset": Asset("restricted", access_control=access_control)} + ti = create_task_instance( + task_id="test_next_kwargs_access_control", + state=State.QUEUED, + session=session, + start_date=TIMESTAMP, + ) + ti.next_method = "execute_complete" + ti.next_kwargs = serde_serialize(original) + session.commit() + + constructed = [] + for cls in (Asset, AssetAccessControl): + original_init = cls.__init__ + + def spy(self, *args, _original_init=original_init, **kwargs): + constructed.append(type(self)) + _original_init(self, *args, **kwargs) + + monkeypatch.setattr(cls, "__init__", spy) + + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + + assert response.status_code == 200 + assert constructed == [] + assert response.json()["next_method"] == "execute_complete" + next_kwargs = response.json()["next_kwargs"] + assert next_kwargs == BaseSerialization.serialize(original) + assert BaseSerialization.deserialize(next_kwargs)["asset"].access_control == expected_access_control + + @pytest.mark.parametrize( + "assets", + [ + pytest.param((Asset("first"), Asset("second")), id="tuple-of-assets"), + pytest.param({Asset("first"), Asset("second")}, id="set-of-assets"), + pytest.param({AssetAlias("resume-alias")}, id="set-with-alias"), + ], + ) + def test_old_version_translates_asset_collections_without_constructing_them( + self, old_ver_client, session, create_task_instance, monkeypatch, assets + ): + """Tuples and sets of assets reach old workers in the legacy format without being rebuilt here.""" + from airflow.sdk.serde import serialize as serde_serialize + + original = {"assets": assets} + expected = BaseSerialization.serialize(original) + + ti = create_task_instance( + task_id="test_next_kwargs_asset_collections", + state=State.QUEUED, + session=session, + start_date=TIMESTAMP, + ) + ti.next_method = "execute_complete" + ti.next_kwargs = serde_serialize(original) + session.commit() + + constructed = [] + for cls in (Asset, AssetAlias): + original_init = cls.__init__ + + def spy(self, *args, _original_init=original_init, **kwargs): + constructed.append(type(self)) + _original_init(self, *args, **kwargs) + + monkeypatch.setattr(cls, "__init__", spy) + + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + + assert response.status_code == 200 + assert constructed == [] + assert response.json()["next_method"] == "execute_complete" + next_kwargs = response.json()["next_kwargs"] + assert next_kwargs == expected + # Decoded per item: core's decoder builds unhashable SerializedAssetAlias, so it cannot rebuild + # a set of aliases; older workers decode into hashable Task SDK classes. + decoded = [BaseSerialization.deserialize(item) for item in next_kwargs["__var"]["assets"]["__var"]] + assert sorted((type(a).__name__, a.name) for a in decoded) == sorted( + (f"Serialized{type(a).__name__}", a.name) for a in assets + ) + + @pytest.mark.parametrize( + "original", + [ + pytest.param( + {"__type": "record", "__var": "resume", "asset": Asset("first")}, id="top-level-envelope-keys" + ), + pytest.param({"payload": {"__type": "record", "asset": Asset("first")}}, id="nested-type-key"), + pytest.param( + {"payload": {"__type": "dict", "__var": {"x": 1}, "asset": Asset("first")}}, + id="nested-envelope-keys", + ), + ], + ) + def test_old_version_converts_dicts_with_legacy_keys( + self, old_ver_client, session, create_task_instance, original + ): + """Dicts that merely contain ``__type``/``__var`` keys are user data, not legacy envelopes.""" + from airflow.sdk.serde import serialize as serde_serialize + + ti = create_task_instance( + task_id="test_next_kwargs_legacy_keys", + state=State.QUEUED, + session=session, + start_date=TIMESTAMP, + ) + ti.next_method = "execute_complete" + ti.next_kwargs = serde_serialize(original) + session.commit() + + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + + assert response.status_code == 200 + assert response.json()["next_method"] == "execute_complete" + next_kwargs = response.json()["next_kwargs"] + assert next_kwargs == BaseSerialization.serialize(original) + assert BaseSerialization.deserialize(next_kwargs).keys() == original.keys() + + @pytest.mark.parametrize( + "assets", + [ + pytest.param(frozenset({Asset("first")}), id="frozenset"), + pytest.param({frozenset({Asset("first")})}, id="frozenset-in-set"), + ], + ) + def test_old_version_fails_task_for_frozenset_of_assets( + self, old_ver_client, session, create_task_instance, monkeypatch, assets + ): + """The legacy format has no frozenset, so a frozenset of assets resumes through a clear failure.""" + from airflow.sdk.serde import serialize as serde_serialize + + ti = create_task_instance( + task_id="test_next_kwargs_frozenset_assets", + state=State.QUEUED, + session=session, + start_date=TIMESTAMP, + ) + ti.next_method = "execute_complete" + ti.next_kwargs = serde_serialize({"assets": assets}) + session.commit() + + constructed = [] + original_init = Asset.__init__ + + def spy(self, *args, **kwargs): + constructed.append(True) + original_init(self, *args, **kwargs) + + monkeypatch.setattr(Asset, "__init__", spy) + + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + + assert response.status_code == 200 + assert constructed == [] + assert response.json()["next_method"] == "__fail__" + error = BaseSerialization.deserialize(response.json()["next_kwargs"])["error"] + assert "'builtins.frozenset'" in error + + @pytest.mark.parametrize("nested", [False, True], ids=["top-level", "nested"]) + def test_old_version_does_not_rebuild_classes_behind_legacy_envelopes( + self, old_ver_client, session, create_task_instance, monkeypatch, nested + ): + """A class envelope inside a legacy dict envelope is never decoded in the API server.""" + payload = { + "__type": "dict", + "__var": { + "__type": "dict", + "__var": {}, + "__classname__": "airflow.sdk.definitions.asset.AssetAlias", + "__version__": 0, + "__data__": {"name": "hidden", "group": "asset"}, + }, + } + ti = create_task_instance( + task_id="test_next_kwargs_hidden_class", + state=State.QUEUED, + session=session, + start_date=TIMESTAMP, + ) + ti.next_method = "execute_complete" + ti.next_kwargs = {"event": payload} if nested else payload + session.commit() + + constructed = [] + for cls in (Asset, AssetAlias): + original_init = cls.__init__ + + def spy(self, *args, _original_init=original_init, **kwargs): + constructed.append(type(self)) + _original_init(self, *args, **kwargs) + + monkeypatch.setattr(cls, "__init__", spy) + + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + + assert response.status_code == 200 + assert constructed == [] + assert response.json()["next_method"] == "__fail__" + error = BaseSerialization.deserialize(response.json()["next_kwargs"])["error"] + assert error.startswith("The deferred task cannot be resumed") + + def test_old_version_fails_task_for_unconvertible_values( + self, old_ver_client, session, create_task_instance, monkeypatch + ): + """Values of other classes are not rebuilt; old workers get a trigger failure naming the type.""" + from airflow.sdk.serde import serialize as serde_serialize + + ti = create_task_instance( + task_id="test_next_kwargs_unconvertible", + state=State.QUEUED, + session=session, + start_date=TIMESTAMP, + ) + ti.next_method = "execute_complete" + ti.next_kwargs = serde_serialize({"payload": _ResumePayload(value=1), "event": TIMESTAMP}) + session.commit() + + constructed = [] + original_init = _ResumePayload.__init__ + + def spy(self, *args, **kwargs): + constructed.append(True) + original_init(self, *args, **kwargs) + + monkeypatch.setattr(_ResumePayload, "__init__", spy) + + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + + assert response.status_code == 200 + assert constructed == [] + assert response.json()["next_method"] == "__fail__" + error = BaseSerialization.deserialize(response.json()["next_kwargs"])["error"] + assert f"{__name__}._ResumePayload" in error + + def test_old_version_passes_base_serialization_assets_through( + self, old_ver_client, session, create_task_instance + ): + """Rolling upgrade: BaseSerialization-format assets already in the DB are returned as stored.""" + + stored = BaseSerialization.serialize({"asset": Asset("resume-data")}) + ti = create_task_instance( + task_id="test_next_kwargs_base_assets", + state=State.QUEUED, + session=session, + start_date=TIMESTAMP, + ) + ti.next_method = "execute_complete" + ti.next_kwargs = stored + session.commit() + + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + + assert response.status_code == 200 + assert response.json()["next_method"] == "execute_complete" + assert response.json()["next_kwargs"] == stored + def test_old_version_handles_already_base_serialization_in_db( self, old_ver_client, session, create_task_instance ):