feat(export): support safe background query exports

This commit is contained in:
zhangyaxi 2026-07-29 03:18:06 +08:00 committed by GitHub
parent 2f44d36018
commit a07def6042
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 577 additions and 195 deletions

View File

@ -1,6 +1,6 @@
<script setup lang="ts">
import { ref, onMounted, onBeforeUnmount, onActivated, onDeactivated, watch, shallowRef, computed, nextTick } from "vue";
import { CaseLower, CaseUpper, Code2, FileCode, Pencil, PencilRuler, Play, Copy, List, Search, Sparkles, Table2, TextSelect, Trash2 } from "@lucide/vue";
import { CaseLower, CaseUpper, Code2, Download, FileCode, Pencil, PencilRuler, Play, Copy, List, Search, Sparkles, Table2, TextSelect, Trash2 } from "@lucide/vue";
import { useI18n } from "vue-i18n";
import type { CompletionContext } from "@codemirror/autocomplete";
import { Transaction, StateEffect } from "@codemirror/state";
@ -138,6 +138,7 @@ const emit = defineEmits<{
formatError: [message: string];
execute: [source: SqlExecutionOverride];
executeInNewResultTab: [source: SqlExecutionOverride];
exportQuery: [payload: { sql: string; format: "csv" | "xlsx" | "txt" }];
save: [];
clickTable: [target: SqlObjectNavigationTarget];
viewTableData: [target: SqlObjectNavigationTarget];
@ -1000,6 +1001,12 @@ function executeInNewResultTabFromContextMenu() {
focusEditor();
}
function exportQueryFromContextMenu(format: "csv" | "xlsx" | "txt") {
const sql = executableSql.value;
if (!sql.trim()) return;
emit("exportQuery", { sql, format });
}
async function copySelectedSqlFromContextMenu() {
if (!canCopySelectedSql.value) return;
try {
@ -1274,6 +1281,16 @@ const contextMenuItems = computed<ContextMenuItem[]>(() => {
icon: Play,
shortcut: shortcuts.executeSqlInNewResultTab,
},
{
label: t("editor.contextMenu.export"),
icon: Download,
disabled: !canExecuteContextSql.value,
children: [
{ label: t("editor.contextMenu.exportQueryResultTo", { format: "CSV" }), action: () => exportQueryFromContextMenu("csv") },
{ label: t("editor.contextMenu.exportQueryResultTo", { format: "XLSX" }), action: () => exportQueryFromContextMenu("xlsx") },
{ label: t("editor.contextMenu.exportQueryResultTo", { format: "TXT" }), action: () => exportQueryFromContextMenu("txt") },
],
},
]),
...queryContextObjectActions(contextObjectTarget.value?.type).map(contextObjectMenuItem),
{ label: "", separator: true },

View File

@ -3,7 +3,7 @@ import { computed, ref } from "vue";
import { useI18n } from "vue-i18n";
import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogFooter } from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Loader2, CheckCircle2, XCircle, AlertCircle, FolderOpen, X } from "@lucide/vue";
import { Loader2, CheckCircle2, XCircle, AlertCircle, FolderOpen, Minimize2, X } from "@lucide/vue";
import { useToast } from "@/composables/useToast";
import { isTauriRuntime } from "@/lib/backend/tauriRuntime";
import * as api from "@/lib/backend/api";
@ -24,10 +24,12 @@ const props = defineProps<{
errorMessage: string | null;
filePath?: string | null;
disableCancel?: boolean;
canMinimize?: boolean;
}>();
const emit = defineEmits<{
cancel: [];
minimize: [];
"update:open": [value: boolean];
}>();
@ -123,6 +125,10 @@ async function revealExportFile() {
<DialogFooter>
<template v-if="isActive">
<Button v-if="canMinimize" variant="ghost" size="sm" @click="emit('minimize')">
<Minimize2 class="h-3.5 w-3.5 mr-1" />
{{ t("exportProgress.minimize") }}
</Button>
<Button variant="outline" size="sm" :disabled="disableCancel" @click="emit('cancel')">
<X class="h-3.5 w-3.5 mr-1" />
{{ t("exportProgress.cancel") }}

View File

@ -5121,6 +5121,7 @@ const exportProgressState = ref({
filePath: null as string | null,
});
const exportCancelHandler = ref<(() => Promise<void>) | null>(null);
const exportCanMinimize = ref(false);
async function cancelActiveExport() {
await exportCancelHandler.value?.();
@ -5196,6 +5197,7 @@ const {
exportProgressDialog,
exportProgressState,
exportCancelHandler,
exportCanMinimize,
});
function copyExtractorLabel(extractor: DataGridCopyExtractorId): string {
@ -9302,7 +9304,7 @@ const gridContextMenuItems = computed<ContextMenuItem[]>(() => {
/>
<ImagePreviewDialog v-if="imagePreviewMounted" v-model:open="imagePreviewOpen" :src="imagePreviewSrc" :title="imagePreviewTitle" />
<component v-if="previewDialogOpen && previewDialogConfig" :is="previewDialogConfig.component" v-model:open="previewDialogOpen" v-bind="previewDialogConfig.props" />
<ExportProgressDialog v-if="exportProgressDialogMounted" v-model:open="exportProgressDialog" v-bind="exportProgressState" :disable-cancel="!exportCancelHandler" @cancel="cancelActiveExport" />
<ExportProgressDialog v-if="exportProgressDialogMounted" v-model:open="exportProgressDialog" v-bind="exportProgressState" :disable-cancel="!exportCancelHandler" :can-minimize="exportCanMinimize" @cancel="cancelActiveExport" @minimize="exportProgressDialog = false" />
</div>
</template>

View File

@ -68,6 +68,7 @@ import { useQueryStore } from "@/stores/queryStore";
import { useConnectionStore } from "@/stores/connectionStore";
import { TABLE_FONT_SIZE_MAX, TABLE_FONT_SIZE_MIN, useSettingsStore, type DataGridSearchMode, type ResultRunDisplayMode } from "@/stores/settingsStore";
import { useToast } from "@/composables/useToast";
import { isTauriRuntime } from "@/lib/backend/tauriRuntime";
import { canCancelQueryExecution, queryExecutionLabelKey } from "@/lib/sql/queryExecutionState";
import { databaseDisplayNameForTab, executionSummaryItems, queryResultExecutionSql, resultGridCacheKey, resultRunItems, resultSourceRange, resultSqlForGrid, statementExecutionMarkers, tabularResultItems } from "@/lib/tabs/tabPresentation";
import { defaultQueryResultArchiveFileName } from "@/lib/query/queryResultArchive";
@ -865,6 +866,20 @@ function requestQueryEditorExecuteInNewResultTab() {
return queryEditorRef.value?.requestExecuteInNewResultTab();
}
async function handleExportQuery(payload: { sql: string; format: "csv" | "xlsx" | "txt" }) {
const tab = props.activeTab;
if (!tab || tab.mode !== "query") return;
let filePath = `query-result.${payload.format}`;
if (isTauriRuntime()) {
const { save } = await import("@tauri-apps/plugin-dialog");
const filterName = payload.format === "csv" ? "CSV" : payload.format === "xlsx" ? "Excel" : "Text";
const picked = await save({ defaultPath: filePath, filters: [{ name: filterName, extensions: [payload.format] }] });
if (!picked) return;
filePath = picked as string;
}
await queryStore.exportQuerySqlDirect(tab.id, payload.sql, payload.format, filePath);
}
function pasteClipboardAsSqlInCondition() {
return queryEditorRef.value?.pasteClipboardAsSqlInCondition();
}
@ -932,6 +947,7 @@ defineExpose({ focusSearch, refreshData, refreshQueryEditorCompletionCache, hand
@format-error="emit('formatError')"
@execute="emit('execute', $event)"
@execute-in-new-result-tab="emit('executeInNewResultTab', $event)"
@export-query="handleExportQuery"
@save="emit('saveSql')"
@click-table="onHandleClickTable"
@view-table-data="onHandleViewTableData"

View File

@ -128,6 +128,7 @@ export interface UseDataGridExportOptions {
filePath: string | null;
}>;
exportCancelHandler?: Ref<(() => Promise<void>) | null>;
exportCanMinimize?: Ref<boolean>;
}
interface CopyInsertData {
@ -140,6 +141,7 @@ interface CopyInsertData {
export function useDataGridExport(options: UseDataGridExportOptions) {
const { t } = useI18n();
const { toast } = useToast();
const tracker = useExportTracker();
const exportGuard: ActionActivationGuard = {};
const { addTask, updateTableExportTask, registerTaskCancelHandler, unregisterTaskCancelHandler, removeTask } = useExportTracker();
@ -185,6 +187,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) {
exportProgressDialog,
exportProgressState,
exportCancelHandler,
exportCanMinimize,
} = options;
const selectedCellMatrix = selectedCellMatrixOption;
const allColumns = allColumnsOption ?? columns;
@ -906,11 +909,14 @@ export function useDataGridExport(options: UseDataGridExportOptions) {
};
}
if (exportProgressDialog) exportProgressDialog.value = true;
if (exportCanMinimize) exportCanMinimize.value = true;
const exportId = uuid();
const task = tracker.addTask(meta.tableName, format, outputPath);
const exportId = task.exportId;
if (exportCancelHandler) {
exportCancelHandler.value = () => api.cancelTableExport(exportId);
}
tracker.registerTaskCancelHandler(exportId, () => api.cancelTableExport(exportId));
const editorSettings = useSettingsStore().editorSettings;
const rowLimit = editorSettings.exportRowLimitEnabled ? editorSettings.exportRowLimit : null;
@ -946,6 +952,7 @@ export function useDataGridExport(options: UseDataGridExportOptions) {
errorMessage: progress.errorMessage || null,
};
}
tracker.updateTableExportTask(exportId, progress);
},
);
if (progress.status === "Done") {
@ -953,6 +960,8 @@ export function useDataGridExport(options: UseDataGridExportOptions) {
}
} finally {
if (exportCancelHandler) exportCancelHandler.value = null;
tracker.unregisterTaskCancelHandler(exportId);
if (exportCanMinimize) exportCanMinimize.value = false;
}
return true;
}
@ -997,9 +1006,12 @@ export function useDataGridExport(options: UseDataGridExportOptions) {
};
}
if (exportProgressDialog) exportProgressDialog.value = true;
if (exportCanMinimize) exportCanMinimize.value = true;
tracker.addTask("Query Result", format, outputPath, exportId);
if (exportCancelHandler) {
exportCancelHandler.value = () => api.cancelQueryResultExport(exportId, request.executionId);
}
tracker.registerTaskCancelHandler(exportId, () => api.cancelQueryResultExport(exportId, request.executionId));
try {
const terminalProgress = await api.startQueryResultExport(request, (progress) => {
@ -1014,12 +1026,15 @@ export function useDataGridExport(options: UseDataGridExportOptions) {
errorMessage: progress.errorMessage || null,
};
}
tracker.updateTableExportTask(exportId, progress);
});
if (terminalProgress.status === "Done") {
toast(t("grid.exported"));
}
} finally {
if (exportCancelHandler) exportCancelHandler.value = null;
tracker.unregisterTaskCancelHandler(exportId);
if (exportCanMinimize) exportCanMinimize.value = false;
}
return true;
}

View File

@ -229,10 +229,10 @@ export function useExportTracker() {
.replace(/^(.{8})(.{4})(.{4})(.{4})(.{12})$/, "$1-$2-$3-$4-$5");
}
function addTask(tableName: string, format: string, filePath: string): ExportTask {
const exportId = generateUUID();
function addTask(tableName: string, format: string, filePath: string, exportId?: string): ExportTask {
const id = exportId ?? generateUUID();
const task = reactive<ExportTask>({
exportId,
exportId: id,
kind: "table-export",
tableName,
format,
@ -242,7 +242,7 @@ export function useExportTracker() {
status: "Running",
errorMessage: null,
});
taskMap.set(exportId, task);
taskMap.set(id, task);
return task;
}

View File

@ -819,6 +819,8 @@ export default {
contextMenu: {
executeSelection: "Execute selection",
executeCurrent: "Execute SQL",
export: "Export",
exportQueryResultTo: "Export current query result to {format}",
copySelection: "Copy selection",
sendToAi: "Send to AI",
uppercaseSelection: "Convert to uppercase",
@ -1399,6 +1401,7 @@ export default {
rowsExported: "{count} rows exported",
rowsShort: "rows",
cancel: "Cancel",
minimize: "Minimize to background",
openFolder: "Open containing folder",
openFolderFailed: "Failed to open folder: {message}",
close: "Close",

View File

@ -799,6 +799,8 @@ export default withEnglishFallback({
contextMenu: {
executeSelection: "Ejecutar seleccion",
executeCurrent: "Ejecutar SQL",
export: "Exportar",
exportQueryResultTo: "Exportar el conjunto de resultados de la consulta actual a {format}",
copySelection: "Copiar seleccion",
sendToAi: "Enviar a IA",
uppercaseSelection: "Convertir a mayusculas",
@ -1342,6 +1344,7 @@ export default withEnglishFallback({
rowsExported: "{count} filas exportadas",
rowsShort: "filas",
cancel: "Cancelar",
minimize: "Minimizar a segundo plano",
openFolder: "Abrir carpeta contenedora",
openFolderFailed: "No se pudo abrir la carpeta: {message}",
close: "Cerrar",

View File

@ -797,6 +797,8 @@ export default withEnglishFallback({
contextMenu: {
executeSelection: "Esegui selezione",
executeCurrent: "Esegui SQL",
export: "Esporta",
exportQueryResultTo: "Esporta il set di risultati della query corrente in {format}",
copySelection: "Copia selezione",
sendToAi: "Invia ad AI",
uppercaseSelection: "Converti in maiuscolo",
@ -1340,6 +1342,7 @@ export default withEnglishFallback({
rowsExported: "{count} righe esportate",
rowsShort: "righe",
cancel: "Annulla",
minimize: "Minimizza in background",
openFolder: "Apri cartella contenente",
openFolderFailed: "Impossibile aprire la cartella: {message}",
close: "Chiudi",

View File

@ -796,6 +796,8 @@ export default withEnglishFallback({
contextMenu: {
executeSelection: "選択範囲を実行",
executeCurrent: "SQLを実行",
export: "エクスポート",
exportQueryResultTo: "現在のクエリ結果セットを {format} にエクスポート",
copySelection: "選択範囲をコピー",
sendToAi: "AIに送信",
uppercaseSelection: "大文字に変換",
@ -1341,6 +1343,7 @@ export default withEnglishFallback({
rowsExported: "{count}行をエクスポートしました",
rowsShort: "行",
cancel: "キャンセル",
minimize: "バックグラウンドに最小化",
openFolder: "保存先フォルダーを開く",
openFolderFailed: "フォルダーを開けませんでした:{message}",
close: "閉じる",

View File

@ -798,6 +798,8 @@ export default withEnglishFallback({
contextMenu: {
executeSelection: "Executar seleção",
executeCurrent: "Executar SQL",
export: "Exportar",
exportQueryResultTo: "Exportar conjunto de resultados da consulta atual para {format}",
copySelection: "Copiar seleção",
sendToAi: "Enviar para IA",
uppercaseSelection: "Converter para maiúsculas",
@ -1342,6 +1344,7 @@ export default withEnglishFallback({
rowsExported: "{count} linhas exportadas",
rowsShort: "linhas",
cancel: "Cancelar",
minimize: "Minimizar para segundo plano",
openFolder: "Abrir pasta do arquivo",
openFolderFailed: "Falha ao abrir a pasta: {message}",
close: "Fechar",

View File

@ -820,6 +820,8 @@ export default withEnglishFallback({
contextMenu: {
executeSelection: "执行选中 SQL",
executeCurrent: "执行 SQL",
export: "导出",
exportQueryResultTo: "导出当前查询结果集到 {format}",
copySelection: "复制选中内容",
sendToAi: "发送到 AI",
uppercaseSelection: "转为大写",
@ -1400,6 +1402,7 @@ export default withEnglishFallback({
rowsExported: "已导出 {count} 行",
rowsShort: "行",
cancel: "取消",
minimize: "最小化到后台",
openFolder: "打开所在文件夹",
openFolderFailed: "打开文件夹失败:{message}",
close: "关闭",

View File

@ -797,6 +797,8 @@ export default withEnglishFallback({
contextMenu: {
executeSelection: "執行選取 SQL",
executeCurrent: "執行 SQL",
export: "匯出",
exportQueryResultTo: "匯出目前查詢結果集到 {format}",
copySelection: "複製選取內容",
sendToAi: "傳送至 AI",
uppercaseSelection: "轉為大寫",
@ -1341,6 +1343,7 @@ export default withEnglishFallback({
rowsExported: "已匯出 {count} 列",
rowsShort: "列",
cancel: "取消",
minimize: "最小化到背景",
openFolder: "開啟所在資料夾",
openFolderFailed: "開啟資料夾失敗:{message}",
close: "關閉",

View File

@ -58,6 +58,7 @@ import * as api from "@/lib/backend/api";
import { useConnectionStore } from "@/stores/connectionStore";
import { useSettingsStore } from "@/stores/settingsStore";
import { useSavedSqlStore } from "@/stores/savedSqlStore";
import { useExportTracker } from "@/composables/useExportTracker";
import { recordQueryCancellationLatency, resourceLifecycleDiagnostics } from "@/lib/diagnostics/resourceLifecycleDiagnostics";
import { appendDebugLog } from "@/lib/backend/debugLog";
import { formatError } from "@/lib/backend/errorUtils";
@ -4770,7 +4771,7 @@ export const useQueryStore = defineStore("query", () => {
const setupSql = batchStatements[resultStatementIndex!]?.sql === tab.result.sourceStatement ? batchStatements.slice(0, resultStatementIndex).map((statement) => statement.sql) : undefined;
const rowLimit = settings.exportRowLimitEnabled ? settings.exportRowLimit : null;
const totalRows = typeof tab.resultTotalRowCount === "number" ? (rowLimit === null ? tab.resultTotalRowCount : Math.min(tab.resultTotalRowCount, rowLimit)) : null;
const clientSessionId = tabClientSessionId(tab, "export");
const clientSessionId = `${tabClientSessionId(tab, "export")}:${options.exportId}`;
return {
exportId: options.exportId,
@ -4798,6 +4799,58 @@ export const useQueryStore = defineStore("query", () => {
};
}
async function exportQuerySqlDirect(id: string, sql: string, format: "csv" | "xlsx" | "txt", filePath: string) {
const tab = tabs.value.find((item) => item.id === id);
if (!tab || tab.mode !== "query" || !sql.trim()) return;
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;
const exportId = uuid();
const request: api.QueryResultExportRequest = {
exportId,
connectionId: tab.connectionId,
database: tab.database,
schema: tab.schema,
sql,
queryBaseSql: sql,
databaseType: effectiveDbType,
useAgentCursor: usesAgentCursorForQuery(conn?.db_type),
filePath,
format,
pageSize: settings.exportBatchSize,
rowLimit: settings.exportRowLimitEnabled ? settings.exportRowLimit : null,
totalRows: null,
timeoutSecs: queryTimeoutSecsForConnection(conn),
keysetOptimizationEnabled: settings.queryExportKeysetOptimizationEnabled,
clientSessionId: `${tabClientSessionId(tab, "export")}:${exportId}`,
executionId: uuid(),
numericColumnRightAlign: settings.numericColumnRightAlign,
};
const tracker = useExportTracker();
tracker.addTask("Query Result", format, filePath, request.exportId);
tracker.registerTaskCancelHandler(request.exportId, () => api.cancelQueryResultExport(request.exportId, request.executionId));
void (async () => {
try {
await api.startQueryResultExport(request, (progress) => tracker.updateTableExportTask(request.exportId, progress));
} catch (error: any) {
const task = tracker.tasks.value.find((item) => item.exportId === request.exportId);
if (task) {
task.status = "Error";
task.errorMessage = error?.message || String(error);
}
} finally {
tracker.unregisterTaskCancelHandler(request.exportId);
}
})();
}
return {
tabs,
activeTabId,
@ -4900,6 +4953,7 @@ export const useQueryStore = defineStore("query", () => {
importResultArchive,
fetchTabResultForExport,
buildQueryResultExportRequest,
exportQuerySqlDirect,
getResourceLifecycleDiagnostics: () => resourceLifecycleDiagnostics(tabs.value),
notifyConnectionMayBeLost,
};

View File

@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize};
use std::fs::File;
use std::io::{BufWriter, Seek, Write};
use std::mem;
use std::path::{Path, PathBuf};
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
@ -102,6 +103,55 @@ pub struct QueryResultExportRequest {
pub numeric_column_right_align: bool,
}
pub struct StagedExportTarget {
destination: PathBuf,
temporary: tempfile::TempPath,
}
impl StagedExportTarget {
pub fn new(destination: &str) -> Result<Self, String> {
let destination = PathBuf::from(destination);
let parent = destination.parent().filter(|path| !path.as_os_str().is_empty()).unwrap_or(Path::new("."));
let file_name = destination.file_name().and_then(|name| name.to_str()).unwrap_or("query-result");
let temporary = tempfile::Builder::new()
.prefix(&format!(".{file_name}.dbx-export-"))
.tempfile_in(parent)
.map_err(|error| format!("Failed to create temporary export file: {error}"))?
.into_temp_path();
Ok(Self { destination, temporary })
}
pub fn path(&self) -> &Path {
self.temporary.as_ref()
}
pub fn path_string(&self) -> Result<String, String> {
self.path()
.to_str()
.map(ToOwned::to_owned)
.ok_or_else(|| "Temporary export path is not valid UTF-8".to_string())
}
pub fn commit(self) -> Result<(), String> {
let staged_file = std::fs::OpenOptions::new()
.read(true)
.write(true)
.open(self.path())
.map_err(|error| format!("Failed to open staged export file: {error}"))?;
if let Ok(metadata) = std::fs::metadata(&self.destination) {
staged_file
.set_permissions(metadata.permissions())
.map_err(|error| format!("Failed to preserve export destination permissions: {error}"))?;
}
staged_file.sync_all().map_err(|error| format!("Failed to synchronize export file: {error}"))?;
drop(staged_file);
self.temporary
.persist(&self.destination)
.map(|_| ())
.map_err(|error| format!("Failed to replace export destination: {}", error.error))
}
}
fn safe_postgres_temp_setup_sql(setup_sql: &[String]) -> Option<Vec<String>> {
if setup_sql.is_empty() {
return None;
@ -231,41 +281,30 @@ fn sql_insert_column_types(request: &QueryResultExportRequest, column_types: &[S
}
}
/// Bounded SQL INSERT writer with temp-file + atomic-rename safety.
/// Bounded SQL INSERT writer with staged-file replacement safety.
///
/// Rows are buffered and flushed to a temp file every [`SQL_INSERT_BATCH_SIZE`]
/// rows, so memory stays bounded regardless of the query page size. The temp file
/// lives alongside the user-chosen target (`<target>.dbx-export-tmp`) and is
/// renamed onto the target only when [`SqlInsertWriter::finish`] succeeds. If the
/// writer is dropped without finishing (export error or cancellation), the temp
/// file is removed and the user's target file is never truncated or deleted.
/// rows, so memory stays bounded regardless of the query page size. The unique
/// temp file lives alongside the target and replaces it only after
/// [`SqlInsertWriter::finish`] flushes and synchronizes the complete output.
struct SqlInsertWriter {
file: Option<BufWriter<File>>,
temp_path: std::path::PathBuf,
target_path: std::path::PathBuf,
target: Option<StagedExportTarget>,
pending_rows: Vec<Vec<Value>>,
columns: Vec<String>,
column_types: Vec<Option<String>>,
database_type: DatabaseType,
schema: Option<String>,
table_name: String,
finished: bool,
}
impl SqlInsertWriter {
/// Create the writer and open the temp file. Column metadata is supplied later
/// via [`SqlInsertWriter::set_columns`] once the executed result is known.
fn create(request: &QueryResultExportRequest) -> Result<Self, String> {
let target_path = std::path::PathBuf::from(&request.file_path);
let mut temp_path = target_path.clone();
let mut temp_name = target_path
.file_name()
.map(|name| name.to_os_string())
.unwrap_or_else(|| std::ffi::OsString::from("export.sql"));
temp_name.push(".dbx-export-tmp");
temp_path.set_file_name(temp_name);
let target = StagedExportTarget::new(&request.file_path)?;
let file = BufWriter::new(
File::create(&temp_path).map_err(|e| format!("Failed to create SQL export temp file: {e}"))?,
File::create(target.path()).map_err(|e| format!("Failed to create SQL export temp file: {e}"))?,
);
let table_name = request
.export_table_name
@ -275,15 +314,13 @@ impl SqlInsertWriter {
.to_string();
Ok(Self {
file: Some(file),
temp_path,
target_path,
target: Some(target),
pending_rows: Vec::new(),
columns: Vec::new(),
column_types: Vec::new(),
database_type: request.database_type,
schema: request.schema.clone(),
table_name,
finished: false,
})
}
@ -330,30 +367,18 @@ impl SqlInsertWriter {
Ok(())
}
/// Flush remaining rows, close the temp file, and atomically rename it onto the
/// target path. After this succeeds the writer no longer removes the file on drop.
/// Flush remaining rows, close the temp file, and atomically replace the target.
fn finish(mut self) -> Result<(), String> {
self.flush_batch()?;
if let Some(file) = self.file.as_mut() {
file.flush().map_err(|e| format!("Failed to flush SQL file: {e}"))?;
}
// Close the file handle before rename (required on Windows).
self.file.take();
self.finished = true;
std::fs::rename(&self.temp_path, &self.target_path)
.map_err(|e| format!("Failed to finalize SQL export file: {e}"))?;
Ok(())
}
}
impl Drop for SqlInsertWriter {
fn drop(&mut self) {
if !self.finished {
// Best-effort: close the handle then remove the temp file so the user's
// chosen target is never left truncated or deleted by a failed export.
self.file.take();
let _ = std::fs::remove_file(&self.temp_path);
}
self.target
.take()
.ok_or_else(|| "SQL export target already finalized".to_string())?
.commit()
.map_err(|error| format!("Failed to finalize SQL export file: {error}"))
}
}
@ -1730,6 +1755,42 @@ async fn try_export_sqlserver_query_result_stream(
mod tests {
use super::*;
#[test]
fn staged_export_target_preserves_existing_destination_on_discard_and_replace_failure() {
let dir = tempfile::tempdir().expect("temp dir");
let destination = dir.path().join("result.csv");
std::fs::write(&destination, "original").expect("write destination");
let discarded = StagedExportTarget::new(destination.to_str().expect("destination path")).expect("target");
std::fs::write(discarded.path(), "partial").expect("write partial export");
drop(discarded);
assert_eq!(std::fs::read_to_string(&destination).expect("read destination"), "original");
let failed = StagedExportTarget::new(destination.to_str().expect("destination path")).expect("target");
std::fs::write(failed.path(), "replacement").expect("write replacement");
std::fs::remove_file(failed.path()).expect("remove staged path");
assert!(failed.commit().expect_err("replace should fail").contains("open staged export file"));
assert_eq!(std::fs::read_to_string(destination).expect("read destination"), "original");
}
#[test]
fn staged_export_targets_are_unique_same_directory_and_replace_existing_destination() {
let dir = tempfile::tempdir().expect("temp dir");
let destination = dir.path().join("result.csv");
std::fs::write(&destination, "original").expect("write destination");
let first = StagedExportTarget::new(destination.to_str().expect("destination path")).expect("first target");
let second = StagedExportTarget::new(destination.to_str().expect("destination path")).expect("second target");
assert_eq!(first.path().parent(), destination.parent());
assert_eq!(second.path().parent(), destination.parent());
assert_ne!(first.path(), second.path());
std::fs::write(first.path(), "replacement").expect("write replacement");
first.commit().expect("commit export");
drop(second);
assert_eq!(std::fs::read_to_string(destination).expect("read destination"), "replacement");
}
#[test]
fn stream_cancel_detection_covers_driver_token_and_export_flags() {
assert!(stream_export_was_cancelled(QUERY_CANCELED, false, false));

View File

@ -6,6 +6,7 @@ import { decodeQueryResultArchive } from "../../apps/desktop/src/lib/query/query
import { analyzeEditableQueryEditability } from "../../apps/desktop/src/lib/sql/sqlAnalysis.ts";
import { resultSqlForGrid } from "../../apps/desktop/src/lib/tabs/tabPresentation.ts";
import { parseMongoCommand } from "../../apps/desktop/src/lib/mongo/mongoShellCommand.ts";
import { useExportTracker } from "../../apps/desktop/src/composables/useExportTracker.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";
@ -4505,14 +4506,117 @@ test("buildQueryResultExportRequest uses sorted SQL and independent row-limit se
assert.equal(request?.rowLimit, null);
assert.equal(request?.totalRows, 123456);
assert.equal(request?.keysetOptimizationEnabled, false);
assert.equal(request?.clientSessionId, `${tabId}:export`);
assert.equal(request?.clientSessionId, `${tabId}:export:export-1`);
assert.match(request?.executionId ?? "", /^[0-9a-f-]{36}$/i);
const concurrentRequest = await store.buildQueryResultExportRequest(tabId, {
exportId: "export-2",
filePath: "C:\\tmp\\events-2.csv",
format: "csv",
});
assert.equal(concurrentRequest?.clientSessionId, `${tabId}:export:export-2`);
assert.notEqual(concurrentRequest?.clientSessionId, request?.clientSessionId);
} finally {
globalThis.fetch = originalFetch;
restoreStorage();
}
});
test("same-tab direct exports use isolated sessions and cancel handlers", async () => {
const restoreStorage = installMemoryStorage();
const originalFetch = globalThis.fetch;
const originalEventSource = Object.getOwnPropertyDescriptor(globalThis, "EventSource");
const startedRequests: Array<Record<string, any>> = [];
const cancelRequests: Array<Record<string, any>> = [];
const createdExportIds: string[] = [];
class FakeEventSource {
static instances: FakeEventSource[] = [];
onopen: ((event: Event) => void) | null = null;
onmessage: ((event: MessageEvent) => void) | null = null;
onerror: ((event: Event) => void) | null = null;
constructor(readonly url: string) {
FakeEventSource.instances.push(this);
}
close() {}
emitOpen() {
this.onopen?.({} as Event);
}
emitProgress(progress: Record<string, unknown>) {
this.onmessage?.({ data: JSON.stringify(progress) } as MessageEvent);
}
}
Object.defineProperty(globalThis, "EventSource", { configurable: true, value: FakeEventSource });
setActivePinia(createPinia());
const connectionStore = useConnectionStore();
const store = useQueryStore();
const tracker = useExportTracker();
connectionStore.addEphemeralConnection(conn("conn-1"));
const tabId = store.createTab("conn-1", "analytics", "Query", "query", "public");
globalThis.fetch = withConnectionHealthMock(async (input, init) => {
const path = String(input);
if (path === "/api/export/query-result") {
startedRequests.push(JSON.parse(String(init?.body ?? "{}")));
return new Response("null", { status: 200, headers: { "Content-Type": "application/json" } });
}
if (path === "/api/export/query-result/cancel") {
cancelRequests.push(JSON.parse(String(init?.body ?? "{}")));
return new Response("null", { status: 200, headers: { "Content-Type": "application/json" } });
}
return new Response("unexpected request", { status: 500 });
});
try {
await store.exportQuerySqlDirect(tabId, "SELECT 1", "csv", "/tmp/first.csv");
await store.exportQuerySqlDirect(tabId, "SELECT 2", "csv", "/tmp/second.csv");
await waitFor(() => FakeEventSource.instances.length === 2);
FakeEventSource.instances.forEach((source) => source.emitOpen());
await waitFor(() => startedRequests.length === 2);
const first = startedRequests[0].request;
const second = startedRequests[1].request;
createdExportIds.push(first.exportId, second.exportId);
assert.equal(first.clientSessionId, `${tabId}:export:${first.exportId}`);
assert.equal(second.clientSessionId, `${tabId}:export:${second.exportId}`);
assert.notEqual(first.clientSessionId, second.clientSessionId);
assert.notEqual(first.executionId, second.executionId);
await tracker.cancelTask(first.exportId);
assert.deepEqual(cancelRequests, [{ exportId: first.exportId, executionId: first.executionId }]);
assert.equal(tracker.tasks.value.find((task) => task.exportId === second.exportId)?.status, "Running");
FakeEventSource.instances[0].emitProgress({
exportId: first.exportId,
tableName: "",
rowsExported: 0,
totalRows: null,
status: "Cancelled",
errorMessage: "Export cancelled",
});
FakeEventSource.instances[1].emitProgress({
exportId: second.exportId,
tableName: "",
rowsExported: 0,
totalRows: null,
status: "Cancelled",
errorMessage: "Export cancelled",
});
await waitFor(() => tracker.tasks.value.filter((task) => task.exportId === first.exportId || task.exportId === second.exportId).every((task) => task.status === "Cancelled"));
} finally {
createdExportIds.forEach((exportId) => tracker.removeTask(exportId));
globalThis.fetch = originalFetch;
if (originalEventSource) Object.defineProperty(globalThis, "EventSource", originalEventSource);
else Reflect.deleteProperty(globalThis, "EventSource");
restoreStorage();
}
});
test("buildQueryResultExportRequest uses exportRowLimit when enabled", async () => {
const restoreStorage = installMemoryStorage();
setActivePinia(createPinia());

View File

@ -1,6 +1,5 @@
use std::path::PathBuf;
use std::sync::{
atomic::{AtomicBool, Ordering},
atomic::{AtomicU64, Ordering},
Arc, Mutex,
};
@ -10,6 +9,7 @@ use crate::commands::connection::AppState;
use dbx_core::query_cancel::RunningTaskMetadata;
pub use dbx_core::query_result_export::QueryResultExportRequest;
use dbx_core::query_result_export::StagedExportTarget;
use dbx_core::table_export::ExportStatus;
pub use dbx_core::table_export::TableExportProgress;
@ -17,120 +17,244 @@ fn emit_progress(app: &AppHandle, progress: TableExportProgress) {
let _ = app.emit("query-result-export-progress", progress);
}
#[derive(Default)]
struct RoutedExportProgress {
terminal: Mutex<Option<TableExportProgress>>,
rows_exported: AtomicU64,
}
impl RoutedExportProgress {
fn take_terminal(&self) -> Option<TableExportProgress> {
self.terminal.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).take()
}
fn rows_exported(&self) -> u64 {
self.rows_exported.load(Ordering::SeqCst)
}
}
fn route_core_progress(
progress: TableExportProgress,
deferred_done: &Mutex<Option<TableExportProgress>>,
cancelled: &AtomicBool,
routed: &RoutedExportProgress,
emit: impl FnOnce(TableExportProgress),
) {
routed.rows_exported.store(progress.rows_exported, Ordering::SeqCst);
match progress.status {
ExportStatus::Done => {
*deferred_done.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(progress);
let mut terminal = routed.terminal.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
if !terminal.as_ref().is_some_and(|progress| matches!(progress.status, ExportStatus::Cancelled)) {
*terminal = Some(progress);
}
}
ExportStatus::Cancelled => {
cancelled.store(true, Ordering::SeqCst);
emit(progress);
*routed.terminal.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(progress);
}
_ => emit(progress),
}
}
/// Build a temp-file path alongside the target path.
/// The temp file has a `.dbx-export-tmp` suffix so it's never confused with the
/// user's chosen file. On success the temp file is renamed atomically onto the
/// target; on error/cancel only the temp file is cleaned up, leaving the user's
/// chosen path untouched.
fn temp_file_path(target: &str) -> (PathBuf, PathBuf) {
let target_path = PathBuf::from(target);
let mut temp_path = target_path.clone();
let mut temp_name =
target_path.file_name().map(|name| name.to_os_string()).unwrap_or_else(|| std::ffi::OsString::from("export"));
temp_name.push(".dbx-export-tmp");
temp_path.set_file_name(temp_name);
(target_path, temp_path)
fn export_error_progress(export_id: &str, rows_exported: u64, error: String) -> TableExportProgress {
TableExportProgress {
export_id: export_id.to_string(),
table_name: String::new(),
rows_exported,
total_rows: None,
status: ExportStatus::Error,
error_message: Some(error),
}
}
fn finalize_staged_export(
target: StagedExportTarget,
export_id: &str,
result: Result<(), String>,
terminal: Option<TableExportProgress>,
cancellation_requested: bool,
rows_exported: u64,
emit: impl Fn(TableExportProgress),
) {
if cancellation_requested
|| terminal.as_ref().is_some_and(|progress| matches!(progress.status, ExportStatus::Cancelled))
{
drop(target);
let cancelled =
terminal.filter(|progress| matches!(progress.status, ExportStatus::Cancelled)).unwrap_or_else(|| {
TableExportProgress {
export_id: export_id.to_string(),
table_name: String::new(),
rows_exported,
total_rows: None,
status: ExportStatus::Cancelled,
error_message: Some("Export cancelled".to_string()),
}
});
emit(cancelled);
return;
}
if let Err(error) = result {
drop(target);
emit(export_error_progress(export_id, rows_exported, error));
return;
}
let Some(done) = terminal.filter(|progress| matches!(progress.status, ExportStatus::Done)) else {
drop(target);
emit(export_error_progress(
export_id,
rows_exported,
"Export finished without a completion status".to_string(),
));
return;
};
match target.commit() {
Ok(()) => emit(done),
Err(error) => emit(export_error_progress(export_id, rows_exported, error)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::{Path, PathBuf};
#[test]
fn temp_file_path_adds_suffix_to_simple_filename() {
let (target, temp) = temp_file_path("C:\\exports\\backup.sql");
assert_eq!(target, PathBuf::from("C:\\exports\\backup.sql"));
assert_eq!(temp, PathBuf::from("C:\\exports\\backup.sql.dbx-export-tmp"));
static NEXT_TEST_DIR: AtomicU64 = AtomicU64::new(0);
struct TestDir(PathBuf);
impl TestDir {
fn new() -> Self {
let path = std::env::temp_dir().join(format!(
"dbx-query-result-export-{}-{}",
std::process::id(),
NEXT_TEST_DIR.fetch_add(1, Ordering::SeqCst)
));
std::fs::create_dir_all(&path).expect("create test dir");
Self(path)
}
fn path(&self) -> &Path {
&self.0
}
}
#[test]
fn temp_file_path_adds_suffix_to_filename_with_multiple_extensions() {
let (target, temp) = temp_file_path("/home/user/query-result.xlsx");
assert_eq!(target, PathBuf::from("/home/user/query-result.xlsx"));
assert_eq!(temp, PathBuf::from("/home/user/query-result.xlsx.dbx-export-tmp"));
impl Drop for TestDir {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.0);
}
}
#[test]
fn temp_file_path_handles_filename_without_extension() {
let (target, temp) = temp_file_path("/tmp/export");
assert_eq!(target, PathBuf::from("/tmp/export"));
assert_eq!(temp, PathBuf::from("/tmp/export.dbx-export-tmp"));
}
#[test]
fn temp_file_path_uses_export_fallback_for_empty_target() {
let (target, temp) = temp_file_path("");
assert_eq!(target, PathBuf::from(""));
// file_name() on a path without a filename returns None, so we fall
// back to "export" as the base name.
assert!(temp.to_string_lossy().ends_with("export.dbx-export-tmp"));
}
#[test]
fn route_core_progress_defers_done_until_finalization() {
let deferred_done = Mutex::new(None);
let cancelled = AtomicBool::new(false);
let emitted = Mutex::new(Vec::new());
let done = TableExportProgress {
fn terminal_progress(status: ExportStatus, rows_exported: u64) -> TableExportProgress {
TableExportProgress {
export_id: "export-1".to_string(),
table_name: String::new(),
rows_exported: 42,
total_rows: Some(42),
status: ExportStatus::Done,
rows_exported,
total_rows: Some(rows_exported),
status,
error_message: None,
};
route_core_progress(done, &deferred_done, &cancelled, |progress| {
emitted.lock().unwrap().push(progress);
});
assert!(emitted.lock().unwrap().is_empty());
assert_eq!(deferred_done.lock().unwrap().as_ref().map(|progress| progress.rows_exported), Some(42));
assert!(!cancelled.load(Ordering::SeqCst));
}
}
#[test]
fn route_core_progress_keeps_cancelled_terminal_state() {
let deferred_done = Mutex::new(None);
let cancelled = AtomicBool::new(false);
fn route_core_progress_buffers_terminal_status_until_finalization() {
let routed = RoutedExportProgress::default();
let emitted = Mutex::new(Vec::new());
let progress = TableExportProgress {
export_id: "export-1".to_string(),
table_name: String::new(),
rows_exported: 7,
total_rows: None,
status: ExportStatus::Cancelled,
error_message: Some("Export cancelled".to_string()),
};
route_core_progress(progress, &deferred_done, &cancelled, |progress| {
emitted.lock().unwrap().push(progress);
route_core_progress(terminal_progress(ExportStatus::Done, 42), &routed, |progress| {
emitted.lock().expect("emitted lock").push(progress);
});
assert!(deferred_done.lock().unwrap().is_none());
assert!(cancelled.load(Ordering::SeqCst));
assert!(emitted.lock().expect("emitted lock").is_empty());
assert_eq!(routed.rows_exported(), 42);
assert!(matches!(routed.take_terminal(), Some(TableExportProgress { status: ExportStatus::Done, .. })));
}
fn assert_native_stream_cancellation_emits_one_cancelled_terminal() {
let dir = TestDir::new();
let destination = dir.path().join("result.csv");
std::fs::write(&destination, "original").expect("write destination");
let target = StagedExportTarget::new(destination.to_str().expect("destination path")).expect("target");
std::fs::write(target.path(), "partial").expect("write partial export");
let routed = RoutedExportProgress::default();
let emitted = Mutex::new(Vec::new());
route_core_progress(terminal_progress(ExportStatus::Cancelled, 7), &routed, |progress| {
emitted.lock().expect("emitted lock").push(progress);
});
route_core_progress(terminal_progress(ExportStatus::Done, 7), &routed, |progress| {
emitted.lock().expect("emitted lock").push(progress);
});
finalize_staged_export(
target,
"export-1",
Err("native stream cancelled".to_string()),
routed.take_terminal(),
true,
routed.rows_exported(),
|progress| emitted.lock().expect("emitted lock").push(progress),
);
assert!(matches!(
emitted.lock().unwrap().as_slice(),
emitted.lock().expect("emitted lock").as_slice(),
[TableExportProgress { status: ExportStatus::Cancelled, .. }]
));
assert_eq!(std::fs::read_to_string(destination).expect("read destination"), "original");
}
#[test]
fn postgres_native_stream_cancellation_emits_one_cancelled_terminal() {
assert_native_stream_cancellation_emits_one_cancelled_terminal();
}
#[test]
fn sqlserver_native_stream_cancellation_emits_one_cancelled_terminal() {
assert_native_stream_cancellation_emits_one_cancelled_terminal();
}
#[test]
fn export_failure_preserves_existing_destination_and_emits_one_error() {
let dir = TestDir::new();
let destination = dir.path().join("result.csv");
std::fs::write(&destination, "original").expect("write destination");
let target = StagedExportTarget::new(destination.to_str().expect("destination path")).expect("target");
std::fs::write(target.path(), "partial").expect("write partial export");
let emitted = Mutex::new(Vec::new());
finalize_staged_export(target, "export-1", Err("write failed".to_string()), None, false, 3, |progress| {
emitted.lock().expect("emitted lock").push(progress);
});
assert!(matches!(
emitted.lock().expect("emitted lock").as_slice(),
[TableExportProgress { status: ExportStatus::Error, .. }]
));
assert_eq!(std::fs::read_to_string(destination).expect("read destination"), "original");
}
#[test]
fn replace_failure_preserves_existing_destination_and_suppresses_done() {
let dir = TestDir::new();
let destination = dir.path().join("result.csv");
std::fs::write(&destination, "original").expect("write destination");
let target = StagedExportTarget::new(destination.to_str().expect("destination path")).expect("target");
std::fs::write(target.path(), "replacement").expect("write replacement");
std::fs::remove_file(target.path()).expect("remove staged path");
let emitted = Mutex::new(Vec::new());
finalize_staged_export(
target,
"export-1",
Ok(()),
Some(terminal_progress(ExportStatus::Done, 4)),
false,
4,
|progress| emitted.lock().expect("emitted lock").push(progress),
);
assert!(matches!(
emitted.lock().expect("emitted lock").as_slice(),
[TableExportProgress { status: ExportStatus::Error, .. }]
));
assert_eq!(std::fs::read_to_string(destination).expect("read destination"), "original");
}
}
@ -142,12 +266,8 @@ pub async fn start_query_result_export(
) -> Result<(), String> {
let state = state.inner().clone();
let export_id = request.export_id.clone();
// Redirect file I/O to a temp file so the user's chosen path is never
// truncated before the query completes. On success the temp file is
// renamed onto the target; on error/cancel only the temp file is removed.
let (target_path, temp_path) = temp_file_path(&request.file_path);
request.file_path = temp_path.to_string_lossy().to_string();
let target = StagedExportTarget::new(&request.file_path)?;
request.file_path = target.path_string()?;
tokio::spawn(async move {
let execution_id = request.execution_id.clone().filter(|id| !id.trim().is_empty());
@ -162,68 +282,34 @@ pub async fn start_query_result_export(
)
});
let cancel_token = registered_query.as_ref().map(|query| query.token());
let cancelled = Arc::new(AtomicBool::new(false));
let cancelled_progress = cancelled.clone();
let deferred_done = Arc::new(Mutex::new(None));
let deferred_done_progress = deferred_done.clone();
let result =
dbx_core::query_result_export::export_query_result_core(&state, &request, cancel_token, |progress| {
route_core_progress(progress, &deferred_done_progress, &cancelled_progress, |progress| {
let routed_progress = Arc::new(RoutedExportProgress::default());
let routed_progress_handler = routed_progress.clone();
let result = dbx_core::query_result_export::export_query_result_core(
&state,
&request,
cancel_token.clone(),
|progress| {
route_core_progress(progress, &routed_progress_handler, |progress| {
emit_progress(&app, progress);
});
})
.await;
},
)
.await;
drop(registered_query);
let completed_progress = deferred_done.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).take();
if let Err(e) = result {
let _ = tokio::fs::remove_file(&request.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(&request.file_path).await;
} else if let Some(progress) = completed_progress {
// Success: atomically rename temp → target so the user's chosen
// file is only created/replaced when the export fully completes.
if let Err(e) = std::fs::rename(&request.file_path, &target_path) {
let _ = tokio::fs::remove_file(&request.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(format!("Failed to finalize export file: {e}")),
},
);
} else {
emit_progress(&app, progress);
}
} else {
let _ = tokio::fs::remove_file(&request.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("Export finished without a completion status".to_string()),
},
);
}
let terminal = routed_progress.take_terminal();
let cancellation_requested =
terminal.as_ref().is_some_and(|progress| matches!(progress.status, ExportStatus::Cancelled))
|| cancel_token.as_ref().is_some_and(|token| token.is_cancelled())
|| dbx_core::database_export::is_export_cancelled(&export_id).await;
finalize_staged_export(
target,
&export_id,
result,
terminal,
cancellation_requested,
routed_progress.rows_exported(),
|progress| emit_progress(&app, progress),
);
dbx_core::database_export::clear_export_cancelled(&export_id).await;
});