feat: add support for bare MySQL connections and improve SQL query handling for doris

This commit is contained in:
t8y2 2026-05-02 21:51:50 +08:00
parent eab5556b6f
commit 44c98ba335
4 changed files with 100 additions and 62 deletions

View File

@ -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?),

View File

@ -400,7 +400,8 @@ pub async fn get_table_ddl(
async fn mysql_ddl(pool: &sqlx::mysql::MySqlPool, table: &str) -> Result<String, String> {
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::<String, _>(1)
.or_else(|_| row.try_get::<Vec<u8>, _>(1).map(|b| String::from_utf8_lossy(&b).to_string()))

View File

@ -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::<String, _>(idx)
.or_else(|_| row.try_get::<Vec<u8>, _>(idx).map(|b| String::from_utf8_lossy(&b).to_string()))
@ -168,7 +172,7 @@ pub async fn connect_bare(url: &str) -> Result<MySqlPool, String> {
}
pub async fn list_databases(pool: &MySqlPool) -> Result<Vec<DatabaseInfo>, String> {
let rows: Vec<MySqlRow> = sqlx::query("SELECT SCHEMA_NAME FROM information_schema.SCHEMATA ORDER BY SCHEMA_NAME")
let rows: Vec<MySqlRow> = 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<Vec<DatabaseInfo>, Strin
}
pub async fn list_tables(pool: &MySqlPool, database: &str) -> Result<Vec<TableInfo>, String> {
let rows: Vec<MySqlRow> = 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<MySqlRow> = 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<Vec<ColumnInfo>, String> {
let rows: Vec<MySqlRow> = 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<MySqlRow> = 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<QueryResult, S
let trimmed = sql.trim().to_uppercase();
if trimmed.starts_with("SELECT") || trimmed.starts_with("SHOW") || trimmed.starts_with("DESCRIBE") || trimmed.starts_with("EXPLAIN") {
let desc = pool.describe(sql).await.map_err(|e| e.to_string())?;
let columns: Vec<String> = desc.columns().iter().map(|c| c.name().to_string()).collect();
let column_types: Vec<String> = desc
.columns()
.iter()
.map(|c| c.type_info().name().to_string())
.collect();
let rows: Vec<MySqlRow> = sqlx::query(sql)
let rows: Vec<MySqlRow> = 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<String> = first.columns().iter().map(|c| c.name().to_string()).collect();
let types: Vec<String> = first.columns().iter().map(|c| c.type_info().name().to_string()).collect();
(cols, types)
} else {
(vec![], vec![])
};
let result_rows: Vec<Vec<serde_json::Value>> = rows
.iter()
.map(|row| {
@ -269,7 +275,7 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str) -> Result<QueryResult, S
truncated: false,
})
} else {
let result = sqlx::query(sql)
let result = sqlx::raw_sql(sql)
.execute(pool)
.await
.map_err(|e| e.to_string())?;
@ -285,19 +291,20 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str) -> Result<QueryResult, S
}
pub async fn list_indexes(pool: &MySqlPool, database: &str, table: &str) -> Result<Vec<IndexInfo>, String> {
let rows: Vec<MySqlRow> = 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<MySqlRow> = 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<Vec<ForeignKeyInfo>, String> {
let rows: Vec<MySqlRow> = 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<MySqlRow> = 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<Vec<TriggerInfo>, String> {
let rows: Vec<MySqlRow> = 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<MySqlRow> = sqlx::raw_sql(&sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()

View File

@ -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";