diff --git a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts index d4093274bc..627e3fc891 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts @@ -6,6 +6,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; const mocks = vi.hoisted(() => ({ defaultAutoKeepResults: false, + continueOnErrorOnBatch: false, analyzeEditableQueryEditability: vi.fn(), cancelQuery: vi.fn(), closeClientConnectionSession: vi.fn(), @@ -44,7 +45,7 @@ vi.mock("@/stores/connectionStore", () => ({ vi.mock("@/stores/settingsStore", () => ({ useSettingsStore: () => ({ - editorSettings: { autoCalculateTotalRows: false, pageSize: 100, continueOnErrorOnBatch: false, defaultAutoKeepResults: mocks.defaultAutoKeepResults }, + editorSettings: { autoCalculateTotalRows: false, pageSize: 100, continueOnErrorOnBatch: mocks.continueOnErrorOnBatch, defaultAutoKeepResults: mocks.defaultAutoKeepResults }, }), })); @@ -115,6 +116,7 @@ describe("queryStore multi-statement errors", () => { beforeEach(() => { vi.clearAllMocks(); mocks.defaultAutoKeepResults = false; + mocks.continueOnErrorOnBatch = false; vi.unstubAllGlobals(); mocks.tabResultSnapshots.clear(); installLocalStorage(); @@ -496,6 +498,52 @@ describe("queryStore multi-statement errors", () => { await execution; }); + it("settles the statements between coalesced progress events without overwriting a failure", async () => { + mocks.continueOnErrorOnBatch = true; + const pendingExecution = deferred(); + let reportProgress!: (progress: any) => void; + mocks.executeMultiWithProgress.mockImplementationOnce((_connectionId, _database, _sql, onProgress) => { + reportProgress = onProgress; + return pendingExecution.promise; + }); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const sql = Array.from({ length: 6 }, (_, index) => `INSERT INTO t VALUES (${index + 1});`).join("\n"); + const tabId = store.createTab("mysql-1", "app", "Query", "query", undefined, sql); + const batch = () => store.tabs.find((item) => item.id === tabId)?.batchSqlExecution; + + const execution = store.executeTabSql(tabId, sql); + await vi.waitFor(() => expect(batch()?.items[0]?.status).toBe("running")); + const executionId = store.tabs.find((item) => item.id === tabId)!.executionId!; + const report = (statementIndex: number, success: boolean) => + reportProgress({ + executionId, + statementIndex, + completed: statementIndex + 1, + total: 6, + success, + executionTimeMs: 1, + affectedRows: success ? 1 : 0, + error: success ? undefined : structuredSqlError(), + }); + + report(1, true); + expect(batch()).toMatchObject({ completed: 2, items: [{ status: "success" }, { status: "success" }, { status: "running" }, { status: "pending" }, { status: "pending" }, { status: "pending" }] }); + report(3, false); + expect(batch()).toMatchObject({ completed: 4, items: [{ status: "success" }, { status: "success" }, { status: "success" }, { status: "error" }, { status: "running" }, { status: "pending" }] }); + report(5, true); + expect(batch()).toMatchObject({ completed: 6, items: [{ status: "success" }, { status: "success" }, { status: "success" }, { status: "error" }, { status: "success" }, { status: "success" }] }); + + pendingExecution.resolve( + Array.from({ length: 6 }, (_, index) => + index === 3 ? { columns: ["Error"], rows: [["duplicate key"]], affected_rows: 0, execution_time_ms: 1, statement_index: index, execution_error: true, error: structuredSqlError() } : { columns: [], rows: [], affected_rows: 1, execution_time_ms: 1, statement_index: index }, + ), + ); + await execution; + + expect(batch()).toMatchObject({ completed: 6, items: [{ status: "success" }, { status: "success" }, { status: "success" }, { status: "error" }, { status: "success" }, { status: "success" }] }); + }); + it("records a top-level batch failure on the current statement", async () => { mocks.executeMultiWithProgress.mockRejectedValueOnce(new Error("transport failed")); const { useQueryStore } = await import("@/stores/queryStore"); diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 3546660b24..df993bf925 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -2,13 +2,29 @@ import { UPDATE_RESTORE_KEY, assertUpdateAllowsInteraction } from "@/lib/app/upd import { defineStore } from "pinia"; import { isRedisMonitorCommand, startRedisMonitor } from "@/lib/redis/redisMonitor"; import { uuid } from "@/lib/common/utils"; -import { computed, markRaw, nextTick, onScopeDispose, reactive, ref, watch } from "vue"; +import { computed, markRaw, nextTick, onScopeDispose, reactive, ref, toRaw, watch } from "vue"; import { useI18n } from "vue-i18n"; import { useToast } from "@/composables/useToast"; import { savedSqlErrorMessage } from "@/lib/savedSql/savedSqlErrors"; import { sanitizeTabPageUiState } from "@/lib/tabs/tabUiState"; import type { DeletedConnectionTabKeepMode } from "@/lib/tabs/deletedConnectionTabs"; -import type { BatchSqlExecution, ConnectionConfig, DatabaseType, IndexInfo, NacosConfigEditorViewport, ObjectBrowserFilter, ObjectBrowserViewport, ObjectSource, ObjectSourceKind, QueryResult, QueryResultSourceColumnRef, QueryTab, TableInfoTab, TableStructureEditorTarget } from "@/types/database"; +import type { + BatchSqlExecution, + BatchStatementExecutionItem, + ConnectionConfig, + DatabaseType, + IndexInfo, + NacosConfigEditorViewport, + ObjectBrowserFilter, + ObjectBrowserViewport, + ObjectSource, + ObjectSourceKind, + QueryResult, + QueryResultSourceColumnRef, + QueryTab, + TableInfoTab, + TableStructureEditorTarget, +} from "@/types/database"; import { orderPinnedFirst } from "@/lib/app/pinnedItems"; import { canCancelQueryExecution } from "@/lib/sql/queryExecutionState"; import { isSqlErrorPositionDebugEnabled, logSqlErrorPosition, sqlErrorHasMessagePosition, sqlErrorMessageText } from "@/lib/sql/errorPosition"; @@ -444,6 +460,13 @@ function annotateSingleStatementErrorResult(errorResult: QueryResult, sourceSql: const NON_STREAMING_BATCH_DATABASE_TYPES = new Set(["sqlserver", "turso", "cloudflare-d1"]); const liveBatchSqlExecutions = new WeakMap(); +// Per running batch: statements before this index are already settled by an +// earlier progress event of the same run. +const batchSqlProgressSettledEnds = new WeakMap(); + +function isCompletedBatchStatement(item: BatchStatementExecutionItem): boolean { + return item.status === "success" || item.status === "error"; +} function cloneBatchSqlExecution(batch: BatchSqlExecution | undefined): BatchSqlExecution | undefined { return batch ? { ...batch, executionTarget: batch.executionTarget ? { ...batch.executionTarget } : undefined, items: batch.items.map((item) => ({ ...item })) } : undefined; @@ -503,20 +526,29 @@ function applyBatchSqlProgress( const statementIndex = statementOffset + progress.statementIndex; const item = batch.items[statementIndex]; if (!item) return; - if (progress.completed > 1) { - for (let index = statementOffset; index < statementOffset + progress.completed - 1; index += 1) { - const completedItem = batch.items[index]; - if (completedItem && (completedItem.status === "pending" || completedItem.status === "running")) { - completedItem.status = "success"; - } + // Statements complete in order and the desktop backend coalesces successful + // ones, so an event also settles every statement since the previous event. + // Walking only that gap and counting incrementally keeps a batch linear; + // rescanning every statement per event made large scripts O(N²). + const rawBatch = toRaw(batch); + const settledStart = Math.max(statementOffset, batchSqlProgressSettledEnds.get(rawBatch) ?? statementOffset); + const settledEnd = statementOffset + progress.completed - 1; + let completed = batch.completed; + for (let index = settledStart; index < settledEnd; index += 1) { + const completedItem = batch.items[index]; + if (completedItem && (completedItem.status === "pending" || completedItem.status === "running")) { + completedItem.status = "success"; + completed += 1; } } + if (settledEnd > settledStart) batchSqlProgressSettledEnds.set(rawBatch, settledEnd); + if (!isCompletedBatchStatement(item)) completed += 1; item.status = progress.success ? "success" : "error"; item.executionTimeMs = progress.executionTimeMs; item.affectedRows = progress.affectedRows; item.errorDetails = progress.error; item.error = progress.error ? translateBackendError(i18n.global.t, progress.error) : undefined; - batch.completed = batch.items.filter((candidate) => candidate.status === "success" || candidate.status === "error").length; + batch.completed = completed; if ((progress.success || continueOnError) && progress.completed < progress.total) { const next = batch.items[statementOffset + progress.completed]; if (next?.status === "pending") next.status = "running"; @@ -539,7 +571,7 @@ function reconcileBatchSqlResults(tab: QueryTab, executionId: string, results: Q item.errorDetails = failed ? result.error : undefined; item.error = failed ? (result.error ? translateBackendError(i18n.global.t, result.error, result.rows[0]?.[0]) : String(result.rows[0]?.[0] ?? "")) : undefined; } - batch.completed = batch.items.filter((item) => item.status === "success" || item.status === "error").length; + batch.completed = batch.items.filter(isCompletedBatchStatement).length; } function failBatchSqlExecution(tab: QueryTab, executionId: string, error: unknown, cancelled: boolean) { @@ -550,7 +582,7 @@ function failBatchSqlExecution(tab: QueryTab, executionId: string, error: unknow item.status = cancelled ? "cancelled" : "error"; item.errorDetails = cancelled ? undefined : (normalizeBackendError(error) ?? undefined); item.error = cancelled ? undefined : translateBackendError(i18n.global.t, error, error instanceof Error ? error.message : undefined); - batch.completed = batch.items.filter((candidate) => candidate.status === "success" || candidate.status === "error").length; + batch.completed = batch.items.filter(isCompletedBatchStatement).length; } function finishBatchSqlExecution(tab: QueryTab, executionId: string, cancelled: boolean) { @@ -572,7 +604,7 @@ function finishBatchSqlExecution(tab: QueryTab, executionId: string, cancelled: item.status = "skipped"; } } - batch.completed = batch.items.filter((item) => item.status === "success" || item.status === "error").length; + batch.completed = batch.items.filter(isCompletedBatchStatement).length; batch.finishedAt = Date.now(); } diff --git a/crates/dbx-core/src/query/batch_progress.rs b/crates/dbx-core/src/query/batch_progress.rs new file mode 100644 index 0000000000..35ff0037ec --- /dev/null +++ b/crates/dbx-core/src/query/batch_progress.rs @@ -0,0 +1,189 @@ +use std::future::Future; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use super::{ExecuteMultiProgress, ExecuteMultiProgressCallback}; + +/// Same cadence as the SQL file import progress. +const EXECUTE_MULTI_PROGRESS_EMIT_INTERVAL: Duration = Duration::from_millis(100); + +/// Runs a SQL batch while forwarding its per-statement progress to `emit` at +/// most once per interval. +/// +/// The core reports every completed statement, but forwarding each one across +/// the desktop event bridge floods the webview on scripts with thousands of +/// statements, and the UI re-derives the whole batch on every event. Successful +/// statements are therefore coalesced into the latest one: statements complete +/// in order, so the UI settles every statement up to it. A failure cannot be +/// inferred from a later event, so it is forwarded immediately, after the +/// successes that preceded it. Whatever is still pending when `run` finishes is +/// forwarded before returning, so the UI knows the last completed statement +/// before the caller reports the batch outcome. +pub async fn with_coalesced_execute_multi_progress(emit: Option, run: F) -> T +where + F: FnOnce(Option) -> Fut, + Fut: Future, +{ + let Some(emit) = emit else { + return run(None).await; + }; + let coalescer = Arc::new(ProgressCoalescer { emit, pending: Mutex::new(None) }); + let progress: ExecuteMultiProgressCallback = { + let coalescer = Arc::clone(&coalescer); + Arc::new(move |progress| coalescer.report(progress)) + }; + let mut run = std::pin::pin!(run(Some(progress))); + loop { + let output = tokio::time::timeout(EXECUTE_MULTI_PROGRESS_EMIT_INTERVAL, run.as_mut()).await; + coalescer.flush(); + if let Ok(output) = output { + return output; + } + } +} + +struct ProgressCoalescer { + emit: ExecuteMultiProgressCallback, + pending: Mutex>, +} + +impl ProgressCoalescer { + // Both paths emit while holding the lock so forwarded events stay in + // statement order. + fn report(&self, progress: ExecuteMultiProgress) { + let mut pending = self.pending.lock().unwrap_or_else(|e| e.into_inner()); + if progress.success { + *pending = Some(progress); + return; + } + if let Some(previous) = pending.take() { + (self.emit)(previous); + } + (self.emit)(progress); + } + + fn flush(&self) { + let mut pending = self.pending.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(progress) = pending.take() { + (self.emit)(progress); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend_error::BackendError; + use tokio::time::{sleep, Instant}; + + type Emitted = Arc>>; + + fn recorder() -> (ExecuteMultiProgressCallback, Emitted) { + let emitted: Emitted = Arc::new(Mutex::new(Vec::new())); + let started_at = Instant::now(); + let callback: ExecuteMultiProgressCallback = { + let emitted = Arc::clone(&emitted); + Arc::new(move |progress| emitted.lock().unwrap().push((progress, started_at.elapsed()))) + }; + (callback, emitted) + } + + fn statement(statement_index: usize, total: usize, success: bool) -> ExecuteMultiProgress { + ExecuteMultiProgress { + statement_index, + completed: statement_index + 1, + total, + success, + execution_time_ms: 1, + affected_rows: 1, + error: (!success).then(|| BackendError::from_sql_detail("Duplicate entry")), + } + } + + fn completed(emitted: &Emitted) -> Vec { + emitted.lock().unwrap().iter().map(|(progress, _)| progress.completed).collect() + } + + #[tokio::test(start_paused = true)] + async fn forwards_at_most_one_success_per_interval() { + let (emit, emitted) = recorder(); + + with_coalesced_execute_multi_progress(Some(emit), |progress| async move { + let progress = progress.unwrap(); + for index in 0..1000 { + progress(statement(index, 1000, true)); + sleep(Duration::from_millis(1)).await; + } + }) + .await; + + let completed = completed(&emitted); + assert!(completed.len() <= 11, "forwarded {} events for one second of statements", completed.len()); + assert!(completed.windows(2).all(|pair| pair[0] < pair[1])); + assert_eq!(completed.last(), Some(&1000)); + } + + #[tokio::test(start_paused = true)] + async fn forwards_failures_immediately_after_pending_successes() { + let (emit, emitted) = recorder(); + + with_coalesced_execute_multi_progress(Some(emit), |progress| async move { + let progress = progress.unwrap(); + progress(statement(0, 4, true)); + progress(statement(1, 4, true)); + progress(statement(2, 4, false)); + progress(statement(3, 4, false)); + sleep(Duration::from_millis(30)).await; + }) + .await; + + let emitted = emitted.lock().unwrap(); + assert_eq!( + emitted.iter().map(|(progress, at)| (progress.completed, progress.success, *at)).collect::>(), + vec![(2, true, Duration::ZERO), (3, false, Duration::ZERO), (4, false, Duration::ZERO)] + ); + assert_eq!(emitted[1].0.error, Some(BackendError::from_sql_detail("Duplicate entry"))); + } + + #[tokio::test(start_paused = true)] + async fn forwards_pending_progress_while_a_later_statement_is_still_running() { + let (emit, emitted) = recorder(); + + with_coalesced_execute_multi_progress(Some(emit), |progress| async move { + let progress = progress.unwrap(); + progress(statement(0, 2, true)); + sleep(Duration::from_secs(30)).await; + progress(statement(1, 2, true)); + }) + .await; + + let emitted = emitted.lock().unwrap(); + assert_eq!(emitted.len(), 2); + assert_eq!((emitted[0].0.completed, emitted[0].1), (1, EXECUTE_MULTI_PROGRESS_EMIT_INTERVAL)); + assert_eq!((emitted[1].0.completed, emitted[1].1), (2, Duration::from_secs(30))); + } + + #[tokio::test(start_paused = true)] + async fn forwards_the_last_completed_statement_before_returning() { + let (emit, emitted) = recorder(); + + let result: Result<(), &str> = with_coalesced_execute_multi_progress(Some(emit), |progress| async move { + let progress = progress.unwrap(); + progress(statement(0, 3, true)); + progress(statement(1, 3, true)); + Err("connection lost") + }) + .await; + + assert_eq!(result, Err("connection lost")); + assert_eq!(completed(&emitted), vec![2]); + } + + #[tokio::test] + async fn runs_without_a_progress_callback_when_nobody_listens() { + let received_callback = + with_coalesced_execute_multi_progress(None, |progress| async move { progress.is_some() }).await; + + assert!(!received_callback); + } +} diff --git a/crates/dbx-core/src/query/mod.rs b/crates/dbx-core/src/query/mod.rs index c9dc8f3e7d..11f079f37f 100644 --- a/crates/dbx-core/src/query/mod.rs +++ b/crates/dbx-core/src/query/mod.rs @@ -1,3 +1,4 @@ +pub mod batch_progress; pub mod document_ops; pub mod hbase_ops; pub mod mongo_ops; diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index df3ffb4432..79b7b56c3e 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -233,34 +233,36 @@ pub async fn execute_multi( schema ); - let result = dbx_core::query::execute_multi_core_with_options_for_client_and_progress_typed( - &state, - &connection_id, - &database, - &sql, - schema.as_deref(), - cancel_token, - dbx_core::query::QueryExecutionOptions { - max_rows, - fetch_size, - page_size, - row_offset, - max_result_bytes, - result_key_columns: result_key_columns.unwrap_or_default(), - table_data_preview: table_data_preview.unwrap_or(false), - catalog, - result_session_id, - client_session_id, - timeout_secs, - await_cancel_completion: false, - execution_id, - use_transaction, - continue_on_error: continue_on_error.unwrap_or(false), - execution_mode: execution_mode.unwrap_or_default(), - preserve_explicit_transaction: preserve_explicit_transaction.unwrap_or(false), - }, - progress, - ) + let result = dbx_core::query::batch_progress::with_coalesced_execute_multi_progress(progress, |progress| { + dbx_core::query::execute_multi_core_with_options_for_client_and_progress_typed( + &state, + &connection_id, + &database, + &sql, + schema.as_deref(), + cancel_token, + dbx_core::query::QueryExecutionOptions { + max_rows, + fetch_size, + page_size, + row_offset, + max_result_bytes, + result_key_columns: result_key_columns.unwrap_or_default(), + table_data_preview: table_data_preview.unwrap_or(false), + catalog, + result_session_id, + client_session_id, + timeout_secs, + await_cancel_completion: false, + execution_id, + use_transaction, + continue_on_error: continue_on_error.unwrap_or(false), + execution_mode: execution_mode.unwrap_or_default(), + preserve_explicit_transaction: preserve_explicit_transaction.unwrap_or(false), + }, + progress, + ) + }) .await; match &result { Ok(results) => log::info!(