AIP-38 Clear Tasks (apache#45559)

* POC working

* Add clear task instance action button

* Harmonize actions columns in table

* Simplify state management by using useEffect

* Rename ClearRun folder to Clear

* Factorize columns

* Better file structure

* Refactor SegmentedControl Component to handle multiple values

airflow/ui/src/components/ClearRun/ClearRunTaskAccordion.tsx → airflow/ui/src/components/Clear/ClearAccordion.tsx

@@ -17,60 +17,14 @@
  * under the License.
  */
 import { Box, Editable, Text, VStack } from "@chakra-ui/react";
-import { Link } from "@chakra-ui/react";
-import type { ColumnDef } from "@tanstack/react-table";
 import type { ChangeEvent } from "react";
-import { Link as RouterLink } from "react-router-dom";
 
-import type {
-  DAGRunResponse,
-  TaskInstanceCollectionResponse,
-  TaskInstanceResponse,
-} from "openapi/requests/types.gen";
+import type { DAGRunResponse, TaskInstanceCollectionResponse } from "openapi/requests/types.gen";
 import { DataTable } from "src/components/DataTable";
 import ReactMarkdown from "src/components/ReactMarkdown";
-import { Status, Tooltip } from "src/components/ui";
 import { Accordion } from "src/components/ui";
-import { getTaskInstanceLink } from "src/utils/links";
-import { trimText } from "src/utils/trimTextFn";
 
-const columns: Array<ColumnDef<TaskInstanceResponse>> = [
-  {
-    accessorKey: "task_display_name",
-    cell: ({ row: { original } }) => (
-      <Tooltip content={original.task_display_name}>
-        <Link asChild color="fg.info" fontWeight="bold" maxWidth="200px" overflow="hidden">
-          <RouterLink to={getTaskInstanceLink(original)}>
-            {trimText(original.task_display_name, 25).trimmedText}
-          </RouterLink>
-        </Link>
-      </Tooltip>
-    ),
-    enableSorting: false,
-    header: "Task ID",
-  },
-  {
-    accessorKey: "state",
-    cell: ({
-      row: {
-        original: { state },
-      },
-    }) => <Status state={state}>{state}</Status>,
-    enableSorting: false,
-    header: () => "State",
-  },
-  {
-    accessorFn: (row: TaskInstanceResponse) => row.rendered_map_index ?? row.map_index,
-    enableSorting: false,
-    header: "Map Index",
-  },
-
-  {
-    accessorKey: "dag_run_id",
-    enableSorting: false,
-    header: "Run Id",
-  },
-];
+import { columns } from "./columns";
 
 type Props = {
   readonly affectedTasks?: TaskInstanceCollectionResponse;
@@ -80,7 +34,7 @@ type Props = {
 
 // Table is in memory, pagination and sorting are disabled.
 // TODO: Make a front-end only unconnected table component with client side ordering and pagination
-const ClearRunTasksAccordion = ({ affectedTasks, note, setNote }: Props) => (
+const ClearAccordion = ({ affectedTasks, note, setNote }: Props) => (
   <Accordion.Root collapsible defaultValue={["note"]} multiple={false} variant="enclosed">
     <Accordion.Item key="tasks" value="tasks">
       <Accordion.ItemTrigger>
@@ -144,4 +98,4 @@ const ClearRunTasksAccordion = ({ affectedTasks, note, setNote }: Props) => (
   </Accordion.Root>
 );
 
-export default ClearRunTasksAccordion;
+export default ClearAccordion;

airflow/ui/src/components/ClearRun/ClearRunButton.tsx → airflow/ui/src/components/Clear/Run/ClearRunButton.tsx

@@ -21,9 +21,9 @@ import { useState } from "react";
 import { FiRefreshCw } from "react-icons/fi";
 
 import type { DAGRunResponse, TaskInstanceCollectionResponse } from "openapi/requests/types.gen";
+import ActionButton from "src/components/ui/ActionButton";
 import { useClearDagRun } from "src/queries/useClearRun";
 
-import ActionButton from "../ui/ActionButton";
 import ClearRunDialog from "./ClearRunDialog";
 
 type Props = {
@@ -34,8 +34,6 @@ type Props = {
 const ClearRunButton = ({ dagRun, withText = true }: Props) => {
   const { onClose, onOpen, open } = useDisclosure();
 
-  const [onlyFailed, setOnlyFailed] = useState(false);
-
   const [affectedTasks, setAffectedTasks] = useState<TaskInstanceCollectionResponse>({
     task_instances: [],
     total_entries: 0,
@@ -56,14 +54,7 @@ const ClearRunButton = ({ dagRun, withText = true }: Props) => {
       <ActionButton
         actionName="Clear Dag Run"
         icon={<FiRefreshCw />}
-        onClick={() => {
-          onOpen();
-          mutate({
-            dagId,
-            dagRunId,
-            requestBody: { dry_run: true, only_failed: onlyFailed },
-          });
-        }}
+        onClick={onOpen}
         text="Clear Run"
         withText={withText}
       />
@@ -74,9 +65,7 @@ const ClearRunButton = ({ dagRun, withText = true }: Props) => {
         isPending={isPending}
         mutate={mutate}
         onClose={onClose}
-        onlyFailed={onlyFailed}
         open={open}
-        setOnlyFailed={setOnlyFailed}
       />
     </Box>
   );

airflow/ui/src/components/ClearRun/ClearRunDialog.tsx → airflow/ui/src/components/Clear/Run/ClearRunDialog.tsx

@@ -17,7 +17,7 @@
  * under the License.
  */
 import { Flex, Heading, VStack } from "@chakra-ui/react";
-import { useState } from "react";
+import { useEffect, useState } from "react";
 import { FiRefreshCw } from "react-icons/fi";
 
 import type {
@@ -26,10 +26,10 @@ import type {
   TaskInstanceCollectionResponse,
 } from "openapi/requests/types.gen";
 import { Button, Dialog } from "src/components/ui";
+import SegmentedControl from "src/components/ui/SegmentedControl";
 import { usePatchDagRun } from "src/queries/usePatchDagRun";
 
-import SegmentedControl from "../ui/SegmentedControl";
-import ClearRunTasksAccordion from "./ClearRunTaskAccordion";
+import ClearAccordion from "../ClearAccordion";
 
 type Props = {
   readonly affectedTasks: TaskInstanceCollectionResponse;
@@ -45,57 +45,36 @@ type Props = {
     requestBody: DAGRunClearBody;
   }) => void;
   readonly onClose: () => void;
-  readonly onlyFailed: boolean;
   readonly open: boolean;
-  readonly setOnlyFailed: (value: boolean) => void;
 };
 
-const ClearRunDialog = ({
-  affectedTasks,
-  dagRun,
-  isPending,
-  mutate,
-  onClose,
-  onlyFailed,
-  open,
-  setOnlyFailed,
-}: Props) => {
+const ClearRunDialog = ({ affectedTasks, dagRun, isPending, mutate, onClose, open }: Props) => {
+  const [selectedOptions, setSelectedOptions] = useState<Array<string>>([]);
+
+  const onlyFailed = selectedOptions.includes("onlyFailed");
+
   const dagId = dagRun.dag_id;
   const dagRunId = dagRun.dag_run_id;
 
   const [note, setNote] = useState<string | null>(dagRun.note);
   const { isPending: isPendingPatchDagRun, mutate: mutatePatchDagRun } = usePatchDagRun({ dagId, dagRunId });
 
-  const onChange = (value: string) => {
-    switch (value) {
-      case "existing_tasks":
-        setOnlyFailed(false);
-        mutate({
-          dagId,
-          dagRunId,
-          requestBody: { dry_run: true, only_failed: false },
-        });
-        break;
-      case "only_failed":
-        setOnlyFailed(true);
-        mutate({
-          dagId,
-          dagRunId,
-          requestBody: { dry_run: true, only_failed: true },
-        });
-        break;
-      default:
-        // TODO: Handle this `new_tasks` case
-        break;
-    }
-  };
+  useEffect(() => {
+    mutate({
+      dagId,
+      dagRunId,
+      requestBody: { dry_run: true, only_failed: onlyFailed },
+    });
+  }, [dagId, dagRunId, mutate, onlyFailed]);
 
   return (
     <Dialog.Root onOpenChange={onClose} open={open} size="xl">
       <Dialog.Content backdrop>
         <Dialog.Header>
           <VStack align="start" gap={4}>
-            <Heading size="xl">Clear DagRun - {dagRunId} </Heading>
+            <Heading size="xl">
+              <strong>Clear DagRun: </strong> {dagRunId}
+            </Heading>
           </VStack>
         </Dialog.Header>
 
@@ -104,21 +83,20 @@ const ClearRunDialog = ({
         <Dialog.Body width="full">
           <Flex justifyContent="center">
             <SegmentedControl
-              mb={3}
-              onValueChange={onChange}
+              defaultValues={["existingTasks"]}
+              onChange={setSelectedOptions}
               options={[
-                { label: "Clear existing tasks", value: "existing_tasks" },
-                { label: "Clear only failed tasks", value: "only_failed" },
+                { label: "Clear existing tasks", value: "existingTasks" },
+                { label: "Clear only failed tasks", value: "onlyFailed" },
                 {
                   disabled: true,
                   label: "Queue up new tasks",
                   value: "new_tasks",
                 },
               ]}
-              value={onlyFailed ? "only_failed" : "existing_tasks"}
             />
           </Flex>
-          <ClearRunTasksAccordion affectedTasks={affectedTasks} note={note} setNote={setNote} />
+          <ClearAccordion affectedTasks={affectedTasks} note={note} setNote={setNote} />
           <Flex justifyContent="end" mt={3}>
             <Button
               colorPalette="blue"

airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceButton.tsx

@@ -0,0 +1,74 @@
+/*!
+ * 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.
+ */
+import { Box, useDisclosure } from "@chakra-ui/react";
+import { useState } from "react";
+import { FiRefreshCw } from "react-icons/fi";
+
+import type { TaskInstanceCollectionResponse, TaskInstanceResponse } from "openapi/requests/types.gen";
+import ActionButton from "src/components/ui/ActionButton";
+import { useClearTaskInstances } from "src/queries/useClearTaskInstances";
+
+import ClearTaskInstanceDialog from "./ClearTaskInstanceDialog";
+
+type Props = {
+  readonly taskInstance: TaskInstanceResponse;
+  readonly withText?: boolean;
+};
+
+const ClearTaskInstanceButton = ({ taskInstance, withText = true }: Props) => {
+  const { onClose, onOpen, open } = useDisclosure();
+
+  const [affectedTasks, setAffectedTasks] = useState<TaskInstanceCollectionResponse>({
+    task_instances: [],
+    total_entries: 0,
+  });
+
+  const dagId = taskInstance.dag_id;
+  const dagRunId = taskInstance.dag_run_id;
+
+  const { isPending, mutate } = useClearTaskInstances({
+    dagId,
+    dagRunId,
+    onSuccessConfirm: onClose,
+    onSuccessDryRun: setAffectedTasks,
+  });
+
+  return (
+    <Box>
+      <ActionButton
+        actionName="Clear Task Instance"
+        icon={<FiRefreshCw />}
+        onClick={onOpen}
+        text="Clear Task Instance"
+        withText={withText}
+      />
+
+      <ClearTaskInstanceDialog
+        affectedTasks={affectedTasks}
+        isPending={isPending}
+        mutate={mutate}
+        onClose={onClose}
+        open={open}
+        taskInstance={taskInstance}
+      />
+    </Box>
+  );
+};
+
+export default ClearTaskInstanceButton;