github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-22 11:02:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mysql): make admin show handling dialect aware

Browse Source
This commit is contained in:
t8y2
2026-06-02 17:35:14 +08:00
parent ac3052c361
commit b154dcceef
2 changed files with 144 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/db/mysql.rs
+74 -26
View File
@@ -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<QueryResult, String> {
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<usize>,
dialect: MySqlQueryDialect,
) -> Result<QueryResult, String> {
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<usize>,
dialect: MySqlQueryDialect,
) -> Result<QueryResult, String> {
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]
crates/dbx-core/src/query.rs
+70 -17
View File
@@ -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<usize>,
@@ -447,6 +455,7 @@ fn resolve_query_timeout(timeout_secs: Option<u64>) -> Option<Duration> {
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<CancellationToken>,
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);