feat(webui): Store only 1000 search results in the results cache for Presto queries. (#1209)

Author: davemarco

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

@@ -15,6 +15,7 @@ import {
     PrestoQueryJobCreationSchema,
     PrestoQueryJobSchema,
 } from "../../../schemas/presto-search.js";
+import {MAX_PRESTO_SEARCH_RESULTS} from "./typings.js";
 import {insertPrestoRowsToMongo} from "./utils.js";
 
 
@@ -61,18 +62,20 @@ const plugin: FastifyPluginAsyncTypebox = async (fastify) => {
             const {queryString} = request.body;
 
             let searchJobId: string;
+            let totalResultsCount = 0;
+            let storedResultsCount = 0;
 
             try {
                 // eslint-disable-next-line max-lines-per-function
                 searchJobId = await new Promise<string>((resolve, reject) => {
                     let isResolved = false;
                     Presto.client.execute({
-                        // eslint-disable-next-line no-warning-comments
-                        // TODO: Error, and success handlers are dummy implementations
-                        // and will be replaced with proper implementations.
                         data: (_, data, columns) => {
+                            totalResultsCount += data.length;
+
                             request.log.info(
-                                `Received ${data.length} rows from Presto query`
+                                `Received ${data.length} rows from Presto query ` +
+                                `(total: ${totalResultsCount})`
                             );
 
                             if (false === isResolved) {
@@ -88,15 +91,35 @@ const plugin: FastifyPluginAsyncTypebox = async (fastify) => {
                                 return;
                             }
 
-                            insertPrestoRowsToMongo(
-                                data,
-                                columns,
-                                searchJobId,
-                                mongoDb
+                            if (storedResultsCount < MAX_PRESTO_SEARCH_RESULTS) {
+                                const remainingSlots =
+                                    MAX_PRESTO_SEARCH_RESULTS - storedResultsCount;
+                                const dataToInsert = data.slice(0, remainingSlots);
+
+                                if (0 < dataToInsert.length) {
+                                    storedResultsCount += dataToInsert.length;
+                                    insertPrestoRowsToMongo(
+                                        dataToInsert,
+                                        columns,
+                                        searchJobId,
+                                        mongoDb
+                                    ).catch((err: unknown) => {
+                                        request.log.error(
+                                            err,
+                                            "Failed to insert Presto results into MongoDB"
+                                        );
+                                    });
+                                }
+                            }
+
+                            // Always update metadata with total count
+                            searchResultsMetadataCollection.updateOne(
+                                {_id: searchJobId},
+                                {$set: {numTotalResults: totalResultsCount}}
                             ).catch((err: unknown) => {
                                 request.log.error(
                                     err,
-                                    "Failed to insert Presto results into MongoDB"
+                                    "Failed to update total results count in metadata"
                                 );
                             });
                         },

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

@@ -0,0 +1,4 @@
+/**
+ * Maximum number of Presto search results to store in MongoDB.
+ */
+export const MAX_PRESTO_SEARCH_RESULTS = 1000;