From 588ba3352d9651dca3358dbf5020b22c940b3245 Mon Sep 17 00:00:00 2001 From: miracle Date: Wed, 24 Jun 2026 21:18:11 +0800 Subject: [PATCH] feat: stream query result CSV/XLSX exports * docs: add query result streaming export design * feat(core): stream query result exports * feat: expose query result streaming export APIs * feat(desktop): stream query result exports from grid * fix: route query result export cancellation through its command * fix: reduce query export batch size and cancel latency * opt: del local test files --- .../editor/EditorSettingsDialog.vue | 41 ++ apps/desktop/src/components/grid/DataGrid.vue | 10 +- .../src/components/layout/ContentArea.vue | 1 + .../src/composables/useDataGridExport.ts | 123 +++- apps/desktop/src/i18n/locales/en.ts | 9 +- apps/desktop/src/i18n/locales/zh-CN.ts | 9 +- apps/desktop/src/lib/api.ts | 3 + apps/desktop/src/lib/http.ts | 59 ++ apps/desktop/src/lib/tauri.ts | 66 ++ apps/desktop/src/stores/queryStore.ts | 57 +- apps/desktop/src/stores/settingsStore.ts | 23 +- crates/dbx-core/src/csv_export.rs | 14 +- crates/dbx-core/src/lib.rs | 1 + crates/dbx-core/src/query_result_export.rs | 621 ++++++++++++++++++ crates/dbx-core/src/query_result_sql.rs | 26 +- crates/dbx-core/src/xlsx_export.rs | 60 +- crates/dbx-web/src/main.rs | 11 + crates/dbx-web/src/routes/mod.rs | 1 + .../dbx-web/src/routes/query_result_export.rs | 168 +++++ packages/app-tests/queryStore.test.ts | 177 +++++ packages/app-tests/settingsStore.test.ts | 47 +- packages/app-tests/useDataGridExport.test.ts | 216 ++++++ src-tauri/src/commands/mod.rs | 1 + src-tauri/src/commands/query_result_export.rs | 88 +++ src-tauri/src/lib.rs | 2 + 25 files changed, 1787 insertions(+), 47 deletions(-) create mode 100644 crates/dbx-core/src/query_result_export.rs create mode 100644 crates/dbx-web/src/routes/query_result_export.rs create mode 100644 packages/app-tests/useDataGridExport.test.ts create mode 100644 src-tauri/src/commands/query_result_export.rs diff --git a/apps/desktop/src/components/editor/EditorSettingsDialog.vue b/apps/desktop/src/components/editor/EditorSettingsDialog.vue index c2aa6fd20c..38225860fe 100644 --- a/apps/desktop/src/components/editor/EditorSettingsDialog.vue +++ b/apps/desktop/src/components/editor/EditorSettingsDialog.vue @@ -136,6 +136,9 @@ const editSidebarHiddenTablePrefixes = ref(settingsStore.editorSettings.sidebarH const editSidebarHideTableComments = ref(settingsStore.editorSettings.sidebarHideTableComments); const editSidebarAllowHorizontalScroll = ref(settingsStore.editorSettings.sidebarAllowHorizontalScroll); const editExportBatchSize = ref(settingsStore.editorSettings.exportBatchSize); +const editExportRowLimitEnabled = ref(settingsStore.editorSettings.exportRowLimitEnabled); +const editExportRowLimit = ref(settingsStore.editorSettings.exportRowLimit); +const editQueryExportKeysetOptimizationEnabled = ref(settingsStore.editorSettings.queryExportKeysetOptimizationEnabled); const editToolbarItems = ref({ ...settingsStore.editorSettings.toolbarItems }); const redisScanPageSizeOptions = [200, 1000, 5000, 10000]; const systemFonts = ref([]); @@ -394,6 +397,9 @@ watch( editSidebarHideTableComments.value = settingsStore.editorSettings.sidebarHideTableComments; editSidebarAllowHorizontalScroll.value = settingsStore.editorSettings.sidebarAllowHorizontalScroll; editExportBatchSize.value = settingsStore.editorSettings.exportBatchSize; + editExportRowLimitEnabled.value = settingsStore.editorSettings.exportRowLimitEnabled; + editExportRowLimit.value = settingsStore.editorSettings.exportRowLimit; + editQueryExportKeysetOptimizationEnabled.value = settingsStore.editorSettings.queryExportKeysetOptimizationEnabled; editToolbarItems.value = { ...settingsStore.editorSettings.toolbarItems }; editSnippets.value = settingsStore.editorSettings.snippets.map((s) => ({ ...s })); } @@ -451,6 +457,9 @@ function hasChanges(): boolean { editSidebarHideTableComments.value !== settingsStore.editorSettings.sidebarHideTableComments || editSidebarAllowHorizontalScroll.value !== settingsStore.editorSettings.sidebarAllowHorizontalScroll || editExportBatchSize.value !== settingsStore.editorSettings.exportBatchSize || + editExportRowLimitEnabled.value !== settingsStore.editorSettings.exportRowLimitEnabled || + editExportRowLimit.value !== settingsStore.editorSettings.exportRowLimit || + editQueryExportKeysetOptimizationEnabled.value !== settingsStore.editorSettings.queryExportKeysetOptimizationEnabled || JSON.stringify(editToolbarItems.value) !== JSON.stringify(settingsStore.editorSettings.toolbarItems) || JSON.stringify(normalizeSidebarHiddenTablePrefixes(editSidebarHiddenTablePrefixes.value)) !== JSON.stringify(settingsStore.editorSettings.sidebarHiddenTablePrefixes) || JSON.stringify(editSnippets.value) !== JSON.stringify(settingsStore.editorSettings.snippets) @@ -491,6 +500,9 @@ async function persistSettings() { sidebarAllowHorizontalScroll: editSidebarAllowHorizontalScroll.value, sidebarHiddenTablePrefixes: normalizeSidebarHiddenTablePrefixes(editSidebarHiddenTablePrefixes.value), exportBatchSize: editExportBatchSize.value, + exportRowLimitEnabled: editExportRowLimitEnabled.value, + exportRowLimit: editExportRowLimit.value, + queryExportKeysetOptimizationEnabled: editQueryExportKeysetOptimizationEnabled.value, toolbarItems: { ...editToolbarItems.value }, snippets: editSnippets.value, }); @@ -559,6 +571,9 @@ function resetDefaultsForTab(tab: SettingsCategory) { editInfiniteScroll.value = DEFAULT_EDITOR_SETTINGS.infiniteScroll; editInfiniteScrollMaxRows.value = DEFAULT_EDITOR_SETTINGS.infiniteScrollMaxRows; editExportBatchSize.value = DEFAULT_EDITOR_SETTINGS.exportBatchSize; + editExportRowLimitEnabled.value = DEFAULT_EDITOR_SETTINGS.exportRowLimitEnabled; + editExportRowLimit.value = DEFAULT_EDITOR_SETTINGS.exportRowLimit; + editQueryExportKeysetOptimizationEnabled.value = DEFAULT_EDITOR_SETTINGS.queryExportKeysetOptimizationEnabled; } else if (tab === "redis") { editRedisScanPageSize.value = DEFAULT_EDITOR_SETTINGS.redisScanPageSize; } else if (tab === "shortcuts") { @@ -606,6 +621,9 @@ function resetAllDefaults() { editSidebarAllowHorizontalScroll.value = DEFAULT_EDITOR_SETTINGS.sidebarAllowHorizontalScroll; editSidebarHiddenTablePrefixes.value = DEFAULT_EDITOR_SETTINGS.sidebarHiddenTablePrefixes.join("\n"); editExportBatchSize.value = DEFAULT_EDITOR_SETTINGS.exportBatchSize; + editExportRowLimitEnabled.value = DEFAULT_EDITOR_SETTINGS.exportRowLimitEnabled; + editExportRowLimit.value = DEFAULT_EDITOR_SETTINGS.exportRowLimit; + editQueryExportKeysetOptimizationEnabled.value = DEFAULT_EDITOR_SETTINGS.queryExportKeysetOptimizationEnabled; editToolbarItems.value = { ...DEFAULT_EDITOR_SETTINGS.toolbarItems }; editSnippets.value = DEFAULT_SQL_SNIPPETS.map((s) => ({ ...s })); } @@ -2282,6 +2300,29 @@ watch( {{ t("settings.exportBatchSizeDescription") }} +
+
+ +

{{ t("settings.exportRowLimitEnabledDescription") }}

+
+ +
+
+ +
+ + + {{ editExportRowLimitEnabled ? t("settings.exportRowLimitDescription") : t("settings.exportRowLimitUnlimited") }} + +
+
+
+
+ +

{{ t("settings.queryExportKeysetOptimizationEnabledDescription") }}

+
+ +
diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index a20b446856..c333da46f5 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -210,6 +210,7 @@ const props = defineProps<{ cacheKey?: string; onExecuteSql?: (sql: string) => Promise; fullExportResult?: (onProgress?: (info: { rowsExported: number; totalRows: number | null }) => void) => Promise; + queryResultExportRequest?: (options: { exportId: string; filePath: string; format: "csv" | "xlsx" }) => Promise; allExportResults?: Array<{ sheetName: string; result: QueryResult }>; customSaveHandler?: import("@/composables/useDataGridEditor").CustomSaveHandler; }>(); @@ -4558,6 +4559,11 @@ const exportProgressState = ref({ status: "", errorMessage: null as string | null, }); +const exportCancelHandler = ref<(() => Promise) | null>(null); + +async function cancelActiveExport() { + await exportCancelHandler.value?.(); +} // --- Export composable --- const { @@ -4615,9 +4621,11 @@ const { selectedRowIds, hasRowSelection, fullExportResult: props.fullExportResult, + queryResultExportRequest: props.queryResultExportRequest, allExportResults: computed(() => props.allExportResults), exportProgressDialog, exportProgressState, + exportCancelHandler, }); const pageSizeMenuItems = computed(() => @@ -8430,7 +8438,7 @@ const gridContextMenuItems = computed(() => { /> - + diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index a0b9be189d..d940a335b9 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -797,6 +797,7 @@ defineExpose({ focusSearch, refreshData, handleModRTarget, requestQueryEditorExe :total-row-count-loading="activeTab.resultTotalRowCountLoading" :on-execute-sql="async (sql: string) => emit('executeSql', sql)" :full-export-result="(onProgress?: (info: { rowsExported: number; totalRows: number | null }) => void) => queryStore.fetchTabResultForExport(activeTab.id, onProgress)" + :query-result-export-request="(options: { exportId: string; filePath: string; format: 'csv' | 'xlsx' }) => queryStore.buildQueryResultExportRequest(activeTab.id, options)" :all-export-results="allResultExportSheets" @update:order-by-input="(v: string) => (activeTab.orderByInput = v)" @reload="(sql?: string, searchText?: string, whereInput?: string, orderBy?: string, limit?: number, offset?: number) => emit('reload', sql, searchText, whereInput, orderBy, limit, offset)" diff --git a/apps/desktop/src/composables/useDataGridExport.ts b/apps/desktop/src/composables/useDataGridExport.ts index a4abc39d23..13fc41aa59 100644 --- a/apps/desktop/src/composables/useDataGridExport.ts +++ b/apps/desktop/src/composables/useDataGridExport.ts @@ -11,6 +11,7 @@ import { buildDataGridCopyInsertStatement, buildDataGridCopyUpdateStatements, ty import { formatSqlInsert } from "@/lib/exportFormats"; import { uuid } from "@/lib/utils"; import type { DatabaseType, QueryResult } from "@/types/database"; +import type { QueryResultExportRequest } from "@/lib/api"; interface RowItem { id: number; @@ -45,6 +46,7 @@ export interface UseDataGridExportOptions { selectedRowIds: Ref> | ComputedRef>; hasRowSelection: ComputedRef; fullExportResult?: (onProgress?: (info: { rowsExported: number; totalRows: number | null }) => void) => Promise; + queryResultExportRequest?: (options: { exportId: string; filePath: string; format: "csv" | "xlsx" }) => Promise; allExportResults?: ComputedRef | undefined>; exportProgressDialog?: Ref; exportProgressState?: Ref<{ @@ -56,6 +58,7 @@ export interface UseDataGridExportOptions { status: string; errorMessage: string | null; }>; + exportCancelHandler?: Ref<(() => Promise) | null>; } interface CopyStatementCache { @@ -110,9 +113,11 @@ export function useDataGridExport(options: UseDataGridExportOptions) { selectedRowIds, hasRowSelection, fullExportResult, + queryResultExportRequest, allExportResults, exportProgressDialog, exportProgressState, + exportCancelHandler, } = options; async function copyText(text: string) { @@ -469,6 +474,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { async function exportCsv(rowIds?: number[]) { await runExclusiveExport(async () => { try { + if (await exportQueryResultViaBackend("csv", rowIds)) return; if (await exportFullTableDataViaBackend("csv", rowIds)) return; const needsFullExport = !rowIds?.length && !!fullExportResult; @@ -662,6 +668,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) { async function exportXlsx(rowIds?: number[]) { await runExclusiveExport(async () => { try { + if (await exportQueryResultViaBackend("xlsx", rowIds)) return; if (await exportFullTableDataViaBackend("xlsx", rowIds)) return; let outputPath = "export.xlsx"; @@ -814,37 +821,109 @@ export function useDataGridExport(options: UseDataGridExportOptions) { } if (exportProgressDialog) exportProgressDialog.value = true; - await api.startTableExport( - { - exportId: uuid(), - connectionId: connectionId.value, - database: database.value, - schema: meta.schema, - tableName: meta.tableName, - filePath: outputPath, + const exportId = uuid(); + if (exportCancelHandler) { + exportCancelHandler.value = () => api.cancelTableExport(exportId); + } + + try { + const progress = await api.startTableExport( + { + exportId, + connectionId: connectionId.value, + database: database.value, + schema: meta.schema, + tableName: meta.tableName, + filePath: outputPath, + format, + columns: columns.value, + columnTypes: columnTypes.value, + primaryKeys: meta.primaryKeys, + whereInput: whereInput.value, + orderBy: orderBy.value, + skipCount: true, + batchSize: exportBatchSize.value, + }, + (progress) => { + if (exportProgressState) { + exportProgressState.value = { + ...exportProgressState.value, + tableName: progress.tableName || meta.tableName, + rowsExported: progress.rowsExported, + totalRows: progress.totalRows, + status: progress.status, + errorMessage: progress.errorMessage || null, + }; + } + }, + ); + if (progress.status === "Done") { + toast(t("grid.exported")); + } + } finally { + if (exportCancelHandler) exportCancelHandler.value = null; + } + return true; + } + + async function exportQueryResultViaBackend(format: "csv" | "xlsx", rowIds?: number[]): Promise { + if (rowIds?.length || context.value !== "results" || !queryResultExportRequest) { + return false; + } + + const extension = format; + const filterName = format === "csv" ? "CSV" : "Excel"; + let outputPath = `query-result.${extension}`; + if (isTauriRuntime()) { + const { save } = await import("@tauri-apps/plugin-dialog"); + const path = await save({ + defaultPath: outputPath, + filters: [{ name: filterName, extensions: [extension] }], + }); + if (!path) return true; + outputPath = path as string; + } + + const exportId = uuid(); + const request = await queryResultExportRequest({ exportId, filePath: outputPath, format }); + if (!request) throw new Error("Unable to build query result export request"); + + if (exportProgressState) { + exportProgressState.value = { + title: t("exportProgress.title"), + tableName: "Query Result", format, - columns: columns.value, - columnTypes: columnTypes.value, - primaryKeys: meta.primaryKeys, - whereInput: whereInput.value, - orderBy: orderBy.value, - skipCount: true, - batchSize: exportBatchSize.value, - }, - (progress) => { + rowsExported: 0, + totalRows: request.totalRows ?? null, + status: "Running", + errorMessage: null, + }; + } + if (exportProgressDialog) exportProgressDialog.value = true; + if (exportCancelHandler) { + exportCancelHandler.value = () => api.cancelQueryResultExport(exportId, request.executionId); + } + + try { + const terminalProgress = await api.startQueryResultExport(request, (progress) => { if (exportProgressState) { + const adjustedTotal = progress.totalRows !== null && progress.rowsExported > progress.totalRows ? progress.rowsExported : progress.totalRows; exportProgressState.value = { ...exportProgressState.value, - tableName: progress.tableName || meta.tableName, + tableName: progress.tableName || "Query Result", rowsExported: progress.rowsExported, - totalRows: progress.totalRows, + totalRows: adjustedTotal, status: progress.status, errorMessage: progress.errorMessage || null, }; } - }, - ); - toast(t("grid.exported")); + }); + if (terminalProgress.status === "Done") { + toast(t("grid.exported")); + } + } finally { + if (exportCancelHandler) exportCancelHandler.value = null; + } return true; } diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 291aabe0b9..02effcf4a8 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -2472,7 +2472,14 @@ export default { redisScanPageSizeDescription: "Keys requested per Redis SCAN page when browsing keys.", redisScanPageSizeOption: "{count} keys", exportBatchSize: "Export batch size", - exportBatchSizeDescription: "Rows fetched per batch when exporting table data (100-100000).", + exportBatchSizeDescription: "Rows fetched per batch when exporting data (100-100000).", + exportRowLimitEnabled: "Limit exported rows", + exportRowLimitEnabledDescription: "When on, query result CSV/XLSX exports stop at the row limit below.", + exportRowLimit: "Query result export row limit", + exportRowLimitDescription: "Maximum rows for query result exports.", + exportRowLimitUnlimited: "Unlimited", + queryExportKeysetOptimizationEnabled: "Try keyset optimization", + queryExportKeysetOptimizationEnabledDescription: "Only applies to safely recognized single-table queries; complex queries automatically fall back.", exportSection: "Export", shortcutExecuteSql: "Execute SQL", shortcutFind: "Find", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index ba089ef93b..e56f3c10c3 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -2496,7 +2496,14 @@ export default { redisScanPageSizeDescription: "浏览 Redis Key 时每次 SCAN 请求的 Key 数量。", redisScanPageSizeOption: "{count} 个 Key", exportBatchSize: "导出批次大小", - exportBatchSizeDescription: "导出表数据时每批次获取的行数（100-100000）。", + exportBatchSizeDescription: "导出数据时每批读取的行数（100-100000）。", + exportRowLimitEnabled: "限制导出行数", + exportRowLimitEnabledDescription: "开启时，查询结果 CSV/XLSX 导出最多为下方设置的行数。", + exportRowLimit: "查询结果导出最多行数", + exportRowLimitDescription: "查询结果导出的总行数上限。", + exportRowLimitUnlimited: "无限制", + queryExportKeysetOptimizationEnabled: "尝试 keyset 优化", + queryExportKeysetOptimizationEnabledDescription: "仅对可安全识别的单表查询生效，复杂查询会自动回退。", exportSection: "导出", shortcutExecuteSql: "执行 SQL", shortcutFind: "查找", diff --git a/apps/desktop/src/lib/api.ts b/apps/desktop/src/lib/api.ts index f3977c0a13..1204b80f96 100644 --- a/apps/desktop/src/lib/api.ts +++ b/apps/desktop/src/lib/api.ts @@ -267,6 +267,8 @@ export const exportQueryResultJson = forward("exportQueryResultJson"); export const exportQueryResultMarkdown = forward("exportQueryResultMarkdown"); export const startTableExport = forward("startTableExport"); export const cancelTableExport = forward("cancelTableExport"); +export const startQueryResultExport = forward("startQueryResultExport"); +export const cancelQueryResultExport = forward("cancelQueryResultExport"); // Redis export const redisListDatabases = forward("redisListDatabases"); @@ -456,5 +458,6 @@ export type { TableExportProgress, TableExportStatus, TableExportRequest, + QueryResultExportRequest, AgentEvent, } from "./tauri"; diff --git a/apps/desktop/src/lib/http.ts b/apps/desktop/src/lib/http.ts index 8aaed4dd93..65c0725131 100644 --- a/apps/desktop/src/lib/http.ts +++ b/apps/desktop/src/lib/http.ts @@ -74,6 +74,7 @@ import type { ExportProgress, TableExportRequest, TableExportProgress, + QueryResultExportRequest, TableCsvExportOptions, XlsxCellValue, QueryPaginationExecutionPlanOptions, @@ -1324,6 +1325,64 @@ export async function cancelTableExport(exportId: string): Promise { return post("/api/export/table/cancel", { exportId }); } +export async function startQueryResultExport(request: QueryResultExportRequest, onProgress: (progress: TableExportProgress) => void): Promise { + const { exportId } = request; + + return new Promise((resolve, reject) => { + let started = false; + let settled = false; + const eventSource = new EventSource(`/api/export/query-result/progress/${exportId}`); + + const finish = (callback: () => void) => { + if (settled) return; + settled = true; + eventSource.close(); + callback(); + }; + + eventSource.onopen = () => { + if (started) return; + started = true; + post("/api/export/query-result", { request }).catch((error) => { + finish(() => reject(error)); + }); + }; + + eventSource.onmessage = (event) => { + const progress: TableExportProgress = JSON.parse(event.data); + onProgress(progress); + if (progress.status === "Done" || progress.status === "Error" || progress.status === "Cancelled") { + if (progress.status === "Error") { + finish(() => reject(new Error(progress.errorMessage || "Export failed"))); + } else if (progress.status === "Done") { + downloadQueryResultExportFile(exportId, request.format); + finish(() => resolve(progress)); + } else { + finish(() => resolve(progress)); + } + } + }; + + eventSource.onerror = () => { + finish(() => reject(new Error("Export progress connection lost"))); + }; + }); +} + +function downloadQueryResultExportFile(exportId: string, format: string): void { + const a = document.createElement("a"); + a.href = `/api/export/query-result/download/${exportId}`; + a.download = `query_result_export_${exportId}.${format}`; + a.click(); +} + +export async function cancelQueryResultExport(exportId: string, executionId?: string): Promise { + return post("/api/export/query-result/cancel", { + exportId, + ...(executionId ? { executionId } : {}), + }); +} + export async function exportQueryResultCsv(filePath: string, columns: string[], rows: readonly (readonly XlsxCellValue[])[]): Promise { const { formatCsv } = await import("./exportFormats"); const content = formatCsv(columns, rows as (string | number | boolean | null)[][]); diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index d01ca94635..17fb5b23d9 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -190,6 +190,7 @@ export interface QueryPaginationExecutionPlanOptions { databaseType?: DatabaseType; pagination: QueryPagination; useAgentCursor: boolean; + firstPageUsesActualSql?: boolean; } export interface QueryPaginationExecutionPlan { @@ -1727,6 +1728,26 @@ export interface TableExportProgress { errorMessage?: string; } +export interface QueryResultExportRequest { + exportId: string; + connectionId: string; + database: string; + schema?: string; + sql: string; + queryBaseSql: string; + databaseType: DatabaseType; + useAgentCursor: boolean; + filePath: string; + format: "csv" | "xlsx"; + pageSize: number; + rowLimit?: number | null; + totalRows?: number | null; + timeoutSecs?: number; + keysetOptimizationEnabled: boolean; + clientSessionId?: string; + executionId?: string; +} + export async function startTableExport(request: TableExportRequest, onProgress: (progress: TableExportProgress) => void): Promise { let unlisten: UnlistenFn | undefined; let settled = false; @@ -1772,6 +1793,51 @@ export async function cancelTableExport(exportId: string): Promise { return invoke("cancel_table_export", { exportId }); } +export async function startQueryResultExport(request: QueryResultExportRequest, onProgress: (progress: TableExportProgress) => void): Promise { + let unlisten: UnlistenFn | undefined; + let settled = false; + let resolveTerminal: (progress: TableExportProgress) => void = () => {}; + let rejectTerminal: (error: unknown) => void = () => {}; + + const terminalProgress = new Promise((resolve, reject) => { + resolveTerminal = resolve; + rejectTerminal = reject; + }); + + const finish = (callback: () => void) => { + if (settled) return; + settled = true; + unlisten?.(); + callback(); + }; + + try { + unlisten = await listen("query-result-export-progress", (event) => { + if (event.payload.exportId !== request.exportId) return; + onProgress(event.payload); + if (event.payload.status === "Done" || event.payload.status === "Error" || event.payload.status === "Cancelled") { + if (event.payload.status === "Error") { + finish(() => rejectTerminal(new Error(event.payload.errorMessage || "Export failed"))); + } else { + finish(() => resolveTerminal(event.payload)); + } + } + }); + await invoke("start_query_result_export", { request }); + return await terminalProgress; + } catch (error) { + if (!settled) { + settled = true; + unlisten?.(); + } + throw error; + } +} + +export async function cancelQueryResultExport(exportId: string, executionId?: string): Promise { + return invoke("cancel_query_result_export", { exportId, executionId: executionId || null }); +} + export async function exportDatabaseSql(request: DatabaseExportRequest, onProgress: (progress: ExportProgress) => void): Promise { const unlisten: UnlistenFn = await listen("database-export-progress", (event) => { if (event.payload.exportId === request.exportId) { diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 27ba637a99..78d019ba31 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -49,6 +49,12 @@ const ACTIVE_TAB_KEY = "dbx-active-tab"; const ORACLE_LIKE_METADATA_TYPES = new Set(["oracle", "dameng", "oceanbase-oracle"]); const BACKGROUND_CLIENT_SESSION_SUFFIXES = ["count", "explain", "export"] as const; +interface BuildQueryResultExportRequestOptions { + exportId: string; + filePath: string; + format: "csv" | "xlsx"; +} + function tabClientSessionId(tab: Pick, suffix?: (typeof BACKGROUND_CLIENT_SESSION_SUFFIXES)[number]): string { return suffix ? `${tab.id}:${suffix}` : tab.id; } @@ -2186,6 +2192,7 @@ export const useQueryStore = defineStore("query", () => { const queryTimeoutSecs = queryTimeoutSecsForConnection(conn); const useAgentCursor = conn?.db_type === "jdbc" || supportsDatabaseFeature(conn?.db_type, "driverManagement"); const queryBaseSql = tab.resultBaseSql ?? sql; + const exportRowLimit = useSettingsStore().editorSettings.exportRowLimit; // 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. const totalRows = typeof tab.resultTotalRowCount === "number" ? tab.resultTotalRowCount : null; @@ -2199,17 +2206,21 @@ export const useQueryStore = defineStore("query", () => { const exportExecutionId = uuid(); try { - while (true) { + while (rows.length < exportRowLimit) { + const remaining = exportRowLimit - rows.length; + const effectivePageLimit = Math.min(pageLimit, remaining); const plan = await api.prepareQueryPaginationExecutionPlan({ sql, queryBaseSql, databaseType: effectiveDbType, - pagination: { limit: pageLimit, offset, sessionId }, + pagination: { limit: effectivePageLimit, offset, sessionId }, useAgentCursor, + firstPageUsesActualSql: true, }); if (typeof plan.pageLimit !== "number" || typeof plan.pageOffset !== "number") return tab.result; const executionOptions = plan.useAgentResultSession ? { + maxRows: exportRowLimit, fetchSize: plan.pageLimit, pageSize: plan.pageLimit, resultSessionId: sessionId, @@ -2226,7 +2237,7 @@ export const useQueryStore = defineStore("query", () => { onProgress?.({ rowsExported: rows.length, totalRows }); sessionId = result.session_id ?? undefined; const shouldFetchNextPage = plan.useAgentResultSession ? result.has_more === true : result.rows.length >= plan.pageLimit; - if (!shouldFetchNextPage) break; + if (!shouldFetchNextPage || rows.length >= exportRowLimit) break; offset += result.rows.length; } } finally { @@ -2244,6 +2255,45 @@ export const useQueryStore = defineStore("query", () => { }; } + async function buildQueryResultExportRequest(id: string, options: BuildQueryResultExportRequestOptions) { + const tab = tabs.value.find((t) => t.id === id); + if (!tab?.result || tab.mode !== "query") return undefined; + + const sql = tab.resultSortedSql ?? tab.resultBaseSql ?? tab.lastExecutedSql ?? tab.sql; + if (!sql.trim()) return undefined; + + const connStore = useConnectionStore(); + await connStore.ensureConnected(tab.connectionId); + const conn = connStore.getConfig(tab.connectionId); + const settings = useSettingsStore().editorSettings; + const effectiveDbType = effectiveDatabaseTypeForConnection(conn); + if (!effectiveDbType) return undefined; + const useAgentCursor = conn?.db_type === "jdbc" || supportsDatabaseFeature(conn?.db_type, "driverManagement"); + const queryBaseSql = tab.resultBaseSql ?? sql; + const totalRows = typeof tab.resultTotalRowCount === "number" ? tab.resultTotalRowCount : null; + const clientSessionId = tabClientSessionId(tab, "export"); + + return { + exportId: options.exportId, + connectionId: tab.connectionId, + database: tab.database, + schema: tab.schema, + sql, + queryBaseSql, + databaseType: effectiveDbType, + useAgentCursor, + filePath: options.filePath, + format: options.format, + pageSize: settings.exportBatchSize, + rowLimit: settings.exportRowLimitEnabled ? settings.exportRowLimit : null, + totalRows, + timeoutSecs: queryTimeoutSecsForConnection(conn), + keysetOptimizationEnabled: settings.queryExportKeysetOptimizationEnabled, + clientSessionId, + executionId: uuid(), + }; + } + return { tabs, activeTabId, @@ -2301,6 +2351,7 @@ export const useQueryStore = defineStore("query", () => { exportResultArchive, importResultArchive, fetchTabResultForExport, + buildQueryResultExportRequest, notifyConnectionMayBeLost, }; }); diff --git a/apps/desktop/src/stores/settingsStore.ts b/apps/desktop/src/stores/settingsStore.ts index b3e93592e0..b23d9a7c8a 100644 --- a/apps/desktop/src/stores/settingsStore.ts +++ b/apps/desktop/src/stores/settingsStore.ts @@ -341,6 +341,9 @@ export interface EditorSettings { customColumnFormatters: Record; snippets: SqlSnippet[]; exportBatchSize: number; + exportRowLimitEnabled: boolean; + exportRowLimit: number; + queryExportKeysetOptimizationEnabled: boolean; toolbarItems: ToolbarItems; } @@ -440,12 +443,17 @@ export const DEFAULT_EDITOR_SETTINGS: EditorSettings = { columnFormatters: {}, customColumnFormatters: {}, snippets: DEFAULT_SQL_SNIPPETS, - exportBatchSize: 10000, + exportBatchSize: 2000, + exportRowLimitEnabled: true, + exportRowLimit: 100000, + queryExportKeysetOptimizationEnabled: true, toolbarItems: { ...DEFAULT_TOOLBAR_ITEMS }, }; export const STORAGE_KEY = "dbx-editor-settings"; const OLD_FONT_SIZE_KEY = "dbx-query-editor-font-size"; +const EXPORT_BATCH_SIZE_DEFAULT_MIGRATION_KEY = "dbx-export-batch-size-default-migrated-v1"; +const LEGACY_DEFAULT_EXPORT_BATCH_SIZE = 10000; const MIN_UI_SCALE = 0.75; const MAX_UI_SCALE = 2; @@ -606,6 +614,9 @@ export function normalizeEditorSettings(settings: Partial, exist customColumnFormatters: normalizeCustomColumnFormatters(settings.customColumnFormatters), snippets: normalizeSqlSnippets(settings.snippets, existing?.snippets), exportBatchSize: typeof settings.exportBatchSize === "number" && settings.exportBatchSize >= 100 && settings.exportBatchSize <= 100000 ? Math.round(settings.exportBatchSize) : DEFAULT_EDITOR_SETTINGS.exportBatchSize, + exportRowLimitEnabled: typeof settings.exportRowLimitEnabled === "boolean" ? settings.exportRowLimitEnabled : DEFAULT_EDITOR_SETTINGS.exportRowLimitEnabled, + exportRowLimit: typeof settings.exportRowLimit === "number" && settings.exportRowLimit >= 100 && settings.exportRowLimit <= 2147483647 ? Math.round(settings.exportRowLimit) : DEFAULT_EDITOR_SETTINGS.exportRowLimit, + queryExportKeysetOptimizationEnabled: typeof settings.queryExportKeysetOptimizationEnabled === "boolean" ? settings.queryExportKeysetOptimizationEnabled : DEFAULT_EDITOR_SETTINGS.queryExportKeysetOptimizationEnabled, toolbarItems: normalizeToolbarItems(settings.toolbarItems), }; } @@ -616,6 +627,13 @@ function loadEditorSettings(): EditorSettings { const raw = localStorage.getItem(STORAGE_KEY); if (raw) { const parsed = JSON.parse(raw) as Partial; + if (parsed.exportBatchSize === LEGACY_DEFAULT_EXPORT_BATCH_SIZE && localStorage.getItem(EXPORT_BATCH_SIZE_DEFAULT_MIGRATION_KEY) !== "1") { + parsed.exportBatchSize = DEFAULT_EDITOR_SETTINGS.exportBatchSize; + localStorage.setItem(EXPORT_BATCH_SIZE_DEFAULT_MIGRATION_KEY, "1"); + const migrated = normalizeEditorSettings(parsed); + saveEditorSettings(migrated); + return migrated; + } return normalizeEditorSettings(parsed); } } catch { @@ -767,6 +785,9 @@ export const useSettingsStore = defineStore("settings", () => { if (partial.customColumnFormatters !== undefined) editorSettings.value.customColumnFormatters = partial.customColumnFormatters; if (partial.snippets !== undefined) editorSettings.value.snippets = normalizeSqlSnippets(partial.snippets); if (partial.exportBatchSize !== undefined) editorSettings.value.exportBatchSize = Math.min(100000, Math.max(100, Math.round(partial.exportBatchSize))); + if (partial.exportRowLimitEnabled !== undefined) editorSettings.value.exportRowLimitEnabled = partial.exportRowLimitEnabled; + if (partial.exportRowLimit !== undefined) editorSettings.value.exportRowLimit = Math.min(2147483647, Math.max(100, Math.round(partial.exportRowLimit))); + if (partial.queryExportKeysetOptimizationEnabled !== undefined) editorSettings.value.queryExportKeysetOptimizationEnabled = partial.queryExportKeysetOptimizationEnabled; if (partial.toolbarItems !== undefined) editorSettings.value.toolbarItems = normalizeToolbarItems(partial.toolbarItems); saveEditorSettings(editorSettings.value); } diff --git a/crates/dbx-core/src/csv_export.rs b/crates/dbx-core/src/csv_export.rs index 9800d49e12..d15afbfcb3 100644 --- a/crates/dbx-core/src/csv_export.rs +++ b/crates/dbx-core/src/csv_export.rs @@ -41,13 +41,25 @@ pub(crate) fn value_to_csv_text(value: &Value) -> String { } } -fn value_to_query_result_csv_text(value: &Value) -> String { +pub(crate) fn value_to_query_result_csv_text(value: &Value) -> String { match value { Value::Null => "NULL".to_string(), other => value_to_csv_text(other), } } +/// Format query-result rows as CSV text without a header row, using the +/// query-result NULL semantics (NULL → "NULL"). Used by the streaming +/// query-result export for batches after the first. +pub fn format_query_result_csv_rows(rows: &[Vec]) -> String { + rows.iter() + .map(|row| { + row.iter().map(|cell| escape_csv(&value_to_query_result_csv_text(cell))).collect::>().join(",") + }) + .collect::>() + .join("\n") +} + fn format_csv_with_value_formatter( columns: &[String], rows: &[Vec], diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index fc6f63042d..c74c4801a4 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -36,6 +36,7 @@ pub mod plugins; pub mod query; pub mod query_cancel; pub mod query_execution_sql; +pub mod query_result_export; pub mod query_result_sql; pub mod redis_ops; pub mod saved_sql; diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs new file mode 100644 index 0000000000..9e54073326 --- /dev/null +++ b/crates/dbx-core/src/query_result_export.rs @@ -0,0 +1,621 @@ +use serde::{Deserialize, Serialize}; +use std::fs::File; +use std::io::{BufWriter, Write}; + +use crate::connection::AppState; +use crate::csv_export::{format_query_result_csv, format_query_result_csv_rows}; +use crate::database_export::is_export_cancelled; +pub use crate::database_export::ExportStatus; +use crate::models::connection::DatabaseType; +use crate::query::{close_query_session, execute_sql_statement_with_options, QueryExecutionOptions, QUERY_CANCELED}; +use crate::query_result_sql::{ + build_query_pagination_execution_plan, QueryPagination, QueryPaginationExecutionPlanOptions, +}; +use crate::table_export::TableExportProgress; +use crate::transfer::keyset_pagination_sql; +use crate::xlsx_export::{finish_streaming_xlsx_workbook, start_streaming_xlsx_workbook}; +use serde_json::Value; +use sqlparser::ast::{GroupByExpr, ObjectNamePart, OrderByKind, SelectItem, SetExpr, Statement, TableFactor}; +use sqlparser::dialect::GenericDialect; +use sqlparser::parser::Parser; +use tokio_util::sync::CancellationToken; + +const AGENT_UNBOUNDED_ROW_LIMIT: usize = i32::MAX as usize; +pub const XLSX_MAX_DATA_ROWS: usize = 1_048_575; +const XLSX_ROW_LIMIT_ERROR: &str = "XLSX 最多支持 1,048,575 行数据，请改用 CSV 导出完整结果。"; +const STREAMING_PAGINATION_UNSUPPORTED_ERROR: &str = "当前查询暂不支持流式导出，请简化查询或使用受支持的驱动。"; +const AGENT_SESSION_MISSING_ERROR: &str = "查询结果流式导出需要驱动返回结果集会话，但当前驱动未返回 session_id。"; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct QueryResultExportRequest { + pub export_id: String, + pub connection_id: String, + pub database: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub schema: Option, + pub sql: String, + pub query_base_sql: String, + pub database_type: DatabaseType, + #[serde(default)] + pub use_agent_cursor: bool, + pub file_path: String, + pub format: String, + pub page_size: usize, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub row_limit: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub total_rows: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub timeout_secs: Option, + #[serde(default)] + pub keyset_optimization_enabled: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_session_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub execution_id: Option, +} + +fn progress( + request: &QueryResultExportRequest, + rows_exported: u64, + status: ExportStatus, + error_message: Option, +) -> TableExportProgress { + let total_rows = request.total_rows.map(|total| { + let format = request.format.to_lowercase(); + let limit = effective_row_limit(&format, request); + limit.map_or(total, |limit| total.min(limit as u64)) + }); + TableExportProgress { + export_id: request.export_id.clone(), + table_name: String::new(), + rows_exported, + total_rows, + status, + error_message, + } +} + +fn effective_row_limit(format: &str, request: &QueryResultExportRequest) -> Option { + if format == "xlsx" { + Some(request.row_limit.map_or(XLSX_MAX_DATA_ROWS, |limit| limit.min(XLSX_MAX_DATA_ROWS))) + } else { + request.row_limit + } +} + +fn xlsx_hard_limit_active(format: &str, request: &QueryResultExportRequest) -> bool { + format == "xlsx" && request.row_limit.map_or(true, |limit| limit > XLSX_MAX_DATA_ROWS) +} + +fn should_fetch_next_page( + use_agent_result_session: bool, + has_more: bool, + fetched_row_count: usize, + written_row_count: usize, + requested_page_size: usize, +) -> bool { + if use_agent_result_session { + has_more + } else { + fetched_row_count > written_row_count || written_row_count >= requested_page_size + } +} + +fn supports_streaming_offset_pagination(request: &QueryResultExportRequest, page_size: usize) -> bool { + let first_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: request.sql.clone(), + query_base_sql: request.query_base_sql.clone(), + database_type: Some(request.database_type), + pagination: QueryPagination { limit: page_size, offset: 0, session_id: None }, + use_agent_cursor: false, + first_page_uses_actual_sql: true, + }); + let second_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: request.sql.clone(), + query_base_sql: request.query_base_sql.clone(), + database_type: Some(request.database_type), + pagination: QueryPagination { limit: page_size, offset: page_size, session_id: None }, + use_agent_cursor: false, + first_page_uses_actual_sql: true, + }); + + let (Some(first_sql), Some(second_sql)) = (first_page.page_sql.as_deref(), second_page.page_sql.as_deref()) else { + return false; + }; + first_page.page_limit.is_some() + && second_page.page_limit.is_some() + && !first_sql.trim().eq_ignore_ascii_case(second_sql.trim()) +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct SafeKeysetCandidate { + schema: Option, + table: String, +} + +struct KeysetPlan { + columns: Vec, + primary_keys: Vec, + pk_indices: Vec, + schema: String, + table: String, + last_pk_values: Vec, +} + +fn object_name_parts(name: &sqlparser::ast::ObjectName) -> Option> { + name.0 + .iter() + .map(|part| match part { + ObjectNamePart::Identifier(ident) => Some(ident.value.clone()), + _ => None, + }) + .collect() +} + +fn safe_keyset_candidate(sql: &str) -> Option { + let dialect = GenericDialect {}; + let statements = Parser::parse_sql(&dialect, sql).ok()?; + let [Statement::Query(query)] = statements.as_slice() else { + return None; + }; + if query.with.is_some() + || query.limit_clause.is_some() + || query.fetch.is_some() + || !query.locks.is_empty() + || query.for_clause.is_some() + || query.settings.is_some() + || query.format_clause.is_some() + || !query.pipe_operators.is_empty() + || query + .order_by + .as_ref() + .is_some_and(|order_by| !matches!(&order_by.kind, OrderByKind::Expressions(exprs) if exprs.is_empty())) + { + return None; + } + let SetExpr::Select(select) = query.body.as_ref() else { + return None; + }; + if select.distinct.is_some() + || !matches!(&select.group_by, GroupByExpr::Expressions(exprs, _) if exprs.is_empty()) + || select.having.is_some() + || select.selection.is_some() + || select.from.len() != 1 + || !(select.projection.len() == 1 && matches!(select.projection.first(), Some(SelectItem::Wildcard(_)))) + { + return None; + } + let source = &select.from[0]; + if !source.joins.is_empty() { + return None; + } + let TableFactor::Table { name, .. } = &source.relation else { + return None; + }; + let parts = object_name_parts(name)?; + match parts.as_slice() { + [table] => Some(SafeKeysetCandidate { schema: None, table: table.clone() }), + [schema, table] => Some(SafeKeysetCandidate { schema: Some(schema.clone()), table: table.clone() }), + _ => None, + } +} + +async fn build_keyset_plan(state: &AppState, request: &QueryResultExportRequest) -> Option { + if !request.keyset_optimization_enabled || request.use_agent_cursor { + return None; + } + let candidate = safe_keyset_candidate(&request.sql)?; + let schema = candidate.schema.or_else(|| request.schema.clone()).unwrap_or_default(); + let columns = + crate::schema::get_columns_core(state, &request.connection_id, &request.database, &schema, &candidate.table) + .await + .ok()?; + let col_names: Vec = columns.iter().map(|column| column.name.clone()).collect(); + let primary_keys: Vec = + columns.iter().filter(|column| column.is_primary_key).map(|column| column.name.clone()).collect(); + if col_names.is_empty() || primary_keys.is_empty() { + return None; + } + let pk_indices: Vec = primary_keys + .iter() + .filter_map(|pk| col_names.iter().position(|column| column.eq_ignore_ascii_case(pk))) + .collect(); + if pk_indices.len() != primary_keys.len() { + return None; + } + Some(KeysetPlan { + columns: col_names, + primary_keys, + pk_indices, + schema, + table: candidate.table, + last_pk_values: Vec::new(), + }) +} + +pub async fn export_query_result_core( + state: &AppState, + request: &QueryResultExportRequest, + cancel_token: Option, + on_progress: impl Fn(TableExportProgress), +) -> Result<(), String> { + let mut session_id: Option = None; + let result = export_query_result_core_inner(state, request, cancel_token, &on_progress, &mut session_id).await; + + if let Some(session_id) = session_id { + let _ = close_query_session( + state, + &request.connection_id, + &request.database, + &session_id, + request.client_session_id.as_deref(), + ) + .await; + } + if let Some(client_session_id) = request.client_session_id.as_deref() { + let database = request.database.trim(); + let database = if database.is_empty() { None } else { Some(database) }; + let _ = state.close_client_session_pool(&request.connection_id, database, client_session_id).await; + } + + result +} + +async fn export_query_result_core_inner( + state: &AppState, + request: &QueryResultExportRequest, + cancel_token: Option, + on_progress: &impl Fn(TableExportProgress), + session_id: &mut Option, +) -> Result<(), String> { + let format = request.format.to_lowercase(); + if format != "csv" && format != "xlsx" { + return Err(format!("Unsupported streaming query-result export format: {format}")); + } + + let page_size = request.page_size.max(1); + let effective_row_limit = effective_row_limit(&format, request); + let xlsx_hard_limit_active = xlsx_hard_limit_active(&format, request); + if xlsx_hard_limit_active && request.total_rows.is_some_and(|total| total > XLSX_MAX_DATA_ROWS as u64) { + return Err(XLSX_ROW_LIMIT_ERROR.to_string()); + } + + let agent_max_rows = if xlsx_hard_limit_active { + XLSX_MAX_DATA_ROWS + 1 + } else { + effective_row_limit.unwrap_or(AGENT_UNBOUNDED_ROW_LIMIT) + } + .max(1); + + on_progress(progress(request, 0, ExportStatus::Running, None)); + + let mut csv_file = if format == "csv" { + Some(BufWriter::new(File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?)) + } else { + None + }; + if let Some(file) = csv_file.as_mut() { + file.write_all(b"\xEF\xBB\xBF").map_err(|e| format!("Failed to write BOM: {e}"))?; + } + + let mut xlsx = None; + let mut columns: Vec = Vec::new(); + let mut rows_exported: u64 = 0; + let mut offset: usize = 0; + let mut wrote_csv_header = false; + let mut keyset_plan = build_keyset_plan(state, request).await; + if keyset_plan.is_none() && !request.use_agent_cursor && !supports_streaming_offset_pagination(request, page_size) { + return Err(STREAMING_PAGINATION_UNSUPPORTED_ERROR.to_string()); + } + + loop { + if cancel_token.as_ref().is_some_and(|token| token.is_cancelled()) + || is_export_cancelled(&request.export_id).await + { + on_progress(progress( + request, + rows_exported, + ExportStatus::Cancelled, + Some("Export cancelled".to_string()), + )); + return Ok(()); + } + + let remaining = effective_row_limit.map(|limit| limit.saturating_sub(rows_exported as usize)); + if matches!(remaining, Some(0)) { + break; + } + let this_page = remaining.map_or(page_size, |rem| rem.min(page_size)).max(1); + let fetch_limit = if xlsx_hard_limit_active && remaining.is_some_and(|rem| rem <= page_size) { + this_page.saturating_add(1) + } else { + this_page + }; + + let (sql_to_execute, plan_limit, use_agent_result_session) = if let Some(plan) = keyset_plan.as_ref() { + ( + keyset_pagination_sql( + &plan.columns, + &plan.table, + &plan.schema, + &request.database_type, + &plan.primary_keys, + &plan.last_pk_values, + fetch_limit, + ), + fetch_limit, + false, + ) + } else { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: request.sql.clone(), + query_base_sql: request.query_base_sql.clone(), + database_type: Some(request.database_type), + pagination: QueryPagination { limit: fetch_limit, offset, session_id: session_id.clone() }, + use_agent_cursor: request.use_agent_cursor, + first_page_uses_actual_sql: true, + }); + let Some(plan_limit) = plan.page_limit else { + return Err("Failed to build query pagination plan for export".to_string()); + }; + (plan.sql_to_execute, plan_limit, plan.use_agent_result_session) + }; + + let options = if use_agent_result_session { + QueryExecutionOptions { + page_size: Some(plan_limit), + fetch_size: Some(plan_limit), + result_session_id: session_id.clone(), + max_rows: Some(agent_max_rows), + timeout_secs: request.timeout_secs, + client_session_id: request.client_session_id.clone(), + execution_id: request.execution_id.clone(), + } + } else { + QueryExecutionOptions { + max_rows: Some(plan_limit), + fetch_size: Some(plan_limit), + timeout_secs: request.timeout_secs, + client_session_id: request.client_session_id.clone(), + execution_id: request.execution_id.clone(), + ..Default::default() + } + }; + + let mut result = match execute_sql_statement_with_options( + state, + &request.connection_id, + &request.database, + &sql_to_execute, + request.schema.as_deref(), + cancel_token.clone(), + options, + ) + .await + { + Ok(result) => result, + Err(error) => { + if error == QUERY_CANCELED + || cancel_token.as_ref().is_some_and(|token| token.is_cancelled()) + || is_export_cancelled(&request.export_id).await + { + on_progress(progress( + request, + rows_exported, + ExportStatus::Cancelled, + Some("Export cancelled".to_string()), + )); + return Ok(()); + } + return Err(error); + } + }; + + if cancel_token.as_ref().is_some_and(|token| token.is_cancelled()) + || is_export_cancelled(&request.export_id).await + { + on_progress(progress( + request, + rows_exported, + ExportStatus::Cancelled, + Some("Export cancelled".to_string()), + )); + return Ok(()); + } + + if columns.is_empty() { + columns = result.columns.clone(); + } + let fetched_row_count = result.rows.len(); + if xlsx_hard_limit_active { + let remaining_rows = XLSX_MAX_DATA_ROWS.saturating_sub(rows_exported as usize); + if fetched_row_count > remaining_rows { + return Err(XLSX_ROW_LIMIT_ERROR.to_string()); + } + } + if result.rows.len() > this_page { + result.rows.truncate(this_page); + } + let row_count = result.rows.len(); + + if format == "csv" { + if let Some(file) = csv_file.as_mut() { + if !wrote_csv_header { + let csv = format_query_result_csv(&columns, &result.rows); + file.write_all(csv.as_bytes()).map_err(|e| format!("Failed to write CSV: {e}"))?; + wrote_csv_header = true; + } else if row_count > 0 { + let rows_csv = format_query_result_csv_rows(&result.rows); + write!(file, "\n{rows_csv}").map_err(|e| format!("Failed to write CSV rows: {e}"))?; + } + } + } else { + if xlsx.is_none() { + let xlsx_file = + File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; + xlsx = Some(start_streaming_xlsx_workbook(BufWriter::new(xlsx_file), Some("Result"), &columns)?); + } + if let Some(writer) = xlsx.as_mut() { + for row in &result.rows { + writer.write_row(row).map_err(|e| format!("Failed to write XLSX row: {e}"))?; + } + } + } + + rows_exported += row_count as u64; + on_progress(progress(request, rows_exported, ExportStatus::Running, None)); + + if result.session_id.is_some() { + *session_id = result.session_id.clone(); + } + if use_agent_result_session && result.has_more && session_id.is_none() { + return Err(AGENT_SESSION_MISSING_ERROR.to_string()); + } + if let Some(plan) = keyset_plan.as_mut() { + if let Some(last_row) = result.rows.last() { + plan.last_pk_values = + plan.pk_indices.iter().map(|&index| last_row.get(index).cloned().unwrap_or(Value::Null)).collect(); + } + } + let should_continue = + should_fetch_next_page(use_agent_result_session, result.has_more, fetched_row_count, row_count, plan_limit); + if cancel_token.as_ref().is_some_and(|token| token.is_cancelled()) + || is_export_cancelled(&request.export_id).await + { + on_progress(progress( + request, + rows_exported, + ExportStatus::Cancelled, + Some("Export cancelled".to_string()), + )); + return Ok(()); + } + if !should_continue || row_count == 0 { + break; + } + offset += row_count; + } + + on_progress(progress(request, rows_exported, ExportStatus::Writing, None)); + + if format == "csv" { + if !wrote_csv_header { + let csv = format_query_result_csv(&columns, &[]); + if let Some(file) = csv_file.as_mut() { + file.write_all(csv.as_bytes()).map_err(|e| format!("Failed to write CSV: {e}"))?; + } + } + if let Some(file) = csv_file.as_mut() { + file.flush().map_err(|e| format!("Failed to flush CSV file: {e}"))?; + } + } else if let Some(writer) = xlsx { + let mut buf = + finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; + buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; + } else { + let xlsx_file = File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; + let writer = start_streaming_xlsx_workbook(BufWriter::new(xlsx_file), Some("Result"), &columns)?; + let mut buf = + finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; + buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; + } + + on_progress(progress(request, rows_exported, ExportStatus::Done, None)); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn request(format: &str, row_limit: Option, total_rows: Option) -> QueryResultExportRequest { + QueryResultExportRequest { + export_id: "export-1".to_string(), + connection_id: "conn-1".to_string(), + database: "db".to_string(), + schema: None, + sql: "SELECT * FROM users".to_string(), + query_base_sql: "SELECT * FROM users".to_string(), + database_type: DatabaseType::Postgres, + use_agent_cursor: false, + file_path: "out.csv".to_string(), + format: format.to_string(), + page_size: 1000, + row_limit, + total_rows, + timeout_secs: None, + keyset_optimization_enabled: true, + client_session_id: None, + execution_id: None, + } + } + + #[test] + fn csv_unlimited_export_has_no_effective_row_limit() { + assert_eq!(effective_row_limit("csv", &request("csv", None, None)), None); + } + + #[test] + fn xlsx_unlimited_export_uses_excel_hard_limit() { + assert_eq!(effective_row_limit("xlsx", &request("xlsx", None, None)), Some(XLSX_MAX_DATA_ROWS)); + } + + #[test] + fn xlsx_row_limit_caps_to_excel_hard_limit() { + assert_eq!( + effective_row_limit("xlsx", &request("xlsx", Some(XLSX_MAX_DATA_ROWS + 10), None)), + Some(XLSX_MAX_DATA_ROWS) + ); + } + + #[test] + fn xlsx_known_total_above_hard_limit_errors_before_export() { + let req = request("xlsx", None, Some(XLSX_MAX_DATA_ROWS as u64 + 1)); + assert!(xlsx_hard_limit_active("xlsx", &req)); + assert!(req.total_rows.is_some_and(|total| total > XLSX_MAX_DATA_ROWS as u64)); + } + + #[test] + fn non_agent_pages_continue_after_trimming_probe_row() { + assert!(should_fetch_next_page(false, false, 101, 100, 100)); + assert!(should_fetch_next_page(false, false, 100, 100, 100)); + assert!(!should_fetch_next_page(false, false, 42, 42, 100)); + } + + #[test] + fn agent_pages_follow_has_more_flag() { + assert!(should_fetch_next_page(true, true, 42, 42, 100)); + assert!(!should_fetch_next_page(true, false, 100, 100, 100)); + } + + #[test] + fn streaming_offset_pagination_requires_distinct_followup_page_sql() { + let req = request("csv", Some(1000), None); + assert!(supports_streaming_offset_pagination(&req, 100)); + + let oracle_req = + QueryResultExportRequest { database_type: DatabaseType::Oracle, ..request("csv", Some(1000), None) }; + assert!(!supports_streaming_offset_pagination(&oracle_req, 100)); + } + + #[test] + fn keyset_candidate_accepts_simple_single_table_wildcard_query() { + let candidate = safe_keyset_candidate("SELECT * FROM public.users").expect("safe keyset candidate"); + assert_eq!(candidate.schema.as_deref(), Some("public")); + assert_eq!(candidate.table, "users"); + } + + #[test] + fn keyset_candidate_rejects_join_and_sorted_queries() { + assert!(safe_keyset_candidate("SELECT * FROM users u JOIN orders o ON o.user_id = u.id").is_none()); + assert!(safe_keyset_candidate("SELECT * FROM users ORDER BY name").is_none()); + } + + #[test] + fn keyset_candidate_rejects_filters_and_projection_changes() { + assert!(safe_keyset_candidate("SELECT * FROM users WHERE active = true").is_none()); + assert!(safe_keyset_candidate("SELECT id, name FROM users").is_none()); + } +} diff --git a/crates/dbx-core/src/query_result_sql.rs b/crates/dbx-core/src/query_result_sql.rs index 7069f5acfd..9abd634363 100644 --- a/crates/dbx-core/src/query_result_sql.rs +++ b/crates/dbx-core/src/query_result_sql.rs @@ -32,6 +32,8 @@ pub struct QueryPaginationExecutionPlanOptions { pub database_type: Option, pub pagination: QueryPagination, pub use_agent_cursor: bool, + #[serde(default)] + pub first_page_uses_actual_sql: bool, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -128,7 +130,9 @@ pub fn build_query_pagination_execution_plan( } if options.use_agent_cursor && options.pagination.offset == 0 { - plan.sql_to_execute = options.query_base_sql; + if !options.first_page_uses_actual_sql { + plan.sql_to_execute = options.query_base_sql; + } plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); plan.use_agent_result_session = true; @@ -922,6 +926,7 @@ mod tests { database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, + first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, sql); @@ -1282,6 +1287,7 @@ WHERE u.id = picked.id; database_type: Some(DatabaseType::Oracle), pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, use_agent_cursor: true, + first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, "SELECT * FROM events"); @@ -1291,6 +1297,24 @@ WHERE u.id = picked.id; assert!(plan.use_agent_result_session); } + #[test] + fn export_agent_cursor_first_page_keeps_actual_sorted_sql() { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT * FROM (SELECT * FROM events) t ORDER BY created_at DESC".to_string(), + query_base_sql: "SELECT * FROM events".to_string(), + database_type: Some(DatabaseType::Oracle), + pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: true, + }); + + assert_eq!(plan.sql_to_execute, "SELECT * FROM (SELECT * FROM events) t ORDER BY created_at DESC"); + assert_eq!(plan.page_limit, Some(500)); + assert_eq!(plan.page_offset, Some(0)); + assert!(plan.page_sql.is_none()); + assert!(plan.use_agent_result_session); + } + #[test] fn builds_sorted_query_sql() { let result = build_sorted_query_sql(SortedQuerySqlOptions { diff --git a/crates/dbx-core/src/xlsx_export.rs b/crates/dbx-core/src/xlsx_export.rs index 015c966395..1df13a15d9 100644 --- a/crates/dbx-core/src/xlsx_export.rs +++ b/crates/dbx-core/src/xlsx_export.rs @@ -59,9 +59,15 @@ pub(crate) fn data_row_xml(row_number: usize, columns: &[String], row: &[Value]) format!("{cells}") } +/// Shared ZIP entry options for all XLSX parts. XLSX is a ZIP of XML, and the +/// `inlineStr` cell encoding is highly repetitive, so Deflate typically shrinks +/// the file several-fold over `Stored` (matching what Excel/Navicat produce). +fn xlsx_zip_options() -> zip::write::SimpleFileOptions { + zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Deflated) +} + fn write_zip_entry(zip: &mut zip::ZipWriter, path: &str, content: &str) -> Result<(), String> { - let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored); - zip.start_file(path, options).map_err(|err| err.to_string())?; + zip.start_file(path, xlsx_zip_options()).map_err(|err| err.to_string())?; zip.write_all(content.as_bytes()).map_err(|err| err.to_string()) } @@ -86,7 +92,7 @@ pub(crate) fn start_streaming_xlsx_workbook( // Begin the sheet1.xml entry with header, frozen pane, column widths and // the header row. - let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored); + let options = xlsx_zip_options(); zip.start_file("xl/worksheets/sheet1.xml", options).map_err(|err| err.to_string())?; let sheet_header = format!( @@ -434,7 +440,7 @@ pub fn build_xlsx_workbook_multi(sheets: &[XlsxWorksheetData]) -> Result let cursor = Cursor::new(Vec::::new()); let mut zip = zip::ZipWriter::new(cursor); - let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored); + let options = xlsx_zip_options(); for (path, content) in files { zip.start_file(path, options).map_err(|err| err.to_string())?; @@ -455,6 +461,31 @@ mod tests { use calamine::{open_workbook_auto, Reader}; use serde_json::json; use std::fs; + use std::io::Read; + + /// Read and decompress a single entry from an in-memory XLSX (ZIP) buffer. + fn read_zip_entry(bytes: &[u8], path: &str) -> String { + let mut archive = zip::ZipArchive::new(std::io::Cursor::new(bytes.to_vec())).expect("open xlsx as zip archive"); + let mut entry = archive.by_name(path).unwrap_or_else(|_| panic!("missing zip entry: {path}")); + let mut content = String::new(); + entry.read_to_string(&mut content).expect("read zip entry"); + content + } + + /// Assert every entry in the XLSX (ZIP) buffer is Deflate-compressed, which + /// is what keeps exported workbooks small (see `xlsx_zip_options`). + fn assert_all_entries_deflated(bytes: &[u8]) { + let mut archive = zip::ZipArchive::new(std::io::Cursor::new(bytes.to_vec())).expect("open xlsx as zip archive"); + for index in 0..archive.len() { + let entry = archive.by_index(index).expect("zip entry"); + assert_eq!( + entry.compression(), + zip::CompressionMethod::Deflated, + "entry {} should be Deflate-compressed", + entry.name() + ); + } + } #[test] fn builds_xlsx_zip_with_sheet_data() { @@ -464,16 +495,19 @@ mod tests { rows: vec![vec![json!(1), json!("Ada & Bob"), json!(true)], vec![json!(2), json!(null), json!(false)]], }) .expect("build workbook"); - let text = String::from_utf8_lossy(&workbook); + // ZIP magic bytes. assert_eq!(workbook[0], 0x50); assert_eq!(workbook[1], 0x4b); - assert!(text.contains("[Content_Types].xml")); - assert!(text.contains("xl/worksheets/sheet1.xml")); - assert!(text.contains("name=\"Users\"")); - assert!(text.contains("1")); - assert!(text.contains("Ada & Bob")); - assert!(text.contains("1")); + + // Entries are stored compressed; assert on their decompressed contents. + let sheet = read_zip_entry(&workbook, "xl/worksheets/sheet1.xml"); + let workbook_xml = read_zip_entry(&workbook, "xl/workbook.xml"); + assert!(workbook_xml.contains("name=\"Users\"")); + assert!(sheet.contains("1")); + assert!(sheet.contains("Ada & Bob")); + assert!(sheet.contains("1")); + assert_all_entries_deflated(&workbook); } #[test] @@ -484,8 +518,8 @@ mod tests { rows: vec![vec![json!("ok")]], }) .expect("build workbook"); - let text = String::from_utf8_lossy(&workbook); - assert!(text.contains("name=\"bad name with chars and-a-very-\"")); + let workbook_xml = read_zip_entry(&workbook, "xl/workbook.xml"); + assert!(workbook_xml.contains("name=\"bad name with chars and-a-very-\"")); } #[test] diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 223e1cea96..cfecc8656f 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -404,6 +404,17 @@ async fn main() { .route("/export/table/progress/{exportId}", get(routes::table_export::table_export_progress)) .route("/export/table/download/{exportId}", get(routes::table_export::table_export_download)) .route("/export/table/cancel", post(routes::table_export::cancel_table_export)) + // Query result export + .route("/export/query-result", post(routes::query_result_export::start_query_result_export)) + .route( + "/export/query-result/progress/{exportId}", + get(routes::query_result_export::query_result_export_progress), + ) + .route( + "/export/query-result/download/{exportId}", + get(routes::query_result_export::query_result_export_download), + ) + .route("/export/query-result/cancel", post(routes::query_result_export::cancel_query_result_export)) // SQL file .route("/sql-file/preview", post(routes::sql_file::preview_sql_file)) .route("/sql-file/execute", post(routes::sql_file::execute_sql_file)) diff --git a/crates/dbx-web/src/routes/mod.rs b/crates/dbx-web/src/routes/mod.rs index df3ec58381..cd5c332d08 100644 --- a/crates/dbx-web/src/routes/mod.rs +++ b/crates/dbx-web/src/routes/mod.rs @@ -15,6 +15,7 @@ pub mod mq; pub mod nacos; pub mod plugins; pub mod query; +pub mod query_result_export; pub mod redis; pub mod redis_pubsub_ws; pub mod saved_sql; diff --git a/crates/dbx-web/src/routes/query_result_export.rs b/crates/dbx-web/src/routes/query_result_export.rs new file mode 100644 index 0000000000..557977949d --- /dev/null +++ b/crates/dbx-web/src/routes/query_result_export.rs @@ -0,0 +1,168 @@ +use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + time::Duration, +}; + +use axum::body::Body; +use axum::extract::{Path, State}; +use axum::http::{header, StatusCode}; +use axum::response::{Response, Sse}; +use axum::Json; +use dbx_core::query_cancel::RunningTaskMetadata; +use dbx_core::query_result_export::{self, QueryResultExportRequest}; +use dbx_core::table_export::{ExportStatus, TableExportProgress}; +use futures::stream::Stream; +use serde::Deserialize; + +use crate::error::AppError; +use crate::state::WebState; + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct StartQueryResultExportRequest { + pub request: QueryResultExportRequest, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CancelQueryResultExportRequest { + pub export_id: String, + pub execution_id: Option, +} + +pub async fn start_query_result_export( + State(state): State>, + Json(body): Json, +) -> Result, AppError> { + let mut req = body.request; + let export_id = req.export_id.clone(); + + let ext = match req.format.as_str() { + "csv" => "csv", + "xlsx" => "xlsx", + _ => return Err(AppError(format!("Unsupported query result export format: {}", req.format))), + }; + let tmp_dir = state.data_dir.join("tmp"); + std::fs::create_dir_all(&tmp_dir).map_err(|e| AppError(e.to_string()))?; + let tmp_file = tmp_dir.join(format!("query_result_export_{export_id}.{ext}")); + let file_path = tmp_file.to_string_lossy().to_string(); + req.file_path = file_path.clone(); + + state.export_files.write().await.insert(export_id.clone(), (file_path, req.format.clone())); + + let tx = { + let mut channels = state.sse_channels.write().await; + channels.entry(export_id.clone()).or_insert_with(|| tokio::sync::broadcast::channel::(256).0).clone() + }; + + let app = state.app.clone(); + let state_clone = state.clone(); + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_progress = cancelled.clone(); + + tokio::spawn(async move { + let execution_id = req.execution_id.clone().filter(|id| !id.trim().is_empty()); + let registered_query = execution_id.as_ref().map(|id| { + app.running_queries.register_task( + id.clone(), + RunningTaskMetadata::query( + req.connection_id.clone(), + req.database.clone(), + req.client_session_id.clone(), + ), + ) + }); + let cancel_token = registered_query.as_ref().map(|query| query.token()); + let result = query_result_export::export_query_result_core(&app, &req, cancel_token, |progress| { + if matches!(progress.status, ExportStatus::Cancelled) { + cancelled_progress.store(true, Ordering::SeqCst); + } + if let Ok(json) = serde_json::to_string(&progress) { + let _ = tx.send(json); + } + }) + .await; + drop(registered_query); + + if let Err(e) = result { + let _ = tokio::fs::remove_file(&req.file_path).await; + state_clone.export_files.write().await.remove(&req.export_id); + let progress = TableExportProgress { + export_id: req.export_id.clone(), + table_name: String::new(), + rows_exported: 0, + total_rows: None, + status: ExportStatus::Error, + error_message: Some(e), + }; + if let Ok(json) = serde_json::to_string(&progress) { + let _ = tx.send(json); + } + } else if cancelled.load(Ordering::SeqCst) { + let _ = tokio::fs::remove_file(&req.file_path).await; + state_clone.export_files.write().await.remove(&req.export_id); + } + + dbx_core::database_export::clear_export_cancelled(&req.export_id).await; + tokio::time::sleep(Duration::from_secs(5)).await; + state_clone.remove_sse_channel(&req.export_id).await; + }); + + Ok(Json(serde_json::json!({ "exportId": export_id }))) +} + +pub async fn query_result_export_progress( + State(state): State>, + Path(export_id): Path, +) -> Result>>, AppError> { + let tx = { + let mut channels = state.sse_channels.write().await; + channels.entry(export_id).or_insert_with(|| tokio::sync::broadcast::channel::(256).0).clone() + }; + let rx = tx.subscribe(); + Ok(crate::sse::sse_from_channel(rx)) +} + +pub async fn cancel_query_result_export( + State(state): State>, + Json(req): Json, +) -> Json { + dbx_core::database_export::set_export_cancelled(&req.export_id).await; + if let Some(execution_id) = req.execution_id.filter(|id| !id.trim().is_empty()) { + state.app.running_queries.cancel(&execution_id); + } + Json(serde_json::json!({ "cancelled": true })) +} + +pub async fn query_result_export_download( + State(state): State>, + Path(export_id): Path, +) -> Result { + let (file_path, format) = state + .export_files + .write() + .await + .remove(&export_id) + .ok_or_else(|| AppError("Export file not found".to_string()))?; + + let data = tokio::fs::read(&file_path).await.map_err(|e| AppError(e.to_string()))?; + let _ = tokio::fs::remove_file(&file_path).await; + + let (content_type, file_ext) = match format.as_str() { + "csv" => ("text/csv; charset=utf-8", "csv"), + "xlsx" => ("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "xlsx"), + _ => return Err(AppError(format!("Unknown format: {format}"))), + }; + + let filename = format!("query_result_export_{export_id}.{file_ext}"); + + Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, content_type) + .header(header::CONTENT_DISPOSITION, format!("attachment; filename=\"{filename}\"")) + .body(Body::from(data)) + .map_err(|e| AppError(e.to_string())) +} diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index b5741c8eeb..d069113303 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -5,6 +5,7 @@ import { isReactive } from "vue"; import { decodeQueryResultArchive } from "../../apps/desktop/src/lib/queryResultArchive.ts"; import { useConnectionStore } from "../../apps/desktop/src/stores/connectionStore.ts"; import { useQueryStore } from "../../apps/desktop/src/stores/queryStore.ts"; +import { useSettingsStore } from "../../apps/desktop/src/stores/settingsStore.ts"; import type { ConnectionConfig } from "../../apps/desktop/src/types/database.ts"; import type { QueryResult } from "../../apps/desktop/src/types/database.ts"; @@ -1686,6 +1687,182 @@ test("table data export fetches every filtered page", async () => { } }); +test("agent-session query export raises maxRows to the configured export limit", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + settingsStore.updateEditorSettings({ exportRowLimit: 50_000 }); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + const executeBodies: any[] = []; + + connectionStore.addEphemeralConnection({ + ...conn("jdbc-1"), + db_type: "jdbc", + connection_string: "jdbc:sqlserver://127.0.0.1:1433;databaseName=db", + jdbc_driver_class: "com.microsoft.sqlserver.jdbc.SQLServerDriver", + }); + const tabId = store.createTab("jdbc-1", "db", "Query", "query", "dbo"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + tab.lastExecutedSql = "SELECT * FROM big_table"; + tab.result = { + columns: ["id"], + rows: [[1]], + affected_rows: 0, + execution_time_ms: 1, + truncated: false, + has_more: true, + }; + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + return new Response(JSON.stringify({ sqlToExecute: "SELECT * FROM big_table", pageLimit: 10_000, pageOffset: 0, useAgentResultSession: true }), { status: 200, headers: { "Content-Type": "application/json" } }); + } + if (url === "/api/query/execute-multi") { + const body = JSON.parse(String(init?.body ?? "{}")); + executeBodies.push(body); + const page = executeBodies.length; + return new Response( + JSON.stringify([ + { + columns: ["id"], + rows: Array.from({ length: 10_000 }, (_, index) => [(page - 1) * 10_000 + index + 1]), + affected_rows: 0, + execution_time_ms: 1, + session_id: "session-1", + has_more: page < 2, + }, + ]), + { status: 200, headers: { "Content-Type": "application/json" } }, + ); + } + if (url === "/api/query/close-session") { + return new Response(JSON.stringify(true), { status: 200, headers: { "Content-Type": "application/json" } }); + } + if (url === "/api/query/close-client-session") { + return new Response(JSON.stringify(true), { status: 200, headers: { "Content-Type": "application/json" } }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const exported = await store.fetchTabResultForExport(tabId); + + // Every agent-session execute call must carry the configured export limit + // as maxRows so the agent's cumulative cap doesn't truncate at 10000. + assert.ok(executeBodies.length >= 2); + for (const body of executeBodies) { + assert.equal(body.maxRows, 50_000); + } + assert.equal(exported?.rows.length, 20_000); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("buildQueryResultExportRequest uses sorted SQL and independent row-limit settings", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + settingsStore.updateEditorSettings({ + exportBatchSize: 2500, + exportRowLimit: 200000, + exportRowLimitEnabled: false, + queryExportKeysetOptimizationEnabled: false, + }); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + + connectionStore.addEphemeralConnection(conn("conn-1")); + const tabId = store.createTab("conn-1", "analytics", "Query", "query", "public"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + tab.sql = "SELECT * FROM events"; + tab.lastExecutedSql = "SELECT * FROM events"; + tab.resultBaseSql = "SELECT * FROM events"; + tab.resultSortedSql = "SELECT * FROM (SELECT * FROM events) t ORDER BY created_at DESC"; + tab.resultTotalRowCount = 123456; + tab.result = { + columns: ["id", "created_at"], + rows: [[1, "2026-06-24"]], + affected_rows: 0, + execution_time_ms: 1, + }; + + globalThis.fetch = withConnectionHealthMock(async () => new Response("unexpected request", { status: 500 })); + + try { + const request = await store.buildQueryResultExportRequest(tabId, { + exportId: "export-1", + filePath: "C:\\tmp\\events.csv", + format: "csv", + }); + + assert.equal(request?.exportId, "export-1"); + assert.equal(request?.connectionId, "conn-1"); + assert.equal(request?.database, "analytics"); + assert.equal(request?.schema, "public"); + assert.equal(request?.sql, "SELECT * FROM (SELECT * FROM events) t ORDER BY created_at DESC"); + assert.equal(request?.queryBaseSql, "SELECT * FROM events"); + assert.equal(request?.databaseType, "postgres"); + assert.equal(request?.useAgentCursor, false); + assert.equal(request?.filePath, "C:\\tmp\\events.csv"); + assert.equal(request?.format, "csv"); + assert.equal(request?.pageSize, 2500); + assert.equal(request?.rowLimit, null); + assert.equal(request?.totalRows, 123456); + assert.equal(request?.keysetOptimizationEnabled, false); + assert.equal(request?.clientSessionId, `${tabId}:export`); + assert.match(request?.executionId ?? "", /^[0-9a-f-]{36}$/i); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("buildQueryResultExportRequest uses exportRowLimit when enabled", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + settingsStore.updateEditorSettings({ exportBatchSize: 2500, exportRowLimit: 200000, exportRowLimitEnabled: true }); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + + connectionStore.addEphemeralConnection(conn("conn-1")); + const tabId = store.createTab("conn-1", "analytics", "Query", "query", "public"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + tab.lastExecutedSql = "SELECT * FROM events"; + tab.result = { + columns: ["id"], + rows: [[1]], + affected_rows: 0, + execution_time_ms: 1, + }; + + globalThis.fetch = withConnectionHealthMock(async () => new Response("unexpected request", { status: 500 })); + + try { + const request = await store.buildQueryResultExportRequest(tabId, { + exportId: "export-2", + filePath: "C:\\tmp\\events.xlsx", + format: "xlsx", + }); + + assert.equal(request?.pageSize, 2500); + assert.equal(request?.rowLimit, 200000); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + test("query execution finishes without waiting for metadata analysis", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia()); diff --git a/packages/app-tests/settingsStore.test.ts b/packages/app-tests/settingsStore.test.ts index 2396291d86..93fb374df7 100644 --- a/packages/app-tests/settingsStore.test.ts +++ b/packages/app-tests/settingsStore.test.ts @@ -54,12 +54,53 @@ test("normalizes saved query result page size", () => { assert.equal(normalizeEditorSettings({ pageSize: 0 }).pageSize, 100); }); -test("defaults export batch size to 10000 rows", () => { - assert.equal(DEFAULT_EDITOR_SETTINGS.exportBatchSize, 10000); - assert.equal(normalizeEditorSettings({}).exportBatchSize, 10000); +test("defaults export batch size to 2000 rows", () => { + assert.equal(DEFAULT_EDITOR_SETTINGS.exportBatchSize, 2000); + assert.equal(normalizeEditorSettings({}).exportBatchSize, 2000); assert.equal(normalizeEditorSettings({ exportBatchSize: 2000 }).exportBatchSize, 2000); }); +test("migrates the legacy saved export batch default to 2000 once", () => { + withMockLocalStorage({ "dbx-editor-settings": JSON.stringify({ exportBatchSize: 10000 }) }, () => { + setActivePinia(createPinia()); + const store = useSettingsStore(); + + assert.equal(store.editorSettings.exportBatchSize, 2000); + assert.equal(localStorage.getItem("dbx-export-batch-size-default-migrated-v1"), "1"); + assert.equal(JSON.parse(localStorage.getItem("dbx-editor-settings") || "{}").exportBatchSize, 2000); + }); +}); + +test("keeps a manually saved 10000 export batch size after migration", () => { + withMockLocalStorage( + { + "dbx-editor-settings": JSON.stringify({ exportBatchSize: 10000 }), + "dbx-export-batch-size-default-migrated-v1": "1", + }, + () => { + setActivePinia(createPinia()); + const store = useSettingsStore(); + + assert.equal(store.editorSettings.exportBatchSize, 10000); + }, + ); +}); + +test("defaults query-result export row limit settings", () => { + assert.equal(DEFAULT_EDITOR_SETTINGS.exportRowLimitEnabled, true); + assert.equal(DEFAULT_EDITOR_SETTINGS.exportRowLimit, 100000); + assert.equal(DEFAULT_EDITOR_SETTINGS.queryExportKeysetOptimizationEnabled, true); + assert.equal(normalizeEditorSettings({}).exportRowLimitEnabled, true); + assert.equal(normalizeEditorSettings({}).exportRowLimit, 100000); + assert.equal(normalizeEditorSettings({}).queryExportKeysetOptimizationEnabled, true); + assert.equal(normalizeEditorSettings({ exportRowLimitEnabled: false }).exportRowLimitEnabled, false); + assert.equal(normalizeEditorSettings({ exportRowLimitEnabled: "nope" as any }).exportRowLimitEnabled, true); + assert.equal(normalizeEditorSettings({ exportRowLimit: 250000 }).exportRowLimit, 250000); + assert.equal(normalizeEditorSettings({ exportRowLimit: 10 }).exportRowLimit, 100000); + assert.equal(normalizeEditorSettings({ queryExportKeysetOptimizationEnabled: false }).queryExportKeysetOptimizationEnabled, false); + assert.equal(normalizeEditorSettings({ queryExportKeysetOptimizationEnabled: "nope" as any }).queryExportKeysetOptimizationEnabled, true); +}); + test("normalizes editor theme settings", () => { assert.equal(DEFAULT_EDITOR_SETTINGS.theme, "app"); assert.equal(normalizeEditorSettings({}).theme, "app"); diff --git a/packages/app-tests/useDataGridExport.test.ts b/packages/app-tests/useDataGridExport.test.ts new file mode 100644 index 0000000000..26e05c2cd8 --- /dev/null +++ b/packages/app-tests/useDataGridExport.test.ts @@ -0,0 +1,216 @@ +import { strict as assert } from "node:assert"; +import { computed, ref } from "vue"; +import { beforeEach, test, vi } from "vitest"; + +const apiMock = vi.hoisted(() => ({ + startQueryResultExport: vi.fn(), + cancelQueryResultExport: vi.fn(), + startTableExport: vi.fn(), + cancelTableExport: vi.fn(), + exportQueryResultCsv: vi.fn(), + exportQueryResultXlsx: vi.fn(), + exportQueryResultJson: vi.fn(), + exportQueryResultMarkdown: vi.fn(), + exportQueryResultsXlsx: vi.fn(), +})); + +vi.mock("@/lib/api", () => apiMock); +vi.mock("@/lib/tauriRuntime", () => ({ isTauriRuntime: () => false })); +vi.mock("@/composables/useToast", () => ({ useToast: () => ({ toast: vi.fn() }) })); +vi.mock("vue-i18n", () => ({ useI18n: () => ({ t: (key: string) => key }) })); + +const { useDataGridExport } = await import("../../apps/desktop/src/composables/useDataGridExport.ts"); + +function buildExportHarness() { + const exportProgressDialog = ref(false); + const exportProgressState = ref({ + title: "", + tableName: "", + format: "csv", + rowsExported: 0, + totalRows: null as number | null, + status: "", + errorMessage: null as string | null, + }); + const exportCancelHandler = ref<(() => Promise) | null>(null); + const fullExportResult = vi.fn(async () => { + throw new Error("fullExportResult should not be called for streaming CSV/XLSX query exports"); + }); + const queryResultExportRequest = vi.fn(async (options: { exportId: string; filePath: string; format: "csv" | "xlsx" }) => ({ + exportId: options.exportId, + connectionId: "conn-1", + database: "db", + schema: "public", + sql: "SELECT * FROM users", + queryBaseSql: "SELECT * FROM users", + databaseType: "postgres" as const, + useAgentCursor: false, + filePath: options.filePath, + format: options.format, + pageSize: 1000, + rowLimit: 100000, + totalRows: 2, + timeoutSecs: 30, + keysetOptimizationEnabled: true, + clientSessionId: "tab-1:export", + executionId: "exec-1", + })); + + const composable = useDataGridExport({ + columns: computed(() => ["id", "name"]), + displayItems: computed(() => [ + { id: 1, data: [1, "Ada"], isNew: false, isDeleted: false, isDirtyCol: [false, false], status: "" }, + { id: 2, data: [2, "Lin"], isNew: false, isDeleted: false, isDirtyCol: [false, false], status: "" }, + ]), + sql: computed(() => "SELECT * FROM users"), + tableMeta: computed(() => undefined), + databaseType: computed(() => "postgres"), + connectionId: computed(() => "conn-1"), + database: computed(() => "db"), + context: computed(() => "results"), + sourceColumns: computed(() => undefined), + columnTypes: computed(() => undefined), + whereInput: computed(() => undefined), + orderBy: computed(() => undefined), + exportBatchSize: computed(() => 1000), + hasCellSelection: computed(() => false), + selectedCells: computed(() => ({ columns: [], rows: [] })), + selectedRange: computed(() => null), + contextCell: ref(null), + getRowItem: (rowId: number) => + [ + { id: 1, data: [1, "Ada"], isNew: false, isDeleted: false, isDirtyCol: [false, false], status: "" }, + { id: 2, data: [2, "Lin"], isNew: false, isDeleted: false, isDirtyCol: [false, false], status: "" }, + ].find((item) => item.id === rowId), + selectedRowIds: ref(new Set()), + hasRowSelection: computed(() => false), + fullExportResult, + queryResultExportRequest, + exportProgressDialog, + exportProgressState, + exportCancelHandler, + }); + + return { + composable, + fullExportResult, + queryResultExportRequest, + exportProgressDialog, + exportProgressState, + exportCancelHandler, + }; +} + +beforeEach(() => { + vi.clearAllMocks(); + apiMock.startQueryResultExport.mockImplementation(async (_request, onProgress) => { + onProgress({ exportId: _request.exportId, tableName: "", rowsExported: 2, totalRows: 2, status: "Done" }); + return { exportId: _request.exportId, tableName: "", rowsExported: 2, totalRows: 2, status: "Done" }; + }); +}); + +test("full query result CSV export streams through the backend without loading all rows", async () => { + const { composable, fullExportResult, queryResultExportRequest, exportProgressDialog, exportProgressState } = buildExportHarness(); + + await composable.exportCsv(); + + assert.equal(fullExportResult.mock.calls.length, 0); + assert.equal(queryResultExportRequest.mock.calls.length, 1); + assert.equal(apiMock.startQueryResultExport.mock.calls.length, 1); + assert.equal(apiMock.exportQueryResultCsv.mock.calls.length, 0); + assert.equal(exportProgressDialog.value, true); + assert.equal(exportProgressState.value.status, "Done"); +}); + +test("query result CSV cancel handler passes export and execution ids", async () => { + const { composable, exportCancelHandler } = buildExportHarness(); + let resolveExport!: () => void; + apiMock.startQueryResultExport.mockImplementationOnce(async (_request, onProgress) => { + await new Promise((resolve) => { + resolveExport = () => { + onProgress({ + exportId: _request.exportId, + tableName: "", + rowsExported: 1, + totalRows: 2, + status: "Cancelled", + errorMessage: "Export cancelled", + }); + resolve(); + }; + }); + return { + exportId: _request.exportId, + tableName: "", + rowsExported: 1, + totalRows: 2, + status: "Cancelled", + errorMessage: "Export cancelled", + }; + }); + + const exportPromise = composable.exportCsv(); + await vi.waitFor(() => assert.ok(exportCancelHandler.value)); + await exportCancelHandler.value?.(); + + const request = apiMock.startQueryResultExport.mock.calls[0][0]; + assert.deepEqual(apiMock.cancelQueryResultExport.mock.calls[0], [request.exportId, "exec-1"]); + + resolveExport(); + await exportPromise; +}); + +test("missing query result export request does not fall back to the in-memory path", async () => { + const { composable, fullExportResult, queryResultExportRequest } = buildExportHarness(); + queryResultExportRequest.mockResolvedValueOnce(undefined); + + await composable.exportCsv(); + + assert.equal(queryResultExportRequest.mock.calls.length, 1); + assert.equal(fullExportResult.mock.calls.length, 0); + assert.equal(apiMock.startQueryResultExport.mock.calls.length, 0); + assert.equal(apiMock.exportQueryResultCsv.mock.calls.length, 0); +}); + +test("selected query result CSV export keeps the existing in-memory path", async () => { + const { composable, queryResultExportRequest } = buildExportHarness(); + + await composable.exportCsv([1]); + + assert.equal(queryResultExportRequest.mock.calls.length, 0); + assert.equal(apiMock.startQueryResultExport.mock.calls.length, 0); + assert.equal(apiMock.exportQueryResultCsv.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("cancelled query result CSV export clears the cancel handler without using the in-memory path", async () => { + const { composable, fullExportResult, exportProgressState, exportCancelHandler } = buildExportHarness(); + apiMock.startQueryResultExport.mockImplementationOnce(async (_request, onProgress) => { + onProgress({ + exportId: _request.exportId, + tableName: "", + rowsExported: 1, + totalRows: 2, + status: "Cancelled", + errorMessage: "Export cancelled", + }); + return { + exportId: _request.exportId, + tableName: "", + rowsExported: 1, + totalRows: 2, + status: "Cancelled", + errorMessage: "Export cancelled", + }; + }); + + await composable.exportCsv(); + + assert.equal(fullExportResult.mock.calls.length, 0); + assert.equal(apiMock.startQueryResultExport.mock.calls.length, 1); + assert.equal(apiMock.exportQueryResultCsv.mock.calls.length, 0); + assert.equal(exportProgressState.value.status, "Cancelled"); + assert.equal(exportProgressState.value.errorMessage, "Export cancelled"); + assert.equal(exportCancelHandler.value, null); +}); diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index 0fbaa969d6..34cd1e26f5 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -24,6 +24,7 @@ pub mod nacos_cmd; pub mod plugins; pub mod query; pub mod query_cancel; +pub mod query_result_export; pub mod redis_cmd; pub mod redis_pubsub_server; pub mod saved_sql; diff --git a/src-tauri/src/commands/query_result_export.rs b/src-tauri/src/commands/query_result_export.rs new file mode 100644 index 0000000000..15eb46418e --- /dev/null +++ b/src-tauri/src/commands/query_result_export.rs @@ -0,0 +1,88 @@ +use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, +}; + +use tauri::{AppHandle, Emitter, State}; + +use crate::commands::connection::AppState; + +use dbx_core::query_cancel::RunningTaskMetadata; +pub use dbx_core::query_result_export::QueryResultExportRequest; +use dbx_core::table_export::ExportStatus; +pub use dbx_core::table_export::TableExportProgress; + +fn emit_progress(app: &AppHandle, progress: TableExportProgress) { + let _ = app.emit("query-result-export-progress", progress); +} + +#[tauri::command] +pub async fn start_query_result_export( + app: AppHandle, + state: State<'_, Arc>, + request: QueryResultExportRequest, +) -> Result<(), String> { + let state = state.inner().clone(); + let export_id = request.export_id.clone(); + let file_path = request.file_path.clone(); + + tokio::spawn(async move { + let execution_id = request.execution_id.clone().filter(|id| !id.trim().is_empty()); + let registered_query = execution_id.as_ref().map(|id| { + state.running_queries.register_task( + id.clone(), + RunningTaskMetadata::query( + request.connection_id.clone(), + request.database.clone(), + request.client_session_id.clone(), + ), + ) + }); + let cancel_token = registered_query.as_ref().map(|query| query.token()); + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_progress = cancelled.clone(); + let result = + dbx_core::query_result_export::export_query_result_core(&state, &request, cancel_token, |progress| { + if matches!(progress.status, ExportStatus::Cancelled) { + cancelled_progress.store(true, Ordering::SeqCst); + } + emit_progress(&app, progress); + }) + .await; + drop(registered_query); + + if let Err(e) = result { + let _ = tokio::fs::remove_file(&file_path).await; + emit_progress( + &app, + TableExportProgress { + export_id: export_id.clone(), + table_name: String::new(), + rows_exported: 0, + total_rows: None, + status: ExportStatus::Error, + error_message: Some(e), + }, + ); + } else if cancelled.load(Ordering::SeqCst) { + let _ = tokio::fs::remove_file(&file_path).await; + } + + dbx_core::database_export::clear_export_cancelled(&export_id).await; + }); + + Ok(()) +} + +#[tauri::command] +pub async fn cancel_query_result_export( + state: State<'_, Arc>, + export_id: String, + execution_id: Option, +) -> Result<(), String> { + dbx_core::database_export::set_export_cancelled(&export_id).await; + if let Some(execution_id) = execution_id.filter(|id| !id.trim().is_empty()) { + state.running_queries.cancel(&execution_id); + } + Ok(()) +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index a4f9622f8f..370bdb4f76 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -707,6 +707,8 @@ pub fn run() { commands::database_export::cancel_database_export, commands::table_export::start_table_export, commands::table_export::cancel_table_export, + commands::query_result_export::start_query_result_export, + commands::query_result_export::cancel_query_result_export, commands::csv_export::export_query_result_csv, commands::csv_export::export_table_data_csv, commands::xlsx_export::export_query_result_xlsx,