Fix run task/run clearing (apache#45987)

* Fix clear modals * Reduce clear task/run API calls * Fix and add tests to datatable
aws-mwaa · Jan 24, 2025 · f5f5200 · f5f5200
1 parent 0639343
commit f5f5200
Show file tree

Hide file tree

Showing 13 changed files with 247 additions and 185 deletions.
diff --git a/airflow/ui/src/components/Clear/ClearAccordion.tsx b/airflow/ui/src/components/Clear/ClearAccordion.tsx
@@ -35,7 +35,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 ClearAccordion = ({ affectedTasks, note, setNote }: Props) => (
-  <Accordion.Root collapsible defaultValue={["note"]} multiple={false} variant="enclosed">
+  <Accordion.Root collapsible defaultValue={["tasks"]} multiple={false} variant="enclosed">
     <Accordion.Item key="tasks" value="tasks">
       <Accordion.ItemTrigger>
         <Text fontWeight="bold">Affected Tasks: {affectedTasks?.total_entries ?? 0}</Text>

diff --git a/airflow/ui/src/components/Clear/Run/ClearRunButton.tsx b/airflow/ui/src/components/Clear/Run/ClearRunButton.tsx
@@ -17,12 +17,10 @@
  * under the License.
  */
 import { Box, useDisclosure } from "@chakra-ui/react";
-import { useState } from "react";
 import { FiRefreshCw } from "react-icons/fi";
 
-import type { DAGRunResponse, TaskInstanceCollectionResponse } from "openapi/requests/types.gen";
+import type { DAGRunResponse } from "openapi/requests/types.gen";
 import ActionButton from "src/components/ui/ActionButton";
-import { useClearDagRun } from "src/queries/useClearRun";
 
 import ClearRunDialog from "./ClearRunDialog";
 
@@ -34,21 +32,6 @@ type Props = {
 const ClearRunButton = ({ dagRun, withText = true }: Props) => {
   const { onClose, onOpen, open } = useDisclosure();
 
-  const [affectedTasks, setAffectedTasks] = useState<TaskInstanceCollectionResponse>({
-    task_instances: [],
-    total_entries: 0,
-  });
-
-  const dagId = dagRun.dag_id;
-  const dagRunId = dagRun.dag_run_id;
-
-  const { isPending, mutate } = useClearDagRun({
-    dagId,
-    dagRunId,
-    onSuccessConfirm: onClose,
-    onSuccessDryRun: setAffectedTasks,
-  });
-
   return (
     <Box>
       <ActionButton
@@ -59,14 +42,7 @@ const ClearRunButton = ({ dagRun, withText = true }: Props) => {
         withText={withText}
       />
 
-      <ClearRunDialog
-        affectedTasks={affectedTasks}
-        dagRun={dagRun}
-        isPending={isPending}
-        mutate={mutate}
-        onClose={onClose}
-        open={open}
-      />
+      <ClearRunDialog dagRun={dagRun} onClose={onClose} open={open} />
     </Box>
   );
 };

diff --git a/airflow/ui/src/components/Clear/Run/ClearRunDialog.tsx b/airflow/ui/src/components/Clear/Run/ClearRunDialog.tsx
@@ -17,58 +17,59 @@
  * under the License.
  */
 import { Flex, Heading, VStack } from "@chakra-ui/react";
-import { useEffect, useState } from "react";
+import { useState } from "react";
 import { FiRefreshCw } from "react-icons/fi";
 
-import type {
-  DAGRunClearBody,
-  DAGRunResponse,
-  TaskInstanceCollectionResponse,
-} from "openapi/requests/types.gen";
+import type { DAGRunResponse } from "openapi/requests/types.gen";
 import { Button, Dialog } from "src/components/ui";
 import SegmentedControl from "src/components/ui/SegmentedControl";
+import { useClearDagRunDryRun } from "src/queries/useClearDagRunDryRun";
+import { useClearDagRun } from "src/queries/useClearRun";
 import { usePatchDagRun } from "src/queries/usePatchDagRun";
 
 import ClearAccordion from "../ClearAccordion";
 
 type Props = {
-  readonly affectedTasks: TaskInstanceCollectionResponse;
   readonly dagRun: DAGRunResponse;
-  readonly isPending: boolean;
-  readonly mutate: ({
-    dagId,
-    dagRunId,
-    requestBody,
-  }: {
-    dagId: string;
-    dagRunId: string;
-    requestBody: DAGRunClearBody;
-  }) => void;
   readonly onClose: () => void;
   readonly open: boolean;
 };
 
-const ClearRunDialog = ({ affectedTasks, dagRun, isPending, mutate, onClose, open }: Props) => {
+const ClearRunDialog = ({ dagRun, 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 { isPending, mutate } = useClearDagRun({
+    dagId,
+    dagRunId,
+    onSuccessConfirm: onClose,
+  });
+
+  const onlyFailed = selectedOptions.includes("onlyFailed");
+
   const [note, setNote] = useState<string | null>(dagRun.note);
   const { isPending: isPendingPatchDagRun, mutate: mutatePatchDagRun } = usePatchDagRun({ dagId, dagRunId });
 
-  useEffect(() => {
-    mutate({
-      dagId,
-      dagRunId,
-      requestBody: { dry_run: true, only_failed: onlyFailed },
-    });
-  }, [dagId, dagRunId, mutate, onlyFailed]);
+  const { data } = useClearDagRunDryRun({
+    dagId,
+    dagRunId,
+    options: {
+      enabled: open,
+    },
+    requestBody: {
+      only_failed: onlyFailed,
+    },
+  });
+
+  const affectedTasks = data ?? {
+    task_instances: [],
+    total_entries: 0,
+  };
 
   return (
-    <Dialog.Root onOpenChange={onClose} open={open} size="xl">
+    <Dialog.Root lazyMount onOpenChange={onClose} open={open} size="xl">
       <Dialog.Content backdrop>
         <Dialog.Header>
           <VStack align="start" gap={4}>
@@ -100,18 +101,21 @@ const ClearRunDialog = ({ affectedTasks, dagRun, isPending, mutate, onClose, ope
           <Flex justifyContent="end" mt={3}>
             <Button
               colorPalette="blue"
+              disabled={affectedTasks.total_entries === 0}
               loading={isPending || isPendingPatchDagRun}
               onClick={() => {
                 mutate({
                   dagId,
                   dagRunId,
                   requestBody: { dry_run: false, only_failed: onlyFailed },
                 });
-                mutatePatchDagRun({
-                  dagId,
-                  dagRunId,
-                  requestBody: { note },
-                });
+                if (note !== dagRun.note) {
+                  mutatePatchDagRun({
+                    dagId,
+                    dagRunId,
+                    requestBody: { note },
+                  });
+                }
               }}
             >
               <FiRefreshCw /> Confirm

diff --git a/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceButton.tsx b/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceButton.tsx
@@ -17,12 +17,10 @@
  * 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 type { TaskInstanceResponse } from "openapi/requests/types.gen";
 import ActionButton from "src/components/ui/ActionButton";
-import { useClearTaskInstances } from "src/queries/useClearTaskInstances";
 
 import ClearTaskInstanceDialog from "./ClearTaskInstanceDialog";
 
@@ -34,21 +32,6 @@ type Props = {
 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
@@ -59,14 +42,7 @@ const ClearTaskInstanceButton = ({ taskInstance, withText = true }: Props) => {
         withText={withText}
       />
 
-      <ClearTaskInstanceDialog
-        affectedTasks={affectedTasks}
-        isPending={isPending}
-        mutate={mutate}
-        onClose={onClose}
-        open={open}
-        taskInstance={taskInstance}
-      />
+      <ClearTaskInstanceDialog onClose={onClose} open={open} taskInstance={taskInstance} />
     </Box>
   );
 };

diff --git a/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx b/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx
@@ -17,43 +17,38 @@
  * under the License.
  */
 import { Flex, Heading, VStack } from "@chakra-ui/react";
-import { useEffect, useState } from "react";
+import { useState } from "react";
 import { FiRefreshCw } from "react-icons/fi";
 
-import type {
-  ClearTaskInstancesBody,
-  TaskInstanceCollectionResponse,
-  TaskInstanceResponse,
-} from "openapi/requests/types.gen";
+import type { TaskInstanceResponse } from "openapi/requests/types.gen";
 import Time from "src/components/Time";
 import { Button, Dialog } from "src/components/ui";
 import SegmentedControl from "src/components/ui/SegmentedControl";
+import { useClearTaskInstances } from "src/queries/useClearTaskInstances";
+import { useClearTaskInstancesDryRun } from "src/queries/useClearTaskInstancesDryRun";
 import { usePatchTaskInstance } from "src/queries/usePatchTaskInstance";
 
 import ClearAccordion from "../ClearAccordion";
 
 type Props = {
-  readonly affectedTasks: TaskInstanceCollectionResponse;
-  readonly isPending: boolean;
-  readonly mutate: ({ dagId, requestBody }: { dagId: string; requestBody: ClearTaskInstancesBody }) => void;
   readonly onClose: () => void;
   readonly open: boolean;
   readonly taskInstance: TaskInstanceResponse;
 };
 
-const ClearTaskInstanceDialog = ({
-  affectedTasks,
-  isPending,
-  mutate,
-  onClose,
-  open,
-  taskInstance,
-}: Props) => {
-  const dagId = taskInstance.dag_id;
-  const dagRunId = taskInstance.dag_run_id;
+const ClearTaskInstanceDialog = ({ onClose, open, taskInstance }: Props) => {
   const taskId = taskInstance.task_id;
   const mapIndex = taskInstance.map_index;
 
+  const dagId = taskInstance.dag_id;
+  const dagRunId = taskInstance.dag_run_id;
+
+  const { isPending, mutate } = useClearTaskInstances({
+    dagId,
+    dagRunId,
+    onSuccessConfirm: onClose,
+  });
+
   const [selectedOptions, setSelectedOptions] = useState<Array<string>>([]);
 
   const onlyFailed = selectedOptions.includes("onlyFailed");
@@ -70,24 +65,29 @@ const ClearTaskInstanceDialog = ({
     taskId,
   });
 
-  useEffect(() => {
-    mutate({
-      dagId,
-      requestBody: {
-        dag_run_id: dagRunId,
-        dry_run: true,
-        include_downstream: downstream,
-        include_future: future,
-        include_past: past,
-        include_upstream: upstream,
-        only_failed: onlyFailed,
-        task_ids: [taskId],
-      },
-    });
-  }, [dagId, dagRunId, downstream, future, mutate, onlyFailed, past, taskId, upstream]);
+  const { data } = useClearTaskInstancesDryRun({
+    dagId,
+    options: {
+      enabled: open,
+    },
+    requestBody: {
+      dag_run_id: dagRunId,
+      include_downstream: downstream,
+      include_future: future,
+      include_past: past,
+      include_upstream: upstream,
+      only_failed: onlyFailed,
+      task_ids: [taskId],
+    },
+  });
+
+  const affectedTasks = data ?? {
+    task_instances: [],
+    total_entries: 0,
+  };
 
   return (
-    <Dialog.Root onOpenChange={onClose} open={open} size="xl">
+    <Dialog.Root lazyMount onOpenChange={onClose} open={open} size="xl">
       <Dialog.Content backdrop>
         <Dialog.Header>
           <VStack align="start" gap={4}>
@@ -118,6 +118,7 @@ const ClearTaskInstanceDialog = ({
           <Flex justifyContent="end" mt={3}>
             <Button
               colorPalette="blue"
+              disabled={affectedTasks.total_entries === 0}
               loading={isPending || isPendingPatchDagRun}
               onClick={() => {
                 mutate({
@@ -133,13 +134,15 @@ const ClearTaskInstanceDialog = ({
                     task_ids: [taskId],
                   },
                 });
-                mutatePatchTaskInstance({
-                  dagId,
-                  dagRunId,
-                  mapIndex,
-                  requestBody: { note },
-                  taskId,
-                });
+                if (note !== taskInstance.note) {
+                  mutatePatchTaskInstance({
+                    dagId,
+                    dagRunId,
+                    mapIndex,
+                    requestBody: { note },
+                    taskId,
+                  });
+                }
               }}
             >
               <FiRefreshCw /> Confirm