diff --git a/airflow/ui/src/components/Graph/DagNode.tsx b/airflow/ui/src/components/Graph/DagNode.tsx index 800411b9cec00..73dfaf537f46e 100644 --- a/airflow/ui/src/components/Graph/DagNode.tsx +++ b/airflow/ui/src/components/Graph/DagNode.tsx @@ -47,11 +47,7 @@ export const DagNode = ({ > - + {dag?.dag_display_name ?? label} diff --git a/airflow/ui/src/components/Menu/RunBackfillForm.tsx b/airflow/ui/src/components/Menu/RunBackfillForm.tsx index 063b1c8e2926a..a5da47640d5a0 100644 --- a/airflow/ui/src/components/Menu/RunBackfillForm.tsx +++ b/airflow/ui/src/components/Menu/RunBackfillForm.tsx @@ -25,6 +25,7 @@ import { Alert, Button } from "src/components/ui"; import { reprocessBehaviors } from "src/constants/reprocessBehaviourParams"; import { useCreateBackfill } from "src/queries/useCreateBackfill"; import { useCreateBackfillDryRun } from "src/queries/useCreateBackfillDryRun"; +import { useTogglePause } from "src/queries/useTogglePause"; import { ErrorAlert } from "../ErrorAlert"; import { Checkbox } from "../ui/Checkbox"; @@ -38,6 +39,7 @@ const today = new Date().toISOString().slice(0, 16); const RunBackfillForm = ({ dag, onClose }: RunBackfillFormProps) => { const [errors, setErrors] = useState<{ conf?: string; date?: unknown }>({}); + const [unpause, setUnpause] = useState(true); const { control, handleSubmit, reset, watch } = useForm({ defaultValues: { @@ -69,6 +71,8 @@ const RunBackfillForm = ({ dag, onClose }: RunBackfillFormProps) => { }, }); + const { mutate: togglePause } = useTogglePause({ dagId: dag.dag_id }); + const { createBackfill, dateValidationError, error, isPending } = useCreateBackfill({ onSuccessConfirm: onClose, }); @@ -83,6 +87,14 @@ const RunBackfillForm = ({ dag, onClose }: RunBackfillFormProps) => { const dataIntervalEnd = watch("to_date"); const onSubmit = (fdata: BackfillPostBody) => { + if (unpause && dag.is_paused) { + togglePause({ + dagId: dag.dag_id, + requestBody: { + is_paused: false, + }, + }); + } createBackfill({ requestBody: fdata, }); @@ -195,6 +207,11 @@ const RunBackfillForm = ({ dag, onClose }: RunBackfillFormProps) => { {affectedTasks.total_entries} runs will be triggered ) : undefined} + {dag.is_paused ? ( + setUnpause(!unpause)}> + Unpause {dag.dag_display_name} on trigger + + ) : undefined} diff --git a/airflow/ui/src/components/TogglePause.tsx b/airflow/ui/src/components/TogglePause.tsx index 0cee25c39ca77..15dd28e74865e 100644 --- a/airflow/ui/src/components/TogglePause.tsx +++ b/airflow/ui/src/components/TogglePause.tsx @@ -17,19 +17,10 @@ * under the License. */ import { useDisclosure } from "@chakra-ui/react"; -import { useQueryClient } from "@tanstack/react-query"; import { useCallback } from "react"; -import { - UseDagRunServiceGetDagRunsKeyFn, - UseDagServiceGetDagDetailsKeyFn, - UseDagServiceGetDagKeyFn, - useDagServiceGetDagsKey, - useDagServicePatchDag, - useDagsServiceRecentDagRunsKey, - UseTaskInstanceServiceGetTaskInstancesKeyFn, -} from "openapi/queries"; import { useConfig } from "src/queries/useConfig"; +import { useTogglePause } from "src/queries/useTogglePause"; import { ConfirmationModal } from "./ConfirmationModal"; import { Switch, type SwitchProps } from "./ui"; @@ -37,41 +28,24 @@ import { Switch, type SwitchProps } from "./ui"; type Props = { readonly dagDisplayName?: string; readonly dagId: string; - readonly isPaused: boolean; + readonly isPaused?: boolean; readonly skipConfirm?: boolean; } & SwitchProps; export const TogglePause = ({ dagDisplayName, dagId, isPaused, skipConfirm, ...rest }: Props) => { - const queryClient = useQueryClient(); const { onClose, onOpen, open } = useDisclosure(); - const onSuccess = async () => { - const queryKeys = [ - [useDagServiceGetDagsKey], - [useDagsServiceRecentDagRunsKey], - UseDagServiceGetDagKeyFn({ dagId }, [{ dagId }]), - UseDagServiceGetDagDetailsKeyFn({ dagId }, [{ dagId }]), - UseDagRunServiceGetDagRunsKeyFn({ dagId }, [{ dagId }]), - UseTaskInstanceServiceGetTaskInstancesKeyFn({ dagId, dagRunId: "~" }, [{ dagId, dagRunId: "~" }]), - ]; - - await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); - }; - - const { mutate } = useDagServicePatchDag({ - onSuccess, - }); - + const { mutate: togglePause } = useTogglePause({ dagId }); const showConfirmation = Boolean(useConfig("require_confirmation_dag_change")); const onToggle = useCallback(() => { - mutate({ + togglePause({ dagId, requestBody: { is_paused: !isPaused, }, }); - }, [dagId, isPaused, mutate]); + }, [dagId, isPaused, togglePause]); const onChange = () => { if (showConfirmation && skipConfirm !== true) { @@ -83,7 +57,13 @@ export const TogglePause = ({ dagDisplayName, dagId, isPaused, skipConfirm, ...r return ( <> - + void; readonly open: boolean; }; @@ -44,11 +47,14 @@ export type DagRunTriggerParams = { note: string; }; -const TriggerDAGForm = ({ dagId, onClose, open }: TriggerDAGFormProps) => { +const TriggerDAGForm = ({ dagId, isPaused, onClose, open }: TriggerDAGFormProps) => { const [errors, setErrors] = useState<{ conf?: string; date?: unknown }>({}); const initialParamsDict = useDagParams(dagId, open); const { error: errorTrigger, isPending, triggerDagRun } = useTrigger({ dagId, onSuccessConfirm: onClose }); const { conf, setConf } = useParamStore(); + const [unpause, setUnpause] = useState(true); + + const { mutate: togglePause } = useTogglePause({ dagId }); const { control, handleSubmit, reset } = useForm({ defaultValues: { @@ -67,6 +73,14 @@ const TriggerDAGForm = ({ dagId, onClose, open }: TriggerDAGFormProps) => { }, [conf, reset]); const onSubmit = (data: DagRunTriggerParams) => { + if (unpause && isPaused) { + togglePause({ + dagId, + requestBody: { + is_paused: false, + }, + }); + } triggerDagRun(data); }; @@ -186,6 +200,11 @@ const TriggerDAGForm = ({ dagId, onClose, open }: TriggerDAGFormProps) => { + {isPaused ? ( + setUnpause(!unpause)}> + Unpause {dagId} on trigger + + ) : undefined} diff --git a/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx b/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx index f241fada325eb..717087bffa68a 100644 --- a/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx +++ b/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx @@ -19,9 +19,8 @@ import { Heading, VStack } from "@chakra-ui/react"; import React from "react"; -import { Alert, Dialog } from "src/components/ui"; +import { Dialog } from "src/components/ui"; -import { TogglePause } from "../TogglePause"; import TriggerDAGForm from "./TriggerDAGForm"; type TriggerDAGModalProps = { @@ -43,21 +42,14 @@ const TriggerDAGModal: React.FC = ({ - - Trigger Dag - {dagDisplayName} - - {isPaused ? ( - - Triggering will create a DAG run, but it will not start until the Dag is unpaused. - - ) : undefined} + Trigger Dag - {dagDisplayName} - + diff --git a/airflow/ui/src/pages/Asset/CreateAssetEventModal.tsx b/airflow/ui/src/pages/Asset/CreateAssetEventModal.tsx index 18a9046b8c950..61846903ab400 100644 --- a/airflow/ui/src/pages/Asset/CreateAssetEventModal.tsx +++ b/airflow/ui/src/pages/Asset/CreateAssetEventModal.tsx @@ -26,6 +26,7 @@ import { UseAssetServiceGetAssetEventsKeyFn, useAssetServiceMaterializeAsset, UseDagRunServiceGetDagRunsKeyFn, + useDagServiceGetDagDetails, useDagsServiceRecentDagRunsKey, useDependenciesServiceGetDependencies, UseGridServiceGridDataKeyFn, @@ -40,7 +41,9 @@ import type { import { ErrorAlert } from "src/components/ErrorAlert"; import { JsonEditor } from "src/components/JsonEditor"; import { Dialog, toaster } from "src/components/ui"; +import { Checkbox } from "src/components/ui/Checkbox"; import { RadioCardItem, RadioCardRoot } from "src/components/ui/RadioCard"; +import { useTogglePause } from "src/queries/useTogglePause"; type Props = { readonly asset: AssetResponse; @@ -51,6 +54,7 @@ type Props = { export const CreateAssetEventModal = ({ asset, onClose, open }: Props) => { const [eventType, setEventType] = useState("manual"); const [extraError, setExtraError] = useState(); + const [unpause, setUnpause] = useState(true); const [extra, setExtra] = useState("{}"); const queryClient = useQueryClient(); @@ -118,6 +122,12 @@ export const CreateAssetEventModal = ({ asset, onClose, open }: Props) => { await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); }; + const { data: dag } = useDagServiceGetDagDetails({ dagId: upstreamDagId ?? "" }, undefined, { + enabled: Boolean(upstreamDagId), + }); + + const { mutate: togglePause } = useTogglePause({ dagId: dag?.dag_id ?? upstreamDagId ?? "" }); + const { error: manualError, isPending, @@ -133,6 +143,14 @@ export const CreateAssetEventModal = ({ asset, onClose, open }: Props) => { const handleSubmit = () => { if (eventType === "materialize") { + if (unpause && dag?.is_paused) { + togglePause({ + dagId: dag.dag_id, + requestBody: { + is_paused: false, + }, + }); + } materializeAsset({ assetId: asset.id }); } else { createAssetEvent({ @@ -162,7 +180,7 @@ export const CreateAssetEventModal = ({ asset, onClose, open }: Props) => { > { {extraError} ) : undefined} + {eventType === "materialize" && dag?.is_paused ? ( + setUnpause(!unpause)}> + Unpause {dag.dag_display_name} on trigger + + ) : undefined} diff --git a/airflow/ui/src/queries/useTogglePause.ts b/airflow/ui/src/queries/useTogglePause.ts new file mode 100644 index 0000000000000..c4e378e7932bc --- /dev/null +++ b/airflow/ui/src/queries/useTogglePause.ts @@ -0,0 +1,50 @@ +/*! + * 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 { useQueryClient } from "@tanstack/react-query"; + +import { + UseDagRunServiceGetDagRunsKeyFn, + UseDagServiceGetDagDetailsKeyFn, + UseDagServiceGetDagKeyFn, + useDagServiceGetDagsKey, + useDagServicePatchDag, + useDagsServiceRecentDagRunsKey, + UseTaskInstanceServiceGetTaskInstancesKeyFn, +} from "openapi/queries"; + +export const useTogglePause = ({ dagId }: { dagId: string }) => { + const queryClient = useQueryClient(); + + const onSuccess = async () => { + const queryKeys = [ + [useDagServiceGetDagsKey], + [useDagsServiceRecentDagRunsKey], + UseDagServiceGetDagKeyFn({ dagId }, [{ dagId }]), + UseDagServiceGetDagDetailsKeyFn({ dagId }, [{ dagId }]), + UseDagRunServiceGetDagRunsKeyFn({ dagId }, [{ dagId }]), + UseTaskInstanceServiceGetTaskInstancesKeyFn({ dagId, dagRunId: "~" }, [{ dagId, dagRunId: "~" }]), + ]; + + await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); + }; + + return useDagServicePatchDag({ + onSuccess, + }); +};