use std::sync::Arc; use axum::extract::State; use axum::Json; use dbx_core::models::connection::ConnectionConfig; use serde::Deserialize; use crate::error::AppError; use crate::state::WebState; #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct ConnectRequest { pub config: ConnectionConfig, } #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct DisconnectRequest { pub connection_id: String, } #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct SaveConnectionsRequest { pub configs: Vec, } pub async fn test_connection( State(state): State>, Json(body): Json, ) -> Result, AppError> { let config = body.config; let app = &state.app; // Store config temporarily let temp_id = format!("__test_{}", uuid::Uuid::new_v4()); app.configs.lock().await.insert(temp_id.clone(), config.clone()); // Try to connect let result = app.get_or_create_pool(&temp_id, config.database.as_deref()).await; // Clean up any pool keys created for the temporary connection, including // database-scoped keys like "__test_uuid:database". let mut connections = app.connections.lock().await; let temp_keys: Vec = connections.keys().filter(|key| key.starts_with(&temp_id)).cloned().collect(); for key in temp_keys { connections.remove(&key); } drop(connections); app.configs.lock().await.remove(&temp_id); match result { Ok(_) => Ok(Json("Connection successful".to_string())), Err(e) => Err(AppError(e)), } } pub async fn connect_db( State(state): State>, Json(body): Json, ) -> Result, AppError> { let config = body.config; let app = &state.app; let connection_id = config.id.clone(); app.configs.lock().await.insert(connection_id.clone(), config.clone()); app.get_or_create_pool(&connection_id, None).await.map_err(AppError)?; Ok(Json(connection_id)) } pub async fn disconnect_db( State(state): State>, Json(body): Json, ) -> Result, AppError> { let app = &state.app; let mut connections = app.connections.lock().await; let pool_prefix = format!("{}:", body.connection_id); let keys_to_remove: Vec = connections.keys().filter(|k| *k == &body.connection_id || k.starts_with(&pool_prefix)).cloned().collect(); for key in keys_to_remove { connections.remove(&key); } drop(connections); app.configs.lock().await.remove(&body.connection_id); app.tunnels.stop_tunnel(&body.connection_id).await; Ok(Json(())) } pub async fn save_connections( State(state): State>, Json(body): Json, ) -> Result, AppError> { state.app.storage.save_connections(&body.configs).await.map_err(AppError)?; cache_connection_configs(&state, &body.configs).await; Ok(Json(())) } pub async fn load_connections(State(state): State>) -> Result>, AppError> { let configs = state.app.storage.load_connections().await.map_err(AppError)?; cache_connection_configs(&state, &configs).await; Ok(Json(configs)) } async fn cache_connection_configs(state: &WebState, configs: &[ConnectionConfig]) { let mut runtime_configs = state.app.configs.lock().await; for config in configs { runtime_configs.insert(config.id.clone(), config.clone()); } } #[cfg(test)] mod tests { use super::{disconnect_db, save_connections, DisconnectRequest, SaveConnectionsRequest}; use crate::state::{LoginRateLimit, WebState}; use axum::extract::State; use axum::Json; use dbx_core::connection::{AppState, PoolKind}; use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; use dbx_core::storage::Storage; use std::collections::{HashMap, HashSet}; use std::sync::Arc; use tokio::sync::{Mutex, RwLock}; fn sqlite_config(id: &str, path: &str) -> ConnectionConfig { ConnectionConfig { id: id.to_string(), name: "SQLite".to_string(), db_type: DatabaseType::Sqlite, driver_profile: None, driver_label: None, url_params: None, host: path.to_string(), port: 0, username: String::new(), password: String::new(), database: None, color: None, ssh_enabled: false, ssh_host: String::new(), ssh_port: 22, ssh_user: String::new(), ssh_password: String::new(), ssh_key_path: String::new(), ssh_key_passphrase: String::new(), ssh_expose_lan: false, ssh_connect_timeout_secs: dbx_core::models::connection::default_ssh_connect_timeout_secs(), ssl: false, sysdba: false, connection_string: None, external_config: None, jdbc_driver_class: None, jdbc_driver_paths: Vec::new(), } } async fn test_web_state() -> (Arc, std::path::PathBuf) { let dir = std::env::temp_dir().join(format!("dbx-web-test-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); let app = Arc::new(AppState::new_with_plugin_dir(storage, dir.join("plugins"))); let state = Arc::new(WebState { app, data_dir: dir.clone(), password_hash: RwLock::new(None), sessions: RwLock::new(HashSet::new()), sse_channels: RwLock::new(HashMap::new()), login_rate_limit: Mutex::new(LoginRateLimit { fail_count: 0, locked_until: None }), }); (state, dir) } #[tokio::test] async fn save_connections_updates_runtime_config_cache() { let (state, dir) = test_web_state().await; let db_path = dir.join("app.db"); let config = sqlite_config("sqlite-conn", &db_path.to_string_lossy()); let result = save_connections(State(state.clone()), Json(SaveConnectionsRequest { configs: vec![config.clone()] })) .await; assert!(result.is_ok()); let configs = state.app.configs.lock().await; assert_eq!(configs.get("sqlite-conn").map(|c| c.host.as_str()), Some(config.host.as_str())); let _ = std::fs::remove_dir_all(dir); } #[tokio::test] async fn disconnect_db_keeps_connections_with_similar_prefixes() { let (state, dir) = test_web_state().await; let conn_path = dir.join("conn.db"); let conn2_path = dir.join("conn2.db"); std::fs::File::create(&conn_path).unwrap(); std::fs::File::create(&conn2_path).unwrap(); let conn_pool = dbx_core::db::sqlite::connect_path(&conn_path.to_string_lossy()).await.unwrap(); let conn2_pool = dbx_core::db::sqlite::connect_path(&conn2_path.to_string_lossy()).await.unwrap(); { let mut connections = state.app.connections.lock().await; connections.insert("conn".to_string(), PoolKind::Sqlite(conn_pool)); connections.insert("conn2".to_string(), PoolKind::Sqlite(conn2_pool)); } let result = disconnect_db(State(state.clone()), Json(DisconnectRequest { connection_id: "conn".to_string() })).await; assert!(result.is_ok()); let connections = state.app.connections.lock().await; assert!(!connections.contains_key("conn")); assert!(connections.contains_key("conn2")); let _ = std::fs::remove_dir_all(dir); } }