fix: release duckdb session locks on pool close

This commit is contained in:
t8y2 2026-05-29 00:16:37 +08:00
parent f8ade78e97
commit a1e00fcf1a
2 changed files with 69 additions and 25 deletions

View File

@ -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;

View File

@ -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);