diff --git a/crates/dbx-core/src/connection/mod.rs b/crates/dbx-core/src/connection/mod.rs index e79d95522b..efe52e886b 100644 --- a/crates/dbx-core/src/connection/mod.rs +++ b/crates/dbx-core/src/connection/mod.rs @@ -19,7 +19,8 @@ use crate::agent_connection::{ agent_connect_params, agent_connect_params_with_role, h2_file_path_from_jdbc_url, hive_uses_zookeeper_discovery, is_h2_file_connection, mongo_legacy_error_with_auth_hint, mongo_uses_legacy_driver, oracle_alternate_connect_config_labels, oracle_alternate_connect_configs, oracle_error_with_driver_hint, - should_retry_mongo_with_legacy_driver, trino_like_jdbc_connection_string, AgentSessionRole, + pick_legacy_postgres_like_database, should_retry_mongo_with_legacy_driver, trino_like_jdbc_connection_string, + AgentSessionRole, }; use crate::agent_manager::{AgentManager, JavaRuntimeMode, DEFAULT_JRE_KEY}; use crate::agent_recovery::{RecoveryDecision, RecoveryPolicy, RecoveryScope}; @@ -966,6 +967,19 @@ fn metadata_pool_database<'a>(config: Option<&ConnectionConfig>, database: Optio } } +/// Always-present KingbaseES/Vastbase catalog used to discover a default database for +/// legacy connections that were saved without one (see issue #9491). +const LEGACY_POSTGRES_LIKE_PROBE_DATABASE: &str = "template1"; +const LEGACY_POSTGRES_LIKE_PROBE_TIMEOUT: Duration = Duration::from_secs(30); + +/// Whether a connection relies on the historical `postgres` default that +/// `ConnectionConfig::default_database()` applies to KingbaseES/Vastbase connections +/// saved without an explicit database. +fn needs_legacy_postgres_like_database_probe(config: &ConnectionConfig) -> bool { + matches!(config.db_type, DatabaseType::Kingbase | DatabaseType::Vastbase) + && config.database.as_deref().is_none_or(|database| database.trim().is_empty()) +} + pub fn sqlserver_legacy_driver_error(agent_error: &str) -> String { // This mapper handles both AgentManager launch strings and Agent call errors, so context // must remain before any structured-error compatibility marker. @@ -2184,6 +2198,91 @@ impl AppState { .await } + /// `KingbaseES`/`Vastbase` do not guarantee the `postgres` database that + /// `ConnectionConfig::default_database()` assumes for saved connections without an + /// explicit database (see issue #9491). When such a connection is opened — an + /// upgraded legacy record or a cleared "default database" — probe the + /// always-present `template1` catalog and open a database that actually exists + /// instead of failing the whole connection with `database "postgres" does not exist`. + async fn resolve_legacy_postgres_like_database( + &self, + connection_id: &str, + db_config: &ConnectionConfig, + ) -> Option { + if !needs_legacy_postgres_like_database_probe(db_config) { + return None; + } + // Reuse the metadata pool machinery: it applies the active session credentials and + // the connection's transport layers, and it is cached by pool key, so repeated + // connects and the object browser share one probe connection. The pool stays + // registered so a concurrent reader of the object browser is never torn down; + // `remove_connection_pools` closes it together with the connection's other pools. + // + // The probe always passes `template1` explicitly, so the resolver cannot re-enter + // itself; `Box::pin` only satisfies the compiler's async-recursion requirement. + let pool_key = match Box::pin(self.get_or_create_metadata_pool_for_session( + connection_id, + Some(LEGACY_POSTGRES_LIKE_PROBE_DATABASE), + None, + )) + .await + { + Ok(pool_key) => pool_key, + Err(error) => { + log::warn!( + "Failed to open a '{LEGACY_POSTGRES_LIKE_PROBE_DATABASE}' probe connection for '{connection_id}' while resolving a legacy default database: {error}" + ); + return None; + } + }; + let databases = { + let pool_handle = self.pool_handle(&pool_key).await; + let client = pool_handle.as_ref().and_then(|pool| match pool { + PoolKind::Agent(client) => Some(client.clone()), + _ => None, + }); + match client { + Some(client) => client + .lock() + .await + .list_databases::>(Some(LEGACY_POSTGRES_LIKE_PROBE_TIMEOUT)) + .await + .map(|databases| databases.into_iter().map(|database| database.name).collect::>()), + None => Err("Legacy default database probe pool is not an Agent pool".to_string()), + } + }; + match databases { + Ok(databases) => { + let resolved = pick_legacy_postgres_like_database(&databases); + match &resolved { + Some(database) => { + log::info!( + "Resolved legacy default database '{database}' for '{connection_id}' from {databases:?}" + ); + // Write the discovered database back so the legacy record stops relying on + // the `postgres` default and later starts skip the probe. Temporary + // connection-test ids are not persisted and are simply ignored. + if let Err(error) = self.save_connection_database(connection_id, database).await { + log::warn!( + "Failed to persist the resolved legacy default database '{database}' for '{connection_id}': {error}" + ); + } + } + None => log::warn!( + "No usable database found for the legacy default database of '{connection_id}' (candidates: {databases:?})" + ), + } + resolved + } + Err(error) => { + log::warn!( + "Failed to list databases for '{connection_id}' while resolving a legacy default database: {error}" + ); + None + } + } + } + async fn get_or_create_pool_for_session_inner( &self, connection_id: &str, @@ -2240,6 +2339,9 @@ impl AppState { self.reset_connection_transport_for_config(connection_id, &db_config).await; return Err(err); } + if let Some(database) = self.resolve_legacy_postgres_like_database(connection_id, &db_config).await { + db_config.database = Some(database); + } if db_config.db_type != DatabaseType::Plugin { probe_connection_endpoint(&db_config, &host, port).await?; } @@ -4680,6 +4782,16 @@ impl AppState { } } + /// Persist the database resolved for a legacy empty-database connection and keep the + /// runtime config in sync so peer pool creations reuse the discovered database. + pub async fn save_connection_database(&self, connection_id: &str, database: &str) -> Result<(), String> { + self.storage.save_connection_database(connection_id, database).await?; + if let Some(config) = self.configs.write().await.get_mut(connection_id) { + config.database = Some(database.to_string()); + } + Ok(()) + } + pub async fn save_connection_database_info( &self, connection_id: &str, @@ -7873,6 +7985,32 @@ mod tests { ); } + #[test] + fn legacy_postgres_like_database_probe_only_applies_to_unconfigured_kingbase() { + let mut config = mysql_config(Some("SAMPLES")); + config.db_type = DatabaseType::Kingbase; + assert!(!super::needs_legacy_postgres_like_database_probe(&config)); + + config.database = None; + assert!(super::needs_legacy_postgres_like_database_probe(&config)); + + config.database = Some(" ".to_string()); + assert!(super::needs_legacy_postgres_like_database_probe(&config)); + + config.database = Some("application".to_string()); + assert!(!super::needs_legacy_postgres_like_database_probe(&config)); + + config.database = None; + config.db_type = DatabaseType::Vastbase; + assert!(super::needs_legacy_postgres_like_database_probe(&config)); + + config.db_type = DatabaseType::Postgres; + assert!(!super::needs_legacy_postgres_like_database_probe(&config)); + + config.db_type = DatabaseType::Mysql; + assert!(!super::needs_legacy_postgres_like_database_probe(&config)); + } + #[test] fn database_scoped_pool_keys_preserve_identifier_whitespace() { assert_eq!( diff --git a/crates/dbx-core/src/persistence/storage.rs b/crates/dbx-core/src/persistence/storage.rs index 95c42e538e..d02c07af62 100644 --- a/crates/dbx-core/src/persistence/storage.rs +++ b/crates/dbx-core/src/persistence/storage.rs @@ -3597,6 +3597,36 @@ impl Storage { .await } + /// Update only the persisted `database` of one saved connection. + /// + /// Unlike `save_connections`, this is an in-place update for a value DBX discovered + /// while connecting (a legacy KingbaseES/Vastbase connection without a database) and + /// must not replace the saved connection list or touch separately stored secrets. + /// Returns `false` when the connection is not persisted, for example a temporary + /// connection-test id. + pub async fn save_connection_database(&self, connection_id: &str, database: &str) -> Result { + let connection_id = connection_id.to_string(); + let database = database.to_string(); + self.with_conn(move |conn| { + let json = conn + .query_row("SELECT config_json FROM connections WHERE id = ?1", [&connection_id], |row| { + row.get::<_, String>(0) + }) + .optional() + .map_err(|error| error.to_string())?; + let Some(json) = json else { + return Ok(false); + }; + let mut config: ConnectionConfig = serde_json::from_str(&json).map_err(|error| error.to_string())?; + config.database = Some(database); + let json = serde_json::to_string(&config).map_err(|error| error.to_string())?; + conn.execute("UPDATE connections SET config_json = ?1 WHERE id = ?2", params![json, connection_id]) + .map(|_| true) + .map_err(|error| error.to_string()) + }) + .await + } + /// Update only the persisted MQTT saved-topic metadata for one connection. /// /// Unlike `save_connections`, this is an in-place update and must not replace @@ -6600,6 +6630,26 @@ mod tests { let _ = std::fs::remove_file(path); } + #[tokio::test] + async fn save_connection_database_updates_only_the_target_and_preserves_secrets() { + let path = temp_db_path("connection-database"); + let storage = Storage::open(&path).await.unwrap(); + let target = mq_connection("target", "target-secret"); + let untouched = mq_connection("untouched", "untouched-secret"); + storage.save_connections(&[target.clone(), untouched.clone()]).await.unwrap(); + + assert!(storage.save_connection_database("target", "dbx_demo").await.unwrap()); + assert!(!storage.save_connection_database("missing", "dbx_demo").await.unwrap()); + + let loaded = storage.load_connections().await.unwrap(); + let stored_target = loaded.iter().find(|config| config.id == "target").unwrap(); + assert_eq!(stored_target.database.as_deref(), Some("dbx_demo")); + assert_eq!(mq_token(stored_target), Some("target-secret")); + assert_eq!(loaded.iter().find(|config| config.id == "untouched"), Some(&untouched)); + + let _ = std::fs::remove_file(path); + } + #[tokio::test] async fn save_connection_driver_profile_rejects_a_stale_connection_config() { let path = temp_db_path("connection-driver-profile-stale"); diff --git a/crates/dbx-drivers/src/agent_connection.rs b/crates/dbx-drivers/src/agent_connection.rs index c396cd945d..e360fa392f 100644 --- a/crates/dbx-drivers/src/agent_connection.rs +++ b/crates/dbx-drivers/src/agent_connection.rs @@ -521,6 +521,26 @@ fn postgres_like_agent_database<'a>(config: &'a ConnectionConfig, database: &'a config.effective_database().unwrap_or("") } +/// Picks the database to open for a legacy KingbaseES/Vastbase connection that was +/// saved without an explicit database. Both vendor servers may lack the `postgres` +/// database that `ConnectionConfig::default_database()` falls back to, so the caller +/// probes the always-present `template1` catalog and hands the found databases here. +/// The historical `postgres` default still wins when it exists so upgraded +/// connections keep their previous behavior. +pub fn pick_legacy_postgres_like_database(databases: &[String]) -> Option { + let names: Vec<&str> = databases.iter().map(|name| name.trim()).filter(|name| !name.is_empty()).collect(); + let preferred = + |wanted: &str| names.iter().find(|name| name.eq_ignore_ascii_case(wanted)).map(|name| (*name).to_string()); + if let Some(database) = preferred("postgres").or_else(|| preferred("kingbase")).or_else(|| preferred("test")) { + return Some(database); + } + names.iter().find(|name| !is_template_database(name)).or_else(|| names.first()).map(|name| (*name).to_string()) +} + +fn is_template_database(name: &str) -> bool { + name.eq_ignore_ascii_case("template0") || name.eq_ignore_ascii_case("template1") +} + pub fn oracle_alternate_connect_config(config: &ConnectionConfig, err: &str) -> Option { oracle_alternate_connect_configs(config, err).into_iter().next() } @@ -992,6 +1012,33 @@ mod tests { assert_eq!(params["connection_string"], "jdbc:kingbase8://kingbase.example.com:54321/application"); } + #[test] + fn legacy_postgres_like_database_prefers_postgres_then_kingbase() { + let databases = vec!["SAMPLES".to_string(), "kingbase".to_string(), "postgres".to_string()]; + assert_eq!(pick_legacy_postgres_like_database(&databases).as_deref(), Some("postgres")); + + let databases = vec!["SAMPLES".to_string(), "KingBase".to_string()]; + assert_eq!(pick_legacy_postgres_like_database(&databases).as_deref(), Some("KingBase")); + + let databases = vec!["security".to_string(), "test".to_string()]; + assert_eq!(pick_legacy_postgres_like_database(&databases).as_deref(), Some("test")); + } + + #[test] + fn legacy_postgres_like_database_skips_templates_and_blank_names() { + let databases = + vec![" ".to_string(), "TEMPLATE0".to_string(), "template1".to_string(), "dbx_demo".to_string()]; + assert_eq!(pick_legacy_postgres_like_database(&databases).as_deref(), Some("dbx_demo")); + } + + #[test] + fn legacy_postgres_like_database_falls_back_to_templates_then_none() { + let databases = vec!["template1".to_string(), "TEMPLATE0".to_string()]; + assert_eq!(pick_legacy_postgres_like_database(&databases).as_deref(), Some("template1")); + assert_eq!(pick_legacy_postgres_like_database(&[]), None); + assert_eq!(pick_legacy_postgres_like_database(&["".to_string()]), None); + } + #[test] fn uxdb_agent_params_use_vendor_jdbc_url() { let cfg = config(DatabaseType::Uxdb, Some("uxdb")); diff --git a/docs/screenshots/issue-9491-kingbase-default-database-after.png b/docs/screenshots/issue-9491-kingbase-default-database-after.png new file mode 100644 index 0000000000..84fa4c8306 Binary files /dev/null and b/docs/screenshots/issue-9491-kingbase-default-database-after.png differ diff --git a/docs/screenshots/issue-9491-kingbase-default-database-before.png b/docs/screenshots/issue-9491-kingbase-default-database-before.png new file mode 100644 index 0000000000..ef327e624a Binary files /dev/null and b/docs/screenshots/issue-9491-kingbase-default-database-before.png differ diff --git a/docs/screenshots/issue-9491-kingbase-default-database-writeback.png b/docs/screenshots/issue-9491-kingbase-default-database-writeback.png new file mode 100644 index 0000000000..3174e7ec02 Binary files /dev/null and b/docs/screenshots/issue-9491-kingbase-default-database-writeback.png differ