diff --git a/components/webui/server/src/routes/api/presto-search/index.ts b/components/webui/server/src/routes/api/presto-search/index.ts index 1f714835fb..ad0dd3f361 100644 --- a/components/webui/server/src/routes/api/presto-search/index.ts +++ b/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((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" ); }); }, diff --git a/components/webui/server/src/routes/api/presto-search/typings.ts b/components/webui/server/src/routes/api/presto-search/typings.ts new file mode 100644 index 0000000000..bdf72a8b87 --- /dev/null +++ b/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;