fix(taskinstnace): set internal asset_type of all Asset subclass to "Asset"

airflow/models/taskinstance.py

@@ -362,7 +362,12 @@ def _run_raw_task(
                 events = context["outlet_events"]
                 for obj in ti.task.outlets or []:
                     # Lineage can have other types of objects besides assets
-                    asset_type = type(obj).__name__
+                    #
+                    # The asset_type here is not Asset.asset_type but the obj type instead.
+                    # Only Asset is expected to be subclassed so
+                    # any subclass of Asset should be set as Asset to be handled correctly.
+                    # TODO: We should rename it or probably rework the logic later.
+                    asset_type = "Asset" if isinstance(obj, Asset) else type(obj).__name__
                     if isinstance(obj, Asset):
                         task_outlets.append(AssetProfile(name=obj.name, uri=obj.uri, asset_type=asset_type))
                         outlet_events.append(attrs.asdict(events[obj]))  # type: ignore

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

@@ -523,7 +523,12 @@ def _process_outlets(context: Context, outlets: list[AssetProfile]):
 
     for obj in outlets or []:
         # Lineage can have other types of objects besides assets
-        asset_type = type(obj).__name__
+        #
+        # The asset_type here is not Asset.asset_type but the obj type instead.
+        # Only Asset is expected to be subclassed so
+        # any subclass of Asset should be set as Asset to be handled correctly.
+        # TODO: We should rename it or probably rework the logic later.
+        asset_type = "Asset" if isinstance(obj, Asset) else type(obj).__name__
         if isinstance(obj, Asset):
             task_outlets.append(AssetProfile(name=obj.name, uri=obj.uri, asset_type=asset_type))
             outlet_events.append(attrs.asdict(events[obj]))  # type: ignore

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

@@ -52,7 +52,9 @@
     TaskInstance,
     TerminalTIState,
 )
-from airflow.sdk.definitions.asset import Asset, AssetAlias
+from airflow.sdk.definitions.asset import Asset, AssetAlias, Dataset, Model
+
+# from airflow.sdk.definitions.asset.decorators import AssetDefinition
 from airflow.sdk.definitions.param import DagParam
 from airflow.sdk.definitions.variable import Variable
 from airflow.sdk.execution_time.comms import (
@@ -703,6 +705,43 @@ def test_dag_parsing_context(make_ti_context, mock_supervisor_comms, monkeypatch
             ),
             id="asset",
         ),
+        pytest.param(
+            [Dataset(name="s3://bucket/my-task", uri="s3://bucket/my-task")],
+            SucceedTask(
+                state="success",
+                end_date=timezone.datetime(2024, 12, 3, 10, 0),
+                task_outlets=[
+                    AssetProfile(name="s3://bucket/my-task", uri="s3://bucket/my-task", asset_type="Asset")
+                ],
+                outlet_events=[
+                    {
+                        "key": {"name": "s3://bucket/my-task", "uri": "s3://bucket/my-task"},
+                        "extra": {},
+                        "asset_alias_events": [],
+                    }
+                ],
+            ),
+            id="dataset",
+        ),
+        pytest.param(
+            [Model(name="s3://bucket/my-task", uri="s3://bucket/my-task")],
+            SucceedTask(
+                state="success",
+                end_date=timezone.datetime(2024, 12, 3, 10, 0),
+                task_outlets=[
+                    AssetProfile(name="s3://bucket/my-task", uri="s3://bucket/my-task", asset_type="Asset")
+                ],
+                outlet_events=[
+                    {
+                        "key": {"name": "s3://bucket/my-task", "uri": "s3://bucket/my-task"},
+                        "extra": {},
+                        "asset_alias_events": [],
+                    }
+                ],
+            ),
+            id="model",
+        ),
+        # TODO: asset name ref and asset uri ref are not checked and should be added later
         pytest.param(
             [AssetAlias(name="example-alias", group="asset")],
             SucceedTask(