From b154dcceef8287545dda53c941ff1ae4babe95d7 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 2 Jun 2026 17:35:04 +0800 Subject: [PATCH] fix(mysql): make admin show handling dialect aware --- crates/dbx-core/src/db/mysql.rs | 100 +++++++++++++++++++++++--------- crates/dbx-core/src/query.rs | 87 +++++++++++++++++++++------ 2 files changed, 144 insertions(+), 43 deletions(-) diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 512a876b7f..24c5525930 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -9,6 +9,7 @@ use std::path::PathBuf; use std::time::Duration; use std::time::Instant; +use crate::models::connection::DatabaseType; use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, ObjectInfo, QueryResult, TableInfo, TriggerInfo, @@ -18,6 +19,21 @@ use super::file_validator::validate_file_path; pub type MySqlPool = mysql_async::Pool; +#[derive(Clone, Copy, Debug, Default)] +pub struct MySqlQueryDialect { + supports_admin_show_results: bool, +} + +impl MySqlQueryDialect { + pub fn for_connection(db_type: DatabaseType, driver_profile: Option<&str>) -> Self { + let profile = driver_profile.map(str::to_ascii_lowercase); + Self { + supports_admin_show_results: matches!(db_type, DatabaseType::Doris | DatabaseType::StarRocks) + || profile.as_deref().is_some_and(|profile| matches!(profile, "doris" | "selectdb" | "starrocks")), + } + } +} + fn quote_value(s: &str) -> String { format!("'{}'", s.replace('\\', "\\\\").replace('\'', "\\'")) } @@ -936,7 +952,7 @@ async fn execute_result_set_with_prepared_protocol_on_conn( } pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result { - execute_query_with_max_rows(pool, sql, bare, None).await + execute_query_with_max_rows(pool, sql, bare, None, MySqlQueryDialect::default()).await } pub async fn execute_query_with_max_rows( @@ -944,9 +960,10 @@ pub async fn execute_query_with_max_rows( sql: &str, bare: bool, max_rows: Option, + dialect: MySqlQueryDialect, ) -> Result { let mut conn = get_conn_with_health_check(pool).await?; - execute_query_on_conn_with_max_rows(&mut conn, sql, bare, max_rows).await + execute_query_on_conn_with_max_rows(&mut conn, sql, bare, max_rows, dialect).await } pub async fn execute_query_on_conn_with_max_rows( @@ -954,12 +971,13 @@ pub async fn execute_query_on_conn_with_max_rows( sql: &str, bare: bool, max_rows: Option, + dialect: MySqlQueryDialect, ) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); - if is_result_set_query(sql) { - if bare || requires_text_protocol_query(sql) { + if is_result_set_query(sql, dialect) { + if bare || requires_text_protocol_query(sql, dialect) { execute_result_set_with_text_protocol_on_conn(conn, sql, row_limit, start).await } else { match execute_result_set_with_prepared_protocol_on_conn(conn, sql, row_limit, start).await { @@ -996,13 +1014,13 @@ pub async fn execute_query_on_conn_with_max_rows( } } -fn is_result_set_query(sql: &str) -> bool { +fn is_result_set_query(sql: &str, dialect: MySqlQueryDialect) -> bool { starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH"]) - || is_admin_show_query(sql) + || dialect.supports_admin_show_results && is_admin_show_query(sql) } -fn requires_text_protocol_query(sql: &str) -> bool { - if is_admin_show_query(sql) { +fn requires_text_protocol_query(sql: &str, dialect: MySqlQueryDialect) -> bool { + if dialect.supports_admin_show_results && is_admin_show_query(sql) { return true; } @@ -1176,33 +1194,63 @@ mod tests { #[test] fn mysql_with_queries_are_treated_as_result_sets() { let sql = "WITH RECURSIVE org_tree AS (SELECT 1 AS id) SELECT id FROM org_tree"; - assert!(is_result_set_query(sql)); + assert!(is_result_set_query(sql, MySqlQueryDialect::default())); } #[test] fn mysql_desc_queries_are_treated_as_result_sets() { - assert!(is_result_set_query("DESC users")); + assert!(is_result_set_query("DESC users", MySqlQueryDialect::default())); } #[test] - fn mysql_admin_show_queries_are_treated_as_result_sets() { + fn starrocks_admin_show_queries_are_treated_as_result_sets() { let sql = "ADMIN SHOW FRONTEND CONFIG LIKE '%default_replication_num%'"; + let dialect = MySqlQueryDialect::for_connection(DatabaseType::StarRocks, None); - assert!(is_result_set_query(sql)); - assert!(requires_text_protocol_query(sql)); + assert!(is_result_set_query(sql, dialect)); + assert!(requires_text_protocol_query(sql, dialect)); } #[test] - fn mysql_admin_show_detection_skips_leading_comments() { + fn doris_admin_show_queries_are_treated_as_result_sets() { + let sql = "ADMIN SHOW FRONTEND CONFIG LIKE '%default_replication_num%'"; + let dialect = MySqlQueryDialect::for_connection(DatabaseType::Doris, None); + + assert!(is_result_set_query(sql, dialect)); + assert!(requires_text_protocol_query(sql, dialect)); + } + + #[test] + fn mysql_starrocks_profile_admin_show_queries_are_treated_as_result_sets() { + let sql = "ADMIN SHOW FRONTEND CONFIG LIKE '%default_replication_num%'"; + let dialect = MySqlQueryDialect::for_connection(DatabaseType::Mysql, Some("starrocks")); + + assert!(is_result_set_query(sql, dialect)); + assert!(requires_text_protocol_query(sql, dialect)); + } + + #[test] + fn mysql_admin_show_queries_are_not_treated_as_result_sets() { + let sql = "ADMIN SHOW FRONTEND CONFIG LIKE '%default_replication_num%'"; + let dialect = MySqlQueryDialect::for_connection(DatabaseType::Mysql, None); + + assert!(!is_result_set_query(sql, dialect)); + assert!(!requires_text_protocol_query(sql, dialect)); + } + + #[test] + fn admin_show_detection_skips_leading_comments() { let sql = "-- inspect FE config\nADMIN /* StarRocks */ SHOW FRONTEND CONFIG"; + let dialect = MySqlQueryDialect::for_connection(DatabaseType::StarRocks, None); - assert!(is_result_set_query(sql)); - assert!(requires_text_protocol_query(sql)); + assert!(is_result_set_query(sql, dialect)); + assert!(requires_text_protocol_query(sql, dialect)); } #[test] - fn mysql_admin_set_queries_are_not_treated_as_result_sets() { - assert!(!is_result_set_query("ADMIN SET FRONTEND CONFIG ('default_replication_num' = '1')")); + fn admin_set_queries_are_not_treated_as_result_sets() { + let dialect = MySqlQueryDialect::for_connection(DatabaseType::StarRocks, None); + assert!(!is_result_set_query("ADMIN SET FRONTEND CONFIG ('default_replication_num' = '1')", dialect)); } #[test] @@ -1327,14 +1375,14 @@ mod tests { #[test] fn mysql_management_show_queries_use_text_protocol() { - assert!(requires_text_protocol_query("SHOW PROCESSLIST")); - assert!(requires_text_protocol_query("show full processlist")); - assert!(requires_text_protocol_query("SHOW SLAVE STATUS")); - assert!(requires_text_protocol_query("show replica status")); - assert!(requires_text_protocol_query("SHOW GRANTS")); - assert!(requires_text_protocol_query("SHOW GRANTS FOR 'repl'@'%'")); - assert!(!requires_text_protocol_query("SHOW TABLES")); - assert!(!requires_text_protocol_query("SELECT * FROM users")); + assert!(requires_text_protocol_query("SHOW PROCESSLIST", MySqlQueryDialect::default())); + assert!(requires_text_protocol_query("show full processlist", MySqlQueryDialect::default())); + assert!(requires_text_protocol_query("SHOW SLAVE STATUS", MySqlQueryDialect::default())); + assert!(requires_text_protocol_query("show replica status", MySqlQueryDialect::default())); + assert!(requires_text_protocol_query("SHOW GRANTS", MySqlQueryDialect::default())); + assert!(requires_text_protocol_query("SHOW GRANTS FOR 'repl'@'%'", MySqlQueryDialect::default())); + assert!(!requires_text_protocol_query("SHOW TABLES", MySqlQueryDialect::default())); + assert!(!requires_text_protocol_query("SELECT * FROM users", MySqlQueryDialect::default())); } #[test] diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index eb864f3d69..7f86c6d824 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -20,6 +20,14 @@ async fn connection_database_type(state: &AppState, connection_id: &str) -> Opti configs.get(connection_id).map(|config| config.db_type) } +async fn connection_mysql_query_dialect(state: &AppState, connection_id: &str) -> db::mysql::MySqlQueryDialect { + let configs = state.configs.read().await; + configs + .get(connection_id) + .map(|config| db::mysql::MySqlQueryDialect::for_connection(config.db_type, config.driver_profile.as_deref())) + .unwrap_or_default() +} + #[derive(Clone, Debug, Default)] pub struct QueryExecutionOptions { pub max_rows: Option, @@ -447,6 +455,7 @@ fn resolve_query_timeout(timeout_secs: Option) -> Option { pub async fn do_execute( state: &AppState, pool_key: &str, + mysql_dialect: db::mysql::MySqlQueryDialect, database: Option<&str>, sql: &str, schema: Option<&str>, @@ -489,7 +498,7 @@ pub async fn do_execute( wait_for_query_opt( cancel_token, query_timeout, - db::mysql::execute_query_with_max_rows(&p, sql, bare, max_rows), + db::mysql::execute_query_with_max_rows(&p, sql, bare, max_rows, mysql_dialect), ) .await } @@ -682,7 +691,10 @@ pub async fn execute_sql_statement_with_options( return Err(canceled_error()); } - let result = do_execute(state, &pool_key, Some(database), sql, schema, cancel_token.clone(), options.clone()).await; + let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; + let result = + do_execute(state, &pool_key, mysql_dialect, Some(database), sql, schema, cancel_token.clone(), options.clone()) + .await; match &result { Err(e) if is_connection_error(e) && !is_canceled(&cancel_token) => { @@ -692,7 +704,7 @@ pub async fn execute_sql_statement_with_options( } else { state.reconnect_pool_for_session(connection_id, db_opt, options.client_session_id.as_deref()).await? }; - do_execute(state, &new_key, Some(database), sql, schema, cancel_token, options).await + do_execute(state, &new_key, mysql_dialect, Some(database), sql, schema, cancel_token, options).await } _ => result, } @@ -800,7 +812,8 @@ pub async fn execute_multi_core_with_options( } if let Some((pool, mode)) = mysql_pool { - return execute_multi_mysql(&pool, mode, &statements, cancel_token, options).await; + let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; + return execute_multi_mysql(&pool, mode, mysql_dialect, &statements, cancel_token, options).await; } let mut results = Vec::with_capacity(statements.len()); @@ -833,6 +846,7 @@ pub async fn execute_multi_core_with_options( async fn execute_multi_mysql( pool: &db::mysql::MySqlPool, mode: crate::connection::MysqlMode, + dialect: db::mysql::MySqlQueryDialect, statements: &[String], cancel_token: Option, options: QueryExecutionOptions, @@ -855,7 +869,7 @@ async fn execute_multi_mysql( match wait_for_query_opt( cancel_token.clone(), query_timeout, - db::mysql::execute_query_on_conn_with_max_rows(&mut conn, stmt, bare, max_rows), + db::mysql::execute_query_on_conn_with_max_rows(&mut conn, stmt, bare, max_rows, dialect), ) .await { @@ -967,9 +981,21 @@ pub async fn execute_statements( let mut total_affected: u64 = 0; let start = std::time::Instant::now(); + let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; for (i, sql) in statements.iter().enumerate() { - match do_execute(state, &pool_key, Some(database), sql, schema, None, QueryExecutionOptions::default()).await { + match do_execute( + state, + &pool_key, + mysql_dialect, + Some(database), + sql, + schema, + None, + QueryExecutionOptions::default(), + ) + .await + { Ok(result) => { total_affected += result.affected_rows; } @@ -1043,9 +1069,13 @@ pub async fn execute_statements_in_transaction( Some(TxPath::Mysql(pool, _bare)) => exec_tx_mysql_inner(pool, statements, start).await, Some(TxPath::Sqlite(pool)) => exec_tx_sqlite_inner(pool, statements, start).await, Some(TxPath::Explicit) => { - exec_tx_explicit_inner(state, &pool_key, Some(database), statements, schema, start).await + let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; + exec_tx_explicit_inner(state, &pool_key, mysql_dialect, Some(database), statements, schema, start).await + } + Some(TxPath::None) => { + let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; + exec_tx_none_inner(state, &pool_key, mysql_dialect, Some(database), statements, schema, start).await } - Some(TxPath::None) => exec_tx_none_inner(state, &pool_key, Some(database), statements, schema, start).await, None => Err("Connection not found for transaction".to_string()), } } @@ -1180,6 +1210,7 @@ async fn exec_tx_sqlite_inner( async fn exec_tx_explicit_inner( state: &AppState, pool_key: &str, + mysql_dialect: db::mysql::MySqlQueryDialect, database: Option<&str>, statements: &[String], schema: Option<&str>, @@ -1193,20 +1224,39 @@ async fn exec_tx_explicit_inner( } drop(conns); - do_execute(state, pool_key, database, "BEGIN TRANSACTION", schema, None, QueryExecutionOptions::default()) - .await - .map_err(|e| format!("Failed to begin transaction: {}", e))?; + do_execute( + state, + pool_key, + mysql_dialect, + database, + "BEGIN TRANSACTION", + schema, + None, + QueryExecutionOptions::default(), + ) + .await + .map_err(|e| format!("Failed to begin transaction: {}", e))?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { - match do_execute(state, pool_key, database, sql, schema, None, QueryExecutionOptions::default()).await { + match do_execute(state, pool_key, mysql_dialect, database, sql, schema, None, QueryExecutionOptions::default()) + .await + { Ok(result) => { total_affected += result.affected_rows; } Err(e) => { - if let Err(rb_err) = - do_execute(state, pool_key, database, "ROLLBACK", schema, None, QueryExecutionOptions::default()) - .await + if let Err(rb_err) = do_execute( + state, + pool_key, + mysql_dialect, + database, + "ROLLBACK", + schema, + None, + QueryExecutionOptions::default(), + ) + .await { log::error!("ROLLBACK failed after statement {} error: {}", i + 1, rb_err); } @@ -1215,7 +1265,7 @@ async fn exec_tx_explicit_inner( } } - do_execute(state, pool_key, database, "COMMIT", schema, None, QueryExecutionOptions::default()) + do_execute(state, pool_key, mysql_dialect, database, "COMMIT", schema, None, QueryExecutionOptions::default()) .await .map_err(|e| format!("COMMIT failed: {}", e))?; @@ -1233,6 +1283,7 @@ async fn exec_tx_explicit_inner( async fn exec_tx_none_inner( state: &AppState, pool_key: &str, + mysql_dialect: db::mysql::MySqlQueryDialect, database: Option<&str>, statements: &[String], schema: Option<&str>, @@ -1241,7 +1292,9 @@ async fn exec_tx_none_inner( let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { log::info!("[query][tx-none:statement:start] index={} sql={}", i + 1, sql); - match do_execute(state, pool_key, database, sql, schema, None, QueryExecutionOptions::default()).await { + match do_execute(state, pool_key, mysql_dialect, database, sql, schema, None, QueryExecutionOptions::default()) + .await + { Ok(result) => { total_affected += result.affected_rows; log::info!("[query][tx-none:statement:done] index={} affected_rows={}", i + 1, result.affected_rows);