github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 06:00:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mysql): report batch DML progress without phantom failures

Browse Source
This commit is contained in:
zipg
2026-08-20 21:00:41 +08:00
committed by GitHub
parent 1a3f20f6c3
commit 8031a9e651
3 changed files with 284 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/db/mysql.rs
+22 -7
View File
@@ -5200,6 +5200,7 @@ pub(crate) async fn execute_non_result_batch_on_conn(
conn: &mut mysql_async::Conn,
sql: &str,
expected_results: usize,
on_result: &mut (dyn FnMut(usize, &QueryResult) + Send),
) -> Result<MySqlNonResultBatchOutcome, String> {
if expected_results == 0 {
return Ok(MySqlNonResultBatchOutcome { results: Vec::new(), error: None });
@@ -5218,9 +5219,19 @@ pub(crate) async fn execute_non_result_batch_on_conn(
let mut results = Vec::with_capacity(expected_results);
let mut warning_counts = Vec::with_capacity(expected_results);
let mut error = None;
let mut previous_elapsed_ms = 0;
for statement_index in 0..expected_results {
if !result.columns_ref().is_empty() {
let Some(columns) = result.columns() else {
error = Some(match result.collect::<mysql_async::Row>().await {
Err(next_error) => next_error.to_string(),
Ok(_) => format!(
"MySQL batch returned fewer results than expected (expected {expected_results}, received {statement_index})."
),
});
break;
};
if !columns.is_empty() {
error = Some(format!(
"MySQL batch statement {} unexpectedly returned rows; retry the statement separately.",
statement_index + 1
@@ -5232,6 +5243,7 @@ pub(crate) async fn execute_non_result_batch_on_conn(
let info = result.info().into_owned();
let messages =
if capture_per_set_messages { mysql_info_message(&info).into_iter().collect() } else { Vec::new() };
let elapsed_ms = start.elapsed().as_millis();
let current_result = QueryResult {
columns: vec![],
column_types: Vec::new(),
@@ -5240,7 +5252,7 @@ pub(crate) async fn execute_non_result_batch_on_conn(
spatial_values: vec![],
rows: vec![],
affected_rows: result.affected_rows(),
execution_time_ms: start.elapsed().as_millis(),
execution_time_ms: elapsed_ms.saturating_sub(previous_elapsed_ms),
truncated: false,
session_id: None,
has_more: false,
@@ -5248,15 +5260,18 @@ pub(crate) async fn execute_non_result_batch_on_conn(
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.
// query_iter/next_set only expose this metadata after the server has
// returned the current statement's OK packet. Report it before collect,
// because collect also waits for the next statement's result.
on_result(statement_index, &current_result);
previous_elapsed_ms = elapsed_ms;
results.push(current_result);
warning_counts.push(warnings);
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 {
crates/dbx-core/src/query.rs
+102 -17
View File
@@ -2774,12 +2774,18 @@ trait MysqlBatchStatementExecutor {
async fn execute_statement(&mut self, statement: &str) -> Result<Vec<db::mysql::MySqlQueryResult>, String>;
async fn execute_non_result_batch(&mut self, statements: &[String]) -> db::mysql::MySqlNonResultBatchOutcome {
async fn execute_non_result_batch(
&mut self,
statements: &[String],
on_result: &mut (dyn FnMut(usize, &db::QueryResult) + Send),
) -> db::mysql::MySqlNonResultBatchOutcome {
let mut results = Vec::with_capacity(statements.len());
for statement in statements {
for (statement_index, statement) in statements.iter().enumerate() {
match self.execute_statement(statement).await {
Ok(statement_results) if statement_results.len() == 1 => {
results.push(statement_results.into_iter().next().expect("single MySQL batch result").result);
let result = statement_results.into_iter().next().expect("single MySQL batch result").result;
on_result(statement_index, &result);
results.push(result);
}
Ok(_) => {
return db::mysql::MySqlNonResultBatchOutcome {
@@ -2830,17 +2836,26 @@ impl MysqlBatchStatementExecutor for MysqlBatchConnection<'_> {
.await
}
async fn execute_non_result_batch(&mut self, statements: &[String]) -> db::mysql::MySqlNonResultBatchOutcome {
async fn execute_non_result_batch(
&mut self,
statements: &[String],
on_result: &mut (dyn FnMut(usize, &db::QueryResult) + Send),
) -> db::mysql::MySqlNonResultBatchOutcome {
let sql = statements.join(";\n");
let mut completed_results = Vec::with_capacity(statements.len());
let mut record_result = |statement_index: usize, result: &db::QueryResult| {
completed_results.push(result.clone());
on_result(statement_index, result);
};
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()),
db::mysql::execute_non_result_batch_on_conn(&mut *self.conn, &sql, statements.len(), &mut record_result),
)
.await
{
Ok(outcome) => outcome,
Err(error) => db::mysql::MySqlNonResultBatchOutcome { results: Vec::new(), error: Some(error) },
Err(error) => db::mysql::MySqlNonResultBatchOutcome { results: completed_results, error: Some(error) },
}
}
}
@@ -2925,19 +2940,21 @@ where
statement_index + 1
};
if batch_end > statement_index + 1 {
let outcome = executor.execute_non_result_batch(&statements[statement_index..batch_end]).await;
let mut report_result = |offset: usize, result: &db::QueryResult| {
report_execute_multi_progress(progress, statement_index + offset, statements.len(), result, true, None);
};
let outcome =
executor.execute_non_result_batch(&statements[statement_index..batch_end], &mut report_result).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);
if completed >= batch_end - statement_index {
return (results, Some(action));
}
let failed_index = statement_index + completed;
let result = error_query_result(error.clone());
let backend_error = crate::backend_error::BackendError::from_legacy_backend(&error);
report_execute_multi_progress(
@@ -5569,9 +5586,17 @@ for line in sys.stdin:
self.statement_outcomes.pop_front().expect("test outcome for single statement")
}
async fn execute_non_result_batch(&mut self, statements: &[String]) -> db::mysql::MySqlNonResultBatchOutcome {
async fn execute_non_result_batch(
&mut self,
statements: &[String],
on_result: &mut (dyn FnMut(usize, &db::QueryResult) + Send),
) -> db::mysql::MySqlNonResultBatchOutcome {
self.batches.push(statements.to_vec());
self.batch_outcomes.pop_front().expect("test outcome for pipelined statements")
let outcome = self.batch_outcomes.pop_front().expect("test outcome for pipelined statements");
for (statement_index, result) in outcome.results.iter().enumerate() {
on_result(statement_index, result);
}
outcome
}
}
@@ -5819,7 +5844,7 @@ for line in sys.stdin:
);
assert_eq!(
progress_events.lock().unwrap().iter().map(|event| event.completed).collect::<Vec<_>>(),
vec![1, 3, 4]
vec![1, 2, 3, 4]
);
assert_eq!(error_action, None);
}
@@ -5840,6 +5865,11 @@ for line in sys.stdin:
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,
@@ -5849,7 +5879,7 @@ for line in sys.stdin:
None,
false,
Some(MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES),
None,
Some(&progress),
)
.await;
@@ -5859,6 +5889,15 @@ for line in sys.stdin:
assert_eq!(results[0].statement_index, Some(0));
assert_eq!(results[1].statement_index, Some(1));
assert!(results[1].execution_error);
assert_eq!(
progress_events
.lock()
.unwrap()
.iter()
.map(|event| (event.statement_index, event.success))
.collect::<Vec<_>>(),
vec![(0, true), (1, false)]
);
assert_eq!(error_action, Some(PoolErrorAction::Keep));
}
@@ -5894,6 +5933,52 @@ for line in sys.stdin:
assert_eq!(error_action, Some(PoolErrorAction::Discard));
}
#[tokio::test]
async fn mysql_pipelined_batch_does_not_add_a_failure_after_every_statement_completed() {
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![empty_query_result(1), empty_query_result(1)],
error: Some(QUERY_CANCELED.to_string()),
}]),
statement_outcomes: std::collections::VecDeque::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!(results.len(), statements.len());
assert!(results.iter().all(|result| !result.execution_error));
assert_eq!(results.iter().map(|result| result.statement_index).collect::<Vec<_>>(), vec![Some(0), Some(1)]);
assert_eq!(
progress_events
.lock()
.unwrap()
.iter()
.map(|event| (event.statement_index, event.completed, event.total, event.success))
.collect::<Vec<_>>(),
vec![(0, 1, 2, true), (1, 2, 2, true)]
);
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<_>>();
crates/dbx-core/tests/live_mysql57.rs
+160 -1
View File
@@ -1,11 +1,14 @@
use std::collections::BTreeSet;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use dbx_core::connection::AppState;
use dbx_core::models::connection::{ConnectionConfig, DatabaseType};
use dbx_core::query::{
execute_multi_core, execute_multi_core_with_options_for_client_typed, execute_sql_statement, QueryExecutionOptions,
execute_multi_core, execute_multi_core_with_options_for_client_and_progress,
execute_multi_core_with_options_for_client_typed, execute_sql_statement, ExecuteMultiProgressCallback,
QueryExecutionOptions,
};
use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest};
use dbx_core::sql::{split_sql_statements_for_database, SqlFileRequest};
@@ -864,6 +867,162 @@ async fn live_mysql_multi_statement_stops_after_first_error() {
assert_eq!(rows.rows, vec![vec![serde_json::json!("1"), serde_json::json!("first")]]);
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_SQL_FILE_MYSQL_* env vars pointing at a writable MySQL connection"]
async fn live_mysql_pipelined_dml_reports_each_statement_before_the_next_finishes() {
let suffix = uuid::Uuid::new_v4().simple().to_string();
let connection_id = format!("live-mysql-progress-{suffix}");
let config = live_mysql_sql_file_config(&connection_id);
let database = std::env::var("DBX_LIVE_SQL_FILE_MYSQL_DATABASE").expect("DBX_LIVE_SQL_FILE_MYSQL_DATABASE");
let (state, db_path) = app_state_with_config(config).await;
let table_name = format!("dbx_batch_progress_{}", &suffix[..8]);
execute_sql_statement(
&state,
&connection_id,
&database,
&format!("CREATE TABLE `{table_name}` (id INT PRIMARY KEY)"),
None,
None,
)
.await
.expect("create live progress table");
execute_sql_statement(
&state,
&connection_id,
&database,
&format!("INSERT INTO `{table_name}` VALUES (1)"),
None,
None,
)
.await
.expect("seed live progress table");
let started = Instant::now();
let progress_events = Arc::new(Mutex::new(Vec::new()));
let progress: ExecuteMultiProgressCallback = {
let progress_events = Arc::clone(&progress_events);
Arc::new(move |event| progress_events.lock().unwrap().push((started.elapsed(), event)))
};
let sql = format!("DELETE FROM `{table_name}`;\nINSERT INTO `{table_name}` SELECT 2 WHERE SLEEP(3) = 0;");
let result = execute_multi_core_with_options_for_client_and_progress(
&state,
&connection_id,
&database,
&sql,
None,
None,
QueryExecutionOptions::default(),
Some(progress),
)
.await;
let _ = execute_sql_statement(
&state,
&connection_id,
&database,
&format!("DROP TABLE IF EXISTS `{table_name}`"),
None,
None,
)
.await;
let _ = std::fs::remove_file(db_path);
let results = result.expect("pipelined DML should succeed");
let events = progress_events.lock().unwrap();
assert_eq!(events.len(), 2);
assert_eq!(events[0].1.statement_index, 0);
assert!(events[0].0 < Duration::from_secs(2), "DELETE progress arrived after the slow INSERT");
assert_eq!(events[1].1.statement_index, 1);
assert!(events[1].0 >= Duration::from_millis(2_500));
assert!(results[0].result.execution_time_ms < 2_000);
assert!(results[1].result.execution_time_ms >= 2_500);
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_SQL_FILE_MYSQL_* env vars pointing at a writable MySQL connection"]
async fn live_mysql_pipelined_dml_keeps_completed_results_before_error() {
let suffix = uuid::Uuid::new_v4().simple().to_string();
let connection_id = format!("live-mysql-progress-error-{suffix}");
let config = live_mysql_sql_file_config(&connection_id);
let database = std::env::var("DBX_LIVE_SQL_FILE_MYSQL_DATABASE").expect("DBX_LIVE_SQL_FILE_MYSQL_DATABASE");
let (state, db_path) = app_state_with_config(config).await;
let table_name = format!("dbx_batch_error_{}", &suffix[..8]);
execute_sql_statement(
&state,
&connection_id,
&database,
&format!("CREATE TABLE `{table_name}` (id INT PRIMARY KEY)"),
None,
None,
)
.await
.expect("create live batch error table");
let progress_events = Arc::new(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 sql = format!(
"INSERT INTO `{table_name}` VALUES (1);\n\
INSERT INTO `{table_name}` VALUES (1);\n\
INSERT INTO `{table_name}` VALUES (2);"
);
let result = execute_multi_core_with_options_for_client_and_progress(
&state,
&connection_id,
&database,
&sql,
None,
None,
QueryExecutionOptions::default(),
Some(progress),
)
.await;
let rows = execute_sql_statement(
&state,
&connection_id,
&database,
&format!("SELECT id FROM `{table_name}` ORDER BY id"),
None,
None,
)
.await;
let _ = execute_sql_statement(
&state,
&connection_id,
&database,
&format!("DROP TABLE IF EXISTS `{table_name}`"),
None,
None,
)
.await;
let _ = std::fs::remove_file(db_path);
let results = result.expect("pipelined DML should return statement errors as results");
assert_eq!(
results.len(),
2,
"results: {:?}",
results
.iter()
.map(|result| (result.statement_index, result.execution_error, result.result.affected_rows))
.collect::<Vec<_>>()
);
assert_eq!(results[0].statement_index, Some(0));
assert!(!results[0].execution_error);
assert_eq!(results[1].statement_index, Some(1));
assert!(results[1].execution_error);
assert_eq!(
progress_events.lock().unwrap().iter().map(|event| (event.statement_index, event.success)).collect::<Vec<_>>(),
vec![(0, true), (1, false)]
);
assert_eq!(rows.expect("read rows after batch error").rows, vec![vec![serde_json::json!("1")]]);
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_SQL_FILE_MYSQL_* env vars pointing at a writable MySQL connection without a required default database"]
async fn live_sql_file_import_creates_database_and_switches_context_without_default_database() {