fix(asset): handle asset ref cases

Author: Lee-W

task-sdk/src/airflow/sdk/execution_time/context.py

@@ -428,9 +428,9 @@ def __hash__(self):
 
 
 class _AssetRefResolutionMixin:
-    _asset_ref_cache: dict[AssetRef, AssetUniqueKey] = {}
+    _asset_ref_cache: dict[AssetRef, tuple[AssetUniqueKey, dict[str, JsonValue]]] = {}
 
-    def _resolve_asset_ref(self, ref: AssetRef) -> AssetUniqueKey:
+    def _resolve_asset_ref(self, ref: AssetRef) -> tuple[AssetUniqueKey, dict[str, JsonValue]]:
         with contextlib.suppress(KeyError):
             return self._asset_ref_cache[ref]
 
@@ -445,8 +445,8 @@ def _resolve_asset_ref(self, ref: AssetRef) -> AssetUniqueKey:
             raise TypeError(f"Unimplemented asset ref: {type(ref)}")
         unique_key = AssetUniqueKey.from_asset(asset)
         for ref in refs_to_cache:
-            self._asset_ref_cache[ref] = unique_key
-        return unique_key
+            self._asset_ref_cache[ref] = (unique_key, asset.extra)
+        return (unique_key, asset.extra)
 
     # TODO: This is temporary to avoid code duplication between here & airflow/models/taskinstance.py
     @staticmethod
@@ -491,15 +491,16 @@ def add(self, asset: Asset | AssetRef, extra: dict[str, JsonValue] | None = None
             return
 
         if isinstance(asset, AssetRef):
-            asset_key = self._resolve_asset_ref(asset)
+            asset_key, asset_extra = self._resolve_asset_ref(asset)
         else:
             asset_key = AssetUniqueKey.from_asset(asset)
+            asset_extra = asset.extra
 
         asset_alias_name = self.key.name
         event = AssetAliasEvent(
             source_alias_name=asset_alias_name,
             dest_asset_key=asset_key,
-            dest_asset_extra=asset.extra if isinstance(asset, Asset) else {},
+            dest_asset_extra=asset_extra,
             extra=extra or {},
         )
         self.asset_alias_events.append(event)
@@ -560,7 +561,7 @@ def __getitem__(self, key: Asset | AssetAlias | AssetRef) -> OutletEventAccessor
         elif isinstance(key, AssetAlias):
             hashable_key = AssetAliasUniqueKey.from_asset_alias(key)
         elif isinstance(key, AssetRef):
-            hashable_key = self._resolve_asset_ref(key)
+            hashable_key, _ = self._resolve_asset_ref(key)
         else:
             raise TypeError(f"Key should be either an asset or an asset alias, not {type(key)}")
 
@@ -769,7 +770,7 @@ def __getitem__(self, key: Asset | AssetAlias | AssetRef) -> Sequence[AssetEvent
         if isinstance(key, Asset):
             hashable_key = AssetUniqueKey.from_asset(key)
         elif isinstance(key, AssetRef):
-            hashable_key = self._resolve_asset_ref(key)
+            hashable_key, _ = self._resolve_asset_ref(key)
         elif isinstance(key, AssetAlias):
             hashable_key = AssetAliasUniqueKey.from_asset_alias(key)
         else:

task-sdk/tests/task_sdk/execution_time/test_context.py

@@ -346,12 +346,13 @@ def test_nested_context(self):
 
 class TestOutletEventAccessor:
     @pytest.mark.parametrize(
-        "add_arg",
+        "add_args",
         [
-            Asset("name", "uri"),
-            Asset.ref(name="name"),
-            Asset.ref(uri="uri"),
+            (Asset("name", "uri", extra={"extra": "from asset itself"}), {"extra": "from event"}),
+            (Asset.ref(name="name"), {"extra": "from event"}),
+            (Asset.ref(uri="uri"), {"extra": "from event"}),
         ],
+        ids=["asset", "asset name ref", "asset uri ref"],
     )
     @pytest.mark.parametrize(
         ("key", "asset_alias_events"),
@@ -363,18 +364,21 @@ class TestOutletEventAccessor:
                     AssetAliasEvent(
                         source_alias_name="test_alias",
                         dest_asset_key=AssetUniqueKey(name="name", uri="uri"),
-                        dest_asset_extra={},
-                        extra={},
+                        dest_asset_extra={"extra": "from asset itself"},
+                        extra={"extra": "from event"},
                     )
                 ],
             ),
         ),
+        ids=["inactive asset", "active asset"],
     )
-    def test_add(self, add_arg, key, asset_alias_events, mock_supervisor_comms):
-        mock_supervisor_comms.send.return_value = AssetResponse(name="name", uri="uri", group="")
+    def test_add(self, add_args, key, asset_alias_events, mock_supervisor_comms):
+        mock_supervisor_comms.send.return_value = AssetResponse(
+            name="name", uri="uri", group="", extra={"extra": "from asset itself"}
+        )
 
         outlet_event_accessor = OutletEventAccessor(key=key, extra={})
-        outlet_event_accessor.add(add_arg)
+        outlet_event_accessor.add(*add_args)
         assert outlet_event_accessor.asset_alias_events == asset_alias_events
 
     @pytest.mark.parametrize(