From 1672d646193c2e783cade35e29657eaa54ada121 Mon Sep 17 00:00:00 2001 From: zipg Date: Tue, 28 Jul 2026 01:02:24 +0800 Subject: [PATCH] fix(mongodb): support query result pagination --- .../src/lib/mongo/mongoShellCommand.ts | 45 ++++ apps/desktop/src/stores/queryStore.ts | 44 +++- packages/app-tests/mongoShellCommand.test.ts | 35 +++ packages/app-tests/queryStore.test.ts | 246 ++++++++++++++++++ 4 files changed, 367 insertions(+), 3 deletions(-) diff --git a/apps/desktop/src/lib/mongo/mongoShellCommand.ts b/apps/desktop/src/lib/mongo/mongoShellCommand.ts index 6f4f4a9b5..5dedf6e70 100644 --- a/apps/desktop/src/lib/mongo/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongo/mongoShellCommand.ts @@ -26,6 +26,15 @@ export interface MongoFindCommand { sort?: string; } +export interface MongoFindPaginationPlan { + pageOffset: number; + pageLimit: number; + requestSkip: number; + requestLimit: number; + logicalSkip: number; + logicalLimit?: number; +} + export interface MongoFindOneCommand { collection: string; filter: string; @@ -167,6 +176,42 @@ export function parseMongoFindCommand(input: string): MongoFindCommand | null { }; } +export function planMongoFindPagination(input: string, command: MongoFindCommand, pageOffset: number, pageLimit: number): MongoFindPaginationPlan | null { + const source = input.trim().replace(/;$/, "").trim(); + const target = parseFindTarget(source); + if (!target) return null; + + const findOpenIndex = source.indexOf("(", target.findCallIndex); + const findCloseIndex = findMatchingParen(source, findOpenIndex); + if (findCloseIndex < 0) return null; + const chain = source.slice(findCloseIndex + 1).trim(); + if (chain && !chain.startsWith(".")) return null; + + const normalizedPageOffset = Math.max(0, Math.trunc(pageOffset)); + const normalizedPageLimit = Math.max(1, Math.trunc(pageLimit)); + const hasExplicitSkip = findChainedMethodCallIndex(chain, "skip") >= 0; + const hasExplicitLimit = findChainedMethodCallIndex(chain, "limit") >= 0; + const logicalSkip = hasExplicitSkip ? Math.max(0, Math.trunc(command.skip)) : 0; + // limit(0) is unbounded in MongoDB; a negative limit keeps the same row + // bound while requesting single-batch cursor semantics. + const logicalLimit = hasExplicitLimit && command.limit !== 0 ? Math.abs(Math.trunc(command.limit)) : undefined; + const remaining = logicalLimit === undefined ? normalizedPageLimit : Math.max(0, logicalLimit - normalizedPageOffset); + + return { + pageOffset: normalizedPageOffset, + pageLimit: normalizedPageLimit, + requestSkip: logicalSkip + normalizedPageOffset, + requestLimit: Math.min(normalizedPageLimit, remaining), + logicalSkip, + logicalLimit, + }; +} + +export function mongoFindLogicalTotal(total: number, plan: Pick): number { + const afterSkip = Math.max(0, Math.trunc(total) - plan.logicalSkip); + return plan.logicalLimit === undefined ? afterSkip : Math.min(afterSkip, plan.logicalLimit); +} + export function parseMongoFindOneCommand(input: string): MongoFindOneCommand | null { const source = input.trim().replace(/;$/, "").trim(); const target = parseCollectionMethodTarget(source, "findOne"); diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index d2e9bcccc..5c8cf298f 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -20,7 +20,9 @@ import { mongoDocumentsToQueryResult, describeMongoCommandParseFailure, mongoDroppedIndexesToQueryResult, + mongoFindLogicalTotal, mongoIndexesToQueryResult, + planMongoFindPagination, mongoUseToQueryResult, mongoVersionToQueryResult, mongoWriteToQueryResult, @@ -3219,6 +3221,7 @@ export const useQueryStore = defineStore("query", () => { // earlier `use ...` statements in the same editor selection. let currentDatabase = tab.database; let mongoEditTarget: QueryTab["mongoEditTarget"] | undefined; + let mongoFindPageState: { pageLimit: number; pageOffset: number; total: number; totalIsExact: boolean } | undefined; for (const parsedCommand of mongoCommands) { let mongoCommand = parsedCommand.command; @@ -3239,9 +3242,25 @@ export const useQueryStore = defineStore("query", () => { switch (mongoCommand.kind) { case "find": { queryExecutionLog("info", "mongo-find:start", { traceId, collection: mongoCommand.collection, database: currentDatabase }); - const result = await api.mongoFindDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, mongoCommand.skip, mongoCommand.limit, mongoCommand.filter, mongoCommand.projection, mongoCommand.sort, executionId); - const queryResult = markQueryResultRowsRaw(annotateMongoResult(mongoDocumentsToQueryResult(result.documents, performance.now() - commandStartedAt, result.total, result.extended_documents, result.total_is_exact !== false))); + const pagePlan = planMongoFindPagination(sourceStatement, mongoCommand, options?.pagination?.offset ?? 0, normalizeResultPageSize(options?.pagination?.limit ?? settingsStore.editorSettings.pageSize)); + if (!pagePlan) throw new Error(describeMongoCommandParseFailure(sourceStatement)); + // A stale request can point past an explicit .limit() bound. Keep + // the backend call bounded so limit(0) cannot become unbounded. + const result = await api.mongoFindDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, pagePlan.requestSkip, Math.max(1, pagePlan.requestLimit), mongoCommand.filter, mongoCommand.projection, mongoCommand.sort, executionId); + const documents = pagePlan.requestLimit === 0 ? [] : result.documents; + const extendedDocuments = pagePlan.requestLimit === 0 ? [] : result.extended_documents; + const totalIsExact = result.total_is_exact !== false; + const reportedTotal = mongoFindLogicalTotal(result.total, pagePlan); + const loadedLowerBound = pagePlan.pageOffset + documents.length; + const total = totalIsExact ? reportedTotal : Math.max(reportedTotal, loadedLowerBound); + const hasMore = totalIsExact ? loadedLowerBound < total : pagePlan.requestLimit > 0 && documents.length >= pagePlan.requestLimit && (pagePlan.logicalLimit === undefined || loadedLowerBound < pagePlan.logicalLimit); + const queryResult = markQueryResultRowsRaw(annotateMongoResult(mongoDocumentsToQueryResult(documents, performance.now() - commandStartedAt, total, extendedDocuments, totalIsExact))); + queryResult.truncated = hasMore; + queryResult.has_more = hasMore; allResults.push(queryResult); + if (mongoCommands.length === 1) { + mongoFindPageState = { pageLimit: pagePlan.pageLimit, pageOffset: pagePlan.pageOffset, total, totalIsExact }; + } mongoEditTarget = mongoCommands.length === 1 && !mongoCommand.projection && queryResult.columns.includes("_id") ? { collection: mongoCommand.collection, idColumn: "_id" } : undefined; queryExecutionLog("info", "mongo-find:done", { traceId, @@ -3472,8 +3491,20 @@ export const useQueryStore = defineStore("query", () => { if (openInNewResultTab && current.isCancelling && restorePendingResultRun(current, executionId)) return false; const activeGroupIndex = current.activeResultIndex; const activeGroupResults = current.results; + const findPageState = mongoFindPageState; + const shouldAppendResult = !!findPageState && !!options?.appendResult && !!current.result && allResults.length === 1; const shouldReplaceActiveResultInGroup = options?.replaceActiveResultInGroup === true && allResults.length === 1 && Array.isArray(activeGroupResults) && typeof activeGroupIndex === "number" && activeGroupIndex >= 0 && activeGroupIndex < activeGroupResults.length; - if (shouldReplaceActiveResultInGroup) { + if (shouldAppendResult) { + if (findPageState!.pageOffset !== current.result!.rows.length) { + throw new Error("Ignoring a stale MongoDB result segment whose offset no longer matches the loaded rows"); + } + const appendedResult = appendQueryResultSegment(current.result!, allResults[0]!, options!.appendResult!.maxRows); + if (Array.isArray(activeGroupResults) && typeof activeGroupIndex === "number" && activeGroupIndex >= 0 && activeGroupIndex < activeGroupResults.length) { + current.results = activeGroupResults.slice(); + current.results[activeGroupIndex] = appendedResult; + } + current.result = appendedResult; + } else if (shouldReplaceActiveResultInGroup) { current.results = activeGroupResults.slice(); current.results[activeGroupIndex] = allResults[0]; current.result = allResults[0]; @@ -3499,6 +3530,13 @@ export const useQueryStore = defineStore("query", () => { current.tableMeta = undefined; current.resultBaseSql = shouldReplaceActiveResultInGroup ? (current.resultBaseSql ?? options?.resultBaseSql ?? sql) : (options?.resultBaseSql ?? sql); current.resultSortedSql = options?.resultSortedSql; + current.resultPageSql = undefined; + current.resultPageLimit = mongoFindPageState?.pageLimit; + current.resultPageOffset = shouldAppendResult ? (current.resultPageOffset ?? 0) : mongoFindPageState?.pageOffset; + current.resultCountSql = undefined; + current.resultSessionId = undefined; + current.resultTotalRowCount = mongoFindPageState?.totalIsExact ? mongoFindPageState.total : undefined; + current.resultTotalRowCountLoading = false; syncDisplayedResultRun(current, current.resultBaseSql ?? options?.resultBaseSql ?? sql, openInNewResultTab); if (current.database !== currentDatabase) current.database = currentDatabase; } diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index 352d0cbe7..be0f2df2a 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -9,6 +9,7 @@ import { mongoCountToQueryResult, mongoDistinctToQueryResult, mongoDocumentsToQueryResult, + mongoFindLogicalTotal, mongoIndexesToQueryResult, normalizeRustMongoCommand, parseMongoAggregateCommand, @@ -23,6 +24,7 @@ import { parseMongoFindOneAndDeleteCommand, parseMongoGetIndexesCommand, parseMongoVersionCommand, + planMongoFindPagination, parseMongoWriteCommand, splitMongoCommands, splitMongoCommandRanges, @@ -56,6 +58,39 @@ test("parseMongoFindCommand parses getCollection find with chained sort skip and }); }); +test("planMongoFindPagination pages unbounded find queries", () => { + const command = parseMongoFindCommand("db.users.find({})"); + assert.ok(command); + const plan = planMongoFindPagination("db.users.find({})", command, 100, 100); + + assert.deepEqual(plan, { + pageOffset: 100, + pageLimit: 100, + requestSkip: 100, + requestLimit: 100, + logicalSkip: 0, + logicalLimit: undefined, + }); + assert.equal(mongoFindLogicalTotal(824, plan!), 824); +}); + +test("planMongoFindPagination preserves explicit skip and limit bounds", () => { + const source = "db.users.find({ active: true }).skip(20).limit(150)"; + const command = parseMongoFindCommand(source); + assert.ok(command); + const plan = planMongoFindPagination(source, command, 100, 100); + + assert.deepEqual(plan, { + pageOffset: 100, + pageLimit: 100, + requestSkip: 120, + requestLimit: 50, + logicalSkip: 20, + logicalLimit: 150, + }); + assert.equal(mongoFindLogicalTotal(824, plan!), 150); +}); + test("parseMongoFindCommand accepts line breaks before find and chained calls", () => { const command = parseMongoFindCommand(`db.getCollection("accounting_reconciliations") .find({ diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index 3bf551d85..2b42036ed 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -3260,6 +3260,252 @@ test("mongo aggregate execution uses editor page size when pagination plan has n } }); +test("mongo find execution uses editor page size and supports server pagination", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const findBodies: any[] = []; + + settingsStore.updateEditorSettings({ pageSize: 100 }); + connectionStore.addEphemeralConnection({ + ...conn("mongo-page-1"), + db_type: "mongodb", + port: 27017, + }); + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + if (String(input) === "/api/document-store/find-documents") { + const body = JSON.parse(String(init?.body ?? "{}")); + findBodies.push(body); + const available = Math.max(0, 824 - body.skip); + const rowCount = Math.min(body.limit, available); + return new Response( + JSON.stringify({ + documents: Array.from({ length: rowCount }, (_, index) => ({ _id: body.skip + index + 1 })), + total: 824, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const tabId = store.createTab("mongo-page-1", "dbx_test", "Query", "query", ""); + await store.executeTabSql(tabId, "db.issue_4566.find({})"); + const tab = store.tabs.find((item) => item.id === tabId); + + assert.equal(findBodies[0]?.collection, "issue_4566"); + assert.equal(findBodies[0]?.skip, 0); + assert.equal(findBodies[0]?.limit, 100); + assert.equal(tab?.result?.rows.length, 100); + assert.equal(tab?.resultPageLimit, 100); + assert.equal(tab?.resultPageOffset, 0); + assert.equal(tab?.resultTotalRowCount, 824); + + await store.executeTabSql(tabId, "db.issue_4566.find({})", { + pagination: { offset: 100, limit: 100 }, + preserveResultDuringExecution: true, + preserveTotalRowCountDuringExecution: true, + }); + + assert.equal(findBodies[1]?.collection, "issue_4566"); + assert.equal(findBodies[1]?.skip, 100); + assert.equal(findBodies[1]?.limit, 100); + assert.equal(tab?.result?.rows[0]?.[0], 101); + assert.equal(tab?.resultPageOffset, 100); + assert.equal(tab?.resultTotalRowCount, 824); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("mongo find pagination does not use an estimated total as a hard limit", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const findBodies: any[] = []; + + settingsStore.updateEditorSettings({ pageSize: 100 }); + connectionStore.addEphemeralConnection({ + ...conn("mongo-estimated-total-1"), + db_type: "mongodb", + port: 27017, + }); + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + if (String(input) === "/api/document-store/find-documents") { + const body = JSON.parse(String(init?.body ?? "{}")); + findBodies.push(body); + const rowCount = body.skip === 0 ? body.limit : 20; + return Response.json({ + documents: Array.from({ length: rowCount }, (_, index) => ({ _id: body.skip + index + 1 })), + total: 50, + total_is_exact: false, + }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const sql = "db.issue_4566.find({})"; + const tabId = store.createTab("mongo-estimated-total-1", "dbx_test", "Query", "query", ""); + await store.executeTabSql(tabId, sql); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + + assert.equal(tab.result?.total_is_exact, false); + assert.equal(tab.result?.affected_rows, 100, "loaded rows raise the displayed lower bound above a stale estimate"); + assert.equal(tab.result?.truncated, true); + assert.equal(tab.result?.has_more, true); + assert.equal(tab.resultTotalRowCount, undefined, "an estimate must not become the pagination upper bound"); + + await store.executeTabSql(tabId, sql, { + pagination: { offset: 100, limit: 100 }, + preserveResultDuringExecution: true, + preserveTotalRowCountDuringExecution: true, + }); + + assert.equal(findBodies[1]?.skip, 100); + assert.equal(tab.result?.rows.length, 20); + assert.equal(tab.result?.affected_rows, 120); + assert.equal(tab.result?.truncated, false); + assert.equal(tab.result?.has_more, false); + assert.equal(tab.resultPageOffset, 100); + assert.equal(tab.resultTotalRowCount, undefined); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("mongo find execution appends server pages for infinite scroll", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const findBodies: any[] = []; + + settingsStore.updateEditorSettings({ pageSize: 100 }); + connectionStore.addEphemeralConnection({ + ...conn("mongo-append-1"), + db_type: "mongodb", + port: 27017, + }); + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + if (String(input) === "/api/document-store/find-documents") { + const body = JSON.parse(String(init?.body ?? "{}")); + findBodies.push(body); + const available = Math.max(0, 824 - body.skip); + const rowCount = Math.min(body.limit, available); + const documents = Array.from({ length: rowCount }, (_, index) => ({ _id: body.skip + index + 1 })); + return Response.json({ documents, extended_documents: documents, total: 824 }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const sql = "db.issue_4566.find({})"; + const tabId = store.createTab("mongo-append-1", "dbx_test", "Query", "query", ""); + await store.executeTabSql(tabId, sql); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + + await store.executeTabSql(tabId, sql, { + pagination: { offset: 100, limit: 100 }, + appendResult: { maxRows: 5000 }, + preserveResultDuringExecution: true, + preserveTotalRowCountDuringExecution: true, + }); + + assert.equal(findBodies[1]?.skip, 100); + assert.equal(findBodies[1]?.limit, 100); + assert.equal(tab.result?.rows.length, 200); + assert.equal(tab.result?.rows[0]?.[0], 1); + assert.equal(tab.result?.rows[199]?.[0], 200); + assert.equal(tab.result?.mongo_documents?.length, 200); + assert.equal(tab.result?.mongo_copy_documents?.length, 200); + assert.equal(tab.result?.appended_from_row_count, 100); + assert.equal(tab.resultPageOffset, 0); + assert.equal(tab.resultPageLimit, 100); + assert.equal(tab.resultTotalRowCount, 824); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("mongo find pagination preserves explicit skip and limit semantics", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const findBodies: any[] = []; + + settingsStore.updateEditorSettings({ pageSize: 100 }); + connectionStore.addEphemeralConnection({ + ...conn("mongo-bounded-1"), + db_type: "mongodb", + port: 27017, + }); + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + if (String(input) === "/api/document-store/find-documents") { + const body = JSON.parse(String(init?.body ?? "{}")); + findBodies.push(body); + const rowCount = Math.min(body.limit, Math.max(0, 824 - body.skip)); + return Response.json({ + documents: Array.from({ length: rowCount }, (_, index) => ({ _id: body.skip + index + 1 })), + total: 824, + }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const sql = "db.issue_4566.find({}).skip(20).limit(150)"; + const tabId = store.createTab("mongo-bounded-1", "dbx_test", "Query", "query", ""); + await store.executeTabSql(tabId, sql); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + + assert.equal(findBodies[0]?.skip, 20); + assert.equal(findBodies[0]?.limit, 100); + assert.equal(tab.result?.rows.length, 100); + assert.equal(tab.resultTotalRowCount, 150); + + await store.executeTabSql(tabId, sql, { + pagination: { offset: 100, limit: 100 }, + preserveResultDuringExecution: true, + preserveTotalRowCountDuringExecution: true, + }); + + assert.equal(findBodies[1]?.skip, 120); + assert.equal(findBodies[1]?.limit, 50); + assert.equal(tab.result?.rows.length, 50); + assert.equal(tab.result?.rows[0]?.[0], 121); + assert.equal(tab.result?.rows[49]?.[0], 170); + assert.equal(tab.result?.truncated, false); + assert.equal(tab.resultPageOffset, 100); + assert.equal(tab.resultTotalRowCount, 150); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + test("mongo multi-find results use database and collection source labels", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia());