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

Author: hoophalab

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<null>> => {
+    console.log("Cancelling query:", JSON.stringify(payload));
+
+    return axios.post("/api/presto-search/cancel", payload);
+};
+
+
 export type {
     PrestoQueryJobCreationSchema,
     PrestoQueryJobSchema,
 };
 
-export {submitQuery};
+export {
+    cancelQuery,
+    submitQuery,
+};

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

@@ -0,0 +1,3 @@
+.cancelButton {
+    width: 100%;
+}

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

@@ -0,0 +1,47 @@
+import {useCallback} from "react";
+
+import {CloseOutlined} from "@ant-design/icons";
+import {
+    Button,
+    Tooltip,
+} from "antd";
+
+import useSearchStore from "../../../../SearchState/index";
+import {handlePrestoQueryCancel} from "../../presto-search-requests";
+import styles from "./index.module.css";
+
+
+/**
+ * Renders a button to cancel the SQL query.
+ *
+ * @return
+ */
+const CancelButton = () => {
+    const searchJobId = useSearchStore((state) => state.searchJobId);
+
+    const handleClick = useCallback(() => {
+        if (null === searchJobId) {
+            console.error("Cannot cancel query: searchJobId is not set.");
+
+            return;
+        }
+        handlePrestoQueryCancel({searchJobId});
+    }, [searchJobId]);
+
+    return (
+        <Tooltip title={"Cancel query"}>
+            <Button
+                className={styles["cancelButton"] || ""}
+                color={"red"}
+                icon={<CloseOutlined/>}
+                size={"large"}
+                variant={"solid"}
+                onClick={handleClick}
+            >
+                Cancel
+            </Button>
+        </Tooltip>
+    );
+};
+
+export default CancelButton;

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

@@ -0,0 +1,3 @@
+.runButton {
+    width: 100%;
+}

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

@@ -6,8 +6,10 @@ import {
     Tooltip,
 } from "antd";
 
-import useSearchStore from "../../../SearchState/index";
-import {handlePrestoQuerySubmit} from "../presto-search-requests";
+import useSearchStore from "../../../../SearchState/index";
+import {SEARCH_UI_STATE} from "../../../../SearchState/typings";
+import {handlePrestoQuerySubmit} from "../../presto-search-requests";
+import styles from "./index.module.css";
 
 
 /**
@@ -16,6 +18,7 @@ import {handlePrestoQuerySubmit} from "../presto-search-requests";
  * @return
  */
 const RunButton = () => {
+    const searchUiState = useSearchStore((state) => state.searchUiState);
     const queryString = useSearchStore((state) => state.queryString);
 
     const isQueryStringEmpty = "" === queryString;
@@ -30,11 +33,13 @@ const RunButton = () => {
     return (
         <Tooltip title={tooltipTitle}>
             <Button
+                className={styles["runButton"] || ""}
                 color={"green"}
-                disabled={isQueryStringEmpty}
                 icon={<CaretRightOutlined/>}
                 size={"large"}
                 variant={"solid"}
+                disabled={isQueryStringEmpty ||
+                    searchUiState === SEARCH_UI_STATE.QUERY_ID_PENDING}
                 onClick={handleClick}
             >
                 Run

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

@@ -0,0 +1,3 @@
+.runButtonContainer {
+    width: 130px;
+}

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

@@ -0,0 +1,25 @@
+import useSearchStore from "../../../SearchState/index";
+import {SEARCH_UI_STATE} from "../../../SearchState/typings";
+import CancelButton from "./CancelButton";
+import styles from "./index.module.css";
+import RunButton from "./RunButton";
+
+
+/**
+ * Renders a button to submit or cancel the SQL query.
+ *
+ * @return
+ */
+const SqlSearchButton = () => {
+    const searchUiState = useSearchStore((state) => state.searchUiState);
+
+    return (
+        <div className={styles["runButtonContainer"] || ""}>
+            { (searchUiState === SEARCH_UI_STATE.QUERYING) ?
+                <CancelButton/> :
+                <RunButton/>}
+        </div>
+    );
+};
+
+export default SqlSearchButton;

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,9 +14,30 @@ 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);
+
+            // eslint-disable-next-line no-warning-comments
+            // TODO: Delete previous query results when the backend is ready
+
             console.debug(
                 "Presto search job created - ",
                 "Search job ID:",
@@ -24,4 +49,29 @@ const handlePrestoQuerySubmit = (payload: PrestoQueryJobCreationSchema) => {
         });
 };
 
-export {handlePrestoQuerySubmit};
+/**
+ * Cancels an ongoing Presto search query on server.
+ *
+ * @param payload
+ */
+const handlePrestoQueryCancel = (payload: PrestoQueryJobSchema) => {
+    const {searchUiState, updateSearchUiState} = useSearchStore.getState();
+    if (searchUiState !== SEARCH_UI_STATE.QUERYING) {
+        console.error("Cannot cancel query if there is no ongoing query.");
+
+        return;
+    }
+
+    updateSearchUiState(SEARCH_UI_STATE.DONE);
+    cancelQuery(payload)
+        .then(() => {
+            console.debug("Query cancelled successfully");
+        })
+        .catch((err: unknown) => {
+            console.error("Failed to cancel query:", err);
+        });
+};
+
+export {
+    handlePrestoQueryCancel, handlePrestoQuerySubmit,
+};

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

@@ -6,8 +6,8 @@ import {
 } from "../../../config";
 import Dataset from "./Dataset";
 import styles from "./index.module.css";
-import RunButton from "./Presto/RunButton";
 import SqlQueryInput from "./Presto/SqlQueryInput";
+import SqlSearchButton from "./Presto/SqlSearchButton";
 import QueryInput from "./QueryInput";
 import SearchButton from "./SearchButton";
 import TimeRangeInput from "./TimeRangeInput";
@@ -43,7 +43,7 @@ const SearchControls = () => {
                     (
                         <>
                             <SqlQueryInput/>
-                            <RunButton/>
+                            <SqlSearchButton/>
                         </>
                     )}
             </div>

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

@@ -1,4 +1,7 @@
-import {FastifyPluginAsyncTypebox} from "@fastify/type-provider-typebox";
+import {
+    FastifyPluginAsyncTypebox,
+    Type,
+} from "@fastify/type-provider-typebox";
 import {StatusCodes} from "http-status-codes";
 
 import {
@@ -148,6 +151,35 @@ const plugin: FastifyPluginAsyncTypebox = async (fastify) => {
             return {searchJobId};
         }
     );
+
+    fastify.post(
+        "/cancel",
+        {
+            schema: {
+                body: PrestoQueryJobSchema,
+                response: {
+                    [StatusCodes.NO_CONTENT]: Type.Null(),
+                    [StatusCodes.INTERNAL_SERVER_ERROR]: ErrorSchema,
+                },
+                tags: ["Presto Search"],
+            },
+        },
+        async (request, reply) => {
+            const {searchJobId} = request.body;
+            await new Promise<void>((resolve, reject) => {
+                Presto.client.kill(searchJobId, (error) => {
+                    if (null !== error) {
+                        reject(new Error("Failed to kill the Presto query job.", {cause: error}));
+                    }
+                    resolve();
+                });
+            });
+            request.log.info(searchJobId, "Presto search cancelled");
+            reply.code(StatusCodes.NO_CONTENT);
+
+            return null;
+        }
+    );
 };
 
 export default plugin;