fix(webui): Improve Presto search metadata handling:

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

@@ -5,7 +5,7 @@ import {
     type PrestoQueryJobSchema,
     submitQuery,
 } from "../../../../api/presto-search";
-import useSearchStore from "../../SearchState";
+import useSearchStore, {SEARCH_STATE_DEFAULT} from "../../SearchState";
 import {SEARCH_UI_STATE} from "../../SearchState/typings";
 
 
@@ -40,7 +40,13 @@ const handlePrestoClearResults = () => {
  * @param payload
  */
 const handlePrestoQuerySubmit = (payload: PrestoQueryJobCreationSchema) => {
-    const {updateSearchJobId, updateSearchUiState, searchUiState} = useSearchStore.getState();
+    const {
+        updateNumSearchResultsTable,
+        updateNumSearchResultsMetadata,
+        updateSearchJobId,
+        updateSearchUiState,
+        searchUiState,
+    } = useSearchStore.getState();
 
     // User should NOT be able to submit a new query while an existing query is in progress.
     if (
@@ -55,6 +61,8 @@ const handlePrestoQuerySubmit = (payload: PrestoQueryJobCreationSchema) => {
 
     handlePrestoClearResults();
 
+    updateNumSearchResultsTable(SEARCH_STATE_DEFAULT.numSearchResultsTable);
+    updateNumSearchResultsMetadata(SEARCH_STATE_DEFAULT.numSearchResultsMetadata);
     updateSearchUiState(SEARCH_UI_STATE.QUERY_ID_PENDING);
 
     submitQuery(payload)

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

@@ -32,7 +32,7 @@ const handleSubmit = (ev: React.FormEvent<HTMLFormElement>) => {
 const SearchControls = () => {
     /* eslint-disable-next-line no-warning-comments */
     // TODO: Remove flag and related logic when the new guide UI is fully implemented.
-    const isGuidedEnabled = "true" === import.meta.env[`VITE_GUIDED_DEV`];
+    const isGuidedEnabled = "true" === import.meta.env["VITE_GUIDED_DEV"];
 
     return (
         <form onSubmit={handleSubmit}>

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

@@ -90,6 +90,7 @@ const handleQuerySubmit = (payload: QueryJobCreationSchema) => {
 
     store.updateNumSearchResultsTable(SEARCH_STATE_DEFAULT.numSearchResultsTable);
     store.updateNumSearchResultsTimeline(SEARCH_STATE_DEFAULT.numSearchResultsTimeline);
+    store.updateNumSearchResultsMetadata(SEARCH_STATE_DEFAULT.numSearchResultsMetadata);
     store.updateSearchUiState(SEARCH_UI_STATE.QUERY_ID_PENDING);
 
     submitQuery(payload)

components/webui/client/src/pages/SearchPage/SearchState/useUpdateStateWithMetadata.ts

@@ -38,7 +38,7 @@ const useUpdateStateWithMetadata = () => {
 
         switch (resultsMetadata.lastSignal) {
             case SEARCH_SIGNAL.RESP_DONE:
-            case PRESTO_SEARCH_SIGNAL.FINISHED:
+            case PRESTO_SEARCH_SIGNAL.DONE:
                 updateSearchUiState(SEARCH_UI_STATE.DONE);
                 break;
             case PRESTO_SEARCH_SIGNAL.FAILED:

components/webui/common/src/index.ts

@@ -107,6 +107,10 @@ enum PRESTO_SEARCH_SIGNAL {
     FINISHED = "FINISHED",
     CANCELED = "CANCELED",
     FAILED = "FAILED",
+
+    // Used internally by the UI to mark when all data has been received, since Presto may report
+    // `FINISHED` before the result stream is fully delivered. `DONE` state is never set by Presto.
+    DONE = "DONE",
 }
 
 /**

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

@@ -170,17 +170,23 @@ const plugin: FastifyPluginAsyncTypebox = async (fastify) => {
                                 });
                                 isResolved = true;
                                 resolve(queryId);
-                            } else {
-                                // Update metadata on subsequent calls
-                                searchResultsMetadataCollection.updateOne(
-                                    {_id: queryId},
-                                    {$set: {lastSignal: newState}}
-                                ).catch((err: unknown) => {
-                                    request.log.error(err, "Failed to update Presto metadata");
-                                });
                             }
                         },
                         success: () => {
+                            if (false === isResolved) {
+                                request.log.error(
+                                    "Presto query finished before searchJobId was resolved; "
+                                );
+
+                                return;
+                            }
+                            searchResultsMetadataCollection.updateOne(
+                                {_id: searchJobId},
+                                {$set: {lastSignal: PRESTO_SEARCH_SIGNAL.DONE}}
+                            ).catch((err: unknown) => {
+                                request.log.error(err, "Failed to update Presto metadata");
+                            });
+
                             request.log.info("Presto search succeeded");
                         },
                         timeout: null,