[v2-10-test] Avoid grouping task instance stats by try_number for dyn…

…amic mapped tasks (#44300) (#44319)

* [v2-10-test] Avoid grouping task instance stats by try_number for dynamic mapped tasks (#44300)
(cherry picked from commit 5e52bd2)

Co-authored-by: Shahar Epstein <60007259+shahar1@users.noreply.github.com>

* Update test_views_grid.py

---------

Co-authored-by: Shahar Epstein <60007259+shahar1@users.noreply.github.com>

airflow/www/views.py

@@ -316,7 +316,10 @@ def dag_to_grid(dag: DagModel, dag_runs: Sequence[DagRun], session: Session) ->
             TaskInstance.task_id,
             TaskInstance.run_id,
             TaskInstance.state,
-            TaskInstance.try_number,
+            case(
+                (TaskInstance.map_index == -1, TaskInstance.try_number),
+                else_=None,
+            ).label("try_number"),
             func.min(TaskInstanceNote.content).label("note"),
             func.count(func.coalesce(TaskInstance.state, sqla.literal("no_status"))).label("state_count"),
             func.min(TaskInstance.queued_dttm).label("queued_dttm"),
@@ -328,7 +331,15 @@ def dag_to_grid(dag: DagModel, dag_runs: Sequence[DagRun], session: Session) ->
             TaskInstance.dag_id == dag.dag_id,
             TaskInstance.run_id.in_([dag_run.run_id for dag_run in dag_runs]),
         )
-        .group_by(TaskInstance.task_id, TaskInstance.run_id, TaskInstance.state, TaskInstance.try_number)
+        .group_by(
+            TaskInstance.task_id,
+            TaskInstance.run_id,
+            TaskInstance.state,
+            case(
+                (TaskInstance.map_index == -1, TaskInstance.try_number),
+                else_=None,
+            ),
+        )
         .order_by(TaskInstance.task_id, TaskInstance.run_id)
     )
 

newsfragments/44300.bugfix.rst

@@ -0,0 +1 @@
+Fix stats of dynamic mapped tasks after automatic retries of failed tasks

tests/www/views/test_views_grid.py

@@ -517,3 +517,47 @@ def test_next_run_datasets_404(admin_client):
     resp = admin_client.get("/object/next_run_datasets/missingdag", follow_redirects=True)
     assert resp.status_code == 404, resp.json
     assert resp.json == {"error": "can't find dag missingdag"}
+
+
+@pytest.mark.usefixtures("freeze_time_for_dagruns")
+def test_dynamic_mapped_task_with_retries(admin_client, dag_with_runs: list[DagRun], session):
+    """
+    Test a DAG with a dynamic mapped task with retries
+    """
+    run1, run2 = dag_with_runs
+
+    for ti in run1.task_instances:
+        ti.state = TaskInstanceState.SUCCESS
+    for ti in sorted(run2.task_instances, key=lambda ti: (ti.task_id, ti.map_index)):
+        if ti.task_id == "task1":
+            ti.state = TaskInstanceState.SUCCESS
+        elif ti.task_id == "group.mapped":
+            if ti.map_index == 0:
+                ti.state = TaskInstanceState.FAILED
+                ti.start_date = pendulum.DateTime(2021, 7, 1, 1, 0, 0, tzinfo=pendulum.UTC)
+                ti.end_date = pendulum.DateTime(2021, 7, 1, 1, 2, 3, tzinfo=pendulum.UTC)
+            elif ti.map_index == 1:
+                ti.try_number = 1
+                ti.state = TaskInstanceState.SUCCESS
+                ti.start_date = pendulum.DateTime(2021, 7, 1, 2, 3, 4, tzinfo=pendulum.UTC)
+                ti.end_date = None
+            elif ti.map_index == 2:
+                ti.try_number = 2
+                ti.state = TaskInstanceState.FAILED
+                ti.start_date = pendulum.DateTime(2021, 7, 1, 2, 3, 4, tzinfo=pendulum.UTC)
+                ti.end_date = None
+            elif ti.map_index == 3:
+                ti.try_number = 3
+                ti.state = TaskInstanceState.SUCCESS
+                ti.start_date = pendulum.DateTime(2021, 7, 1, 2, 3, 4, tzinfo=pendulum.UTC)
+                ti.end_date = None
+    session.flush()
+
+    resp = admin_client.get(f"/object/grid_data?dag_id={DAG_ID}", follow_redirects=True)
+
+    assert resp.status_code == 200, resp.json
+
+    assert resp.json["groups"]["children"][-1]["children"][-1]["instances"][-1]["mapped_states"] == {
+        "failed": 2,
+        "success": 2,
+    }