+
+
{{ 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") }}
-
-
-
-
#{{ item.statementIndex + 1 }}
-
-
{{ item.sql || t("executionSummary.noSql") }}
-
- {{ item.error }}
-
-
-
-
-
-
-
-
-
-
-
- {{ 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") }}
+
+
+
+
+
#{{ item.statementIndex + 1 }}
+
+
{{ item.sql || t("executionSummary.noSql") }}
+
+ {{ item.error }}
+
+
+
+
+
+
+
+
+
+
+
+ {{ 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);