From 209806497e2cc2487ba41e5154a1cca53057020e Mon Sep 17 00:00:00 2001 From: nothinmin Date: Thu, 11 Sep 2025 19:16:36 +0900 Subject: [PATCH 1/4] Fix Grid cache invalidation for multi-run task operations --- .../ui/src/queries/useClearTaskInstances.ts | 13 ++++++++--- .../ui/src/queries/usePatchTaskInstance.ts | 23 +++++++++++++++---- 2 files changed, 29 insertions(+), 7 deletions(-) diff --git a/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts b/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts index ed3431bba08e3..0b81aed63d9c6 100644 --- a/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts +++ b/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts @@ -26,6 +26,7 @@ import { useTaskInstanceServicePostClearTaskInstances, UseGridServiceGetGridRunsKeyFn, UseGridServiceGetGridTiSummariesKeyFn, + useGridServiceGetGridTiSummariesKey, } from "openapi/queries"; import type { ClearTaskInstancesBody, TaskInstanceCollectionResponse } from "openapi/requests/types.gen"; import { toaster } from "src/components/ui"; @@ -78,17 +79,23 @@ export const useClearTaskInstances = ({ ), ]; - const queryKeys = [ + // Check if this clear operation affects multiple DAG runs + const { include_future: includeFuture, include_past: includePast } = variables.requestBody; + const affectsMultipleRuns = includeFuture === true || includePast === true; + + const invalidateKeys = [ ...taskInstanceKeys, UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }), [useDagRunServiceGetDagRunsKey], [useClearTaskInstancesDryRunKey, dagId], [usePatchTaskInstanceDryRunKey, dagId, dagRunId], UseGridServiceGetGridRunsKeyFn({ dagId }, [{ dagId }]), - UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: dagRunId }, [{ dagId, runId: dagRunId }]), + affectsMultipleRuns + ? [useGridServiceGetGridTiSummariesKey, { dagId }] + : UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: dagRunId }), ]; - await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); + await Promise.all(invalidateKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); onSuccessConfirm(); }; diff --git a/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts b/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts index bc2747b1a575b..44e8dcac4f327 100644 --- a/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts +++ b/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts @@ -26,6 +26,7 @@ import { useTaskInstanceServicePatchTaskInstance, UseGridServiceGetGridRunsKeyFn, UseGridServiceGetGridTiSummariesKeyFn, + useGridServiceGetGridTiSummariesKey, } from "openapi/queries"; import { toaster } from "src/components/ui"; @@ -58,18 +59,32 @@ export const usePatchTaskInstance = ({ }); }; - const onSuccessFn = async () => { - const queryKeys = [ + const onSuccessFn = async ( + _: unknown, + variables: { + dagId: string; + dagRunId: string; + requestBody: { include_future?: boolean; include_past?: boolean }; + taskId: string; + }, + ) => { + // Check if this patch operation affects multiple DAG runs + const { include_future: includeFuture, include_past: includePast } = variables.requestBody; + const affectsMultipleRuns = includeFuture === true || includePast === true; + + const invalidateKeys = [ UseTaskInstanceServiceGetTaskInstanceKeyFn({ dagId, dagRunId, taskId }), UseTaskInstanceServiceGetMappedTaskInstanceKeyFn({ dagId, dagRunId, mapIndex, taskId }), [useTaskInstanceServiceGetTaskInstancesKey], [usePatchTaskInstanceDryRunKey, dagId, dagRunId, { mapIndex, taskId }], [useClearTaskInstancesDryRunKey, dagId], UseGridServiceGetGridRunsKeyFn({ dagId }, [{ dagId }]), - UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: dagRunId }, [{ dagId, runId: dagRunId }]), + affectsMultipleRuns + ? [useGridServiceGetGridTiSummariesKey, { dagId }] + : UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: dagRunId }), ]; - await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); + await Promise.all(invalidateKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); if (onSuccess) { onSuccess(); From f184b08ada1c6a5b7a8f4e7efb4fef31ca4d707c Mon Sep 17 00:00:00 2001 From: Gwak Beomgyu <67548026+nothingmin@users.noreply.github.com> Date: Fri, 12 Sep 2025 17:31:49 +0900 Subject: [PATCH 2/4] made the code simpler Co-authored-by: Brent Bovenzi --- .../src/airflow/ui/src/queries/useClearTaskInstances.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts b/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts index 0b81aed63d9c6..998ae3bf6fa4f 100644 --- a/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts +++ b/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts @@ -90,9 +90,7 @@ export const useClearTaskInstances = ({ [useClearTaskInstancesDryRunKey, dagId], [usePatchTaskInstanceDryRunKey, dagId, dagRunId], UseGridServiceGetGridRunsKeyFn({ dagId }, [{ dagId }]), - affectsMultipleRuns - ? [useGridServiceGetGridTiSummariesKey, { dagId }] - : UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: dagRunId }), + UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: affectsMultipleRuns ? undefined : dagRunId }), ]; await Promise.all(invalidateKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); From c44e6c2e7275cbe2a9b8fad0984ad711187f9eef Mon Sep 17 00:00:00 2001 From: nothinmin Date: Fri, 12 Sep 2025 23:34:01 +0900 Subject: [PATCH 3/4] Restore variable name --- .../src/airflow/ui/src/queries/useClearTaskInstances.ts | 4 ++-- .../src/airflow/ui/src/queries/usePatchTaskInstance.ts | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts b/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts index 998ae3bf6fa4f..1dfcfdf94ffdb 100644 --- a/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts +++ b/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts @@ -83,7 +83,7 @@ export const useClearTaskInstances = ({ const { include_future: includeFuture, include_past: includePast } = variables.requestBody; const affectsMultipleRuns = includeFuture === true || includePast === true; - const invalidateKeys = [ + const queryKeys = [ ...taskInstanceKeys, UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }), [useDagRunServiceGetDagRunsKey], @@ -93,7 +93,7 @@ export const useClearTaskInstances = ({ UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: affectsMultipleRuns ? undefined : dagRunId }), ]; - await Promise.all(invalidateKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); + await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); onSuccessConfirm(); }; diff --git a/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts b/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts index 44e8dcac4f327..c1ae88cacb166 100644 --- a/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts +++ b/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts @@ -72,7 +72,7 @@ export const usePatchTaskInstance = ({ const { include_future: includeFuture, include_past: includePast } = variables.requestBody; const affectsMultipleRuns = includeFuture === true || includePast === true; - const invalidateKeys = [ + const queryKeys = [ UseTaskInstanceServiceGetTaskInstanceKeyFn({ dagId, dagRunId, taskId }), UseTaskInstanceServiceGetMappedTaskInstanceKeyFn({ dagId, dagRunId, mapIndex, taskId }), [useTaskInstanceServiceGetTaskInstancesKey], @@ -84,7 +84,7 @@ export const usePatchTaskInstance = ({ : UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: dagRunId }), ]; - await Promise.all(invalidateKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); + await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); if (onSuccess) { onSuccess(); From 303326ac34be2dc0e5fc434743e1816fb887f372 Mon Sep 17 00:00:00 2001 From: nothinmin Date: Sat, 13 Sep 2025 00:08:01 +0900 Subject: [PATCH 4/4] fix typescript error --- .../src/airflow/ui/src/queries/useClearTaskInstances.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts b/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts index 1dfcfdf94ffdb..b20dae1075f07 100644 --- a/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts +++ b/airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.ts @@ -90,7 +90,9 @@ export const useClearTaskInstances = ({ [useClearTaskInstancesDryRunKey, dagId], [usePatchTaskInstanceDryRunKey, dagId, dagRunId], UseGridServiceGetGridRunsKeyFn({ dagId }, [{ dagId }]), - UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: affectsMultipleRuns ? undefined : dagRunId }), + affectsMultipleRuns + ? [useGridServiceGetGridTiSummariesKey, { dagId }] + : UseGridServiceGetGridTiSummariesKeyFn({ dagId, runId: dagRunId }), ]; await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key })));