diff --git a/apps/desktop/src/components/connection/ConnectionDialog.vue b/apps/desktop/src/components/connection/ConnectionDialog.vue index 4f782a4de..4a35fa93c 100644 --- a/apps/desktop/src/components/connection/ConnectionDialog.vue +++ b/apps/desktop/src/components/connection/ConnectionDialog.vue @@ -87,6 +87,7 @@ const defaultForm = (): Omit => ({ ssh_key_passphrase: "", ssh_expose_lan: false, ssh_connect_timeout_secs: 5, + connect_timeout_secs: 5, proxy_enabled: false, proxy_type: "socks5", proxy_host: "", @@ -345,6 +346,7 @@ watch( ssh_key_passphrase: config.ssh_key_passphrase || "", ssh_expose_lan: config.ssh_expose_lan || false, ssh_connect_timeout_secs: config.ssh_connect_timeout_secs || 5, + connect_timeout_secs: config.connect_timeout_secs || 5, proxy_enabled: config.proxy_enabled || false, proxy_type: config.proxy_type || "socks5", proxy_host: config.proxy_host || "", @@ -737,6 +739,8 @@ function connectionConfigForSubmit(id: string): ConnectionConfig { } const sshTimeout = Number(config.ssh_connect_timeout_secs); config.ssh_connect_timeout_secs = Number.isFinite(sshTimeout) && sshTimeout > 0 ? sshTimeout : 5; + const connectTimeout = Number(config.connect_timeout_secs); + config.connect_timeout_secs = Number.isFinite(connectTimeout) && connectTimeout > 0 ? connectTimeout : 5; const proxyPort = Number(config.proxy_port); config.proxy_port = Number.isFinite(proxyPort) && proxyPort > 0 ? proxyPort : 1080; if (!config.one_time) config.one_time = undefined; @@ -2314,6 +2318,17 @@ function openExternalUrl(url: string) { :disabled="!form.ssh_enabled" /> +
+ + +
diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index d187eb587..2ec479ef7 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -218,6 +218,7 @@ export default { sshKeyPathBrowse: "Browse", sshExposeLan: "Expose tunnel to LAN", sshConnectTimeout: "SSH Timeout (seconds)", + connectTimeout: "Connection Timeout (seconds)", proxy: "Proxy", proxyEnable: "Connect database through proxy", proxyType: "Proxy Type", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index d875e7c8d..6b3ed52d0 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -210,6 +210,7 @@ export default { sshKeyPathBrowse: "Examinar", sshExposeLan: "Exponer túnel a la red local", sshConnectTimeout: "Tiempo de espera SSH (segundos)", + connectTimeout: "Tiempo de espera de conexión (segundos)", proxy: "Proxy", proxyEnable: "Conectar la base de datos mediante proxy", proxyType: "Tipo de proxy", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index 19458aee1..f0e23bbbf 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -214,6 +214,7 @@ export default { sshKeyPathBrowse: "浏览", sshExposeLan: "允许局域网访问隧道", sshConnectTimeout: "SSH 超时时间(秒)", + connectTimeout: "连接超时(秒)", proxy: "代理", proxyEnable: "通过代理连接数据库", proxyType: "代理类型", diff --git a/apps/desktop/src/lib/dbeaverImport.ts b/apps/desktop/src/lib/dbeaverImport.ts index cad330230..5ccb792e9 100644 --- a/apps/desktop/src/lib/dbeaverImport.ts +++ b/apps/desktop/src/lib/dbeaverImport.ts @@ -230,6 +230,7 @@ function buildConnection( ssh_key_passphrase: "", ssh_expose_lan: false, ssh_connect_timeout_secs: 5, + connect_timeout_secs: 5, ssl: false, oracle_connection_type: profile.dbType === "oracle" ? parsedUrl.oracleConnectionType || "service_name" : undefined, connection_string: profile.dbType === "jdbc" || profile.dbType === "mongodb" ? url || undefined : undefined, diff --git a/apps/desktop/src/lib/navicatImport.ts b/apps/desktop/src/lib/navicatImport.ts index 8461cc7f7..5569583f8 100644 --- a/apps/desktop/src/lib/navicatImport.ts +++ b/apps/desktop/src/lib/navicatImport.ts @@ -164,6 +164,7 @@ async function parseConnection(node: ParsedNode): Promise { ? config.attached_databases.filter((database) => database.name?.trim() && database.path?.trim()) : [], ssh_connect_timeout_secs: config.ssh_connect_timeout_secs || 5, + connect_timeout_secs: config.connect_timeout_secs || 5, proxy_type: config.proxy_type || "socks5", proxy_port: config.proxy_port || 1080, }; diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index 0c9223f8f..73c98c947 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -74,6 +74,7 @@ export interface ConnectionConfig { ssh_key_passphrase?: string; ssh_expose_lan?: boolean; ssh_connect_timeout_secs?: number; + connect_timeout_secs?: number; proxy_enabled?: boolean; proxy_type?: "socks5" | "http"; proxy_host?: string; diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 395185cc4..7ba0b6db6 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -170,20 +170,21 @@ impl AppState { let (host, port) = self.connection_host_port(connection_id, &db_config).await?; probe_connection_endpoint(&db_config, &host, port).await?; let url = connection_url_for_endpoint(&db_config, &host, port); + let connect_timeout = std::time::Duration::from_secs(db_config.effective_connect_timeout_secs()); let pool = match db_config.db_type { DatabaseType::Mysql if db_config.needs_bare_mysql() => { - PoolKind::Mysql(db::mysql::connect_bare(&url).await?, MysqlMode::Bare) + PoolKind::Mysql(db::mysql::connect_bare(&url, connect_timeout).await?, MysqlMode::Bare) } DatabaseType::Mysql => { - let pool = db::mysql::connect_with_ca_cert(&url, Some(&db_config.ca_cert_path)).await?; + let pool = db::mysql::connect_with_ca_cert(&url, Some(&db_config.ca_cert_path), connect_timeout).await?; let mode = detect_ob_oracle_mode(&db_config, &pool).await; PoolKind::Mysql(pool, mode) } DatabaseType::Doris | DatabaseType::StarRocks => { - PoolKind::Mysql(db::mysql::connect_bare(&url).await?, MysqlMode::Bare) + PoolKind::Mysql(db::mysql::connect_bare(&url, connect_timeout).await?, MysqlMode::Bare) } DatabaseType::Postgres | DatabaseType::Redshift | DatabaseType::Gaussdb | DatabaseType::OpenGauss => { - PoolKind::Postgres(db::postgres::connect(&url).await?) + PoolKind::Postgres(db::postgres::connect(&url, connect_timeout).await?) } DatabaseType::Sqlite => PoolKind::Sqlite(db::sqlite::connect_path(&expand_tilde(&db_config.host)).await?), DatabaseType::Redis => { @@ -195,7 +196,7 @@ impl AppState { )) } else { db::redis_driver::RedisConnection::Direct(tokio::sync::Mutex::new( - db::redis_driver::connect(&url).await?, + db::redis_driver::connect(&url, connect_timeout).await?, )) }; PoolKind::Redis(con) @@ -211,8 +212,8 @@ impl AppState { PoolKind::DuckDb(con) } DatabaseType::MongoDb => { - let native_err = match db::mongo_driver::connect(&url).await { - Ok(client) => match db::mongo_driver::test_connection(&client).await { + let native_err = match db::mongo_driver::connect(&url, connect_timeout).await { + Ok(client) => match db::mongo_driver::test_connection(&client, connect_timeout).await { Ok(()) => { self.connections.write().await.insert(pool_key.clone(), PoolKind::MongoDb(client)); return Ok(pool_key); @@ -239,8 +240,9 @@ impl AppState { username, password, Some(&db_config.ca_cert_path), + connect_timeout, )?; - db::clickhouse_driver::test_connection(&client).await?; + db::clickhouse_driver::test_connection(&client, connect_timeout).await?; PoolKind::ClickHouse(client) } DatabaseType::SqlServer => { @@ -250,6 +252,7 @@ impl AppState { &db_config.username, &db_config.password, db_config.database.as_deref(), + connect_timeout, ) .await?; PoolKind::SqlServer(Arc::new(tokio::sync::Mutex::new(client))) @@ -261,8 +264,9 @@ impl AppState { Some(&db_config.username), Some(&db_config.password), accept_invalid_certs, + connect_timeout, ); - db::elasticsearch_driver::test_connection(&client).await?; + db::elasticsearch_driver::test_connection(&client, connect_timeout).await?; PoolKind::Elasticsearch(client) } DatabaseType::Dameng @@ -820,7 +824,8 @@ pub async fn probe_connection_endpoint(config: &ConnectionConfig, host: &str, po if !uses_tcp_probe(config, host, port) { return Ok(()); } - db::probe_tcp_endpoint(&format!("{:?}", config.db_type), host, port).await + let timeout = std::time::Duration::from_secs(config.effective_connect_timeout_secs()); + db::probe_tcp_endpoint(&format!("{:?}", config.db_type), host, port, timeout).await } fn uses_tcp_probe(config: &ConnectionConfig, host: &str, port: u16) -> bool { diff --git a/crates/dbx-core/src/db/clickhouse_driver.rs b/crates/dbx-core/src/db/clickhouse_driver.rs index 8c557838a..d6ead12b4 100644 --- a/crates/dbx-core/src/db/clickhouse_driver.rs +++ b/crates/dbx-core/src/db/clickhouse_driver.rs @@ -1,9 +1,9 @@ use reqwest::{Certificate, Client as HttpClient}; use serde::Deserialize; use std::fs; -use std::time::Instant; +use std::time::{Duration, Instant}; -use super::{connection_timeout, with_connection_timeout}; +use super::with_connection_timeout; use crate::query::MAX_ROWS; use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, QueryResult, TableInfo}; @@ -16,9 +16,8 @@ pub struct ChClient { } impl ChClient { - pub fn new(url: &str, username: Option, password: Option) -> Self { - let http = - HttpClient::builder().connect_timeout(connection_timeout()).build().unwrap_or_else(|_| HttpClient::new()); + pub fn new(url: &str, username: Option, password: Option, timeout: Duration) -> Self { + let http = HttpClient::builder().connect_timeout(timeout).build().unwrap_or_else(|_| HttpClient::new()); Self { http, base_url: url.trim_end_matches('/').to_string(), username, password } } @@ -27,8 +26,9 @@ impl ChClient { username: Option, password: Option, ca_cert_path: Option<&str>, + timeout: Duration, ) -> Result { - let mut builder = HttpClient::builder().connect_timeout(connection_timeout()); + let mut builder = HttpClient::builder().connect_timeout(timeout); if let Some(path) = ca_cert_path.map(str::trim).filter(|path| !path.is_empty()) { let path = expand_cert_path(path); let cert_bytes = @@ -157,10 +157,10 @@ fn limited_query_result(result: ChJsonResult, execution_time_ms: u128, max_rows: QueryResult { columns, rows, affected_rows: 0, execution_time_ms, truncated, session_id: None, has_more: false } } -pub async fn test_connection(client: &ChClient) -> Result<(), String> { +pub async fn test_connection(client: &ChClient, timeout: Duration) -> Result<(), String> { let url = format!("{}/?query=SELECT%201", client.base_url); let req = build_request(client, client.http.get(&url)); - let resp = with_connection_timeout("ClickHouse", connection_timeout(), async { + let resp = with_connection_timeout("ClickHouse", timeout, async { req.send().await.map_err(|e| format!("ClickHouse connection failed: {e}")) }) .await?; diff --git a/crates/dbx-core/src/db/elasticsearch_driver.rs b/crates/dbx-core/src/db/elasticsearch_driver.rs index 448bbc338..8cc03084b 100644 --- a/crates/dbx-core/src/db/elasticsearch_driver.rs +++ b/crates/dbx-core/src/db/elasticsearch_driver.rs @@ -1,7 +1,8 @@ use reqwest::Client as HttpClient; use serde::Deserialize; +use std::time::Duration; -use super::{connection_timeout, with_connection_timeout}; +use super::with_connection_timeout; use crate::db::mongo_driver::MongoDocumentResult; pub struct EsClient { @@ -11,13 +12,19 @@ pub struct EsClient { } impl EsClient { - pub fn new(url: &str, username: Option<&str>, password: Option<&str>, accept_invalid_certs: bool) -> Self { + pub fn new( + url: &str, + username: Option<&str>, + password: Option<&str>, + accept_invalid_certs: bool, + timeout: Duration, + ) -> Self { let auth = match (username, password) { (Some(u), Some(p)) if !u.is_empty() => Some((u.to_string(), p.to_string())), _ => None, }; let http = HttpClient::builder() - .connect_timeout(connection_timeout()) + .connect_timeout(timeout) .danger_accept_invalid_certs(accept_invalid_certs) .build() .unwrap_or_else(|_| HttpClient::new()); @@ -59,8 +66,8 @@ impl Clone for EsClient { } } -pub async fn test_connection(client: &EsClient) -> Result<(), String> { - let resp = with_connection_timeout("Elasticsearch", connection_timeout(), async { +pub async fn test_connection(client: &EsClient, timeout: Duration) -> Result<(), String> { + let resp = with_connection_timeout("Elasticsearch", timeout, async { client.get("/").send().await.map_err(|e| format!("Elasticsearch connection failed: {e}")) }) .await?; diff --git a/crates/dbx-core/src/db/mod.rs b/crates/dbx-core/src/db/mod.rs index f1001f82e..4d58a8f12 100644 --- a/crates/dbx-core/src/db/mod.rs +++ b/crates/dbx-core/src/db/mod.rs @@ -64,8 +64,12 @@ pub fn tcp_probe_timeout() -> Duration { } pub fn parse_connect_timeout(url: &str) -> Duration { + parse_connect_timeout_with_fallback(url, connection_timeout()) +} + +pub fn parse_connect_timeout_with_fallback(url: &str, fallback: Duration) -> Duration { let Some(query) = url.split('?').nth(1) else { - return connection_timeout(); + return fallback; }; for param in query.split('&') { let trimmed = param.trim(); @@ -76,7 +80,11 @@ pub fn parse_connect_timeout(url: &str) -> Duration { Some(pair) => pair, None => continue, }; - if key.eq_ignore_ascii_case("connect_timeout") || key.eq_ignore_ascii_case("connectTimeout") { + if key.eq_ignore_ascii_case("connect_timeout") + || key.eq_ignore_ascii_case("connectTimeout") + || key.eq_ignore_ascii_case("connection_timeout") + || key.eq_ignore_ascii_case("connectionTimeout") + { if let Ok(v) = value.parse::() { if v >= 1 && v <= 300 { return Duration::from_secs(v); @@ -84,7 +92,7 @@ pub fn parse_connect_timeout(url: &str) -> Duration { } } } - connection_timeout() + fallback } pub async fn with_connection_timeout(label: &str, timeout: Duration, future: F) -> Result @@ -96,10 +104,10 @@ where .map_err(|_| format!("{label} connection timed out ({}s)", timeout.as_secs()))? } -pub async fn probe_tcp_endpoint(label: &str, host: &str, port: u16) -> Result<(), String> { - tokio::time::timeout(tcp_probe_timeout(), tokio::net::TcpStream::connect((host, port))) +pub async fn probe_tcp_endpoint(label: &str, host: &str, port: u16, timeout: Duration) -> Result<(), String> { + tokio::time::timeout(timeout, tokio::net::TcpStream::connect((host, port))) .await - .map_err(|_| format!("{label} TCP connection timed out ({TCP_PROBE_TIMEOUT_SECS}s)"))? + .map_err(|_| format!("{label} TCP connection timed out ({}s)", timeout.as_secs()))? .map(|_| ()) .map_err(|e| format!("{label} TCP connection failed: {e}")) } diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index 19cb7f0a6..c8200d531 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -4,7 +4,8 @@ use mongodb::{ }; use serde::{Deserialize, Serialize}; -use super::{connection_timeout, with_connection_timeout, CONNECTION_TIMEOUT_SECS}; +use super::with_connection_timeout; +use std::time::Duration; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MongoDocumentResult { @@ -12,17 +13,17 @@ pub struct MongoDocumentResult { pub total: u64, } -pub async fn connect(url: &str) -> Result { - with_connection_timeout("MongoDB", connection_timeout(), async { +pub async fn connect(url: &str, timeout: Duration) -> Result { + with_connection_timeout("MongoDB", timeout, async { Client::with_uri_str(url).await.map_err(|e| format!("MongoDB connection failed: {e}")) }) .await } -pub async fn test_connection(client: &Client) -> Result<(), String> { - tokio::time::timeout(connection_timeout(), client.list_database_names()) +pub async fn test_connection(client: &Client, timeout: Duration) -> Result<(), String> { + tokio::time::timeout(timeout, client.list_database_names()) .await - .map_err(|_| format!("MongoDB connection timed out ({CONNECTION_TIMEOUT_SECS}s)"))? + .map_err(|_| format!("MongoDB connection timed out ({}s)", timeout.as_secs()))? .map(|_| ()) .map_err(|e| format!("MongoDB connection failed: {e}")) } diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 149e2ac2c..5c90ddc7d 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -258,12 +258,12 @@ fn mysql_value_to_json(row: &mysql_async::Row, idx: usize) -> serde_json::Value .unwrap_or(serde_json::Value::Null) } -pub async fn connect(url: &str) -> Result { - connect_with_ca_cert(url, None).await +pub async fn connect(url: &str, fallback_timeout: Duration) -> Result { + connect_with_ca_cert(url, None, fallback_timeout).await } -pub async fn connect_with_ca_cert(url: &str, ca_cert_path: Option<&str>) -> Result { - let timeout = super::parse_connect_timeout(url); +pub async fn connect_with_ca_cert(url: &str, ca_cert_path: Option<&str>, fallback_timeout: Duration) -> Result { + let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); let pool = create_pool(url, ca_cert_path)?; let result = verify_pool_connection(&pool, timeout).await; @@ -606,8 +606,8 @@ fn mysql_async_url(url: &str) -> Cow<'_, str> { } } -pub async fn connect_bare(url: &str) -> Result { - let timeout = super::parse_connect_timeout(url); +pub async fn connect_bare(url: &str, fallback_timeout: Duration) -> Result { + let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); let pool = create_pool(url, None)?; verify_pool_connection(&pool, timeout).await.map(|_| pool) } diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 8d1f248e5..218456fee 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -12,7 +12,7 @@ use std::fs::File; use std::io::BufReader; use std::str::FromStr; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use tokio_postgres::config::SslMode; use tokio_postgres::types::{FromSql, Type}; use tokio_postgres::{Row, SimpleQueryMessage}; @@ -344,13 +344,14 @@ async fn execute_select_query( } } -pub async fn connect(url: &str) -> Result { +pub async fn connect(url: &str, fallback_timeout: Duration) -> Result { let postgres_url = postgres_connection_url(url)?; let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); let tz = iana_time_zone::get_timezone().unwrap_or_else(|_| "UTC".to_string()); - super::with_connection_timeout("PostgreSQL", super::connection_timeout(), async { + super::with_connection_timeout("PostgreSQL", timeout, async { let pg_config = tokio_postgres::Config::from_str(&postgres_url.url) .map_err(|e| format!("Invalid PostgreSQL connection URL: {e}"))?; @@ -369,7 +370,7 @@ pub async fn connect(url: &str) -> Result { let pool = Pool::builder(mgr) .max_size(1) .runtime(Runtime::Tokio1) - .wait_timeout(Some(super::connection_timeout())) + .wait_timeout(Some(timeout)) .build() .map_err(|e| format!("Failed to create PostgreSQL pool: {e}"))?; diff --git a/crates/dbx-core/src/db/redis_driver.rs b/crates/dbx-core/src/db/redis_driver.rs index 7bbc6f5ad..2b3581404 100644 --- a/crates/dbx-core/src/db/redis_driver.rs +++ b/crates/dbx-core/src/db/redis_driver.rs @@ -87,9 +87,19 @@ pub struct RedisNodeEndpoint { pub port: u16, } -pub async fn connect(url: &str) -> Result { +pub async fn connect(url: &str, timeout: std::time::Duration) -> Result { let client = redis::Client::open(url).map_err(|e| format!("Redis connection failed: {e}"))?; - connect_client(client).await + let mut con = tokio::time::timeout(timeout, client.get_multiplexed_async_connection()) + .await + .map_err(|_| format!("Redis connection timed out ({}s)", timeout.as_secs()))? + .map_err(|e| format!("Redis connection failed: {e}"))?; + + tokio::time::timeout(timeout, redis::cmd("PING").query_async::(&mut con)) + .await + .map_err(|_| format!("Redis ping timed out ({}s)", timeout.as_secs()))? + .map_err(|e| format!("Redis authentication failed or command rejected: {e}"))?; + + Ok(con) } pub async fn connect_sentinel(config: &ConnectionConfig) -> Result { diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index d1f1dda5c..651fd7262 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -1,11 +1,9 @@ use futures::TryStreamExt; use rust_decimal::Decimal; -use std::time::Instant; +use std::time::{Duration, Instant}; use tiberius::{AuthMethod, Client, ColumnData, Config, FromSql, QueryItem, QueryStream, SqlBrowser}; use tokio::net::TcpStream; use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt}; - -use super::{connection_timeout, CONNECTION_TIMEOUT_SECS}; use crate::query::MAX_ROWS; use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; @@ -39,10 +37,11 @@ pub async fn connect( user: &str, pass: &str, database: Option<&str>, + timeout: Duration, ) -> Result { - match try_connect(host, port, user, pass, database, true).await { + match try_connect(host, port, user, pass, database, true, timeout).await { Ok(client) => Ok(client), - Err(_) => try_connect(host, port, user, pass, database, false).await, + Err(_) => try_connect(host, port, user, pass, database, false, timeout).await, } } @@ -53,6 +52,7 @@ async fn try_connect( pass: &str, database: Option<&str>, use_encryption: bool, + timeout: Duration, ) -> Result { let mut config = Config::new(); let endpoint = sqlserver_endpoint(host); @@ -72,19 +72,19 @@ async fn try_connect( } let tcp = if endpoint.instance_name.is_some() { - tokio::time::timeout(connection_timeout(), TcpStream::connect_named(&config)) + tokio::time::timeout(timeout, TcpStream::connect_named(&config)) .await - .map_err(|_| format!("SQL Server connection timed out ({CONNECTION_TIMEOUT_SECS}s)"))? + .map_err(|_| format!("SQL Server connection timed out ({}s)", timeout.as_secs()))? .map_err(|e| format!("SQL Server connection failed: {e}"))? } else { - tokio::time::timeout(connection_timeout(), TcpStream::connect(config.get_addr())) + tokio::time::timeout(timeout, TcpStream::connect(config.get_addr())) .await - .map_err(|_| format!("SQL Server connection timed out ({CONNECTION_TIMEOUT_SECS}s)"))? + .map_err(|_| format!("SQL Server connection timed out ({}s)", timeout.as_secs()))? .map_err(|e| format!("SQL Server connection failed: {e}"))? }; - tokio::time::timeout(connection_timeout(), Client::connect(config, tcp.compat_write())) + tokio::time::timeout(timeout, Client::connect(config, tcp.compat_write())) .await - .map_err(|_| format!("SQL Server handshake timed out ({CONNECTION_TIMEOUT_SECS}s)"))? + .map_err(|_| format!("SQL Server handshake timed out ({}s)", timeout.as_secs()))? .map_err(|e| format!("SQL Server connection failed: {e}")) } diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs index 63fc453e9..d163e985b 100644 --- a/crates/dbx-core/src/models/connection.rs +++ b/crates/dbx-core/src/models/connection.rs @@ -41,6 +41,8 @@ pub struct ConnectionConfig { pub ssh_expose_lan: bool, #[serde(default = "default_ssh_connect_timeout_secs")] pub ssh_connect_timeout_secs: u64, + #[serde(default = "default_connect_timeout_secs")] + pub connect_timeout_secs: u64, #[serde(default)] pub proxy_enabled: bool, #[serde(default)] @@ -102,6 +104,10 @@ pub fn default_ssh_connect_timeout_secs() -> u64 { 5 } +pub fn default_connect_timeout_secs() -> u64 { + 5 +} + fn default_proxy_port() -> u16 { 1080 } @@ -194,6 +200,14 @@ impl ConnectionConfig { } } + pub fn effective_connect_timeout_secs(&self) -> u64 { + if self.connect_timeout_secs == 0 { + default_connect_timeout_secs() + } else { + self.connect_timeout_secs.clamp(1, 300) + } + } + pub fn effective_database(&self) -> Option<&str> { self.database .as_deref() diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index b840c6e40..9fc8f58c4 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -219,33 +219,40 @@ pub async fn test_connection(state: State<'_, Arc>, config: Connection let probe_result = probe_connection_endpoint(&config, &host, port).await; let url = connection_url_for_endpoint(&config, &host, port); let target = redacted_connection_url_for_endpoint(&config, &host, port); + let connect_timeout = std::time::Duration::from_secs(config.effective_connect_timeout_secs()); log::info!("[test_connection] db_type={:?} target={}", config.db_type, target); let result = match probe_result { Err(e) => Err(e), Ok(()) => match config.db_type { - DatabaseType::Mysql if config.needs_bare_mysql() => match db::mysql::connect_bare(&url).await { - Ok(pool) => { - let _ = pool.disconnect().await; - Ok("Connection successful".to_string()) + DatabaseType::Mysql if config.needs_bare_mysql() => { + match db::mysql::connect_bare(&url, connect_timeout).await { + Ok(pool) => { + let _ = pool.disconnect().await; + Ok("Connection successful".to_string()) + } + Err(e) => Err(e), } - Err(e) => Err(e), - }, - DatabaseType::Mysql => match db::mysql::connect_with_ca_cert(&url, Some(&config.ca_cert_path)).await { - Ok(pool) => { - let _ = pool.disconnect().await; - Ok("Connection successful".to_string()) + } + DatabaseType::Mysql => { + match db::mysql::connect_with_ca_cert(&url, Some(&config.ca_cert_path), connect_timeout).await { + Ok(pool) => { + let _ = pool.disconnect().await; + Ok("Connection successful".to_string()) + } + Err(e) => Err(e), } - Err(e) => Err(e), - }, - DatabaseType::Doris | DatabaseType::StarRocks => match db::mysql::connect_bare(&url).await { - Ok(pool) => { - let _ = pool.disconnect().await; - Ok("Connection successful".to_string()) + } + DatabaseType::Doris | DatabaseType::StarRocks => { + match db::mysql::connect_bare(&url, connect_timeout).await { + Ok(pool) => { + let _ = pool.disconnect().await; + Ok("Connection successful".to_string()) + } + Err(e) => Err(e), } - Err(e) => Err(e), - }, + } DatabaseType::Postgres | DatabaseType::Redshift | DatabaseType::Gaussdb | DatabaseType::OpenGauss => { - match db::postgres::connect(&url).await { + match db::postgres::connect(&url, connect_timeout).await { Ok(pool) => { pool.close(); Ok("Connection successful".to_string()) @@ -264,7 +271,7 @@ pub async fn test_connection(state: State<'_, Arc>, config: Connection } else if config.uses_redis_sentinel() { db::redis_driver::connect_sentinel(&config).await? } else { - db::redis_driver::connect(&url).await? + db::redis_driver::connect(&url, connect_timeout).await? }; drop(con); Ok("Connection successful".to_string()) @@ -279,8 +286,8 @@ pub async fn test_connection(state: State<'_, Arc>, config: Connection } } DatabaseType::MongoDb => { - let native_err = match db::mongo_driver::connect(&url).await { - Ok(client) => match db::mongo_driver::test_connection(&client).await { + let native_err = match db::mongo_driver::connect(&url, connect_timeout).await { + Ok(client) => match db::mongo_driver::test_connection(&client, connect_timeout).await { Ok(()) => return Ok("Connection successful".to_string()), Err(e) => e, }, @@ -307,22 +314,35 @@ pub async fn test_connection(state: State<'_, Arc>, config: Connection username, password, Some(&config.ca_cert_path), + connect_timeout, )?; - db::clickhouse_driver::test_connection(&client).await.map(|_| "Connection successful".to_string()) - } - DatabaseType::SqlServer => { - db::sqlserver::connect(&host, port, &config.username, &config.password, config.database.as_deref()) + db::clickhouse_driver::test_connection(&client, connect_timeout) .await .map(|_| "Connection successful".to_string()) } + DatabaseType::SqlServer => { + db::sqlserver::connect( + &host, + port, + &config.username, + &config.password, + config.database.as_deref(), + connect_timeout, + ) + .await + .map(|_| "Connection successful".to_string()) + } DatabaseType::Elasticsearch => { let client = db::elasticsearch_driver::EsClient::new( &url, Some(&config.username), Some(&config.password), config.ssl, + connect_timeout, ); - db::elasticsearch_driver::test_connection(&client).await.map(|_| "Connection successful".to_string()) + db::elasticsearch_driver::test_connection(&client, connect_timeout) + .await + .map(|_| "Connection successful".to_string()) } db_type if database_capabilities::is_agent_type(&db_type) => { test_agent_connection(state.inner(), &config, &host, port).await @@ -362,20 +382,21 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi let (host, port) = state.connection_host_port(&id, &db_config).await?; probe_connection_endpoint(&db_config, &host, port).await?; let url = connection_url_for_endpoint(&db_config, &host, port); + let connect_timeout = std::time::Duration::from_secs(db_config.effective_connect_timeout_secs()); let pool = match db_config.db_type { DatabaseType::Mysql if db_config.needs_bare_mysql() => { - PoolKind::Mysql(db::mysql::connect_bare(&url).await?, MysqlMode::Bare) + PoolKind::Mysql(db::mysql::connect_bare(&url, connect_timeout).await?, MysqlMode::Bare) } DatabaseType::Mysql => PoolKind::Mysql( - db::mysql::connect_with_ca_cert(&url, Some(&db_config.ca_cert_path)).await?, + db::mysql::connect_with_ca_cert(&url, Some(&db_config.ca_cert_path), connect_timeout).await?, MysqlMode::Normal, ), DatabaseType::Doris | DatabaseType::StarRocks => { - PoolKind::Mysql(db::mysql::connect_bare(&url).await?, MysqlMode::Bare) + PoolKind::Mysql(db::mysql::connect_bare(&url, connect_timeout).await?, MysqlMode::Bare) } DatabaseType::Postgres | DatabaseType::Redshift | DatabaseType::Gaussdb | DatabaseType::OpenGauss => { - PoolKind::Postgres(db::postgres::connect(&url).await?) + PoolKind::Postgres(db::postgres::connect(&url, connect_timeout).await?) } DatabaseType::Sqlite => PoolKind::Sqlite(db::sqlite::connect_path(&expand_tilde(&db_config.host)).await?), DatabaseType::Redis => { @@ -389,7 +410,7 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi ))) } else { PoolKind::Redis(db::redis_driver::RedisConnection::Direct(tokio::sync::Mutex::new( - db::redis_driver::connect(&url).await?, + db::redis_driver::connect(&url, connect_timeout).await?, ))) }; con @@ -405,8 +426,8 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi PoolKind::DuckDb(con) } DatabaseType::MongoDb => { - let native_err = match db::mongo_driver::connect(&url).await { - Ok(client) => match db::mongo_driver::test_connection(&client).await { + let native_err = match db::mongo_driver::connect(&url, connect_timeout).await { + Ok(client) => match db::mongo_driver::test_connection(&client, connect_timeout).await { Ok(()) => { state.configs.write().await.insert(id.clone(), config); state.connections.write().await.insert(id.clone(), PoolKind::MongoDb(client)); @@ -435,8 +456,9 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi username, password, Some(&db_config.ca_cert_path), + connect_timeout, )?; - db::clickhouse_driver::test_connection(&client).await?; + db::clickhouse_driver::test_connection(&client, connect_timeout).await?; PoolKind::ClickHouse(client) } DatabaseType::SqlServer => { @@ -446,6 +468,7 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi &db_config.username, &db_config.password, db_config.database.as_deref(), + connect_timeout, ) .await?; PoolKind::SqlServer(std::sync::Arc::new(tokio::sync::Mutex::new(client))) @@ -456,8 +479,9 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi Some(&db_config.username), Some(&db_config.password), db_config.ssl, + connect_timeout, ); - db::elasticsearch_driver::test_connection(&client).await?; + db::elasticsearch_driver::test_connection(&client, connect_timeout).await?; PoolKind::Elasticsearch(client) } db_type if database_capabilities::is_agent_type(&db_type) => {