Fall back when a task cannot be expanded (#49772)

Author: uranusjr

airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import contextlib
 from functools import cache
 from operator import methodcaller
 from typing import Callable
@@ -42,6 +43,8 @@
 from airflow.models.dag_version import DagVersion
 from airflow.models.taskmap import TaskMap
 from airflow.sdk import BaseOperator
+from airflow.sdk.definitions._internal.abstractoperator import NotMapped
+from airflow.sdk.definitions._internal.expandinput import NotFullyPopulated
 from airflow.sdk.definitions.mappedoperator import MappedOperator
 from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup
 from airflow.serialization.serialized_objects import SerializedDAG
@@ -140,19 +143,14 @@ def get_child_task_map(parent_task_id: str, task_node_map: dict[str, dict[str, A
     return [task_id for task_id, task_map in task_node_map.items() if task_map["parent_id"] == parent_task_id]
 
 
-def _get_total_task_count(
-    run_id: str, task_count: list[int | MappedTaskGroup | MappedOperator], session: SessionDep
-) -> int:
-    return sum(
-        node
-        if isinstance(node, int)
-        else (
-            DBBaseOperator.get_mapped_ti_count(node, run_id=run_id, session=session) or 0
-            if isinstance(node, (MappedTaskGroup, MappedOperator))
-            else node
-        )
-        for node in task_count
-    )
+def _count_tis(node: int | MappedTaskGroup | MappedOperator, run_id: str, session: SessionDep) -> int:
+    if not isinstance(node, (MappedTaskGroup, MappedOperator)):
+        return node
+    with contextlib.suppress(NotFullyPopulated, NotMapped):
+        return DBBaseOperator.get_mapped_ti_count(node, run_id=run_id, session=session)
+    # If the downstream is not actually mapped, or we don't have information to
+    # determine the length yet, simply return 1 to represent the stand-in ti.
+    return 1
 
 
 def fill_task_instance_summaries(
@@ -253,7 +251,7 @@ def fill_task_instance_summaries(
                 end_date=ti_end_date,
                 queued_dttm=ti_queued_dttm,
                 child_states=child_states,
-                task_count=_get_total_task_count(run_id, task_node_map[task_id]["task_count"], session),
+                task_count=sum(_count_tis(n, run_id, session) for n in task_node_map[task_id]["task_count"]),
                 state=TaskInstanceState[overall_ti_state.upper()]
                 if overall_ti_state != "no_status"
                 else None,

airflow-core/src/airflow/models/baseoperator.py

@@ -621,7 +621,7 @@ def get_mapped_ti_count(cls, task: DAGNode, run_id: str, *, session: Session) ->
             raise NotImplementedError(f"Not implemented for {type(task)}")
 
         # https://github.com/python/cpython/issues/86153
-        # WHile we support Python 3.9 we can't rely on the type hint, we need to pass the type explicitly to
+        # While we support Python 3.9 we can't rely on the type hint, we need to pass the type explicitly to
         # register.
         @get_mapped_ti_count.register(TaskSDKAbstractOperator)
         @classmethod

airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py

@@ -48,6 +48,7 @@
 TASK_ID_4 = "task4"
 SUB_TASK_ID = "subtask"
 MAPPED_TASK_ID = "mapped_task"
+MAPPED_TASK_ID_2 = "mapped_task_2"
 TASK_GROUP_ID = "task_group"
 INNER_TASK_GROUP = "inner_task_group"
 INNER_TASK_GROUP_SUB_TASK = "inner_task_group_sub_task"
@@ -140,6 +141,31 @@
             "task_id": "task_group",
             "try_number": 0,
         },
+        {
+            "child_states": {
+                "deferred": 0,
+                "failed": 0,
+                "no_status": 0,
+                "queued": 0,
+                "removed": 0,
+                "restarting": 0,
+                "running": 0,
+                "scheduled": 0,
+                "skipped": 0,
+                "success": 1,
+                "up_for_reschedule": 0,
+                "up_for_retry": 0,
+                "upstream_failed": 0,
+            },
+            "end_date": None,
+            "note": None,
+            "queued_dttm": None,
+            "start_date": None,
+            "state": "success",
+            "task_count": 1,
+            "task_id": "mapped_task_2",
+            "try_number": 0,
+        },
         {
             "child_states": {
                 "deferred": 0,
@@ -331,6 +357,31 @@
             "task_id": "task_group",
             "try_number": 0,
         },
+        {
+            "child_states": {
+                "deferred": 0,
+                "failed": 0,
+                "no_status": 1,
+                "queued": 0,
+                "removed": 0,
+                "restarting": 0,
+                "running": 0,
+                "scheduled": 0,
+                "skipped": 0,
+                "success": 0,
+                "up_for_reschedule": 0,
+                "up_for_retry": 0,
+                "upstream_failed": 0,
+            },
+            "end_date": None,
+            "note": None,
+            "queued_dttm": None,
+            "start_date": None,
+            "state": None,
+            "task_count": 1,
+            "task_id": "mapped_task_2",
+            "try_number": 0,
+        },
         {
             "child_states": {
                 "deferred": 0,
@@ -517,6 +568,17 @@
             "tooltip": "",
             "type": "task",
         },
+        {
+            "asset_condition_type": None,
+            "children": None,
+            "id": "mapped_task_2",
+            "is_mapped": True,
+            "label": "mapped_task_2",
+            "operator": "MockOperator",
+            "setup_teardown_type": None,
+            "tooltip": None,
+            "type": "task",
+        },
     ],
 }
 
@@ -536,7 +598,7 @@ def setup(dag_maker, session=None):
 
     # DAG 1
     with dag_maker(dag_id=DAG_ID, serialized=True, session=session) as dag:
-        EmptyOperator(task_id=TASK_ID)
+        task = EmptyOperator(task_id=TASK_ID)
 
         @task_group
         def mapped_task_group(arg1):
@@ -549,6 +611,9 @@ def mapped_task_group(arg1):
             with TaskGroup(group_id=INNER_TASK_GROUP):
                 MockOperator.partial(task_id=INNER_TASK_GROUP_SUB_TASK).expand(arg1=["a", "b"])
 
+        # Mapped but never expanded. API should not crash, but count this as one no-status ti.
+        MockOperator.partial(task_id=MAPPED_TASK_ID_2).expand(arg1=task.output)
+
     triggered_by_kwargs = {"triggered_by": DagRunTriggeredByType.TEST}
     logical_date = timezone.datetime(2024, 11, 30)
     data_interval = dag.timetable.infer_manual_data_interval(run_after=logical_date)