diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index 73d33dde48..cf743c563f 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -9,6 +9,7 @@ import type { CSSProperties } from "vue"; import { useI18n } from "vue-i18n"; import { Check, CheckSquare2, Columns3Cog, Copy, EyeOff, Loader2, Search, TableProperties, ChevronDown, ChevronUp, Inbox, RefreshCcw, Wrench, Toolbox, Database, Download, Upload, X, Pin, Rows3, SquareDashed, Minus, Plus, ShieldAlert, AlignLeft, AlignRight, PanelsTopLeft } from "@lucide/vue"; import { Splitpanes, Pane } from "splitpanes"; +import { DynamicScroller, DynamicScrollerItem } from "vue-virtual-scroller"; import "splitpanes/dist/splitpanes.css"; import { Button } from "@/components/ui/button"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; @@ -1412,14 +1413,14 @@ defineExpose({ focusSearch, refreshData, refreshQueryEditorCompletionCache, hand -
+
-
-
+
+
{{ activeTab.isExecuting ? t("executionSummary.executing") : t("executionSummary.finished") }} {{ batchExecutionProgress.completed }} / {{ batchExecutionProgress.total }} @@ -1429,59 +1430,61 @@ defineExpose({ focusSearch, refreshData, refreshQueryEditorCompletionCache, hand
-
-
-
{{ t("executionSummary.statement") }}
-
{{ t("executionSummary.sql") }}
-
{{ t("executionSummary.status") }}
-
{{ t("executionSummary.rows") }}
-
{{ t("executionSummary.time") }}
-
-
- - -
-
-
- - - - - - {{ t(`executionSummary.statuses.${item.status}`) }} - -
-
{{ item.status === "pending" || item.status === "running" || item.status === "skipped" ? "—" : item.rowCount.toLocaleString() }}
-
{{ item.executionTimeMs > 0 || item.status === "success" || item.status === "error" ? `${item.executionTimeMs}ms` : "—" }}
-
+
+
{{ t("executionSummary.statement") }}
+
{{ t("executionSummary.sql") }}
+
{{ t("executionSummary.status") }}
+
{{ t("executionSummary.rows") }}
+
{{ t("executionSummary.time") }}
-
{{ t("executionSummary.navigationHint") }}
+ + +
+ + +
+
+
+ + + + + + {{ t(`executionSummary.statuses.${item.status}`) }} + +
+
{{ item.status === "pending" || item.status === "running" || item.status === "skipped" ? "—" : item.rowCount.toLocaleString() }}
+
{{ item.executionTimeMs > 0 || item.status === "success" || item.status === "error" ? `${item.executionTimeMs}ms` : "—" }}
+
+ + +
{{ t("executionSummary.navigationHint") }}
diff --git a/apps/desktop/src/lib/__tests__/sql/queryTimeout.spec.ts b/apps/desktop/src/lib/__tests__/sql/queryTimeout.spec.ts index c20380528e..0eb33dac14 100644 --- a/apps/desktop/src/lib/__tests__/sql/queryTimeout.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/queryTimeout.spec.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { frontendQueryTimeoutSecsForSql, queryTimeoutSecsForConnection } from "@/lib/sql/queryTimeout"; +import { frontendQueryTimeoutDelayMs, frontendQueryTimeoutSecsForSql, queryTimeoutSecsForConnection } from "@/lib/sql/queryTimeout"; describe("queryTimeout", () => { it("lets PostgreSQL row queries use the backend inactivity timeout", () => { @@ -19,6 +19,12 @@ describe("queryTimeout", () => { expect(queryTimeoutSecsForConnection({ query_timeout_secs: undefined })).toBe(30); }); + it("does not schedule frontend timeouts beyond the browser timer limit", () => { + expect(frontendQueryTimeoutDelayMs(60)).toBe(60_000); + expect(frontendQueryTimeoutDelayMs(11_401_200)).toBeUndefined(); + expect(frontendQueryTimeoutDelayMs(0)).toBeUndefined(); + }); + it("uses the global timeout only for inheriting connections", () => { expect(queryTimeoutSecsForConnection({ query_timeout_secs: 30, query_timeout_inherit: true }, 12)).toBe(12); expect(queryTimeoutSecsForConnection({ query_timeout_secs: 30, query_timeout_inherit: false }, 12)).toBe(30); diff --git a/apps/desktop/src/lib/__tests__/tabs/tabPresentation.spec.ts b/apps/desktop/src/lib/__tests__/tabs/tabPresentation.spec.ts index 2e09e381b0..4f58081a81 100644 --- a/apps/desktop/src/lib/__tests__/tabs/tabPresentation.spec.ts +++ b/apps/desktop/src/lib/__tests__/tabs/tabPresentation.spec.ts @@ -234,6 +234,35 @@ describe("execution summary", () => { { status: "error", isError: true }, ]); }); + + it("maps out-of-order results to their explicit statement indexes", () => { + const items = executionSummaryItems({ + results: [ + { columns: ["value"], rows: [["third"]], affected_rows: 0, execution_time_ms: 3, statement_index: 2 }, + { columns: ["value"], rows: [["first"]], affected_rows: 0, execution_time_ms: 1, statement_index: 0 }, + ], + batchSqlExecution: { + executionId: "run-out-of-order", + submittedSql: "SELECT 'first'; SELECT 'second'; SELECT 'third'", + editorFingerprint: "fingerprint", + sourceOffset: 0, + completed: 2, + total: 3, + startedAt: 1, + items: [ + { statementIndex: 0, sql: "SELECT 'first'", from: 0, to: 14, status: "success" }, + { statementIndex: 1, sql: "SELECT 'second'", from: 16, to: 31, status: "skipped" }, + { statementIndex: 2, sql: "SELECT 'third'", from: 33, to: 47, status: "success" }, + ], + }, + }); + + expect(items.map((item) => [item.statementIndex, item.result?.rows[0]?.[0]])).toEqual([ + [0, "first"], + [1, undefined], + [2, "third"], + ]); + }); }); describe("statement execution markers", () => { diff --git a/apps/desktop/src/lib/sql/queryTimeout.ts b/apps/desktop/src/lib/sql/queryTimeout.ts index c4d44cafac..ccb3adfdef 100644 --- a/apps/desktop/src/lib/sql/queryTimeout.ts +++ b/apps/desktop/src/lib/sql/queryTimeout.ts @@ -3,6 +3,7 @@ import { tokenizeSqlSemantic } from "@/lib/sql/semantic/tokens"; import type { ConnectionConfig, DatabaseType } from "@/types/database"; export const DEFAULT_QUERY_TIMEOUT_SECS = 30; +const MAX_BROWSER_TIMER_DELAY_MS = 2_147_483_647; const POSTGRES_ROW_STATEMENT_KEYWORDS = new Set(["select", "show", "explain", "table", "with"]); const POSTGRES_RETURNING_STATEMENT_KEYWORDS = new Set(["insert", "update", "delete", "merge"]); @@ -30,6 +31,12 @@ export function frontendQueryTimeoutSecsForSql(sql: string, databaseType: Databa return baseTimeoutSecs * statementCount; } +export function frontendQueryTimeoutDelayMs(timeoutSecs: number): number | undefined { + const timeoutMs = timeoutSecs * 1000; + if (!Number.isFinite(timeoutMs) || timeoutMs <= 0 || timeoutMs > MAX_BROWSER_TIMER_DELAY_MS) return undefined; + return timeoutMs; +} + function postgresQueryMayReturnRows(sql: string, databaseType: DatabaseType): boolean { return splitSqlStatementRanges(sql, databaseType).some(({ sql: statement }) => { const tokens = tokenizeSqlSemantic(statement, "postgres"); diff --git a/apps/desktop/src/lib/tabs/tabPresentation.ts b/apps/desktop/src/lib/tabs/tabPresentation.ts index 6f6a06e262..32de87f134 100644 --- a/apps/desktop/src/lib/tabs/tabPresentation.ts +++ b/apps/desktop/src/lib/tabs/tabPresentation.ts @@ -368,8 +368,12 @@ export interface ExecutionSummaryItem { export function executionSummaryItems(tab: Pick): ExecutionSummaryItem[] { const results = tab.results?.length ? tab.results : tab.result ? [tab.result] : []; if (tab.batchSqlExecution?.items.length) { + const resultsByStatementIndex = new Map(); + results.forEach((result, resultIndex) => { + resultsByStatementIndex.set(result.statement_index ?? resultIndex, result); + }); return tab.batchSqlExecution.items.map((item, index) => { - const result = results.find((candidate, resultIndex) => (candidate.statement_index ?? resultIndex) === item.statementIndex); + const result = resultsByStatementIndex.get(item.statementIndex); const returnedRows = result?.rows.length ?? 0; const affectedRows = item.affectedRows ?? result?.affected_rows ?? 0; const hasTabularResult = (result?.columns.length ?? 0) > 0; diff --git a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts index c95dcff7c1..6029cfa28c 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts @@ -197,6 +197,47 @@ describe("queryStore multi-statement errors", () => { }); }); + it("marks every statement completed by a pipelined progress event", async () => { + 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 = "SET @n = 0;\nINSERT INTO t VALUES (1);\nINSERT INTO t VALUES (2);\nSELECT COUNT(*) FROM t"; + const tabId = store.createTab("mysql-1", "app", "Query", "query", undefined, sql); + + const execution = store.executeTabSql(tabId, sql); + await vi.waitFor(() => expect(store.tabs.find((item) => item.id === tabId)?.batchSqlExecution?.items[0]?.status).toBe("running")); + const executionId = store.tabs.find((item) => item.id === tabId)!.executionId!; + + reportProgress({ + executionId, + statementIndex: 2, + completed: 3, + total: 4, + success: true, + executionTimeMs: 8, + affectedRows: 1, + }); + + expect(store.tabs.find((item) => item.id === tabId)?.batchSqlExecution).toMatchObject({ + completed: 3, + total: 4, + items: [{ status: "success" }, { status: "success" }, { status: "success" }, { status: "running" }], + }); + + pendingExecution.resolve([ + { columns: [], rows: [], affected_rows: 0, execution_time_ms: 2, statement_index: 0 }, + { columns: [], rows: [], affected_rows: 1, execution_time_ms: 4, statement_index: 1 }, + { columns: [], rows: [], affected_rows: 1, execution_time_ms: 8, statement_index: 2 }, + { columns: ["COUNT(*)"], rows: [[2]], affected_rows: 0, execution_time_ms: 2, statement_index: 3 }, + ]); + await execution; + }); + 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 0ae77ceb31..d270a359ee 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -43,7 +43,7 @@ import { tableOpenPageLimit } from "@/lib/table/tableOpenPageLimit"; import { getCachedTableMetadata, loadTableIndexes, loadTableMetadata, type TableMetadataRequest } from "@/lib/metadata/tableMetadataCache"; import { buildTableSelectSql, quoteTableDataIdentifier } from "@/lib/table/tableSelectSql"; import { connectionObjectTreeNodeSchema, connectionQueryExecutionSchema, connectionUsesDatabaseObjectTreeMode, effectiveDatabaseTypeForConnection, metadataSchemaForConnection } from "@/lib/database/jdbcDialect"; -import { frontendQueryTimeoutSecsForSql, queryTimeoutSecsForConnection } from "@/lib/sql/queryTimeout"; +import { frontendQueryTimeoutDelayMs, frontendQueryTimeoutSecsForSql, queryTimeoutSecsForConnection } from "@/lib/sql/queryTimeout"; import { queryResultNameFromPreamble, queryResultSourceLabel } from "@/lib/sql/queryResultSource"; import { beginDataGridNativeSelectionBlock, finishDataGridNativeSelectionBlock } from "@/lib/dataGrid/dataGridNativeSelection"; import { simpleDataGridOrderByReferencesMissingColumn, sortDataGridRowIndexes, type DataGridSortDirection } from "@/lib/dataGrid/dataGridSort"; @@ -295,8 +295,17 @@ function applyBatchSqlProgress( ) { const batch = batchSqlExecutionFor(tab, progress.executionId); if (!batch) return; + const previousCompleted = batch.completed; const item = batch.items[progress.statementIndex]; if (!item) return; + if (progress.success && progress.completed > previousCompleted + 1) { + for (let index = previousCompleted; index < progress.completed - 1; index += 1) { + const completedItem = batch.items[index]; + if (completedItem && (completedItem.status === "pending" || completedItem.status === "running")) { + completedItem.status = "success"; + } + } + } item.status = progress.success ? "success" : "error"; item.executionTimeMs = progress.executionTimeMs; item.affectedRows = progress.affectedRows; @@ -304,7 +313,7 @@ function applyBatchSqlProgress( item.error = progress.error ? translateBackendError(i18n.global.t, progress.error) : undefined; batch.completed = Math.max(batch.completed, progress.completed); if ((progress.success || continueOnError) && progress.completed < batch.total) { - const next = batch.items[progress.statementIndex + 1]; + const next = batch.items[progress.completed]; if (next?.status === "pending") next.status = "running"; } } @@ -430,14 +439,15 @@ function displayedQueryMetadataSql(tab: QueryTab, fallbackSql: string): string { } async function withFrontendQueryTimeout(promise: Promise, timeoutSecs: number, message: string): Promise { - if (timeoutSecs === 0) return promise; + const timeoutMs = frontendQueryTimeoutDelayMs(timeoutSecs); + if (timeoutMs === undefined) return promise; let timer: ReturnType | undefined; try { return await Promise.race([ promise, new Promise((_, reject) => { - timer = setTimeout(() => reject(new Error(message)), timeoutSecs * 1000); + timer = setTimeout(() => reject(new Error(message)), timeoutMs); }), ]); } finally { diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 040ad7a9b0..01b0b99957 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -4207,7 +4207,11 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result Result { let mut conn = get_conn_with_health_check(pool).await?; - query_first_column::(&mut conn, "SELECT @@max_allowed_packet") + max_allowed_packet_on_conn(&mut conn).await +} + +pub(crate) async fn max_allowed_packet_on_conn(conn: &mut mysql_async::Conn) -> Result { + query_first_column::(conn, "SELECT @@max_allowed_packet") .await .map_err(|error| error.to_string())? .ok_or_else(|| "MySQL did not return @@max_allowed_packet".to_string()) @@ -4493,6 +4497,103 @@ pub async fn execute_query_results_on_conn_with_max_rows( } } +#[derive(Debug)] +pub(crate) struct MySqlNonResultBatchOutcome { + pub results: Vec, + pub error: Option, +} + +/// Executes a chunk of statements that are known not to return rows using one +/// COM_QUERY packet. MySQL still executes every statement in order and exposes +/// one OK packet per statement, so affected-row counts and failure positions +/// remain attributable to the original statements without one network round +/// trip per statement. +pub(crate) async fn execute_non_result_batch_on_conn( + conn: &mut mysql_async::Conn, + sql: &str, + expected_results: usize, +) -> Result { + if expected_results == 0 { + return Ok(MySqlNonResultBatchOutcome { results: Vec::new(), error: None }); + } + + let start = Instant::now(); + let capture_per_set_messages = conn.opts().deprecate_eof(); + let previous_explicit_timestamp_defaults = enable_explicit_timestamp_defaults_for_query(conn, sql).await; + let mut result = match conn.query_iter(sql).await { + Ok(result) => result, + Err(error) => { + restore_explicit_timestamp_defaults_for_query(conn, previous_explicit_timestamp_defaults).await; + return Ok(MySqlNonResultBatchOutcome { results: Vec::new(), error: Some(error.to_string()) }); + } + }; + let mut results = Vec::with_capacity(expected_results); + let mut warning_counts = Vec::with_capacity(expected_results); + let mut error = None; + + for statement_index in 0..expected_results { + if !result.columns_ref().is_empty() { + error = Some(format!( + "MySQL batch statement {} unexpectedly returned rows; retry the statement separately.", + statement_index + 1 + )); + break; + } + + let warnings = result.warnings(); + let info = result.info().into_owned(); + let messages = + if capture_per_set_messages { mysql_info_message(&info).into_iter().collect() } else { Vec::new() }; + let current_result = QueryResult { + columns: vec![], + column_types: Vec::new(), + column_sortables: vec![], + spatial_columns: vec![], + spatial_values: vec![], + rows: vec![], + affected_rows: result.affected_rows(), + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages, + }; + + // mysql_async can expose a failed next statement only when the current + // result set is consumed. Do not mark the current metadata slot as a + // successful statement until that consumption succeeds. + if let Err(next_error) = result.collect::().await { + error = Some(next_error.to_string()); + break; + } + results.push(current_result); + warning_counts.push(warnings); + } + + if let Err(drop_error) = result.drop_result().await { + if error.is_none() { + error = Some(drop_error.to_string()); + } + } + if error.is_none() && !results.is_empty() { + let last_index = results.len() - 1; + for (index, warnings) in warning_counts.into_iter().enumerate() { + if warnings == 0 { + continue; + } + if index == last_index { + results[index].messages.extend(collect_mysql_server_messages(conn, warnings, "").await); + } else if results[index].messages.is_empty() { + results[index].messages.push(mysql_warnings_fallback_message(warnings)); + } + } + } + restore_explicit_timestamp_defaults_for_query(conn, previous_explicit_timestamp_defaults).await; + + Ok(MySqlNonResultBatchOutcome { results, error }) +} + fn prefers_text_protocol_query(sql: &str, dialect: MySqlQueryDialect) -> bool { // User-entered result-set queries are not parameterized in DBX. Text protocol // avoids binary result decoding bugs in MySQL-compatible servers and proxies. @@ -4508,6 +4609,15 @@ pub(crate) fn is_result_set_query(sql: &str, dialect: MySqlQueryDialect) -> bool || dialect.supports_admin_show_results && is_admin_show_query(sql) } +pub(crate) fn is_batchable_non_result_query(sql: &str, dialect: MySqlQueryDialect) -> bool { + !is_result_set_query(sql, dialect) + && starts_with_executable_sql_keyword_for_database( + sql, + &["INSERT", "REPLACE", "UPDATE", "DELETE"], + DatabaseType::Mysql, + ) +} + /// MariaDB 10.5+ returns a result set for INSERT/DELETE/REPLACE ... RETURNING. /// Route it through the existing query path; MySQL servers that do not support /// the syntax still return their normal SQL syntax error. @@ -5538,6 +5648,17 @@ mod tests { assert!(!is_result_set_query("UPDATE users SET name = 'Ada'", dialect)); } + #[test] + fn mysql_multi_statement_pipeline_only_accepts_known_dml() { + let dialect = MySqlQueryDialect::default(); + + assert!(is_batchable_non_result_query("INSERT INTO users(id) VALUES (1)", dialect)); + assert!(is_batchable_non_result_query("UPDATE users SET active = 1", dialect)); + assert!(!is_batchable_non_result_query("SELECT * FROM users", dialect)); + assert!(!is_batchable_non_result_query("ANALYZE TABLE users", dialect)); + assert!(!is_batchable_non_result_query("CREATE TABLE users(id INT)", dialect)); + } + #[test] fn mysql_hash_comments_before_queries_preserve_result_sets_per_issue_3830() { let dialect = MySqlQueryDialect::default(); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 6d255d9056..3dd16f87a3 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -2319,6 +2319,23 @@ fn single_statement_multi_result( trait MysqlBatchStatementExecutor { async fn execute_statement(&mut self, statement: &str) -> Result, String>; + + async fn execute_non_result_batch(&mut self, statements: &[String]) -> db::mysql::MySqlNonResultBatchOutcome { + let mut results = Vec::with_capacity(statements.len()); + for statement in statements { + match self.execute_statement(statement).await { + Ok(mut statement_results) if statement_results.len() == 1 => results.append(&mut statement_results), + Ok(_) => { + return db::mysql::MySqlNonResultBatchOutcome { + results, + error: Some("A non-result MySQL batch statement returned multiple results.".to_string()), + }; + } + Err(error) => return db::mysql::MySqlNonResultBatchOutcome { results, error: Some(error) }, + } + } + db::mysql::MySqlNonResultBatchOutcome { results, error: None } + } } struct MysqlBatchConnection<'a> { @@ -2345,26 +2362,129 @@ impl MysqlBatchStatementExecutor for MysqlBatchConnection<'_> { ) .await } + + async fn execute_non_result_batch(&mut self, statements: &[String]) -> db::mysql::MySqlNonResultBatchOutcome { + let sql = statements.join(";\n"); + match wait_for_result_opt( + self.cancel_token.clone(), + self.query_timeout, + db::mysql::execute_non_result_batch_on_conn(&mut *self.conn, &sql, statements.len()), + ) + .await + { + Ok(outcome) => outcome, + Err(error) => db::mysql::MySqlNonResultBatchOutcome { results: Vec::new(), error: Some(error) }, + } + } +} + +const MYSQL_MULTI_STATEMENT_BATCH_MAX_STATEMENTS: usize = 50; +const MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES: usize = 4 * 1024 * 1024; + +fn mysql_batch_pool_error_action(db_type: Option, error: &str) -> PoolErrorAction { + if error == QUERY_CANCELED { + // Dropping an in-flight COM_QUERY future can leave unread packets on the + // connection. Do not return that connection to the pool. + PoolErrorAction::Discard + } else { + pool_error_action(db_type, error) + } +} + +fn mysql_non_result_batch_end( + statements: &[String], + start: usize, + dialect: db::mysql::MySqlQueryDialect, + max_bytes: usize, +) -> usize { + let Some(first) = statements.get(start) else { + return start; + }; + if !db::mysql::is_batchable_non_result_query(first, dialect) { + return start + 1; + } + + let mut end = start; + let mut byte_len = 0usize; + while let Some(statement) = statements.get(end) { + if end - start >= MYSQL_MULTI_STATEMENT_BATCH_MAX_STATEMENTS + || !db::mysql::is_batchable_non_result_query(statement, dialect) + { + break; + } + let next_len = byte_len.saturating_add(statement.len()).saturating_add(2); + if end > start && next_len > max_bytes { + break; + } + byte_len = next_len; + end += 1; + } + end.max(start + 1) } async fn execute_mysql_batch_statements( executor: &mut E, statements: &[String], db_type: Option, + mysql_dialect: db::mysql::MySqlQueryDialect, cancel_token: Option, continue_on_error: bool, + pipeline_non_result_max_bytes: Option, progress: Option<&ExecuteMultiProgressCallback>, ) -> (Vec, Option) where E: MysqlBatchStatementExecutor, { let mut results = Vec::with_capacity(statements.len()); - for (statement_index, statement) in statements.iter().enumerate() { + let mut statement_index = 0usize; + while statement_index < statements.len() { if is_canceled(&cancel_token) { results.push(ExecuteMultiResult::execution_error(error_query_result(canceled_error()))); return (results, None); } + let batch_end = if let Some(max_bytes) = pipeline_non_result_max_bytes { + mysql_non_result_batch_end(statements, statement_index, mysql_dialect, max_bytes) + } else { + statement_index + 1 + }; + if batch_end > statement_index + 1 { + let outcome = executor.execute_non_result_batch(&statements[statement_index..batch_end]).await; + let completed = outcome.results.len(); + for (offset, result) in outcome.results.into_iter().enumerate() { + results.push(ExecuteMultiResult::success_with_index(result, statement_index + offset)); + } + if completed > 0 { + let last_index = statement_index + completed - 1; + let last_result = &results.last().expect("completed MySQL batch result").result; + report_execute_multi_progress(progress, last_index, statements.len(), last_result, true, None); + } + if let Some(error) = outcome.error { + let failed_index = statement_index + completed; + let action = mysql_batch_pool_error_action(db_type, &error); + let result = error_query_result(error.clone()); + let backend_error = crate::backend_error::BackendError::from_legacy_backend(&error); + report_execute_multi_progress( + progress, + failed_index, + statements.len(), + &result, + false, + Some(backend_error.clone()), + ); + results.push(ExecuteMultiResult::execution_error_with_backend( + result, + Some(failed_index), + backend_error, + )); + return (results, Some(action)); + } + statement_index = batch_end; + continue; + } + + let statement = &statements[statement_index]; + match executor.execute_statement(statement).await { Ok(statement_results) => { if let Some(result) = statement_results.last() { @@ -2377,7 +2497,7 @@ where ); } Err(err) => { - let action = pool_error_action(db_type, &err); + let action = mysql_batch_pool_error_action(db_type, &err); let result = error_query_result(err.clone()); report_execute_multi_progress( progress, @@ -2395,6 +2515,7 @@ where } } } + statement_index += 1; } (results, None) @@ -2419,6 +2540,7 @@ async fn execute_multi_mysql( let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await; let bare = mode == crate::connection::MysqlMode::Bare; let max_rows = options.max_rows; + let pipeline_non_result_statements = !options.continue_on_error && mode == crate::connection::MysqlMode::Normal; let mut conn = match db::mysql::get_conn_with_health_check_with_cancel( pool, operation_budget.checkout_timeout, @@ -2444,6 +2566,15 @@ async fn execute_multi_mysql( db::mysql::apply_catalog_database_context(&mut conn, catalog_dialect, options.catalog.as_deref(), database), ) .await?; + let pipeline_non_result_max_bytes = if pipeline_non_result_statements { + db::mysql::max_allowed_packet_on_conn(&mut conn) + .await + .ok() + .and_then(db::mysql::mysql_sql_statement_hard_limit) + .map(|limit| limit.min(MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES)) + } else { + None + }; let mut executor = MysqlBatchConnection { conn: &mut conn, @@ -2457,8 +2588,10 @@ async fn execute_multi_mysql( &mut executor, statements, db_type, + dialect, cancel_token, options.continue_on_error, + pipeline_non_result_max_bytes, progress, ) .await; @@ -4795,6 +4928,25 @@ for line in sys.stdin: } } + struct FakePipelinedMysqlBatchExecutor { + batch_outcomes: std::collections::VecDeque, + statement_outcomes: std::collections::VecDeque, String>>, + batches: Vec>, + statements: Vec, + } + + impl MysqlBatchStatementExecutor for FakePipelinedMysqlBatchExecutor { + async fn execute_statement(&mut self, statement: &str) -> Result, String> { + self.statements.push(statement.to_string()); + self.statement_outcomes.pop_front().expect("test outcome for single statement") + } + + async fn execute_non_result_batch(&mut self, statements: &[String]) -> db::mysql::MySqlNonResultBatchOutcome { + self.batches.push(statements.to_vec()); + self.batch_outcomes.pop_front().expect("test outcome for pipelined statements") + } + } + fn mysql_batch_result(result: db::QueryResult) -> Result, String> { Ok(vec![result]) } @@ -4906,9 +5058,17 @@ for line in sys.stdin: executed: Vec::new(), }; - let (results, error_action) = - execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, false, None) - .await; + let (results, error_action) = execute_mysql_batch_statements( + &mut executor, + &statements, + Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), + None, + false, + None, + None, + ) + .await; assert_eq!(executor.executed, vec!["first", "fails"]); assert_eq!(results.len(), 2); @@ -4938,8 +5098,10 @@ for line in sys.stdin: &mut executor, &statements, Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), None, false, + None, Some(&progress), ) .await; @@ -4972,6 +5134,184 @@ for line in sys.stdin: assert_eq!(error_action, Some(PoolErrorAction::Keep)); } + #[tokio::test] + async fn mysql_batch_pipelines_adjacent_non_result_statements() { + let statements = vec![ + "SET @batch = 1".to_string(), + "INSERT INTO users(id) VALUES (1)".to_string(), + "INSERT INTO users(id) VALUES (2)".to_string(), + "SELECT COUNT(*) FROM users".to_string(), + ]; + let mut executor = FakePipelinedMysqlBatchExecutor { + batch_outcomes: std::collections::VecDeque::from([db::mysql::MySqlNonResultBatchOutcome { + results: vec![empty_query_result(2), empty_query_result(3)], + error: None, + }]), + statement_outcomes: std::collections::VecDeque::from([ + mysql_batch_result(empty_query_result(1)), + mysql_batch_result(db::QueryResult { + columns: vec!["COUNT(*)".to_string()], + column_types: vec!["BIGINT".to_string()], + column_sortables: vec![], + spatial_columns: vec![], + spatial_values: vec![], + rows: vec![vec![serde_json::json!(2)]], + affected_rows: 0, + execution_time_ms: 4, + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages: Vec::new(), + }), + ]), + batches: Vec::new(), + statements: Vec::new(), + }; + let progress_events = Arc::new(std::sync::Mutex::new(Vec::new())); + let progress: ExecuteMultiProgressCallback = { + let progress_events = Arc::clone(&progress_events); + Arc::new(move |event| progress_events.lock().unwrap().push(event)) + }; + + let (results, error_action) = execute_mysql_batch_statements( + &mut executor, + &statements, + Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), + None, + false, + Some(MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES), + Some(&progress), + ) + .await; + + assert_eq!(executor.batches, vec![statements[1..3].to_vec()]); + assert_eq!(executor.statements, vec![statements[0].clone(), statements[3].clone()]); + assert_eq!(results.len(), 4); + assert_eq!( + results.iter().map(|result| result.statement_index).collect::>(), + vec![Some(0), Some(1), Some(2), Some(3)] + ); + assert_eq!( + progress_events.lock().unwrap().iter().map(|event| event.completed).collect::>(), + vec![1, 3, 4] + ); + assert_eq!(error_action, None); + } + + #[tokio::test] + async fn mysql_pipelined_batch_reports_the_first_failed_statement() { + let statements = vec![ + "INSERT INTO users(id) VALUES (1)".to_string(), + "INSERT INTO users(id) VALUES (1)".to_string(), + "INSERT INTO users(id) VALUES (2)".to_string(), + ]; + let mut executor = FakePipelinedMysqlBatchExecutor { + batch_outcomes: std::collections::VecDeque::from([db::mysql::MySqlNonResultBatchOutcome { + results: vec![empty_query_result(1)], + error: Some("Duplicate entry".to_string()), + }]), + statement_outcomes: std::collections::VecDeque::new(), + batches: Vec::new(), + statements: Vec::new(), + }; + + let (results, error_action) = execute_mysql_batch_statements( + &mut executor, + &statements, + Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), + None, + false, + Some(MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES), + None, + ) + .await; + + assert_eq!(executor.batches, vec![statements]); + assert!(executor.statements.is_empty()); + assert_eq!(results.len(), 2); + assert_eq!(results[0].statement_index, Some(0)); + assert_eq!(results[1].statement_index, Some(1)); + assert!(results[1].execution_error); + assert_eq!(error_action, Some(PoolErrorAction::Keep)); + } + + #[tokio::test] + async fn mysql_pipelined_batch_discards_a_cancelled_connection() { + let statements = + vec!["INSERT INTO users(id) VALUES (1)".to_string(), "INSERT INTO users(id) VALUES (2)".to_string()]; + let mut executor = FakePipelinedMysqlBatchExecutor { + batch_outcomes: std::collections::VecDeque::from([db::mysql::MySqlNonResultBatchOutcome { + results: Vec::new(), + error: Some(QUERY_CANCELED.to_string()), + }]), + statement_outcomes: std::collections::VecDeque::new(), + batches: Vec::new(), + statements: Vec::new(), + }; + + let (results, error_action) = execute_mysql_batch_statements( + &mut executor, + &statements, + Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), + None, + false, + Some(MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES), + None, + ) + .await; + + assert_eq!(results.len(), 1); + assert_eq!(results[0].statement_index, Some(0)); + assert!(results[0].execution_error); + assert_eq!(error_action, Some(PoolErrorAction::Discard)); + } + + #[test] + fn mysql_non_result_batches_respect_the_statement_limit() { + let statements = (0..51).map(|index| format!("INSERT INTO users(id) VALUES ({index})")).collect::>(); + + assert_eq!( + mysql_non_result_batch_end( + &statements, + 0, + db::mysql::MySqlQueryDialect::default(), + MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES, + ), + MYSQL_MULTI_STATEMENT_BATCH_MAX_STATEMENTS + ); + assert_eq!( + mysql_non_result_batch_end( + &statements, + MYSQL_MULTI_STATEMENT_BATCH_MAX_STATEMENTS, + db::mysql::MySqlQueryDialect::default(), + MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES, + ), + 51 + ); + } + + #[test] + fn mysql_non_result_batches_respect_the_byte_limit() { + let payload = "x".repeat(1_500_000); + let statements = (0..3) + .map(|index| format!("INSERT INTO users(id, payload) VALUES ({index}, '{payload}')")) + .collect::>(); + + assert_eq!( + mysql_non_result_batch_end( + &statements, + 0, + db::mysql::MySqlQueryDialect::default(), + MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES, + ), + 2 + ); + } + #[tokio::test] async fn mysql_batch_preserves_multiple_result_sets_from_one_statement() { let statements = vec!["CALL testA()".to_string(), "UPDATE users SET active = 1".to_string()]; @@ -4998,9 +5338,17 @@ for line in sys.stdin: executed: Vec::new(), }; - let (results, error_action) = - execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, false, None) - .await; + let (results, error_action) = execute_mysql_batch_statements( + &mut executor, + &statements, + Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), + None, + false, + None, + None, + ) + .await; assert_eq!(executor.executed, statements); assert_eq!(results.len(), 4); @@ -5026,9 +5374,17 @@ for line in sys.stdin: executed: Vec::new(), }; - let (results, error_action) = - execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, false, None) - .await; + let (results, error_action) = execute_mysql_batch_statements( + &mut executor, + &statements, + Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), + None, + false, + None, + None, + ) + .await; assert_eq!(executor.executed, vec!["fails"]); assert_eq!(results.len(), 1); @@ -5048,9 +5404,17 @@ for line in sys.stdin: executed: Vec::new(), }; - let (results, error_action) = - execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, true, None) - .await; + let (results, error_action) = execute_mysql_batch_statements( + &mut executor, + &statements, + Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), + None, + true, + None, + None, + ) + .await; assert_eq!(executor.executed, statements); assert_eq!(results.len(), 3); @@ -5073,9 +5437,17 @@ for line in sys.stdin: executed: Vec::new(), }; - let (results, error_action) = - execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, true, None) - .await; + let (results, error_action) = execute_mysql_batch_statements( + &mut executor, + &statements, + Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), + None, + true, + None, + None, + ) + .await; assert_eq!(executor.executed, statements); assert_eq!(results.len(), 2); @@ -5095,9 +5467,17 @@ for line in sys.stdin: executed: Vec::new(), }; - let (results, error_action) = - execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, true, None) - .await; + let (results, error_action) = execute_mysql_batch_statements( + &mut executor, + &statements, + Some(DatabaseType::Mysql), + db::mysql::MySqlQueryDialect::default(), + None, + true, + None, + None, + ) + .await; assert_eq!(executor.executed, vec!["first", "disconnects"]); assert_eq!(results.len(), 2);