From c1be3da134ec040387d85ae8120ceacb6bc63ff6 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 31 Jul 2026 22:58:55 +0800 Subject: [PATCH] fix(mongodb): export full find query results --- .../src/composables/useDataGridExport.ts | 2 + apps/desktop/src/stores/queryStore.ts | 60 ++++ packages/app-tests/queryStore.test.ts | 302 ++++++++++++++++++ packages/app-tests/useDataGridExport.test.ts | 18 ++ 4 files changed, 382 insertions(+) diff --git a/apps/desktop/src/composables/useDataGridExport.ts b/apps/desktop/src/composables/useDataGridExport.ts index 7dc95eb20..9edd45404 100644 --- a/apps/desktop/src/composables/useDataGridExport.ts +++ b/apps/desktop/src/composables/useDataGridExport.ts @@ -1049,6 +1049,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { if (rowIds !== undefined || context.value !== "results" || !queryResultExportRequest) { return false; } + if (databaseType.value === "mongodb") return false; // The full result is already in memory — don't re-execute the query on the // backend just to stream the same rows back to a file. if (hasCompleteLocalResult?.value) return false; @@ -1122,6 +1123,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { async function exportQueryResultSqlViaBackend(rowIds?: number[]): Promise { // Guard: only for query-result context without complete local result, desktop only if (rowIds !== undefined || context.value !== "results" || !queryResultExportRequest) return false; + if (databaseType.value === "mongodb") return false; if (hasCompleteLocalResult?.value) return false; if (!isTauriRuntime()) return false; // Web → local export fallback diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index ebc8ab7bb..666ad5e6b 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -73,6 +73,7 @@ import { translateBackendError } from "@/i18n/backend-errors"; const ORACLE_LIKE_METADATA_TYPES = new Set(["oracle", "dameng", "oceanbase-oracle"]); const HIDDEN_QUERY_KEY_DATABASE_TYPES = new Set(["mysql", "postgres", "sqlserver", "oracle"]); +const QUERY_RESULT_EXPORT_UNSUPPORTED_ERROR = "Streaming export is unsupported for this query. Simplify it or use a supported driver."; const BACKGROUND_CLIENT_SESSION_SUFFIXES = ["count", "explain", "export"] as const; const CANCEL_QUERY_TIMEOUT_MS = 10_000; const CANCEL_ACK_SETTLE_TIMEOUT_MS = 2_000; @@ -4973,6 +4974,64 @@ export const useQueryStore = defineStore("query", () => { const queryBaseSql = queryResultBaseSql(tab); const exportSettings = useSettingsStore().editorSettings; const exportRowLimit = exportSettings.exportRowLimitEnabled ? exportSettings.exportRowLimit : Number.POSITIVE_INFINITY; + + if (effectiveDbType === "mongodb") { + let mongoCommand; + try { + mongoCommand = await api.mongoParseShellCommand(sql); + } catch { + throw new Error(QUERY_RESULT_EXPORT_UNSUPPORTED_ERROR); + } + if (mongoCommand.kind !== "find") throw new Error(QUERY_RESULT_EXPORT_UNSUPPORTED_ERROR); + + const pageLimit = Math.max(1, Math.trunc(exportSettings.exportBatchSize)); + const documents: unknown[] = []; + let copyDocuments: unknown[] | undefined = []; + let pageOffset = 0; + let totalRows = typeof tab.resultTotalRowCount === "number" ? Math.min(tab.resultTotalRowCount, exportRowLimit) : null; + const exportStartedAt = performance.now(); + const exportExecutionId = uuid(); + + while (documents.length < exportRowLimit) { + const remaining = exportRowLimit - documents.length; + const plan = planMongoFindPagination(sql, mongoCommand, pageOffset, Math.min(pageLimit, remaining)); + if (!plan) throw new Error(QUERY_RESULT_EXPORT_UNSUPPORTED_ERROR); + if (plan.requestLimit === 0) break; + + const result = await api.mongoFindDocuments(tab.connectionId, tab.database, mongoCommand.collection, plan.requestSkip, plan.requestLimit, mongoCommand.filter, mongoCommand.projection, mongoCommand.sort, exportExecutionId); + const pageDocuments = result.documents.slice(0, plan.requestLimit); + documents.push(...pageDocuments); + + if (copyDocuments) { + if (result.extended_documents?.length === result.documents.length) { + copyDocuments.push(...result.extended_documents.slice(0, pageDocuments.length)); + } else { + copyDocuments = undefined; + } + } + + if (result.total_is_exact !== false) { + totalRows = Math.min(mongoFindLogicalTotal(result.total, plan), exportRowLimit); + } + onProgress?.({ rowsExported: documents.length, totalRows }); + + pageOffset += pageDocuments.length; + const reachedLogicalLimit = plan.logicalLimit !== undefined && pageOffset >= plan.logicalLimit; + const reachedExactTotal = result.total_is_exact !== false && pageOffset >= mongoFindLogicalTotal(result.total, plan); + if (pageDocuments.length === 0 || pageDocuments.length < plan.requestLimit || reachedLogicalLimit || reachedExactTotal) break; + } + + const result = mongoDocumentsToQueryResult(documents, performance.now() - exportStartedAt, totalRows ?? documents.length, copyDocuments, totalRows !== null); + if (result.columns.length === 0) { + result.columns = tab.result.columns; + result.column_types = tab.result.column_types; + } + result.affected_rows = documents.length; + result.truncated = false; + result.has_more = false; + return result; + } + const agentExportMaxRows = exportSettings.exportRowLimitEnabled ? exportSettings.exportRowLimit : 2_147_483_647; // Use the already-computed total row count as a progress estimate so the // export dialog shows a moving bar instead of a stuck 0 while paginating. @@ -5050,6 +5109,7 @@ export const useQueryStore = defineStore("query", () => { const settings = useSettingsStore().editorSettings; const effectiveDbType = effectiveDatabaseTypeForConnection(conn); if (!effectiveDbType) return undefined; + if (effectiveDbType === "mongodb") return undefined; const useAgentCursor = usesAgentCursorForQuery(conn?.db_type); const queryBaseSql = queryResultBaseSql(tab); const resultStatementIndex = tab.result.statement_index; diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index 53c146afe..962e3255d 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -3381,6 +3381,308 @@ test("query result export treats the known query total as a progress estimate", } }); +test("MongoDB query result export pages find commands through the document API", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const findBodies: any[] = []; + const progress: Array<{ rowsExported: number; totalRows: number | null }> = []; + const command = 'db.permissions.find({"role":"admin"},{"name":1,"active":1}).sort({"createdTime":-1}).skip(3).limit(205)'; + const documents = Array.from({ length: 205 }, (_, index) => (index === 100 ? { _id: index + 4, active: true } : { _id: index + 4, name: `user-${index + 4}` })); + const copyDocuments = documents.map((document) => ({ ...document, _id: { $numberInt: String(document._id) } })); + + settingsStore.updateEditorSettings({ exportBatchSize: 100, exportRowLimitEnabled: false }); + connectionStore.addEphemeralConnection({ ...conn("mongo-export-find-1"), db_type: "mongodb", port: 27017 }); + const tabId = store.createTab("mongo-export-find-1", "dbx_test"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + tab.lastExecutedSql = command; + tab.resultTotalRowCount = 205; + tab.result = { + columns: ["_id", "name"], + rows: [[4, "user-4"]], + mongo_documents: [documents[0]], + mongo_copy_documents: [copyDocuments[0]], + affected_rows: 205, + execution_time_ms: 1, + sourceStatement: command, + truncated: true, + has_more: true, + }; + + 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 start = body.skip - 3; + return Response.json({ + documents: documents.slice(start, start + body.limit), + extended_documents: copyDocuments.slice(start, start + body.limit), + total: 500, + total_is_exact: true, + }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const exported = await store.fetchTabResultForExport(tabId, (info) => progress.push(info)); + + assert.deepEqual( + findBodies.map(({ skip, limit }) => ({ skip, limit })), + [ + { skip: 3, limit: 100 }, + { skip: 103, limit: 100 }, + { skip: 203, limit: 5 }, + ], + ); + assert.ok(findBodies.every((body) => body.collection === "permissions" && body.filter === '{"role":"admin"}' && body.projection === '{"name":1,"active":1}' && body.sort === '{"createdTime":-1}')); + assert.equal(new Set(findBodies.map((body) => body.executionId)).size, 1); + assert.ok(findBodies[0]?.executionId); + assert.deepEqual(exported?.columns, ["_id", "name", "active"]); + assert.equal(exported?.rows.length, 205); + assert.deepEqual(exported?.rows[0], [4, "user-4", null]); + assert.deepEqual(exported?.rows[100], [104, null, true]); + assert.deepEqual(exported?.rows.at(-1), [208, "user-208", null]); + assert.deepEqual(exported?.mongo_documents, documents); + assert.deepEqual(exported?.mongo_copy_documents, copyDocuments); + assert.equal(exported?.truncated, false); + assert.equal(exported?.has_more, false); + assert.deepEqual(progress, [ + { rowsExported: 100, totalRows: 205 }, + { rowsExported: 200, totalRows: 205 }, + { rowsExported: 205, totalRows: 205 }, + ]); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("MongoDB query result export keeps limit(0) unbounded and stops on a short page", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const findBodies: any[] = []; + const progress: Array<{ rowsExported: number; totalRows: number | null }> = []; + const command = "db.permissions.find({}).limit(0)"; + const documents = Array.from({ length: 104 }, (_, index) => ({ _id: index + 1 })); + + settingsStore.updateEditorSettings({ exportBatchSize: 100, exportRowLimitEnabled: false }); + connectionStore.addEphemeralConnection({ ...conn("mongo-export-unbounded-1"), db_type: "mongodb", port: 27017 }); + const tabId = store.createTab("mongo-export-unbounded-1", "dbx_test"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + tab.lastExecutedSql = command; + tab.result = { columns: ["_id"], rows: [[1]], affected_rows: 4, execution_time_ms: 1, sourceStatement: command, truncated: true, has_more: true }; + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + if (String(input) === "/api/document-store/find-documents") { + const body = JSON.parse(String(init?.body ?? "{}")); + findBodies.push(body); + return Response.json({ documents: documents.slice(body.skip, body.skip + body.limit), total: documents.length, total_is_exact: true }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const exported = await store.fetchTabResultForExport(tabId, (info) => progress.push(info)); + + assert.deepEqual( + findBodies.map(({ skip, limit }) => ({ skip, limit })), + [ + { skip: 0, limit: 100 }, + { skip: 100, limit: 100 }, + ], + ); + assert.equal(exported?.rows.length, 104); + assert.deepEqual(exported?.rows.at(-1), [104]); + assert.deepEqual(progress, [ + { rowsExported: 100, totalRows: 104 }, + { rowsExported: 104, totalRows: 104 }, + ]); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("MongoDB query result export combines negative limits with the export row limit", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const findBodies: any[] = []; + const progress: Array<{ rowsExported: number; totalRows: number | null }> = []; + const command = "db.permissions.find({}).limit(-150)"; + const documents = Array.from({ length: 200 }, (_, index) => ({ _id: index + 1 })); + + settingsStore.updateEditorSettings({ exportBatchSize: 100, exportRowLimit: 120, exportRowLimitEnabled: true }); + connectionStore.addEphemeralConnection({ ...conn("mongo-export-negative-1"), db_type: "mongodb", port: 27017 }); + const tabId = store.createTab("mongo-export-negative-1", "dbx_test"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + tab.lastExecutedSql = command; + tab.result = { columns: ["_id"], rows: [[1]], affected_rows: 150, execution_time_ms: 1, sourceStatement: command, truncated: true, has_more: true }; + + globalThis.fetch = async (input, init) => { + if (String(input) === "/api/connection/check-health") { + return new Response("null", { status: 200, headers: { "Content-Type": "application/json" } }); + } + if (String(input) === "/api/mongo/parse-shell-command") { + return Response.json({ kind: "find", collection: "permissions", filter: "{}", skip: 0, limit: -150 }); + } + if (String(input) === "/api/document-store/find-documents") { + const body = JSON.parse(String(init?.body ?? "{}")); + findBodies.push(body); + return Response.json({ documents: documents.slice(body.skip, body.skip + body.limit), total: documents.length, total_is_exact: true }); + } + return new Response("unexpected request", { status: 500 }); + }; + + try { + const exported = await store.fetchTabResultForExport(tabId, (info) => progress.push(info)); + + assert.deepEqual( + findBodies.map(({ skip, limit }) => ({ skip, limit })), + [ + { skip: 0, limit: 100 }, + { skip: 100, limit: 20 }, + ], + ); + assert.equal(exported?.rows.length, 120); + assert.deepEqual(exported?.rows.at(-1), [120]); + assert.deepEqual(progress, [ + { rowsExported: 100, totalRows: 120 }, + { rowsExported: 120, totalRows: 120 }, + ]); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("MongoDB query result export preserves columns when a find command returns no documents", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const findBodies: any[] = []; + + settingsStore.updateEditorSettings({ exportBatchSize: 100, exportRowLimitEnabled: false }); + connectionStore.addEphemeralConnection({ ...conn("mongo-export-safe-1"), db_type: "mongodb", port: 27017 }); + const findCommand = "db.permissions.find({})"; + const tabId = store.createTab("mongo-export-safe-1", "dbx_test"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + tab.lastExecutedSql = findCommand; + tab.result = { columns: ["_id", "name"], column_types: ["objectId", "string"], rows: [["old", "old"]], affected_rows: 1, execution_time_ms: 1, sourceStatement: findCommand, truncated: true, has_more: true }; + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + if (String(input) === "/api/document-store/find-documents") { + findBodies.push(JSON.parse(String(init?.body ?? "{}"))); + return Response.json({ documents: [], total: 0, total_is_exact: true }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const empty = await store.fetchTabResultForExport(tabId); + assert.equal(findBodies.length, 1); + assert.deepEqual(empty?.columns, ["_id", "name"]); + assert.deepEqual(empty?.column_types, ["objectId", "string"]); + assert.deepEqual(empty?.rows, []); + assert.deepEqual(empty?.mongo_documents, []); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("MongoDB query result export rejects non-find and parse failures without replay", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const replayRequests: string[] = []; + const unsupportedError = "Streaming export is unsupported for this query. Simplify it or use a supported driver."; + + connectionStore.addEphemeralConnection({ ...conn("mongo-export-unsupported-1"), db_type: "mongodb", port: 27017 }); + const tabId = store.createTab("mongo-export-unsupported-1", "dbx_test"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + + globalThis.fetch = withConnectionHealthMock(async (input) => { + replayRequests.push(String(input)); + return new Response("unexpected request", { status: 500 }); + }); + + try { + const aggregateCommand = "db.permissions.aggregate([])"; + tab.lastExecutedSql = aggregateCommand; + tab.result = { columns: ["count"], rows: [[7]], affected_rows: 1, execution_time_ms: 1, sourceStatement: aggregateCommand, truncated: true, has_more: true }; + + await assert.rejects(store.fetchTabResultForExport(tabId), { message: unsupportedError }); + const backendRequest = await store.buildQueryResultExportRequest(tabId, { exportId: "mongo-export", filePath: "/tmp/mongo.csv", format: "csv" }); + assert.equal(backendRequest, undefined); + + const invalidCommand = "db.permissions.find("; + tab.lastExecutedSql = invalidCommand; + tab.result = { columns: ["_id"], rows: [["partial"]], affected_rows: 1, execution_time_ms: 1, sourceStatement: invalidCommand, truncated: true, has_more: true }; + + await assert.rejects(store.fetchTabResultForExport(tabId), { message: unsupportedError }); + assert.deepEqual(replayRequests, []); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("MongoDB query result export rejects pagination-plan failures without replay", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const replayRequests: string[] = []; + const unsupportedError = "Streaming export is unsupported for this query. Simplify it or use a supported driver."; + const command = "db.permissions.find({}) trailing"; + + connectionStore.addEphemeralConnection({ ...conn("mongo-export-plan-failure-1"), db_type: "mongodb", port: 27017 }); + const tabId = store.createTab("mongo-export-plan-failure-1", "dbx_test"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + tab.lastExecutedSql = command; + tab.result = { columns: ["_id"], rows: [["partial"]], affected_rows: 1, execution_time_ms: 1, sourceStatement: command, truncated: true, has_more: true }; + + globalThis.fetch = async (input) => { + const url = String(input); + if (url === "/api/connection/check-health") return Response.json(null); + if (url === "/api/mongo/parse-shell-command") return Response.json({ kind: "find", collection: "permissions", filter: "{}", skip: 0, limit: 100 }); + replayRequests.push(url); + return new Response("unexpected request", { status: 500 }); + }; + + try { + await assert.rejects(store.fetchTabResultForExport(tabId), { message: unsupportedError }); + assert.deepEqual(replayRequests, []); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + test("jdbc query pagination uses result sessions without capping max rows to one page", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia()); diff --git a/packages/app-tests/useDataGridExport.test.ts b/packages/app-tests/useDataGridExport.test.ts index d15a95a2f..0ccc00438 100644 --- a/packages/app-tests/useDataGridExport.test.ts +++ b/packages/app-tests/useDataGridExport.test.ts @@ -397,6 +397,24 @@ test("full query result CSV export streams through the backend without loading a assert.equal(exportProgressState.value.filePath, apiMock.startQueryResultExport.mock.calls[0][0].filePath); }); +test("MongoDB full query result CSV export uses the full-result fallback", async () => { + const { composable, fullExportResult, queryResultExportRequest } = buildExportHarness({ databaseType: "mongodb" }); + fullExportResult.mockResolvedValueOnce({ + columns: ["_id", "name"], + rows: [["1", "Ada"]], + affected_rows: 1, + execution_time_ms: 1, + }); + + await composable.exportCsv(); + + assert.equal(queryResultExportRequest.mock.calls.length, 0); + assert.equal(apiMock.startQueryResultExport.mock.calls.length, 0); + assert.equal(fullExportResult.mock.calls.length, 1); + assert.deepEqual(apiMock.exportQueryResultCsv.mock.calls[0][1], ["_id", "name"]); + assert.deepEqual(apiMock.exportQueryResultCsv.mock.calls[0][2], [["1", "Ada"]]); +}); + test("streaming query result export translates streaming unsupported error before the toast", async () => { const rawMessage = "Streaming export is unsupported for this query. Simplify it or use a supported driver."; apiMock.startQueryResultExport.mockImplementationOnce(async (request, onProgress) => {