UI: Bulk mark Dag runs as success/failed from multi-select

airflow-core/src/airflow/ui/src/pages/DagRuns/BulkMarkDagRunsAsButton.tsx

@@ -0,0 +1,130 @@
+/*!
+ * 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 { Badge, Box, Button, Flex, Heading, HStack, VStack, useDisclosure } from "@chakra-ui/react";
+import { useState } from "react";
+import { useTranslation } from "react-i18next";
+import { FiX } from "react-icons/fi";
+import { LuCheck } from "react-icons/lu";
+
+import type { DagRunMutableStates, DAGRunResponse } from "openapi/requests/types.gen";
+import { ActionAccordion } from "src/components/ActionAccordion";
+import { ActionErrors } from "src/components/ActionErrors";
+import { allowedStates } from "src/components/MarkAs/utils";
+import { StateBadge } from "src/components/StateBadge";
+import { Dialog, Menu } from "src/components/ui";
+import { useBulkPatchDagRun } from "src/queries/useBulkPatchDagRun";
+
+type Props = {
+  readonly deselectKeys: (keys: Array<string>) => void;
+  readonly selectedDagRuns: Array<DAGRunResponse>;
+};
+
+const BulkMarkDagRunsAsButton = ({ deselectKeys, selectedDagRuns }: Props) => {
+  const { t: translate } = useTranslation(["common", "dags"]);
+  const { onClose, onOpen, open } = useDisclosure();
+  const [state, setState] = useState<DagRunMutableStates>("success");
+  const [note, setNote] = useState<string | null>(null);
+  const { bulkAction, data, error, isPending, reset } = useBulkPatchDagRun({
+    deselectKeys,
+    onSuccessConfirm: onClose,
+  });
+
+  const handleOpen = (newState: DagRunMutableStates) => {
+    setState(newState);
+    setNote(null);
+    reset();
+    onOpen();
+  };
+
+  return (
+    <Box>
+      <Menu.Root positioning={{ gutter: 0, placement: "top" }}>
+        <Menu.Trigger asChild>
+          <Button variant="outline">
+            <HStack gap={1} mx={1}>
+              <LuCheck />
+              <span>/</span>
+              <FiX />
+            </HStack>
+            {translate("dags:runAndTaskActions.markAs.button", { type: translate("dagRun_other") })}
+          </Button>
+        </Menu.Trigger>
+        <Menu.Content>
+          {allowedStates.map((menuState) => (
+            <Menu.Item key={menuState} onClick={() => handleOpen(menuState)} value={menuState}>
+              <HStack justify="space-between" width="full">
+                <StateBadge state={menuState}>{translate(`common:states.${menuState}`)}</StateBadge>
+                <Badge colorPalette="gray" variant="subtle">
+                  {selectedDagRuns.length}
+                </Badge>
+              </HStack>
+            </Menu.Item>
+          ))}
+        </Menu.Content>
+      </Menu.Root>
+
+      <Dialog.Root onOpenChange={onClose} open={open}>
+        <Dialog.Content backdrop>
+          <Dialog.Header>
+            <VStack align="start" gap={4}>
+              <Heading size="xl">
+                {translate("dags:runAndTaskActions.markAs.title", {
+                  state,
+                  type: translate("dagRun_other"),
+                })}{" "}
+                <StateBadge state={state} />
+              </Heading>
+            </VStack>
+          </Dialog.Header>
+
+          <Dialog.CloseTrigger />
+          <Dialog.Body width="full">
+            <ActionAccordion note={note} setNote={setNote} />
+            <ActionErrors actionResponse={data?.update} error={error} />
+            <Flex justifyContent="end" mt={3}>
+              <Button
+                loading={isPending}
+                onClick={() => {
+                  bulkAction({
+                    actions: [
+                      {
+                        action: "update" as const,
+                        action_on_non_existence: "skip",
+                        entities: selectedDagRuns.map((dagRun) => ({
+                          dag_id: dagRun.dag_id,
+                          dag_run_id: dagRun.dag_run_id,
+                          note,
+                          state,
+                        })),
+                      },
+                    ],
+                  });
+                }}
+              >
+                {translate("modal.confirm")}
+              </Button>
+            </Flex>
+          </Dialog.Body>
+        </Dialog.Content>
+      </Dialog.Root>
+    </Box>
+  );
+};
+
+export default BulkMarkDagRunsAsButton;

airflow-core/src/airflow/ui/src/pages/DagRuns/DagRuns.tsx

@@ -51,6 +51,7 @@ import { renderDuration, useAutoRefresh, isStatePending } from "src/utils";
 
 import BulkClearDagRunsButton from "./BulkClearDagRunsButton";
 import BulkDeleteDagRunsButton from "./BulkDeleteDagRunsButton";
+import BulkMarkDagRunsAsButton from "./BulkMarkDagRunsAsButton";
 import { DagRunsFilters } from "./DagRunsFilters";
 import DeleteRunButton from "./DeleteRunButton";
 
@@ -360,6 +361,7 @@ export const DagRuns = () => {
           </ActionBar.SelectionTrigger>
           <ActionBar.Separator />
           <BulkClearDagRunsButton deselectKeys={deselectKeys} selectedDagRuns={selectedDagRuns} />
+          <BulkMarkDagRunsAsButton deselectKeys={deselectKeys} selectedDagRuns={selectedDagRuns} />
           <BulkDeleteDagRunsButton deselectKeys={deselectKeys} selectedDagRuns={selectedDagRuns} />
           <ActionBar.CloseTrigger onClick={clearSelections} />
         </ActionBar.Content>

airflow-core/src/airflow/ui/src/queries/useBulkPatchDagRun.ts

@@ -0,0 +1,115 @@
+/*!
+ * 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 { useQueryClient } from "@tanstack/react-query";
+import { useRef } from "react";
+import { useTranslation } from "react-i18next";
+
+import {
+  useDagRunServiceBulkDagRuns,
+  useDagRunServiceGetDagRunKey,
+  useDagRunServiceGetDagRunsKey,
+  useGanttServiceGetGanttDataKey,
+  useTaskInstanceServiceGetMappedTaskInstanceKey,
+  useTaskInstanceServiceGetTaskInstanceKey,
+  useTaskInstanceServiceGetTaskInstancesKey,
+} from "openapi/queries";
+import type { BulkBody_BulkDAGRunBody_, BulkResponse } from "openapi/requests/types.gen";
+import { toaster } from "src/components/ui";
+
+import { gridQueryKeys, tiPerAttemptQueryKeys } from "./gridViewQueryKeys";
+import { useClearDagRunDryRunKey } from "./useClearDagRunDryRun";
+
+type Props = {
+  readonly deselectKeys: (keys: Array<string>) => void;
+  readonly onSuccessConfirm: VoidFunction;
+};
+
+export const useBulkPatchDagRun = ({ deselectKeys, onSuccessConfirm }: Props) => {
+  const queryClient = useQueryClient();
+  const affectedDagIds = useRef<Set<string>>(new Set());
+  const { t: translate } = useTranslation(["common", "dags"]);
+
+  const onSuccess = async (responseData: BulkResponse) => {
+    // Marking runs changes run + task instance states, so invalidate the same surfaces the
+    // single-run mark (usePatchDagRun) does: run/TI lists, run + single/mapped TI details, gantt,
+    // grid, and the clear dry-run preview.
+    await Promise.all([
+      queryClient.invalidateQueries({ queryKey: [useDagRunServiceGetDagRunsKey] }),
+      queryClient.invalidateQueries({ queryKey: [useDagRunServiceGetDagRunKey] }),
+      queryClient.invalidateQueries({ queryKey: [useTaskInstanceServiceGetTaskInstancesKey] }),
+      queryClient.invalidateQueries({ queryKey: [useTaskInstanceServiceGetTaskInstanceKey] }),
+      queryClient.invalidateQueries({ queryKey: [useTaskInstanceServiceGetMappedTaskInstanceKey] }),
+      queryClient.invalidateQueries({ queryKey: [useGanttServiceGetGanttDataKey] }),
+      ...tiPerAttemptQueryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key })),
+      ...[...affectedDagIds.current].flatMap((dagId) => [
+        ...gridQueryKeys(dagId).map((key) => queryClient.invalidateQueries({ queryKey: key })),
+        queryClient.invalidateQueries({ queryKey: [useClearDagRunDryRunKey, dagId] }),
+      ]),
+    ]);
+
+    const updateResult = responseData.update;
+
+    if (!updateResult) {
+      return;
+    }
+
+    const successKeys = updateResult.success ?? [];
+    const actionErrors = updateResult.errors ?? [];
+
+    if (successKeys.length > 0) {
+      toaster.create({
+        description: translate("toaster.bulkUpdate.success.description", {
+          count: successKeys.length,
+          keys: successKeys.join(", "),
+          resourceName: translate("dagRun_other"),
+        }),
+        title: translate("toaster.bulkUpdate.success.title", {
+          resourceName: translate("dagRun_other"),
+        }),
+        type: "success",
+      });
+      deselectKeys(successKeys);
+    }
+
+    // Per-entity failures (status 200 with items in ``errors``) keep the dialog open
+    // so the user can see what failed; the consumer renders ``data.update.errors``.
+    if (actionErrors.length === 0) {
+      onSuccessConfirm();
+    }
+  };
+
+  const { data, error, isPending, mutate, reset } = useDagRunServiceBulkDagRuns({ onSuccess });
+
+  const bulkAction = (requestBody: BulkBody_BulkDAGRunBody_) => {
+    reset();
+    const dagIds = new Set<string>();
+
+    for (const action of requestBody.actions) {
+      for (const entity of action.entities) {
+        if (typeof entity !== "string" && entity.dag_id !== null && entity.dag_id !== undefined) {
+          dagIds.add(entity.dag_id);
+        }
+      }
+    }
+    affectedDagIds.current = dagIds;
+    mutate({ dagId: "~", requestBody });
+  };
+
+  return { bulkAction, data, error, isPending, reset };
+};