perf(mysql): batch adjacent DML into multi-statement COM_QUERY for large scripts

This commit is contained in:
zipg authored and GitHub committed 2026-08-11 13:06:21 +08:00
1 parent adb2b4c84b
commit 7c1d094946
9 files changed
+683 -82

No files matched your search

@@ -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
<QueryChart v-else-if="activeOutputView === 'chart' && activeTab.result && !activeElasticsearchJsonResponse" class="flex-1 min-h-0" :result="activeTab.result" />
<div v-else-if="activeOutputView === 'summary'" class="flex-1 min-h-0 overflow-auto bg-background">
<div v-else-if="activeOutputView === 'summary'" class="flex flex-1 min-h-0 flex-col bg-background">
<div v-if="summaryItems.length === 0" class="flex h-full items-center justify-center text-sm text-muted-foreground">
<Loader2 v-if="activeTab.isExecuting" class="mr-2 h-4 w-4 animate-spin" />
<template v-if="activeTab.isExecuting">{{ t("executionSummary.executing") }}</template>
<template v-else>{{ t("executionSummary.empty") }}</template>
</div>
<div v-else class="min-w-[46rem]">
<div v-if="batchExecutionProgress" class="sticky top-0 z-10 border-b bg-background/95 px-3 py-2 backdrop-blur">
<div v-else class="flex h-full min-h-0 min-w-[46rem] flex-col">
<div v-if="batchExecutionProgress" class="z-10 shrink-0 border-b bg-background/95 px-3 py-2 backdrop-blur">
<div class="mb-1.5 flex items-center gap-3 text-xs">
<span class="font-medium">{{ activeTab.isExecuting ? t("executionSummary.executing") : t("executionSummary.finished") }}</span>
<span class="tabular-nums text-muted-foreground">{{ batchExecutionProgress.completed }} / {{ batchExecutionProgress.total }}</span>
@@ -1429,59 +1430,61 @@ defineExpose({ focusSearch, refreshData, refreshQueryEditorCompletionCache, hand
<div class="h-full rounded-full bg-primary transition-[width] duration-200" :style="{ width: `${batchExecutionPercent}%` }" />
</div>
</div>
<div class="overflow-hidden border-b">
<div class="grid grid-cols-[4rem_minmax(14rem,1fr)_7rem_7rem_6rem] border-b bg-muted/30 px-3 py-2 text-xs font-medium text-muted-foreground">
<div>{{ t("executionSummary.statement") }}</div>
<div>{{ t("executionSummary.sql") }}</div>
<div>{{ t("executionSummary.status") }}</div>
<div class="text-right">{{ t("executionSummary.rows") }}</div>
<div class="text-right">{{ t("executionSummary.time") }}</div>
</div>
<div v-for="item in summaryItems" :key="item.statementIndex" class="relative grid w-full grid-cols-[4rem_minmax(14rem,1fr)_7rem_7rem_6rem] items-center border-b px-3 py-2 text-left text-xs last:border-b-0">
<button
type="button"
class="absolute inset-0 z-0 cursor-pointer text-left transition-colors hover:bg-muted/35 focus-visible:bg-muted/40 focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-inset focus-visible:ring-primary"
:title="item.error || item.sql"
:aria-label="item.error || item.sql || t('executionSummary.noSql')"
@click="previewExecutionSummaryItem(item)"
@dblclick="focusExecutionSummaryItem(item)"
@keydown.enter.prevent="focusExecutionSummaryItem(item)"
/>
<div class="pointer-events-none relative z-[1] font-mono text-muted-foreground">#{{ item.statementIndex + 1 }}</div>
<div class="relative z-[1] min-w-0 cursor-pointer" @click="previewExecutionSummaryItem(item)" @dblclick="focusExecutionSummaryItem(item)">
<div class="truncate font-mono text-[11px] text-foreground">{{ item.sql || t("executionSummary.noSql") }}</div>
<div v-if="item.error" class="mt-0.5 flex min-w-0 items-center gap-1 text-[11px] text-destructive">
<span data-native-clipboard class="min-w-0 flex-1 cursor-text select-text truncate" :title="item.error" @mousedown.stop @click.stop @dblclick.stop>{{ item.error }}</span>
<LightTooltip :text="t('grid.copy')" side="bottom" :delay="0" :close-delay="0" nowrap>
<button type="button" class="pointer-events-auto flex h-5 w-5 shrink-0 items-center justify-center rounded text-destructive/70 hover:bg-destructive/10 hover:text-destructive" :aria-label="t('grid.copy')" @mousedown.stop @click.stop="copyExecutionSummaryError(item.error)">
<Copy class="h-3 w-3" />
</button>
</LightTooltip>
</div>
</div>
<div class="pointer-events-none relative z-[1]">
<span
class="inline-flex h-5 items-center gap-1 rounded-full border px-2 text-[10px]"
:class="{
'border-primary/35 bg-primary/10 text-primary': item.status === 'running',
'border-emerald-500/30 bg-emerald-500/10 text-emerald-700 dark:text-emerald-300': item.status === 'success',
'border-destructive/40 bg-destructive/10 text-destructive': item.status === 'error',
'border-border bg-muted/40 text-muted-foreground': item.status === 'pending' || item.status === 'skipped',
'border-amber-500/35 bg-amber-500/10 text-amber-700 dark:text-amber-300': item.status === 'cancelled',
}"
>
<Loader2 v-if="item.status === 'running'" class="h-3 w-3 animate-spin" />
<Check v-else-if="item.status === 'success'" class="h-3 w-3" />
<X v-else-if="item.status === 'error'" class="h-3 w-3" />
<SquareDashed v-else class="h-3 w-3" />
{{ t(`executionSummary.statuses.${item.status}`) }}
</span>
</div>
<div class="pointer-events-none relative z-[1] text-right tabular-nums">{{ item.status === "pending" || item.status === "running" || item.status === "skipped" ? "—" : item.rowCount.toLocaleString() }}</div>
<div class="pointer-events-none relative z-[1] text-right tabular-nums">{{ item.executionTimeMs > 0 || item.status === "success" || item.status === "error" ? `${item.executionTimeMs}ms` : "—" }}</div>
</div>
<div class="grid shrink-0 grid-cols-[4rem_minmax(14rem,1fr)_7rem_7rem_6rem] border-b bg-muted/30 px-3 py-2 text-xs font-medium text-muted-foreground">
<div>{{ t("executionSummary.statement") }}</div>
<div>{{ t("executionSummary.sql") }}</div>
<div>{{ t("executionSummary.status") }}</div>
<div class="text-right">{{ t("executionSummary.rows") }}</div>
<div class="text-right">{{ t("executionSummary.time") }}</div>
</div>
<div class="px-3 py-2 text-[11px] text-muted-foreground">{{ t("executionSummary.navigationHint") }}</div>
<DynamicScroller v-slot="{ item, index, active }" class="min-h-0 flex-1 border-b" :items="summaryItems" :min-item-size="37" :buffer="600" :skip-hover="true" key-field="statementIndex">
<DynamicScrollerItem :item="item" :active="active" :data-index="index" :size-dependencies="[item.error]">
<div class="relative grid w-full grid-cols-[4rem_minmax(14rem,1fr)_7rem_7rem_6rem] items-center border-b px-3 py-2 text-left text-xs last:border-b-0">
<button
type="button"
class="absolute inset-0 z-0 cursor-pointer text-left transition-colors hover:bg-muted/35 focus-visible:bg-muted/40 focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-inset focus-visible:ring-primary"
:title="item.error || item.sql"
:aria-label="item.error || item.sql || t('executionSummary.noSql')"
@click="previewExecutionSummaryItem(item)"
@dblclick="focusExecutionSummaryItem(item)"
@keydown.enter.prevent="focusExecutionSummaryItem(item)"
/>
<div class="pointer-events-none relative z-[1] font-mono text-muted-foreground">#{{ item.statementIndex + 1 }}</div>
<div class="relative z-[1] min-w-0 cursor-pointer" @click="previewExecutionSummaryItem(item)" @dblclick="focusExecutionSummaryItem(item)">
<div class="truncate font-mono text-[11px] text-foreground">{{ item.sql || t("executionSummary.noSql") }}</div>
<div v-if="item.error" class="mt-0.5 flex min-w-0 items-center gap-1 text-[11px] text-destructive">
<span data-native-clipboard class="min-w-0 flex-1 cursor-text select-text truncate" :title="item.error" @mousedown.stop @click.stop @dblclick.stop>{{ item.error }}</span>
<LightTooltip :text="t('grid.copy')" side="bottom" :delay="0" :close-delay="0" nowrap>
<button type="button" class="pointer-events-auto flex h-5 w-5 shrink-0 items-center justify-center rounded text-destructive/70 hover:bg-destructive/10 hover:text-destructive" :aria-label="t('grid.copy')" @mousedown.stop @click.stop="copyExecutionSummaryError(item.error)">
<Copy class="h-3 w-3" />
</button>
</LightTooltip>
</div>
</div>
<div class="pointer-events-none relative z-[1]">
<span
class="inline-flex h-5 items-center gap-1 rounded-full border px-2 text-[10px]"
:class="{
'border-primary/35 bg-primary/10 text-primary': item.status === 'running',
'border-emerald-500/30 bg-emerald-500/10 text-emerald-700 dark:text-emerald-300': item.status === 'success',
'border-destructive/40 bg-destructive/10 text-destructive': item.status === 'error',
'border-border bg-muted/40 text-muted-foreground': item.status === 'pending' || item.status === 'skipped',
'border-amber-500/35 bg-amber-500/10 text-amber-700 dark:text-amber-300': item.status === 'cancelled',
}"
>
<Loader2 v-if="item.status === 'running'" class="h-3 w-3 animate-spin" />
<Check v-else-if="item.status === 'success'" class="h-3 w-3" />
<X v-else-if="item.status === 'error'" class="h-3 w-3" />
<SquareDashed v-else class="h-3 w-3" />
{{ t(`executionSummary.statuses.${item.status}`) }}
</span>
</div>
<div class="pointer-events-none relative z-[1] text-right tabular-nums">{{ item.status === "pending" || item.status === "running" || item.status === "skipped" ? "—" : item.rowCount.toLocaleString() }}</div>
<div class="pointer-events-none relative z-[1] text-right tabular-nums">{{ item.executionTimeMs > 0 || item.status === "success" || item.status === "error" ? `${item.executionTimeMs}ms` : "—" }}</div>
</div>
</DynamicScrollerItem>
</DynamicScroller>
<div class="shrink-0 px-3 py-2 text-[11px] text-muted-foreground">{{ t("executionSummary.navigationHint") }}</div>
</div>
</div>
+7 -1
View File
@@ -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);
@@ -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", () => {
+7
View File
@@ -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");
+5 -1
View File
@@ -368,8 +368,12 @@ export interface ExecutionSummaryItem {
export function executionSummaryItems(tab: Pick<QueryTab, "result" | "results" | "batchSqlExecution">): ExecutionSummaryItem[] {
const results = tab.results?.length ? tab.results : tab.result ? [tab.result] : [];
if (tab.batchSqlExecution?.items.length) {
const resultsByStatementIndex = new Map<number, QueryResult>();
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;
@@ -197,6 +197,47 @@ describe("queryStore multi-statement errors", () => {
});
});
it("marks every statement completed by a pipelined progress event", async () => {
const pendingExecution = deferred<any[]>();
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");
+14 -4
View File
@@ -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<T>(promise: Promise<T>, timeoutSecs: number, message: string): Promise<T> {
if (timeoutSecs === 0) return promise;
const timeoutMs = frontendQueryTimeoutDelayMs(timeoutSecs);
if (timeoutMs === undefined) return promise;
let timer: ReturnType<typeof setTimeout> | undefined;
try {
return await Promise.race([
promise,
new Promise<never>((_, reject) => {
timer = setTimeout(() => reject(new Error(message)), timeoutSecs * 1000);
timer = setTimeout(() => reject(new Error(message)), timeoutMs);
}),
]);
} finally {
+122 -1
View File
@@ -4207,7 +4207,11 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result<Qu
pub async fn max_allowed_packet(pool: &MySqlPool) -> Result<u64, String> {
let mut conn = get_conn_with_health_check(pool).await?;
query_first_column::<u64>(&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<u64, String> {
query_first_column::<u64>(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<QueryResult>,
pub error: Option<String>,
}
/// 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<MySqlNonResultBatchOutcome, String> {
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::<mysql_async::Row>().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();
+400 -20
View File
@@ -2319,6 +2319,23 @@ fn single_statement_multi_result(
trait MysqlBatchStatementExecutor {
async fn execute_statement(&mut self, statement: &str) -> Result<Vec<db::QueryResult>, 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<DatabaseType>, 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<E>(
executor: &mut E,
statements: &[String],
db_type: Option<DatabaseType>,
mysql_dialect: db::mysql::MySqlQueryDialect,
cancel_token: Option<CancellationToken>,
continue_on_error: bool,
pipeline_non_result_max_bytes: Option<usize>,
progress: Option<&ExecuteMultiProgressCallback>,
) -> (Vec<ExecuteMultiResult>, Option<PoolErrorAction>)
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<db::mysql::MySqlNonResultBatchOutcome>,
statement_outcomes: std::collections::VecDeque<Result<Vec<db::QueryResult>, String>>,
batches: Vec<Vec<String>>,
statements: Vec<String>,
}
impl MysqlBatchStatementExecutor for FakePipelinedMysqlBatchExecutor {
async fn execute_statement(&mut self, statement: &str) -> Result<Vec<db::QueryResult>, 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<Vec<db::QueryResult>, 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<_>>(),
vec![Some(0), Some(1), Some(2), Some(3)]
);
assert_eq!(
progress_events.lock().unwrap().iter().map(|event| event.completed).collect::<Vec<_>>(),
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::<Vec<_>>();
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::<Vec<_>>();
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);