Expand/collapse all groups

airflow/www/static/js/grid/Grid.jsx

@@ -41,6 +41,7 @@ import DagRuns from './dagRuns';
 import Details from './details';
 import useSelection from './utils/useSelection';
 import { useAutoRefresh } from './context/autorefresh';
+import ToggleGroups from './ToggleGroups';
 
 const sidePanelKey = 'hideSidePanel';
 
@@ -86,22 +87,25 @@ const Grid = () => {
 
   return (
     <Box>
-      <Flex flexGrow={1} justifyContent="flex-end" alignItems="center">
-        <ResetRoot />
-        <FormControl display="flex" width="auto" mr={2}>
-          {isRefreshOn && <Spinner color="blue.500" speed="1s" mr="4px" />}
-          <FormLabel htmlFor="auto-refresh" mb={0} fontWeight="normal">
-            Auto-refresh
-          </FormLabel>
-          <Switch
-            id="auto-refresh"
-            onChange={() => toggleRefresh(true)}
-            isDisabled={isPaused}
-            isChecked={isRefreshOn}
-            size="lg"
-            title={isPaused ? 'Autorefresh is disabled while the DAG is paused' : ''}
-          />
-        </FormControl>
+      <Flex flexGrow={1} justifyContent="space-between" alignItems="center">
+        <Flex alignItems="center">
+          <FormControl display="flex" width="auto" mr={2}>
+            {isRefreshOn && <Spinner color="blue.500" speed="1s" mr="4px" />}
+            <FormLabel htmlFor="auto-refresh" mb={0} fontWeight="normal">
+              Auto-refresh
+            </FormLabel>
+            <Switch
+              id="auto-refresh"
+              onChange={() => toggleRefresh(true)}
+              isDisabled={isPaused}
+              isChecked={isRefreshOn}
+              size="lg"
+              title={isPaused ? 'Autorefresh is disabled while the DAG is paused' : ''}
+            />
+          </FormControl>
+          <ToggleGroups groups={groups} />
+          <ResetRoot />
+        </Flex>
         <Button
           onClick={toggleSidePanel}
           aria-label={isOpen ? 'Show Details' : 'Hide Details'}

airflow/www/static/js/grid/TaskName.jsx

@@ -25,13 +25,13 @@ import {
 import { FiChevronUp, FiChevronDown } from 'react-icons/fi';
 
 const TaskName = ({
-  isGroup = false, isMapped = false, onToggle, isOpen, level, taskName,
+  isGroup = false, isMapped = false, onToggle, isOpen, level, label,
 }) => (
   <Flex
     as={isGroup ? 'button' : 'div'}
     onClick={onToggle}
-    aria-label={taskName}
-    title={taskName}
+    aria-label={label}
+    title={label}
     mr={4}
     width="100%"
     alignItems="center"
@@ -42,7 +42,7 @@ const TaskName = ({
       ml={level * 4 + 4}
       isTruncated
     >
-      {taskName}
+      {label}
       {isMapped && (
         ' [ ]'
       )}

airflow/www/static/js/grid/ToggleGroups.jsx

@@ -0,0 +1,92 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/* global localStorage, CustomEvent, document */
+
+import React, { useState } from 'react';
+import { Flex, IconButton } from '@chakra-ui/react';
+import { MdExpand, MdCompress } from 'react-icons/md';
+
+import { getMetaValue } from '../utils';
+
+const dagId = getMetaValue('dag_id');
+
+const getGroupIds = (groups) => {
+  const groupIds = [];
+  const checkTasks = (tasks) => tasks.forEach((task) => {
+    if (task.children) {
+      groupIds.push(task.label);
+      checkTasks(task.children);
+    }
+  });
+  checkTasks(groups);
+  return groupIds;
+};
+
+const ToggleGroups = ({ groups }) => {
+  const openGroupsKey = `${dagId}/open-groups`;
+  const allGroupIds = getGroupIds(groups.children);
+  const storedGroups = JSON.parse(localStorage.getItem(openGroupsKey)) || [];
+  const [openGroupIds, setOpenGroupIds] = useState(storedGroups);
+
+  const isExpandDisabled = allGroupIds.length === openGroupIds.length;
+  const isCollapseDisabled = !openGroupIds.length;
+
+  // Don't show button if the DAG has no task groups
+  const hasGroups = groups.children.find((c) => !!c.children);
+  if (!hasGroups) return null;
+
+  const onExpand = () => {
+    const closeEvent = new CustomEvent('toggleGroups', { detail: { dagId, openGroups: true } });
+    document.dispatchEvent(closeEvent);
+    localStorage.setItem(openGroupsKey, JSON.stringify(allGroupIds));
+    setOpenGroupIds(allGroupIds);
+  };
+
+  const onCollapse = () => {
+    const closeEvent = new CustomEvent('toggleGroups', { detail: { dagId, closeGroups: true } });
+    document.dispatchEvent(closeEvent);
+    localStorage.removeItem(openGroupsKey);
+    setOpenGroupIds([]);
+  };
+
+  return (
+    <Flex>
+      <IconButton
+        fontSize="2xl"
+        onClick={onExpand}
+        title="Expand all task groups"
+        aria-label="Expand all task groups"
+        icon={<MdExpand />}
+        isDisabled={isExpandDisabled}
+        mr={2}
+      />
+      <IconButton
+        fontSize="2xl"
+        onClick={onCollapse}
+        title="Collapse all task groups"
+        aria-label="Collapse all task groups"
+        isDisabled={isCollapseDisabled}
+        icon={<MdCompress />}
+      />
+    </Flex>
+  );
+};
+
+export default ToggleGroups;

airflow/www/static/js/grid/ToggleGroups.test.jsx

@@ -0,0 +1,177 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/* global describe, test, expect */
+
+import React from 'react';
+import { Flex, Table, Tbody } from '@chakra-ui/react';
+import { render, fireEvent, waitFor } from '@testing-library/react';
+
+import renderTaskRows from './renderTaskRows';
+import ToggleGroups from './ToggleGroups';
+import { Wrapper } from './utils/testUtils';
+
+const mockGridData = {
+  groups: {
+    id: null,
+    label: null,
+    children: [
+      {
+        extraLinks: [],
+        id: 'group_1',
+        label: 'group_1',
+        instances: [
+          {
+            dagId: 'dagId',
+            duration: 0,
+            endDate: '2021-10-26T15:42:03.391939+00:00',
+            executionDate: '2021-10-25T15:41:09.726436+00:00',
+            operator: 'DummyOperator',
+            runId: 'run1',
+            startDate: '2021-10-26T15:42:03.391917+00:00',
+            state: 'success',
+            taskId: 'group_1',
+            tryNumber: 1,
+          },
+        ],
+        children: [
+          {
+            id: 'group_1.task_1',
+            label: 'task_1',
+            extraLinks: [],
+            instances: [
+              {
+                dagId: 'dagId',
+                duration: 0,
+                endDate: '2021-10-26T15:42:03.391939+00:00',
+                executionDate: '2021-10-25T15:41:09.726436+00:00',
+                operator: 'DummyOperator',
+                runId: 'run1',
+                startDate: '2021-10-26T15:42:03.391917+00:00',
+                state: 'success',
+                taskId: 'group_1.task_1',
+                tryNumber: 1,
+              },
+            ],
+            children: [
+              {
+                id: 'group_1.task_1.sub_task_1',
+                label: 'sub_task_1',
+                extraLinks: [],
+                instances: [
+                  {
+                    dagId: 'dagId',
+                    duration: 0,
+                    endDate: '2021-10-26T15:42:03.391939+00:00',
+                    executionDate: '2021-10-25T15:41:09.726436+00:00',
+                    operator: 'DummyOperator',
+                    runId: 'run1',
+                    startDate: '2021-10-26T15:42:03.391917+00:00',
+                    state: 'success',
+                    taskId: 'group_1.task_1.sub_task_1',
+                    tryNumber: 1,
+                  },
+                ],
+              },
+            ],
+          },
+        ],
+      },
+    ],
+    instances: [],
+  },
+  dagRuns: [
+    {
+      dagId: 'dagId',
+      runId: 'run1',
+      dataIntervalStart: new Date(),
+      dataIntervalEnd: new Date(),
+      startDate: '2021-11-08T21:14:19.704433+00:00',
+      endDate: '2021-11-08T21:17:13.206426+00:00',
+      state: 'failed',
+      runType: 'scheduled',
+      executionDate: '2021-11-08T21:14:19.704433+00:00',
+    },
+  ],
+};
+
+const EXPAND = 'Expand all task groups';
+const COLLAPSE = 'Collapse all task groups';
+
+describe('Test ToggleGroups', () => {
+  test('Buttons are disabled if all groups are expanded or collapsed', () => {
+    const { getByTitle } = render(
+      <ToggleGroups groups={mockGridData.groups} />,
+      { wrapper: Wrapper },
+    );
+
+    const expandButton = getByTitle(EXPAND);
+    const collapseButton = getByTitle(COLLAPSE);
+
+    expect(expandButton).toBeEnabled();
+    expect(collapseButton).toBeDisabled();
+
+    fireEvent.click(expandButton);
+
+    expect(collapseButton).toBeEnabled();
+    expect(expandButton).toBeDisabled();
+  });
+
+  test('Expand/collapse buttons toggle nested groups', async () => {
+    global.gridData = mockGridData;
+    const dagRunIds = mockGridData.dagRuns.map((dr) => dr.runId);
+    const task = mockGridData.groups;
+
+    const { getByText, queryAllByTestId, getByTitle } = render(
+      <Flex>
+        <ToggleGroups groups={task} />
+        <Table>
+          <Tbody>
+            {renderTaskRows({ task, dagRunIds })}
+          </Tbody>
+        </Table>
+      </Flex>,
+      { wrapper: Wrapper },
+    );
+
+    const expandButton = getByTitle(EXPAND);
+    const collapseButton = getByTitle(COLLAPSE);
+
+    const groupName = getByText('group_1');
+
+    expect(queryAllByTestId('task-instance')).toHaveLength(3);
+    expect(groupName).toBeInTheDocument();
+
+    expect(queryAllByTestId('open-group')).toHaveLength(2);
+    expect(queryAllByTestId('closed-group')).toHaveLength(0);
+
+    fireEvent.click(collapseButton);
+
+    await waitFor(() => expect(queryAllByTestId('task-instance')).toHaveLength(1));
+    expect(queryAllByTestId('open-group')).toHaveLength(0);
+    // Since the groups are nested, only the parent row is rendered
+    expect(queryAllByTestId('closed-group')).toHaveLength(1);
+
+    fireEvent.click(expandButton);
+
+    await waitFor(() => expect(queryAllByTestId('task-instance')).toHaveLength(3));
+    expect(queryAllByTestId('open-group')).toHaveLength(2);
+    expect(queryAllByTestId('closed-group')).toHaveLength(0);
+  });
+});