diff --git a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts index 74b4780b4..ae215a834 100644 --- a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts +++ b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts @@ -2,6 +2,7 @@ import { computed, ref } from "vue"; import { createPinia, setActivePinia } from "pinia"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { requiresDatabaseSelection, useSqlExecution } from "../useSqlExecution"; +import { useConnectionStore } from "@/stores/connectionStore"; import { useHistoryStore } from "@/stores/historyStore"; import { useQueryStore } from "@/stores/queryStore"; import { useSettingsStore } from "@/stores/settingsStore"; @@ -140,6 +141,88 @@ describe("useSqlExecution", () => { expect(executedSql).toContain("where fp.create_at < @date_start"); }); + it("records a later MySQL batch error and skips metadata refresh", async () => { + const activeTab = ref(queryTab("app")); + const activeConnection = ref(connection("mysql")); + const activeOutputView = ref<"result" | "summary" | "explain" | "chart">("result"); + const queryStore = useQueryStore(); + const historyStore = useHistoryStore(); + const connectionStore = useConnectionStore(); + vi.spyOn(queryStore, "executeCurrentSql").mockImplementation(async () => { + const tab = activeTab.value; + if (!tab) return; + const successfulResult = { columns: ["value"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }; + tab.result = successfulResult; + tab.results = [successfulResult, { columns: ["Error"], execution_error: true, rows: [["Duplicate entry '1'"]], affected_rows: 0, execution_time_ms: 1 }]; + tab.activeResultIndex = 0; + }); + const addHistory = vi.spyOn(historyStore, "add").mockResolvedValue(undefined); + const refreshObjects = vi.spyOn(connectionStore, "refreshObjectListTreeNode").mockResolvedValue(undefined); + + const execution = useSqlExecution({ + activeTab: computed(() => activeTab.value), + activeConnection: computed(() => activeConnection.value), + executableSql: computed(() => "SELECT 1 AS value; CREATE TABLE duplicate_target (id INT)"), + activeOutputView, + }); + + await execution.tryExecute(); + + expect(addHistory).toHaveBeenCalledWith(expect.objectContaining({ success: false, error: "Duplicate entry '1'", affected_rows: undefined })); + expect(refreshObjects).not.toHaveBeenCalled(); + }); + + it("does not treat an unmarked MySQL Error alias as a batch failure", async () => { + const activeTab = ref(queryTab("app")); + const activeConnection = ref(connection("mysql")); + const activeOutputView = ref<"result" | "summary" | "explain" | "chart">("result"); + const queryStore = useQueryStore(); + const historyStore = useHistoryStore(); + vi.spyOn(queryStore, "executeCurrentSql").mockImplementation(async () => { + const tab = activeTab.value; + if (!tab) return; + const successfulResult = { columns: ["value"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }; + tab.result = successfulResult; + tab.results = [successfulResult, { columns: ["Error"], rows: [[2]], affected_rows: 0, execution_time_ms: 1 }]; + tab.activeResultIndex = 0; + }); + const addHistory = vi.spyOn(historyStore, "add").mockResolvedValue(undefined); + + const execution = useSqlExecution({ + activeTab: computed(() => activeTab.value), + activeConnection: computed(() => activeConnection.value), + executableSql: computed(() => "SELECT 1 AS value; SELECT 2 AS Error"), + activeOutputView, + }); + + await execution.tryExecute(); + + expect(addHistory).toHaveBeenCalledWith(expect.objectContaining({ success: true, error: undefined })); + }); + + it("continues to record active non-MySQL errors as failures", async () => { + const activeTab = ref(queryTab("app")); + const activeConnection = ref(connection("postgres")); + const activeOutputView = ref<"result" | "summary" | "explain" | "chart">("result"); + const queryStore = useQueryStore(); + const historyStore = useHistoryStore(); + vi.spyOn(queryStore, "executeCurrentSql").mockImplementation(async () => { + if (activeTab.value) activeTab.value.result = { columns: ["Error"], rows: [["relation does not exist"]], affected_rows: 0, execution_time_ms: 1 }; + }); + const addHistory = vi.spyOn(historyStore, "add").mockResolvedValue(undefined); + + const execution = useSqlExecution({ + activeTab: computed(() => activeTab.value), + activeConnection: computed(() => activeConnection.value), + executableSql: computed(() => "SELECT * FROM missing_table"), + activeOutputView, + }); + + await execution.tryExecute(); + + expect(addHistory).toHaveBeenCalledWith(expect.objectContaining({ success: false, error: "relation does not exist" })); + }); + it("requires production confirmation even when ordinary danger prompts are disabled", async () => { const activeTab = ref(queryTab("prod_app")); const activeConnection = ref({ ...connection("mysql"), production_databases: ["prod_app"] }); diff --git a/apps/desktop/src/composables/useSqlExecution.ts b/apps/desktop/src/composables/useSqlExecution.ts index 1db076f18..5fad05fda 100644 --- a/apps/desktop/src/composables/useSqlExecution.ts +++ b/apps/desktop/src/composables/useSqlExecution.ts @@ -9,6 +9,7 @@ import { isSingleDatabase, usesTreeSchemaMode } from "@/lib/database/databaseCap import { canExecuteWithoutSelectedDatabase } from "@/lib/connection/connectionLevelDatabaseBootstrap"; import { classifySqlActivityKind } from "@/lib/history/historyActivityKind"; import { sqlMetadataRefreshTarget } from "@/lib/sql/sqlMetadataRefresh"; +import { isMysqlExecutionErrorResult, usesMysqlProtocolDatabaseType } from "@/lib/query/queryResultError"; import { classifyRedisCommandSafety, firstRedisCommandToken } from "@/lib/redis/redisCommandSafety"; import { isSqlExecutionSnapshot, resolveExecutableSql, type SqlExecutionOverride, type SqlExecutionSnapshot } from "@/lib/sql/sqlExecutionTarget"; import { extractSqlParameterDescriptors, type SqlParameterDescriptor, type SqlParameterSyntax } from "@/lib/sql/sqlParameters"; @@ -41,6 +42,16 @@ function primarySqlOperation(sql: string): string { return statement?.match(/^([a-z]+)/i)?.[1]?.toUpperCase() || "SQL"; } +function firstQueryExecutionError(tab: Pick, databaseType: DatabaseType | undefined) { + const activeResult = tab.result; + if (activeResult && isMysqlExecutionErrorResult(activeResult, databaseType)) return activeResult; + if (!usesMysqlProtocolDatabaseType(databaseType) && activeResult?.columns.includes("Error")) return activeResult; + if (!usesMysqlProtocolDatabaseType(databaseType)) return undefined; + + const results = tab.results?.length ? tab.results : tab.result ? [tab.result] : []; + return results.find((result) => isMysqlExecutionErrorResult(result, databaseType)); +} + export function useSqlExecution(deps: { activeTab: ComputedRef; activeConnection: ComputedRef; @@ -146,20 +157,23 @@ export function useSqlExecution(deps: { sql ??= await resolvedExecutableSql(); const tab = deps.activeTab.value; if (!tab || !sql.trim()) return; - if (requiresDatabaseSelection(tab, deps.activeConnection.value, sql)) { + const executionConnection = connectionStore.getConfig(tab.connectionId) ?? deps.activeConnection.value; + const executionDatabaseType = executionConnection?.db_type; + if (requiresDatabaseSelection(tab, executionConnection, sql)) { deps.onMissingDatabase?.(); return; } deps.activeOutputView.value = "result"; - const connName = connectionStore.getConfig(tab.connectionId)?.name || ""; + const connName = executionConnection?.name || ""; const start = Date.now(); - const isRedis = deps.activeConnection.value?.db_type === "redis"; + const isRedis = executionDatabaseType === "redis"; await queryStore.executeCurrentSql(sql, isRedis ? { skipRedisSafetyCheck: deps.blockDangerousRedisCommands?.value === false } : undefined); if (tab.result && !tab.result.columns.length && !tab.results?.some((result) => result.columns.length > 0)) { deps.activeOutputView.value = "summary"; } const elapsed = Date.now() - start; - const success = !tab.result?.columns.includes("Error"); + const failure = firstQueryExecutionError(tab, executionDatabaseType); + const success = !failure; historyStore.add({ connection_id: tab.connectionId, connection_name: connName, @@ -167,7 +181,7 @@ export function useSqlExecution(deps: { sql, execution_time_ms: elapsed, success, - error: success ? undefined : String(tab.result?.rows?.[0]?.[0] ?? ""), + error: failure ? String(failure.rows?.[0]?.[0] ?? "") : undefined, activity_kind: classifySqlActivityKind(sql), operation: primarySqlOperation(sql), affected_rows: success ? tab.result?.affected_rows : undefined, diff --git a/apps/desktop/src/lib/__tests__/query/queryResultError.spec.ts b/apps/desktop/src/lib/__tests__/query/queryResultError.spec.ts index de64be64a..4d848da3b 100644 --- a/apps/desktop/src/lib/__tests__/query/queryResultError.spec.ts +++ b/apps/desktop/src/lib/__tests__/query/queryResultError.spec.ts @@ -2,7 +2,7 @@ import { describe, expect, it } from "vitest"; import type { QueryResult } from "@/types/database"; -import { isNoSnapshotErrorResult } from "@/lib/query/queryResultError"; +import { isMysqlExecutionErrorResult, isNoSnapshotErrorResult } from "@/lib/query/queryResultError"; function errorResult(message: string): QueryResult { return { columns: ["Error"], rows: [[message]], affected_rows: 0, execution_time_ms: 0 }; @@ -39,3 +39,23 @@ describe("isNoSnapshotErrorResult", () => { expect(isNoSnapshotErrorResult({ ...errorResult("There is currently no snapshot."), rows: [] })).toBe(false); }); }); + +describe("isMysqlExecutionErrorResult", () => { + it("recognizes an explicitly marked MySQL batch execution error", () => { + expect(isMysqlExecutionErrorResult({ ...errorResult("Duplicate entry '1'"), execution_error: true }, "mysql")).toBe(true); + }); + + it("does not mistake an unmarked Error alias without type metadata for an execution error", () => { + const result: QueryResult = { + columns: ["Error"], + rows: [["2"]], + affected_rows: 0, + execution_time_ms: 1, + }; + expect(isMysqlExecutionErrorResult(result, "mysql")).toBe(false); + }); + + it("does not apply the native MySQL heuristic to JDBC connections", () => { + expect(isMysqlExecutionErrorResult({ ...errorResult("Duplicate entry '1'"), execution_error: true }, "jdbc")).toBe(false); + }); +}); diff --git a/apps/desktop/src/lib/query/queryResultError.ts b/apps/desktop/src/lib/query/queryResultError.ts index e15d0b999..43d6a4ad0 100644 --- a/apps/desktop/src/lib/query/queryResultError.ts +++ b/apps/desktop/src/lib/query/queryResultError.ts @@ -1,4 +1,4 @@ -import type { QueryResult } from "@/types/database"; +import type { DatabaseType, QueryResult } from "@/types/database"; // Lake/external tables (e.g. Paimon in StarRocks) return this error on a data // read when no snapshot exists yet, while metadata reads (DESC/SHOW CREATE) @@ -6,6 +6,17 @@ import type { QueryResult } from "@/types/database"; // callers can detect this case and fall back to a structure-only (LIMIT 0) // preview instead of showing a cryptic server error. const NO_SNAPSHOT_ERROR_PATTERN = /there is currently no snapshot/i; +const MYSQL_PROTOCOL_DATABASE_TYPES = new Set(["mysql", "doris", "starrocks", "manticoresearch"]); + +export function usesMysqlProtocolDatabaseType(databaseType: DatabaseType | undefined): boolean { + return databaseType !== undefined && MYSQL_PROTOCOL_DATABASE_TYPES.has(databaseType); +} + +// The batch executor marks synthesized MySQL-protocol errors explicitly so a +// successful result column named Error is never mistaken for a failure. +export function isMysqlExecutionErrorResult(result: QueryResult, databaseType: DatabaseType | undefined): boolean { + return usesMysqlProtocolDatabaseType(databaseType) && result.execution_error === true; +} export function isNoSnapshotErrorResult(result: QueryResult | undefined | null): boolean { if (!result || !result.columns.includes("Error") || result.rows.length === 0) return false; diff --git a/apps/desktop/src/lib/tabs/tabResultCache.ts b/apps/desktop/src/lib/tabs/tabResultCache.ts index 4a056531b..6c62e2d68 100644 --- a/apps/desktop/src/lib/tabs/tabResultCache.ts +++ b/apps/desktop/src/lib/tabs/tabResultCache.ts @@ -40,6 +40,7 @@ export interface TabResultSnapshot { interface ColumnarQueryResult { columns: string[]; + execution_error?: true; column_types?: string[]; columnValues: CellValue[][]; rowCount: number; @@ -134,6 +135,7 @@ function stripSessionIds(result: QueryResult | undefined): QueryResult | undefin if (!result) return undefined; return { columns: [...result.columns], + execution_error: result.execution_error, column_types: result.column_types ? [...result.column_types] : undefined, rows: result.rows.map((row) => [...row]), mongo_documents: result.mongo_documents ? clonePlain(result.mongo_documents) : undefined, @@ -167,6 +169,7 @@ function toColumnarResult(result: QueryResult | undefined): ColumnarQueryResult const columnValues = result.columns.map((_, colIndex) => result.rows.map((row) => row[colIndex] ?? null)); return removeUndefinedFields({ columns: [...result.columns], + execution_error: result.execution_error, column_types: result.column_types ? [...result.column_types] : undefined, columnValues, rowCount: result.rows.length, @@ -185,6 +188,7 @@ function fromColumnarResult(result: ColumnarQueryResult | undefined): QueryResul const rows = Array.from({ length: result.rowCount }, (_, rowIndex) => result.columnValues.map((values) => values[rowIndex] ?? null)); return { columns: [...result.columns], + execution_error: result.execution_error, column_types: result.column_types ? [...result.column_types] : undefined, rows, mongo_documents: result.mongo_documents ? clonePlain(result.mongo_documents) : undefined, diff --git a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts new file mode 100644 index 000000000..bb758e790 --- /dev/null +++ b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts @@ -0,0 +1,125 @@ +import { createPinia, setActivePinia } from "pinia"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + analyzeEditableQueryEditability: vi.fn(), + closeClientConnectionSession: vi.fn(), + closeQuerySession: vi.fn(), + executeMulti: vi.fn(), + getConnectionConfig: vi.fn(), + prepareQueryPaginationExecutionPlan: vi.fn(), + saveOpenTabsState: vi.fn(), +})); + +vi.mock("@/lib/backend/api", () => ({ + analyzeEditableQueryEditability: mocks.analyzeEditableQueryEditability, + closeClientConnectionSession: mocks.closeClientConnectionSession, + closeQuerySession: mocks.closeQuerySession, + executeMulti: mocks.executeMulti, + prepareQueryPaginationExecutionPlan: mocks.prepareQueryPaginationExecutionPlan, + saveOpenTabsState: mocks.saveOpenTabsState, +})); + +vi.mock("@/stores/connectionStore", () => ({ + useConnectionStore: () => ({ + ensureConnected: vi.fn().mockResolvedValue(undefined), + getConfig: mocks.getConnectionConfig, + recordConnectionLostError: vi.fn(), + }), +})); + +vi.mock("@/stores/settingsStore", () => ({ + useSettingsStore: () => ({ + editorSettings: { autoCalculateTotalRows: false, pageSize: 100 }, + }), +})); + +function installLocalStorage() { + const data = new Map(); + vi.stubGlobal("localStorage", { + getItem: vi.fn((key: string) => data.get(key) ?? null), + setItem: vi.fn((key: string, value: string) => data.set(key, value)), + removeItem: vi.fn((key: string) => data.delete(key)), + }); +} + +describe("queryStore multi-statement errors", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.unstubAllGlobals(); + installLocalStorage(); + setActivePinia(createPinia()); + mocks.getConnectionConfig.mockReturnValue({ + id: "mysql-1", + name: "MySQL", + db_type: "mysql", + database: "app", + query_timeout_secs: 30, + }); + mocks.prepareQueryPaginationExecutionPlan.mockImplementation(async (options) => ({ + sqlToExecute: options.sql, + pageSql: undefined, + pageLimit: undefined, + pageOffset: undefined, + countSql: undefined, + useAgentResultSession: false, + })); + mocks.analyzeEditableQueryEditability.mockResolvedValue({ editable: false, reason: "multiple-statements" }); + }); + + it("opens the first error result from a mixed result batch", async () => { + mocks.executeMulti.mockResolvedValue([ + { columns: ["value"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }, + { columns: ["Error"], execution_error: true, rows: [["no such table: missing"]], affected_rows: 0, execution_time_ms: 1 }, + ]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT 1 AS value; SELECT * FROM missing"); + + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.activeResultIndex).toBe(1); + expect(tab.result?.columns).toEqual(["Error"]); + }); + + it("does not promote an unmarked Error alias without type metadata as a batch failure", async () => { + mocks.executeMulti.mockResolvedValue([ + { columns: ["value"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }, + { columns: ["Error"], rows: [[2]], affected_rows: 0, execution_time_ms: 1 }, + ]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT 1 AS value; SELECT 2 AS Error"); + + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.activeResultIndex).toBe(0); + expect(tab.result?.columns).toEqual(["value"]); + }); + + it("does not apply the MySQL result heuristic to a JDBC MySQL dialect", async () => { + mocks.getConnectionConfig.mockReturnValue({ + id: "mysql-1", + name: "JDBC MySQL", + db_type: "jdbc", + connection_string: "jdbc:mysql://localhost:3306/app", + database: "app", + query_timeout_secs: 30, + }); + mocks.executeMulti.mockResolvedValue([ + { columns: ["value"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }, + { columns: ["Error"], rows: [[2]], affected_rows: 0, execution_time_ms: 1 }, + ]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT 1 AS value; SELECT 2 AS Error"); + + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.activeResultIndex).toBe(0); + expect(tab.result?.columns).toEqual(["value"]); + }); +}); diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index f221101b9..71fb51deb 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -45,6 +45,7 @@ import { externalSqlFileDisplayTitles, normalizeExternalSqlPath } from "@/lib/sq import { clearDataGridPendingSnapshotsForTab } from "@/composables/useDataGridEditor"; import { buildTabResultSnapshot, deleteTabResultSnapshot, readTabResultSnapshot, tabResultCacheKey, writeTabResultSnapshot } from "@/lib/tabs/tabResultCache"; import { queryResultBaseSql, queryResultExecutionSql } from "@/lib/tabs/tabPresentation"; +import { isMysqlExecutionErrorResult } from "@/lib/query/queryResultError"; import { decodeQueryResultArchive, encodeQueryResultArchive, type DecodedQueryResultArchive } from "@/lib/query/queryResultArchive"; import * as api from "@/lib/backend/api"; import { useConnectionStore } from "@/stores/connectionStore"; @@ -1982,6 +1983,7 @@ export const useQueryStore = defineStore("query", () => { const message = e instanceof Error ? e.message : String(e); return markQueryResultRowsRaw({ columns: ["Error"], + execution_error: true, rows: [[message]], affected_rows: 0, execution_time_ms: 0, @@ -2920,8 +2922,9 @@ export const useQueryStore = defineStore("query", () => { current.results[activeGroupIndex] = results[0]; current.result = results[0]; } else if (results.length > 1) { + const errorResultIndex = results.findIndex((result) => isMysqlExecutionErrorResult(result, conn?.db_type)); const activeResultIndex = results.findIndex((result) => result.columns.length > 0); - const resultIndex = preservedResultIndex(results, current.activeResultIndex, options?.preserveActiveResultIndex) ?? (activeResultIndex >= 0 ? activeResultIndex : 0); + const resultIndex = errorResultIndex >= 0 ? errorResultIndex : (preservedResultIndex(results, current.activeResultIndex, options?.preserveActiveResultIndex) ?? (activeResultIndex >= 0 ? activeResultIndex : 0)); current.results = results; current.activeResultIndex = resultIndex; current.result = results[resultIndex]; diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index 607b8b6ca..d8d5a00ab 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -449,6 +449,8 @@ export interface OwnerInfo { export interface QueryResult { columns: string[]; + /** Set for synthesized query execution failures. */ + execution_error?: true; /** Internal row identifiers appended to editable query results. */ hidden_column_indexes?: number[]; /** diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 56a9c1b07..54ed71197 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -4,6 +4,7 @@ use chrono::{DateTime, Duration as ChronoDuration, NaiveDate, NaiveDateTime, Nai use duckdb::types::{TimeUnit, Value, ValueRef}; use futures::StreamExt; use mysql_async::prelude::Queryable; +use serde::Serialize; use sqlparser::ast::{visit_relations_mut, Ident, ObjectName, ObjectNamePart, ObjectType, Statement}; use sqlparser::dialect::{GenericDialect, PostgreSqlDialect}; use sqlparser::parser::Parser; @@ -43,6 +44,38 @@ pub enum PoolErrorAction { ReconnectAndRetry, } +/// A multi-statement result with metadata intended for query clients. +/// +/// `execution_error` is emitted only for synthesized MySQL-protocol errors so +/// clients can distinguish them from a successful result column named `Error`. +#[derive(Debug, Clone, Serialize)] +pub struct ExecuteMultiResult { + #[serde(flatten)] + pub result: db::QueryResult, + #[serde(skip_serializing_if = "is_false")] + pub execution_error: bool, +} + +impl ExecuteMultiResult { + fn execution_error(result: db::QueryResult) -> Self { + Self { result, execution_error: true } + } + + fn into_query_result(self) -> db::QueryResult { + self.result + } +} + +impl From for ExecuteMultiResult { + fn from(result: db::QueryResult) -> Self { + Self { result, execution_error: false } + } +} + +fn is_false(value: &bool) -> bool { + !*value +} + /// Unified database operation execution budget. /// query_timeout = None only means SQL execution has no upper limit; /// checkout/connect/recycle/cancel/cleanup always have hard upper limits and cannot be disabled. @@ -1744,6 +1777,21 @@ pub async fn execute_multi_core_with_options( cancel_token: Option, options: QueryExecutionOptions, ) -> Result, String> { + execute_multi_core_with_options_for_client(state, connection_id, database, sql, schema, cancel_token, options) + .await + .map(|results| results.into_iter().map(ExecuteMultiResult::into_query_result).collect()) +} + +/// Execute a SQL batch and retain client-facing metadata for synthesized errors. +pub async fn execute_multi_core_with_options_for_client( + state: &AppState, + connection_id: &str, + database: &str, + sql: &str, + schema: Option<&str>, + cancel_token: Option, + options: QueryExecutionOptions, +) -> Result, String> { // Reject MongoDB queries that fall through to the generic executor. if connection_is_mongodb(state, connection_id).await { return Err("Use MongoDB-specific commands".to_string()); @@ -1768,7 +1816,9 @@ pub async fn execute_multi_core_with_options( }; if is_sqlserver { - return execute_multi_sqlserver(state, &pool_key, sql, cancel_token, options).await; + return execute_multi_sqlserver(state, &pool_key, sql, cancel_token, options) + .await + .map(|results| results.into_iter().map(Into::into).collect()); } let is_turso = { @@ -1781,7 +1831,7 @@ pub async fn execute_multi_core_with_options( let result = execute_sql_statement_with_options(state, connection_id, database, sql, schema, cancel_token, options) .await?; - return Ok(vec![result]); + return Ok(vec![result.into()]); } let db_type = connection_database_type(state, connection_id).await; @@ -1790,14 +1840,14 @@ pub async fn execute_multi_core_with_options( |db_type| crate::sql::split_sql_statements_for_database(sql, db_type), ); if statements.is_empty() { - return Ok(vec![empty_query_result(0)]); + return Ok(vec![empty_query_result(0).into()]); } // When use_transaction is explicitly true and we have multiple statements, // route through the transaction wrapper instead of the sequential auto-commit loop. if options.use_transaction == Some(true) && statements.len() > 1 { let result = execute_statements_in_transaction(state, connection_id, database, &statements, schema).await?; - return Ok(vec![result]); + return Ok(vec![result.into()]); } let mysql_pool = { @@ -1820,7 +1870,7 @@ pub async fn execute_multi_core_with_options( options, ) .await?; - return Ok(vec![result]); + return Ok(vec![result.into()]); } if let Some((pool, mode)) = mysql_pool { @@ -1865,7 +1915,67 @@ pub async fn execute_multi_core_with_options( } } - Ok(results) + Ok(results.into_iter().map(Into::into).collect()) +} + +trait MysqlBatchStatementExecutor { + async fn execute_statement(&mut self, statement: &str) -> Result; +} + +struct MysqlBatchConnection<'a> { + conn: &'a mut mysql_async::Conn, + cancel_token: Option, + query_timeout: Option, + bare: bool, + max_rows: Option, + dialect: db::mysql::MySqlQueryDialect, +} + +impl MysqlBatchStatementExecutor for MysqlBatchConnection<'_> { + async fn execute_statement(&mut self, statement: &str) -> Result { + wait_for_query_opt( + self.cancel_token.clone(), + self.query_timeout, + db::mysql::execute_query_on_conn_with_max_rows( + &mut *self.conn, + statement, + self.bare, + self.max_rows, + self.dialect, + ), + ) + .await + } +} + +async fn execute_mysql_batch_statements( + executor: &mut E, + statements: &[String], + db_type: Option, + cancel_token: Option, +) -> (Vec, Option) +where + E: MysqlBatchStatementExecutor, +{ + let mut results = Vec::with_capacity(statements.len()); + for statement in statements { + if is_canceled(&cancel_token) { + results.push(ExecuteMultiResult::execution_error(error_query_result(canceled_error()))); + return (results, None); + } + + match executor.execute_statement(statement).await { + Ok(result) => results.push(result.into()), + Err(err) => { + let action = pool_error_action(db_type, &err); + results.push(ExecuteMultiResult::execution_error(error_query_result(err))); + // Do not run dependent statements after any MySQL-protocol statement fails. + return (results, Some(action)); + } + } + } + + (results, None) } async fn execute_multi_mysql( @@ -1878,7 +1988,7 @@ async fn execute_multi_mysql( statements: &[String], cancel_token: Option, options: QueryExecutionOptions, -) -> Result, String> { +) -> Result, String> { let query_timeout = resolve_query_timeout(options.timeout_secs); let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await; let bare = mode == crate::connection::MysqlMode::Bare; @@ -1898,36 +2008,25 @@ async fn execute_multi_mysql( { state.remove_pool_by_key(pool_key).await; } - return Ok(vec![error_query_result(err)]); + return Ok(vec![ExecuteMultiResult::execution_error(error_query_result(err))]); } }; - let mut results = Vec::with_capacity(statements.len()); - apply_oceanbase_mysql_session_timeout(state, pool_key, &mut conn, options.timeout_secs).await?; - for stmt in statements { - if is_canceled(&cancel_token) { - results.push(error_query_result(canceled_error())); - break; - } + let mut executor = MysqlBatchConnection { + conn: &mut conn, + cancel_token: cancel_token.clone(), + query_timeout, + bare, + max_rows, + dialect, + }; + let (results, error_action) = + execute_mysql_batch_statements(&mut executor, statements, db_type, cancel_token).await; + drop(executor); - match wait_for_query_opt( - cancel_token.clone(), - query_timeout, - db::mysql::execute_query_on_conn_with_max_rows(&mut conn, stmt, bare, max_rows, dialect), - ) - .await - { - Ok(result) => results.push(result), - Err(err) => { - let action = pool_error_action(db_type, &err); - results.push(error_query_result(err)); - if matches!(action, PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) { - state.remove_pool_by_key(pool_key).await; - break; - } - } - } + if matches!(error_action, Some(PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry)) { + state.remove_pool_by_key(pool_key).await; } Ok(results) @@ -3065,6 +3164,18 @@ mod tests { } } + struct FakeMysqlBatchExecutor { + outcomes: std::collections::VecDeque>, + executed: Vec, + } + + impl MysqlBatchStatementExecutor for FakeMysqlBatchExecutor { + async fn execute_statement(&mut self, statement: &str) -> Result { + self.executed.push(statement.to_string()); + self.outcomes.pop_front().expect("test outcome for statement") + } + } + #[test] fn agent_execute_batch_unsupported_detects_case_insensitive_method_errors() { assert!(is_agent_execute_batch_unsupported("Agent RPC error (-1): unknown method: execute_batch")); @@ -3094,6 +3205,39 @@ mod tests { ); } + #[tokio::test] + async fn mysql_batch_stops_after_the_first_statement_error() { + let statements = vec!["first".to_string(), "fails".to_string(), "must-not-run".to_string()]; + let mut executor = FakeMysqlBatchExecutor { + outcomes: std::collections::VecDeque::from([ + Ok(empty_query_result(0)), + Err("Duplicate entry".to_string()), + Ok(empty_query_result(0)), + ]), + executed: Vec::new(), + }; + + let (results, error_action) = + execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None).await; + + assert_eq!(executor.executed, vec!["first", "fails"]); + assert_eq!(results.len(), 2); + assert!(results[1].execution_error); + assert_eq!(error_action, Some(PoolErrorAction::Keep)); + } + + #[test] + fn execute_multi_result_serializes_error_marker_only_for_synthesized_errors() { + let success = serde_json::to_value(ExecuteMultiResult::from(empty_query_result(0))).unwrap(); + assert!(success.get("execution_error").is_none()); + + let failure = + serde_json::to_value(ExecuteMultiResult::execution_error(error_query_result("failed".to_string()))) + .unwrap(); + assert_eq!(failure.get("execution_error"), Some(&serde_json::Value::Bool(true))); + assert_eq!(failure.get("columns"), Some(&serde_json::json!(["Error"]))); + } + #[test] fn external_driver_method_unsupported_detects_legacy_plugin_errors() { assert!(is_external_driver_method_unsupported( diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 30d3e37d0..71f073d29 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -3,7 +3,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use dbx_core::connection::AppState; use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; -use dbx_core::query::execute_sql_statement; +use dbx_core::query::{execute_multi_core, execute_sql_statement}; use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest}; use dbx_core::sql::{split_sql_statements_for_database, SqlFileRequest}; use dbx_core::sql_file_import::execute_sql_file_path; @@ -381,6 +381,76 @@ async fn live_mysql_recovers_after_server_idle_disconnect() { assert_eq!(result.rows, vec![vec![serde_json::json!("1")]]); } +#[tokio::test] +#[ignore = "requires DBX_LIVE_SQL_FILE_MYSQL_* env vars pointing at a writable MySQL connection without a required default database"] +async fn live_mysql_multi_statement_stops_after_first_error() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let connection_id = format!("live-mysql-multi-stop-{suffix}"); + let config = live_mysql_sql_file_config(&connection_id); + let (state, db_path) = app_state_with_config(config.clone()).await; + let database_name = format!("dbx_issue_2783_{suffix}"); + let table_name = "statement_order"; + let script = format!( + "INSERT INTO `{table_name}` (id, label) VALUES (1, 'first');\n\ + INSERT INTO `{table_name}` (id, label) VALUES (1, 'duplicate');\n\ + INSERT INTO `{table_name}` (id, label) VALUES (2, 'must-not-run');" + ); + + let result = async { + let _ = execute_sql_statement( + &state, + &config.id, + "", + &format!("DROP DATABASE IF EXISTS `{database_name}`"), + None, + None, + ) + .await; + execute_sql_statement(&state, &config.id, "", &format!("CREATE DATABASE `{database_name}`"), None, None) + .await?; + execute_sql_statement( + &state, + &config.id, + &database_name, + &format!("CREATE TABLE `{table_name}` (id INT PRIMARY KEY, label VARCHAR(32) NOT NULL)"), + None, + None, + ) + .await?; + + let results = execute_multi_core(&state, &config.id, &database_name, &script, None, None).await?; + let rows = execute_sql_statement( + &state, + &config.id, + &database_name, + &format!("SELECT id, label FROM `{table_name}` ORDER BY id"), + None, + None, + ) + .await?; + Ok::<_, String>((results, rows)) + } + .await; + + let _ = execute_sql_statement( + &state, + &config.id, + "", + &format!("DROP DATABASE IF EXISTS `{database_name}`"), + None, + None, + ) + .await; + let _ = std::fs::remove_file(db_path); + + let (results, rows) = result.expect("multi-statement execution should return the first failure"); + assert_eq!(results.len(), 2); + assert_eq!(results[0].affected_rows, 1); + assert_eq!(results[1].columns, vec!["Error"]); + assert_eq!(rows.columns, vec!["id", "label"]); + assert_eq!(rows.rows, vec![vec![serde_json::json!("1"), serde_json::json!("first")]]); +} + #[tokio::test] #[ignore = "requires DBX_LIVE_SQL_FILE_MYSQL_* env vars pointing at a writable MySQL connection without a required default database"] async fn live_sql_file_import_creates_database_and_switches_context_without_default_database() { diff --git a/crates/dbx-web/src/routes/query.rs b/crates/dbx-web/src/routes/query.rs index e9ca5c309..e818c9e1a 100644 --- a/crates/dbx-web/src/routes/query.rs +++ b/crates/dbx-web/src/routes/query.rs @@ -347,7 +347,7 @@ pub async fn execute_query( pub async fn execute_multi( State(state): State>, Json(req): Json, -) -> Result>, AppError> { +) -> Result>, AppError> { let execution_id = req.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let registered = state.app.running_queries.register_task( @@ -356,7 +356,7 @@ pub async fn execute_multi( ); let cancel_token = registered.token(); - let result = dbx_core::query::execute_multi_core_with_options( + let result = dbx_core::query::execute_multi_core_with_options_for_client( &state.app, &req.connection_id, &req.database, diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index 4c7ed21a8..6a26d1f39 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -70,7 +70,7 @@ pub async fn execute_multi( client_session_id: Option, timeout_secs: Option, use_transaction: Option, -) -> Result, String> { +) -> Result, String> { let execution_id = execution_id.filter(|id| !id.trim().is_empty()); let registered_query = execution_id.as_ref().map(|id| { state.running_queries.register_task( @@ -90,7 +90,7 @@ pub async fn execute_multi( sql ); - let result = dbx_core::query::execute_multi_core_with_options( + let result = dbx_core::query::execute_multi_core_with_options_for_client( &state, &connection_id, &database, @@ -115,8 +115,8 @@ pub async fn execute_multi( trace_id, started_at.elapsed().as_millis(), results.len(), - results.iter().map(|result| result.rows.len()).collect::>(), - results.iter().map(|result| result.execution_time_ms).collect::>() + results.iter().map(|result| result.result.rows.len()).collect::>(), + results.iter().map(|result| result.result.execution_time_ms).collect::>() ), Err(error) => log::error!( "[query][execute_multi:error] trace_id={} elapsed_ms={} error={}",