From f598b5a91c75957129efc22e0bcba43d086e259a Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Sat, 2 May 2026 23:25:39 +0800 Subject: [PATCH] feat: enhance MySQL connection handling with bare option and optimize component imports --- src-tauri/src/commands/connection.rs | 16 +++--- src-tauri/src/commands/query.rs | 5 +- src-tauri/src/commands/schema.rs | 14 ++--- src-tauri/src/commands/transfer.rs | 7 ++- src-tauri/src/db/mysql.rs | 82 +++++++++++++++++++--------- src/App.vue | 16 +++--- vite.config.ts | 17 ++++++ 7 files changed, 102 insertions(+), 55 deletions(-) diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 84d918170..823e326c7 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -12,7 +12,7 @@ use crate::db::ssh_tunnel::TunnelManager; use crate::models::connection::{ConnectionConfig, DatabaseType}; pub enum PoolKind { - Mysql(sqlx::mysql::MySqlPool), + Mysql(sqlx::mysql::MySqlPool, bool), Postgres(sqlx::postgres::PgPool), Sqlite(sqlx::sqlite::SqlitePool), Redis(tokio::sync::Mutex), @@ -89,9 +89,9 @@ 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::Mysql if db_config.needs_bare_mysql() => PoolKind::Mysql(db::mysql::connect_bare(&url).await?, true), + DatabaseType::Mysql => PoolKind::Mysql(db::mysql::connect(&url).await?, false), + DatabaseType::Doris | DatabaseType::StarRocks => PoolKind::Mysql(db::mysql::connect_bare(&url).await?, true), DatabaseType::Postgres | DatabaseType::Redshift => PoolKind::Postgres(db::postgres::connect(&url).await?), DatabaseType::Sqlite => PoolKind::Sqlite(db::sqlite::connect(&url).await?), DatabaseType::Redis => { @@ -382,9 +382,9 @@ 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::Mysql if config.needs_bare_mysql() => PoolKind::Mysql(db::mysql::connect_bare(&url).await?, true), + DatabaseType::Mysql => PoolKind::Mysql(db::mysql::connect(&url).await?, false), + DatabaseType::Doris | DatabaseType::StarRocks => PoolKind::Mysql(db::mysql::connect_bare(&url).await?, true), DatabaseType::Postgres | DatabaseType::Redshift => PoolKind::Postgres(db::postgres::connect(&url).await?), DatabaseType::Sqlite => PoolKind::Sqlite(db::sqlite::connect(&url).await?), DatabaseType::Redis => { @@ -458,7 +458,7 @@ pub async fn disconnect_db( for key in keys_to_remove { if let Some(pool) = conns.remove(&key) { match pool { - PoolKind::Mysql(p) => p.close().await, + PoolKind::Mysql(p, _) => p.close().await, PoolKind::Postgres(p) => p.close().await, PoolKind::Sqlite(p) => p.close().await, PoolKind::Redis(_) => {}, diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index c406dce54..0faae57b5 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -140,10 +140,11 @@ async fn do_execute( }) .await } - PoolKind::Mysql(p) => { + PoolKind::Mysql(p, bare) => { let p = p.clone(); + let bare = *bare; drop(connections); - wait_for_query(cancel_token, db::mysql::execute_query(&p, sql)) + wait_for_query(cancel_token, db::mysql::execute_query(&p, sql, bare)) .await .map(truncate_result) } diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index ee578d983..45f88a038 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -111,7 +111,7 @@ pub async fn list_databases( let pool = connections.get(&connection_id).ok_or("Connection not found")?; match pool { - PoolKind::Mysql(p) => db::mysql::list_databases(p).await, + PoolKind::Mysql(p, _) => db::mysql::list_databases(p).await, PoolKind::Postgres(p) => db::postgres::list_databases(p).await, PoolKind::Sqlite(p) => db::sqlite::list_databases(p).await, PoolKind::DuckDb(_) => Ok(vec![db::DatabaseInfo { name: "main".to_string() }]), @@ -186,7 +186,7 @@ pub async fn list_tables( let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { - PoolKind::Mysql(p) => db::mysql::list_tables(p, &schema).await, + PoolKind::Mysql(p, _) => db::mysql::list_tables(p, &schema).await, PoolKind::Postgres(p) => db::postgres::list_tables(p, &schema).await, PoolKind::Sqlite(p) => db::sqlite::list_tables(p, &schema).await, _ => Ok(vec![]), @@ -230,7 +230,7 @@ pub async fn get_columns( let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { - PoolKind::Mysql(p) => db::mysql::get_columns(p, &schema, &table).await, + PoolKind::Mysql(p, _) => db::mysql::get_columns(p, &schema, &table).await, PoolKind::Postgres(p) => db::postgres::get_columns(p, &schema, &table).await, PoolKind::Sqlite(p) => db::sqlite::get_columns(p, &schema, &table).await, _ => Ok(vec![]), @@ -265,7 +265,7 @@ pub async fn list_indexes( let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { - PoolKind::Mysql(p) => db::mysql::list_indexes(p, &schema, &table).await, + PoolKind::Mysql(p, _) => db::mysql::list_indexes(p, &schema, &table).await, PoolKind::Postgres(p) => db::postgres::list_indexes(p, &schema, &table).await, PoolKind::Sqlite(p) => db::sqlite::list_indexes(p, &schema, &table).await, _ => Ok(vec![]), @@ -300,7 +300,7 @@ pub async fn list_foreign_keys( let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { - PoolKind::Mysql(p) => db::mysql::list_foreign_keys(p, &schema, &table).await, + PoolKind::Mysql(p, _) => db::mysql::list_foreign_keys(p, &schema, &table).await, PoolKind::Postgres(p) => db::postgres::list_foreign_keys(p, &schema, &table).await, PoolKind::Sqlite(p) => db::sqlite::list_foreign_keys(p, &schema, &table).await, _ => Ok(vec![]), @@ -335,7 +335,7 @@ pub async fn list_triggers( let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { - PoolKind::Mysql(p) => db::mysql::list_triggers(p, &schema, &table).await, + PoolKind::Mysql(p, _) => db::mysql::list_triggers(p, &schema, &table).await, PoolKind::Postgres(p) => db::postgres::list_triggers(p, &schema, &table).await, PoolKind::Sqlite(p) => db::sqlite::list_triggers(p, &schema, &table).await, _ => Ok(vec![]), @@ -391,7 +391,7 @@ pub async fn get_table_ddl( let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { - PoolKind::Mysql(p) => mysql_ddl(p, &table).await, + PoolKind::Mysql(p, _) => mysql_ddl(p, &table).await, PoolKind::Postgres(p) => pg_ddl(p, &schema, &table).await, PoolKind::Sqlite(p) => sqlite_ddl(p, &table).await, _ => Err("DDL not supported for this database type".to_string()), diff --git a/src-tauri/src/commands/transfer.rs b/src-tauri/src/commands/transfer.rs index 5c03cdd1a..965c69a39 100644 --- a/src-tauri/src/commands/transfer.rs +++ b/src-tauri/src/commands/transfer.rs @@ -321,10 +321,11 @@ pub(crate) async fn execute_on_pool( let pool = connections.get(pool_key).ok_or("Connection not found")?; match pool { - PoolKind::Mysql(p) => { + PoolKind::Mysql(p, bare) => { let p = p.clone(); + let bare = *bare; drop(connections); - db::mysql::execute_query(&p, sql).await + db::mysql::execute_query(&p, sql, bare).await } PoolKind::Postgres(p) => { let p = p.clone(); @@ -475,7 +476,7 @@ async fn get_columns_for_transfer( let schema = schema.to_string(); let table = table.to_string(); match pool { - PoolKind::Mysql(p) => { + PoolKind::Mysql(p, _) => { let p = p.clone(); drop(connections); db::mysql::get_columns(&p, &schema, &table).await diff --git a/src-tauri/src/db/mysql.rs b/src-tauri/src/db/mysql.rs index 2f5ee70c1..3e58f9785 100644 --- a/src-tauri/src/db/mysql.rs +++ b/src-tauri/src/db/mysql.rs @@ -1,7 +1,7 @@ use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc}; use rust_decimal::Decimal; use sqlx::mysql::{MySqlPool, MySqlPoolOptions, MySqlRow}; -use sqlx::{Column, Row, TypeInfo, ValueRef}; +use sqlx::{Column, Executor, Row, TypeInfo, ValueRef}; use std::time::{Duration, Instant}; use super::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; @@ -250,40 +250,68 @@ pub async fn get_columns( .collect()) } -pub async fn execute_query(pool: &MySqlPool, sql: &str) -> Result { +pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result { let start = Instant::now(); let trimmed = sql.trim().to_uppercase(); if trimmed.starts_with("SELECT") || trimmed.starts_with("SHOW") || trimmed.starts_with("DESCRIBE") || trimmed.starts_with("EXPLAIN") { - let rows: Vec = sqlx::raw_sql(sql) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; + if bare { + 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 (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| { - (0..row.len()) - .map(|i| mysql_value_to_json(row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) - .collect() + let result_rows: Vec> = rows + .iter() + .map(|row| { + (0..row.len()) + .map(|i| mysql_value_to_json(row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) + .collect() + }) + .collect(); + + Ok(QueryResult { + columns, + rows: result_rows, + affected_rows: 0, + execution_time_ms: start.elapsed().as_millis(), + truncated: false, }) - .collect(); + } else { + let desc = pool.describe(sql).await.map_err(|e| e.to_string())?; + let columns: Vec = 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(); - Ok(QueryResult { - columns, - rows: result_rows, - affected_rows: 0, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - }) + let rows: Vec = sqlx::query(sql) + .fetch_all(pool) + .await + .map_err(|e| e.to_string())?; + + let result_rows: Vec> = rows + .iter() + .map(|row| { + (0..row.len()) + .map(|i| mysql_value_to_json(row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) + .collect() + }) + .collect(); + + Ok(QueryResult { + columns, + rows: result_rows, + affected_rows: 0, + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + }) + } } else { let result = sqlx::raw_sql(sql) .execute(pool) diff --git a/src/App.vue b/src/App.vue index f8c17704d..5acfff5fb 100644 --- a/src/App.vue +++ b/src/App.vue @@ -1,5 +1,5 @@