Allow selecting loop tasks in the Grid view (#76291)

Loops, the tasks inside them and the task groups around them were left out of
the Grid multi-selection because a looped task's cell sums up every pass while
its link opens only the newest one. The bulk Clear, Mark as and Delete actions
now address each pass on their own, so a loop cell can stand for every pass it
covers, the way a mapped task's cell stands for all its mapped task instances.

The bulk endpoint reports a task instance that has a region by its id rather
than by its coordinates, so the action bar also maps those ids back to the
selected cells to deselect them once they are updated.

airflow-core/docs/ui.rst

@@ -171,8 +171,8 @@ Use Grid View to:
   Shift + click selects every cell between the last one picked and the clicked one. The selected task instances can
   then be cleared, marked, or deleted together, and the actions of the details panel are hidden meanwhile. A mapped
   task's cell stands for all of its mapped task instances, and a task group's cell for every task instance inside the
-  group in that run. Loops, the tasks inside them, and task groups that contain or sit inside a loop cannot be
-  selected this way.
+  group in that run. The cell of a task inside a loop stands for every pass of that task in the run, and a loop's
+  cell for every pass of the loop.
 - **Filter tasks** by name or partial ID
 - **Select a run range**, like "last 25 runs" using the dropdown above the grid
 

airflow-core/src/airflow/ui/public/i18n/locales/en/common.json

@@ -386,12 +386,12 @@
       "scrollBottom": "Scroll to bottom",
       "scrollTop": "Scroll to top",
       "searchDags": "Search $t(dag_other)",
-      "selectGridCellRange": "Select a range of $t(taskInstance_other) or $t(taskGroup_other) (not loop tasks)",
+      "selectGridCellRange": "Select a range of $t(taskInstance_other) or $t(taskGroup_other)",
       "showHelp": "Show $t(shortcuts.title)",
       "toggleExpand": "Expand or collapse all groups",
       "toggleFullscreen": "Toggle fullscreen",
       "toggleGraphGrid": "Toggle graph / grid view",
-      "toggleGridCellSelection": "Add or remove a $t(taskInstance_one) or $t(taskGroup_one) from the selection (not loop tasks)",
+      "toggleGridCellSelection": "Add or remove a $t(taskInstance_one) or $t(taskGroup_one) from the selection",
       "toggleLogLevel": "Toggle log levels",
       "toggleSource": "Toggle source",
       "toggleTaskGroup": "Expand or collapse $t(taskGroup_one)",

airflow-core/src/airflow/ui/src/layouts/Details/Grid/Grid.tsx

@@ -189,7 +189,6 @@ export const Grid = ({
   const { clearSelection, deselectCells, handleSelectionClick, selection } = useGridSelection({
     flatNodes,
     runIds: gridRuns?.map((dr: GridRunsResponse) => dr.run_id) ?? [],
-    selectableGroups,
     summariesByRunId,
   });
 

airflow-core/src/airflow/ui/src/layouts/Details/Grid/GridSelectionActionBar.test.tsx

@@ -26,15 +26,19 @@ import { Wrapper } from "src/utils/Wrapper";
 import { GridSelectionActionBar } from "./GridSelectionActionBar";
 import { getCellKey } from "./useGridSelection";
 
+// `loop.c` in run r2 is a loop pass, which the bulk actions name by its id.
 vi.mock("./useGridSelectedTaskInstances", () => ({
-  useGridSelectedTaskInstances: () => ({ data: [], isLoading: false }),
+  useGridSelectedTaskInstances: () => ({
+    data: [{ dag_id: "dag", dag_run_id: "r2", id: "pass-id", map_index: -1, task_id: "loop.c" }],
+    isLoading: false,
+  }),
 }));
 vi.mock("src/pages/TaskInstances/BulkClearTaskInstancesButton", () => ({ default: () => undefined }));
 vi.mock("src/pages/TaskInstances/BulkDeleteTaskInstancesButton", () => ({ default: () => undefined }));
-// Reports `group.b` in run r1 as updated, the way the bulk actions name task instances.
+// Reports `group.b` in run r1 and the loop pass as updated, the way the bulk actions name them.
 vi.mock("src/pages/TaskInstances/BulkMarkTaskInstancesAsButton", () => ({
   default: ({ deselectKeys }: { readonly deselectKeys: (keys: Array<string>) => void }) => (
-    <button onClick={() => deselectKeys(["dag.r1.group.b[-1]"])} type="button">
+    <button onClick={() => deselectKeys(["dag.r1.group.b[-1]", "pass-id"])} type="button">
       mark
     </button>
   ),
@@ -44,7 +48,7 @@ const selectionOf = (cells: Array<GridCell>) =>
   new Map(cells.map((cell) => [getCellKey(cell.runId, cell.taskId), cell]));
 
 describe("GridSelectionActionBar", () => {
-  it("deselects a group cell once one of its task instances is updated", () => {
+  it("deselects the group cells whose task instances are updated, loop passes included", () => {
     const deselectCells = vi.fn();
 
     render(
@@ -52,11 +56,18 @@ describe("GridSelectionActionBar", () => {
         clearSelection={vi.fn()}
         dagId="dag"
         deselectCells={deselectCells}
-        selectableGroups={new Map([["group", ["group.a", "group.b"]]])}
+        selectableGroups={
+          new Map([
+            ["group", ["group.a", "group.b"]],
+            ["loop", ["loop.c"]],
+          ])
+        }
         selection={selectionOf([
           { runId: "r1", taskId: "group" },
           { runId: "r2", taskId: "group" },
           { runId: "r1", taskId: "group.a" },
+          { runId: "r1", taskId: "loop" },
+          { runId: "r2", taskId: "loop" },
         ])}
       />,
       { wrapper: Wrapper },
@@ -64,6 +75,6 @@ describe("GridSelectionActionBar", () => {
 
     fireEvent.click(screen.getByRole("button", { name: "mark" }));
 
-    expect(deselectCells).toHaveBeenCalledWith([getCellKey("r1", "group")]);
+    expect(deselectCells).toHaveBeenCalledWith([getCellKey("r1", "group"), getCellKey("r2", "loop")]);
   });
 });

airflow-core/src/airflow/ui/src/layouts/Details/Grid/GridSelectionActionBar.tsx

@@ -29,7 +29,7 @@ import type { GridSelection } from "src/context/gridSelection";
 import { formatNumber } from "src/utils";
 
 import { useGridSelectedTaskInstances } from "./useGridSelectedTaskInstances";
-import { getCoveredNodeIds, type SelectableGroups } from "./useGridSelection";
+import { getCellKey, getCoveredNodeIds, type SelectableGroups } from "./useGridSelection";
 
 type Props = {
   readonly clearSelection: VoidFunction;
@@ -53,18 +53,28 @@ export const GridSelectionActionBar = ({
     selectableGroups,
   );
 
-  // Bulk actions report each updated task instance as `dag_id.run_id.task_id[map_index]`. A group cell
-  // is deselected as soon as one of its task instances is updated.
-  const deselectKeys = (keys: Array<string>) =>
+  // Bulk actions report each updated task instance as `dag_id.run_id.task_id[map_index]`, or by its id
+  // when it has a region (a loop pass or a mapped task instance). A group cell is deselected as soon as
+  // one of its task instances is updated.
+  const deselectKeys = (keys: Array<string>) => {
+    const updatedCellKeys = new Set(
+      selectedTaskInstances
+        .filter((ti) => keys.includes(ti.id))
+        .map((ti) => getCellKey(ti.dag_run_id, ti.task_id)),
+    );
+
     deselectCells(
       [...selection.entries()]
         .filter(([, cell]) =>
-          getCoveredNodeIds(cell, selectableGroups).some((taskId) =>
-            keys.some((key) => key.startsWith(`${dagId}.${cell.runId}.${taskId}[`)),
+          getCoveredNodeIds(cell, selectableGroups).some(
+            (taskId) =>
+              updatedCellKeys.has(getCellKey(cell.runId, taskId)) ||
+              keys.some((key) => key.startsWith(`${dagId}.${cell.runId}.${taskId}[`)),
           ),
         )
         .map(([cellKey]) => cellKey),
     );
+  };
 
   return (
     <ActionBar.Root closeOnInteractOutside={false} open={selection.size > 0}>

airflow-core/src/airflow/ui/src/layouts/Details/Grid/useGridSelection.test.tsx

@@ -64,7 +64,7 @@ const FLAT_NODES = [node("group", true), node("a"), node("b"), node("c")];
 const SELECTABLE_GROUPS = new Map([["group", ["a", "b"]]]);
 const LOOPED = { latest_region_id: "11111111-1111-4111-8111-111111111111", latest_region_index: 2 };
 
-// "c" is looped in r2 and "b" has no task instance in r3, so neither can be selected.
+// "c" is looped in r2, and "b" has no task instance in r3 so it cannot be selected there.
 const SUMMARIES = new Map<string, GridTISummaries>(
   RUN_IDS.map((runId) => [
     runId,
@@ -90,7 +90,6 @@ const renderSelection = (route = "/dags/dag") =>
       ...useGridSelection({
         flatNodes: FLAT_NODES,
         runIds: RUN_IDS,
-        selectableGroups: SELECTABLE_GROUPS,
         summariesByRunId: SUMMARIES,
       }),
       hasGridSelection: useHasGridSelection(),
@@ -126,7 +125,7 @@ const groupNode = (id: string, children: Array<GridNodeResponse>, isLoop = false
 const leafNode = (id: string): GridNodeResponse => ({ id, is_mapped: false, label: id });
 
 describe("buildSelectableGroups", () => {
-  it("lists the members of each group, leaving out loops and the groups around or inside them", () => {
+  it("lists the members of each group, loops and the groups around or inside them included", () => {
     const structure = [
       groupNode("outer", [leafNode("outer.a"), groupNode("outer.inner", [leafNode("outer.inner.b")])]),
       groupNode("with_loop", [leafNode("with_loop.c"), groupNode("with_loop.loop", [leafNode("d")], true)]),
@@ -136,8 +135,12 @@ describe("buildSelectableGroups", () => {
 
     expect(buildSelectableGroups(structure)).toEqual(
       new Map([
+        ["loop.group", ["loop.group.e"]],
+        ["loop", ["loop.group", "loop.group.e"]],
         ["outer.inner", ["outer.inner.b"]],
         ["outer", ["outer.a", "outer.inner", "outer.inner.b"]],
+        ["with_loop.loop", ["d"]],
+        ["with_loop", ["with_loop.c", "with_loop.loop", "d"]],
       ]),
     );
   });
@@ -158,24 +161,16 @@ describe("getCoveredCellKeys", () => {
 
 describe("isSelectableCell", () => {
   it.each([
-    { expected: true, name: "a plain task instance", taskInstance: instance("a"), taskNode: node("a") },
-    { expected: true, name: "a group", taskInstance: instance("group"), taskNode: node("group", true) },
-    {
-      expected: false,
-      name: "a group that is or holds a loop",
-      taskInstance: instance("loop"),
-      taskNode: node("loop", true),
-    },
-    { expected: false, name: "a cell without task instance", taskInstance: undefined, taskNode: node("a") },
+    { expected: true, name: "a plain task instance", taskInstance: instance("a") },
+    { expected: true, name: "a looped task", taskInstance: instance("c", LOOPED) },
+    { expected: false, name: "a cell without task instance", taskInstance: undefined },
     {
       expected: false,
       name: "a task that has no task instance in the run yet",
       taskInstance: instance("a", { dag_version_number: null }),
-      taskNode: node("a"),
     },
-    { expected: false, name: "a looped task", taskInstance: instance("c", LOOPED), taskNode: node("c") },
-  ])("returns $expected for $name", ({ expected, taskInstance, taskNode }) => {
-    expect(isSelectableCell(taskNode, taskInstance, SELECTABLE_GROUPS)).toBe(expected);
+  ])("returns $expected for $name", ({ expected, taskInstance }) => {
+    expect(isSelectableCell(taskInstance)).toBe(expected);
   });
 });
 
@@ -248,6 +243,7 @@ describe("useGridSelection", () => {
         getCellKey("r1", "b"),
         getCellKey("r2", "b"),
         getCellKey("r1", "c"),
+        getCellKey("r2", "c"),
         getCellKey("r3", "c"),
       ].sort(),
     );
@@ -273,7 +269,7 @@ describe("useGridSelection", () => {
       result.current.handleSelectionClick(click({ ctrlKey: true }), { runId: "r1", taskId: "a" });
     });
     act(() => {
-      handled = result.current.handleSelectionClick(click({ ctrlKey: true }), { runId: "r2", taskId: "c" });
+      handled = result.current.handleSelectionClick(click({ ctrlKey: true }), { runId: "r3", taskId: "b" });
     });
 
     expect(handled).toBe(false);

airflow-core/src/airflow/ui/src/layouts/Details/Grid/useGridSelection.ts

@@ -38,28 +38,22 @@ export type SelectableGroups = ReadonlyMap<string, Array<string>>;
 
 export const getCellKey = (runId: string, taskId: string) => JSON.stringify([runId, taskId]);
 
-// A selected group cell acts on everything inside it. A loop group, a group that contains a loop
-// or one inside a loop is left out, for the same reason as a looped task.
+// A selected group cell acts on everything inside it, every pass of a loop included.
 export const buildSelectableGroups = (nodes: Array<GridNodeResponse> | undefined): SelectableGroups => {
   const groups = new Map<string, Array<string>>();
-  const visit = (node: GridNodeResponse, inLoop: boolean): { hasLoop: boolean; memberIds: Array<string> } => {
+  const visit = (node: GridNodeResponse): Array<string> => {
     if (!node.children) {
-      return { hasLoop: false, memberIds: [] };
+      return [];
     }
 
-    const isLoop = node.is_loop === true;
-    const children = node.children.map((child) => ({ id: child.id, ...visit(child, inLoop || isLoop) }));
-    const hasLoop = isLoop || children.some((child) => child.hasLoop);
-    const memberIds = children.flatMap((child) => [child.id, ...child.memberIds]);
+    const memberIds = node.children.flatMap((child) => [child.id, ...visit(child)]);
 
-    if (!inLoop && !hasLoop) {
-      groups.set(node.id, memberIds);
-    }
+    groups.set(node.id, memberIds);
 
-    return { hasLoop, memberIds };
+    return memberIds;
   };
 
-  nodes?.forEach((node) => visit(node, false));
+  nodes?.forEach(visit);
 
   return groups;
 };
@@ -78,22 +72,13 @@ export const getCoveredCellKeys = (selection: GridSelection, selectableGroups: S
     ),
   );
 
-// A looped task's cell aggregates every pass while its link opens only the newest one, so a bulk
-// action on it would be ambiguous. Cells without a task instance have nothing to act on.
-export const isSelectableCell = (
-  node: GridTask,
-  taskInstance: LightGridTaskInstanceSummary | undefined,
-  selectableGroups: SelectableGroups,
-) =>
-  (node.isGroup !== true || selectableGroups.has(node.id)) &&
-  taskInstance?.dag_version_number !== null &&
-  taskInstance?.dag_version_number !== undefined &&
-  (taskInstance.latest_region_id === null || taskInstance.latest_region_id === undefined);
+// Cells without a task instance have nothing to act on.
+export const isSelectableCell = (taskInstance: LightGridTaskInstanceSummary | undefined) =>
+  taskInstance?.dag_version_number !== null && taskInstance?.dag_version_number !== undefined;
 
 type GridData = {
   readonly flatNodes: Array<GridTask>;
   readonly runIds: Array<string>;
-  readonly selectableGroups: SelectableGroups;
   readonly summariesByRunId: ReadonlyMap<string, GridTISummaries>;
 };
 
@@ -113,7 +98,7 @@ export const buildRangeSelection = (data: GridData, from: GridCell, to: GridCell
 
   for (const node of data.flatNodes.slice(Math.min(...rows), Math.max(...rows) + 1)) {
     for (const runId of data.runIds.slice(Math.min(...columns), Math.max(...columns) + 1)) {
-      if (isSelectableCell(node, findTaskInstance(data, runId, node.id), data.selectableGroups)) {
+      if (isSelectableCell(findTaskInstance(data, runId, node.id))) {
         selection.set(getCellKey(runId, node.id), { runId, taskId: node.id });
       }
     }
@@ -189,18 +174,9 @@ export const useGridSelection = (data: GridData) => {
   // Returns whether the click changed the selection, in which case it must not navigate.
   const handleSelectionClick = useCallback(
     (event: MouseEvent, cell: GridCell): boolean => {
-      const isSelectable = ({ runId, taskId }: GridCell) => {
-        const node = dataRef.current.flatNodes.find((flatNode) => flatNode.id === taskId);
-
-        return (
-          node !== undefined &&
-          isSelectableCell(
-            node,
-            findTaskInstance(dataRef.current, runId, taskId),
-            dataRef.current.selectableGroups,
-          )
-        );
-      };
+      const isSelectable = ({ runId, taskId }: GridCell) =>
+        dataRef.current.flatNodes.some((flatNode) => flatNode.id === taskId) &&
+        isSelectableCell(findTaskInstance(dataRef.current, runId, taskId));
       const openedId = openedTaskId ?? openedGroupId;
       const openedCell =
         openedRunId !== undefined && openedId !== undefined