From 44c98ba335b22da77fe517dc56441774f4df2f47 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Sat, 2 May 2026 21:51:50 +0800 Subject: [PATCH] feat: add support for bare MySQL connections and improve SQL query handling for doris --- src-tauri/src/commands/connection.rs | 35 ++++++--- src-tauri/src/commands/schema.rs | 3 +- src-tauri/src/db/mysql.rs | 111 +++++++++++++++------------ src-tauri/src/models/connection.rs | 13 ++++ 4 files changed, 100 insertions(+), 62 deletions(-) diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index b6a041596..84d918170 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -89,6 +89,7 @@ impl AppState { let (host, port) = self.connection_host_port(connection_id, &db_config).await?; let url = connection_url_for_endpoint(&db_config, &host, port); let pool = match db_config.db_type { + DatabaseType::Mysql if db_config.needs_bare_mysql() => PoolKind::Mysql(db::mysql::connect_bare(&url).await?), DatabaseType::Mysql => PoolKind::Mysql(db::mysql::connect(&url).await?), DatabaseType::Doris | DatabaseType::StarRocks => PoolKind::Mysql(db::mysql::connect_bare(&url).await?), DatabaseType::Postgres | DatabaseType::Redshift => PoolKind::Postgres(db::postgres::connect(&url).await?), @@ -266,19 +267,32 @@ pub async fn test_connection( target ); let result = match config.db_type { - DatabaseType::Mysql => match db::mysql::connect(&url).await { - Ok(pool) => { - pool.close().await; - Ok("Connection successful".to_string()) + DatabaseType::Mysql if config.needs_bare_mysql() => { + match db::mysql::connect_bare(&url).await { + Ok(pool) => { + pool.close().await; + Ok("Connection successful".to_string()) + } + Err(e) => Err(e), } - Err(e) => Err(e), }, - DatabaseType::Doris | DatabaseType::StarRocks => match db::mysql::connect_bare(&url).await { - Ok(pool) => { - pool.close().await; - Ok("Connection successful".to_string()) + DatabaseType::Mysql => { + match db::mysql::connect(&url).await { + Ok(pool) => { + pool.close().await; + Ok("Connection successful".to_string()) + } + Err(e) => Err(e), + } + }, + DatabaseType::Doris | DatabaseType::StarRocks => { + match db::mysql::connect_bare(&url).await { + Ok(pool) => { + pool.close().await; + Ok("Connection successful".to_string()) + } + Err(e) => Err(e), } - Err(e) => Err(e), }, DatabaseType::Postgres | DatabaseType::Redshift => match db::postgres::connect(&url).await { Ok(pool) => { @@ -368,6 +382,7 @@ pub async fn connect_db( let url = connection_url_for_endpoint(&config, &host, port); let pool = match config.db_type { + DatabaseType::Mysql if config.needs_bare_mysql() => PoolKind::Mysql(db::mysql::connect_bare(&url).await?), DatabaseType::Mysql => PoolKind::Mysql(db::mysql::connect(&url).await?), DatabaseType::Doris | DatabaseType::StarRocks => PoolKind::Mysql(db::mysql::connect_bare(&url).await?), DatabaseType::Postgres | DatabaseType::Redshift => PoolKind::Postgres(db::postgres::connect(&url).await?), diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index 2aab53741..ee578d983 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -400,7 +400,8 @@ pub async fn get_table_ddl( async fn mysql_ddl(pool: &sqlx::mysql::MySqlPool, table: &str) -> Result { use sqlx::Row; - let row: sqlx::mysql::MySqlRow = sqlx::query(&format!("SHOW CREATE TABLE `{}`", table.replace('`', "``"))) + let sql = format!("SHOW CREATE TABLE `{}`", table.replace('`', "``")); + let row: sqlx::mysql::MySqlRow = sqlx::raw_sql(&sql) .fetch_one(pool).await.map_err(|e| e.to_string())?; row.try_get::(1) .or_else(|_| row.try_get::, _>(1).map(|b| String::from_utf8_lossy(&b).to_string())) diff --git a/src-tauri/src/db/mysql.rs b/src-tauri/src/db/mysql.rs index 80556cbd4..39c98d35f 100644 --- a/src-tauri/src/db/mysql.rs +++ b/src-tauri/src/db/mysql.rs @@ -1,11 +1,15 @@ use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc}; use rust_decimal::Decimal; use sqlx::mysql::{MySqlPool, MySqlPoolOptions, MySqlRow}; -use sqlx::{Column, Executor, Row, TypeInfo, ValueRef}; +use sqlx::{Column, Row, TypeInfo, ValueRef}; use std::time::{Duration, Instant}; use super::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; +fn quote_value(s: &str) -> String { + format!("'{}'", s.replace('\\', "\\\\").replace('\'', "\\'")) +} + fn get_str(row: &MySqlRow, idx: usize) -> String { row.try_get::(idx) .or_else(|_| row.try_get::, _>(idx).map(|b| String::from_utf8_lossy(&b).to_string())) @@ -168,7 +172,7 @@ pub async fn connect_bare(url: &str) -> Result { } pub async fn list_databases(pool: &MySqlPool) -> Result, String> { - let rows: Vec = sqlx::query("SELECT SCHEMA_NAME FROM information_schema.SCHEMATA ORDER BY SCHEMA_NAME") + let rows: Vec = sqlx::raw_sql("SELECT SCHEMA_NAME FROM information_schema.SCHEMATA ORDER BY SCHEMA_NAME") .fetch_all(pool) .await .map_err(|e| e.to_string())?; @@ -177,13 +181,14 @@ pub async fn list_databases(pool: &MySqlPool) -> Result, Strin } pub async fn list_tables(pool: &MySqlPool, database: &str) -> Result, String> { - let rows: Vec = sqlx::query( - "SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA = ? ORDER BY TABLE_NAME", - ) - .bind(database) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; + let sql = format!( + "SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA = {} ORDER BY TABLE_NAME", + quote_value(database), + ); + let rows: Vec = sqlx::raw_sql(&sql) + .fetch_all(pool) + .await + .map_err(|e| e.to_string())?; Ok(rows .iter() @@ -199,7 +204,7 @@ pub async fn get_columns( database: &str, table: &str, ) -> Result, String> { - let rows: Vec = sqlx::query( + let sql = format!( "SELECT c.COLUMN_NAME, c.DATA_TYPE, c.IS_NULLABLE, c.COLUMN_DEFAULT, c.EXTRA, c.COLUMN_COMMENT, \ CASE WHEN kcu.COLUMN_NAME IS NOT NULL THEN 1 ELSE 0 END AS IS_PK, \ c.NUMERIC_PRECISION, c.NUMERIC_SCALE \ @@ -209,14 +214,15 @@ pub async fn get_columns( AND c.TABLE_NAME = kcu.TABLE_NAME \ AND c.COLUMN_NAME = kcu.COLUMN_NAME \ AND kcu.CONSTRAINT_NAME = 'PRIMARY' \ - WHERE c.TABLE_SCHEMA = ? AND c.TABLE_NAME = ? \ + WHERE c.TABLE_SCHEMA = {} AND c.TABLE_NAME = {} \ ORDER BY c.ORDINAL_POSITION", - ) - .bind(database) - .bind(table) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; + quote_value(database), + quote_value(table), + ); + let rows: Vec = sqlx::raw_sql(&sql) + .fetch_all(pool) + .await + .map_err(|e| e.to_string())?; Ok(rows .iter() @@ -239,19 +245,19 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str) -> Result = desc.columns().iter().map(|c| c.name().to_string()).collect(); - let column_types: Vec = desc - .columns() - .iter() - .map(|c| c.type_info().name().to_string()) - .collect(); - - let rows: Vec = sqlx::query(sql) + let rows: Vec = sqlx::raw_sql(sql) .fetch_all(pool) .await .map_err(|e| e.to_string())?; + let (columns, column_types) = if let Some(first) = rows.first() { + let cols: Vec = first.columns().iter().map(|c| c.name().to_string()).collect(); + let types: Vec = first.columns().iter().map(|c| c.type_info().name().to_string()).collect(); + (cols, types) + } else { + (vec![], vec![]) + }; + let result_rows: Vec> = rows .iter() .map(|row| { @@ -269,7 +275,7 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str) -> Result Result Result, String> { - let rows: Vec = sqlx::query( + let sql = format!( "SELECT INDEX_NAME, GROUP_CONCAT(COLUMN_NAME ORDER BY SEQ_IN_INDEX) AS columns, \ NOT NON_UNIQUE AS is_unique, INDEX_NAME = 'PRIMARY' AS is_primary \ FROM information_schema.STATISTICS \ - WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? \ + WHERE TABLE_SCHEMA = {} AND TABLE_NAME = {} \ GROUP BY INDEX_NAME, NON_UNIQUE \ ORDER BY INDEX_NAME", - ) - .bind(database) - .bind(table) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; + quote_value(database), + quote_value(table), + ); + let rows: Vec = sqlx::raw_sql(&sql) + .fetch_all(pool) + .await + .map_err(|e| e.to_string())?; Ok(rows .iter() @@ -314,19 +321,20 @@ pub async fn list_indexes(pool: &MySqlPool, database: &str, table: &str) -> Resu } pub async fn list_foreign_keys(pool: &MySqlPool, database: &str, table: &str) -> Result, String> { - let rows: Vec = sqlx::query( + let sql = format!( "SELECT kcu.CONSTRAINT_NAME, kcu.COLUMN_NAME, \ kcu.REFERENCED_TABLE_NAME, kcu.REFERENCED_COLUMN_NAME \ FROM information_schema.KEY_COLUMN_USAGE kcu \ - WHERE kcu.TABLE_SCHEMA = ? AND kcu.TABLE_NAME = ? \ + WHERE kcu.TABLE_SCHEMA = {} AND kcu.TABLE_NAME = {} \ AND kcu.REFERENCED_TABLE_NAME IS NOT NULL \ ORDER BY kcu.CONSTRAINT_NAME", - ) - .bind(database) - .bind(table) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; + quote_value(database), + quote_value(table), + ); + let rows: Vec = sqlx::raw_sql(&sql) + .fetch_all(pool) + .await + .map_err(|e| e.to_string())?; Ok(rows .iter() @@ -340,17 +348,18 @@ pub async fn list_foreign_keys(pool: &MySqlPool, database: &str, table: &str) -> } pub async fn list_triggers(pool: &MySqlPool, database: &str, table: &str) -> Result, String> { - let rows: Vec = sqlx::query( + let sql = format!( "SELECT TRIGGER_NAME, EVENT_MANIPULATION, ACTION_TIMING \ FROM information_schema.TRIGGERS \ - WHERE TRIGGER_SCHEMA = ? AND EVENT_OBJECT_TABLE = ? \ + WHERE TRIGGER_SCHEMA = {} AND EVENT_OBJECT_TABLE = {} \ ORDER BY TRIGGER_NAME", - ) - .bind(database) - .bind(table) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; + quote_value(database), + quote_value(table), + ); + let rows: Vec = sqlx::raw_sql(&sql) + .fetch_all(pool) + .await + .map_err(|e| e.to_string())?; Ok(rows .iter() diff --git a/src-tauri/src/models/connection.rs b/src-tauri/src/models/connection.rs index dca3169a9..2d96ff3eb 100644 --- a/src-tauri/src/models/connection.rs +++ b/src-tauri/src/models/connection.rs @@ -69,6 +69,12 @@ pub enum DatabaseType { } impl ConnectionConfig { + pub fn needs_bare_mysql(&self) -> bool { + matches!(self.db_type, DatabaseType::Doris | DatabaseType::StarRocks) + || self.driver_profile.as_deref().map(|p| p.to_lowercase()) + .is_some_and(|p| matches!(p.as_str(), "doris" | "starrocks" | "selectdb" | "tdengine")) + } + pub fn connection_url(&self) -> String { self.connection_url_with_host(&self.host, self.port) } @@ -187,6 +193,13 @@ impl ConnectionConfig { fn normalized_url_params(&self) -> String { let value = self.url_params.as_deref().unwrap_or("").trim(); + if self.needs_bare_mysql() { + let v = value.trim_start_matches('?'); + let filtered: Vec<&str> = v.split('&') + .filter(|p| !p.is_empty() && !p.starts_with("charset=") && !p.starts_with("ssl-mode=preferred")) + .collect(); + return if filtered.is_empty() { "ssl-mode=disabled".to_string() } else { format!("ssl-mode=disabled&{}", filtered.join("&")) }; + } match self.db_type { DatabaseType::Mysql => { let base = "ssl-mode=preferred&charset=utf8mb4";