From a9299e947fe04ae40e430a279b1ffddc95f424c1 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 12 Aug 2026 10:34:27 +0000 Subject: [PATCH] fix(mysql): preserve results from single procedure calls Closes #5560 --- crates/dbx-core/src/query.rs | 81 ++++++++- crates/dbx-core/tests/live_mysql57.rs | 247 +++++++++++++++++++++++++- 2 files changed, 324 insertions(+), 4 deletions(-) diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 0ff4f6c71b..82a85ec524 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -24,7 +24,7 @@ use crate::db; use crate::db::agent_driver::{AgentCallError, AgentErrorStage, AgentOperationOutcome}; use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::query_execution_sql::is_write_sql; -use crate::sql::{split_sql_batches, split_sql_statements}; +use crate::sql::{split_sql_batches, split_sql_statements, starts_with_executable_sql_keyword_for_database}; use crate::sql_dialect::{resolve_for_db, CAP_TRANSACTIONAL_DDL}; use crate::sql_risk::{classify_sql_risk_for_database, SqlRisk}; @@ -2280,7 +2280,14 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( } }; - if statements.len() <= 1 && !(mysql_pool.is_some() && options.max_result_bytes.is_some_and(|value| value > 0)) { + if statements.len() == 1 + && !mysql_single_statement_uses_batch_route( + db_type, + mysql_pool.is_some(), + &statements[0], + options.max_result_bytes, + ) + { let single_sql = statements.into_iter().next().unwrap_or_default(); return single_statement_multi_result( execute_sql_statement_with_options_typed( @@ -2383,6 +2390,18 @@ fn single_statement_multi_result( result.map(|result| vec![result.into()]) } +fn mysql_single_statement_uses_batch_route( + db_type: Option, + has_mysql_pool: bool, + sql: &str, + max_result_bytes: Option, +) -> bool { + has_mysql_pool + && (max_result_bytes.is_some_and(|value| value > 0) + || (db_type == Some(DatabaseType::Mysql) + && starts_with_executable_sql_keyword_for_database(sql, &["CALL"], DatabaseType::Mysql))) +} + trait MysqlBatchStatementExecutor { async fn execute_statement(&mut self, statement: &str) -> Result, String>; @@ -2494,6 +2513,14 @@ fn mysql_non_result_batch_end( end.max(start + 1) } +fn mysql_non_result_pipeline_enabled( + statement_count: usize, + continue_on_error: bool, + mode: crate::connection::MysqlMode, +) -> bool { + statement_count > 1 && !continue_on_error && mode == crate::connection::MysqlMode::Normal +} + async fn execute_mysql_batch_statements( executor: &mut E, statements: &[String], @@ -2622,7 +2649,8 @@ async fn execute_multi_mysql( let bare = mode == crate::connection::MysqlMode::Bare; let max_rows = options.max_rows; let max_result_bytes = options.max_result_bytes.filter(|value| *value > 0); - let pipeline_non_result_statements = !options.continue_on_error && mode == crate::connection::MysqlMode::Normal; + let pipeline_non_result_statements = + mysql_non_result_pipeline_enabled(statements.len(), options.continue_on_error, mode); let mut conn = match db::mysql::get_conn_with_health_check_with_cancel( pool, operation_budget.checkout_timeout, @@ -5400,6 +5428,53 @@ for line in sys.stdin: ); } + #[test] + fn mysql_single_call_uses_multi_result_route() { + assert!(mysql_single_statement_uses_batch_route( + Some(DatabaseType::Mysql), + true, + "# generated call\nCALL testA()", + None, + )); + } + + #[test] + fn ordinary_mysql_single_statements_keep_singular_route() { + for sql in ["SELECT 1", "SHOW TABLES", "UPDATE users SET active = 1"] { + assert!(!mysql_single_statement_uses_batch_route(Some(DatabaseType::Mysql), true, sql, None)); + } + } + + #[test] + fn mysql_call_route_requires_native_mysql_type_and_pool() { + assert!(!mysql_single_statement_uses_batch_route(Some(DatabaseType::Doris), true, "CALL testA()", None,)); + assert!(!mysql_single_statement_uses_batch_route(Some(DatabaseType::Mysql), false, "CALL testA()", None,)); + } + + #[test] + fn mysql_result_byte_limit_keeps_existing_batch_route() { + assert!(mysql_single_statement_uses_batch_route( + Some(DatabaseType::Mysql), + true, + "SELECT * FROM users", + Some(1024), + )); + assert!(!mysql_single_statement_uses_batch_route( + Some(DatabaseType::Mysql), + true, + "SELECT * FROM users", + Some(0), + )); + } + + #[test] + fn single_mysql_batch_route_never_probes_non_result_pipeline_limits() { + assert!(!mysql_non_result_pipeline_enabled(1, false, crate::connection::MysqlMode::Normal)); + assert!(mysql_non_result_pipeline_enabled(2, false, crate::connection::MysqlMode::Normal)); + assert!(!mysql_non_result_pipeline_enabled(2, true, crate::connection::MysqlMode::Normal)); + assert!(!mysql_non_result_pipeline_enabled(2, false, crate::connection::MysqlMode::Bare)); + } + #[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()]; diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 7e4faddf68..dfd9dd6221 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -4,7 +4,9 @@ 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_sql_statement}; +use dbx_core::query::{ + execute_multi_core, execute_multi_core_with_options_for_client_typed, execute_sql_statement, QueryExecutionOptions, +}; use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest}; use dbx_core::sql::{split_sql_statements_for_database, SqlFileRequest}; use dbx_core::sql_file_import::execute_sql_file_path; @@ -132,6 +134,249 @@ async fn live_mysql_stored_procedure_preserves_all_result_sets() { ); } +#[tokio::test] +#[ignore = "requires a writable MySQL endpoint for stored procedure creation"] +async fn live_mysql_single_call_public_route_preserves_all_result_sets() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let connection_id = format!("live-mysql-multi-result-{suffix}"); + let database = std::env::var("DBX_LIVE_SQL_FILE_MYSQL_DATABASE").expect("DBX_LIVE_SQL_FILE_MYSQL_DATABASE"); + let config = live_mysql_sql_file_config(&connection_id); + let (state, db_path) = app_state_with_config(config).await; + let procedure = format!("dbx_issue_5560_{}", &suffix[..8]); + let empty_procedure = format!("dbx_issue_5560_empty_{}", &suffix[..8]); + let slow_procedure = format!("dbx_issue_5560_slow_{}", &suffix[..8]); + let table = format!("dbx_issue_5560_{}", &suffix[..8]); + let session_id = format!("dbx-issue-5560-{suffix}"); + let create_sql = format!( + "CREATE PROCEDURE `{procedure}`() BEGIN \ + SET @dbx_issue_5560_session = 'kept'; \ + SELECT 11 AS value UNION ALL SELECT 12 AS value; \ + SELECT 21 AS value UNION ALL SELECT 22 AS value; \ + SELECT 31 AS value UNION ALL SELECT 32 AS value; \ + END" + ); + + execute_sql_statement(&state, &connection_id, &database, &create_sql, None, None) + .await + .expect("create live multi-result procedure"); + execute_sql_statement( + &state, + &connection_id, + &database, + &format!("CREATE PROCEDURE `{empty_procedure}`() BEGIN SET @dbx_issue_5560_empty = 1; END"), + None, + None, + ) + .await + .expect("create live no-result procedure"); + execute_sql_statement( + &state, + &connection_id, + &database, + &format!("CREATE PROCEDURE `{slow_procedure}`() BEGIN SELECT SLEEP(5) AS slept; END"), + None, + None, + ) + .await + .expect("create live slow procedure"); + execute_sql_statement( + &state, + &connection_id, + &database, + &format!("CREATE TABLE `{table}` (id INT PRIMARY KEY, value INT NOT NULL)"), + None, + None, + ) + .await + .expect("create live DML table"); + execute_sql_statement( + &state, + &connection_id, + &database, + &format!("INSERT INTO `{table}` (id, value) VALUES (1, 1)"), + None, + None, + ) + .await + .expect("seed live DML table"); + + let options = + QueryExecutionOptions { max_rows: Some(1), client_session_id: Some(session_id.clone()), ..Default::default() }; + let results = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + &format!("CALL `{procedure}`()"), + None, + None, + options.clone(), + ) + .await + .expect("execute single CALL through the public multi-result route"); + let ordinary_select = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + "SELECT 42 AS value", + None, + None, + options.clone(), + ) + .await + .expect("execute ordinary SELECT through the public route"); + let ordinary_dml = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + &format!("UPDATE `{table}` SET value = value + 1 WHERE id = 1"), + None, + None, + options.clone(), + ) + .await + .expect("execute ordinary DML through the public route"); + let session_result = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + "SELECT @dbx_issue_5560_session AS session_value", + None, + None, + options.clone(), + ) + .await + .expect("reuse the CALL client session"); + let empty_results = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + &format!("CALL `{empty_procedure}`()"), + None, + None, + options.clone(), + ) + .await + .expect("execute no-result CALL through the public route"); + let sql_error = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + &format!("CALL `dbx_issue_5560_missing_{}`()", &suffix[..8]), + None, + None, + options.clone(), + ) + .await + .expect("return a structured execution result for a missing procedure"); + + state.configs.write().await.get_mut(&connection_id).expect("live config").read_only = true; + let read_only_error = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + &format!("CALL `{procedure}`()"), + None, + None, + options.clone(), + ) + .await + .expect_err("read-only mode must reject CALL before dispatch"); + state.configs.write().await.get_mut(&connection_id).expect("live config").read_only = false; + + let timeout_results = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + &format!("CALL `{slow_procedure}`()"), + None, + None, + QueryExecutionOptions { timeout_secs: Some(1), ..options.clone() }, + ) + .await + .expect("return a structured timeout result"); + let cancel_token = CancellationToken::new(); + let cancel_task = { + let cancel_token = cancel_token.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(200)).await; + cancel_token.cancel(); + }) + }; + let canceled_results = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + &format!("CALL `{slow_procedure}`()"), + None, + Some(cancel_token), + options.clone(), + ) + .await + .expect("return a structured cancellation result"); + cancel_task.await.unwrap(); + let recovery_result = execute_multi_core_with_options_for_client_typed( + &state, + &connection_id, + &database, + "SELECT 7 AS recovered", + None, + None, + options, + ) + .await + .expect("recreate the discarded session pool after timeout and cancellation"); + + for cleanup_sql in [ + format!("DROP PROCEDURE `{procedure}`"), + format!("DROP PROCEDURE `{empty_procedure}`"), + format!("DROP PROCEDURE `{slow_procedure}`"), + format!("DROP TABLE `{table}`"), + ] { + execute_sql_statement(&state, &connection_id, &database, &cleanup_sql, None, None) + .await + .expect("clean up live issue fixture"); + } + let _ = std::fs::remove_file(db_path); + + assert_eq!(results.len(), 3); + assert_eq!(results.iter().map(|result| result.statement_index).collect::>(), vec![Some(0); 3]); + assert_eq!( + results.iter().map(|result| result.result.rows[0][0].clone()).collect::>(), + vec![serde_json::json!("11"), serde_json::json!("21"), serde_json::json!("31")] + ); + assert!(results.iter().all(|result| result.result.truncated)); + assert_eq!(ordinary_select.len(), 1); + assert_eq!(ordinary_select[0].statement_index, None); + assert_eq!(ordinary_select[0].result.rows, vec![vec![serde_json::json!("42")]]); + assert_eq!(ordinary_dml.len(), 1); + assert_eq!(ordinary_dml[0].statement_index, None); + assert_eq!(ordinary_dml[0].result.affected_rows, 1); + assert_eq!(session_result[0].result.rows, vec![vec![serde_json::json!("kept")]]); + assert_eq!(empty_results.len(), 1); + assert_eq!(empty_results[0].statement_index, Some(0)); + assert!(empty_results[0].result.columns.is_empty()); + assert!(empty_results[0].result.rows.is_empty()); + assert_eq!(sql_error.len(), 1); + assert_eq!(sql_error[0].statement_index, Some(0)); + assert!(sql_error[0].execution_error); + let sql_error_message = sql_error[0].result.rows[0][0].as_str().expect("missing procedure error message"); + assert!(sql_error_message.contains("does not exist"), "unexpected SQL error: {sql_error_message}"); + assert_eq!(sql_error[0].error.as_ref().map(|error| error.code()), Some("DBX-LEGACY-0001")); + assert!(read_only_error.into_legacy_string().to_ascii_lowercase().contains("read-only")); + assert_eq!(timeout_results.len(), 1); + assert!(timeout_results[0].execution_error); + assert_eq!(timeout_results[0].statement_index, Some(0)); + let timeout_message = timeout_results[0].result.rows[0][0].as_str().expect("timeout error message"); + assert_eq!(timeout_message, "Query timed out after 1 seconds"); + assert_eq!(timeout_results[0].error.as_ref().map(|error| error.code()), Some("DBX-LEGACY-0001")); + assert_eq!(canceled_results.len(), 1); + assert!(canceled_results[0].execution_error); + assert_eq!(canceled_results[0].statement_index, Some(0)); + assert_eq!(canceled_results[0].result.rows[0][0], serde_json::json!("Query canceled")); + assert_eq!(canceled_results[0].error.as_ref().map(|error| error.code()), Some("DBX-JDBC-2003")); + assert_eq!(recovery_result[0].result.rows, vec![vec![serde_json::json!("7")]]); +} + #[tokio::test] #[ignore = "requires a remote MySQL-compatible endpoint with a limited result-set query"] async fn live_mysql_compatible_limited_text_protocol_query_succeeds() {