From 421b3443e2213f461636c474826b5116d543f938 Mon Sep 17 00:00:00 2001 From: Pierre Jeambrun Date: Thu, 2 Jul 2026 11:38:26 +0200 Subject: [PATCH] Make Dag pause toggle update immediately on click (#69134) * Make Dag pause toggle update immediately on click The pause/unpause switch on the Dags list page waited for the dags list refetch to settle before flipping, which under fast interaction looked like the first clicks were dropped: clicking once on several rows left them all visually unchanged until enough subsequent clicks accumulated and the list eventually refetched. Apply React Query's optimistic update pattern: write the new is_paused into the dag detail and dags list caches on mutate, roll back if the server rejects, and invalidate the list on settle so filtered views (e.g. paused=true) drop the row that no longer matches the filter. * Test the optimistic pause toggle against a sibling filtered Dag list The test seeded a single dags-list query and "verified" sibling invalidation with isFetching(...) >= 0, which is always true and asserted nothing. Seed a second list under a different filter and assert that both the prefix-keyed optimistic write and the on-settle invalidation reach it, so the test proves the cache update is not limited to the one seeded list variant. (cherry picked from commit 2033e1a57c7fe00fbf4c3f36021c0612e71efd2d) --- .../ui/src/queries/useTogglePause.test.tsx | 183 ++++++++++++++++++ .../airflow/ui/src/queries/useTogglePause.ts | 104 ++++++++-- 2 files changed, 276 insertions(+), 11 deletions(-) create mode 100644 airflow-core/src/airflow/ui/src/queries/useTogglePause.test.tsx diff --git a/airflow-core/src/airflow/ui/src/queries/useTogglePause.test.tsx b/airflow-core/src/airflow/ui/src/queries/useTogglePause.test.tsx new file mode 100644 index 0000000000000..47fa20849bf9c --- /dev/null +++ b/airflow-core/src/airflow/ui/src/queries/useTogglePause.test.tsx @@ -0,0 +1,183 @@ +/*! + * 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 { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import { http, HttpResponse } from "msw"; +import { setupServer } from "msw/node"; +import React from "react"; +import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; + +import { UseDagServiceGetDagsUiKeyFn } from "openapi/queries"; +import type { DAGWithLatestDagRunsCollectionResponse } from "openapi/requests/types.gen"; +import { useTogglePause } from "src/queries/useTogglePause"; + +const DAG_ID = "dag_under_test"; + +const buildDagsList = (isPaused: boolean): DAGWithLatestDagRunsCollectionResponse => ({ + dags: [ + { + asset_expression: null, + bundle_name: null, + bundle_version: null, + dag_display_name: DAG_ID, + dag_id: DAG_ID, + description: null, + file_token: "", + fileloc: "/dags/dag.py", + has_import_errors: false, + has_task_concurrency_limits: false, + is_favorite: false, + is_paused: isPaused, + is_stale: false, + last_expired: null, + last_parsed_time: null, + latest_dag_runs: [], + max_active_runs: 16, + max_active_tasks: 16, + max_consecutive_failed_dag_runs: 0, + next_dagrun_create_after: null, + next_dagrun_data_interval_end: null, + next_dagrun_data_interval_start: null, + next_dagrun_logical_date: null, + next_dagrun_run_after: null, + owners: ["airflow"], + pending_actions: 0, + tags: [], + timetable_description: null, + timetable_partitioned: false, + timetable_summary: null, + } as unknown as DAGWithLatestDagRunsCollectionResponse["dags"][number], + ], + total_entries: 1, +}); + +const server = setupServer(); + +const createWrapper = + (queryClient: QueryClient) => + ({ children }: { readonly children: React.ReactNode }) => ( + {children} + ); + +const seedClient = (initialIsPaused: boolean) => { + const queryClient = new QueryClient({ + defaultOptions: { + mutations: { retry: false }, + queries: { gcTime: Infinity, retry: false, staleTime: Infinity }, + }, + }); + const dagsListKey = UseDagServiceGetDagsUiKeyFn({ dagRunsLimit: 1 }); + // A sibling list query with a different filter (paused=false), so tests can verify the + // prefix-keyed optimistic write and invalidation reach every dags-list cache, not just one. + const filteredListKey = UseDagServiceGetDagsUiKeyFn({ dagRunsLimit: 1, paused: false }); + + queryClient.setQueryData( + dagsListKey, + buildDagsList(initialIsPaused), + ); + queryClient.setQueryData( + filteredListKey, + buildDagsList(initialIsPaused), + ); + + return { dagsListKey, filteredListKey, queryClient }; +}; + +beforeAll(() => server.listen({ onUnhandledRequest: "error" })); +afterEach(() => server.resetHandlers()); +afterAll(() => server.close()); + +describe("useTogglePause", () => { + it("optimistically flips is_paused in the dags list cache as soon as mutate() is called", async () => { + server.use( + http.patch("*/api/v2/dags/:dagId", async () => { + // Make the server slow so the test can observe the optimistic state + // before the response is settled. + await new Promise((resolve) => { + setTimeout(resolve, 50); + }); + + return HttpResponse.json({ dag_id: DAG_ID, is_paused: true }); + }), + ); + + const { dagsListKey, filteredListKey, queryClient } = seedClient(false); + const { result } = renderHook(() => useTogglePause({ dagId: DAG_ID }), { + wrapper: createWrapper(queryClient), + }); + + result.current.mutate({ dagId: DAG_ID, requestBody: { is_paused: true } }); + + // The Switch is bound to is_paused — the optimistic update must apply before the network + // response settles so the user sees the flip on click. It is keyed by the list prefix, so + // every dags-list cache flips, including a sibling query with a different filter. + await waitFor(() => { + expect( + queryClient.getQueryData(dagsListKey)?.dags[0]?.is_paused, + ).toBe(true); + expect( + queryClient.getQueryData(filteredListKey)?.dags[0]?.is_paused, + ).toBe(true); + }); + + await waitFor(() => expect(result.current.isPending).toBe(false)); + // Post-settle the cache still reflects the new value. + const settled = queryClient.getQueryData(dagsListKey); + + expect(settled?.dags[0]?.is_paused).toBe(true); + }); + + it("rolls back the optimistic update when the server rejects the change", async () => { + server.use(http.patch("*/api/v2/dags/:dagId", () => new HttpResponse(null, { status: 500 }))); + + const { dagsListKey, queryClient } = seedClient(false); + const { result } = renderHook(() => useTogglePause({ dagId: DAG_ID }), { + wrapper: createWrapper(queryClient), + }); + + result.current.mutate({ dagId: DAG_ID, requestBody: { is_paused: true } }); + + await waitFor(() => expect(result.current.isError).toBe(true)); + + const cached = queryClient.getQueryData(dagsListKey); + + expect(cached?.dags[0]?.is_paused).toBe(false); + }); + + it("invalidates the dags list query on settle so filtered lists refetch", async () => { + server.use( + http.patch("*/api/v2/dags/:dagId", () => HttpResponse.json({ dag_id: DAG_ID, is_paused: true })), + ); + + const { dagsListKey, filteredListKey, queryClient } = seedClient(false); + const { result } = renderHook(() => useTogglePause({ dagId: DAG_ID }), { + wrapper: createWrapper(queryClient), + }); + + result.current.mutate({ dagId: DAG_ID, requestBody: { is_paused: true } }); + + await waitFor(() => expect(result.current.isPending).toBe(false)); + + // Both the seeded list and a sibling list with a different filter (paused=false) are marked + // stale, because invalidation is keyed by the shared prefix — so a filtered consumer refetches + // and may drop the dag from the visible page. + expect(queryClient.getQueryState(dagsListKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(filteredListKey)?.isInvalidated).toBe(true); + }); +}); diff --git a/airflow-core/src/airflow/ui/src/queries/useTogglePause.ts b/airflow-core/src/airflow/ui/src/queries/useTogglePause.ts index 04a7c510acbd1..9ed05420f2cab 100644 --- a/airflow-core/src/airflow/ui/src/queries/useTogglePause.ts +++ b/airflow-core/src/airflow/ui/src/queries/useTogglePause.ts @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -import { useQueryClient } from "@tanstack/react-query"; +import { useQueryClient, type QueryKey } from "@tanstack/react-query"; import { useTranslation } from "react-i18next"; import { @@ -27,17 +27,102 @@ import { useDagServiceGetDagsUiKey, UseTaskInstanceServiceGetTaskInstancesKeyFn, } from "openapi/queries"; +import type { + DAGDetailsResponse, + DAGPatchBody, + DAGResponse, + DAGWithLatestDagRunsCollectionResponse, +} from "openapi/requests/types.gen"; import { createErrorToaster } from "src/utils"; +type TogglePauseVariables = { dagId: string; requestBody: DAGPatchBody }; + +type TogglePauseContext = { + previousDag: DAGResponse | undefined; + previousDagDetails: DAGDetailsResponse | undefined; + previousDagsLists: Array<[QueryKey, DAGWithLatestDagRunsCollectionResponse | undefined]>; +}; + export const useTogglePause = ({ dagId }: { dagId: string }) => { const queryClient = useQueryClient(); const { t: translate } = useTranslation("common"); - const onSuccess = async () => { - const queryKeys = [ - [useDagServiceGetDagsUiKey], - UseDagServiceGetDagKeyFn({ dagId }, [{ dagId }]), - UseDagServiceGetDagDetailsKeyFn({ dagId }, [{ dagId }]), + const dagKey = UseDagServiceGetDagKeyFn({ dagId }, [{ dagId }]); + const dagDetailsKey = UseDagServiceGetDagDetailsKeyFn({ dagId }, [{ dagId }]); + const dagsListPrefix: QueryKey = [useDagServiceGetDagsUiKey]; + + const onMutate = async ({ requestBody }: TogglePauseVariables): Promise => { + const nextIsPaused = requestBody.is_paused; + + // Cancel in-flight refetches so they cannot overwrite the optimistic update. + await Promise.all([ + queryClient.cancelQueries({ queryKey: dagKey }), + queryClient.cancelQueries({ queryKey: dagDetailsKey }), + queryClient.cancelQueries({ queryKey: dagsListPrefix }), + ]); + + const previousDag = queryClient.getQueryData(dagKey); + const previousDagDetails = queryClient.getQueryData(dagDetailsKey); + const previousDagsLists = queryClient.getQueriesData({ + queryKey: dagsListPrefix, + }); + + // Optimistically reflect the new is_paused value so the Switch flips + // immediately on click rather than waiting for the server round-trip. + if (previousDag !== undefined) { + queryClient.setQueryData(dagKey, { ...previousDag, is_paused: nextIsPaused }); + } + if (previousDagDetails !== undefined) { + queryClient.setQueryData(dagDetailsKey, { + ...previousDagDetails, + is_paused: nextIsPaused, + }); + } + queryClient.setQueriesData( + { queryKey: dagsListPrefix }, + (current) => + current === undefined + ? current + : { + ...current, + dags: current.dags.map((dag) => + dag.dag_id === dagId ? { ...dag, is_paused: nextIsPaused } : dag, + ), + }, + ); + + return { previousDag, previousDagDetails, previousDagsLists }; + }; + + const onError = ( + error: unknown, + _variables: TogglePauseVariables, + context: TogglePauseContext | undefined, + ) => { + // Roll back the optimistic update if the server rejected the change. + if (context !== undefined) { + if (context.previousDag !== undefined) { + queryClient.setQueryData(dagKey, context.previousDag); + } + if (context.previousDagDetails !== undefined) { + queryClient.setQueryData(dagDetailsKey, context.previousDagDetails); + } + context.previousDagsLists.forEach(([key, data]) => { + queryClient.setQueryData(key, data); + }); + } + + createErrorToaster(error, { titleKey: "common:error.title" }, translate); + }; + + const onSettled = async () => { + // Invalidate after the mutation settles (success or error) so filtered list + // queries (e.g. paused=true/false) refetch and may move the dag in or out + // of the visible page. + const queryKeys: Array = [ + dagsListPrefix, + dagKey, + dagDetailsKey, UseDagRunServiceGetDagRunsKeyFn({ dagId }, [{ dagId }]), UseTaskInstanceServiceGetTaskInstancesKeyFn({ dagId, dagRunId: "~" }, [{ dagId, dagRunId: "~" }]), ]; @@ -45,12 +130,9 @@ export const useTogglePause = ({ dagId }: { dagId: string }) => { await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); }; - const onError = (error: unknown) => { - createErrorToaster(error, { titleKey: "common:error.title" }, translate); - }; - return useDagServicePatchDag({ onError, - onSuccess, + onMutate, + onSettled, }); };