Add the bulk support for maint and exit maint mode for edge worker

providers/edge3/src/airflow/providers/edge3/plugins/www/src/components/BulkWorkerOperations.tsx

@@ -16,10 +16,13 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Button, Dialog, HStack, List, Portal, Text, useDisclosure } from "@chakra-ui/react";
+import { Button, Dialog, HStack, List, Portal, Text, Textarea, useDisclosure } from "@chakra-ui/react";
 import type { Worker } from "openapi/requests/types.gen";
+import { useState } from "react";
 import { FaPowerOff } from "react-icons/fa";
 import { FaRegTrashCan } from "react-icons/fa6";
+import { HiOutlineWrenchScrewdriver } from "react-icons/hi2";
+import { IoMdExit } from "react-icons/io";
 
 import { useBulkWorkerActions } from "src/hooks/useBulkWorkerActions";
 
@@ -44,12 +47,29 @@ export const BulkWorkerOperations = ({
     onOpen: onOpenDeleteDialog,
     open: isDeleteDialogOpen,
   } = useDisclosure();
+  const {
+    onClose: onCloseMaintenanceEnterDialog,
+    onOpen: onOpenMaintenanceEnterDialog,
+    open: isMaintenanceEnterDialogOpen,
+  } = useDisclosure();
+  const {
+    onClose: onCloseMaintenanceExitDialog,
+    onOpen: onOpenMaintenanceExitDialog,
+    open: isMaintenanceExitDialogOpen,
+  } = useDisclosure();
+  const [maintenanceComment, setMaintenanceComment] = useState("");
   const {
     deleteWorkers,
     handleBulkDelete,
+    handleBulkMaintenanceEnter,
+    handleBulkMaintenanceExit,
     handleBulkShutdown,
     isBulkDeletePending,
+    isBulkMaintenanceEnterPending,
+    isBulkMaintenanceExitPending,
     isBulkShutdownPending,
+    maintenanceEnterWorkers,
+    maintenanceExitWorkers,
     shutdownWorkers,
   } = useBulkWorkerActions({
     onClearSelection,
@@ -58,24 +78,53 @@ export const BulkWorkerOperations = ({
   });
 
   const onBulkShutdown = async () => {
-    try {
-      await handleBulkShutdown();
-    } finally {
-      onCloseShutdownDialog();
-    }
+    await handleBulkShutdown();
+    onCloseShutdownDialog();
   };
 
   const onBulkDelete = async () => {
-    try {
-      await handleBulkDelete();
-    } finally {
-      onCloseDeleteDialog();
-    }
+    await handleBulkDelete();
+    onCloseDeleteDialog();
+  };
+
+  const onMaintenanceEnterDialogClose = () => {
+    onCloseMaintenanceEnterDialog();
+    setMaintenanceComment("");
+  };
+
+  const onBulkMaintenanceEnter = async () => {
+    await handleBulkMaintenanceEnter(maintenanceComment);
+    onMaintenanceEnterDialogClose();
+  };
+
+  const onBulkMaintenanceExit = async () => {
+    await handleBulkMaintenanceExit();
+    onCloseMaintenanceExitDialog();
   };
 
   return (
     <>
       <HStack>
+        <Button
+          colorPalette="warning"
+          disabled={maintenanceEnterWorkers.length === 0}
+          onClick={onOpenMaintenanceEnterDialog}
+          size="sm"
+          variant="outline"
+        >
+          <HiOutlineWrenchScrewdriver />
+          Send to Maintenance ({maintenanceEnterWorkers.length})
+        </Button>
+        <Button
+          colorPalette="warning"
+          disabled={maintenanceExitWorkers.length === 0}
+          onClick={onOpenMaintenanceExitDialog}
+          size="sm"
+          variant="outline"
+        >
+          <IoMdExit />
+          Exit Maintenance ({maintenanceExitWorkers.length})
+        </Button>
         <Button
           colorPalette="danger"
           disabled={shutdownWorkers.length === 0}
@@ -98,6 +147,94 @@ export const BulkWorkerOperations = ({
         </Button>
       </HStack>
 
+      <Dialog.Root onOpenChange={onMaintenanceEnterDialogClose} open={isMaintenanceEnterDialogOpen} size="lg">
+        <Portal>
+          <Dialog.Backdrop />
+          <Dialog.Positioner>
+            <Dialog.Content>
+              <Dialog.Header>
+                <Dialog.Title>
+                  Send {maintenanceEnterWorkers.length} selected worker(s) to maintenance
+                </Dialog.Title>
+              </Dialog.Header>
+              <Dialog.Body>
+                <Text mb={3}>
+                  Maintenance mode can be requested only for workers in states: idle or running.
+                </Text>
+                <List.Root ps={5} mb={3}>
+                  {maintenanceEnterWorkers.map((worker) => (
+                    <List.Item key={worker.worker_name}>{worker.worker_name}</List.Item>
+                  ))}
+                </List.Root>
+                <Textarea
+                  placeholder="Enter maintenance comment (required)"
+                  value={maintenanceComment}
+                  onChange={(e) => setMaintenanceComment(e.target.value)}
+                  required
+                  maxLength={1024}
+                  size="sm"
+                />
+              </Dialog.Body>
+              <Dialog.Footer>
+                <Dialog.ActionTrigger asChild>
+                  <Button variant="outline">Cancel</Button>
+                </Dialog.ActionTrigger>
+                <Button
+                  colorPalette="warning"
+                  disabled={!maintenanceComment.trim()}
+                  loading={isBulkMaintenanceEnterPending}
+                  loadingText="Sending to maintenance..."
+                  onClick={onBulkMaintenanceEnter}
+                >
+                  <HiOutlineWrenchScrewdriver />
+                  Send to Maintenance
+                </Button>
+              </Dialog.Footer>
+            </Dialog.Content>
+          </Dialog.Positioner>
+        </Portal>
+      </Dialog.Root>
+
+      <Dialog.Root onOpenChange={onCloseMaintenanceExitDialog} open={isMaintenanceExitDialogOpen} size="lg">
+        <Portal>
+          <Dialog.Backdrop />
+          <Dialog.Positioner>
+            <Dialog.Content>
+              <Dialog.Header>
+                <Dialog.Title>
+                  Exit maintenance for {maintenanceExitWorkers.length} selected worker(s)
+                </Dialog.Title>
+              </Dialog.Header>
+              <Dialog.Body>
+                <Text mb={3}>
+                  Exit maintenance is available only for workers in states: maintenance pending, maintenance
+                  mode, maintenance request, or offline maintenance.
+                </Text>
+                <List.Root ps={5}>
+                  {maintenanceExitWorkers.map((worker) => (
+                    <List.Item key={worker.worker_name}>{worker.worker_name}</List.Item>
+                  ))}
+                </List.Root>
+              </Dialog.Body>
+              <Dialog.Footer>
+                <Dialog.ActionTrigger asChild>
+                  <Button variant="outline">Cancel</Button>
+                </Dialog.ActionTrigger>
+                <Button
+                  colorPalette="warning"
+                  loading={isBulkMaintenanceExitPending}
+                  loadingText="Exiting maintenance..."
+                  onClick={onBulkMaintenanceExit}
+                >
+                  <IoMdExit />
+                  Exit Maintenance
+                </Button>
+              </Dialog.Footer>
+            </Dialog.Content>
+          </Dialog.Positioner>
+        </Portal>
+      </Dialog.Root>
+
       <Dialog.Root onOpenChange={onCloseShutdownDialog} open={isShutdownDialogOpen} size="lg">
         <Portal>
           <Dialog.Backdrop />

providers/edge3/src/airflow/providers/edge3/plugins/www/src/constants.ts

@@ -69,4 +69,16 @@ export const bulkWorkerDeleteEligibleStates = new Set<EdgeWorkerState>([
   "offline maintenance",
 ]);
 
+export const bulkWorkerMaintenanceEnterEligibleStates = new Set<EdgeWorkerState>([
+  "idle",
+  "running",
+]);
+
+export const bulkWorkerMaintenanceExitEligibleStates = new Set<EdgeWorkerState>([
+  "maintenance pending",
+  "maintenance mode",
+  "maintenance request",
+  "offline maintenance",
+]);
+
 export const bulkWorkerActionBatchSize = 10;

providers/edge3/src/airflow/providers/edge3/plugins/www/src/hooks/useBulkWorkerActions.ts

@@ -16,14 +16,21 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { useUiServiceDeleteWorker, useUiServiceRequestWorkerShutdown } from "openapi/queries";
+import {
+  useUiServiceDeleteWorker,
+  useUiServiceExitWorkerMaintenance,
+  useUiServiceRequestWorkerMaintenance,
+  useUiServiceRequestWorkerShutdown,
+} from "openapi/queries";
 import type { Worker } from "openapi/requests/types.gen";
-import { useMemo, useState } from "react";
+import { useCallback, useMemo, useState } from "react";
 
 import { toaster } from "src/components/ui";
 import {
   bulkWorkerActionBatchSize,
   bulkWorkerDeleteEligibleStates,
+  bulkWorkerMaintenanceEnterEligibleStates,
+  bulkWorkerMaintenanceExitEligibleStates,
   bulkWorkerShutdownEligibleStates,
 } from "src/constants";
 
@@ -99,9 +106,13 @@ export const useBulkWorkerActions = ({
 }: UseBulkWorkerActionsProps) => {
   const [isBulkDeletePending, setIsBulkDeletePending] = useState(false);
   const [isBulkShutdownPending, setIsBulkShutdownPending] = useState(false);
+  const [isBulkMaintenanceEnterPending, setIsBulkMaintenanceEnterPending] = useState(false);
+  const [isBulkMaintenanceExitPending, setIsBulkMaintenanceExitPending] = useState(false);
 
   const shutdownMutation = useUiServiceRequestWorkerShutdown();
   const deleteMutation = useUiServiceDeleteWorker();
+  const maintenanceEnterMutation = useUiServiceRequestWorkerMaintenance();
+  const maintenanceExitMutation = useUiServiceExitWorkerMaintenance();
 
   const shutdownWorkers = useMemo(
     () => selectedWorkers.filter((worker) => bulkWorkerShutdownEligibleStates.has(worker.state)),
@@ -111,8 +122,16 @@ export const useBulkWorkerActions = ({
     () => selectedWorkers.filter((worker) => bulkWorkerDeleteEligibleStates.has(worker.state)),
     [selectedWorkers],
   );
+  const maintenanceEnterWorkers = useMemo(
+    () => selectedWorkers.filter((worker) => bulkWorkerMaintenanceEnterEligibleStates.has(worker.state)),
+    [selectedWorkers],
+  );
+  const maintenanceExitWorkers = useMemo(
+    () => selectedWorkers.filter((worker) => bulkWorkerMaintenanceExitEligibleStates.has(worker.state)),
+    [selectedWorkers],
+  );
 
-  const handleBulkAction = async ({
+  const handleBulkAction = useCallback(async ({
     failureTitle,
     setPending,
     successToast,
@@ -144,9 +163,9 @@ export const useBulkWorkerActions = ({
     } finally {
       setPending(false);
     }
-  };
+  }, [onClearSelection, onOperations]);
 
-  const handleBulkShutdown = async (): Promise<void> => {
+  const handleBulkShutdown = useCallback(async (): Promise<void> => {
     await handleBulkAction({
       failureTitle: "Bulk Shutdown Partially Failed",
       setPending: setIsBulkShutdownPending,
@@ -158,9 +177,9 @@ export const useBulkWorkerActions = ({
       workers: shutdownWorkers,
       workerMutation: (worker) => shutdownMutation.mutateAsync({ workerName: worker.worker_name }),
     });
-  };
+  }, [handleBulkAction, shutdownMutation, shutdownWorkers]);
 
-  const handleBulkDelete = async (): Promise<void> => {
+  const handleBulkDelete = useCallback(async (): Promise<void> => {
     await handleBulkAction({
       failureTitle: "Bulk Delete Partially Failed",
       setPending: setIsBulkDeletePending,
@@ -172,14 +191,52 @@ export const useBulkWorkerActions = ({
       workers: deleteWorkers,
       workerMutation: (worker) => deleteMutation.mutateAsync({ workerName: worker.worker_name }),
     });
-  };
+  }, [deleteWorkers, deleteMutation, handleBulkAction]);
+
+  const handleBulkMaintenanceEnter = useCallback(async (comment: string): Promise<void> => {
+    await handleBulkAction({
+      failureTitle: "Bulk Maintenance Enter Partially Failed",
+      setPending: setIsBulkMaintenanceEnterPending,
+      successToast: (successCount) => ({
+        description: `Maintenance mode requested for ${successCount} worker(s).`,
+        title: "Bulk Maintenance Enter Requested",
+        type: "success",
+      }),
+      workers: maintenanceEnterWorkers,
+      workerMutation: (worker) =>
+        maintenanceEnterMutation.mutateAsync({
+          requestBody: { maintenance_comment: comment },
+          workerName: worker.worker_name,
+        }),
+    });
+  }, [handleBulkAction, maintenanceEnterMutation, maintenanceEnterWorkers]);
+
+  const handleBulkMaintenanceExit = useCallback(async (): Promise<void> => {
+    await handleBulkAction({
+      failureTitle: "Bulk Maintenance Exit Partially Failed",
+      setPending: setIsBulkMaintenanceExitPending,
+      successToast: (successCount) => ({
+        description: `${successCount} worker(s) requested to exit maintenance mode.`,
+        title: "Bulk Maintenance Exit Requested",
+        type: "success",
+      }),
+      workers: maintenanceExitWorkers,
+      workerMutation: (worker) => maintenanceExitMutation.mutateAsync({ workerName: worker.worker_name }),
+    });
+  }, [handleBulkAction, maintenanceExitMutation, maintenanceExitWorkers]);
 
   return {
     deleteWorkers,
     handleBulkDelete,
+    handleBulkMaintenanceEnter,
+    handleBulkMaintenanceExit,
     handleBulkShutdown,
     isBulkDeletePending,
+    isBulkMaintenanceEnterPending,
+    isBulkMaintenanceExitPending,
     isBulkShutdownPending,
+    maintenanceEnterWorkers,
+    maintenanceExitWorkers,
     shutdownWorkers,
   };
 };