From 86bb88ab66fa36955905d7744269c619c2546ceb Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 14 Aug 2026 15:27:55 +0000 Subject: [PATCH] fix(mongodb): install legacy agent before automatic fallback Closes #6167 --- crates/dbx-core/src/agent_manager.rs | 2 +- crates/dbx-core/src/agent_service.rs | 249 +++++++++++++++++++++++++++ crates/dbx-core/src/connection.rs | 19 +- src-tauri/src/commands/connection.rs | 81 +++++++-- 4 files changed, 336 insertions(+), 15 deletions(-) diff --git a/crates/dbx-core/src/agent_manager.rs b/crates/dbx-core/src/agent_manager.rs index 2c4a23ac16..7a4b9880ff 100644 --- a/crates/dbx-core/src/agent_manager.rs +++ b/crates/dbx-core/src/agent_manager.rs @@ -696,7 +696,7 @@ impl AgentManager { } pub fn is_jre_installed(&self, jre_key: &str) -> bool { - self.jre_java_path(jre_key).exists() + std::fs::metadata(self.jre_java_path(jre_key)).is_ok_and(|metadata| metadata.is_file() && metadata.len() > 0) } pub fn is_driver_installed(&self, db_type: &str) -> bool { diff --git a/crates/dbx-core/src/agent_service.rs b/crates/dbx-core/src/agent_service.rs index fcf0cf6516..4427a08b6e 100644 --- a/crates/dbx-core/src/agent_service.rs +++ b/crates/dbx-core/src/agent_service.rs @@ -432,6 +432,55 @@ pub async fn install_agent_driver_from( install_agent_driver_with_batch(am, db_type, source, &progress, None, None).await } +/// Ensure an already-selected fallback Agent can be launched, installing only +/// missing or invalid runtime artifacts. A usable installed driver is never +/// upgraded implicitly. +pub async fn ensure_agent_driver_ready(am: &AgentManager, db_type: &str) -> Result<(), String> { + ensure_agent_driver_ready_from(am, db_type, DownloadSource::Official).await +} + +async fn ensure_agent_driver_ready_from( + am: &AgentManager, + db_type: &str, + source: DownloadSource, +) -> Result<(), String> { + if agent_driver_runtime_readiness(am, db_type).is_ok() { + return Ok(()); + } + + let _installation_guard = am.installation_operation_lock.read().await; + let driver_lock = driver_operation_lock(am, db_type).await; + let _driver_guard = driver_lock.lock().await; + + // Another fallback may have completed installation while this task waited. + if agent_driver_runtime_readiness(am, db_type).is_ok() { + return Ok(()); + } + + let registry = fetch_registry_from(source).await?; + let progress = |_| {}; + if !am.is_driver_installed(db_type) { + install_agent_driver_from_registry(am, ®istry, source, db_type, &progress, None, None).await?; + } else if am.driver_requires_java_runtime(db_type) { + let state = am.load_state(); + let jre_key = state + .installed_drivers + .get(db_type) + .map(|driver| driver.jre.as_str()) + .or_else(|| agent_registry_driver(®istry, db_type).map(|driver| driver.jre.as_str())) + .unwrap_or(DEFAULT_JRE_KEY); + ensure_jre_from_registry(am, ®istry, source, jre_key, db_type, &progress, None, None).await?; + } + + agent_driver_runtime_readiness(am, db_type) +} + +fn agent_driver_runtime_readiness(am: &AgentManager, db_type: &str) -> Result<(), String> { + let state = am.load_state(); + let jre_key = state.installed_drivers.get(db_type).map(|driver| driver.jre.as_str()).unwrap_or(DEFAULT_JRE_KEY); + am.resolve_agent_launch_spec(&state, db_type, jre_key).map(|_| ()) +} + pub async fn upgrade_all_agent_drivers( am: &AgentManager, progress: impl Fn(AgentProgressEvent), @@ -2605,6 +2654,9 @@ mod agent_registry_install_tests { use super::*; use crate::agent_manager::{ArtifactFormat, ArtifactInfo, DriverInfo, InstalledDriver, JavaRuntimeConfig, JreInfo}; + static ENSURE_AGENT_TEST_LOCK: std::sync::LazyLock> = + std::sync::LazyLock::new(|| tokio::sync::Mutex::new(())); + fn test_manager(name: &str) -> AgentManager { let dir = std::env::temp_dir().join(format!("dbx-agent-registry-install-{name}-{}", uuid::Uuid::new_v4())); AgentManager::new_with_base_dir(dir) @@ -2945,6 +2997,203 @@ mod agent_registry_install_tests { std::fs::write(cache_path, archive).unwrap(); } + async fn cache_test_registry(registry: AgentRegistry) { + REGISTRY_CACHE.lock().await.insert(DownloadSource::Cnb, (std::time::Instant::now(), registry)); + } + + fn mongodb_registry_with_jre( + driver_version: &str, + driver_url: &str, + driver_size: u64, + jre_version: &str, + jre_url: &str, + jre_size: u64, + ) -> AgentRegistry { + let mut registry = registry_with_jar("mongodb", driver_version, driver_url, driver_size); + registry.jres = registry_with_jre(DEFAULT_JRE_KEY, jre_version, jre_url, jre_size).jres; + registry + } + + #[tokio::test] + async fn ensure_agent_runtime_installs_missing_driver_and_jre() { + let _test_guard = ENSURE_AGENT_TEST_LOCK.lock().await; + let manager = test_manager("ensure-missing-driver-and-jre"); + let driver_version = "0.1.47"; + let driver_url = "https://example.com/dbx-agent-mongodb.jar"; + let driver_bytes = test_agent_jar(); + let jre_version = "21.0.12"; + let jre_url = "https://example.com/dbx-jre.tar.gz"; + let jre_archive = build_jre_archive(&manager, DEFAULT_JRE_KEY); + let registry = mongodb_registry_with_jre( + driver_version, + driver_url, + driver_bytes.len() as u64, + jre_version, + jre_url, + jre_archive.len() as u64, + ); + write_cached_driver_download( + &manager, + "mongodb", + driver_version, + driver_url, + &manager.driver_jar_path("mongodb"), + &driver_bytes, + ); + write_cached_jre_download(&manager, DEFAULT_JRE_KEY, jre_version, jre_url, None, None, &jre_archive); + cache_test_registry(registry).await; + + ensure_agent_driver_ready_from(&manager, "mongodb", DownloadSource::Cnb).await.unwrap(); + + assert!(manager.is_driver_jar_valid("mongodb")); + assert!(manager.is_jre_installed(DEFAULT_JRE_KEY)); + let state = manager.load_state(); + assert_eq!(state.installed_drivers["mongodb"].version, driver_version); + assert_eq!(state.jre_versions[DEFAULT_JRE_KEY], jre_version); + } + + #[tokio::test] + async fn ensure_agent_runtime_ready_fast_path_preserves_valid_old_driver() { + let _test_guard = ENSURE_AGENT_TEST_LOCK.lock().await; + let manager = test_manager("ensure-ready-old-driver"); + let driver_path = manager.driver_jar_path("mongodb"); + std::fs::create_dir_all(driver_path.parent().unwrap()).unwrap(); + let old_driver_bytes = test_agent_jar(); + std::fs::write(&driver_path, &old_driver_bytes).unwrap(); + install_jre(&manager); + record_driver(&manager, "mongodb"); + + let latest_version = "9.9.9"; + let latest_url = "https://example.com/dbx-agent-mongodb-latest.jar"; + let latest_driver_bytes = test_agent_jar(); + let cache_path = write_cached_driver_download( + &manager, + "mongodb", + latest_version, + latest_url, + &driver_path, + &latest_driver_bytes, + ); + let registry_guard = REGISTRY_CACHE.lock().await; + + tokio::time::timeout( + std::time::Duration::from_millis(100), + ensure_agent_driver_ready_from(&manager, "mongodb", DownloadSource::Cnb), + ) + .await + .expect("ready Agent path must not wait for the registry cache") + .unwrap(); + drop(registry_guard); + + assert_eq!(std::fs::read(&driver_path).unwrap(), old_driver_bytes); + assert_eq!(manager.load_state().installed_drivers["mongodb"].version, "1.0.0"); + assert!(cache_path.exists(), "ready fast path must not consume a pending driver download"); + } + + #[tokio::test] + async fn ensure_agent_runtime_repairs_damaged_jre_without_upgrading_valid_driver() { + let _test_guard = ENSURE_AGENT_TEST_LOCK.lock().await; + let manager = test_manager("ensure-damaged-jre"); + let driver_path = manager.driver_jar_path("mongodb"); + std::fs::create_dir_all(driver_path.parent().unwrap()).unwrap(); + let old_driver_bytes = test_agent_jar(); + std::fs::write(&driver_path, &old_driver_bytes).unwrap(); + record_driver(&manager, "mongodb"); + let damaged_java = manager.jre_java_path(DEFAULT_JRE_KEY); + std::fs::create_dir_all(&damaged_java).unwrap(); + + let jre_version = "21.0.12"; + let jre_url = "https://example.com/dbx-jre.tar.gz"; + let jre_archive = build_jre_archive(&manager, DEFAULT_JRE_KEY); + let registry = mongodb_registry_with_jre( + "9.9.9", + "https://example.com/dbx-agent-mongodb-latest.jar", + test_agent_jar().len() as u64, + jre_version, + jre_url, + jre_archive.len() as u64, + ); + write_cached_jre_download(&manager, DEFAULT_JRE_KEY, jre_version, jre_url, None, None, &jre_archive); + cache_test_registry(registry).await; + + ensure_agent_driver_ready_from(&manager, "mongodb", DownloadSource::Cnb).await.unwrap(); + + assert!(manager.is_jre_installed(DEFAULT_JRE_KEY)); + assert_eq!(std::fs::read(&driver_path).unwrap(), old_driver_bytes); + assert_eq!(manager.load_state().installed_drivers["mongodb"].version, "1.0.0"); + } + + #[tokio::test] + async fn concurrent_ensure_agent_runtime_installs_once() { + let _test_guard = ENSURE_AGENT_TEST_LOCK.lock().await; + let manager = Arc::new(test_manager("ensure-concurrent-single-install")); + let driver_version = "0.1.47"; + let driver_url = "https://example.com/dbx-agent-mongodb.jar"; + let driver_bytes = test_agent_jar(); + let jre_version = "21.0.12"; + let jre_url = "https://example.com/dbx-jre.tar.gz"; + let jre_archive = build_jre_archive(&manager, DEFAULT_JRE_KEY); + let registry = mongodb_registry_with_jre( + driver_version, + driver_url, + driver_bytes.len() as u64, + jre_version, + jre_url, + jre_archive.len() as u64, + ); + let driver_cache_path = write_cached_driver_download( + &manager, + "mongodb", + driver_version, + driver_url, + &manager.driver_jar_path("mongodb"), + &driver_bytes, + ); + write_cached_jre_download(&manager, DEFAULT_JRE_KEY, jre_version, jre_url, None, None, &jre_archive); + cache_test_registry(registry).await; + + let first = ensure_agent_driver_ready_from(&manager, "mongodb", DownloadSource::Cnb); + let second = ensure_agent_driver_ready_from(&manager, "mongodb", DownloadSource::Cnb); + let (first, second) = tokio::join!(first, second); + + first.unwrap(); + second.unwrap(); + assert!(manager.is_driver_jar_valid("mongodb")); + assert!(manager.is_jre_installed(DEFAULT_JRE_KEY)); + assert!(!driver_cache_path.exists(), "the single successful install should consume the cached artifact"); + } + + #[tokio::test] + async fn ensure_agent_runtime_propagates_corrupt_install_error() { + let _test_guard = ENSURE_AGENT_TEST_LOCK.lock().await; + let manager = test_manager("ensure-corrupt-install-error"); + manager + .mutate_state(|state| { + state.java_runtime = JavaRuntimeConfig { mode: JavaRuntimeMode::System, custom_java_path: None }; + }) + .unwrap(); + let driver_version = "0.1.47"; + let driver_url = "https://example.com/dbx-agent-mongodb.jar"; + let corrupt_driver = b"not-a-jar"; + let registry = registry_with_jar("mongodb", driver_version, driver_url, corrupt_driver.len() as u64); + write_cached_driver_download( + &manager, + "mongodb", + driver_version, + driver_url, + &manager.driver_jar_path("mongodb"), + corrupt_driver, + ); + cache_test_registry(registry).await; + + let error = ensure_agent_driver_ready_from(&manager, "mongodb", DownloadSource::Cnb) + .await + .expect_err("corrupt driver install must fail"); + + assert!(error.contains("invalid or corrupt"), "unexpected error: {error}"); + assert!(!manager.load_state().installed_drivers.contains_key("mongodb")); + } + #[test] fn artifact_info_deserializes_sha256_metadata() { let expected_sha256 = "a".repeat(64); diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index c9d187f7ef..98d2fd4731 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -14,7 +14,7 @@ use crate::agent_connection::{ oracle_alternate_connect_configs, oracle_error_with_driver_hint, should_retry_mongo_with_legacy_driver, trino_like_jdbc_connection_string, AgentSessionRole, }; -use crate::agent_manager::{JavaRuntimeMode, DEFAULT_JRE_KEY}; +use crate::agent_manager::{AgentManager, JavaRuntimeMode, DEFAULT_JRE_KEY}; use crate::agent_recovery::{RecoveryDecision, RecoveryPolicy, RecoveryScope}; use crate::database_capabilities; use crate::db; @@ -1992,8 +1992,23 @@ impl AppState { if should_retry_mongo_with_legacy_driver(&native_err) { log::info!("Native MongoDB driver failed ({native_err}), falling back to agent driver"); let connect_params = serde_json::json!({ "connection": agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or("")) }); + let legacy_agent_key = + AgentManager::db_type_to_agent_key(&DatabaseType::MongoDb, Some("mongodb-legacy")) + .ok_or_else(|| "MongoDB (Legacy) Agent mapping is unavailable".to_string())?; + crate::agent_service::ensure_agent_driver_ready(&self.agent_manager, legacy_agent_key) + .await + .map_err(|err| { + format!("{native_err}\n\nFailed to prepare MongoDB (Legacy) fallback driver: {err}") + })?; let mut client = - self.agent_manager.spawn(&DatabaseType::MongoDb, Some("mongodb-legacy")).await?; + self.agent_manager.spawn(&DatabaseType::MongoDb, Some("mongodb-legacy")).await.map_err( + |err| { + format!( + "{native_err}\n\nFallback with MongoDB (Legacy) driver failed: {}", + mongo_legacy_error_with_auth_hint(&err) + ) + }, + )?; client.connect(connect_params).await.map_err(|err| { format!( "{native_err}\n\nFallback with MongoDB (Legacy) driver failed: {}", diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 8749330570..a425fc9ec3 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -14,7 +14,7 @@ pub use dbx_core::connection::{ }; use dbx_core::database_capabilities; use dbx_core::db; -use dbx_core::db::agent_driver::AgentMethod; +use dbx_core::db::agent_driver::{AgentDriverClient, AgentMethod}; use dbx_core::models::connection::{ database_info_from_protocol_value, rewrite_jdbc_url_host, ConnectionConfig, ConnectionTestResult, DatabaseConnectionInfo, DatabaseType, @@ -34,6 +34,36 @@ fn mongo_legacy_connect_params(config: &ConnectionConfig, host: &str, port: u16) }) } +fn mongo_legacy_fallback_error(native_error: &str, stage: &str, fallback_error: &str) -> String { + format!("{native_error}\n\n{stage}: {fallback_error}") +} + +async fn spawn_mongo_legacy_fallback_agent( + state: &AppState, + db_type: &DatabaseType, + native_error: &str, +) -> Result { + let agent_key = + dbx_core::agent_manager::AgentManager::db_type_to_agent_key(db_type, Some(MONGO_LEGACY_DRIVER_PROFILE)) + .ok_or_else(|| { + mongo_legacy_fallback_error( + native_error, + "Failed to prepare MongoDB (Legacy) fallback driver", + "Agent mapping is unavailable", + ) + })?; + dbx_core::agent_service::ensure_agent_driver_ready(&state.agent_manager, agent_key).await.map_err(|error| { + mongo_legacy_fallback_error(native_error, "Failed to prepare MongoDB (Legacy) fallback driver", &error) + })?; + state.agent_manager.spawn(db_type, Some(MONGO_LEGACY_DRIVER_PROFILE)).await.map_err(|error| { + mongo_legacy_fallback_error( + native_error, + "Fallback with MongoDB (Legacy) driver failed", + &mongo_legacy_error_with_auth_hint(&error), + ) + }) +} + fn mark_mongo_legacy_driver(config: &mut ConnectionConfig) -> bool { if config.db_type != DatabaseType::MongoDb { return false; @@ -169,8 +199,8 @@ mod tests { use super::load_connection_configs; use super::{ connect_sqlite_from_config, gaussdb_m_jdbc_command_config, mark_mongo_legacy_driver, - mongo_legacy_connect_params, persist_mongo_legacy_driver_profile, save_connection_configs, - sync_connection_configs, MONGO_LEGACY_DRIVER_LABEL, MONGO_LEGACY_DRIVER_PROFILE, + mongo_legacy_connect_params, mongo_legacy_fallback_error, persist_mongo_legacy_driver_profile, + save_connection_configs, sync_connection_configs, MONGO_LEGACY_DRIVER_LABEL, MONGO_LEGACY_DRIVER_PROFILE, }; use dbx_core::connection::{AppState, PoolKind}; use dbx_core::models::connection::{AttachedDatabaseConfig, ConnectionConfig, DatabaseType}; @@ -410,6 +440,30 @@ mod tests { assert!(!mark_mongo_legacy_driver(&mut config)); } + #[test] + fn automatic_mongo_desktop_paths_share_the_ensure_helper() { + let source = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/commands/connection.rs")); + let helper_call = ["spawn_mongo_legacy_fallback_", "agent("].concat(); + let ensure_call = ["ensure_agent_driver_", "ready("].concat(); + + assert_eq!(source.matches(&helper_call).count(), 3); + assert_eq!(source.matches(&ensure_call).count(), 1); + } + + #[test] + fn mongo_legacy_fallback_error_preserves_native_and_agent_errors() { + let error = mongo_legacy_fallback_error( + "native wire version error", + "Failed to prepare MongoDB (Legacy) fallback driver", + "registry unavailable", + ); + + assert_eq!( + error, + "native wire version error\n\nFailed to prepare MongoDB (Legacy) fallback driver: registry unavailable" + ); + } + #[tokio::test] async fn persist_mongo_legacy_driver_profile_updates_only_the_target_connection() { let dir = std::env::temp_dir().join(format!("dbx-tauri-mongo-profile-{}", uuid::Uuid::new_v4())); @@ -975,12 +1029,13 @@ async fn test_connection_with_info_inner( Err(e) => e, }; if should_retry_mongo_with_legacy_driver(&native_err) { - let am = &state.agent_manager; - let mut client = am.spawn(&config.db_type, Some("mongodb-legacy")).await?; + let mut client = + spawn_mongo_legacy_fallback_agent(state.as_ref(), &config.db_type, &native_err).await?; client.connect(mongo_legacy_connect_params(&config, &host, port)).await.map_err(|err| { - format!( - "{native_err}\n\nFallback with MongoDB (Legacy) driver failed: {}", - mongo_legacy_error_with_auth_hint(&err) + mongo_legacy_fallback_error( + &native_err, + "Fallback with MongoDB (Legacy) driver failed", + &mongo_legacy_error_with_auth_hint(&err), ) })?; client.disconnect().await.ok(); @@ -1380,12 +1435,14 @@ pub async fn connect_db( if should_retry_mongo_with_legacy_driver(&native_err) { log::info!("Native MongoDB driver failed ({native_err}), falling back to agent driver"); let mut client = - state.agent_manager.spawn(&db_config.db_type, Some(MONGO_LEGACY_DRIVER_PROFILE)).await?; + spawn_mongo_legacy_fallback_agent(state.inner().as_ref(), &db_config.db_type, &native_err) + .await?; state.ensure_current_connection_attempt(&id, Some(attempt)).await?; client.connect(mongo_legacy_connect_params(&db_config, &host, port)).await.map_err(|err| { - format!( - "{native_err}\n\nFallback with MongoDB (Legacy) driver failed: {}", - mongo_legacy_error_with_auth_hint(&err) + mongo_legacy_fallback_error( + &native_err, + "Fallback with MongoDB (Legacy) driver failed", + &mongo_legacy_error_with_auth_hint(&err), ) })?; state.ensure_current_connection_attempt(&id, Some(attempt)).await?;