From 66e6ca121121843d84af9d734c6bed4c0846b6dc Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 17 Jun 2026 11:14:09 +0800 Subject: [PATCH] fix(db): refresh stale native pools --- crates/dbx-core/src/agent_kv.rs | 6 ++++++ crates/dbx-core/src/connection.rs | 32 ++++++++++++++++++++++++++++++ crates/dbx-core/src/mongo_ops.rs | 14 +++++++++++++ crates/dbx-core/src/storage.rs | 4 +++- crates/dbx-web/src/routes/query.rs | 3 +++ src-tauri/src/commands/query.rs | 3 +++ 6 files changed, 61 insertions(+), 1 deletion(-) diff --git a/crates/dbx-core/src/agent_kv.rs b/crates/dbx-core/src/agent_kv.rs index 8f2a6a2cc5..a83f4e37c6 100644 --- a/crates/dbx-core/src/agent_kv.rs +++ b/crates/dbx-core/src/agent_kv.rs @@ -145,12 +145,18 @@ pub async fn kv_delete_core(state: &AppState, connection_id: &str, key: &str) -> call_agent_kv(state, connection_id, AgentKvMethod::Delete, kv_delete_params(key)).await } +async fn ensure_agent_kv_pool(state: &AppState, connection_id: &str) -> Result<(), String> { + state.get_or_create_pool(connection_id, None).await.map(|_| ()) +} + async fn call_agent_kv( state: &AppState, connection_id: &str, method: AgentKvMethod, params: serde_json::Value, ) -> Result { + ensure_agent_kv_pool(state, connection_id).await?; + let connections = state.connections.read().await; let pool = connections.get(connection_id).ok_or("Connection not found")?; match pool { diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 3e1789618e..fc1e6152b5 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -823,6 +823,17 @@ impl AppState { return false; }; match pool { + PoolKind::Mysql(pool, _) => { + let pool = pool.clone(); + drop(connections); + match db::mysql::get_conn_with_health_check(&pool).await { + Ok(_) => false, + Err(err) => { + log::warn!("MySQL connection pool '{pool_key}' is stale: {err}"); + true + } + } + } PoolKind::SqlServer(client) => { let client = client.clone(); drop(connections); @@ -842,6 +853,27 @@ impl AppState { true } }, + PoolKind::MongoDb(client) => { + let client = client.clone(); + drop(connections); + let (connect_timeout, database) = { + let configs = self.configs.read().await; + let config = config_for_pool_key(pool_key, &configs); + ( + config + .map(|config| Duration::from_secs(config.effective_connect_timeout_secs().max(1))) + .unwrap_or_else(|| Duration::from_secs(1)), + config.and_then(|config| config.effective_database().map(str::to_string)), + ) + }; + match db::mongo_driver::test_connection(&client, connect_timeout, database.as_deref()).await { + Ok(()) => false, + Err(err) => { + log::warn!("MongoDB connection pool '{pool_key}' is stale: {err}"); + true + } + } + } _ => false, } }; diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 38ecc33f4d..4306c3bac3 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -12,7 +12,12 @@ fn sort_names(mut names: Vec) -> Vec { names } +async fn ensure_document_pool(state: &AppState, connection_id: &str) -> Result<(), String> { + state.get_or_create_pool(connection_id, None).await.map(|_| ()) +} + pub async fn mongo_list_databases_core(state: &AppState, connection_id: &str) -> Result, String> { + ensure_document_pool(state, connection_id).await?; let fallback_database = configured_mongo_database(state, connection_id).await; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { @@ -59,6 +64,7 @@ pub async fn mongo_list_collections_core( connection_id: &str, database: &str, ) -> Result, String> { + ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => mongo_driver::list_collections(client, database).await.map(sort_names), @@ -82,6 +88,7 @@ pub async fn document_find_documents_core( filter: Option<&str>, sort: Option<&str>, ) -> Result { + ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => { @@ -131,6 +138,7 @@ pub async fn mongo_aggregate_documents_core( pipeline_json: &str, max_rows: Option, ) -> Result { + ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => { @@ -148,6 +156,7 @@ pub async fn mongo_insert_document_core( collection: &str, doc_json: &str, ) -> Result { + ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => mongo_driver::insert_document(client, database, collection, doc_json).await, @@ -178,6 +187,7 @@ pub async fn mongo_insert_documents_core( collection: &str, docs_json: &str, ) -> Result { + ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => mongo_driver::insert_documents(client, database, collection, docs_json).await, @@ -194,6 +204,7 @@ pub async fn mongo_update_document_core( id: &str, doc_json: &str, ) -> Result { + ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => mongo_driver::update_document(client, database, collection, id, doc_json).await, @@ -227,6 +238,7 @@ pub async fn mongo_update_documents_core( update_json: &str, many: bool, ) -> Result { + ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => { @@ -244,6 +256,7 @@ pub async fn mongo_delete_document_core( collection: &str, id: &str, ) -> Result { + ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => mongo_driver::delete_document(client, database, collection, id).await, @@ -270,6 +283,7 @@ pub async fn mongo_delete_documents_core( filter_json: &str, many: bool, ) -> Result { + ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => { diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index 82382d4ae0..9e001d2a8f 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -2084,6 +2084,7 @@ mod tests { driver_store_dir: Some("/tmp/dbx-drivers".to_string()), plugin_store_dir: Some("/tmp/dbx-plugins".to_string()), agent_store_dir: Some("/tmp/dbx-agents".to_string()), + sidebar_table_page_size: DesktopSettings::default().sidebar_table_page_size, }) .await .unwrap(); @@ -2098,7 +2099,8 @@ mod tests { saved_sql_sync_dir: None, driver_store_dir: Some("/tmp/dbx-drivers".to_string()), plugin_store_dir: Some("/tmp/dbx-plugins".to_string()), - agent_store_dir: Some("/tmp/dbx-agents".to_string()) + agent_store_dir: Some("/tmp/dbx-agents".to_string()), + sidebar_table_page_size: DesktopSettings::default().sidebar_table_page_size, } ); } diff --git a/crates/dbx-web/src/routes/query.rs b/crates/dbx-web/src/routes/query.rs index 54daaea807..0053da3f60 100644 --- a/crates/dbx-web/src/routes/query.rs +++ b/crates/dbx-web/src/routes/query.rs @@ -482,6 +482,9 @@ pub async fn get_explain_info( State(state): State>, Json(req): Json, ) -> Result, AppError> { + let database_for_pool = req.database.as_deref().filter(|database| !database.trim().is_empty()); + state.app.get_or_create_pool(&req.connection_id, database_for_pool).await.map_err(AppError)?; + let client = { let connections = state.app.connections.read().await; let pool = connections.get(&req.connection_id).ok_or_else(|| AppError("Connection not found".to_string()))?; diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index c69e78d271..0a2c4b7caa 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -462,6 +462,9 @@ pub async fn get_explain_info( sql: String, mode: Option, ) -> Result { + let database_for_pool = database.as_deref().filter(|database| !database.trim().is_empty()); + state.get_or_create_pool(&connection_id, database_for_pool).await?; + let client = { let connections = state.connections.read().await; let pool = connections.get(&connection_id).ok_or_else(|| "Connection not found".to_string())?;