From 9c478d143766b9b2e1578d85f4f0577a10342148 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 15 Jun 2026 22:43:05 +0800 Subject: [PATCH] fix(connection): respect agent connect timeout --- .../connectionAttemptTimeout.spec.ts | 16 ++++++ .../src/lib/__tests__/defaultDatabase.spec.ts | 19 +++++++ .../src/lib/connectionAttemptTimeout.ts | 6 +- apps/desktop/src/lib/defaultDatabase.ts | 8 ++- crates/dbx-core/src/agent_manager.rs | 31 +++++++++++ crates/dbx-core/src/agent_runtime.rs | 45 +++++++++++++++ crates/dbx-core/src/connection.rs | 55 +++++++++++++++++-- src-tauri/src/commands/connection.rs | 31 ++++++++--- 8 files changed, 193 insertions(+), 18 deletions(-) create mode 100644 apps/desktop/src/lib/__tests__/connectionAttemptTimeout.spec.ts create mode 100644 apps/desktop/src/lib/__tests__/defaultDatabase.spec.ts diff --git a/apps/desktop/src/lib/__tests__/connectionAttemptTimeout.spec.ts b/apps/desktop/src/lib/__tests__/connectionAttemptTimeout.spec.ts new file mode 100644 index 0000000000..ebce369cfd --- /dev/null +++ b/apps/desktop/src/lib/__tests__/connectionAttemptTimeout.spec.ts @@ -0,0 +1,16 @@ +import { describe, expect, it } from "vitest"; +import { ACCESS_AGENT_MIN_CONNECT_TIMEOUT_SECS, AGENT_DRIVER_MIN_CONNECT_TIMEOUT_SECS, connectionAttemptTimeoutMs, CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS } from "@/lib/connectionAttemptTimeout"; + +describe("connectionAttemptTimeout", () => { + it("uses a 30s agent startup floor", () => { + expect(connectionAttemptTimeoutMs({ db_type: "oracle", connect_timeout_secs: 5, transport_layers: [] })).toBe(AGENT_DRIVER_MIN_CONNECT_TIMEOUT_SECS * 1000 + CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS); + }); + + it("uses a 30s Access agent startup floor", () => { + expect(connectionAttemptTimeoutMs({ db_type: "access", connect_timeout_secs: 5, transport_layers: [] })).toBe(ACCESS_AGENT_MIN_CONNECT_TIMEOUT_SECS * 1000 + CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS); + }); + + it("respects a user configured Access timeout above the default floor", () => { + expect(connectionAttemptTimeoutMs({ db_type: "access", connect_timeout_secs: 45, transport_layers: [] })).toBe(47_000); + }); +}); diff --git a/apps/desktop/src/lib/__tests__/defaultDatabase.spec.ts b/apps/desktop/src/lib/__tests__/defaultDatabase.spec.ts new file mode 100644 index 0000000000..c0c1784cb3 --- /dev/null +++ b/apps/desktop/src/lib/__tests__/defaultDatabase.spec.ts @@ -0,0 +1,19 @@ +import { describe, expect, it } from "vitest"; +import { decodeSelectableDatabaseValue, EMPTY_DATABASE_SELECT_VALUE, encodeSelectableDatabaseValue, TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE } from "@/lib/defaultDatabase"; + +describe("defaultDatabase selectable values", () => { + it("encodes empty tree-schema databases with the default database sentinel", () => { + expect(encodeSelectableDatabaseValue("postgres", "")).toBe(TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE); + expect(decodeSelectableDatabaseValue("postgres", TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE)).toBe(""); + }); + + it("encodes empty non-tree-schema databases with a non-empty sentinel", () => { + expect(encodeSelectableDatabaseValue("access", "")).toBe(EMPTY_DATABASE_SELECT_VALUE); + expect(decodeSelectableDatabaseValue("access", EMPTY_DATABASE_SELECT_VALUE)).toBe(""); + }); + + it("preserves non-empty database names", () => { + expect(encodeSelectableDatabaseValue("access", "653128SXB.mdb")).toBe("653128SXB.mdb"); + expect(decodeSelectableDatabaseValue("access", "653128SXB.mdb")).toBe("653128SXB.mdb"); + }); +}); diff --git a/apps/desktop/src/lib/connectionAttemptTimeout.ts b/apps/desktop/src/lib/connectionAttemptTimeout.ts index c809996105..4e8030df10 100644 --- a/apps/desktop/src/lib/connectionAttemptTimeout.ts +++ b/apps/desktop/src/lib/connectionAttemptTimeout.ts @@ -2,7 +2,8 @@ import type { ConnectionConfig, DatabaseType } from "@/types/database"; export const CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS = 2_000; export const MONGO_LEGACY_FALLBACK_TIMEOUT_BUFFER_MS = 30_000; -export const AGENT_DRIVER_MIN_CONNECT_TIMEOUT_SECS = 15; +export const AGENT_DRIVER_MIN_CONNECT_TIMEOUT_SECS = 30; +export const ACCESS_AGENT_MIN_CONNECT_TIMEOUT_SECS = 30; const DEFAULT_CONNECT_TIMEOUT_SECS = 5; const AGENT_DRIVER_TYPES = new Set([ @@ -46,7 +47,8 @@ function positiveSeconds(value: unknown, fallback: number): number { export function connectionAttemptTimeoutMs(config: Pick & Partial>): number { const baseTimeoutSecs = positiveSeconds(config.connect_timeout_secs, DEFAULT_CONNECT_TIMEOUT_SECS); - const timeouts = [AGENT_DRIVER_TYPES.has(config.db_type as DatabaseType) ? Math.max(baseTimeoutSecs, AGENT_DRIVER_MIN_CONNECT_TIMEOUT_SECS) : baseTimeoutSecs]; + const agentMinTimeoutSecs = config.db_type === "access" ? ACCESS_AGENT_MIN_CONNECT_TIMEOUT_SECS : AGENT_DRIVER_MIN_CONNECT_TIMEOUT_SECS; + const timeouts = [AGENT_DRIVER_TYPES.has(config.db_type as DatabaseType) ? Math.max(baseTimeoutSecs, agentMinTimeoutSecs) : baseTimeoutSecs]; for (const layer of config.transport_layers ?? []) { if (layer.type === "ssh") { timeouts.push(positiveSeconds(layer.connect_timeout_secs, DEFAULT_CONNECT_TIMEOUT_SECS)); diff --git a/apps/desktop/src/lib/defaultDatabase.ts b/apps/desktop/src/lib/defaultDatabase.ts index ef081898e3..3ffed632af 100644 --- a/apps/desktop/src/lib/defaultDatabase.ts +++ b/apps/desktop/src/lib/defaultDatabase.ts @@ -2,6 +2,7 @@ import type { ConnectionConfig, DatabaseType } from "@/types/database"; import { usesTreeSchemaMode } from "@/lib/databaseCapabilities"; export const TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE = "__dbx_tree_schema_default_database__"; +export const EMPTY_DATABASE_SELECT_VALUE = "__dbx_empty_database__"; export function resolveDefaultDatabase(connection: Pick, options: string[]): string { return connection.database || options[0] || ""; @@ -12,11 +13,14 @@ export function isTreeSchemaDefaultDatabase(dbType: DatabaseType | undefined, da } export function encodeSelectableDatabaseValue(dbType: DatabaseType | undefined, database: string): string { - return isTreeSchemaDefaultDatabase(dbType, database) ? TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE : database; + if (isTreeSchemaDefaultDatabase(dbType, database)) return TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE; + return database === "" ? EMPTY_DATABASE_SELECT_VALUE : database; } export function decodeSelectableDatabaseValue(dbType: DatabaseType | undefined, value: string): string { - return value === TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE && usesTreeSchemaMode(dbType) ? "" : value; + if (value === TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE && usesTreeSchemaMode(dbType)) return ""; + if (value === EMPTY_DATABASE_SELECT_VALUE) return ""; + return value; } export function formatDatabaseLabel(connection: Pick | undefined, database: string, labels: { defaultDatabase: string; noDatabase: string }): string { diff --git a/crates/dbx-core/src/agent_manager.rs b/crates/dbx-core/src/agent_manager.rs index 96f35c2ef2..554548c2b5 100644 --- a/crates/dbx-core/src/agent_manager.rs +++ b/crates/dbx-core/src/agent_manager.rs @@ -751,6 +751,18 @@ impl AgentManager { crate::agent_runtime::call_daemon(self, db_type, driver_profile, method, params).await } + pub async fn call_daemon_with_timeout( + &self, + db_type: &DatabaseType, + driver_profile: Option<&str>, + method: &str, + params: serde_json::Value, + timeout_duration: Option, + ) -> Result { + crate::agent_runtime::call_daemon_with_timeout(self, db_type, driver_profile, method, params, timeout_duration) + .await + } + pub async fn call_daemon_method( &self, db_type: &DatabaseType, @@ -761,6 +773,25 @@ impl AgentManager { crate::agent_runtime::call_daemon_method(self, db_type, driver_profile, method, params).await } + pub async fn call_daemon_method_with_timeout( + &self, + db_type: &DatabaseType, + driver_profile: Option<&str>, + method: AgentMethod, + params: serde_json::Value, + timeout_duration: Option, + ) -> Result { + crate::agent_runtime::call_daemon_method_with_timeout( + self, + db_type, + driver_profile, + method, + params, + timeout_duration, + ) + .await + } + pub async fn download_file(url: &str, dest: &Path) -> Result<(), String> { if let Some(parent) = dest.parent() { std::fs::create_dir_all(parent).map_err(|e| e.to_string())?; diff --git a/crates/dbx-core/src/agent_runtime.rs b/crates/dbx-core/src/agent_runtime.rs index 26fa011e9d..3a6d59ef32 100644 --- a/crates/dbx-core/src/agent_runtime.rs +++ b/crates/dbx-core/src/agent_runtime.rs @@ -1,4 +1,5 @@ use serde::de::DeserializeOwned; +use std::time::Duration; use crate::agent_manager::{AgentManager, DEFAULT_JRE_KEY}; use crate::database_capabilities; @@ -70,6 +71,39 @@ pub async fn call_daemon( } } +pub async fn call_daemon_with_timeout( + manager: &AgentManager, + db_type: &DatabaseType, + driver_profile: Option<&str>, + method: &str, + params: serde_json::Value, + timeout_duration: Option, +) -> Result { + let keys = runtime_agent_key_candidates(db_type, driver_profile) + .ok_or_else(|| format!("{:?} is not an agent-driven database type", db_type))?; + let key = first_installed_agent_key(manager, &keys).unwrap_or(keys[0]).to_string(); + + let mut daemons = manager.daemons.lock().await; + + if !daemons.contains_key(&key) { + let client = spawn_client_for_key(manager, &key).await?; + daemons.insert(key.clone(), client); + } + + let client = daemons.get_mut(&key).unwrap(); + match client.call_with_timeout::(method, params.clone(), timeout_duration).await { + Ok(result) => Ok(result), + Err(err) => { + log::warn!("[agent] daemon call failed, respawning: {err}"); + daemons.remove(&key); + let mut new_client = spawn_client_for_key(manager, &key).await?; + let result = new_client.call_with_timeout::(method, params, timeout_duration).await?; + daemons.insert(key, new_client); + Ok(result) + } + } +} + pub async fn call_daemon_method( manager: &AgentManager, db_type: &DatabaseType, @@ -80,6 +114,17 @@ pub async fn call_daemon_method( call_daemon(manager, db_type, driver_profile, method.as_str(), params).await } +pub async fn call_daemon_method_with_timeout( + manager: &AgentManager, + db_type: &DatabaseType, + driver_profile: Option<&str>, + method: AgentMethod, + params: serde_json::Value, + timeout_duration: Option, +) -> Result { + call_daemon_with_timeout(manager, db_type, driver_profile, method.as_str(), params, timeout_duration).await +} + fn runtime_agent_key_candidates(db_type: &DatabaseType, driver_profile: Option<&str>) -> Option> { let primary = db_type_to_agent_key(db_type, driver_profile)?; if *db_type == DatabaseType::Oracle { diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 6eeb8421e7..6500e55931 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -28,6 +28,8 @@ use crate::storage::Storage; pub const JDBC_PLUGIN_NOT_INSTALLED: &str = "JDBC plugin is not installed. Install the optional JDBC plugin to use this connection."; +const DEFAULT_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30; +const ACCESS_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30; #[cfg(feature = "duckdb-bundled")] mod duckdb_types { @@ -570,8 +572,13 @@ impl AppState { agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or("")); let mut client = self.agent_manager.spawn(&db_config.db_type, db_config.driver_profile.as_deref()).await?; - let connect_result = - client.call_method::(AgentMethod::Connect, connect_params.clone()).await; + let connect_result = client + .call_method_with_timeout::( + AgentMethod::Connect, + connect_params.clone(), + Some(agent_connect_timeout(&db_config)), + ) + .await; if let Err(err) = connect_result { let alternate_configs = oracle_alternate_connect_configs(&db_config, &err); if !alternate_configs.is_empty() { @@ -589,7 +596,7 @@ impl AppState { .next() .unwrap_or_else(|| "alternate".to_string()); match client - .call_method::( + .call_method_with_timeout::( AgentMethod::Connect, agent_connect_params( &alternate_config, @@ -597,6 +604,7 @@ impl AppState { port, alternate_config.effective_database().unwrap_or(""), ), + Some(agent_connect_timeout(&alternate_config)), ) .await { @@ -627,7 +635,11 @@ impl AppState { match self.agent_manager.spawn(&db_config.db_type, Some(profile)).await { Ok(mut fallback_client) => { match fallback_client - .call_method::(AgentMethod::Connect, connect_params.clone()) + .call_method_with_timeout::( + AgentMethod::Connect, + connect_params.clone(), + Some(agent_connect_timeout(&db_config)), + ) .await { Ok(_) => { @@ -1170,8 +1182,17 @@ pub fn redacted_connection_url_for_endpoint(config: &ConnectionConfig, host: &st } } +pub fn agent_connect_timeout(config: &ConnectionConfig) -> std::time::Duration { + let min_timeout = if config.db_type == DatabaseType::Access { + ACCESS_AGENT_CONNECT_TIMEOUT_SECS + } else { + DEFAULT_AGENT_CONNECT_TIMEOUT_SECS + }; + std::time::Duration::from_secs(config.effective_connect_timeout_secs().max(min_timeout)) +} + fn external_driver_connect_timeout(config: &ConnectionConfig) -> std::time::Duration { - std::time::Duration::from_secs(config.effective_connect_timeout_secs().max(30)) + agent_connect_timeout(config) } fn native_postgres_url_config(config: &ConnectionConfig) -> Option { @@ -1313,7 +1334,7 @@ async fn detect_ob_oracle_mode(config: &ConnectionConfig, pool: &db::mysql::MySq #[cfg(test)] mod tests { use super::{ - connection_url_for_endpoint, database_connection_config, metadata_connection_config, + agent_connect_timeout, connection_url_for_endpoint, database_connection_config, metadata_connection_config, mysql_metadata_fallback_url, redacted_connection_url_for_endpoint, uses_bare_mysql_pool, uses_tcp_probe, validate_h2_database_path, AppState, PoolKind, }; @@ -1377,6 +1398,28 @@ mod tests { } } + #[test] + fn access_agent_connect_timeout_has_longer_default_floor() { + let mut config = mysql_config(None); + config.db_type = DatabaseType::Access; + config.connect_timeout_secs = 5; + assert_eq!(agent_connect_timeout(&config).as_secs(), 30); + + config.connect_timeout_secs = 45; + assert_eq!(agent_connect_timeout(&config).as_secs(), 45); + } + + #[test] + fn non_access_agent_connect_timeout_uses_standard_floor() { + let mut config = mysql_config(None); + config.db_type = DatabaseType::Oracle; + config.connect_timeout_secs = 5; + assert_eq!(agent_connect_timeout(&config).as_secs(), 30); + + config.connect_timeout_secs = 45; + assert_eq!(agent_connect_timeout(&config).as_secs(), 45); + } + #[test] fn agent_connect_params_include_url_params() { let mut config = mysql_config(Some("testdb")); diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 196ca0eb93..fd33e16461 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -7,8 +7,9 @@ pub use dbx_core::agent_connection::{ should_retry_oracle_with_10g_driver, }; pub use dbx_core::connection::{ - connect_bare_metadata_pool, connect_mysql_metadata_pool, connection_url_for_endpoint, metadata_connection_config, - probe_connection_endpoint, redacted_connection_url_for_endpoint, AppState, MysqlMode, PoolKind, + agent_connect_timeout, connect_bare_metadata_pool, connect_mysql_metadata_pool, connection_url_for_endpoint, + metadata_connection_config, probe_connection_endpoint, redacted_connection_url_for_endpoint, AppState, MysqlMode, + PoolKind, }; use dbx_core::database_capabilities; use dbx_core::db; @@ -31,11 +32,12 @@ async fn test_agent_connection( let connect_params = agent_connect_params(config, host, port, config.database.as_deref().unwrap_or("")); let result = state .agent_manager - .call_daemon_method::( + .call_daemon_method_with_timeout::( &config.db_type, config.driver_profile.as_deref(), AgentMethod::TestConnection, connect_params.clone(), + Some(agent_connect_timeout(config)), ) .await; @@ -43,7 +45,7 @@ async fn test_agent_connection( if let Some(alternate_config) = oracle_alternate_connect_config(config, &err) { state .agent_manager - .call_daemon_method::( + .call_daemon_method_with_timeout::( &alternate_config.db_type, alternate_config.driver_profile.as_deref(), AgentMethod::TestConnection, @@ -53,6 +55,7 @@ async fn test_agent_connection( port, alternate_config.database.as_deref().unwrap_or(""), ), + Some(agent_connect_timeout(&alternate_config)), ) .await .map_err(|alternate_err| { @@ -64,11 +67,12 @@ async fn test_agent_connection( for profile in oracle_auth_fallback_profiles(config, &err) { match state .agent_manager - .call_daemon_method::( + .call_daemon_method_with_timeout::( &config.db_type, Some(profile), AgentMethod::TestConnection, connect_params.clone(), + Some(agent_connect_timeout(config)), ) .await { @@ -101,12 +105,18 @@ async fn connect_agent_pool( ) -> Result { let connect_params = agent_connect_params(config, host, port, config.effective_database().unwrap_or("")); let mut client = state.agent_manager.spawn(&config.db_type, config.driver_profile.as_deref()).await?; - let connect_result = client.call_method::(AgentMethod::Connect, connect_params.clone()).await; + let connect_result = client + .call_method_with_timeout::( + AgentMethod::Connect, + connect_params.clone(), + Some(agent_connect_timeout(config)), + ) + .await; if let Err(err) = connect_result { if let Some(alternate_config) = oracle_alternate_connect_config(config, &err) { client - .call_method::( + .call_method_with_timeout::( AgentMethod::Connect, agent_connect_params( &alternate_config, @@ -114,6 +124,7 @@ async fn connect_agent_pool( port, alternate_config.effective_database().unwrap_or(""), ), + Some(agent_connect_timeout(&alternate_config)), ) .await .map_err(|alternate_err| { @@ -126,7 +137,11 @@ async fn connect_agent_pool( match state.agent_manager.spawn(&config.db_type, Some(profile)).await { Ok(mut fallback_client) => { match fallback_client - .call_method::(AgentMethod::Connect, connect_params.clone()) + .call_method_with_timeout::( + AgentMethod::Connect, + connect_params.clone(), + Some(agent_connect_timeout(config)), + ) .await { Ok(_) => {