From a1e00fcf1a0544ef7237c2bcd266d4178d0bdd14 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Fri, 29 May 2026 00:16:37 +0800
Subject: [PATCH] fix: release duckdb session locks on pool close

---
 crates/dbx-core/src/connection.rs    | 72 ++++++++++++++++++++++++++--
 src-tauri/src/commands/connection.rs | 22 +--------
 2 files changed, 69 insertions(+), 25 deletions(-)

diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs
index d08552057..a788af5a8 100644
--- a/crates/dbx-core/src/connection.rs
+++ b/crates/dbx-core/src/connection.rs
@@ -471,7 +471,10 @@ impl AppState {
             self.remove_connection_pools(connection_id).await;
             self.reset_connection_transport(connection_id).await;
         } else {
-            self.connections.write().await.remove(&pool_key);
+            let removed = self.connections.write().await.remove(&pool_key);
+            if let Some(pool) = removed {
+                close_pool_kind(pool).await;
+            }
         }
         self.get_or_create_pool_for_session(connection_id, database, client_session_id).await
     }
@@ -492,7 +495,13 @@ impl AppState {
         };
         let base_pool_key = base_pool_key_for(db_type, connection_id, database, false);
         let pool_key = session_scoped_pool_key(base_pool_key, Some(&session));
-        Ok(self.connections.write().await.remove(&pool_key).is_some())
+        let removed = self.connections.write().await.remove(&pool_key);
+        if let Some(pool) = removed {
+            close_pool_kind(pool).await;
+            Ok(true)
+        } else {
+            Ok(false)
+        }
     }
 
     pub async fn duckdb_existing_pool_is_usable_for_config(&self, config: &ConnectionConfig) -> Result<bool, String> {
@@ -539,11 +548,16 @@ impl AppState {
             .filter(|k| *k == connection_id || k.starts_with(&format!("{connection_id}:")))
             .cloned()
             .collect();
+        let mut removed = Vec::with_capacity(keys_to_remove.len());
         for key in keys_to_remove {
-            if let Some(PoolKind::DuckDb(con)) = conns.remove(&key) {
-                crate::db::duckdb_driver::close_connection(con);
+            if let Some(pool) = conns.remove(&key) {
+                removed.push(pool);
             }
         }
+        drop(conns);
+        for pool in removed {
+            close_pool_kind(pool).await;
+        }
     }
 
     async fn uses_forwarded_transport(&self, connection_id: &str) -> bool {
@@ -565,6 +579,30 @@ fn session_scoped_pool_key(base_pool_key: String, client_session_id: Option<&str
         .unwrap_or(base_pool_key)
 }
 
+pub async fn close_pool_kind(pool: PoolKind) {
+    match pool {
+        PoolKind::Mysql(p, _) => {
+            let _ = p.disconnect().await;
+        }
+        PoolKind::Postgres(p) => p.close(),
+        PoolKind::Sqlite(_) => {}
+        PoolKind::Redis(_) => {}
+        PoolKind::DuckDb(con) => {
+            crate::db::duckdb_driver::close_connection(con);
+        }
+        PoolKind::MongoDb(_) => {}
+        PoolKind::ClickHouse(_) => {}
+        PoolKind::SqlServer(_) => {}
+        PoolKind::Elasticsearch(_) => {}
+        PoolKind::Agent(client) => {
+            let mut client = client.lock().await;
+            let _ = client.disconnect().await;
+        }
+        PoolKind::ExternalTabular(_) => {}
+        PoolKind::ExternalDriver { .. } => {}
+    }
+}
+
 fn base_pool_key_for(
     db_type: Option<DatabaseType>,
     connection_id: &str,
@@ -1442,6 +1480,32 @@ mod tests {
         let _ = std::fs::remove_dir_all(dir);
     }
 
+    #[tokio::test]
+    async fn close_client_session_pool_removes_session_scoped_duckdb_pool() {
+        let (state, dir) = test_app_state().await;
+        let db_path = dir.join("session.duckdb");
+        let mut config = mysql_config(None);
+        config.id = "duckdb-conn".to_string();
+        config.name = "DuckDB".to_string();
+        config.db_type = DatabaseType::DuckDb;
+        config.host = db_path.to_string_lossy().to_string();
+        config.port = 0;
+
+        state.configs.write().await.insert(config.id.clone(), config.clone());
+        let pool_key = state
+            .get_or_create_pool_for_session("duckdb-conn", Some("main"), Some("tab-1"))
+            .await
+            .unwrap();
+        assert_eq!(pool_key, "duckdb-conn:session:tab-1");
+
+        assert!(state.close_client_session_pool("duckdb-conn", Some("main"), "tab-1").await.unwrap());
+
+        let conns = state.connections.read().await;
+        assert!(!conns.contains_key("duckdb-conn:session:tab-1"));
+
+        let _ = std::fs::remove_dir_all(dir);
+    }
+
     #[tokio::test]
     async fn proxy_connection_uses_local_forward_endpoint() {
         let (state, dir) = test_app_state().await;
diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs
index 3f5a21ea7..d44104254 100644
--- a/src-tauri/src/commands/connection.rs
+++ b/src-tauri/src/commands/connection.rs
@@ -504,27 +504,7 @@ pub async fn disconnect_db(state: State<'_, Arc<AppState>>, connection_id: Strin
         conns.keys().filter(|k| *k == &connection_id || k.starts_with(&format!("{connection_id}:"))).cloned().collect();
     for key in keys_to_remove {
         if let Some(pool) = conns.remove(&key) {
-            match pool {
-                PoolKind::Mysql(p, _) => {
-                    let _ = p.disconnect().await;
-                }
-                PoolKind::Postgres(p) => p.close(),
-                PoolKind::Sqlite(_) => {}
-                PoolKind::Redis(_) => {}
-                PoolKind::DuckDb(con) => {
-                    dbx_core::db::duckdb_driver::close_connection(con);
-                }
-                PoolKind::MongoDb(_) => {}
-                PoolKind::ClickHouse(_) => {}
-                PoolKind::SqlServer(_) => {}
-                PoolKind::Elasticsearch(_) => {}
-                PoolKind::Agent(client) => {
-                    let mut client = client.lock().await;
-                    let _ = client.disconnect().await;
-                }
-                PoolKind::ExternalTabular(_) => {}
-                PoolKind::ExternalDriver { .. } => {}
-            }
+            dbx_core::connection::close_pool_kind(pool).await;
         }
     }
     drop(conns);