feat(webui): Add support for cancelling Presto SQL queries.

components/webui/client/src/api/presto-search/index.ts

@@ -28,9 +28,28 @@ const submitQuery = async (
     return axios.post<PrestoQueryJobSchema>("/api/presto-search/query", payload);
 };
 
+
+/**
+ * Sends post request to server to cancel presto query.
+ *
+ * @param payload
+ * @return
+ */
+const cancelQuery = async (
+    payload: PrestoQueryJobSchema
+): Promise<AxiosResponse<PrestoQueryJobSchema>> => {
+    console.log("Cancelling query:", JSON.stringify(payload));
+
+    return axios.post<PrestoQueryJobSchema>("/api/presto-search/cancel", payload);
+};
+
+
 export type {
     PrestoQueryJobCreationSchema,
     PrestoQueryJobSchema,
 };
 
-export {submitQuery};
+export {
+    cancelQuery,
+    submitQuery,
+};

components/webui/client/src/pages/SearchPage/SearchControls/Presto/RunButton/index.module.css

@@ -0,0 +1,11 @@
+.runButtonContainer {
+    width: 120px;
+}
+
+.cancelButton {
+    width: 100%;
+}
+
+.runButton {
+    width: 100%;
+}

components/webui/client/src/pages/SearchPage/SearchControls/Presto/RunButton/index.tsx

@@ -1,13 +1,21 @@
 import {useCallback} from "react";
 
-import {CaretRightOutlined} from "@ant-design/icons";
+import {
+    CaretRightOutlined,
+    CloseOutlined,
+} from "@ant-design/icons";
 import {
     Button,
     Tooltip,
 } from "antd";
 
 import useSearchStore from "../../../SearchState/index";
-import {handlePrestoQuerySubmit} from "../presto-search-requests";
+import {SEARCH_UI_STATE} from "../../../SearchState/typings";
+import {
+    handlePrestoQueryCancel,
+    handlePrestoQuerySubmit,
+} from "../presto-search-requests";
+import styles from "./index.module.css";
 
 
 /**
@@ -16,7 +24,9 @@ import {handlePrestoQuerySubmit} from "../presto-search-requests";
  * @return
  */
 const RunButton = () => {
+    const searchUiState = useSearchStore((state) => state.searchUiState);
     const queryString = useSearchStore((state) => state.queryString);
+    const searchJobId = useSearchStore((state) => state.searchJobId);
 
     const isQueryStringEmpty = "" === queryString;
     const tooltipTitle = isQueryStringEmpty ?
@@ -27,19 +37,47 @@ const RunButton = () => {
         handlePrestoQuerySubmit({queryString});
     }, [queryString]);
 
+    const handleCancel = useCallback(() => {
+        if (null === searchJobId) {
+            console.error("Cannot cancel query: searchJobId is not set.");
+
+            return;
+        }
+        handlePrestoQueryCancel({searchJobId});
+    }, [searchJobId]);
+
     return (
-        <Tooltip title={tooltipTitle}>
-            <Button
-                color={"green"}
-                disabled={isQueryStringEmpty}
-                icon={<CaretRightOutlined/>}
-                size={"large"}
-                variant={"solid"}
-                onClick={handleClick}
-            >
-                Run
-            </Button>
-        </Tooltip>
+        <div className={styles["runButtonContainer"] || ""}>
+            { (searchUiState === SEARCH_UI_STATE.QUERYING) ?
+
+                <Tooltip title={"Cancel query"}>
+                    <Button
+                        className={styles["cancelButton"] || ""}
+                        color={"red"}
+                        icon={<CloseOutlined/>}
+                        size={"large"}
+                        variant={"solid"}
+                        onClick={handleCancel}
+                    >
+                        Cancel
+                    </Button>
+                </Tooltip> :
+
+                <Tooltip title={tooltipTitle}>
+                    <Button
+                        className={styles["runButton"] || ""}
+                        color={"green"}
+                        icon={<CaretRightOutlined/>}
+                        size={"large"}
+                        variant={"solid"}
+                        disabled={isQueryStringEmpty ||
+                            searchUiState === SEARCH_UI_STATE.QUERY_ID_PENDING}
+                        onClick={handleClick}
+                    >
+                        Run
+                    </Button>
+                </Tooltip>}
+        </div>
     );
 };
 

components/webui/client/src/pages/SearchPage/SearchControls/Presto/presto-search-requests.ts

@@ -1,7 +1,11 @@
 import {
+    cancelQuery,
     type PrestoQueryJobCreationSchema,
+    type PrestoQueryJobSchema,
     submitQuery,
 } from "../../../../api/presto-search";
+import useSearchStore from "../../SearchState";
+import {SEARCH_UI_STATE} from "../../SearchState/typings";
 
 
 /**
@@ -10,18 +14,69 @@ import {
  * @param payload
  */
 const handlePrestoQuerySubmit = (payload: PrestoQueryJobCreationSchema) => {
+    const {updateSearchJobId, updateSearchUiState, searchUiState} = useSearchStore.getState();
+
+    // User should NOT be able to submit a new query while an existing query is in progress.
+    if (
+        searchUiState !== SEARCH_UI_STATE.DEFAULT &&
+        searchUiState !== SEARCH_UI_STATE.DONE
+    ) {
+        console.error("Cannot submit query while existing query is in progress.");
+
+        return;
+    }
+
+    updateSearchUiState(SEARCH_UI_STATE.QUERY_ID_PENDING);
+
     submitQuery(payload)
         .then((result) => {
             const {searchJobId} = result.data;
+
+            updateSearchJobId(searchJobId);
+            updateSearchUiState(SEARCH_UI_STATE.QUERYING);
+
             console.debug(
                 "Presto search job created - ",
                 "Search job ID:",
                 searchJobId
             );
+
+            // eslint-disable-next-line no-warning-comments
+            // TODO: Remove this timeout after we can show the search results in the table
+            // and set ui state to DONE.
+            setTimeout(() => {
+                updateSearchUiState(SEARCH_UI_STATE.DONE);
+            // eslint-disable-next-line no-magic-numbers
+            }, 5000);
         })
         .catch((err: unknown) => {
             console.error("Failed to submit query:", err);
         });
 };
 
-export {handlePrestoQuerySubmit};
+/**
+ * Cancels a new Presto query to server.
+ *
+ * @param payload
+ */
+const handlePrestoQueryCancel = (payload: PrestoQueryJobSchema) => {
+    const {searchUiState} = useSearchStore.getState();
+    if (searchUiState !== SEARCH_UI_STATE.QUERYING) {
+        console.error("Cannot cancel query if there is no ongoing query.");
+
+        return;
+    }
+
+    cancelQuery(payload)
+        .then(() => {
+            const {updateSearchUiState} = useSearchStore.getState();
+            updateSearchUiState(SEARCH_UI_STATE.DONE);
+        })
+        .catch((err: unknown) => {
+            console.error("Failed to cancel query:", err);
+        });
+};
+
+export {
+    handlePrestoQueryCancel, handlePrestoQuerySubmit,
+};

components/webui/server/src/routes/api/presto-search/index.ts

@@ -1,5 +1,10 @@
-import {FastifyPluginAsyncTypebox} from "@fastify/type-provider-typebox";
+import {
+    FastifyPluginAsyncTypebox,
+    Type,
+} from "@fastify/type-provider-typebox";
 import {StatusCodes} from "http-status-codes";
+import {PrestoRequestError} from "presto-client";
+import {Nullable} from "src/typings/common.js";
 
 import {ErrorSchema} from "../../../schemas/error.js";
 import {
@@ -120,6 +125,39 @@ const plugin: FastifyPluginAsyncTypebox = async (fastify) => {
             return {searchJobId};
         }
     );
+
+    fastify.post(
+        "/cancel",
+        {
+            schema: {
+                body: PrestoQueryJobSchema,
+                response: {
+                    [StatusCodes.OK]: Type.Null(),
+                    [StatusCodes.INTERNAL_SERVER_ERROR]: ErrorSchema,
+                },
+                tags: ["Presto Search"],
+            },
+        },
+        async (request, reply) => {
+            const {searchJobId} = request.body;
+            const error : Nullable<PrestoRequestError> = await new Promise((resolve) => {
+                Presto.client.kill(searchJobId, (err) => {
+                    resolve(err);
+                });
+            });
+
+            if (null !== error) {
+                reply.code(StatusCodes.INTERNAL_SERVER_ERROR);
+                request.log.error(error, "Failed to cancel Presto query");
+
+                return {message: error.message};
+            }
+            reply.code(StatusCodes.OK);
+            request.log.info(searchJobId, "Presto search cancelled");
+
+            return null;
+        }
+    );
 };
 
 export default plugin;