From b423967c2bb4b8bb4424fb97233957ed0eae48f9 Mon Sep 17 00:00:00 2001
From: Pierre Jeambrun <pierrejbrun@gmail.com>
Date: Thu, 19 Dec 2024 23:55:37 +0800
Subject: [PATCH] Add clear dag run list action (#45045)

* Add Clear Run list action

* Update following code review

* Add aria label
---
 .../components/ClearRun/ClearRunButton.tsx    | 49 ++++++++++++-------
 airflow/ui/src/components/ui/Tooltip.tsx      |  8 +--
 airflow/ui/src/pages/Dag/Runs/Runs.tsx        | 15 ++++++
 airflow/ui/src/queries/useClearRun.ts         |  2 +
 4 files changed, 52 insertions(+), 22 deletions(-)

diff --git a/airflow/ui/src/components/ClearRun/ClearRunButton.tsx b/airflow/ui/src/components/ClearRun/ClearRunButton.tsx
index 7c3060765ccf8..ae0386612c3d5 100644
--- a/airflow/ui/src/components/ClearRun/ClearRunButton.tsx
+++ b/airflow/ui/src/components/ClearRun/ClearRunButton.tsx
@@ -16,12 +16,17 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Box, useDisclosure } from "@chakra-ui/react";
-import { useState } from "react";
+import {
+  Box,
+  type ButtonProps,
+  IconButton,
+  useDisclosure,
+} from "@chakra-ui/react";
+import { type FC, useState } from "react";
 import { FiRefreshCw } from "react-icons/fi";
 
 import type { TaskInstanceCollectionResponse } from "openapi/requests/types.gen";
-import { Button } from "src/components/ui";
+import { Button, Tooltip } from "src/components/ui";
 import { useClearDagRun } from "src/queries/useClearRun";
 
 import ClearRunDialog from "./ClearRunDialog";
@@ -29,9 +34,10 @@ import ClearRunDialog from "./ClearRunDialog";
 type Props = {
   readonly dagId: string;
   readonly dagRunId: string;
+  readonly withText?: boolean;
 };
 
-const ClearRunButton = ({ dagId, dagRunId }: Props) => {
+const ClearRunButton = ({ dagId, dagRunId, withText = true }: Props) => {
   const { onClose, onOpen, open } = useDisclosure();
 
   const [onlyFailed, setOnlyFailed] = useState(false);
@@ -42,6 +48,8 @@ const ClearRunButton = ({ dagId, dagRunId }: Props) => {
       total_entries: 0,
     });
 
+  const ButtonComponent: FC<ButtonProps> = withText ? Button : IconButton;
+
   const { isPending, mutate } = useClearDagRun({
     dagId,
     dagRunId,
@@ -51,20 +59,25 @@ const ClearRunButton = ({ dagId, dagRunId }: Props) => {
 
   return (
     <Box>
-      <Button
-        onClick={() => {
-          onOpen();
-          mutate({
-            dagId,
-            dagRunId,
-            requestBody: { dry_run: true, only_failed: onlyFailed },
-          });
-        }}
-        variant="outline"
-      >
-        <FiRefreshCw height={5} width={5} />
-        Clear Run
-      </Button>
+      <Tooltip content="Clear Dag Run" disabled={Boolean(withText)}>
+        <ButtonComponent
+          aria-label="Clear Dag Run"
+          colorPalette={withText ? undefined : "blue"}
+          onClick={() => {
+            onOpen();
+            mutate({
+              dagId,
+              dagRunId,
+              requestBody: { dry_run: true, only_failed: onlyFailed },
+            });
+          }}
+          size={withText ? "md" : "sm"}
+          variant={withText ? "outline" : "ghost"}
+        >
+          <FiRefreshCw />
+          {withText ? "Clear Run" : ""}
+        </ButtonComponent>
+      </Tooltip>
 
       <ClearRunDialog
         affectedTasks={affectedTasks}
diff --git a/airflow/ui/src/components/ui/Tooltip.tsx b/airflow/ui/src/components/ui/Tooltip.tsx
index 2edb43c8b12f1..3be15cf35e0f8 100644
--- a/airflow/ui/src/components/ui/Tooltip.tsx
+++ b/airflow/ui/src/components/ui/Tooltip.tsx
@@ -17,7 +17,7 @@
  * under the License.
  */
 import { Tooltip as ChakraTooltip, Portal } from "@chakra-ui/react";
-import { forwardRef } from "react";
+import * as React from "react";
 
 export type TooltipProps = {
   content: React.ReactNode;
@@ -28,7 +28,7 @@ export type TooltipProps = {
   showArrow?: boolean;
 } & ChakraTooltip.RootProps;
 
-export const Tooltip = forwardRef<HTMLDivElement, TooltipProps>(
+export const Tooltip = React.forwardRef<HTMLDivElement, TooltipProps>(
   (props, ref) => {
     const {
       children,
@@ -46,7 +46,7 @@ export const Tooltip = forwardRef<HTMLDivElement, TooltipProps>(
     }
 
     return (
-      <ChakraTooltip.Root closeDelay={100} openDelay={100} {...rest}>
+      <ChakraTooltip.Root {...rest}>
         <ChakraTooltip.Trigger asChild>{children}</ChakraTooltip.Trigger>
         <Portal container={portalRef} disabled={!portalled}>
           <ChakraTooltip.Positioner>
@@ -55,7 +55,7 @@ export const Tooltip = forwardRef<HTMLDivElement, TooltipProps>(
                 <ChakraTooltip.Arrow>
                   <ChakraTooltip.ArrowTip />
                 </ChakraTooltip.Arrow>
-              ) : undefined}
+              ) : null}
               {content}
             </ChakraTooltip.Content>
           </ChakraTooltip.Positioner>
diff --git a/airflow/ui/src/pages/Dag/Runs/Runs.tsx b/airflow/ui/src/pages/Dag/Runs/Runs.tsx
index ae2cf2201c0ff..3d8be5adf6826 100644
--- a/airflow/ui/src/pages/Dag/Runs/Runs.tsx
+++ b/airflow/ui/src/pages/Dag/Runs/Runs.tsx
@@ -36,6 +36,7 @@ import {
 
 import { useDagRunServiceGetDagRuns } from "openapi/queries";
 import type { DAGRunResponse, DagRunState } from "openapi/requests/types.gen";
+import ClearRunButton from "src/components/ClearRun/ClearRunButton";
 import { DataTable } from "src/components/DataTable";
 import { useTableURLState } from "src/components/DataTable/useTableUrlState";
 import { ErrorAlert } from "src/components/ErrorAlert";
@@ -92,6 +93,20 @@ const columns: Array<ColumnDef<DAGRunResponse>> = [
       `${dayjs.duration(dayjs(original.end_date).diff(original.start_date)).asSeconds().toFixed(2)}s`,
     header: "Duration",
   },
+  {
+    accessorKey: "clear_dag_run",
+    cell: ({ row }) => (
+      <Flex justifyContent="end">
+        <ClearRunButton
+          dagId={row.original.dag_id}
+          dagRunId={row.original.dag_run_id}
+          withText={false}
+        />
+      </Flex>
+    ),
+    enableSorting: false,
+    header: "",
+  },
 ];
 
 const stateOptions = createListCollection({
diff --git a/airflow/ui/src/queries/useClearRun.ts b/airflow/ui/src/queries/useClearRun.ts
index 9efcb7ae03700..e1c7227ab9d2b 100644
--- a/airflow/ui/src/queries/useClearRun.ts
+++ b/airflow/ui/src/queries/useClearRun.ts
@@ -21,6 +21,7 @@ import { useQueryClient } from "@tanstack/react-query";
 import {
   useDagRunServiceClearDagRun,
   UseDagRunServiceGetDagRunKeyFn,
+  useDagRunServiceGetDagRunsKey,
   UseDagServiceGetDagDetailsKeyFn,
   useTaskInstanceServiceGetTaskInstancesKey,
 } from "openapi/queries";
@@ -66,6 +67,7 @@ export const useClearDagRun = ({
         [useTaskInstanceServiceGetTaskInstancesKey],
         UseDagServiceGetDagDetailsKeyFn({ dagId }),
         UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }),
+        [useDagRunServiceGetDagRunsKey],
       ];
 
       await Promise.all(