From 8031a9e6518db1978ec4bcd873889c18ff895f0b Mon Sep 17 00:00:00 2001 From: zipg Date: Thu, 20 Aug 2026 21:00:41 +0800 Subject: [PATCH] fix(mysql): report batch DML progress without phantom failures --- crates/dbx-core/src/db/mysql.rs | 29 +++-- crates/dbx-core/src/query.rs | 119 ++++++++++++++++--- crates/dbx-core/tests/live_mysql57.rs | 161 +++++++++++++++++++++++++- 3 files changed, 284 insertions(+), 25 deletions(-) diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 6be201158c..961eb059dc 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -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 { 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::().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, ¤t_result); + previous_elapsed_ms = elapsed_ms; + results.push(current_result); + warning_counts.push(warnings); + 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 { diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index dcc7b47971..b0f83797c8 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -2774,12 +2774,18 @@ 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 { + 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![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![(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![Some(0), Some(1)]); + assert_eq!( + progress_events + .lock() + .unwrap() + .iter() + .map(|event| (event.statement_index, event.completed, event.total, event.success)) + .collect::>(), + 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::>(); diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index b8c1905a03..f38d03c768 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -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::>() + ); + 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![(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() {