From 9f6c35ee78191cfeeb19e36d81b4d3879435960c Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 25 May 2026 16:23:25 +0800 Subject: [PATCH] fix: make MySQL connection timeout configurable via URL param connect_timeout Previously the connection timeout was hardcoded to 5s via tokio::time::timeout, ignoring any connect_timeout in the URL. Also mysql_async rejected camelCase connectTimeout as unknown. Now parse_connect_timeout extracts connect_timeout/connectTimeout from the URL (case-insensitive, 1-300s range), removes them before passing to mysql_async, and feeds the value to with_connection_timeout. --- crates/dbx-core/src/db/mod.rs | 30 ++++++++++++++++++-- crates/dbx-core/src/db/mysql.rs | 49 +++++++++++++++++++++++++++++---- 2 files changed, 71 insertions(+), 8 deletions(-) diff --git a/crates/dbx-core/src/db/mod.rs b/crates/dbx-core/src/db/mod.rs index 46b46cb6b6..cf1236aa30 100644 --- a/crates/dbx-core/src/db/mod.rs +++ b/crates/dbx-core/src/db/mod.rs @@ -49,13 +49,37 @@ pub fn tcp_probe_timeout() -> Duration { Duration::from_secs(TCP_PROBE_TIMEOUT_SECS) } -pub async fn with_connection_timeout(label: &str, future: F) -> Result +pub fn parse_connect_timeout(url: &str) -> Duration { + let Some(query) = url.split('?').nth(1) else { + return connection_timeout(); + }; + for param in query.split('&') { + let trimmed = param.trim(); + if trimmed.is_empty() { + continue; + } + let (key, value) = match trimmed.split_once('=') { + Some(pair) => pair, + None => continue, + }; + if key.eq_ignore_ascii_case("connect_timeout") || key.eq_ignore_ascii_case("connectTimeout") { + if let Ok(v) = value.parse::() { + if v >= 1 && v <= 300 { + return Duration::from_secs(v); + } + } + } + } + connection_timeout() +} + +pub async fn with_connection_timeout(label: &str, timeout: Duration, future: F) -> Result where F: Future>, { - tokio::time::timeout(connection_timeout(), future) + tokio::time::timeout(timeout, future) .await - .map_err(|_| format!("{label} connection timed out ({CONNECTION_TIMEOUT_SECS}s)"))? + .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> { diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 307ebfcba9..6460b6541a 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -164,15 +164,16 @@ fn mysql_value_to_json(row: &mysql_async::Row, idx: usize) -> serde_json::Value } pub async fn connect(url: &str) -> Result { + let timeout = super::parse_connect_timeout(url); let pool = create_pool(url)?; - let result = verify_pool_connection(&pool).await; + let result = verify_pool_connection(&pool, timeout).await; if let Err(ref e) = result { if mysql_error_should_retry_without_ssl(e) { if let Some(fallback_url) = ssl_fallback_url(url) { log::info!("SSL handshake failed, retrying with ssl-mode=disabled"); let fallback_pool = create_pool(&fallback_url)?; - return match verify_pool_connection(&fallback_pool).await { + return match verify_pool_connection(&fallback_pool, timeout).await { Ok(()) => Ok(fallback_pool), Err(e) => Err(e), }; @@ -193,8 +194,8 @@ fn create_pool(url: &str) -> Result { Ok(MySqlPool::new(builder)) } -async fn verify_pool_connection(pool: &MySqlPool) -> Result<(), String> { - super::with_connection_timeout("MySQL", async { +async fn verify_pool_connection(pool: &MySqlPool, timeout: Duration) -> Result<(), String> { + super::with_connection_timeout("MySQL", timeout, async { let mut conn = pool.get_conn().await.map_err(|e| format!("MySQL connection failed: {e}"))?; conn.ping().await.map_err(|e| format!("MySQL ping failed: {e}"))?; Ok(()) @@ -242,6 +243,8 @@ fn mysql_async_url(url: &str) -> Cow<'_, str> { && !segment.starts_with("charset=") && !segment.starts_with("time_zone=") && !segment.starts_with("time-zone=") + && !segment.to_ascii_lowercase().starts_with("connect_timeout=") + && !segment.to_ascii_lowercase().starts_with("connecttimeout=") }) .collect(); @@ -255,8 +258,9 @@ fn mysql_async_url(url: &str) -> Cow<'_, str> { } pub async fn connect_bare(url: &str) -> Result { + let timeout = super::parse_connect_timeout(url); let pool = create_pool(url)?; - verify_pool_connection(&pool).await.map(|_| pool) + verify_pool_connection(&pool, timeout).await.map(|_| pool) } pub async fn list_databases(pool: &MySqlPool) -> Result, String> { @@ -789,4 +793,39 @@ mod tests { let _ = conn.disconnect().await; let _ = pool.disconnect().await; } + + #[test] + fn parse_connect_timeout_extracts_underscore_form() { + let url = "mysql://host:3306/db?connect_timeout=30"; + assert_eq!(super::parse_connect_timeout(url), Duration::from_secs(30)); + } + + #[test] + fn parse_connect_timeout_extracts_camelcase_form() { + let url = "mysql://host:3306/db?connectTimeout=60"; + assert_eq!(super::parse_connect_timeout(url), Duration::from_secs(60)); + } + + #[test] + fn parse_connect_timeout_ignores_out_of_range() { + let default = super::connection_timeout(); + let url = "mysql://host:3306/db?connect_timeout=999"; + assert_eq!(super::parse_connect_timeout(url), default); + let url2 = "mysql://host:3306/db?connect_timeout=0"; + assert_eq!(super::parse_connect_timeout(url2), default); + } + + #[test] + fn parse_connect_timeout_returns_default_when_missing() { + let default = super::connection_timeout(); + let url = "mysql://host:3306/db?ssl-mode=preferred&charset=utf8mb4"; + assert_eq!(super::parse_connect_timeout(url), default); + } + + #[test] + fn parse_connect_timeout_returns_default_when_no_query() { + let default = super::connection_timeout(); + let url = "mysql://host:3306/db"; + assert_eq!(super::parse_connect_timeout(url), default); + } }