use std::sync::Arc; use tauri::State; pub use dbx_core::connection::{ connection_url_for_endpoint, expand_tilde, metadata_connection_config, probe_connection_endpoint, redacted_connection_url_for_endpoint, AppState, MysqlMode, OraclePool, PoolKind, }; use dbx_core::db; use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; #[tauri::command] pub async fn save_connections(state: State<'_, Arc>, configs: Vec) -> Result<(), String> { state.storage.save_connections(&configs).await } #[tauri::command] pub async fn load_connections(state: State<'_, Arc>) -> Result, String> { state.storage.load_connections().await } #[tauri::command] pub async fn save_sidebar_layout(state: State<'_, Arc>, layout: serde_json::Value) -> Result<(), String> { state.storage.save_sidebar_layout(&layout).await } #[tauri::command] pub async fn load_sidebar_layout(state: State<'_, Arc>) -> Result, String> { state.storage.load_sidebar_layout().await } #[tauri::command] pub async fn test_connection(state: State<'_, Arc>, config: ConnectionConfig) -> Result { let tunnel_id = format!("{}:test", config.id); let connection_id = if config.ssh_enabled && !config.ssh_host.is_empty() { tunnel_id.as_str() } else { config.id.as_str() }; let (host, port) = state.connection_host_port(connection_id, &config).await?; let probe_result = probe_connection_endpoint(&config, &host, port).await; let url = connection_url_for_endpoint(&config, &host, port); let target = redacted_connection_url_for_endpoint(&config, &host, port); log::info!("[test_connection] db_type={:?} target={}", config.db_type, target); let result = match probe_result { Err(e) => Err(e), Ok(()) => match config.db_type { 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), }, 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), }, DatabaseType::Postgres | DatabaseType::Redshift => match db::postgres::connect(&url).await { Ok(pool) => { pool.close().await; Ok("Connection successful".to_string()) } Err(e) => Err(e), }, DatabaseType::Sqlite => match db::sqlite::connect_path(&expand_tilde(&config.host)).await { Ok(pool) => { pool.close().await; Ok("Connection successful".to_string()) } Err(e) => Err(e), }, DatabaseType::Redis => db::redis_driver::connect(&url).await.map(|_| "Connection successful".to_string()), DatabaseType::DuckDb => duckdb::Connection::open(&expand_tilde(&config.host)) .map(|_| "Connection successful".to_string()) .map_err(|e| e.to_string()), DatabaseType::MongoDb => match db::mongo_driver::connect(&url).await { Ok(client) => { db::mongo_driver::test_connection(&client).await.map(|_| "Connection successful".to_string()) } Err(e) => Err(e.to_string()), }, DatabaseType::ClickHouse => { let username = if config.username.is_empty() { None } else { Some(config.username.clone()) }; let password = if config.password.is_empty() { None } else { Some(config.password.clone()) }; let client = db::clickhouse_driver::ChClient::new(&url, username, password); db::clickhouse_driver::test_connection(&client).await.map(|_| "Connection successful".to_string()) } DatabaseType::SqlServer => { db::sqlserver::connect(&host, port, &config.username, &config.password, config.database.as_deref()) .await .map(|_| "Connection successful".to_string()) } DatabaseType::Oracle => db::oracle_driver::connect( &host, port, config.database.as_deref().unwrap_or("ORCL"), &config.username, &config.password, config.sysdba, ) .await .map(|_| "Connection successful".to_string()), DatabaseType::Elasticsearch => { let client = db::elasticsearch_driver::EsClient::new(&url, Some(&config.username), Some(&config.password)); db::elasticsearch_driver::test_connection(&client).await.map(|_| "Connection successful".to_string()) } DatabaseType::Dameng => db::dm_driver::connect( &host, port, config.database.as_deref().unwrap_or(""), &config.username, &config.password, ) .await .map(|_| "Connection successful".to_string()), DatabaseType::Gaussdb => db::gaussdb_driver::connect( &host, port, config.database.as_deref().unwrap_or(""), &config.username, &config.password, ) .await .map(|_| "Connection successful".to_string()), DatabaseType::Jdbc => state.test_external_driver("jdbc", &config).await, }, }; if config.ssh_enabled && !config.ssh_host.is_empty() { state.tunnels.stop_tunnel(&tunnel_id).await; } result } #[tauri::command] pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfig) -> Result { let id = config.id.clone(); let db_config = metadata_connection_config(&config); let (host, port) = state.connection_host_port(&id, &db_config).await?; probe_connection_endpoint(&db_config, &host, port).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?, MysqlMode::Bare) } DatabaseType::Mysql => PoolKind::Mysql(db::mysql::connect(&url).await?, MysqlMode::Normal), DatabaseType::Doris | DatabaseType::StarRocks => { PoolKind::Mysql(db::mysql::connect_bare(&url).await?, MysqlMode::Bare) } DatabaseType::Postgres | DatabaseType::Redshift => PoolKind::Postgres(db::postgres::connect(&url).await?), DatabaseType::Sqlite => PoolKind::Sqlite(db::sqlite::connect_path(&expand_tilde(&db_config.host)).await?), DatabaseType::Redis => { let con = db::redis_driver::connect(&url).await?; PoolKind::Redis(tokio::sync::Mutex::new(con)) } DatabaseType::DuckDb => { let con = duckdb::Connection::open(&expand_tilde(&db_config.host)).map_err(|e| e.to_string())?; PoolKind::DuckDb(std::sync::Arc::new(std::sync::Mutex::new(con))) } DatabaseType::MongoDb => { let client = db::mongo_driver::connect(&url).await?; db::mongo_driver::test_connection(&client).await?; PoolKind::MongoDb(client) } DatabaseType::ClickHouse => { let username = if db_config.username.is_empty() { None } else { Some(db_config.username.clone()) }; let password = if db_config.password.is_empty() { None } else { Some(db_config.password.clone()) }; let client = db::clickhouse_driver::ChClient::new(&url, username, password); db::clickhouse_driver::test_connection(&client).await?; PoolKind::ClickHouse(client) } DatabaseType::SqlServer => { let client = db::sqlserver::connect( &host, port, &db_config.username, &db_config.password, db_config.database.as_deref(), ) .await?; PoolKind::SqlServer(std::sync::Arc::new(tokio::sync::Mutex::new(client))) } DatabaseType::Oracle => { let client = db::oracle_driver::connect( &host, port, db_config.database.as_deref().unwrap_or("ORCL"), &db_config.username, &db_config.password, db_config.sysdba, ) .await?; PoolKind::Oracle(std::sync::Arc::new(OraclePool::new(vec![client]))) } DatabaseType::Elasticsearch => { let client = db::elasticsearch_driver::EsClient::new(&url, Some(&db_config.username), Some(&db_config.password)); db::elasticsearch_driver::test_connection(&client).await?; PoolKind::Elasticsearch(client) } DatabaseType::Dameng => { let client = db::dm_driver::connect( &host, port, db_config.database.as_deref().unwrap_or(""), &db_config.username, &db_config.password, ) .await?; PoolKind::Dameng(std::sync::Arc::new(std::sync::Mutex::new(client))) } DatabaseType::Gaussdb => { let client = db::gaussdb_driver::connect( &host, port, db_config.database.as_deref().unwrap_or(""), &db_config.username, &db_config.password, ) .await?; PoolKind::Gaussdb(std::sync::Arc::new(tokio::sync::Mutex::new(client))) } DatabaseType::Jdbc => state.external_driver_pool("jdbc", &db_config).await?, }; state.connections.lock().await.insert(id.clone(), pool); state.configs.lock().await.insert(id.clone(), config); Ok(id) } #[tauri::command] pub async fn disconnect_db(state: State<'_, Arc>, connection_id: String) -> Result<(), String> { let mut conns = state.connections.lock().await; let keys_to_remove: Vec = conns.keys().filter(|k| *k == &connection_id || k.starts_with(&format!("{connection_id}:"))).cloned().collect(); for key in keys_to_remove { if let Some(pool) = conns.remove(&key) { match pool { PoolKind::Mysql(p, _) => p.close().await, PoolKind::Postgres(p) => p.close().await, PoolKind::Sqlite(p) => p.close().await, PoolKind::Redis(_) => {} PoolKind::DuckDb(_) => {} PoolKind::MongoDb(_) => {} PoolKind::ClickHouse(_) => {} PoolKind::SqlServer(_) => {} PoolKind::Oracle(_) => {} PoolKind::Elasticsearch(_) => {} PoolKind::Dameng(_) => {} PoolKind::Gaussdb(_) => {} PoolKind::ExternalDriver { .. } => {} } } } drop(conns); state.configs.lock().await.remove(&connection_id); state.tunnels.stop_tunnel(&connection_id).await; Ok(()) }