Fix bug in query invalidation and remove custom predicate logic (#46485)

* Fix bug in query invalidation and remove custom predicate logic * Address PR feedback * Only use key fn when we have props to pass
apache · Feb 6, 2025 · fd6cbed · fd6cbed
1 parent aaaea35
commit fd6cbed
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 31 deletions.
diff --git a/airflow/ui/src/queries/useTrigger.ts b/airflow/ui/src/queries/useTrigger.ts
@@ -20,15 +20,14 @@ import { useQueryClient } from "@tanstack/react-query";
 import { useState } from "react";
 
 import {
-  useDagRunServiceGetDagRunsKey,
+  UseDagRunServiceGetDagRunsKeyFn,
   useDagRunServiceTriggerDagRun,
   useDagServiceGetDagsKey,
   useDagsServiceRecentDagRunsKey,
-  useTaskInstanceServiceGetTaskInstancesKey,
+  UseTaskInstanceServiceGetTaskInstancesKeyFn,
 } from "openapi/queries";
 import type { DagRunTriggerParams } from "src/components/TriggerDag/TriggerDAGForm";
 import { toaster } from "src/components/ui";
-import { doQueryKeysMatch, type PartialQueryKey } from "src/utils";
 
 export const useTrigger = ({ dagId, onSuccessConfirm }: { dagId: string; onSuccessConfirm: () => void }) => {
   const queryClient = useQueryClient();
@@ -37,14 +36,14 @@ export const useTrigger = ({ dagId, onSuccessConfirm }: { dagId: string; onSucce
   const [dateValidationError, setDateValidationError] = useState<unknown>(undefined);
 
   const onSuccess = async () => {
-    const queryKeys: Array<PartialQueryKey> = [
-      { baseKey: useDagServiceGetDagsKey },
-      { baseKey: useDagsServiceRecentDagRunsKey },
-      { baseKey: useDagRunServiceGetDagRunsKey, options: { dagIds: [dagId] } },
-      { baseKey: useTaskInstanceServiceGetTaskInstancesKey, options: { dagId, dagRunId: "~" } },
+    const queryKeys = [
+      [useDagServiceGetDagsKey],
+      [useDagsServiceRecentDagRunsKey],
+      UseDagRunServiceGetDagRunsKeyFn({ dagId }, [{ dagId }]),
+      UseTaskInstanceServiceGetTaskInstancesKeyFn({ dagId, dagRunId: "~" }, [{ dagId, dagRunId: "~" }]),
     ];
 
-    await queryClient.invalidateQueries({ predicate: (query) => doQueryKeysMatch(query, queryKeys) });
+    await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key })));
 
     toaster.create({
       description: "DAG run has been successfully triggered.",

diff --git a/airflow/ui/src/utils/query.ts b/airflow/ui/src/utils/query.ts
@@ -16,8 +16,6 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import type { Query } from "@tanstack/react-query";
-
 import { useDagServiceGetDagDetails } from "openapi/queries";
 import type { TaskInstanceState } from "openapi/requests/types.gen";
 import { useConfig } from "src/queries/useConfig";
@@ -32,26 +30,6 @@ export const isStatePending = (state?: TaskInstanceState | null) =>
   state === "restarting" ||
   !Boolean(state);
 
-export type PartialQueryKey = { baseKey: string; options?: Record<string, unknown> };
-
-// This allows us to specify what query key values we actually care about and ignore the rest
-// ex: match everything with this dagId and dagRunId but ignore anything related to pagination
-export const doQueryKeysMatch = (query: Query, queryKeysToMatch: Array<PartialQueryKey>) => {
-  const [baseKey, options] = query.queryKey;
-
-  const matchedKey = queryKeysToMatch.find((qk) => qk.baseKey === baseKey);
-
-  if (!matchedKey) {
-    return false;
-  }
-
-  return matchedKey.options
-    ? Object.entries(matchedKey.options).every(
-        ([key, value]) => typeof options === "object" && (options as Record<string, unknown>)[key] === value,
-      )
-    : true;
-};
-
 export const useAutoRefresh = ({ dagId, isPaused }: { dagId?: string; isPaused?: boolean }) => {
   const autoRefreshInterval = useConfig("auto_refresh_interval") as number | undefined;
   const { data: dag } = useDagServiceGetDagDetails(