From 7f5de70687ab3abb365ebaee374befccfcd03778 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 1 Jun 2026 14:32:16 +0800 Subject: [PATCH] fix(export): include all paginated query result rows --- apps/desktop/src/components/grid/DataGrid.vue | 2 + .../src/components/layout/ContentArea.vue | 2 + .../src/composables/useDataGridExport.ts | 39 ++++++---- apps/desktop/src/stores/queryStore.ts | 76 +++++++++++++++++++ packages/app-tests/queryStore.test.ts | 70 +++++++++++++++++ 5 files changed, 174 insertions(+), 15 deletions(-) diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index 23c40e3788..31279fada8 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -216,6 +216,7 @@ const props = defineProps<{ loading?: boolean; cacheKey?: string; onExecuteSql?: (sql: string) => Promise; + fullExportResult?: () => Promise; customSave?: (changes: { dirtyRows: Map>; newRows: (string | number | boolean | null)[][]; @@ -2944,6 +2945,7 @@ const { getRowItem: (rowId: number) => visibleDisplayItems.value.find((item) => item.id === rowId), selectedRowIds, hasRowSelection, + fullExportResult: props.fullExportResult, }); const pageSizeMenuItems = computed(() => diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index 203339f6e7..d3d8cd5ca0 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -465,6 +465,7 @@ defineExpose({ focusSearch, refreshData, handleModRTarget }); :count-sql="activeTab.resultCountSql" :total-row-count="activeTab.resultTotalRowCount" :on-execute-sql="async (sql: string) => emit('executeSql', sql)" + :full-export-result="() => queryStore.fetchTabResultForExport(activeTab.id)" @reload=" ( sql?: string, @@ -667,6 +668,7 @@ defineExpose({ focusSearch, refreshData, handleModRTarget }); :page-offset="activeTab.resultPageOffset" :page-limit="activeTab.resultPageLimit" :on-execute-sql="async (sql: string) => emit('executeSql', sql)" + :full-export-result="() => queryStore.fetchTabResultForExport(activeTab.id)" @update:where-input="(v: string) => (activeTab.whereInput = v)" @reload=" ( diff --git a/apps/desktop/src/composables/useDataGridExport.ts b/apps/desktop/src/composables/useDataGridExport.ts index d79dbc2611..102f5ef70b 100644 --- a/apps/desktop/src/composables/useDataGridExport.ts +++ b/apps/desktop/src/composables/useDataGridExport.ts @@ -15,7 +15,7 @@ import { displayCellValue, type CellValue } from "@/lib/cellValue"; import { tryStartExclusiveActivation, type ActionActivationGuard } from "@/lib/actionActivation"; import { copyToClipboard } from "@/lib/clipboard"; import { buildDataGridCopyInsertStatement, buildDataGridCopyUpdateStatements } from "@/lib/dataGridSql"; -import type { DatabaseType } from "@/types/database"; +import type { DatabaseType, QueryResult } from "@/types/database"; interface RowItem { id: number; @@ -44,6 +44,7 @@ export interface UseDataGridExportOptions { getRowItem: (rowId: number) => RowItem | undefined; selectedRowIds: Ref> | ComputedRef>; hasRowSelection: ComputedRef; + fullExportResult?: () => Promise; } interface CopyStatementCache { @@ -90,6 +91,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { getRowItem, selectedRowIds, hasRowSelection, + fullExportResult, } = options; async function copyText(text: string) { @@ -107,6 +109,17 @@ export function useDataGridExport(options: UseDataGridExportOptions) { return displayItems.value.filter((item) => rowIdSet.has(item.id)); } + async function resultToExport(rowIds?: number[]): Promise<{ columns: string[]; rows: CellValue[][] }> { + if (!rowIds?.length && fullExportResult) { + const result = await fullExportResult(); + if (result) return { columns: result.columns, rows: result.rows }; + } + return { + columns: columns.value, + rows: rowsToExport(rowIds).map((item) => item.data), + }; + } + function targetedRows(): RowItem[] { if (hasRowSelection.value && selectedRowIds.value.size > 0) { return displayItems.value.filter((item) => selectedRowIds.value.has(item.id)); @@ -427,7 +440,8 @@ export function useDataGridExport(options: UseDataGridExportOptions) { async function exportCsv(rowIds?: number[]) { await runExclusiveExport(async () => { try { - const rows = rowsToExport(rowIds).map((item) => item.data.map((c) => displayCellValue(c))); + const result = await resultToExport(rowIds); + const rows = result.rows.map((row) => row.map((c) => displayCellValue(c))); let outputPath = "export.csv"; if (isTauriRuntime()) { const { save } = await import("@tauri-apps/plugin-dialog"); @@ -438,7 +452,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { if (!path) return; outputPath = path as string; } - await api.exportQueryResultCsv(outputPath, columns.value, rows); + await api.exportQueryResultCsv(outputPath, result.columns, rows); toast(t("grid.exported")); } catch (e: any) { toast(t("grid.exportFailed", { message: e?.message || String(e) }), 5000); @@ -459,11 +473,8 @@ export function useDataGridExport(options: UseDataGridExportOptions) { if (!path) return; outputPath = path as string; } - await api.exportQueryResultJson( - outputPath, - columns.value, - rowsToExport(rowIds).map((item) => item.data), - ); + const result = await resultToExport(rowIds); + await api.exportQueryResultJson(outputPath, result.columns, result.rows); toast(t("grid.exported")); } catch (e: any) { toast(t("grid.exportFailed", { message: e?.message || String(e) }), 5000); @@ -484,11 +495,8 @@ export function useDataGridExport(options: UseDataGridExportOptions) { if (!path) return; outputPath = path as string; } - await api.exportQueryResultMarkdown( - outputPath, - columns.value, - rowsToExport(rowIds).map((item) => item.data), - ); + const result = await resultToExport(rowIds); + await api.exportQueryResultMarkdown(outputPath, result.columns, result.rows); toast(t("grid.exported")); } catch (e: any) { toast(t("grid.exportFailed", { message: e?.message || String(e) }), 5000); @@ -509,11 +517,12 @@ export function useDataGridExport(options: UseDataGridExportOptions) { if (!path) return; outputPath = path as string; } + const result = await resultToExport(rowIds); await api.exportQueryResultXlsx( outputPath, tableMeta.value?.tableName || "Export", - columns.value, - rowsToExport(rowIds).map((item) => item.data), + result.columns, + result.rows, ); toast(t("grid.exported")); } catch (e: any) { diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 79766f3a25..04bf1f250e 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -20,6 +20,7 @@ import { } from "@/lib/mongoShellCommand"; import { AGENT_DRIVER_TYPES } from "@/lib/databaseCapabilities"; import { editablePrimaryKeys } from "@/lib/tableEditing"; +import { TABLE_DATA_EXPORT_PAGE_SIZE } from "@/lib/tableDataExport"; import * as api from "@/lib/api"; import { useConnectionStore } from "@/stores/connectionStore"; import { useSettingsStore } from "@/stores/settingsStore"; @@ -1073,6 +1074,80 @@ export const useQueryStore = defineStore("query", () => { }); } + async function fetchTabResultForExport(id: string): Promise { + const tab = tabs.value.find((t) => t.id === id); + if (!tab?.result) return undefined; + if (tab.mode !== "query") return tab.result; + + const sql = tab.resultSortedSql ?? tab.resultBaseSql ?? tab.lastExecutedSql ?? tab.sql; + if (!sql.trim()) return tab.result; + + const connStore = useConnectionStore(); + await connStore.ensureConnected(tab.connectionId); + const conn = connStore.getConfig(tab.connectionId); + const useAgentCursor = !!conn?.db_type && AGENT_DRIVER_TYPES.has(conn.db_type); + const queryBaseSql = tab.resultBaseSql ?? sql; + const pageLimit = Math.max(tab.resultPageLimit ?? 0, TABLE_DATA_EXPORT_PAGE_SIZE); + const rows: QueryResult["rows"] = []; + let columns: string[] = []; + let executionTimeMs = 0; + let offset = 0; + let sessionId: string | undefined; + const clientSessionId = `${tab.id}:export`; + + try { + while (true) { + const plan = await api.prepareQueryPaginationExecutionPlan({ + sql, + queryBaseSql, + databaseType: conn?.db_type, + pagination: { limit: pageLimit, offset, sessionId }, + useAgentCursor, + }); + if (typeof plan.pageLimit !== "number" || typeof plan.pageOffset !== "number") return tab.result; + const executionOptions = plan.useAgentResultSession + ? { + maxRows: plan.pageLimit, + fetchSize: plan.pageLimit, + pageSize: plan.pageLimit, + resultSessionId: sessionId, + clientSessionId, + } + : { maxRows: plan.pageLimit, fetchSize: plan.pageLimit }; + const results = await api.executeMulti( + tab.connectionId, + tab.database, + plan.sqlToExecute, + tab.schema, + undefined, + executionOptions, + ); + const result = results[0]; + if (!result) break; + if (columns.length === 0) columns = result.columns; + rows.push(...result.rows); + executionTimeMs += result.execution_time_ms ?? 0; + sessionId = result.session_id ?? undefined; + const shouldFetchNextPage = plan.useAgentResultSession + ? result.has_more === true + : result.rows.length >= plan.pageLimit; + if (!shouldFetchNextPage) break; + offset += result.rows.length; + } + } finally { + if (sessionId) void api.closeQuerySession(tab.connectionId, tab.database, sessionId, clientSessionId); + } + + return { + columns: columns.length ? columns : tab.result.columns, + rows, + affected_rows: 0, + execution_time_ms: executionTimeMs, + truncated: false, + has_more: false, + }; + } + return { tabs, activeTabId, @@ -1102,6 +1177,7 @@ export const useQueryStore = defineStore("query", () => { cancelTabExecution, cancelTabExplain, reloadEvictedTab, + fetchTabResultForExport, notifyConnectionMayBeLost, }; }); diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index f0e3ae504a..9ac84d3840 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -364,6 +364,76 @@ test("data tab execution preserves pagination offset metadata", async () => { } }); +test("query result export fetches every paginated page", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const preparedOffsets: number[] = []; + const executedSqls: string[] = []; + + connectionStore.addEphemeralConnection(conn("conn-1")); + const tabId = store.createTab("conn-1", "db"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + tab.lastExecutedSql = "select id from users"; + tab.resultPageLimit = 100; + tab.resultPageOffset = 0; + tab.result = { + columns: ["id"], + rows: [[1]], + affected_rows: 0, + execution_time_ms: 1, + truncated: false, + has_more: true, + }; + + globalThis.fetch = (async (input, init) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + const body = JSON.parse(String(init?.body ?? "{}")); + const offset = Number(body.options.pagination.offset); + const limit = Number(body.options.pagination.limit); + preparedOffsets.push(offset); + return new Response( + JSON.stringify({ + sqlToExecute: `select id from users /* offset:${offset} */`, + pageSql: `select id from users /* offset:${offset} */`, + pageLimit: limit, + pageOffset: offset, + useAgentResultSession: false, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ); + } + if (url === "/api/query/execute-multi") { + const body = JSON.parse(String(init?.body ?? "{}")); + executedSqls.push(body.sql); + const rows = String(body.sql).includes("offset:0") + ? Array.from({ length: 10_000 }, (_, index) => [index + 1]) + : [[10_001], [10_002]]; + return new Response(JSON.stringify([{ columns: ["id"], rows, affected_rows: 0, execution_time_ms: 1 }]), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response("unexpected request", { status: 500 }); + }) as typeof fetch; + + try { + const exported = await store.fetchTabResultForExport(tabId); + + assert.deepEqual(preparedOffsets, [0, 10_000]); + assert.deepEqual(executedSqls, ["select id from users /* offset:0 */", "select id from users /* offset:10000 */"]); + assert.equal(exported?.rows.length, 10_002); + assert.deepEqual(exported?.rows.at(-1), [10_002]); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + test("query execution finishes without waiting for metadata analysis", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia());