Add edit pool button (apache#46998)

Author: shubhamraj-git

airflow/ui/src/pages/Pools/EditPoolButton.tsx

@@ -0,0 +1,85 @@
+/*!
+ * 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 { Heading, useDisclosure } from "@chakra-ui/react";
+import { FiEdit } from "react-icons/fi";
+
+import type { PoolResponse } from "openapi/requests/types.gen";
+import { Dialog } from "src/components/ui";
+import ActionButton from "src/components/ui/ActionButton";
+import { useEditPool } from "src/queries/useEditPool";
+
+import PoolForm, { type PoolBody } from "./PoolForm";
+
+type Props = {
+  readonly pool: PoolResponse;
+};
+
+const EditPoolButton = ({ pool }: Props) => {
+  const { onClose, onOpen, open } = useDisclosure();
+  const initialPoolValue: PoolBody = {
+    description: pool.description ?? "",
+    include_deferred: pool.include_deferred,
+    name: pool.name,
+    slots: pool.slots,
+  };
+  const { editPool, error, isPending, setError } = useEditPool(initialPoolValue, {
+    onSuccessConfirm: onClose,
+  });
+
+  const handleClose = () => {
+    setError(undefined);
+    onClose();
+  };
+
+  return (
+    <>
+      <ActionButton
+        actionName="Edit Pool"
+        icon={<FiEdit />}
+        onClick={() => {
+          onOpen();
+        }}
+        text="Edit Pool"
+        withText={false}
+      />
+
+      <Dialog.Root onOpenChange={handleClose} open={open} size="xl">
+        <Dialog.Content backdrop>
+          <Dialog.Header>
+            <Heading size="xl">Edit Pool</Heading>
+          </Dialog.Header>
+
+          <Dialog.CloseTrigger />
+
+          <Dialog.Body>
+            <PoolForm
+              error={error}
+              initialPool={initialPoolValue}
+              isPending={isPending}
+              manageMutate={editPool}
+              setError={setError}
+            />
+          </Dialog.Body>
+        </Dialog.Content>
+      </Dialog.Root>
+    </>
+  );
+};
+
+export default EditPoolButton;

airflow/ui/src/pages/Pools/PoolBar.tsx

@@ -27,6 +27,7 @@ import { Tooltip } from "src/components/ui";
 import { capitalize } from "src/utils";
 
 import DeletePoolButton from "./DeletePoolButton";
+import EditPoolButton from "./EditPoolButton";
 
 const slots = {
   open_slots: { color: "success", icon: <StateIcon color="white" state="success" /> },
@@ -54,7 +55,10 @@ const PoolBar = ({ pool }: PoolBarProps) => (
               </Tooltip>
             ) : undefined}
           </Text>
-          {pool.name === "default_pool" ? undefined : <DeletePoolButton poolName={pool.name} />}
+          <HStack gap={0}>
+            <EditPoolButton pool={pool} />
+            {pool.name === "default_pool" ? undefined : <DeletePoolButton poolName={pool.name} />}
+          </HStack>
         </HStack>
         {pool.description ?? (
           <Text color="gray.fg" fontSize="sm">

airflow/ui/src/pages/Pools/PoolForm.tsx

@@ -96,7 +96,7 @@ const PoolForm = ({ error, initialPool, isPending, manageMutate, setError }: Poo
         render={({ field }) => (
           <Field.Root mb={4} mt={4}>
             <Field.Label fontSize="md">Description</Field.Label>
-            <Textarea {...field} size="sm" />
+            <Textarea {...field} disabled={initialPool.name === "default_pool"} size="sm" />
           </Field.Root>
         )}
       />

airflow/ui/src/queries/useEditPool.ts

@@ -0,0 +1,89 @@
+/*!
+ * 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 { useState } from "react";
+
+import { usePoolServiceGetPoolsKey, usePoolServicePatchPool } from "openapi/queries";
+import type { PoolBody } from "openapi/requests/types.gen";
+import { toaster } from "src/components/ui";
+
+export const useEditPool = (
+  initialPool: PoolBody,
+  {
+    onSuccessConfirm,
+  }: {
+    onSuccessConfirm: () => void;
+  },
+) => {
+  const queryClient = useQueryClient();
+  const [error, setError] = useState<unknown>(undefined);
+
+  const onSuccess = async () => {
+    await queryClient.invalidateQueries({
+      queryKey: [usePoolServiceGetPoolsKey],
+    });
+
+    toaster.create({
+      description: "Pool has been edited successfully",
+      title: "Pool Edit Request Submitted",
+      type: "success",
+    });
+
+    onSuccessConfirm();
+  };
+
+  const onError = (_error: unknown) => {
+    setError(_error);
+  };
+
+  const { isPending, mutate } = usePoolServicePatchPool({
+    onError,
+    onSuccess,
+  });
+
+  const editPool = (editPoolRequestBody: PoolBody) => {
+    const updateMask: Array<string> = [];
+
+    if (editPoolRequestBody.slots !== initialPool.slots) {
+      updateMask.push("slots");
+    }
+    if (editPoolRequestBody.description !== initialPool.description) {
+      updateMask.push("description");
+    }
+    if (editPoolRequestBody.include_deferred !== initialPool.include_deferred) {
+      updateMask.push("include_deferred");
+    }
+
+    const parsedDescription =
+      editPoolRequestBody.description === "" ? undefined : editPoolRequestBody.description;
+
+    mutate({
+      poolName: initialPool.name,
+      requestBody: {
+        description: parsedDescription,
+        include_deferred: editPoolRequestBody.include_deferred,
+        pool: editPoolRequestBody.name,
+        slots: editPoolRequestBody.slots,
+      },
+      updateMask,
+    });
+  };
+
+  return { editPool, error, isPending, setError };
+};