(airflow) 17/22: Fix Grid cache invalidation for multi-run task operations (#55504)

[kaxilnaik]

This is an automated email from the ASF dual-hosted git repository.

kaxilnaik pushed a commit to branch v3-1-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit dd8690c4feafd45c66e4d170eae2ba48530693df
Author: Gwak Beomgyu <67548026+nothing...@users.noreply.github.com>
AuthorDate: Sat Sep 13 02:08:10 2025 +0900

    Fix Grid cache invalidation for multi-run task operations (#55504)
    
    * Fix Grid cache invalidation for multi-run task operations
    
    * made the code simpler
    
    Co-authored-by: Brent Bovenzi <brent.bove...@gmail.com>
    
    * Restore variable name
    
    * fix typescript error
    
    ---------
    
    Co-authored-by: Brent Bovenzi <brent.bove...@gmail.com>
    (cherry picked from commit 07dd277985fe64ba32633603d7427fc3a2fa5ac1)
---
 .../airflow/ui/src/queries/useClearTaskInstances.ts   |  9 ++++++++-
 .../airflow/ui/src/queries/usePatchTaskInstance.ts    | 19 +++++++++++++++++--
 2 files changed, 25 insertions(+), 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 ed3431bba08..b20dae1075f 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,6 +79,10 @@ export const useClearTaskInstances = ({
       ),
     ];
 
+    // 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 queryKeys = [
       ...taskInstanceKeys,
       UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }),
@@ -85,7 +90,9 @@ export const useClearTaskInstances = ({
       [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 })));
diff --git a/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts 
b/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts
index bc2747b1a57..c1ae88cacb1 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,7 +59,19 @@ export const usePatchTaskInstance = ({
     });
   };
 
-  const onSuccessFn = async () => {
+  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 queryKeys = [
       UseTaskInstanceServiceGetTaskInstanceKeyFn({ dagId, dagRunId, taskId }),
       UseTaskInstanceServiceGetMappedTaskInstanceKeyFn({ dagId, dagRunId, 
mapIndex, taskId }),
@@ -66,7 +79,9 @@ export const usePatchTaskInstance = ({
       [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 })));