Support using asset ref to emit alias event (#48922)

Author: uranusjr

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

@@ -371,22 +371,27 @@ def __getitem__(self, key: Asset | AssetAlias | AssetRef) -> Sequence[AssetEvent
 
 
 @attrs.define
-class OutletEventAccessor:
+class OutletEventAccessor(_AssetRefResolutionMixin):
     """Wrapper to access an outlet asset event in template."""
 
     key: BaseAssetUniqueKey
     extra: dict[str, Any] = attrs.Factory(dict)
     asset_alias_events: list[AssetAliasEvent] = attrs.field(factory=list)
 
-    def add(self, asset: Asset, extra: dict[str, Any] | None = None) -> None:
+    def add(self, asset: Asset | AssetRef, extra: dict[str, Any] | None = None) -> None:
         """Add an AssetEvent to an existing Asset."""
         if not isinstance(self.key, AssetAliasUniqueKey):
             return
 
+        if isinstance(asset, AssetRef):
+            asset_key = self._resolve_asset_ref(asset)
+        else:
+            asset_key = AssetUniqueKey.from_asset(asset)
+
         asset_alias_name = self.key.name
         event = AssetAliasEvent(
             source_alias_name=asset_alias_name,
-            dest_asset_key=AssetUniqueKey.from_asset(asset),
+            dest_asset_key=asset_key,
             extra=extra or {},
         )
         self.asset_alias_events.append(event)

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

@@ -351,6 +351,14 @@ def test_nested_context(self):
 
 
 class TestOutletEventAccessor:
+    @pytest.mark.parametrize(
+        "add_arg",
+        [
+            Asset("name", "uri"),
+            Asset.ref(name="name"),
+            Asset.ref(uri="uri"),
+        ],
+    )
     @pytest.mark.parametrize(
         "key, asset_alias_events",
         (
@@ -360,21 +368,28 @@ class TestOutletEventAccessor:
                 [
                     AssetAliasEvent(
                         source_alias_name="test_alias",
-                        dest_asset_key=AssetUniqueKey(uri="test_uri", name="test_uri"),
+                        dest_asset_key=AssetUniqueKey(name="name", uri="uri"),
                         extra={},
                     )
                 ],
             ),
         ),
     )
-    def test_add(self, key, asset_alias_events, mock_supervisor_comms):
-        asset = Asset("test_uri")
-        mock_supervisor_comms.get_message.return_value = asset
+    def test_add(self, add_arg, key, asset_alias_events, mock_supervisor_comms):
+        mock_supervisor_comms.get_message.return_value = AssetResponse(name="name", uri="uri", group="")
 
         outlet_event_accessor = OutletEventAccessor(key=key, extra={})
-        outlet_event_accessor.add(asset)
+        outlet_event_accessor.add(add_arg)
         assert outlet_event_accessor.asset_alias_events == asset_alias_events
 
+    @pytest.mark.parametrize(
+        "add_arg",
+        [
+            Asset("name", "uri"),
+            Asset.ref(name="name"),
+            Asset.ref(uri="uri"),
+        ],
+    )
     @pytest.mark.parametrize(
         "key, asset_alias_events",
         (
@@ -384,19 +399,18 @@ def test_add(self, key, asset_alias_events, mock_supervisor_comms):
                 [
                     AssetAliasEvent(
                         source_alias_name="test_alias",
-                        dest_asset_key=AssetUniqueKey(name="test-asset", uri="test://asset-uri/"),
+                        dest_asset_key=AssetUniqueKey(name="name", uri="uri"),
                         extra={},
                     )
                 ],
             ),
         ),
     )
-    def test_add_with_db(self, key, asset_alias_events, mock_supervisor_comms):
-        asset = Asset(uri="test://asset-uri", name="test-asset")
-        mock_supervisor_comms.get_message.return_value = asset
+    def test_add_with_db(self, add_arg, key, asset_alias_events, mock_supervisor_comms):
+        mock_supervisor_comms.get_message.return_value = AssetResponse(name="name", uri="uri", group="")
 
         outlet_event_accessor = OutletEventAccessor(key=key, extra={"not": ""})
-        outlet_event_accessor.add(asset, extra={})
+        outlet_event_accessor.add(add_arg, extra={})
         assert outlet_event_accessor.asset_alias_events == asset_alias_events