This is an automated email from the ASF dual-hosted git repository.
bbovenzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f5f5200a909 Fix run task/run clearing (#45987)
f5f5200a909 is described below
commit f5f5200a909af434e4f31a2aa03567f051e4438c
Author: Brent Bovenzi <[email protected]>
AuthorDate: Fri Jan 24 11:46:17 2025 -0500
Fix run task/run clearing (#45987)
* Fix clear modals
* Reduce clear task/run API calls
* Fix and add tests to datatable
---
airflow/ui/src/components/Clear/ClearAccordion.tsx | 2 +-
.../ui/src/components/Clear/Run/ClearRunButton.tsx | 28 +------
.../ui/src/components/Clear/Run/ClearRunDialog.tsx | 70 +++++++++--------
.../Clear/TaskInstance/ClearTaskInstanceButton.tsx | 28 +------
.../Clear/TaskInstance/ClearTaskInstanceDialog.tsx | 87 +++++++++++-----------
.../ui/src/components/DataTable/DataTable.test.tsx | 22 +++++-
airflow/ui/src/components/DataTable/DataTable.tsx | 5 +-
airflow/ui/src/components/ui/Button.tsx | 2 +-
airflow/ui/src/layouts/Details/DagVizModal.tsx | 2 +-
airflow/ui/src/queries/useClearDagRunDryRun.ts | 49 ++++++++++++
airflow/ui/src/queries/useClearRun.ts | 32 +++-----
airflow/ui/src/queries/useClearTaskInstances.ts | 49 +++++-------
.../ui/src/queries/useClearTaskInstancesDryRun.ts | 56 ++++++++++++++
13 files changed, 247 insertions(+), 185 deletions(-)
diff --git a/airflow/ui/src/components/Clear/ClearAccordion.tsx
b/airflow/ui/src/components/Clear/ClearAccordion.tsx
index 74be262da75..8c5bc7e9184 100644
--- 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
index 151e2f3bd9e..30a3f2aa3df 100644
--- 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
index 4cb48669b55..47e9b95b47d 100644
--- 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,6 +101,7 @@ 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({
@@ -107,11 +109,13 @@ const ClearRunDialog = ({ affectedTasks, dagRun,
isPending, mutate, onClose, ope
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
index 8f4a6275667..fb39b45e0bd 100644
--- 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
index 45c5a1c2d0f..b2344bf90d2 100644
--- 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
diff --git a/airflow/ui/src/components/DataTable/DataTable.test.tsx
b/airflow/ui/src/components/DataTable/DataTable.test.tsx
index 3d593869d85..e38bcc2de8a 100644
--- a/airflow/ui/src/components/DataTable/DataTable.test.tsx
+++ b/airflow/ui/src/components/DataTable/DataTable.test.tsx
@@ -67,7 +67,7 @@ describe("DataTable", () => {
render(
<DataTable
columns={columns}
- data={data}
+ data={[{ name: "John Doe" }]}
initialState={{ pagination, sorting: [] }}
onStateChange={onStateChange}
total={2}
@@ -84,7 +84,7 @@ describe("DataTable", () => {
render(
<DataTable
columns={columns}
- data={data}
+ data={[{ name: "John Doe" }]}
initialState={{
pagination: { pageIndex: 0, pageSize: 10 },
sorting: [],
@@ -100,6 +100,24 @@ describe("DataTable", () => {
expect(screen.getByTestId("next")).toBeDisabled();
});
+ it("renders no pagination when not needed", () => {
+ render(
+ <DataTable
+ columns={columns}
+ data={data}
+ initialState={{ pagination, sorting: [] }}
+ onStateChange={onStateChange}
+ total={2}
+ />,
+ {
+ wrapper: ChakraWrapper,
+ },
+ );
+
+ expect(screen.queryByTestId("prev")).toBeNull();
+ expect(screen.queryByTestId("next")).toBeNull();
+ });
+
it("when isLoading renders skeleton columns", () => {
render(<DataTable columns={columns} data={data} isLoading />, {
wrapper: ChakraWrapper,
diff --git a/airflow/ui/src/components/DataTable/DataTable.tsx
b/airflow/ui/src/components/DataTable/DataTable.tsx
index 4f02b57cb8a..86182edb0b8 100644
--- a/airflow/ui/src/components/DataTable/DataTable.tsx
+++ b/airflow/ui/src/components/DataTable/DataTable.tsx
@@ -116,6 +116,9 @@ export const DataTable = <TData,>({
const display = displayMode === "card" && Boolean(cardDef) ? "card" :
"table";
const hasRows = rows.length > 0;
+ const hasPagination =
+ table.getState().pagination.pageIndex !== 0 ||
+ (table.getState().pagination.pageIndex === 0 && rows.length !== total);
return (
<>
@@ -127,7 +130,7 @@ export const DataTable = <TData,>({
<CardList cardDef={cardDef} isLoading={isLoading} table={table} />
) : undefined}
{!hasRows && !Boolean(isLoading) && <Text pt={1}>{noRowsMessage ?? `No
${modelName}s found.`}</Text>}
- {hasRows ? (
+ {hasPagination ? (
<Pagination.Root
count={table.getRowCount()}
my={2}
diff --git a/airflow/ui/src/components/ui/Button.tsx
b/airflow/ui/src/components/ui/Button.tsx
index d4d41836696..1e0372bc5a9 100644
--- a/airflow/ui/src/components/ui/Button.tsx
+++ b/airflow/ui/src/components/ui/Button.tsx
@@ -31,7 +31,7 @@ export const Button = React.forwardRef<HTMLButtonElement,
ButtonProps>((props, r
const { children, disabled, loading, loadingText, ...rest } = props;
return (
- <ChakraButton disabled={loading ?? disabled} ref={ref} {...rest}>
+ <ChakraButton disabled={disabled ?? loading} ref={ref} {...rest}>
{loading && !Boolean(loadingText) ? (
<>
<AbsoluteCenter display="inline-flex">
diff --git a/airflow/ui/src/layouts/Details/DagVizModal.tsx
b/airflow/ui/src/layouts/Details/DagVizModal.tsx
index 864be59ca5f..a4286c6f736 100644
--- a/airflow/ui/src/layouts/Details/DagVizModal.tsx
+++ b/airflow/ui/src/layouts/Details/DagVizModal.tsx
@@ -63,7 +63,7 @@ export const DagVizModal: React.FC<DAGVizModalProps> = ({
dagDisplayName, dagId,
params.delete("modal");
return (
- <Dialog.Root motionPreset="none" onOpenChange={onClose} open={open}
size="full">
+ <Dialog.Root lazyMount motionPreset="none" onOpenChange={onClose}
open={open} size="full">
<Dialog.Content backdrop ref={contentRef}>
<Dialog.Header bg="blue.muted" pr={16}>
<HStack>
diff --git a/airflow/ui/src/queries/useClearDagRunDryRun.ts
b/airflow/ui/src/queries/useClearDagRunDryRun.ts
new file mode 100644
index 00000000000..a15099e69ec
--- /dev/null
+++ b/airflow/ui/src/queries/useClearDagRunDryRun.ts
@@ -0,0 +1,49 @@
+/*!
+ * 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 { useQuery, type UseQueryOptions } from "@tanstack/react-query";
+
+import { DagRunService } from "openapi/requests/services.gen";
+import type { DAGRunClearBody, TaskInstanceCollectionResponse } from
"openapi/requests/types.gen";
+
+type Props<TData, TError> = {
+ dagId: string;
+ dagRunId: string;
+ options?: Omit<UseQueryOptions<TData, TError>, "queryFn" | "queryKey">;
+ requestBody: DAGRunClearBody;
+};
+
+export const useClearDagRunDryRun = <TData = TaskInstanceCollectionResponse,
TError = unknown>({
+ dagId,
+ dagRunId,
+ options,
+ requestBody,
+}: Props<TData, TError>) =>
+ useQuery<TData, TError>({
+ ...options,
+ queryFn: () =>
+ DagRunService.clearDagRun({
+ dagId,
+ dagRunId,
+ requestBody: {
+ dry_run: true,
+ ...requestBody,
+ },
+ }) as TData,
+ queryKey: ["clearDagRun", dagId, requestBody.only_failed],
+ });
diff --git a/airflow/ui/src/queries/useClearRun.ts
b/airflow/ui/src/queries/useClearRun.ts
index 336b42e74d8..526662f0b3e 100644
--- a/airflow/ui/src/queries/useClearRun.ts
+++ b/airflow/ui/src/queries/useClearRun.ts
@@ -25,7 +25,6 @@ import {
UseDagServiceGetDagDetailsKeyFn,
useTaskInstanceServiceGetTaskInstancesKey,
} from "openapi/queries";
-import type { DAGRunClearBody, TaskInstanceCollectionResponse } from
"openapi/requests/types.gen";
import { toaster } from "src/components/ui";
const onError = () => {
@@ -40,37 +39,24 @@ export const useClearDagRun = ({
dagId,
dagRunId,
onSuccessConfirm,
- onSuccessDryRun,
}: {
dagId: string;
dagRunId: string;
onSuccessConfirm: () => void;
- onSuccessDryRun: (date: TaskInstanceCollectionResponse) => void;
}) => {
const queryClient = useQueryClient();
- const onSuccess = async (
- data: TaskInstanceCollectionResponse,
- variables: {
- dagId: string;
- dagRunId: string;
- requestBody: DAGRunClearBody;
- },
- ) => {
- if (variables.requestBody.dry_run) {
- onSuccessDryRun(data);
- } else {
- const queryKeys = [
- [useTaskInstanceServiceGetTaskInstancesKey],
- UseDagServiceGetDagDetailsKeyFn({ dagId }),
- UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }),
- [useDagRunServiceGetDagRunsKey],
- ];
+ const onSuccess = async () => {
+ const queryKeys = [
+ [useTaskInstanceServiceGetTaskInstancesKey],
+ UseDagServiceGetDagDetailsKeyFn({ dagId }),
+ UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }),
+ [useDagRunServiceGetDagRunsKey],
+ ];
- await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({
queryKey: key })));
+ await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({
queryKey: key })));
- onSuccessConfirm();
- }
+ onSuccessConfirm();
};
return useDagRunServiceClearDagRun({
diff --git a/airflow/ui/src/queries/useClearTaskInstances.ts
b/airflow/ui/src/queries/useClearTaskInstances.ts
index 3770a88616a..d03def7d693 100644
--- a/airflow/ui/src/queries/useClearTaskInstances.ts
+++ b/airflow/ui/src/queries/useClearTaskInstances.ts
@@ -40,49 +40,40 @@ export const useClearTaskInstances = ({
dagId,
dagRunId,
onSuccessConfirm,
- onSuccessDryRun,
}: {
dagId: string;
dagRunId: string;
onSuccessConfirm: () => void;
- onSuccessDryRun: (date: TaskInstanceCollectionResponse) => void;
}) => {
const queryClient = useQueryClient();
const onSuccess = async (
- data: TaskInstanceCollectionResponse,
- variables: {
- dagId: string;
- requestBody: ClearTaskInstancesBody;
- },
+ _: TaskInstanceCollectionResponse,
+ variables: { dagId: string; requestBody: ClearTaskInstancesBody },
) => {
- if (variables.requestBody.dry_run) {
- onSuccessDryRun(data);
- } else {
- const taskInstanceKeys = (variables.requestBody.task_ids ?? [])
- .map((taskId) => {
- const runId = variables.requestBody.dag_run_id;
+ const taskInstanceKeys = (variables.requestBody.task_ids ?? [])
+ .map((taskId) => {
+ const runId = variables.requestBody.dag_run_id;
- if (runId === null || runId === undefined) {
- return undefined;
- }
- const params = { dagId, dagRunId: runId, taskId };
+ if (runId === null || runId === undefined) {
+ return undefined;
+ }
+ const params = { dagId, dagRunId: runId, taskId };
- return UseTaskInstanceServiceGetTaskInstanceKeyFn(params);
- })
- .filter((key) => key !== undefined);
+ return UseTaskInstanceServiceGetTaskInstanceKeyFn(params);
+ })
+ .filter((key) => key !== undefined);
- const queryKeys = [
- [useTaskInstanceServiceGetTaskInstancesKey],
- ...taskInstanceKeys,
- UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }),
- [useDagRunServiceGetDagRunsKey],
- ];
+ const queryKeys = [
+ [useTaskInstanceServiceGetTaskInstancesKey],
+ ...taskInstanceKeys,
+ UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }),
+ [useDagRunServiceGetDagRunsKey],
+ ];
- await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({
queryKey: key })));
+ await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({
queryKey: key })));
- onSuccessConfirm();
- }
+ onSuccessConfirm();
};
return useTaskInstanceServicePostClearTaskInstances({
diff --git a/airflow/ui/src/queries/useClearTaskInstancesDryRun.ts
b/airflow/ui/src/queries/useClearTaskInstancesDryRun.ts
new file mode 100644
index 00000000000..e85fae59a57
--- /dev/null
+++ b/airflow/ui/src/queries/useClearTaskInstancesDryRun.ts
@@ -0,0 +1,56 @@
+/*!
+ * 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 { useQuery, type UseQueryOptions } from "@tanstack/react-query";
+
+import { TaskInstanceService } from "openapi/requests/services.gen";
+import type { ClearTaskInstancesBody, PostClearTaskInstancesResponse } from
"openapi/requests/types.gen";
+
+type Props<TData, TError> = {
+ dagId: string;
+ options?: Omit<UseQueryOptions<TData, TError>, "queryFn" | "queryKey">;
+ requestBody: ClearTaskInstancesBody;
+};
+
+export const useClearTaskInstancesDryRun = <TData =
PostClearTaskInstancesResponse, TError = unknown>({
+ dagId,
+ options,
+ requestBody,
+}: Props<TData, TError>) =>
+ useQuery<TData, TError>({
+ ...options,
+ queryFn: () =>
+ TaskInstanceService.postClearTaskInstances({
+ dagId,
+ requestBody: {
+ dry_run: true,
+ ...requestBody,
+ },
+ }) as TData,
+ queryKey: [
+ "clearTaskInstance",
+ dagId,
+ requestBody.dag_run_id,
+ requestBody.only_failed,
+ requestBody.task_ids,
+ requestBody.include_downstream,
+ requestBody.include_future,
+ requestBody.include_past,
+ requestBody.include_upstream,
+ ],
+ });
