diff --git a/src-tauri/src/capabilities/commands.rs b/src-tauri/src/capabilities/commands.rs index e82d143..dca80fa 100644 --- a/src-tauri/src/capabilities/commands.rs +++ b/src-tauri/src/capabilities/commands.rs @@ -58,10 +58,7 @@ fn to_metadata(cap: &Capability) -> Value { async fn resolve_connection_config(app: &AppHandle, connection_id: &str) -> Result { let state: tauri::State<'_, crate::state::AppState> = app.state(); - let conns = state.connections.read().await; - let _active = conns - .get(connection_id) - .ok_or_else(|| format!("Connection not found: {}", connection_id))?; + state.ensure_connection(connection_id).await?; Ok(json!({ "connectionId": connection_id })) } diff --git a/src-tauri/src/capabilities/sql.rs b/src-tauri/src/capabilities/sql.rs index a5dba00..6ce0ce0 100644 --- a/src-tauri/src/capabilities/sql.rs +++ b/src-tauri/src/capabilities/sql.rs @@ -65,8 +65,8 @@ async fn resolve_adapter(connection_id: &str) -> Result = app.state(); - let mut conns = state.connections.write().await; - conns.insert(connection_id.to_string(), adapter.clone()); + state.connections.write().await.insert(connection_id.to_string(), adapter.clone()); + state.configs.write().await.insert(connection_id.to_string(), server_config); } Ok(adapter) diff --git a/src-tauri/src/commands/browse.rs b/src-tauri/src/commands/browse.rs index b40fef6..5f5421a 100644 --- a/src-tauri/src/commands/browse.rs +++ b/src-tauri/src/commands/browse.rs @@ -147,13 +147,7 @@ pub async fn list_databases( connection_id: String, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let databases = connection .list_databases() .await @@ -179,13 +173,7 @@ pub async fn list_schemas( database: String, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let schemas = connection .list_schemas(Some(&database)) @@ -214,13 +202,7 @@ pub async fn list_tables( schema: Option, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let tables = connection .list_tables(Some(&database), schema.as_deref()) @@ -251,13 +233,7 @@ pub async fn get_table_info( table_name: String, state: State<'_, AppState>, ) -> Result { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let table_info = connection .get_table_info(Some(&database), schema.as_deref(), &table_name) @@ -276,13 +252,7 @@ pub async fn list_columns( table_name: String, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let columns = match &connection { ActiveConnection::Postgres(adapter) => { @@ -344,13 +314,7 @@ pub async fn get_foreign_keys( schema: Option, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let result = connection .get_foreign_keys(Some(&database), schema.as_deref()) @@ -377,13 +341,7 @@ pub async fn get_table_data( query: TableDataQuery, state: State<'_, AppState>, ) -> Result { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let limit_val = query.limit.unwrap_or(100); let offset_val = query.offset.unwrap_or(0); @@ -510,13 +468,7 @@ pub async fn get_table_count( filter: Option, state: State<'_, AppState>, ) -> Result { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let filter_ref = filter.as_deref(); let db_type = get_db_type_string(&connection); @@ -685,13 +637,7 @@ pub async fn update_table_row( return Ok(()); } - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let db_type = get_db_type_string(&connection); @@ -823,13 +769,7 @@ pub async fn delete_table_row( return Err("Cannot delete row: no primary key values provided".to_string()); } - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; let db_type = get_db_type_string(&connection); @@ -930,13 +870,7 @@ pub async fn list_views( schema: Option, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; connection .list_views(Some(&database), schema.as_deref()) @@ -952,13 +886,7 @@ pub async fn list_procedures( schema: Option, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; connection .list_procedures(Some(&database), schema.as_deref()) @@ -974,13 +902,7 @@ pub async fn list_functions( schema: Option, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; connection .list_functions(Some(&database), schema.as_deref()) @@ -997,13 +919,7 @@ pub async fn list_triggers( table: String, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; connection .list_triggers(Some(&database), schema.as_deref(), &table) @@ -1020,13 +936,7 @@ pub async fn list_indexes( table: String, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; connection .list_indexes(Some(&database), schema.as_deref(), &table) @@ -1043,13 +953,7 @@ pub async fn list_foreign_keys( table: String, state: State<'_, AppState>, ) -> Result, String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; connection .list_foreign_keys_for_table(Some(&database), schema.as_deref(), &table) @@ -1067,13 +971,7 @@ pub async fn get_object_ddl( object_type: String, state: State<'_, AppState>, ) -> Result { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; connection .get_object_ddl( @@ -1096,13 +994,7 @@ pub async fn drop_object( object_type: String, state: State<'_, AppState>, ) -> Result<(), String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; connection .drop_object( @@ -1126,13 +1018,7 @@ pub async fn rename_object( new_name: String, state: State<'_, AppState>, ) -> Result<(), String> { - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + let connection = state.ensure_connection(&connection_id).await?; connection .rename_object( @@ -1205,14 +1091,11 @@ pub async fn build_table_search_filter( search_term: String, state: State<'_, AppState>, ) -> Result { - let connections = state.connections.read().await; - let connection = connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))?; + let connection = state.ensure_connection(&connection_id).await?; - let db_type = get_db_type_string(connection); + let db_type = get_db_type_string(&connection); - let columns = match connection { + let columns = match &connection { ActiveConnection::Postgres(adapter) => { let adapter = adapter.lock().await; if Some(database.as_str()) != adapter.config.database.as_deref() { diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 64da7b5..7997fb6 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -28,6 +28,8 @@ pub async fn connect_server( let mut connections = state.connections.write().await; connections.insert(id.clone(), connection.clone()); drop(connections); + // Keep the config so a lost session can be transparently recreated later. + state.configs.write().await.insert(id.clone(), config); let status = connection .test_connection() .await @@ -46,6 +48,8 @@ pub async fn disconnect_server(id: String, state: State<'_, AppState>) -> Result .remove(&id) .ok_or_else(|| format!("No active connection found for server '{}'", id))?; + state.configs.write().await.remove(&id); + let disconnect_result = connection.disconnect().await; if let Err(e) = disconnect_result { diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index 6641afa..930b0a3 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -129,10 +129,7 @@ pub async fn execute_query( } let temp_kind: Option = { - let connections = state.connections.read().await; - let connection = connections - .get(&connection_id) - .ok_or_else(|| "No active connection found".to_string())?; + let connection = state.ensure_connection(&connection_id).await?; match connection { ActiveConnection::Postgres(adapter) => { @@ -250,14 +247,8 @@ pub async fn execute_query( } } - // Common path: use the already-connected adapter - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| "No active connection found".to_string())? - .clone() - }; + // Common path: use the already-connected adapter (reconnecting if needed) + let connection = state.ensure_connection(&connection_id).await?; // Cache this handle for future cross-database lookups state @@ -334,11 +325,8 @@ pub async fn execute_sorted_query( ) -> Result, String> { // Derive database type from the active connection let db_type = { - let connections = state.connections.read().await; - let connection = connections - .get(&connection_id) - .ok_or_else(|| "No active connection found".to_string())?; - crate::commands::browse::get_db_type_string(connection).to_string() + let connection = state.ensure_connection(&connection_id).await?; + crate::commands::browse::get_db_type_string(&connection).to_string() }; // Inject the database type into options and build the wrapped SQL @@ -434,10 +422,7 @@ pub async fn explain_query( } let temp_kind: Option = { - let connections = state.connections.read().await; - let connection = connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))?; + let connection = state.ensure_connection(&connection_id).await?; match connection { ActiveConnection::Postgres(adapter) => { @@ -717,14 +702,8 @@ pub async fn explain_query( } } - // Common path: use the already-connected adapter - let connection = { - let connections = state.connections.read().await; - connections - .get(&connection_id) - .ok_or_else(|| format!("No active connection found for ID '{}'", connection_id))? - .clone() - }; + // Common path: use the already-connected adapter (reconnecting if needed) + let connection = state.ensure_connection(&connection_id).await?; let database_type = match &connection { ActiveConnection::Postgres(_) => "postgresql", diff --git a/src-tauri/src/commands/transfer.rs b/src-tauri/src/commands/transfer.rs index e77914a..70d96c8 100644 --- a/src-tauri/src/commands/transfer.rs +++ b/src-tauri/src/commands/transfer.rs @@ -15,10 +15,7 @@ pub async fn preview_export_data( preview_rows: u32, state: State<'_, AppState>, ) -> Result { - let connections = state.connections.read().await; - let connection = connections - .get(&request.connection_id) - .ok_or_else(|| "No active connection found".to_string())?; + let connection = state.ensure_connection(&request.connection_id).await?; match connection { ActiveConnection::Postgres(adapter) => { @@ -47,10 +44,7 @@ pub async fn execute_export_data( app_handle: AppHandle, state: State<'_, AppState>, ) -> Result { - let connections = state.connections.read().await; - let connection = connections - .get(&request.connection_id) - .ok_or_else(|| "No active connection found".to_string())?; + let connection = state.ensure_connection(&request.connection_id).await?; match connection { ActiveConnection::Postgres(adapter) => { @@ -93,10 +87,7 @@ pub async fn execute_import_data( app_handle: AppHandle, state: State<'_, AppState>, ) -> Result { - let connections = state.connections.read().await; - let connection = connections - .get(&request.connection_id) - .ok_or_else(|| "No active connection found".to_string())?; + let connection = state.ensure_connection(&request.connection_id).await?; match connection { ActiveConnection::Postgres(adapter) => { @@ -124,11 +115,7 @@ pub async fn preview_migration_data( request: MigrationRequest, state: State<'_, AppState>, ) -> Result { - let connections = state.connections.read().await; - - let source_connection = connections - .get(&request.source_connection_id) - .ok_or_else(|| "No source connection found".to_string())?; + let source_connection = state.ensure_connection(&request.source_connection_id).await?; match source_connection { ActiveConnection::Postgres(adapter) => { @@ -157,15 +144,9 @@ pub async fn execute_migration_data( app_handle: AppHandle, state: State<'_, AppState>, ) -> Result { - let connections = state.connections.read().await; - - let source_connection = connections - .get(&request.source_connection_id) - .ok_or_else(|| "No source connection found".to_string())?; + let source_connection = state.ensure_connection(&request.source_connection_id).await?; - let target_connection = connections - .get(&request.target_connection_id) - .ok_or_else(|| "No target connection found".to_string())?; + let target_connection = state.ensure_connection(&request.target_connection_id).await?; macro_rules! run_migration { ($source_adapter:expr, $target_adapter:expr) => { @@ -269,10 +250,7 @@ pub async fn auto_map_migration_columns( ) -> Result, String> { use crate::database::DatabaseAdapter; - let connections = state.connections.read().await; - let connection = connections - .get(&connection_id) - .ok_or_else(|| "No active connection found".to_string())?; + let connection = state.ensure_connection(&connection_id).await?; let target_db_type = match target_engine.to_lowercase().as_str() { "postgresql" | "postgres" => DatabaseType::PostgreSQL, @@ -307,10 +285,7 @@ pub async fn generate_ddl_for_objects( request: DdlRequest, state: State<'_, AppState>, ) -> Result { - let connections = state.connections.read().await; - let connection = connections - .get(&request.connection_id) - .ok_or_else(|| "No active connection found".to_string())?; + let connection = state.ensure_connection(&request.connection_id).await?; let engine = match connection { ActiveConnection::Postgres(_) => DatabaseType::PostgreSQL, @@ -441,10 +416,7 @@ pub async fn execute_sql_content( on_error: Option, state: State<'_, AppState>, ) -> Result { - let connections = state.connections.read().await; - let connection = connections - .get(&connection_id) - .ok_or_else(|| "No active connection found".to_string())?; + let connection = state.ensure_connection(&connection_id).await?; let strategy = on_error.as_deref().unwrap_or("stop"); let statements = split_sql_statements(&content); diff --git a/src-tauri/src/connection/cache.rs b/src-tauri/src/connection/cache.rs index 94ad1e9..c8dde42 100644 --- a/src-tauri/src/connection/cache.rs +++ b/src-tauri/src/connection/cache.rs @@ -85,13 +85,9 @@ impl ConnectionCache { } } - // Slow path: need to create connection - let conns = app_state.connections.read().await; - let connection = conns - .get(&key.connection_id) - .cloned() - .ok_or_else(|| format!("No active connection found for '{}'", key.connection_id))?; - drop(conns); + // Slow path: need to create connection — transparently reconnect when + // the session was lost or idle-evicted instead of erroring out. + let connection = app_state.ensure_connection(&key.connection_id).await?; // Evict if at capacity { diff --git a/src-tauri/src/connection/guardian.rs b/src-tauri/src/connection/guardian.rs index 5edc4f4..e495b04 100644 --- a/src-tauri/src/connection/guardian.rs +++ b/src-tauri/src/connection/guardian.rs @@ -359,18 +359,15 @@ impl ConnectionGuardian { } self.emit_state_change(connection_id, HealthState::Reconnecting, None); - // Try to reconnect by calling test_connection on the active connection + // Try to reconnect by recreating the adapter from the stored config. + // Pinging the stale adapter is not enough — the underlying session may + // be irrecoverable (server restart, dropped connection). let success = { let state = APP_HANDLE .get() .expect("APP_HANDLE not initialized") .state::(); - let conns = state.connections.read().await; - let conn = conns.get(connection_id); - match conn { - Some(c) => self.ping_connection(c).await, - None => false, - } + state.ensure_connection(connection_id).await.is_ok() }; if success { diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs index b253f23..72b6c07 100644 --- a/src-tauri/src/mcp_bridge.rs +++ b/src-tauri/src/mcp_bridge.rs @@ -511,10 +511,7 @@ async fn resolve_connection(connection_id: &str) -> Result { .ok_or_else(|| "AppHandle not initialized".to_string())?; use tauri::State; let state: State<'_, crate::state::AppState> = handle.state(); - let conns = state.connections.read().await; - let _active = conns - .get(connection_id) - .ok_or_else(|| format!("Connection not found: {}", connection_id))?; + state.ensure_connection(connection_id).await?; Ok(json!({ "connectionId": connection_id })) } diff --git a/src-tauri/src/state.rs b/src-tauri/src/state.rs index 9147268..d0df29f 100644 --- a/src-tauri/src/state.rs +++ b/src-tauri/src/state.rs @@ -287,6 +287,9 @@ pub struct AppConfig { pub struct AppState { /// Active database connections indexed by connection ID. pub connections: Arc>>, + /// Server configs for active connections, kept so a lost/evicted session can + /// be transparently recreated on the next user action. + pub configs: Arc>>, /// LRU cache for cross-database connection handles. pub cache: crate::connection::cache::ConnectionCache, /// SSH tunnel lifecycle manager. @@ -298,10 +301,61 @@ impl AppState { pub fn new() -> Self { Self { connections: Arc::new(RwLock::new(HashMap::new())), + configs: Arc::new(RwLock::new(HashMap::new())), cache: crate::connection::cache::ConnectionCache::default(), tunnels: TunnelManager::new(), } } + + /// Return the active connection for `connection_id`. + /// + /// If the session was lost or idle-evicted, transparently reconnects using + /// the stored server config so user actions never fail with a + /// "No active connection found" error. Errors are returned only when the + /// server cannot be reached or no config was ever stored for this id. + pub async fn ensure_connection(&self, connection_id: &str) -> Result { + if let Some(conn) = self.connections.read().await.get(connection_id) { + return Ok(conn.clone()); + } + + let config = self + .configs + .read() + .await + .get(connection_id) + .cloned() + .ok_or_else(|| format!("No active connection found for server '{}'", connection_id))?; + + let mut conn_config = config.to_connection_config()?; + let (host, port) = crate::commands::helpers::connection_host_port( + connection_id, + &conn_config, + &self.tunnels, + ) + .await?; + conn_config.host = host; + conn_config.port = port; + + let timeout_secs = conn_config.connect_timeout_secs; + let connection = tokio::time::timeout( + std::time::Duration::from_secs(timeout_secs), + crate::commands::helpers::create_and_connect_adapter(&config.db_type, conn_config), + ) + .await + .map_err(|_| format!("Reconnect timed out after {} seconds", timeout_secs))??; + + self.connections + .write() + .await + .insert(connection_id.to_string(), connection.clone()); + + // Notify the guardian so the frontend state flips back to CONNECTED. + if let Some(guardian) = crate::GUARDIAN.get() { + guardian.mark_healthy(connection_id, None).await; + } + + Ok(connection) + } } impl Default for AppState { diff --git a/src/components/connections/ServerCard.vue b/src/components/connections/ServerCard.vue index df8f887..e422358 100644 --- a/src/components/connections/ServerCard.vue +++ b/src/components/connections/ServerCard.vue @@ -27,6 +27,7 @@ const emit = defineEmits<{ (e: 'edit', connection: ServerConnection): void (e: 'delete', connection: ServerConnection): void (e: 'duplicate', connection: ServerConnection): void + (e: 'disconnect', connection: ServerConnection): void }>() const { t } = useI18n() @@ -91,6 +92,12 @@ const handleDoubleClick = () => emit('dblclick', props.connection) const handleEdit = () => emit('edit', props.connection) const handleDelete = () => emit('delete', props.connection) const handleDuplicate = () => emit('duplicate', props.connection) +const handleDisconnect = () => emit('disconnect', props.connection) + +const isConnected = computed(() => props.connectionStatus === ConnectionStatus.CONNECTED) +const connectTooltip = computed(() => + isConnected.value ? t('components.serverCard.actions.openEditor') : t('components.serverCard.actions.connect'), +)