Integrate API with Trigger Dag Run (apache#44850)

* conf * conf check * basic trigger api * retrieve params * refactor the api * fix * fix details page * refactor the logic * refactor * reviews * error on form * data date validation * remove empty runid and note
got686-yandex · Jan 30, 2025 · 85f27f2 · 85f27f2
1 parent adc33e1
commit 85f27f2
Show file tree

Hide file tree

Showing 6 changed files with 254 additions and 142 deletions.
diff --git a/airflow/ui/src/components/DataTable/DataTable.tsx b/airflow/ui/src/components/DataTable/DataTable.tsx
@@ -30,7 +30,7 @@ import {
 } from "@tanstack/react-table";
 import React, { type ReactNode, useCallback, useRef } from "react";
 
-import { ProgressBar, Pagination } from "../ui";
+import { ProgressBar, Pagination, Toaster } from "../ui";
 import { CardList } from "./CardList";
 import { TableList } from "./TableList";
 import { createSkeletonMock } from "./skeleton";
@@ -128,6 +128,7 @@ export const DataTable = <TData,>({
           Boolean(isFetching) && !Boolean(isLoading) ? "visible" : "hidden"
         }
       />
+      <Toaster />
       {errorMessage}
       {!Boolean(isLoading) && !rows.length && (
         <Text pt={1}>{noRowsMessage ?? `No ${modelName}s found.`}</Text>

diff --git a/airflow/ui/src/components/TriggerDag/TriggerDAGForm.tsx b/airflow/ui/src/components/TriggerDag/TriggerDAGForm.tsx
@@ -20,28 +20,50 @@ import { Input, Button, Box, Text, Spacer, HStack } from "@chakra-ui/react";
 import { json } from "@codemirror/lang-json";
 import { githubLight, githubDark } from "@uiw/codemirror-themes-all";
 import CodeMirror from "@uiw/react-codemirror";
-import { useEffect, useState } from "react";
+import { useEffect, useMemo, useState } from "react";
 import { useForm, Controller } from "react-hook-form";
 import { FiPlay } from "react-icons/fi";
 
 import { useColorMode } from "src/context/colorMode";
+import { useDagParams } from "src/queries/useDagParams";
+import { useTrigger } from "src/queries/useTrigger";
 
+import { ErrorAlert } from "../ErrorAlert";
 import { Accordion } from "../ui";
-import type { DagParams } from "./TriggerDag";
 
 type TriggerDAGFormProps = {
-  dagParams: DagParams;
+  dagId: string;
   onClose: () => void;
-  onTrigger: (updatedDagParams: DagParams) => void;
-  setDagParams: React.Dispatch<React.SetStateAction<DagParams>>;
+  open: boolean;
+};
+
+export type DagRunTriggerParams = {
+  conf: string;
+  dagRunId: string;
+  dataIntervalEnd: string;
+  dataIntervalStart: string;
+  note: string;
 };
 
 const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
-  dagParams,
-  onTrigger,
-  setDagParams,
+  dagId,
+  onClose,
+  open,
 }) => {
-  const [jsonError, setJsonError] = useState<string | undefined>();
+  const [errors, setErrors] = useState<{ conf?: string; date?: string }>({});
+  const conf = useDagParams(dagId, open);
+  const { error: errorTrigger, isPending, triggerDagRun } = useTrigger(onClose);
+
+  const dagRunRequestBody: DagRunTriggerParams = useMemo(
+    () => ({
+      conf,
+      dagRunId: "",
+      dataIntervalEnd: "",
+      dataIntervalStart: "",
+      note: "",
+    }),
+    [conf],
+  );
 
   const {
     control,
@@ -50,40 +72,47 @@ const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
     reset,
     setValue,
     watch,
-  } = useForm<DagParams>({
-    defaultValues: dagParams,
-  });
+  } = useForm<DagRunTriggerParams>({ defaultValues: dagRunRequestBody });
 
   const dataIntervalStart = watch("dataIntervalStart");
   const dataIntervalEnd = watch("dataIntervalEnd");
 
   useEffect(() => {
-    reset(dagParams);
-  }, [dagParams, reset]);
-
-  const onSubmit = (data: DagParams) => {
-    onTrigger(data);
-    setDagParams(data);
-    setJsonError(undefined);
-  };
+    reset(dagRunRequestBody);
+  }, [dagRunRequestBody, reset]);
 
   const validateAndPrettifyJson = (value: string) => {
     try {
       const parsedJson = JSON.parse(value) as JSON;
 
-      setJsonError(undefined);
+      setErrors((prev) => ({ ...prev, conf: undefined }));
 
       return JSON.stringify(parsedJson, undefined, 2);
     } catch (error) {
       const errorMessage =
         error instanceof Error ? error.message : "Unknown error occurred.";
 
-      setJsonError(`Invalid JSON format: ${errorMessage}`);
+      setErrors((prev) => ({
+        ...prev,
+        conf: `Invalid JSON format: ${errorMessage}`,
+      }));
 
       return value;
     }
   };
 
+  const onSubmit = (data: DagRunTriggerParams) => {
+    if (Boolean(data.dataIntervalStart) !== Boolean(data.dataIntervalEnd)) {
+      setErrors((prev) => ({
+        ...prev,
+        date: "Either both Data Interval Start and End must be provided, or both must be empty.",
+      }));
+
+      return;
+    }
+    triggerDagRun(dagId, data);
+  };
+
   const validateDates = (
     fieldName: "dataIntervalEnd" | "dataIntervalStart",
   ) => {
@@ -92,6 +121,8 @@ const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
       : undefined;
     const endDate = dataIntervalEnd ? new Date(dataIntervalEnd) : undefined;
 
+    setErrors((prev) => ({ ...prev, date: undefined }));
+
     if (startDate && endDate) {
       if (fieldName === "dataIntervalStart" && startDate > endDate) {
         setValue("dataIntervalStart", dataIntervalEnd);
@@ -105,7 +136,8 @@ const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
 
   return (
     <>
-      <Accordion.Root collapsible size="lg" variant="enclosed">
+      <ErrorAlert error={errorTrigger} />
+      <Accordion.Root collapsible mt={4} size="lg" variant="enclosed">
         <Accordion.Item key="advancedOptions" value="advancedOptions">
           <Accordion.ItemTrigger cursor="button">
             Advanced Options
@@ -153,7 +185,7 @@ const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
               </Text>
               <Controller
                 control={control}
-                name="runId"
+                name="dagRunId"
                 render={({ field }) => (
                   <Input
                     {...field}
@@ -168,7 +200,7 @@ const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
               </Text>
               <Controller
                 control={control}
-                name="configJson"
+                name="conf"
                 render={({ field }) => (
                   <Box mb={4}>
                     <CodeMirror
@@ -196,11 +228,11 @@ const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
                       }}
                       theme={colorMode === "dark" ? githubDark : githubLight}
                     />
-                    {Boolean(jsonError) ? (
+                    {Boolean(errors.conf) && (
                       <Text color="red.500" fontSize="sm" mt={2}>
-                        {jsonError}
+                        {errors.conf}
                       </Text>
-                    ) : undefined}
+                    )}
                   </Box>
                 )}
               />
@@ -210,7 +242,7 @@ const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
               </Text>
               <Controller
                 control={control}
-                name="notes"
+                name="note"
                 render={({ field }) => (
                   <Input {...field} placeholder="Optional" size="sm" />
                 )}
@@ -219,7 +251,11 @@ const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
           </Accordion.ItemContent>
         </Accordion.Item>
       </Accordion.Root>
-
+      {Boolean(errors.date) && (
+        <Text color="red.500" fontSize="sm" mt={2}>
+          {errors.date}
+        </Text>
+      )}
       <Box as="footer" display="flex" justifyContent="flex-end" mt={4}>
         <HStack w="full">
           {isDirty ? (
@@ -230,7 +266,7 @@ const TriggerDAGForm: React.FC<TriggerDAGFormProps> = ({
           <Spacer />
           <Button
             colorPalette="blue"
-            disabled={Boolean(jsonError)}
+            disabled={Boolean(errors.conf) || Boolean(errors.date) || isPending}
             onClick={() => void handleSubmit(onSubmit)()}
           >
             <FiPlay /> Trigger

diff --git a/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx b/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx
@@ -17,14 +17,12 @@
  * under the License.
  */
 import { Heading, VStack } from "@chakra-ui/react";
-import React, { useCallback, useEffect, useMemo, useState } from "react";
+import React from "react";
 
 import { Alert, Dialog } from "src/components/ui";
 
 import { TogglePause } from "../TogglePause";
 import TriggerDAGForm from "./TriggerDAGForm";
-import type { DagParams } from "./TriggerDag";
-import { TriggerDag as triggerDag } from "./TriggerDag";
 
 type TriggerDAGModalProps = {
   dagDisplayName: string;
@@ -40,70 +38,37 @@ const TriggerDAGModal: React.FC<TriggerDAGModalProps> = ({
   isPaused,
   onClose,
   open,
-}) => {
-  const initialDagParams = useMemo(
-    () => ({
-      configJson: "{}",
-      dagId,
-      dataIntervalEnd: "",
-      dataIntervalStart: "",
-      notes: "",
-      runId: "",
-    }),
-    [dagId],
-  );
+}) => (
+  <Dialog.Root
+    lazyMount
+    onOpenChange={onClose}
+    open={open}
+    size="xl"
+    unmountOnExit
+  >
+    <Dialog.Content backdrop>
+      <Dialog.Header>
+        <VStack align="start" gap={4}>
+          <Heading size="xl">
+            Trigger DAG - {dagDisplayName}{" "}
+            <TogglePause dagId={dagId} isPaused={isPaused} skipConfirm />
+          </Heading>
+          {isPaused ? (
+            <Alert status="warning" title="Paused DAG">
+              Triggering will create a DAG run, but it will not start until the
+              DAG is unpaused.
+            </Alert>
+          ) : undefined}
+        </VStack>
+      </Dialog.Header>
 
-  const [dagParams, setDagParams] = useState<DagParams>(initialDagParams);
+      <Dialog.CloseTrigger />
 
-  const handleTrigger = useCallback(
-    (updatedDagParams: DagParams) => {
-      triggerDag(updatedDagParams);
-      onClose();
-    },
-    [onClose],
-  );
-
-  useEffect(() => {
-    if (!open) {
-      setDagParams(initialDagParams);
-    }
-  }, [open, initialDagParams]);
-
-  return (
-    <Dialog.Root onOpenChange={onClose} open={open} size="xl">
-      <Dialog.Content backdrop>
-        <Dialog.Header>
-          <VStack align="start" gap={4}>
-            <Heading size="xl">
-              Trigger DAG - {dagDisplayName}{" "}
-              <TogglePause
-                dagId={dagParams.dagId}
-                isPaused={isPaused}
-                skipConfirm
-              />
-            </Heading>
-            {isPaused ? (
-              <Alert status="warning" title="Paused DAG">
-                Triggering will create a DAG run, but it will not start until
-                the DAG is unpaused.
-              </Alert>
-            ) : undefined}
-          </VStack>
-        </Dialog.Header>
-
-        <Dialog.CloseTrigger />
-
-        <Dialog.Body>
-          <TriggerDAGForm
-            dagParams={dagParams}
-            onClose={onClose}
-            onTrigger={handleTrigger}
-            setDagParams={setDagParams}
-          />
-        </Dialog.Body>
-      </Dialog.Content>
-    </Dialog.Root>
-  );
-};
+      <Dialog.Body>
+        <TriggerDAGForm dagId={dagId} onClose={onClose} open={open} />
+      </Dialog.Body>
+    </Dialog.Content>
+  </Dialog.Root>
+);
 
 export default TriggerDAGModal;
diff --git a/airflow/ui/src/components/TriggerDag/TriggerDag.tsx b/airflow/ui/src/components/TriggerDag/TriggerDag.tsx