diff --git a/crates/dbx-core/src/agent_explain.rs b/crates/dbx-core/src/agent_explain.rs index 3806871c78..955eaceef7 100644 --- a/crates/dbx-core/src/agent_explain.rs +++ b/crates/dbx-core/src/agent_explain.rs @@ -38,8 +38,8 @@ pub async fn get_agent_explain_info_core( } let target = { - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or_else(|| "Connection not found".to_string())?; + let pool_handle = state.pool_handle(&pool_key).await; + let pool = pool_handle.as_ref().ok_or_else(|| "Connection not found".to_string())?; match pool { PoolKind::Agent(client) => ExplainTarget::Agent(client.clone()), PoolKind::ExternalDriver { config, session, .. } => { @@ -227,14 +227,18 @@ mod tests { let storage = crate::storage::Storage::open(&dir.join("storage.db")).await.unwrap(); let state = AppState::new(storage); state.configs.write().await.insert(config.id.clone(), config.clone()); - state.connections.write().await.insert( - "oracle-jdbc".to_string(), - PoolKind::ExternalDriver { - driver_id: "jdbc".to_string(), - config: Arc::new(config), - session: session.clone(), - }, - ); + state + .update_connection_pools(|connections| { + connections.insert( + "oracle-jdbc".to_string(), + PoolKind::ExternalDriver { + driver_id: "jdbc".to_string(), + config: Arc::new(config), + session: session.clone(), + }, + ); + }) + .await; let plan = get_agent_explain_info_core( &state, diff --git a/crates/dbx-core/src/agent_kv.rs b/crates/dbx-core/src/agent_kv.rs index 954d7fa1c1..8251e075f4 100644 --- a/crates/dbx-core/src/agent_kv.rs +++ b/crates/dbx-core/src/agent_kv.rs @@ -1013,8 +1013,8 @@ async fn ensure_etcd_dangerous_action_capability( .get(connection_id) .and_then(|config| crate::agent_catalog::agent_key(&config.db_type, config.driver_profile.as_deref())) }; - let connections = state.connections.read().await; - let PoolKind::Agent(client) = connections.get(connection_id).ok_or("Connection not found")? else { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + let PoolKind::Agent(client) = &pool else { return Err("Not an agent key-value connection".to_string()); }; if !client.lock().await.supports_capability(capability) { @@ -1706,8 +1706,8 @@ fn kv_put_required_capabilities(options: &KvPutOptions) -> Vec pub async fn kv_supports_ttl_core(state: &AppState, connection_id: &str) -> 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")?; + let pool_handle = state.pool_handle(connection_id).await; + let pool = pool_handle.as_ref().ok_or("Connection not found")?; match pool { PoolKind::Agent(client) => Ok(client.lock().await.supports_capability(AgentCapability::KvTtl)), _ => Err("Not an agent key-value connection".to_string()), @@ -1738,8 +1738,8 @@ async fn call_agent_kv( .and_then(|config| crate::agent_catalog::agent_key(&config.db_type, config.driver_profile.as_deref())) }; let client = { - let connections = state.connections.read().await; - match connections.get(connection_id) { + let pool_handle = state.pool_handle(connection_id).await; + match pool_handle.as_ref() { Some(PoolKind::Agent(client)) => client.clone(), Some(_) => return Err("Not an agent key-value connection".to_string()), None => return Err("Connection not found".to_string()), diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index d254cef57d..72c5cd7625 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -1872,7 +1872,11 @@ for line in sys.stdin: let state = Arc::new(AppState::new(storage)); let connection = agent_test_connection("dameng-1", "Dameng", DatabaseType::Dameng, "APPDB"); state.configs.write().await.insert(connection.id.clone(), connection); - state.connections.write().await.insert("dameng-1:APPDB".to_string(), PoolKind::agent(client)); + state + .update_connection_pools(|connections| { + connections.insert("dameng-1:APPDB".to_string(), PoolKind::agent(client)); + }) + .await; let read = ToolCall { id: "read".to_string(), @@ -1936,7 +1940,11 @@ for line in sys.stdin: let state = Arc::new(AppState::new(storage)); let connection = agent_test_connection("mysql-1", "MySQL", DatabaseType::Mysql, "rs_main"); state.configs.write().await.insert(connection.id.clone(), connection); - state.connections.write().await.insert("mysql-1:rs_main".to_string(), PoolKind::agent(client)); + state + .update_connection_pools(|connections| { + connections.insert("mysql-1:rs_main".to_string(), PoolKind::agent(client)); + }) + .await; let sql = "SHOW TRIGGERS FROM `rs_main` LIKE 'trg_order_items_after_%';"; let call = ToolCall { diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 4ff2c71fb7..ff45092f67 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -82,6 +82,7 @@ fn mysql_pool_setup_queries(config: &ConnectionConfig, url: &str) -> Vec queries } +#[derive(Clone)] pub enum PoolKind { Mysql(db::mysql::MySqlPool, MysqlMode), Postgres(deadpool_postgres::Pool), @@ -89,7 +90,7 @@ pub enum PoolKind { Rqlite(db::rqlite_driver::RqliteClient), Turso(db::turso_driver::TursoClient), CloudflareD1(db::cloudflare_d1_driver::CloudflareD1Client), - Redis(db::redis_driver::RedisConnection), + Redis(Arc), DuckDbWorker(DuckDbWorkerHandle), MongoDb(mongodb::Client), DynamoDb(db::dynamodb_driver::DynamoDbClient), @@ -121,46 +122,6 @@ pub enum PoolKind { } impl PoolKind { - /// Clone only the handles used by metadata operations so the global pool - /// map lock can be released before any database I/O begins. - pub(crate) fn clone_for_metadata(&self) -> Option { - match self { - Self::Mysql(pool, mode) => Some(Self::Mysql(pool.clone(), *mode)), - Self::Postgres(pool) => Some(Self::Postgres(pool.clone())), - Self::Sqlite(pool) => Some(Self::Sqlite(pool.clone())), - Self::Rqlite(client) => Some(Self::Rqlite(client.clone())), - Self::Turso(client) => Some(Self::Turso(client.clone())), - Self::CloudflareD1(client) => Some(Self::CloudflareD1(client.clone())), - #[cfg(feature = "duckdb-sidecar")] - Self::DuckDbWorker(client) => Some(Self::DuckDbWorker(client.clone())), - #[cfg(not(feature = "duckdb-sidecar"))] - Self::DuckDbWorker(_) => Some(Self::DuckDbWorker(())), - Self::MongoDb(client) => Some(Self::MongoDb(client.clone())), - Self::ClickHouse(client) => Some(Self::ClickHouse(client.clone())), - Self::SqlServer(client) => Some(Self::SqlServer(client.clone())), - Self::Elasticsearch(client) => Some(Self::Elasticsearch(client.clone())), - Self::Easysearch(client) => Some(Self::Easysearch(client.clone())), - Self::Meilisearch(client) => Some(Self::Meilisearch(client.clone())), - Self::HBase(client) => Some(Self::HBase(client.clone())), - Self::VectorDb(client) => Some(Self::VectorDb(client.clone())), - Self::InfluxDb(client) => Some(Self::InfluxDb(client.clone())), - Self::InfluxDb3(client) => Some(Self::InfluxDb3(client.clone())), - Self::VictoriaMetrics(client) => Some(Self::VictoriaMetrics(client.clone())), - Self::Agent(client) => Some(Self::Agent(client.clone())), - Self::ExternalDriver { driver_id, config, session } => Some(Self::ExternalDriver { - driver_id: driver_id.clone(), - config: config.clone(), - session: session.clone(), - }), - Self::MessageQueue => Some(Self::MessageQueue), - Self::Nacos => Some(Self::Nacos), - Self::Consul(client) => Some(Self::Consul(client.clone())), - #[cfg(feature = "mq-admin")] - Self::Mqtt(client) => Some(Self::Mqtt(client.clone())), - _ => None, - } - } - pub fn agent(client: db::agent_driver::AgentDriverClient) -> Self { Self::Agent(Arc::new(db::agent_driver::PooledAgentClient::new(client))) } @@ -173,6 +134,89 @@ impl PoolKind { } } +#[derive(Clone)] +struct PoolPublication(Arc<()>); + +impl PoolPublication { + fn new() -> Self { + Self(Arc::new(())) + } + + fn is_same(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } +} + +#[derive(Clone)] +struct PoolPublicationSnapshot { + pool: PoolKind, + publication: PoolPublication, +} + +#[cfg(test)] +#[derive(Default)] +struct StalePoolCleanupBarriers { + before_removal: Option<(Arc, Arc)>, + after_removal: Option<(Arc, Arc)>, +} + +/// Internal pool registry that assigns an opaque identity to every published +/// entry. The identity belongs to the registry publication rather than to the +/// driver handle, so replacing an entry always creates a new generation even +/// when the replacement reuses a cloned handle. +#[doc(hidden)] +pub struct ConnectionPoolRegistry { + pools: HashMap, + publications: HashMap, +} + +impl ConnectionPoolRegistry { + fn new() -> Self { + Self { pools: HashMap::new(), publications: HashMap::new() } + } + + pub fn insert(&mut self, pool_key: String, pool: PoolKind) -> Option { + self.publications.insert(pool_key.clone(), PoolPublication::new()); + self.pools.insert(pool_key, pool) + } + + pub fn remove(&mut self, pool_key: &str) -> Option { + self.publications.remove(pool_key); + self.pools.remove(pool_key) + } + + #[cfg(test)] + fn clear(&mut self) { + self.publications.clear(); + self.pools.clear(); + } + + fn drain(&mut self) -> std::collections::hash_map::Drain<'_, String, PoolKind> { + self.publications.clear(); + self.pools.drain() + } + + fn snapshot(&self, pool_key: &str) -> Option { + Some(PoolPublicationSnapshot { + pool: self.pools.get(pool_key)?.clone(), + publication: self.publications.get(pool_key)?.clone(), + }) + } + + fn remove_if_publication(&mut self, pool_key: &str, expected: &PoolPublication) -> Option { + let is_current = self.publications.get(pool_key).is_some_and(|current| current.is_same(expected)); + is_current.then(|| self.remove(pool_key)).flatten() + } +} + +impl std::ops::Deref for ConnectionPoolRegistry { + type Target = HashMap; + + fn deref(&self) -> &Self::Target { + &self.pools + } +} + enum ConnectionDatabaseInfoSource { Agent(Arc), MongoAgent(Arc, Option), @@ -275,7 +319,7 @@ macro_rules! agent_connection_pool_database_type { } pub struct AppState { - pub connections: Arc>>, + connections: Arc>, task_supervisor: TaskSupervisor, pool_activity: Arc>>, draining_pools: Arc>>>, @@ -409,14 +453,14 @@ impl PoolActivity { pub struct PoolActivityTouch { pool_key: String, - connections: Arc>>, + connections: Arc>, pool_activity: Arc>>, task_supervisor: TaskSupervisor, } #[derive(Clone)] struct PoolRoutingControl { - connections: Arc>>, + connections: Arc>, pool_activity: Arc>>, postgres_cancel_contexts: Arc>>, task_supervisor: TaskSupervisor, @@ -507,10 +551,17 @@ impl PoolRoutingControl { &self, pool_key: &str, expected_client: &Arc, + expected_publication: Option<&PoolPublication>, replace_agent_runtime: bool, ) -> bool { let removed = { let mut connections = self.connections.write().await; + let publication_is_current = expected_publication.is_none_or(|expected| { + connections.publications.get(pool_key).is_some_and(|current| current.is_same(expected)) + }); + if !publication_is_current { + return false; + } let is_current = matches!( connections.get(pool_key), Some(PoolKind::Agent(current)) if Arc::ptr_eq(current, expected_client) @@ -583,7 +634,10 @@ impl PoolRoutingControl { }; let protects_manual_txn = match agent_client.as_ref() { Some(client) if client.uses_shared_runtime() => { - let connections = self.connections.read().await; + let connections = { + let registry = self.connections.read().await; + registry.pools.clone() + }; connections.iter().any(|(key, pool)| { is_manual_transaction_pool_key(key) && matches!(pool, PoolKind::Agent(sibling) if client.shares_runtime_with(sibling)) @@ -1123,6 +1177,46 @@ fn mysql_metadata_fallback_url( } impl AppState { + /// Return an owned pool handle. The registry read lock is released before + /// the caller can perform any asynchronous database operation. + pub async fn pool_handle(&self, pool_key: &str) -> Option { + self.connections.read().await.get(pool_key).cloned() + } + + async fn pool_publication_snapshot(&self, pool_key: &str) -> Option { + self.connections.read().await.snapshot(pool_key) + } + + async fn connection_pool_publication_snapshots(&self) -> Vec<(String, PoolPublicationSnapshot)> { + let connections = self.connections.read().await; + connections + .pools + .keys() + .filter_map(|pool_key| connections.snapshot(pool_key).map(|snapshot| (pool_key.clone(), snapshot))) + .collect() + } + + /// Return an owned snapshot for operations that need to inspect multiple + /// entries. Cloning handles is cheap and prevents registry guards from + /// leaking into asynchronous database work. + pub async fn connection_pools_snapshot(&self) -> HashMap { + self.connections.read().await.pools.clone() + } + + /// Inspect the registry while holding its read lock. The callback is + /// deliberately synchronous so no database I/O can run under the lock. + pub async fn with_connection_pools(&self, inspect: impl FnOnce(&HashMap) -> R) -> R { + let connections = self.connections.read().await; + inspect(&connections.pools) + } + + /// Mutate the registry atomically. The callback is deliberately + /// synchronous; asynchronous cleanup must use values returned from it. + pub async fn update_connection_pools(&self, update: impl FnOnce(&mut ConnectionPoolRegistry) -> R) -> R { + let mut connections = self.connections.write().await; + update(&mut connections) + } + fn pool_routing_control(&self) -> PoolRoutingControl { PoolRoutingControl { connections: self.connections.clone(), @@ -1229,7 +1323,7 @@ impl AppState { ) -> Self { let data_dir = storage.data_dir().to_path_buf(); Self { - connections: Arc::new(RwLock::new(HashMap::new())), + connections: Arc::new(RwLock::new(ConnectionPoolRegistry::new())), task_supervisor: TaskSupervisor::new(), pool_activity: Arc::new(RwLock::new(HashMap::new())), draining_pools: Arc::new(std::sync::Mutex::new(HashMap::new())), @@ -1619,8 +1713,11 @@ impl AppState { if let Some(pool) = previous { routing.close_pool_with_timeout(pool_key.clone(), pool).await; } - let route_is_available = - self.connections.read().await.get(&pool_key).is_some_and(PoolKind::is_available_for_routing); + let route_is_available = self + .with_connection_pools(|connections| { + connections.get(&pool_key).is_some_and(PoolKind::is_available_for_routing) + }) + .await; if !route_is_available { routing.detach_pool_by_key(&pool_key, true).await; return Err("Agent runtime is unavailable while publishing the connection pool".to_string()); @@ -2047,9 +2144,7 @@ impl AppState { loop { self.wait_for_pool_drain(&pool_key).await; - let conns = self.connections.read().await; - if conns.contains_key(&pool_key) { - drop(conns); + if self.pool_handle(&pool_key).await.is_some() { if self.remove_pool_if_duckdb_isolation_mismatch(&pool_key).await { // Recreate below using the current DuckDB isolation mode. } else if self.pool_credential_owner_mismatch(&config, &pool_key).await { @@ -2063,12 +2158,10 @@ impl AppState { } break; } - drop(conns); - // A reclaim may have removed the pool after the first drain check. Wait // for its confirmed close or rollback before deciding to create a new one. self.wait_for_pool_drain(&pool_key).await; - if self.connections.read().await.contains_key(&pool_key) { + if self.pool_handle(&pool_key).await.is_some() { continue; } break; @@ -2218,7 +2311,7 @@ impl AppState { db::redis_driver::connect_standalone(&db_config, &host, port, connect_timeout).await?, )) }; - PoolKind::Redis(con) + PoolKind::Redis(Arc::new(con)) } #[cfg(feature = "duckdb-sidecar")] DatabaseType::DuckDb => self.create_duckdb_pool(&db_config).await?, @@ -2253,7 +2346,7 @@ impl AppState { { Ok(()) => { // Re-check: another task may have created the pool while we were connecting. - if self.connections.read().await.contains_key(&pool_key) { + if self.pool_handle(&pool_key).await.is_some() { self.pool_routing_control() .close_pool_with_timeout(pool_key.clone(), PoolKind::MongoDb(client)) .await; @@ -3408,17 +3501,13 @@ impl AppState { return false; } - let mut checked_mysql_pool = None; + let Some(checked) = self.pool_publication_snapshot(pool_key).await else { + return false; + }; let stale = { - let connections = self.connections.read().await; - let Some(pool) = connections.get(pool_key) else { - return false; - }; - match pool { + match &checked.pool { PoolKind::Mysql(pool, _) => { let pool = pool.clone(); - checked_mysql_pool = Some(pool.clone()); - drop(connections); match db::mysql::checkout_mysql_conn(&pool, HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT).await { // The 500 ms probe budget is intentionally shorter than a foreground checkout. A timeout // while waiting, creating, or recycling is inconclusive: slow remote handshakes and active @@ -3452,7 +3541,6 @@ impl AppState { } PoolKind::Postgres(pool) => { let pool = pool.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::postgres::checkout_postgres_client_classified( &pool, @@ -3486,7 +3574,6 @@ impl AppState { } PoolKind::SqlServer(client) => { let client = client.clone(); - drop(connections); let mut client = client.lock().await; match db::sqlserver::test_connection(&mut client).await { Ok(()) => false, @@ -3505,7 +3592,6 @@ impl AppState { }, 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); @@ -3526,7 +3612,6 @@ impl AppState { } PoolKind::DynamoDb(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::dynamodb_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3538,7 +3623,6 @@ impl AppState { } PoolKind::ClickHouse(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::clickhouse_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3550,7 +3634,6 @@ impl AppState { } PoolKind::Elasticsearch(client) => { let mut client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::elasticsearch_driver::test_connection(&mut client, timeout).await { Ok(()) => false, @@ -3562,7 +3645,6 @@ impl AppState { } PoolKind::Easysearch(client) => { let mut client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::easysearch_driver::test_connection(&mut client, timeout).await { Ok(()) => false, @@ -3574,7 +3656,6 @@ impl AppState { } PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::meilisearch_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3586,7 +3667,6 @@ impl AppState { } PoolKind::HBase(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::hbase_driver::test_connection(&client, timeout).await { Ok(_) => false, @@ -3598,7 +3678,6 @@ impl AppState { } PoolKind::VectorDb(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::vector_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3610,7 +3689,6 @@ impl AppState { } PoolKind::InfluxDb(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::influxdb_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3622,7 +3700,6 @@ impl AppState { } PoolKind::InfluxDb3(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::influxdb3_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3634,7 +3711,6 @@ impl AppState { } PoolKind::VictoriaMetrics(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::victoriametrics_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3646,7 +3722,6 @@ impl AppState { } PoolKind::Rqlite(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::rqlite_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3658,7 +3733,6 @@ impl AppState { } PoolKind::Turso(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::turso_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3670,7 +3744,6 @@ impl AppState { } PoolKind::CloudflareD1(client) => { let client = client.clone(); - drop(connections); let timeout = crate::db::connection_timeout(); match db::cloudflare_d1_driver::test_connection(&client, timeout).await { Ok(()) => false, @@ -3682,7 +3755,6 @@ impl AppState { } PoolKind::Agent(client) => { let client = client.clone(); - drop(connections); let Ok(mut agent) = client.try_lock() else { log::debug!("Agent connection pool '{pool_key}' is busy; skipping health probe"); return false; @@ -3701,12 +3773,18 @@ impl AppState { "Agent connection pool '{pool_key}' requested runtime replacement during health probe: {err}" ); drop(agent); - return self.detach_agent_pool_if_current(pool_key, &client, true).await; + return self + .pool_routing_control() + .detach_agent_pool_if_current(pool_key, &client, Some(&checked.publication), true) + .await; } Err(err) => { log::warn!("Agent connection pool '{pool_key}' is stale: {err}"); drop(agent); - return self.detach_agent_pool_if_current(pool_key, &client, false).await; + return self + .pool_routing_control() + .detach_agent_pool_if_current(pool_key, &client, Some(&checked.publication), false) + .await; } } } @@ -3725,52 +3803,40 @@ impl AppState { return false; } - if let Some(checked_pool) = checked_mysql_pool { - return self.remove_stale_mysql_pool_if_current(pool_key, &checked_pool).await; - } - - let removed = self.connections.write().await.remove(pool_key); - - let Some(pool) = removed else { - return false; - }; - - self.stop_keepalive_task(pool_key).await; - self.pool_activity.write().await.remove(pool_key); - self.postgres_cancel_contexts.write().await.remove(pool_key); - self.pool_routing_control().close_pool_with_timeout(pool_key.to_string(), pool).await; - true + self.remove_stale_pool_if_current(pool_key, &checked.publication).await } - - async fn remove_stale_mysql_pool_if_current(&self, pool_key: &str, checked_pool: &db::mysql::MySqlPool) -> bool { - self.remove_stale_mysql_pool_if_current_inner( + async fn remove_stale_pool_if_current(&self, pool_key: &str, checked_publication: &PoolPublication) -> bool { + self.remove_stale_pool_if_current_inner( pool_key, - checked_pool, + checked_publication, #[cfg(test)] None, ) .await } - async fn remove_stale_mysql_pool_if_current_inner( + async fn remove_stale_pool_if_current_inner( &self, pool_key: &str, - checked_pool: &db::mysql::MySqlPool, - #[cfg(test)] cleanup_barriers: Option<( - std::sync::Arc, - std::sync::Arc, - )>, + checked_publication: &PoolPublication, + #[cfg(test)] cleanup_barriers: Option, ) -> bool { + #[cfg(test)] + if let Some((cleanup_ready, continue_cleanup)) = + cleanup_barriers.as_ref().and_then(|barriers| barriers.before_removal.as_ref()) + { + cleanup_ready.wait().await; + continue_cleanup.wait().await; + } + let routing = self.pool_routing_control(); let removed = loop { let mut connections = self.connections.write().await; - let is_current = matches!( - connections.get(pool_key), - Some(PoolKind::Mysql(current, _)) if checked_pool.is_same_pool(current) - ); + let is_current = + connections.publications.get(pool_key).is_some_and(|current| current.is_same(checked_publication)); if !is_current { log::debug!( - "MySQL connection pool '{pool_key}' was replaced while its health check was running; keeping the current route" + "Connection pool '{pool_key}' was replaced while its health check was running; keeping the current route" ); return false; } @@ -3779,10 +3845,13 @@ impl AppState { tokio::task::yield_now().await; continue; }; - let removed = remove_mysql_pool_if_current(&mut connections, pool_key, checked_pool) - .expect("checked MySQL pool must remain current while routing is locked"); + let removed = connections + .remove_if_publication(pool_key, checked_publication) + .expect("checked pool publication must remain current while routing is locked"); #[cfg(test)] - if let Some((route_removed, continue_cleanup)) = cleanup_barriers.as_ref() { + if let Some((route_removed, continue_cleanup)) = + cleanup_barriers.as_ref().and_then(|barriers| barriers.after_removal.as_ref()) + { route_removed.wait().await; continue_cleanup.wait().await; } @@ -4018,8 +4087,7 @@ impl AppState { pool_key_for_session_role(config.as_ref(), base_pool_key, client_session_id, AgentSessionRole::Metadata); if let Some(session_id) = agent_session_id { let expected_client = { - let connections = self.connections.read().await; - match connections.get(&pool_key) { + match self.pool_handle(&pool_key).await.as_ref() { Some(PoolKind::Agent(client)) if client.matches_session_id(session_id) => Some(client.clone()), Some(PoolKind::Agent(_)) => return false, _ => None, @@ -4043,7 +4111,9 @@ impl AppState { expected_client: &Arc, replace_agent_runtime: bool, ) -> bool { - self.pool_routing_control().detach_agent_pool_if_current(pool_key, expected_client, replace_agent_runtime).await + self.pool_routing_control() + .detach_agent_pool_if_current(pool_key, expected_client, None, replace_agent_runtime) + .await } async fn take_client_session_pool( @@ -4070,7 +4140,7 @@ impl AppState { self.stop_keepalive_task(&pool_key).await; self.pool_activity.write().await.remove(&pool_key); self.postgres_cancel_contexts.write().await.remove(&pool_key); - let removed = self.connections.write().await.remove(&pool_key); + let removed = self.update_connection_pools(|connections| connections.remove(&pool_key)).await; Ok(removed.map(|pool| (pool_key, pool))) } @@ -4090,7 +4160,7 @@ impl AppState { async fn reclaim_idle_base_pool_for_session(&self, connection_id: &str, preferred_base_pool_key: &str) -> bool { let pool_prefix = format!("{connection_id}:"); let activity = self.pool_activity.read().await; - let connections = self.connections.read().await; + let connections = self.connection_pools_snapshot().await; let mut candidates: Vec<(String, (usize, u64))> = connections .iter() .filter_map(|(key, pool)| { @@ -4167,8 +4237,7 @@ impl AppState { config_for_pool_key(pool_key, &configs).cloned() }; let client = { - let connections = self.connections.read().await; - match connections.get(pool_key) { + match self.pool_handle(pool_key).await.as_ref() { Some(PoolKind::Agent(client)) => Some(client.clone()), _ => None, } @@ -4247,7 +4316,7 @@ impl AppState { pub async fn active_agent_connection_driver_keys(&self) -> HashSet { let configs = self.configs.read().await; - let connections = self.connections.read().await; + let connections = self.connection_pools_snapshot().await; let mut keys = HashSet::new(); for (pool_key, pool) in connections.iter() { @@ -4316,8 +4385,7 @@ impl AppState { ExternalDriver { config: Arc, session: Arc }, } let source = { - let connections = self.connections.read().await; - match connections.get(&pool_key) { + match self.pool_handle(&pool_key).await.as_ref() { Some(PoolKind::Postgres(pool)) if config.db_type == DatabaseType::Gaussdb => { Some(IdentifierQuoteSource::NativeGaussdb(pool.clone())) } @@ -4372,8 +4440,7 @@ impl AppState { .ok_or_else(|| format!("Connection config not found: {connection_id}"))?; let pool_key = self.get_or_create_pool(connection_id, database).await?; let source = { - let connections = self.connections.read().await; - match connections.get(&pool_key) { + match self.pool_handle(&pool_key).await.as_ref() { Some(PoolKind::Agent(client)) if config.db_type == DatabaseType::MongoDb => { Some(ConnectionDatabaseInfoSource::MongoAgent(client.clone(), database.map(str::to_string))) } @@ -4443,13 +4510,10 @@ impl AppState { Some(ConnectionDatabaseInfoSource::VictoriaMetrics(client)) => { db::victoriametrics_driver::database_connection_info(&client, db::connection_timeout()).await.map(Some) } - Some(ConnectionDatabaseInfoSource::Redis(pool_key)) => { - let connections = self.connections.read().await; - match connections.get(&pool_key) { - Some(PoolKind::Redis(redis)) => db::redis_driver::database_connection_info(redis).await.map(Some), - _ => Ok(None), - } - } + Some(ConnectionDatabaseInfoSource::Redis(pool_key)) => match self.pool_handle(&pool_key).await.as_ref() { + Some(PoolKind::Redis(redis)) => db::redis_driver::database_connection_info(redis).await.map(Some), + _ => Ok(None), + }, Some(ConnectionDatabaseInfoSource::Nacos) => { let admin_config = self.nacos_admin_config_for_connection(connection_id, &config).await?; let admin = self.nacos_registry.get_or_build_config(connection_id, admin_config).await?; @@ -4554,11 +4618,8 @@ impl AppState { let pool_key = base_pool_key_for(db_type, connection_id, None, false); // Check if pool exists first - { - let connections = self.connections.read().await; - if !connections.contains_key(&pool_key) { - return Err("No active connection pool found".to_string()); - } + if self.pool_handle(&pool_key).await.is_none() { + return Err("No active connection pool found".to_string()); } // `remove_stale_connection_pool` returns true if the pool was stale (and removed) @@ -4571,26 +4632,15 @@ impl AppState { pub async fn refresh_connections(&self) { // Clone pool handles under a short-lived read lock, then release it // before performing I/O-heavy health checks to avoid blocking writers. - // Redis pools are handled separately because RedisConnection cannot be cloned. - let (checks, redis_keys): (Vec<(String, PoolKind)>, Vec) = { - let conns = self.connections.read().await; - let mut checks = Vec::with_capacity(conns.len()); - let mut redis_keys = Vec::new(); - for (key, pool) in conns.iter() { - match pool { - PoolKind::Redis(_) => redis_keys.push(key.clone()), - _ => checks.push((key.clone(), clone_pool_kind(pool))), - } - } - (checks, redis_keys) - }; + let checks = self.connection_pool_publication_snapshots().await; let mut failed_agent_checks = Vec::new(); let mut dead_pools = Vec::new(); let timeout = crate::db::connection_timeout(); // Check cloned pools (async I/O, no lock held) - for (key, pool) in &checks { + for (key, checked) in &checks { + let pool = &checked.pool; let healthy = match pool { PoolKind::Mysql(p, _) => match db::mysql::get_conn_with_health_check(p).await { Ok(_) => true, @@ -4765,6 +4815,7 @@ impl AppState { failed_agent_checks.push(( key.clone(), client.clone(), + checked.publication.clone(), RecoveryPolicy::decide(&e, RecoveryScope::Keepalive).replaces_runtime(), )); false @@ -4779,32 +4830,26 @@ impl AppState { | PoolKind::Consul(_) => true, #[cfg(feature = "mq-admin")] PoolKind::Mqtt(_) => true, - PoolKind::Redis(_) => unreachable!("Redis handled separately"), + PoolKind::Redis(redis) => match db::redis_driver::test_connection(redis).await { + Ok(()) => true, + Err(e) => { + log::warn!("Redis connection pool '{key}' is unhealthy: {e}"); + false + } + }, }; if !healthy && !matches!(pool, PoolKind::Agent(_)) { - dead_pools.push((key.clone(), agent_pool_identity(pool))); - } - } - - // Check Redis pools (read lock held briefly, no cloning needed) - { - let conns = self.connections.read().await; - for key in &redis_keys { - if let Some(PoolKind::Redis(redis)) = conns.get(key) { - match db::redis_driver::test_connection(redis).await { - Ok(()) => {} - Err(e) => { - log::warn!("Redis connection pool '{key}' is unhealthy: {e}"); - dead_pools.push((key.clone(), None)); - } - } - } + dead_pools.push((key.clone(), checked.publication.clone())); } } let mut detached_pool_keys = Vec::new(); - for (key, client, replace_runtime) in failed_agent_checks { - if self.detach_agent_pool_if_current(&key, &client, replace_runtime).await { + for (key, client, publication, replace_runtime) in failed_agent_checks { + if self + .pool_routing_control() + .detach_agent_pool_if_current(&key, &client, Some(&publication), replace_runtime) + .await + { detached_pool_keys.push(key); } } @@ -4813,20 +4858,11 @@ impl AppState { if !dead_pools.is_empty() { let mut conns = self.connections.write().await; let mut removed = Vec::with_capacity(dead_pools.len()); - for (key, expected_agent) in &dead_pools { - let still_checked_pool = match expected_agent { - Some(expected) => matches!( - conns.get(key), - Some(PoolKind::Agent(current)) if Arc::ptr_eq(current, expected) - ), - None => true, - }; - if still_checked_pool { - if let Some(pool) = conns.remove(key) { - removed.push((key.clone(), pool)); - } + for (key, publication) in &dead_pools { + if let Some(pool) = conns.remove_if_publication(key, publication) { + removed.push((key.clone(), pool)); } else { - log::debug!("Skipping stale Agent health result for replaced pool '{key}'"); + log::debug!("Skipping stale refresh health result for replaced pool '{key}'"); } } drop(conns); @@ -4872,7 +4908,7 @@ impl AppState { } async fn drain_all_connection_pools(&self) -> Vec<(String, PoolKind)> { - let pool_keys = self.connections.read().await.keys().cloned().collect::>(); + let pool_keys = self.connection_pools_snapshot().await.keys().cloned().collect::>(); self.stop_keepalive_tasks(&pool_keys).await; self.pool_activity.write().await.clear(); self.session_credentials.clear_pool_owners(); @@ -5070,7 +5106,7 @@ impl KeepaliveTarget { } async fn remove_keepalive_pool_if_current( - connections: &Arc>>, + connections: &Arc>, pool_key: &str, target: &KeepaliveTarget, ) -> Option { @@ -5084,13 +5120,13 @@ async fn remove_keepalive_pool_if_current( async fn detach_keepalive_target_if_current( routing: &PoolRoutingControl, - connections: &Arc>>, + connections: &Arc>, pool_key: &str, target: &KeepaliveTarget, replace_agent_runtime: bool, ) -> bool { if let KeepaliveTarget::Agent(expected) = target { - return routing.detach_agent_pool_if_current(pool_key, expected, replace_agent_runtime).await; + return routing.detach_agent_pool_if_current(pool_key, expected, None, replace_agent_runtime).await; } let Some(pool) = remove_keepalive_pool_if_current(connections, pool_key, target).await else { return false; @@ -5503,53 +5539,9 @@ fn pool_key_for_session_role( } } +#[cfg(test)] fn clone_pool_kind(pool: &PoolKind) -> PoolKind { - match pool { - PoolKind::Mysql(p, mode) => PoolKind::Mysql(p.clone(), *mode), - PoolKind::Postgres(p) => PoolKind::Postgres(p.clone()), - PoolKind::Sqlite(p) => PoolKind::Sqlite(p.clone()), - PoolKind::Rqlite(client) => PoolKind::Rqlite(client.clone()), - PoolKind::Turso(client) => PoolKind::Turso(client.clone()), - PoolKind::CloudflareD1(client) => PoolKind::CloudflareD1(client.clone()), - #[cfg(feature = "duckdb-sidecar")] - PoolKind::DuckDbWorker(client) => PoolKind::DuckDbWorker(client.clone()), - #[cfg(not(feature = "duckdb-sidecar"))] - PoolKind::DuckDbWorker(_) => PoolKind::DuckDbWorker(()), - PoolKind::MongoDb(client) => PoolKind::MongoDb(client.clone()), - PoolKind::DynamoDb(client) => PoolKind::DynamoDb(client.clone()), - PoolKind::ClickHouse(client) => PoolKind::ClickHouse(client.clone()), - PoolKind::SqlServer(client) => PoolKind::SqlServer(client.clone()), - PoolKind::Elasticsearch(client) => PoolKind::Elasticsearch(client.clone()), - PoolKind::Easysearch(client) => PoolKind::Easysearch(client.clone()), - PoolKind::Meilisearch(client) => PoolKind::Meilisearch(client.clone()), - PoolKind::HBase(client) => PoolKind::HBase(client.clone()), - PoolKind::VectorDb(client) => PoolKind::VectorDb(client.clone()), - PoolKind::InfluxDb(client) => PoolKind::InfluxDb(client.clone()), - PoolKind::InfluxDb3(client) => PoolKind::InfluxDb3(client.clone()), - PoolKind::VictoriaMetrics(client) => PoolKind::VictoriaMetrics(client.clone()), - PoolKind::Agent(client) => PoolKind::Agent(client.clone()), - PoolKind::ExternalDriver { driver_id, config, session } => { - PoolKind::ExternalDriver { driver_id: driver_id.clone(), config: config.clone(), session: session.clone() } - } - PoolKind::MessageQueue => PoolKind::MessageQueue, - PoolKind::Nacos => PoolKind::Nacos, - PoolKind::Consul(client) => PoolKind::Consul(client.clone()), - #[cfg(feature = "mq-admin")] - PoolKind::Mqtt(client) => PoolKind::Mqtt(Arc::clone(client)), - PoolKind::Redis(_) => panic!("clone_pool_kind not supported for Redis — handled separately"), - } -} - -fn remove_mysql_pool_if_current( - connections: &mut HashMap, - pool_key: &str, - expected: &db::mysql::MySqlPool, -) -> Option { - let is_current = matches!( - connections.get(pool_key), - Some(PoolKind::Mysql(current, _)) if expected.is_same_pool(current) - ); - is_current.then(|| connections.remove(pool_key)).flatten() + pool.clone() } async fn close_pool_kind(pool: PoolKind) -> Result<(), String> { @@ -5726,13 +5718,6 @@ fn should_validate_existing_pool_before_reuse(db_type: DatabaseType) -> bool { db_type != DatabaseType::Postgres && !matches!(db_type, agent_connection_pool_database_type!()) } -fn agent_pool_identity(pool: &PoolKind) -> Option> { - match pool { - PoolKind::Agent(client) => Some(client.clone()), - _ => None, - } -} - #[cfg(test)] fn uses_bare_mysql_pool(db_type: &DatabaseType) -> bool { matches!(db_type, DatabaseType::Doris | DatabaseType::StarRocks | DatabaseType::ManticoreSearch) @@ -5978,6 +5963,7 @@ mod tests { use crate::query; use crate::schema; use crate::storage::Storage; + use std::sync::Arc; use std::time::{Duration, Instant}; fn mysql_config(database: Option<&str>) -> ConnectionConfig { @@ -6862,6 +6848,51 @@ mod tests { PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub()) } + #[tokio::test] + async fn owned_pool_handle_does_not_block_registry_writes() { + let (state, dir) = test_app_state().await; + state + .update_connection_pools(|connections| { + connections.insert("slow-operation".to_string(), agent_pool_stub()); + }) + .await; + + let handle = state.pool_handle("slow-operation").await.expect("pool handle"); + tokio::time::timeout( + Duration::from_millis(100), + state.update_connection_pools(|connections| { + connections.insert("other-connection".to_string(), agent_pool_stub()); + }), + ) + .await + .expect("an owned handle must not retain the registry lock"); + + assert!(matches!(handle, PoolKind::Agent(_))); + assert!(state.pool_handle("other-connection").await.is_some()); + drop(state); + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn redis_pool_handle_survives_registry_removal() { + let (state, dir) = test_app_state().await; + let redis = Arc::new(crate::db::redis_driver::redis_connection_test_stub()); + state + .update_connection_pools(|connections| { + connections.insert("redis".to_string(), PoolKind::Redis(Arc::clone(&redis))); + }) + .await; + + let handle = state.pool_handle("redis").await.expect("Redis handle"); + let removed = state.update_connection_pools(|connections| connections.remove("redis")).await; + + assert!(matches!(handle, PoolKind::Redis(ref current) if Arc::ptr_eq(current, &redis))); + assert!(matches!(removed, Some(PoolKind::Redis(ref current)) if Arc::ptr_eq(current, &redis))); + assert!(state.pool_handle("redis").await.is_none()); + drop(state); + let _ = std::fs::remove_dir_all(dir); + } + #[tokio::test] async fn shutdown_releases_connection_pools_and_agent_daemons() { let (state, dir) = test_app_state().await; @@ -7675,26 +7706,57 @@ mod tests { assert_eq!(super::mysql_pool_max_connections_for_session(Some("tab-1")), 1); } - #[test] - fn stale_mysql_pool_observation_does_not_remove_replacement_generation() { - let checked = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10); - let checked_clone = checked.clone(); - let replacement = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10); - assert!(checked.is_same_pool(&checked_clone)); - assert!(!checked.is_same_pool(&replacement)); + #[tokio::test] + async fn stale_postgres_cleanup_preserves_concurrent_replacement_publication() { + let (state, dir) = test_app_state().await; + let state = std::sync::Arc::new(state); + let pool_key = "conn:app"; + let postgres_pool = || { + let manager = deadpool_postgres::Manager::new(tokio_postgres::Config::new(), tokio_postgres::NoTls); + deadpool_postgres::Pool::builder(manager) + .runtime(deadpool_postgres::Runtime::Tokio1) + .build() + .expect("build PostgreSQL test pool") + }; + let checked = postgres_pool(); + let replacement = postgres_pool(); + state.connections.write().await.insert(pool_key.to_string(), PoolKind::Postgres(checked)); + state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now()); + let checked_publication = state.pool_publication_snapshot(pool_key).await.unwrap().publication; + + let cleanup_ready = std::sync::Arc::new(tokio::sync::Barrier::new(2)); + let continue_cleanup = std::sync::Arc::new(tokio::sync::Barrier::new(2)); + let cleanup_state = state.clone(); + let cleanup_publication = checked_publication.clone(); + let cleanup_ready_for_task = cleanup_ready.clone(); + let cleanup_continue = continue_cleanup.clone(); + let cleanup = tokio::spawn(async move { + cleanup_state + .remove_stale_pool_if_current_inner( + pool_key, + &cleanup_publication, + Some(super::StalePoolCleanupBarriers { + before_removal: Some((cleanup_ready_for_task, cleanup_continue)), + after_removal: None, + }), + ) + .await + }); + cleanup_ready.wait().await; - let mut connections = std::collections::HashMap::from([( - "conn:app".to_string(), - PoolKind::Mysql(replacement.clone(), MysqlMode::Normal), - )]); + state.connections.write().await.insert(pool_key.to_string(), PoolKind::Postgres(replacement)); + let replacement_publication = state.pool_publication_snapshot(pool_key).await.unwrap().publication; + assert!(!checked_publication.is_same(&replacement_publication)); + continue_cleanup.wait().await; + assert!(!cleanup.await.unwrap()); - assert!(super::remove_mysql_pool_if_current(&mut connections, "conn:app", &checked).is_none()); - assert!( - matches!(connections.get("conn:app"), Some(PoolKind::Mysql(current, _)) if replacement.is_same_pool(current)) - ); + let current = state.pool_publication_snapshot(pool_key).await.expect("replacement must remain routable"); + assert!(matches!(current.pool, PoolKind::Postgres(_))); + assert!(current.publication.is_same(&replacement_publication)); + assert!(state.pool_activity.read().await.contains_key(pool_key)); - assert!(super::remove_mysql_pool_if_current(&mut connections, "conn:app", &replacement).is_some()); - assert!(!connections.contains_key("conn:app")); + state.shutdown(Duration::from_secs(1)).await; + let _ = std::fs::remove_dir_all(dir); } #[tokio::test] @@ -7712,9 +7774,10 @@ mod tests { .await .insert(pool_key.to_string(), PoolKind::Mysql(checked.clone(), MysqlMode::Normal)); state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now()); + let checked_publication = state.pool_publication_snapshot(pool_key).await.unwrap().publication; state.start_keepalive_task( pool_key, - &PoolKind::Mysql(checked.clone(), MysqlMode::Normal), + &PoolKind::Mysql(checked, MysqlMode::Normal), &config, #[cfg(feature = "mq-admin")] None, @@ -7724,15 +7787,18 @@ mod tests { let route_removed = std::sync::Arc::new(tokio::sync::Barrier::new(2)); let continue_cleanup = std::sync::Arc::new(tokio::sync::Barrier::new(2)); let cleanup_state = state.clone(); - let cleanup_checked = checked.clone(); + let cleanup_publication = checked_publication.clone(); let cleanup_route_removed = route_removed.clone(); let cleanup_continue = continue_cleanup.clone(); let cleanup = tokio::spawn(async move { cleanup_state - .remove_stale_mysql_pool_if_current_inner( + .remove_stale_pool_if_current_inner( pool_key, - &cleanup_checked, - Some((cleanup_route_removed, cleanup_continue)), + &cleanup_publication, + Some(super::StalePoolCleanupBarriers { + before_removal: None, + after_removal: Some((cleanup_route_removed, cleanup_continue)), + }), ) .await }); @@ -7759,11 +7825,9 @@ mod tests { assert!(cleanup.await.unwrap()); publish.await.unwrap().unwrap(); - let connections = state.connections.read().await; - assert!( - matches!(connections.get(pool_key), Some(PoolKind::Mysql(current, _)) if replacement.is_same_pool(current)) - ); - drop(connections); + let current = state.pool_publication_snapshot(pool_key).await.expect("replacement must remain routable"); + assert!(matches!(current.pool, PoolKind::Mysql(ref pool, _) if replacement.is_same_pool(pool))); + assert!(!current.publication.is_same(&checked_publication)); assert!(state.pool_activity.read().await.contains_key(pool_key)); assert_eq!(state.supervised_task_count(), 1); diff --git a/crates/dbx-core/src/consul/client.rs b/crates/dbx-core/src/consul/client.rs index 62166dcafa..e3d212c235 100644 --- a/crates/dbx-core/src/consul/client.rs +++ b/crates/dbx-core/src/consul/client.rs @@ -224,8 +224,8 @@ impl ConsulClient { pub(super) async fn client_for_state(state: &AppState, connection_id: &str) -> Result { state.get_or_create_pool(connection_id, None).await?; - let connections = state.connections.read().await; - match connections.get(connection_id) { + let pool_handle = state.pool_handle(connection_id).await; + match pool_handle.as_ref() { Some(PoolKind::Consul(client)) => Ok(client.clone()), Some(_) => Err("Connection is not a Consul connection".to_string()), None => Err("Connection not found".to_string()), diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 4bb3c3ec69..2dc0038e21 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -1484,8 +1484,8 @@ async fn list_postgres_extension_members( schema: &str, ) -> Result { let pool = { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(crate::connection::PoolKind::Postgres(pool)) => pool.clone(), _ => return Ok(PostgresExtensionMembers::default()), } @@ -1515,8 +1515,8 @@ async fn list_postgres_export_sequences( fail_on_error: bool, ) -> Result, String> { let pool = { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(crate::connection::PoolKind::Postgres(pool)) => pool.clone(), _ => return Ok(Vec::new()), } @@ -2718,7 +2718,8 @@ async fn export_database_sql_core_inner( let concurrent_prefetch_is_safe = match state.get_or_create_pool(&request.connection_id, Some(&request.database)).await { Ok(metadata_pool_key) => { - concurrent_metadata_prefetch_allowed(state.connections.read().await.get(&metadata_pool_key)) + let pool = state.pool_handle(&metadata_pool_key).await; + concurrent_metadata_prefetch_allowed(pool.as_ref()) } // 建池失败时不预取，让写出循环的直查路径按原有方式报告错误 Err(_) => false, diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 48e9a624c6..cf0e55ab8d 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -60,6 +60,7 @@ impl MySqlPool { /// generation. Pool options are not an identity: a reconnect creates a new /// pool with identical options, and a late health probe for the old pool /// must not be allowed to remove that replacement from routing. + #[cfg(test)] pub(crate) fn is_same_pool(&self, other: &Self) -> bool { std::sync::Arc::ptr_eq(&self.inner.metrics(), &other.inner.metrics()) } diff --git a/crates/dbx-core/src/db/redis_driver.rs b/crates/dbx-core/src/db/redis_driver.rs index 36420d58ac..a7b62a5951 100644 --- a/crates/dbx-core/src/db/redis_driver.rs +++ b/crates/dbx-core/src/db/redis_driver.rs @@ -305,6 +305,22 @@ pub enum RedisConnection { Cluster(RedisClusterPool), } +#[cfg(test)] +pub(crate) fn redis_connection_test_stub() -> RedisConnection { + RedisConnection::Cluster(RedisClusterPool { + connection: None, + seed_nodes: Vec::new(), + seed_routes: Vec::new(), + slot_ranges: Vec::new(), + node_routes: Vec::new(), + tls: false, + tls_insecure: false, + username: String::new(), + password: String::new(), + scan_sessions: Box::new(Mutex::new(RedisClusterScanSessions::default())), + }) +} + pub struct RedisClusterPool { pub connection: Option>, pub seed_nodes: Vec, diff --git a/crates/dbx-core/src/document_ops.rs b/crates/dbx-core/src/document_ops.rs index 11b57604a9..919b5decea 100644 --- a/crates/dbx-core/src/document_ops.rs +++ b/crates/dbx-core/src/document_ops.rs @@ -54,8 +54,8 @@ async fn ensure_document_pool(state: &AppState, connection_id: &str) -> Result<( pub async fn 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => match mongo_driver::list_databases(client).await { Ok(databases) => Ok(sort_names(databases)), Err(error) if mongo_list_databases_unauthorized(&error) => { @@ -247,8 +247,8 @@ pub async fn list_collections_core( 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { let specs = sort_mongo_collection_specs(mongo_driver::list_collection_specs(client, database).await?); let names: Vec = specs.iter().map(|spec| spec.name.clone()).collect(); @@ -258,7 +258,6 @@ pub async fn list_collections_core( } PoolKind::DynamoDb(client) => { let client = client.clone(); - drop(connections); let names = dynamodb_driver::list_tables(&client).await?; Ok(names .into_iter() @@ -306,8 +305,8 @@ pub async fn list_gridfs_files_core( sort: Option<&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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::list_gridfs_files(client, database, bucket, filter, sort).await, PoolKind::Agent(_) => Err("MongoDB legacy agent does not support GridFS file browsing".to_string()), _ => Err("Not a MongoDB connection".to_string()), @@ -322,8 +321,8 @@ pub async fn list_gridfs_buckets_core( sort: Option<&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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { let names = sort_names(mongo_driver::list_collections(client, database).await?); let bucket_names = mongo_gridfs_bucket_names(&names); @@ -345,8 +344,8 @@ pub async fn create_gridfs_bucket_core( bucket: &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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::create_gridfs_bucket(client, database, bucket).await, PoolKind::Agent(_) => Err("MongoDB legacy agent does not support GridFS bucket creation".to_string()), _ => Err("Not a MongoDB connection".to_string()), @@ -360,8 +359,8 @@ pub async fn delete_gridfs_bucket_core( bucket: &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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::delete_gridfs_bucket(client, database, bucket).await, PoolKind::Agent(_) => Err("MongoDB legacy agent does not support GridFS bucket deletion".to_string()), _ => Err("Not a MongoDB connection".to_string()), @@ -376,8 +375,8 @@ pub async fn download_gridfs_file_core( file_id: &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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::download_gridfs_file(client, database, bucket, file_id).await, PoolKind::Agent(_) => Err("MongoDB legacy agent does not support GridFS download".to_string()), _ => Err("Not a MongoDB connection".to_string()), @@ -394,8 +393,8 @@ pub async fn upload_gridfs_file_core( content_type: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::upload_gridfs_file(client, database, bucket, file_name, data, content_type).await } @@ -412,8 +411,8 @@ pub async fn delete_gridfs_file_core( file_id: &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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::delete_gridfs_file(client, database, bucket, file_id).await, PoolKind::Agent(_) => Err("MongoDB legacy agent does not support GridFS file deletion".to_string()), _ => Err("Not a MongoDB connection".to_string()), @@ -436,8 +435,8 @@ pub async fn find_documents_core( cursor_pagination: bool, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { // Document browser responses must retain BSON type metadata so nested filters // can round-trip ObjectId, Date, and int64 values through Extended JSON. @@ -448,13 +447,11 @@ pub async fn find_documents_core( } PoolKind::DynamoDb(client) => { let client = client.clone(); - drop(connections); let _ = (database, skip, projection, collation); dynamodb_driver::find_items(&client, collection, limit, filter, sort, cursor).await } PoolKind::Elasticsearch(client) => { let client = client.clone(); - drop(connections); if cursor_pagination { elasticsearch_driver::find_documents_with_cursor(&client, collection, limit, filter, sort, cursor).await } else { @@ -463,7 +460,6 @@ pub async fn find_documents_core( } PoolKind::Easysearch(client) => { let client = client.clone(); - drop(connections); if cursor_pagination { easysearch_driver::find_documents_with_cursor(&client, collection, limit, filter, sort, cursor).await } else { @@ -472,12 +468,10 @@ pub async fn find_documents_core( } PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::find_documents(&client, collection, skip, limit, filter, sort).await } PoolKind::VectorDb(client) => { let client = client.clone(); - drop(connections); let _ = (filter, sort); vector_driver::find_documents(&client, database, collection, skip, limit).await } @@ -516,21 +510,18 @@ pub async fn count_document_store_documents_core( filter: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::DynamoDb(client) => { let client = client.clone(); - drop(connections); dynamodb_driver::count_items(&client, collection, filter).await } PoolKind::Elasticsearch(client) => { let client = client.clone(); - drop(connections); elasticsearch_driver::count_documents(&client, collection, filter).await } PoolKind::Easysearch(client) => { let client = client.clone(); - drop(connections); easysearch_driver::count_documents(&client, collection, filter).await } _ => Err("Document count is not supported for this connection".to_string()), @@ -543,11 +534,10 @@ pub async fn describe_dynamodb_table_core( table: &str, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::DynamoDb(client) => { let client = client.clone(); - drop(connections); dynamodb_driver::describe_table(&client, table).await } _ => Err("Not a DynamoDB connection".to_string()), @@ -561,16 +551,14 @@ pub async fn count_elasticsearch_documents_core( filter: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Elasticsearch(client) => { let client = client.clone(); - drop(connections); elasticsearch_driver::count_documents(&client, index, filter).await } PoolKind::Easysearch(client) => { let client = client.clone(); - drop(connections); easysearch_driver::count_documents(&client, index, filter).await } _ => Err("Not an Elasticsearch connection".to_string()), @@ -586,11 +574,10 @@ pub async fn elasticsearch_get_index_metadata_core( kind: ElasticsearchIndexMetadataKind, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Elasticsearch(client) => { let client = client.clone(); - drop(connections); match kind { ElasticsearchIndexMetadataKind::Mapping => { elasticsearch_driver::get_index_mapping(&client, index).await @@ -603,7 +590,6 @@ pub async fn elasticsearch_get_index_metadata_core( } PoolKind::Easysearch(client) => { let client = client.clone(); - drop(connections); match kind { ElasticsearchIndexMetadataKind::Mapping => easysearch_driver::get_index_mapping(&client, index).await, ElasticsearchIndexMetadataKind::Settings => easysearch_driver::get_index_settings(&client, index).await, @@ -621,16 +607,14 @@ pub async fn elasticsearch_delete_all_documents_core( index: &str, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Elasticsearch(client) => { let client = client.clone(); - drop(connections); elasticsearch_driver::delete_all_documents(&client, index).await } PoolKind::Easysearch(client) => { let client = client.clone(); - drop(connections); easysearch_driver::delete_all_documents(&client, index).await } _ => Err("Not an Elasticsearch connection".to_string()), @@ -651,27 +635,23 @@ pub async fn insert_document_core( routing: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::insert_document(client, database, collection, doc_json).await, PoolKind::DynamoDb(client) => { let client = client.clone(); - drop(connections); dynamodb_driver::insert_item(&client, collection, doc_json).await } PoolKind::Elasticsearch(client) => { let client = client.clone(); - drop(connections); elasticsearch_driver::insert_document(&client, collection, doc_json, routing).await } PoolKind::Easysearch(client) => { let client = client.clone(); - drop(connections); easysearch_driver::insert_document(&client, collection, doc_json, routing).await } PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::insert_document(&client, collection, doc_json).await } PoolKind::Agent(client) => { @@ -698,15 +678,12 @@ pub async fn insert_document_preserving_bson_types_core( routing: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::insert_document_extended_json(client, database, collection, doc_json).await } - _ => { - drop(connections); - insert_document_core(state, connection_id, database, collection, doc_json, routing).await - } + _ => insert_document_core(state, connection_id, database, collection, doc_json, routing).await, } } @@ -720,29 +697,25 @@ pub async fn update_document_core( routing: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::update_document(client, database, collection, id, doc_json).await, PoolKind::DynamoDb(client) => { let client = client.clone(); - drop(connections); dynamodb_driver::update_item(&client, collection, id, doc_json).await } PoolKind::Elasticsearch(client) => { let client = client.clone(); - drop(connections); // Elasticsearch requires the same custom routing value for writes // as was used to index the document. elasticsearch_driver::update_document(&client, collection, id, doc_json, routing).await } PoolKind::Easysearch(client) => { let client = client.clone(); - drop(connections); easysearch_driver::update_document(&client, collection, id, doc_json, routing).await } PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::update_document(&client, collection, id, doc_json).await } PoolKind::Agent(client) => { @@ -782,29 +755,25 @@ pub async fn delete_document_core_with_type( document_type: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::delete_document(client, database, collection, id).await, PoolKind::DynamoDb(client) => { let client = client.clone(); - drop(connections); dynamodb_driver::delete_item(&client, collection, id).await } PoolKind::Elasticsearch(client) => { let client = client.clone(); - drop(connections); // Elasticsearch requires the same custom routing value for writes // as was used to index the document. elasticsearch_driver::delete_document(&client, collection, id, document_type, routing).await } PoolKind::Easysearch(client) => { let client = client.clone(); - drop(connections); easysearch_driver::delete_document(&client, collection, id, document_type, routing).await } PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::delete_document(&client, collection, id).await } PoolKind::Agent(client) => { @@ -826,11 +795,10 @@ pub async fn save_meilisearch_document_batch_core( inserts: &[String], ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::save_document_batch(&client, collection, updates, delete_ids, inserts).await } _ => Err("Not a Meilisearch connection".to_string()), @@ -853,11 +821,10 @@ pub async fn meilisearch_search_documents_core( ranking_score_threshold: Option, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); let hybrid = hybrid_embedder.map(|embedder| crate::db::meilisearch_driver::MeilisearchHybrid { embedder: embedder.to_string(), semantic_ratio: hybrid_semantic_ratio.unwrap_or(0.5), @@ -890,11 +857,10 @@ pub async fn meilisearch_fetch_document_page_core( offset: u64, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::fetch_document_page(&client, index, offset, limit, filter, sort).await } _ => Err("Not a Meilisearch connection".to_string()), @@ -907,11 +873,10 @@ pub async fn meilisearch_get_index_settings_core( index: &str, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::get_index_settings(&client, index).await } _ => Err("Not a Meilisearch connection".to_string()), @@ -925,11 +890,10 @@ pub async fn meilisearch_get_document_core( 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::get_document(&client, index, id).await } _ => Err("Not a Meilisearch connection".to_string()), @@ -943,11 +907,10 @@ pub async fn meilisearch_update_index_settings_core( settings: &serde_json::Value, ) -> Result<(), String> { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::update_index_settings(&client, index, settings).await } _ => Err("Not a Meilisearch connection".to_string()), @@ -960,11 +923,10 @@ pub async fn meilisearch_get_index_stats_core( index: &str, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::get_index_stats(&client, index).await } _ => Err("Not a Meilisearch connection".to_string()), @@ -977,11 +939,10 @@ pub async fn meilisearch_get_index_overview_core( index: &str, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::get_index_overview(&client, index).await } _ => Err("Not a Meilisearch connection".to_string()), @@ -990,11 +951,10 @@ pub async fn meilisearch_get_index_overview_core( pub async fn meilisearch_delete_index_core(state: &AppState, connection_id: &str, index: &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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::delete_index(&client, index).await } _ => Err("Not a Meilisearch connection".to_string()), @@ -1007,11 +967,10 @@ pub async fn meilisearch_delete_all_documents_core( index: &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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => { let client = client.clone(); - drop(connections); crate::db::meilisearch_driver::delete_all_documents(&client, index).await } _ => Err("Not a Meilisearch connection".to_string()), @@ -1023,8 +982,8 @@ async fn meilisearch_client_core( connection_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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Meilisearch(client) => Ok(client.clone()), _ => Err("Not a Meilisearch connection".to_string()), } diff --git a/crates/dbx-core/src/driver_runtime.rs b/crates/dbx-core/src/driver_runtime.rs index 89e017c6a3..e749a15503 100644 --- a/crates/dbx-core/src/driver_runtime.rs +++ b/crates/dbx-core/src/driver_runtime.rs @@ -147,7 +147,7 @@ async fn collect_runtime_seeds(state: &AppState) -> Vec { } let configs = state.configs.read().await; - let connections = state.connections.read().await; + let connections = state.connection_pools_snapshot().await; for (pool_key, pool) in connections.iter() { match pool { PoolKind::Agent(client) => { diff --git a/crates/dbx-core/src/hbase_ops.rs b/crates/dbx-core/src/hbase_ops.rs index bcacd1b863..bd5e1229a3 100644 --- a/crates/dbx-core/src/hbase_ops.rs +++ b/crates/dbx-core/src/hbase_ops.rs @@ -3,8 +3,8 @@ use crate::db::hbase_driver::{self, HBasePutRowInput, HBaseRow, HBaseScanResult, async fn client(state: &AppState, connection_id: &str) -> Result { let pool_key = state.get_or_create_pool(connection_id, None).await?; - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::HBase(client)) => Ok(client.clone()), _ => Err("Not an HBase connection".to_string()), } diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index e60d75d3b8..a266483ee6 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -29,8 +29,8 @@ pub async fn mongo_list_collections_core( pub async fn mongo_create_database_core(state: &AppState, 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::create_database(client, database).await, PoolKind::Agent(_) => Err("MongoDB legacy agent does not support create database".to_string()), _ => Err("Not a MongoDB connection".to_string()), @@ -40,8 +40,8 @@ pub async fn mongo_create_database_core(state: &AppState, connection_id: &str, d pub async fn mongo_drop_database_core(state: &AppState, connection_id: &str, database: &str) -> Result<(), String> { mongo_driver::validate_mongo_namespace_name(database, "Database")?; ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::drop_database(client, database).await, PoolKind::Agent(client) => { let mut client = client.lock().await; @@ -67,8 +67,8 @@ pub async fn mongo_drop_collection_core( mongo_driver::validate_mongo_namespace_name(database, "Database")?; mongo_driver::validate_mongo_namespace_name(collection, "Collection")?; ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::drop_collection(client, database, collection).await, PoolKind::Agent(client) => { let mut client = client.lock().await; @@ -92,8 +92,8 @@ pub async fn mongo_rename_collection_core( new_name: &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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::rename_collection(client, database, collection, new_name).await, PoolKind::Agent(_) => Err("MongoDB legacy agent does not support rename collection".to_string()), _ => Err("Not a MongoDB connection".to_string()), @@ -109,8 +109,8 @@ pub async fn mongo_clone_collection_core( ) -> Result { mongo_driver::validate_clone_collection_names(database, source_collection, target_collection)?; ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::clone_collection(client, database, source_collection, target_collection).await } @@ -140,8 +140,8 @@ pub async fn mongo_server_version_core( database: &str, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::server_version(client, database).await, PoolKind::Agent(client) => { let mut client = client.lock().await; @@ -158,8 +158,8 @@ pub async fn mongo_run_command_core( command_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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::run_command(client, database, command_json).await, PoolKind::Agent(client) => { let mut client = client.lock().await; @@ -192,8 +192,8 @@ pub async fn mongo_collection_stats_core( scale: Option, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::collection_stats(client, database, collection, scale).await, PoolKind::Agent(_) => Err("MongoDB legacy agent does not support collection stats helpers".to_string()), _ => Err("Not a MongoDB connection".to_string()), @@ -244,8 +244,8 @@ async fn mongo_find_documents_without_total_core( collation: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::find_documents_without_total( client, database, collection, skip, limit, filter, projection, sort, collation, @@ -284,8 +284,8 @@ pub async fn mongo_find_one_core( options: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::find_one(client, database, collection, filter, projection, options).await } @@ -320,8 +320,8 @@ pub async fn mongo_explain_find_core( verbosity: &str, ) -> Result { ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::explain_find( client, database, collection, skip, limit, filter, projection, sort, collation, verbosity, @@ -358,8 +358,8 @@ pub async fn mongo_count_documents_core( ) -> Result { let accurate = mode != Some("legacy"); ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::count_documents(client, database, collection, filter, accurate).await } @@ -406,8 +406,8 @@ pub async fn mongo_find_documents_extended_json_core( 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::find_documents_extended_json( client, database, collection, skip, limit, filter, projection, sort, None, @@ -454,8 +454,8 @@ pub async fn mongo_aggregate_documents_core( options_json: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::aggregate_documents(client, database, collection, pipeline_json, max_rows, options_json).await } @@ -490,8 +490,8 @@ pub async fn mongo_distinct_core( filter: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::distinct(client, database, collection, field, filter).await, // The legacy agent protocol has no distinct method and no read that could stand in for it. PoolKind::Agent(_) => Err("MongoDB legacy agent does not support distinct".to_string()), @@ -515,8 +515,8 @@ pub async fn mongo_list_index_specs_core( mongo_driver::validate_mongo_namespace_name(collection, "Collection")?; ensure_document_pool(state, connection_id).await?; let is_native = { - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(_) => true, PoolKind::Agent(_) => false, _ => return Err("Not a MongoDB connection".to_string()), @@ -529,8 +529,8 @@ pub async fn mongo_list_index_specs_core( return Ok(indexes.iter().map(mongo_driver::index_spec_from_index_info).collect()); } - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::list_index_specs(client, database, collection).await, _ => Err("Not a MongoDB connection".to_string()), } @@ -550,8 +550,8 @@ pub async fn mongo_create_index_core( mongo_driver::validate_mongo_namespace_name(collection, "Collection")?; mongo_driver::validate_create_index_request(keys_json, options_json)?; ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::create_index(client, database, collection, keys_json, options_json).await } @@ -586,8 +586,8 @@ pub async fn mongo_create_user_core( mongo_driver::validate_mongo_namespace_name(database, "Database")?; mongo_driver::validate_create_user_request(user_json, write_concern_json)?; ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::create_user(client, database, user_json, write_concern_json).await?; Ok(1) @@ -662,8 +662,8 @@ async fn mongo_drop_indexes_once_core( // to Native MongoDB or the Legacy Agent. mongo_driver::validate_drop_indexes_request(indexes_json, single)?; ensure_document_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::drop_indexes(client, database, collection, indexes_json, single).await } @@ -700,8 +700,8 @@ pub async fn mongo_insert_documents_core( 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => mongo_driver::insert_documents(client, database, collection, docs_json).await, PoolKind::Agent(client) => { let documents: serde_json::Value = @@ -747,8 +747,8 @@ pub async fn mongo_insert_documents_extended_json_core( 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::insert_documents_extended_json(client, database, collection, docs_json).await } @@ -780,8 +780,8 @@ pub async fn mongo_update_documents_core( options_json: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::update_documents(client, database, collection, filter_json, update_json, many, options_json) .await @@ -824,8 +824,8 @@ pub async fn mongo_delete_documents_core( 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::delete_documents(client, database, collection, filter_json, many).await } @@ -855,8 +855,8 @@ pub async fn mongo_find_one_and_update_core( options_json: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::find_one_and_update(client, database, collection, filter_json, update_json, options_json) .await @@ -876,8 +876,8 @@ pub async fn mongo_find_one_and_replace_core( options_json: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::find_one_and_replace( client, @@ -903,8 +903,8 @@ pub async fn mongo_find_one_and_delete_core( options_json: 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")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::MongoDb(client) => { mongo_driver::find_one_and_delete(client, database, collection, filter_json, options_json).await } @@ -1409,6 +1409,7 @@ mod tests { &[AgentCapability::MongoDropDatabase.as_str()], None, None, + None, ) .await } @@ -1420,8 +1421,16 @@ mod tests { expected_result: serde_json::Value, capabilities: &[&str], ) -> (AppState, tempfile::TempDir) { - legacy_mongo_state_with_options(expected_method, expected_params, expected_result, capabilities, None, None) - .await + legacy_mongo_state_with_options( + expected_method, + expected_params, + expected_result, + capabilities, + None, + None, + None, + ) + .await } #[cfg(unix)] @@ -1438,6 +1447,7 @@ mod tests { &[AgentCapability::MongoDropDatabase.as_str()], Some(server_version), None, + None, ) .await } @@ -1456,6 +1466,7 @@ mod tests { &[AgentCapability::MongoDropDatabase.as_str()], Some(server_version), Some(expected_error), + None, ) .await } @@ -1468,6 +1479,7 @@ mod tests { capabilities: &[&str], server_version: Option<&str>, expected_error: Option<&str>, + response_delay_ms: Option, ) -> (AppState, tempfile::TempDir) { use std::io::Write; @@ -1482,10 +1494,13 @@ mod tests { }; let server_version = python_optional_string(server_version); let expected_error = python_optional_string(expected_error); + let response_delay_ms = response_delay_ms.unwrap_or_default(); + let request_entered_path = serde_json::to_string(&directory.path().join("request-entered")).unwrap(); write!( script, r#"import json import sys +import time EXPECTED_METHOD = {expected_method} EXPECTED_PARAMS = json.loads({expected_params}) @@ -1493,6 +1508,8 @@ EXPECTED_RESULT = json.loads({expected_result}) CAPABILITIES = {capabilities} SERVER_VERSION = {server_version} EXPECTED_ERROR = {expected_error} +RESPONSE_DELAY_SECONDS = {response_delay_ms} / 1000 +REQUEST_ENTERED_PATH = {request_entered_path} expected_calls = 0 print(json.dumps({{"ready": True}}), flush=True) @@ -1515,6 +1532,9 @@ for line in sys.stdin: if expected_calls > 1: print(json.dumps({{"jsonrpc": "2.0", "id": request["id"], "error": {{"code": -1, "message": "duplicate MongoDB RPC"}}}}), flush=True) continue + if RESPONSE_DELAY_SECONDS: + open(REQUEST_ENTERED_PATH, "w").close() + time.sleep(RESPONSE_DELAY_SECONDS) if EXPECTED_ERROR is not None: print(json.dumps({{"jsonrpc": "2.0", "id": request["id"], "error": {{"code": -1, "message": EXPECTED_ERROR}}}}), flush=True) else: @@ -1547,7 +1567,11 @@ for line in sys.stdin: })) .unwrap(); state.configs.write().await.insert("legacy".to_string(), config); - state.connections.write().await.insert("legacy".to_string(), PoolKind::agent(client)); + state + .update_connection_pools(|connections| { + connections.insert("legacy".to_string(), PoolKind::agent(client)); + }) + .await; (state, directory) } @@ -1858,6 +1882,54 @@ for line in sys.stdin: assert_eq!(name, "email_1"); } + #[cfg(unix)] + #[tokio::test] + async fn stalled_legacy_create_index_does_not_block_registry_writes() { + let keys_json = r#"{"email":1}"#; + let options_json = r#"{"name":"email_1"}"#; + let (state, directory) = legacy_mongo_state_with_options( + "create_index", + serde_json::json!({ + "database": "app", + "collection": "users", + "keys_json": keys_json, + "options_json": options_json, + }), + serde_json::json!({ "name": "email_1" }), + &[AgentCapability::MongoDropDatabase.as_str()], + None, + None, + Some(500), + ) + .await; + let state = std::sync::Arc::new(state); + let index_state = std::sync::Arc::clone(&state); + let index_task = tokio::spawn(async move { + mongo_create_index_core(&index_state, "legacy", "app", "users", keys_json, Some(options_json)).await + }); + + let request_entered = directory.path().join("request-entered"); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while !request_entered.exists() { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("legacy createIndex request must reach the Agent"); + + tokio::time::timeout( + std::time::Duration::from_millis(100), + state.update_connection_pools(|connections| { + connections.insert("other-connection".to_string(), PoolKind::MessageQueue); + }), + ) + .await + .expect("a stalled Agent call must not retain the registry lock"); + + assert_eq!(index_task.await.expect("index task join").as_deref(), Ok("email_1")); + assert!(state.pool_handle("other-connection").await.is_some()); + } + #[test] fn mongo_indexes_query_result_matches_desktop_contract_and_limits_rows() { let result = mongo_indexes_query_result( diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 47104670d9..8e128216a7 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -1674,8 +1674,7 @@ async fn sqlserver_pool_is_current( pool_key: &str, client: &Arc>, ) -> bool { - let connections = state.connections.read().await; - matches!(connections.get(pool_key), Some(PoolKind::SqlServer(current)) if Arc::ptr_eq(current, client)) + matches!(state.pool_handle(pool_key).await, Some(PoolKind::SqlServer(current)) if Arc::ptr_eq(¤t, client)) } pub fn query_timeout_duration(timeout_secs: Option) -> Option { @@ -1769,13 +1768,12 @@ async fn do_execute_typed( let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await; let pool_db_type = connection_database_type_for_pool_key(state, pool_key).await; let mysql_catalog_dialect = connection_mysql_catalog_dialect_for_pool_key(state, pool_key).await; - let connections = state.connections.read().await; - let pool = connections.get(pool_key).ok_or("Connection not found")?; + let pool = state.pool_handle(pool_key).await.ok_or("Connection not found")?; let mut typed_agent_error = None; #[cfg(feature = "duckdb-sidecar")] let mut typed_duckdb_error = None; - let result: Result = match pool { + let result: Result = match &pool { #[cfg(feature = "duckdb-sidecar")] PoolKind::DuckDbWorker(client) => { let client = client.clone(); @@ -1793,7 +1791,6 @@ async fn do_execute_typed( let sql = sql.to_string(); let database = database.map(str::to_string); let max_rows = options.max_rows; - drop(connections); match client.execute_typed(database, sql, max_rows, cancel_token, query_timeout).await { Ok(result) => Ok(result), Err(error) => { @@ -1815,7 +1812,6 @@ async fn do_execute_typed( let bare = *mode == crate::connection::MysqlMode::Bare; let max_rows = options.max_rows; let max_result_bytes = options.max_result_bytes.filter(|value| *value > 0); - drop(connections); let mut conn = match db::mysql::get_conn_with_health_check_with_cancel( &p, operation_budget.checkout_timeout, @@ -1880,7 +1876,6 @@ async fn do_execute_typed( let prefer_text_protocol = postgres_prefers_text_protocol(pool_db_type); let execution_mode = options.execution_mode; let cancel_context = state.get_postgres_cancel_context(pool_key).await; - drop(connections); if execution_mode == QueryExecutionMode::PostgresReadOnlyTransaction { db::postgres::execute_query_in_read_only_transaction_with_rollback( &p, @@ -1920,14 +1915,12 @@ async fn do_execute_typed( PoolKind::Sqlite(p) => { let p = p.clone(); let max_rows = options.max_rows; - drop(connections); wait_for_query_opt(cancel_token, query_timeout, db::sqlite::execute_query_with_max_rows(&p, sql, max_rows)) .await } PoolKind::Rqlite(client) => { let client = client.clone(); let max_rows = options.max_rows; - drop(connections); wait_for_query_opt( cancel_token, query_timeout, @@ -1938,7 +1931,6 @@ async fn do_execute_typed( PoolKind::Turso(client) => { let client = client.clone(); let max_rows = options.max_rows; - drop(connections); wait_for_query_opt( cancel_token, query_timeout, @@ -1949,7 +1941,6 @@ async fn do_execute_typed( PoolKind::CloudflareD1(client) => { let client = client.clone(); let max_rows = options.max_rows; - drop(connections); wait_for_query_opt( cancel_token, query_timeout, @@ -1961,7 +1952,6 @@ async fn do_execute_typed( let client = client.clone(); let database = pool_key.split(':').nth(1).unwrap_or("default").to_string(); let max_rows = options.max_rows; - drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, @@ -1978,7 +1968,6 @@ async fn do_execute_typed( let client = client.clone(); let max_rows = options.max_rows; let execution_mode = options.execution_mode; - drop(connections); let (mut client, lock_wait_ms) = match lock_shared_client_with_wait(&client, cancel_token.clone(), None).await { Ok(value) => value, @@ -2010,7 +1999,6 @@ async fn do_execute_typed( let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows; - drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, @@ -2031,7 +2019,6 @@ async fn do_execute_typed( let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows; - drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, @@ -2052,7 +2039,6 @@ async fn do_execute_typed( let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows; - drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, @@ -2069,7 +2055,6 @@ async fn do_execute_typed( let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows; - drop(connections); let result = wait_for_query_opt(cancel_token, query_timeout, db::vector_driver::execute_rest_query(&client, &sql)) .await @@ -2089,7 +2074,6 @@ async fn do_execute_typed( let client = client.clone(); let database = pool_key.split(':').nth(1).unwrap_or("default").to_string(); let max_rows = options.max_rows; - drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, @@ -2106,7 +2090,6 @@ async fn do_execute_typed( let client = client.clone(); let database = pool_key.split(':').nth(1).unwrap_or("default").to_string(); let max_rows = options.max_rows; - drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, @@ -2122,7 +2105,6 @@ async fn do_execute_typed( PoolKind::VictoriaMetrics(client) => { let client = client.clone(); let max_rows = options.max_rows; - drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, @@ -2144,7 +2126,6 @@ async fn do_execute_typed( let schema = schema_for_execution_context(pool_db_type, schema).map(|s| s.to_string()); let max_rows = options.max_rows; let rpc_timeout = query_timeout; - drop(connections); if is_canceled(&cancel_token) { return Err(canceled_error().into()); } @@ -2203,7 +2184,6 @@ async fn do_execute_typed( let database = database.unwrap_or_else(|| config.effective_database().unwrap_or("")).to_string(); let max_rows = options.max_rows; let plugin_timeout = query_timeout; - drop(connections); wait_for_query_opt(cancel_token, query_timeout, async move { if let Some(session_id) = options.result_session_id.as_deref() { let params = external_driver_fetch_query_page_params( @@ -2233,7 +2213,6 @@ async fn do_execute_typed( let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows.unwrap_or(MAX_ROWS); - drop(connections); // Keep the AWS SDK cold-path future off this already-large query dispatcher stack. let execution = Box::pin(db::dynamodb_driver::execute_statement(&client, &sql, max_rows)); wait_for_query_opt(cancel_token, query_timeout, execution).await @@ -2726,8 +2705,8 @@ async fn execute_postgres_drop_database( check_read_only_for_connection(state, &pool_key, sql).await?; let pool = { - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Postgres(pool)) => pool.clone(), Some(_) => return Err("DROP DATABASE reconnect did not create a PostgreSQL connection".to_string()), None => return Err("Connection not found".to_string()), @@ -2786,19 +2765,17 @@ pub async fn close_query_session( let pool_database = query_pool_database(database, catalog); let pool_key = state.get_or_create_pool_for_session(connection_id, pool_database, client_session_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or("Connection not found")?; + let pool_handle = state.pool_handle(&pool_key).await; + let pool = pool_handle.as_ref().ok_or("Connection not found")?; match pool { PoolKind::Agent(client) => { let client = client.clone(); - drop(connections); let mut client = client.lock().await; client.close_query_session(session_id).await } PoolKind::ExternalDriver { config, session, .. } => { let config = config.clone(); let session = session.clone(); - drop(connections); let params = external_driver_fetch_query_page_params(config.as_ref(), session_id, 1); session .invoke::("closeQuerySession", params) @@ -2807,13 +2784,11 @@ pub async fn close_query_session( } PoolKind::Elasticsearch(client) => { let client = client.clone(); - drop(connections); db::elasticsearch_driver::close_cursor(&client, session_id).await?; Ok(true) } PoolKind::Easysearch(client) => { let client = client.clone(); - drop(connections); db::easysearch_driver::close_cursor(&client, session_id).await?; Ok(true) } @@ -2954,10 +2929,7 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( state.touch_pool_activity(&pool_key).await; let _activity_touch = state.pool_activity_touch(pool_key.as_str()); - let is_sqlserver = { - let connections = state.connections.read().await; - matches!(connections.get(&pool_key), Some(PoolKind::SqlServer(_))) - }; + let is_sqlserver = { matches!(state.pool_handle(&pool_key).await, Some(PoolKind::SqlServer(_))) }; if is_sqlserver { return execute_multi_sqlserver(state, &pool_key, sql, cancel_token, options).await.map_err(Into::into); @@ -3015,8 +2987,8 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( } let mysql_pool = { - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Mysql(pool, mode)) => Some((pool.clone(), *mode)), _ => None, } @@ -3595,13 +3567,12 @@ async fn execute_multi_sqlserver( break; } - let connections = state.connections.read().await; - let pool = connections.get(pool_key).ok_or("Connection not found")?; + let pool_handle = state.pool_handle(pool_key).await; + let pool = pool_handle.as_ref().ok_or("Connection not found")?; let client = match pool { PoolKind::SqlServer(c) => c.clone(), _ => return Err("Expected SQL Server connection".to_string()), }; - drop(connections); let (mut client_guard, lock_wait_ms) = match lock_shared_client_with_wait(&client, cancel_token.clone(), query_timeout).await { @@ -3698,8 +3669,8 @@ pub async fn execute_statements( let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; let agent_client = { - let conns = state.connections.read().await; - match conns.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Agent(client)) => Some(client.clone()), _ => None, } @@ -4047,10 +4018,8 @@ pub async fn execute_schema_diff_deploy( } } }; - let has_transactional_path = { - let conns = state.connections.read().await; - conns.get(&pool_key).is_some_and(pool_kind_has_transactional_path) - }; + let has_transactional_path = + { state.pool_handle(&pool_key).await.as_ref().is_some_and(pool_kind_has_transactional_path) }; let atomicity = classify_schema_diff_atomicity(db_type, &parsed, has_transactional_path); match execute_statements_in_transaction_on_pool(state, &pool_key, connection_id, database, &parsed, schema, None) @@ -4190,8 +4159,7 @@ pub async fn execute_statements_in_transaction_on_pool_typed( // Clone the pool handle within the lock, then drop it before any async work. let path = { - let conns = state.connections.read().await; - conns.get(pool_key).map(|p| match p { + state.pool_handle(pool_key).await.as_ref().map(|p| match p { PoolKind::Postgres(pg) => TxPath::Pg(pg.clone()), PoolKind::Mysql(mp, _mode) => TxPath::Mysql(mp.clone(), false), PoolKind::Sqlite(sq) => TxPath::Sqlite(sq.clone()), @@ -4736,8 +4704,8 @@ async fn begin_transaction_session( ExternalDriver, } let pool_handle = { - let connections = state.connections.read().await; - match connections.get(&probe_pool_key).ok_or("Connection not found")? { + let pool = state.pool_handle(&probe_pool_key).await.ok_or("Connection not found")?; + match &pool { PoolKind::Postgres(pg) => TxnPoolHandle::Postgres(pg.clone()), PoolKind::Mysql(mp, _) => TxnPoolHandle::Mysql(mp.clone()), PoolKind::Agent(_) if !consistent_snapshot => TxnPoolHandle::Agent, @@ -4792,8 +4760,8 @@ async fn begin_transaction_session( let agent_pool_key = state.get_or_create_pool_for_session(connection_id, pool_database, Some(&client_session_id)).await?; let client = { - let connections = state.connections.read().await; - match connections.get(&agent_pool_key) { + let pool_handle = state.pool_handle(&agent_pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Agent(client)) => client.clone(), _ => { let _ = state.close_client_session_pool(connection_id, pool_database, &client_session_id).await; @@ -4830,8 +4798,8 @@ async fn begin_transaction_session( let external_pool_key = state.get_or_create_pool_for_session(connection_id, pool_database, Some(&client_session_id)).await?; let (config, session) = { - let connections = state.connections.read().await; - match connections.get(&external_pool_key) { + let pool_handle = state.pool_handle(&external_pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::ExternalDriver { config, session, .. }) => (config.clone(), session.clone()), _ => { let _ = state.close_client_session_pool(connection_id, pool_database, &client_session_id).await; @@ -6215,7 +6183,11 @@ for line in sys.stdin: let client = db::dynamodb_driver::connect(&config, host, config.port).unwrap(); db::dynamodb_driver::test_connection(&client, Duration::from_secs(5)).await.unwrap(); state.configs.write().await.insert(config.id.clone(), config.clone()); - state.connections.write().await.insert(config.id.clone(), PoolKind::DynamoDb(client)); + state + .update_connection_pools(|connections| { + connections.insert(config.id.clone(), PoolKind::DynamoDb(client)); + }) + .await; let results = execute_multi_core_with_options_for_client_and_progress_typed( &state, @@ -6294,13 +6266,17 @@ for line in sys.stdin: let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); let state = AppState::new(storage); state.configs.write().await.insert("conn-1".to_string(), test_connection_config(DatabaseType::Dameng)); - state.connections.write().await.insert( - "conn-1".to_string(), - PoolKind::agent(crate::db::agent_driver::AgentDriverClient::shared_session( - runtime.clone(), - "session-1".to_string(), - )), - ); + state + .update_connection_pools(|connections| { + connections.insert( + "conn-1".to_string(), + PoolKind::agent(crate::db::agent_driver::AgentDriverClient::shared_session( + runtime.clone(), + "session-1".to_string(), + )), + ); + }) + .await; (state, dir, runtime) } @@ -6312,7 +6288,7 @@ for line in sys.stdin: let error = execute_sql_statement(&state, "conn-1", "", "SELECT 1", None, None).await.unwrap_err(); assert!(error.contains("injected Agent failure")); - assert!(!state.connections.read().await.contains_key("conn-1")); + assert!(!state.pool_handle("conn-1").await.is_some()); assert!(runtime.is_failed()); runtime.kill(); @@ -6350,7 +6326,11 @@ for line in sys.stdin: let state = AppState::new(storage); let connection_id = "sqlite-cancel"; let sqlite = db::sqlite::connect_path_create_if_missing(dir.join("query.db").to_str().unwrap()).await.unwrap(); - state.connections.write().await.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); + state + .update_connection_pools(|connections| { + connections.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); + }) + .await; state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Sqlite)); let cancel_token = CancellationToken::new(); cancel_token.cancel(); @@ -6434,7 +6414,7 @@ for line in sys.stdin: .unwrap_err(); assert!(error.contains("injected Agent failure")); - assert!(!state.connections.read().await.contains_key("conn-1")); + assert!(!state.pool_handle("conn-1").await.is_some()); assert!(runtime.is_failed()); runtime.kill(); @@ -6451,7 +6431,7 @@ for line in sys.stdin: .unwrap_err(); assert!(error.contains("injected Agent failure")); - assert!(!state.connections.read().await.contains_key("conn-1")); + assert!(!state.pool_handle("conn-1").await.is_some()); assert!(runtime.is_failed()); runtime.kill(); @@ -6465,7 +6445,7 @@ for line in sys.stdin: let error = execute_sql_statement(&state, "conn-1", "", "SELECT 1", None, None).await.unwrap_err(); assert!(error.contains("injected Agent failure")); - assert!(!state.connections.read().await.contains_key("conn-1")); + assert!(!state.pool_handle("conn-1").await.is_some()); assert!(!runtime.is_failed()); runtime.kill(); @@ -6526,7 +6506,11 @@ for line in sys.stdin: let state = AppState::new(storage); let connection_id = "sqlite-batch"; let sqlite = db::sqlite::connect_path_create_if_missing(dir.join("query.db").to_str().unwrap()).await.unwrap(); - state.connections.write().await.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); + state + .update_connection_pools(|connections| { + connections.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); + }) + .await; state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Sqlite)); let sql = if failure_first { @@ -6622,7 +6606,11 @@ for line in sys.stdin: let state = AppState::new(storage); let connection_id = "gaussdb-on-error-stop"; let sqlite = db::sqlite::connect_path_create_if_missing(dir.join("query.db").to_str().unwrap()).await.unwrap(); - state.connections.write().await.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); + state + .update_connection_pools(|connections| { + connections.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); + }) + .await; state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Gaussdb)); let results = execute_multi_core_with_options( @@ -7766,14 +7754,18 @@ for line in sys.stdin: let client_session_id = "manual-txn-test"; let pool_key = "jdbc-conn:session:manual-txn-test"; let config = Arc::new(config); - state.connections.write().await.insert( - pool_key.to_string(), - PoolKind::ExternalDriver { - driver_id: "jdbc".to_string(), - config: config.clone(), - session: session.clone(), - }, - ); + state + .update_connection_pools(|connections| { + connections.insert( + pool_key.to_string(), + PoolKind::ExternalDriver { + driver_id: "jdbc".to_string(), + config: config.clone(), + session: session.clone(), + }, + ); + }) + .await; let cleanup_guard = state.workload_session_pool_cleanup_guard("jdbc-conn", Some("dbx_test"), client_session_id).await.unwrap(); state.transaction_sessions.write().await.insert( @@ -7799,7 +7791,7 @@ for line in sys.stdin: execute_in_manual_transaction(&state, "txn-test", "SELECT 42", "dbx_test", None, Some(10)).await.unwrap(); assert_eq!(results[0].rows, vec![vec![serde_json::json!(42)]]); commit_manual_transaction(&state, "txn-test").await.unwrap(); - assert!(!state.connections.read().await.contains_key(pool_key)); + assert!(!state.pool_handle(pool_key).await.is_some()); assert_eq!( std::fs::read_to_string(&calls).unwrap(), "beginManualTransaction\nexecuteInManualTransaction\ncommitManualTransaction\n" diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 6961d2405c..4181198bc7 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -1037,14 +1037,13 @@ async fn try_export_postgres_query_result_stream( ) .await? }; - let connections = state.connections.read().await; - let Some(pool) = connections.get(&pool_key).and_then(|pool| match pool { + let pool_handle = state.pool_handle(&pool_key).await; + let Some(pool) = pool_handle.as_ref().and_then(|pool| match pool { PoolKind::Postgres(pool) => Some(pool.clone()), _ => None, }) else { return Ok(false); }; - drop(connections); if let Some(execution_id) = request.execution_id.as_deref() { state.running_queries.set_pool_key(execution_id, pool_key.clone()); @@ -1212,14 +1211,13 @@ async fn try_export_mysql_query_result_stream( ) .await? }; - let connections = state.connections.read().await; - let Some((pool, bare)) = connections.get(&pool_key).and_then(|pool| match pool { + let pool_handle = state.pool_handle(&pool_key).await; + let Some((pool, bare)) = pool_handle.as_ref().and_then(|pool| match pool { PoolKind::Mysql(pool, mode) => Some((pool.clone(), *mode == crate::connection::MysqlMode::Bare)), _ => None, }) else { return Ok(false); }; - drop(connections); if let Some(execution_id) = request.execution_id.as_deref() { state.running_queries.set_pool_key(execution_id, pool_key.clone()); @@ -1487,14 +1485,13 @@ async fn try_export_clickhouse_query_result_stream( ) .await? }; - let connections = state.connections.read().await; - let Some(client) = connections.get(&pool_key).and_then(|pool| match pool { + let pool_handle = state.pool_handle(&pool_key).await; + let Some(client) = pool_handle.as_ref().and_then(|pool| match pool { PoolKind::ClickHouse(client) => Some(client.clone()), _ => None, }) else { return Ok(false); }; - drop(connections); if let Some(execution_id) = request.execution_id.as_deref() { state.running_queries.set_pool_key(execution_id, pool_key.clone()); @@ -1659,14 +1656,13 @@ async fn try_export_sqlserver_query_result_stream( } let pool_key = state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?; - let connections = state.connections.read().await; - let Some(client) = connections.get(&pool_key).and_then(|pool| match pool { + let pool_handle = state.pool_handle(&pool_key).await; + let Some(client) = pool_handle.as_ref().and_then(|pool| match pool { PoolKind::SqlServer(client) => Some(client.clone()), _ => None, }) else { return Ok(false); }; - drop(connections); if let Some(execution_id) = request.execution_id.as_deref() { state.running_queries.set_pool_key(execution_id, pool_key); diff --git a/crates/dbx-core/src/redis_ops.rs b/crates/dbx-core/src/redis_ops.rs index d7420aaf6f..2de097c73b 100644 --- a/crates/dbx-core/src/redis_ops.rs +++ b/crates/dbx-core/src/redis_ops.rs @@ -13,10 +13,9 @@ pub async fn redis_list_databases_core( connection_id: &str, ) -> Result, String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { - PoolKind::Redis(redis) => match redis { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::list_databases(&mut *con).await @@ -54,10 +53,9 @@ pub async fn redis_scan_keys_batch_core( include_types: bool, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { - PoolKind::Redis(redis) => match redis { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -85,10 +83,9 @@ pub async fn redis_scan_values_core( count: usize, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { - PoolKind::Redis(redis) => match redis { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -115,12 +112,11 @@ pub async fn redis_get_value_in_db_core( key_raw: &str, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -146,12 +142,11 @@ pub async fn redis_get_ttl_in_db_core( key_raw: &str, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -176,12 +171,11 @@ pub async fn redis_stream_entries_in_db_core( cursor: Option<&str>, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -205,12 +199,11 @@ pub async fn redis_stream_groups_in_db_core( key_raw: &str, ) -> Result, String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -235,13 +228,12 @@ pub async fn redis_stream_consumers_in_db_core( group_raw: &str, ) -> Result, String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; let group = redis_driver::redis_key_raw_to_bytes(group_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -268,14 +260,13 @@ pub async fn redis_stream_pending_in_db_core( consumer_raw: Option<&str>, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; let group = redis_driver::redis_key_raw_to_bytes(group_raw)?; let consumer = consumer_raw.map(redis_driver::redis_key_raw_to_bytes).transpose()?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -311,12 +302,11 @@ pub async fn redis_set_string_in_db_core( ttl: Option, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -344,12 +334,11 @@ pub async fn redis_delete_key_in_db_core( key_raw: &str, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; - match pool { + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -374,15 +363,14 @@ pub async fn redis_rename_key_in_db_core( new_key_raw: &str, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or("Connection not found")?; + let pool = state.pool_handle(connection_id).await.ok_or("Connection not found")?; let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; let new_key = redis_driver::redis_key_raw_to_bytes(new_key_raw)?; if key == new_key { return Ok(()); } - match pool { - PoolKind::Redis(redis) => match redis { + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -419,11 +407,11 @@ pub async fn redis_hash_set_in_db_core( ttl: Option, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -452,11 +440,11 @@ pub async fn redis_hash_del_in_db_core( field: &str, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -489,11 +477,11 @@ pub async fn redis_hash_field_update_in_db_core( value: &str, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -519,11 +507,11 @@ pub async fn redis_hash_field_set_ttl_in_db_core( ttl: i64, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -549,11 +537,11 @@ pub async fn redis_hash_field_set_expire_at_in_db_core( expire_at: i64, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -589,11 +577,11 @@ pub async fn redis_list_push_in_db_core( ttl: Option, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -619,11 +607,11 @@ pub async fn redis_list_set_in_db_core( value: &str, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -657,11 +645,11 @@ pub async fn redis_list_remove_in_db_core( index: i64, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -697,11 +685,11 @@ pub async fn redis_set_add_in_db_core( ttl: Option, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -735,11 +723,11 @@ pub async fn redis_set_remove_in_db_core( member: &str, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -766,11 +754,11 @@ pub async fn redis_zadd_in_db_core( ttl: Option, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -795,11 +783,11 @@ pub async fn redis_zrem_in_db_core( member: &str, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -827,11 +815,11 @@ pub async fn redis_zset_update_in_db_core( score: &str, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -858,11 +846,11 @@ pub async fn redis_stream_add_in_db_core( ttl: Option, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -888,11 +876,11 @@ pub async fn redis_json_set_in_db_core( ttl: Option, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -915,9 +903,9 @@ pub async fn redis_check_json_module_in_db_core( db: u32, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::Redis(redis) => match redis { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -941,11 +929,11 @@ pub async fn redis_set_ttl_in_db_core( ttl: i64, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -970,11 +958,11 @@ pub async fn redis_set_expire_at_in_db_core( expire_at: i64, ) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -998,13 +986,13 @@ pub async fn redis_delete_keys_in_db_core( key_raws: &[String], ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let keys: Result>, String> = key_raws.iter().map(|k| redis_driver::redis_key_raw_to_bytes(k)).collect(); let keys = keys?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -1027,9 +1015,9 @@ pub async fn redis_delete_keys_in_db_core( pub async fn redis_flush_db_core(state: &AppState, connection_id: &str, db: u32) -> Result<(), String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::Redis(redis) => match redis { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -1052,9 +1040,9 @@ pub async fn redis_execute_command_core( skip_safety_check: bool, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::Redis(redis) => match redis { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::execute_console_command(&mut *con, db, command, skip_safety_check).await @@ -1104,11 +1092,11 @@ pub async fn redis_load_more_in_db_core( sort_direction: Option<&str>, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { PoolKind::Redis(redis) => { let key = redis_driver::redis_key_raw_to_bytes(key_raw)?; - match redis { + match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -1135,9 +1123,9 @@ pub async fn redis_publish_core( message: &str, ) -> Result { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::Redis(redis) => match redis { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; redis_driver::select_db(&mut *con, db).await?; @@ -1175,9 +1163,9 @@ pub async fn redis_slowlog_get_core( node_port: Option, ) -> Result, String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::Redis(redis) => match redis { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Direct(con) => { let mut con = con.lock().await; // SLOWLOG is a server-level command, no select_db needed @@ -1203,9 +1191,9 @@ pub async fn redis_cluster_master_nodes_core( connection_id: &str, ) -> Result, String> { ensure_redis_pool(state, connection_id).await?; - let connections = state.connections.read().await; - match connections.get(connection_id).ok_or("Not found")? { - PoolKind::Redis(redis) => match redis { + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { + PoolKind::Redis(redis) => match redis.as_ref() { RedisConnection::Cluster(cluster) => redis_driver::cluster_master_nodes(cluster).await, _ => Ok(Vec::new()), }, diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 9a9c2d143f..f1820d2ccd 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -15,8 +15,8 @@ use std::time::{Duration, Instant}; mod kingbase; macro_rules! extract_pool { - ($connections:expr, $key:expr, $variant:ident) => { - $connections.get($key).and_then(|v| match v { + ($pool:expr, $variant:ident) => { + $pool.and_then(|v| match v { PoolKind::$variant(val) => Some(val.clone()), _ => None, }) @@ -34,7 +34,7 @@ macro_rules! dispatch_mysql { } async fn clone_metadata_pool(state: &AppState, pool_key: &str) -> Option { - state.connections.read().await.get(pool_key).and_then(PoolKind::clone_for_metadata) + state.pool_handle(pool_key).await } struct EphemeralAgentMetadataSession { @@ -69,9 +69,8 @@ impl EphemeralAgentMetadataSession { } macro_rules! try_sqlserver { - ($connections:expr, $pool_key:expr, $method:ident $(, $arg:expr)*) => { - if let Some(client) = extract_pool!(&$connections, $pool_key, SqlServer) { - drop($connections); + ($pool:expr, $method:ident $(, $arg:expr)*) => { + if let Some(client) = extract_pool!($pool.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::$method(&mut client $(, $arg)*).await; } @@ -218,9 +217,8 @@ pub async fn list_xugu_tablespaces_core( ) .await?; let config = connection_config(state, connection_id).await; - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let mut client = client.lock().await; return client .list_xugu_tablespaces::>(database, agent_metadata_timeout(config.as_ref())) @@ -271,8 +269,8 @@ async fn list_database_storage_once( } let pool = { - let connections = state.connections.read().await; - match connections.get(connection_id) { + let pool_handle = state.pool_handle(connection_id).await; + match pool_handle.as_ref() { Some(PoolKind::Postgres(pool)) => pool.clone(), _ => return Ok(Vec::new()), } @@ -307,9 +305,8 @@ pub async fn list_sqlserver_linked_servers_core( connection_id: &str, ) -> Result, String> { let _metadata_permit = state.acquire_metadata_permit(connection_id, None, DatabaseType::SqlServer, None).await?; - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, connection_id, SqlServer) { - drop(connections); + let pool_handle = state.pool_handle(connection_id).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::list_linked_servers(&mut client).await; } @@ -327,11 +324,10 @@ pub async fn get_sqlserver_completion_context_core( let completion_context_sql = db::sqlserver::completion_context_sql_for_profile( db_config.as_ref().and_then(|config| config.driver_profile.as_deref()), ); - let connections = state.connections.read().await; - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { let config = config.clone(); let session = session.clone(); - drop(connections); let result: db::QueryResult = session .invoke_with_timeout( "executeQuery", @@ -346,9 +342,8 @@ pub async fn get_sqlserver_completion_context_core( .await?; return db::sqlserver::completion_context_from_query_result(result); } - try_sqlserver!(connections, &pool_key, get_completion_context); - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + try_sqlserver!(pool_handle, get_completion_context); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let mut client = client.lock().await; let result = client .execute_query_with_timeout::( @@ -374,9 +369,8 @@ pub async fn list_sqlserver_linked_server_catalogs_core( server: &str, ) -> Result, String> { let _metadata_permit = state.acquire_metadata_permit(connection_id, None, DatabaseType::SqlServer, None).await?; - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, connection_id, SqlServer) { - drop(connections); + let pool_handle = state.pool_handle(connection_id).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::list_linked_server_catalogs(&mut client, server).await; } @@ -390,9 +384,8 @@ pub async fn list_sqlserver_linked_server_schemas_core( catalog: &str, ) -> Result, String> { let _metadata_permit = state.acquire_metadata_permit(connection_id, None, DatabaseType::SqlServer, None).await?; - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, connection_id, SqlServer) { - drop(connections); + let pool_handle = state.pool_handle(connection_id).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::list_linked_server_schemas(&mut client, server, catalog).await; } @@ -410,9 +403,8 @@ pub async fn list_sqlserver_linked_server_tables_core( offset: Option, ) -> Result, String> { let _metadata_permit = state.acquire_metadata_permit(connection_id, None, DatabaseType::SqlServer, None).await?; - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, connection_id, SqlServer) { - drop(connections); + let pool_handle = state.pool_handle(connection_id).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::list_linked_server_tables(&mut client, server, catalog, schema, filter, limit, offset) .await; @@ -720,11 +712,10 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result>( "listDatabases", @@ -733,31 +724,25 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result R { return list_databases_once(state, connection_id).await; } - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, connection_id, SqlServer) { - drop(connections); + let pool_handle = state.pool_handle(connection_id).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::list_database_metadata(&mut client).await; } - if let Some(PoolKind::Mysql(pool, mode)) = connections.get(connection_id) { + if let Some(PoolKind::Mysql(pool, mode)) = pool_handle.as_ref() { let pool = pool.clone(); let mode = *mode; - drop(connections); return if mode == MysqlMode::OceanBaseOracle { db::ob_oracle::list_databases(&pool).await } else { db::mysql::list_database_metadata(&pool).await }; } - if let Some(pool) = extract_pool!(&connections, connection_id, Postgres) { - drop(connections); + if let Some(pool) = extract_pool!(pool_handle.as_ref(), Postgres) { return db::postgres::list_database_metadata(&pool).await; } - drop(connections); list_databases_once(state, connection_id).await } @@ -875,11 +856,10 @@ pub async fn list_data_types_core( retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let db_config = connection_config(state, connection_id).await; - let connections = state.connections.read().await; - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { let config = config.clone(); let session = session.clone(); - drop(connections); return session .invoke_with_timeout::>( "listDataTypes", @@ -889,8 +869,7 @@ pub async fn list_data_types_core( .await .map(deduplicate_data_type_names); } - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let mut client = client.lock().await; return client .list_data_types::>(database, agent_metadata_timeout(db_config.as_ref())) @@ -930,11 +909,10 @@ async fn list_schemas_once( let visible_schema_filter = visible_schema_filter(db_config.as_ref(), database, apply_visible_filter); { - let connections = state.connections.read().await; - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { let config = config.clone(); let session = session.clone(); - drop(connections); return session .invoke_with_timeout::>( "listSchemas", @@ -943,10 +921,9 @@ async fn list_schemas_once( ) .await; } - try_sqlserver!(connections, &pool_key, list_schemas); - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + try_sqlserver!(pool_handle, list_schemas); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let fallback_config = db_config.clone(); - drop(connections); let mut client = client.lock().await; match client .list_schemas_filtered::>( @@ -1097,8 +1074,8 @@ pub async fn list_vector_collections_core( ) .await?; let client = { - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::VectorDb(client)) => client.clone(), _ => return Err("Not a vector database connection".to_string()), } @@ -1121,8 +1098,8 @@ pub async fn get_vector_collection_detail_core( ) .await?; let client = { - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::VectorDb(client)) => client.clone(), _ => return Err("Not a vector database connection".to_string()), } @@ -1133,8 +1110,8 @@ pub async fn get_vector_collection_detail_core( pub async fn drop_vector_database_core(state: &AppState, connection_id: &str, database: &str) -> Result<(), String> { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let client = { - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::VectorDb(client)) => client.clone(), _ => return Err("Not a vector database connection".to_string()), } @@ -1150,8 +1127,8 @@ pub async fn drop_vector_collection_core( ) -> Result<(), String> { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let client = { - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::VectorDb(client)) => client.clone(), _ => return Err("Not a vector database connection".to_string()), } @@ -1168,8 +1145,8 @@ pub async fn rename_vector_collection_core( ) -> Result<(), String> { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let client = { - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::VectorDb(client)) => client.clone(), _ => return Err("Not a vector database connection".to_string()), } @@ -1249,15 +1226,14 @@ async fn get_table_comment_core_for_session( let db_config = connection_config(state, connection_id).await; { - let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, get_table_comment, schema, table); - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + let pool_handle = state.pool_handle(&pool_key).await; + try_sqlserver!(pool_handle, get_table_comment, schema, table); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { if db_config.as_ref().is_some_and(|config| { matches!(config.db_type, DatabaseType::Oracle | DatabaseType::OceanbaseOracle) }) { let sql = oracle_table_comment_sql(schema, table); let timeout = agent_metadata_timeout(db_config.as_ref()); - drop(connections); let mut client = client.lock().await; let result = client .execute_query_with_timeout::( @@ -1274,7 +1250,6 @@ async fn get_table_comment_core_for_session( } if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) { let timeout = agent_metadata_timeout(db_config.as_ref()); - drop(connections); let mut client = client.lock().await; return client.get_table_comment::>(database, schema, table, timeout).await; } @@ -1282,7 +1257,6 @@ async fn get_table_comment_core_for_session( let metadata_database = if schema.trim().is_empty() { database } else { schema }; let sql = tdengine_table_comment_sql(metadata_database, table); let timeout = agent_metadata_timeout(db_config.as_ref()); - drop(connections); let mut client = client.lock().await; let result = client .execute_query_with_timeout::( @@ -2258,12 +2232,11 @@ async fn list_tables_once( } { - let connections = state.connections.read().await; - if let Some(PoolKind::ExternalDriver { driver_id, config, session }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(PoolKind::ExternalDriver { driver_id, config, session }) = pool_handle.as_ref() { let driver_id = driver_id.clone(); let config = config.clone(); let session = session.clone(); - drop(connections); if uses_presto_like_information_schema_tables(&config.db_type) { let force_local_table_name_filter = table_name_filter.is_some_and(|filter| !filter.is_empty()); return external_driver_presto_like_tables( @@ -2309,17 +2282,15 @@ async fn list_tables_once( }); } #[cfg(feature = "duckdb-sidecar")] - if let Some(client) = extract_pool!(&connections, &pool_key, DuckDbWorker) { + if let Some(client) = extract_pool!(pool_handle.as_ref(), DuckDbWorker) { let database = database.to_string(); let schema = schema.to_string(); - drop(connections); return client .list_tables(database, schema) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types, table_name_filter)); } - if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), ClickHouse) { if requests_table_objects_only(object_types) && table_name_filter.is_none_or(TableNameFilter::is_empty) { return db::clickhouse_driver::list_table_objects_filtered( &client, @@ -2335,27 +2306,23 @@ async fn list_tables_once( .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types, table_name_filter)); } - if let Some(client) = extract_pool!(&connections, &pool_key, InfluxDb) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), InfluxDb) { return db::influxdb_driver::list_tables(&client, database) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types, table_name_filter)); } - if let Some(client) = extract_pool!(&connections, &pool_key, InfluxDb3) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), InfluxDb3) { return db::influxdb3_driver::list_tables(&client, database) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types, table_name_filter)); } - if let Some(client) = extract_pool!(&connections, &pool_key, VictoriaMetrics) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), VictoriaMetrics) { return db::victoriametrics_driver::list_tables(&client) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types, table_name_filter)); } if let Some(linked) = crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema) { - if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::list_linked_server_tables( &mut client, @@ -2371,23 +2338,21 @@ async fn list_tables_once( } } if requests_table_objects_only(object_types) && table_name_filter.is_none_or(TableNameFilter::is_empty) { - if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::list_table_objects(&mut client, schema, filter, limit, offset).await; } } if object_types.is_some() || table_name_filter.is_some_and(|filter| !filter.is_empty()) { - if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::list_tables(&mut client, schema, filter, None, None) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types, table_name_filter)); } } - try_sqlserver!(connections, &pool_key, list_tables, schema, filter, limit, offset); - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + try_sqlserver!(pool_handle, list_tables, schema, filter, limit, offset); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let use_mongodb_collection_listing = uses_mongodb_agent_collection_listing(db_config.as_ref()); let is_oracle = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle); let is_tdengine = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Tdengine); @@ -2400,7 +2365,6 @@ async fn list_tables_once( filter_locally_after_oracle_comments || filter_locally_after_tdengine_comments; let timeout_duration = agent_metadata_timeout(db_config.as_ref()); let fallback_config = db_config.clone(); - drop(connections); let mut client = client.lock().await; if use_mongodb_collection_listing { let collection_names = client.mongo_list_collections::>(database).await?; @@ -3255,7 +3219,11 @@ mod tests { config.host = base_url.to_string(); state.configs.write().await.insert(config.id.clone(), config); let client = db::turso_driver::TursoClient::new(base_url, "test-token", false, Duration::from_secs(2)).unwrap(); - state.connections.write().await.insert("test".to_string(), PoolKind::Turso(client)); + state + .update_connection_pools(|connections| { + connections.insert("test".to_string(), PoolKind::Turso(client)); + }) + .await; (state, dir) } @@ -3873,13 +3841,18 @@ mod tests { let storage = crate::storage::Storage::open(&dir.join("storage.db")).await.unwrap(); let state = crate::connection::AppState::new(storage); let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap(); - state.connections.write().await.insert("conn".to_string(), super::PoolKind::Sqlite(pool)); + state + .update_connection_pools(|connections| { + connections.insert("conn".to_string(), super::PoolKind::Sqlite(pool)); + }) + .await; let snapshot = super::clone_metadata_pool(&state, "conn").await.expect("metadata pool snapshot"); - let write_guard = state.connections.try_write().expect("snapshot must not retain the global pool-map lock"); + tokio::time::timeout(std::time::Duration::from_millis(100), state.update_connection_pools(|_| ())) + .await + .expect("snapshot must not retain the global pool-map lock"); assert!(matches!(snapshot, super::PoolKind::Sqlite(_))); - drop(write_guard); drop(snapshot); drop(state); let _ = std::fs::remove_dir_all(dir); @@ -3912,14 +3885,18 @@ mod tests { let mut config = test_connection_config(DatabaseType::Dameng); config.id = "conn".to_string(); state.configs.write().await.insert(config.id.clone(), config); - state.connections.write().await.insert( - "conn:analytics:role:metadata".to_string(), - super::PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub()), - ); + state + .update_connection_pools(|connections| { + connections.insert( + "conn:analytics:role:metadata".to_string(), + super::PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub()), + ); + }) + .await; replace_metadata_runtime(&state, "conn", Some("analytics"), None).await; - assert!(!state.connections.read().await.contains_key("conn:analytics:role:metadata")); + assert!(!state.pool_handle("conn:analytics:role:metadata").await.is_some()); let _ = std::fs::remove_dir_all(dir); } @@ -3933,7 +3910,11 @@ mod tests { config.id = "conn".to_string(); state.configs.write().await.insert(config.id.clone(), config); let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap(); - state.connections.write().await.insert("conn:role:metadata".to_string(), super::PoolKind::Sqlite(pool)); + state + .update_connection_pools(|connections| { + connections.insert("conn:role:metadata".to_string(), super::PoolKind::Sqlite(pool)); + }) + .await; let mut attempts = 0; let result = super::retry_metadata_connection_for_session(&state, "conn", None, None, || { @@ -3944,7 +3925,7 @@ mod tests { assert_eq!(result.unwrap_err(), "Agent RPC call timed out (30s)"); assert_eq!(attempts, 1); - assert!(!state.connections.read().await.contains_key("conn:role:metadata")); + assert!(!state.pool_handle("conn:role:metadata").await.is_some()); let _ = std::fs::remove_dir_all(dir); } @@ -3961,7 +3942,11 @@ mod tests { config.database = None; state.configs.write().await.insert(config.id.clone(), config); let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap(); - state.connections.write().await.insert("conn".to_string(), super::PoolKind::Sqlite(pool)); + state + .update_connection_pools(|connections| { + connections.insert("conn".to_string(), super::PoolKind::Sqlite(pool)); + }) + .await; let mut attempts = 0; let quarantine = "Agent RPC error (-1): connection lost\nDBX_AGENT_ERROR_DATA:{\"category\":\"connection\",\"sessionDisposition\":\"quarantine\"}"; @@ -3973,7 +3958,7 @@ mod tests { assert_eq!(result.unwrap_err(), quarantine); assert_eq!(attempts, 2); - assert!(!state.connections.read().await.contains_key("conn")); + assert!(!state.pool_handle("conn").await.is_some()); let _ = std::fs::remove_dir_all(dir); } @@ -4036,7 +4021,11 @@ for line in sys.stdin: config.id = "conn".to_string(); state.configs.write().await.insert(config.id.clone(), config); let pool_key = "conn:analytics:role:metadata"; - state.connections.write().await.insert(pool_key.to_string(), super::PoolKind::agent(client)); + state + .update_connection_pools(|connections| { + connections.insert(pool_key.to_string(), super::PoolKind::agent(client)); + }) + .await; let error = super::get_table_ddl_core(&state, "conn", "analytics", "APP", "EVENTS", None).await.unwrap_err(); @@ -4045,7 +4034,7 @@ for line in sys.stdin: Some(crate::db::agent_driver::AgentErrorCategory::Timeout) ); assert_eq!(std::fs::read_to_string(call_count_path).unwrap(), "1"); - assert!(!state.connections.read().await.contains_key(pool_key)); + assert!(!state.pool_handle(pool_key).await.is_some()); runtime.kill(); let _ = std::fs::remove_dir_all(dir); } @@ -5497,56 +5486,51 @@ pub async fn completion_assistant_search_core( .await?; log::debug!("[schema][completion_assistant:start] {request_summary}"); { - let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, completion_assistant_search, &request); + let pool_handle = state.pool_handle(&pool_key).await; + try_sqlserver!(pool_handle, completion_assistant_search, &request); } { - let connections = state.connections.read().await; - if let Some(pool) = connections.get(&pool_key).and_then(|pool| match pool { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(pool) = pool_handle.as_ref().and_then(|pool| match pool { PoolKind::Sqlite(pool) => Some(pool.clone()), _ => None, }) { - drop(connections); return db::sqlite::completion_assistant_search(&pool, &request).await; } } #[cfg(feature = "duckdb-sidecar")] { - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, &pool_key, DuckDbWorker) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), DuckDbWorker) { return client.completion_assistant(request.clone()).await; } } { - let connections = state.connections.read().await; - if let Some(pool) = connections.get(&pool_key).and_then(|pool| match pool { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(pool) = pool_handle.as_ref().and_then(|pool| match pool { PoolKind::Postgres(pool) => Some(pool.clone()), _ => None, }) { - drop(connections); return db::postgres::completion_assistant_search(&pool, &request).await; } } { - let connections = state.connections.read().await; - if let Some(pool) = connections.get(&pool_key).and_then(|pool| match pool { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(pool) = pool_handle.as_ref().and_then(|pool| match pool { PoolKind::Mysql(pool, mode) if *mode != MysqlMode::OceanBaseOracle => Some(pool.clone()), _ => None, }) { - drop(connections); return db::mysql::completion_assistant_search(&pool, &request).await; } } { - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let db_config = connection_config(state, &request.connection_id).await; let mut client = client.lock().await; match client @@ -5729,11 +5713,10 @@ async fn list_object_statistics_once( ) -> Result, String> { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let db_config = connection_config(state, connection_id).await; - let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, list_object_statistics, schema); - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + let pool_handle = state.pool_handle(&pool_key).await; + try_sqlserver!(pool_handle, list_object_statistics, schema); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle) { - drop(connections); return oracle_agent_list_object_statistics( client, database, @@ -5743,7 +5726,6 @@ async fn list_object_statistics_once( .await; } if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Dameng) { - drop(connections); return dameng_agent_list_object_statistics( client, database, @@ -5754,7 +5736,6 @@ async fn list_object_statistics_once( } if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) { let sql = kingbase::object_statistics_sql(schema); - drop(connections); return agent_list_object_statistics( client, database, @@ -5768,7 +5749,6 @@ async fn list_object_statistics_once( config.db_type == DatabaseType::Gbase && config.driver_profile.as_deref() != Some("gbase8s") }) { let sql = gbase8a_object_statistics_sql(database); - drop(connections); return agent_list_object_statistics( client, database, @@ -5779,11 +5759,9 @@ async fn list_object_statistics_once( .await; } } - if let Some(client) = extract_pool!(&connections, &pool_key, VictoriaMetrics) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), VictoriaMetrics) { return db::victoriametrics_driver::list_object_statistics(&client).await; } - drop(connections); let pool = clone_metadata_pool(state, &pool_key).await.ok_or("Pool not found")?; match &pool { PoolKind::Mysql(p, mode) => { @@ -5848,11 +5826,10 @@ async fn list_objects_once( }; { - let connections = state.connections.read().await; - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { let config = config.clone(); let session = session.clone(); - drop(connections); if uses_presto_like_information_schema_tables(&config.db_type) { return external_driver_presto_like_objects( session, @@ -5882,19 +5859,17 @@ async fn list_objects_once( .await .map(unpaged_object_list); } - if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::list_objects(&mut client, schema).await.map(unpaged_object_list); } - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let is_oracle = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle); let use_oracle_agent_paging = db_config.as_ref().is_some_and(is_default_oracle_agent_config); let filter_locally_after_oracle_comments = is_oracle && filter.is_some_and(|filter| !filter.trim().is_empty()); let timeout_duration = agent_metadata_timeout(db_config.as_ref()); let fallback_config = db_config.clone(); - drop(connections); if is_oracle && !use_oracle_agent_paging { return oracle_agent_list_objects(client, database, schema, timeout_duration) .await @@ -6079,11 +6054,10 @@ async fn list_completion_objects_once( state.get_or_create_metadata_pool_for_session(connection_id, Some(database), client_session_id).await?; let db_config = connection_config(state, connection_id).await; - let connections = state.connections.read().await; - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { let config = config.clone(); let session = session.clone(); - drop(connections); return session .invoke_with_timeout::>( "listObjects", @@ -6093,10 +6067,9 @@ async fn list_completion_objects_once( .await .map(filter_completion_objects); } - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let is_oracle = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle); let fallback_config = db_config.clone(); - drop(connections); let objects = if is_oracle { oracle_agent_list_objects(client, database, schema, agent_metadata_timeout(db_config.as_ref())).await? } else { @@ -6153,7 +6126,6 @@ async fn list_completion_objects_once( return Ok(filter_completion_objects(objects)); } - drop(connections); let pool = clone_metadata_pool(state, &pool_key).await.ok_or("Pool not found")?; match &pool { PoolKind::Mysql(p, mode) if *mode != MysqlMode::OceanBaseOracle => { @@ -6471,12 +6443,11 @@ async fn get_columns_core_for_session_inner( let db_config = connection_config(state, connection_id).await; { - let connections = state.connections.read().await; - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { let config = config.clone(); let session = session.clone(); - drop(connections); - if uses_presto_like_information_schema_tables(&config.db_type) { + if uses_presto_like_information_schema_tables(&config.db_type) { return external_driver_presto_like_columns(session, config.as_ref(), database, schema, table).await; } let query_oracle_columns_first = @@ -6544,35 +6515,29 @@ async fn get_columns_core_for_session_inner( return Ok(deduplicate_column_infos(columns)); } #[cfg(feature = "duckdb-sidecar")] - if let Some(client) = extract_pool!(&connections, &pool_key, DuckDbWorker) { + if let Some(client) = extract_pool!(pool_handle.as_ref(), DuckDbWorker) { let database = database.to_string(); let schema = schema.to_string(); let table = table.to_string(); - drop(connections); - return client.list_columns(database, schema, table).await; + return client.list_columns(database, schema, table).await; } - if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { - drop(connections); - return db::clickhouse_driver::get_columns(&client, clickhouse_metadata_database(database, schema), table) + if let Some(client) = extract_pool!(pool_handle.as_ref(), ClickHouse) { + return db::clickhouse_driver::get_columns(&client, clickhouse_metadata_database(database, schema), table) .await .map(deduplicate_column_infos); } - if let Some(client) = extract_pool!(&connections, &pool_key, InfluxDb) { - drop(connections); - return db::influxdb_driver::get_columns(&client, database, table).await.map(deduplicate_column_infos); + if let Some(client) = extract_pool!(pool_handle.as_ref(), InfluxDb) { + return db::influxdb_driver::get_columns(&client, database, table).await.map(deduplicate_column_infos); } - if let Some(client) = extract_pool!(&connections, &pool_key, InfluxDb3) { - drop(connections); - return db::influxdb3_driver::get_columns(&client, database, table).await.map(deduplicate_column_infos); + if let Some(client) = extract_pool!(pool_handle.as_ref(), InfluxDb3) { + return db::influxdb3_driver::get_columns(&client, database, table).await.map(deduplicate_column_infos); } - if let Some(client) = extract_pool!(&connections, &pool_key, VictoriaMetrics) { - drop(connections); - return db::victoriametrics_driver::get_columns(&client, table).await.map(deduplicate_column_infos); + if let Some(client) = extract_pool!(pool_handle.as_ref(), VictoriaMetrics) { + return db::victoriametrics_driver::get_columns(&client, table).await.map(deduplicate_column_infos); } if let Some(linked) = crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema) { - if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { - drop(connections); - let mut client = lock_sqlserver_metadata_client(&client).await?; + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { + let mut client = lock_sqlserver_metadata_client(&client).await?; return db::sqlserver::get_linked_server_columns( &mut client, &linked.server, @@ -6584,11 +6549,10 @@ async fn get_columns_core_for_session_inner( .map(deduplicate_column_infos); } } - try_sqlserver!(connections, &pool_key, get_columns, schema, table); - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + try_sqlserver!(pool_handle, get_columns, schema, table); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let fallback_config = db_config.clone(); - drop(connections); - let mut client = client.lock().await; + let mut client = client.lock().await; let oracle_sql_config = fallback_config.as_ref().filter(|config| { should_query_oracle_columns_via_sql_first(&config.db_type, schema, context_session_id) }); @@ -6771,8 +6735,8 @@ pub async fn get_sqlserver_column_metadata_core( ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; - let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, get_column_metadata, schema, table); + let pool_handle = state.pool_handle(&pool_key).await; + try_sqlserver!(pool_handle, get_column_metadata, schema, table); Err("SQL Server column metadata requires a native SQL Server connection".to_string()) }) .await @@ -6949,18 +6913,16 @@ async fn list_indexes_core_for_session( let db_config = connection_config(state, connection_id).await; { - let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, list_indexes, schema, table); - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + try_sqlserver!(pool_handle, list_indexes, schema, table); + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { if external_driver_uses_mysql_ddl(config.as_ref()) { let config = config.clone(); let session = session.clone(); - drop(connections); return external_driver_gaussdb_m_indexes(session, config.as_ref(), database, schema, table).await; } let config = config.clone(); let session = session.clone(); - drop(connections); return session .invoke_with_timeout::>( "listIndexes", @@ -6974,8 +6936,7 @@ async fn list_indexes_core_for_session( ) .await; } - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let mut client = client.lock().await; return client.list_indexes(database, schema, table, agent_metadata_timeout(db_config.as_ref())).await; } @@ -7057,10 +7018,9 @@ async fn list_foreign_keys_core_for_session( let db_config = connection_config(state, connection_id).await; { - let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, list_foreign_keys, schema, table); - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + try_sqlserver!(pool_handle, list_foreign_keys, schema, table); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let mut client = client.lock().await; return client .list_foreign_keys(database, schema, table, agent_metadata_timeout(db_config.as_ref())) @@ -7107,10 +7067,9 @@ pub async fn list_triggers_core( let db_config = connection_config(state, connection_id).await; { - let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, list_triggers, schema, table); - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + try_sqlserver!(pool_handle, list_triggers, schema, table); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let mut client = client.lock().await; return client.list_triggers(database, schema, table, agent_metadata_timeout(db_config.as_ref())).await; } @@ -7152,9 +7111,8 @@ pub async fn list_constraints_core( let db_config = connection_config(state, connection_id).await; { - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let mut client = client.lock().await; return client .list_constraints(database, schema, table, agent_metadata_timeout(db_config.as_ref())) @@ -7194,9 +7152,8 @@ pub async fn list_partitions_core( retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let db_config = connection_config(state, connection_id).await; - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let mut client = client.lock().await; return client.list_partitions(database, schema, table, agent_metadata_timeout(db_config.as_ref())).await; } @@ -7227,8 +7184,8 @@ pub async fn table_partition_status_core( ) -> Result { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Postgres(pool)) => { let info = db::postgres::get_table_partition_info(pool, schema, table).await?; Ok(TablePartitionStatus { @@ -7253,8 +7210,8 @@ pub async fn list_invalid_indexes_core( ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Postgres(pool)) => db::postgres::list_invalid_indexes(pool, schema, table).await, _ => Ok(vec![]), } @@ -7272,9 +7229,8 @@ pub async fn list_subpartitions_core( retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let db_config = connection_config(state, connection_id).await; - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let mut client = client.lock().await; return client .list_subpartitions(database, schema, table, agent_metadata_timeout(db_config.as_ref())) @@ -7354,9 +7310,8 @@ pub async fn list_extensions_core( let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let db_config = connection_config(state, connection_id).await; if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) { - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { return kingbase::list_extensions(client, database, schema, agent_metadata_timeout(db_config.as_ref())) .await; } @@ -7390,9 +7345,8 @@ pub async fn list_available_extensions_core( let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let db_config = connection_config(state, connection_id).await; if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) { - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { return kingbase::list_available_extensions( client, database, @@ -7651,23 +7605,20 @@ async fn get_table_ddl_once( let db_config = connection_config(state, connection_id).await; { - let connections = state.connections.read().await; - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { if external_driver_uses_mysql_ddl(config.as_ref()) { let config = config.clone(); let session = session.clone(); - drop(connections); return external_driver_mysql_ddl(session, config.as_ref(), database, schema, table).await; } } #[cfg(feature = "duckdb-sidecar")] - if let Some(client) = extract_pool!(&connections, &pool_key, DuckDbWorker) { + if let Some(client) = extract_pool!(pool_handle.as_ref(), DuckDbWorker) { let client = client.clone(); - drop(connections); return client.get_table_ddl(database.to_string(), schema.to_string(), table.to_string()).await; } - if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), ClickHouse) { let clickhouse_database = clickhouse_metadata_database(database, schema); let result = db::clickhouse_driver::execute_query( &client, @@ -7683,13 +7634,11 @@ async fn get_table_ddl_once( .map(|s| s.to_string()) .ok_or_else(|| "Table not found".to_string()); } - if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; return build_sqlserver_ddl(&mut client, schema, table).await; } - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { if let Some(config) = db_config.as_ref().filter(|config| is_agent_postgres_metadata_fallback_config(config)) { match native_postgres_metadata_pool(state, connection_id, database, config).await { @@ -8739,10 +8688,9 @@ async fn get_custom_type_details_once( return Err(format!("custom type details are not supported for {:?} connections", config.db_type)); } { - let connections = state.connections.read().await; - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { let timeout_duration = agent_metadata_timeout(db_config.as_ref()); - drop(connections); let mut client = client.lock().await; return client .get_custom_type_details::(database, schema, name, timeout_duration) @@ -8858,11 +8806,10 @@ async fn get_object_source_once( let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let db_config = connection_config(state, connection_id).await; let source = { - let connections = state.connections.read().await; - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { let config = config.clone(); let session = session.clone(); - drop(connections); if let Some(sql) = gaussdb_m_view_object_source_sql(config.as_ref(), database, schema, name, &object_type) { let result: db::QueryResult = session .invoke_with_timeout( @@ -8901,8 +8848,7 @@ async fn get_object_source_once( .await?; return Ok(result); } - if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { - drop(connections); + if let Some(client) = extract_pool!(pool_handle.as_ref(), SqlServer) { let mut client = lock_sqlserver_metadata_client(&client).await?; let result = db::sqlserver::execute_query(&mut client, &sqlserver_object_source_sql(schema, name, &object_type)) @@ -8913,8 +8859,7 @@ async fn get_object_source_once( state.remove_pool_by_key(&pool_key).await; } first_string_cell(result?)? - } else if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - drop(connections); + } else if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { if uses_oracle_metadata_object_source(db_config.as_ref(), &object_type) { oracle_agent_object_source( client, @@ -8940,7 +8885,7 @@ async fn get_object_source_once( return Ok(result); } } else { - match connections.get(&pool_key).ok_or("Pool not found")? { + match pool_handle.as_ref().ok_or("Pool not found")? { PoolKind::Mysql(pool, _) => { mysql_object_source(pool, mysql_table_metadata_catalog(database, schema), name, &object_type) .await? @@ -8974,7 +8919,6 @@ async fn get_object_source_once( let schema = schema.to_string(); let name = name.to_string(); let object_type = object_type.clone(); - drop(connections); client.get_object_source(database, schema, name, object_type).await? } PoolKind::Rqlite(client) => { diff --git a/crates/dbx-core/src/sql_file_import.rs b/crates/dbx-core/src/sql_file_import.rs index 98be138045..708181cb75 100644 --- a/crates/dbx-core/src/sql_file_import.rs +++ b/crates/dbx-core/src/sql_file_import.rs @@ -161,8 +161,8 @@ impl MySqlSqlFileExecutor { let database = (!database.is_empty()).then_some(database); let pool_key = state.get_or_create_pool_for_session(&request.connection_id, database, None).await?; let (db_type, driver_profile, bare) = { - let connections = state.connections.read().await; - let Some(PoolKind::Mysql(_, mode)) = connections.get(&pool_key) else { + let pool_handle = state.pool_handle(&pool_key).await; + let Some(PoolKind::Mysql(_, mode)) = pool_handle.as_ref() else { return Ok(None); }; (Some(target.db_type), target.driver_profile.as_deref(), *mode == crate::connection::MysqlMode::Bare) @@ -303,8 +303,8 @@ impl MySqlSqlFileExecutor { let database = (!database.is_empty()).then_some(database); self.pool_key = state.get_or_create_pool_for_session(&self.connection_id, database, None).await?; let pool = { - let connections = state.connections.read().await; - match connections.get(&self.pool_key) { + let pool_handle = state.pool_handle(&self.pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Mysql(pool, _)) => pool.clone(), Some(_) => return Err("SQL file import expected a MySQL-compatible pooled connection".to_string()), None => return Err("Connection not found".to_string()), @@ -2276,10 +2276,10 @@ mod tests { let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap(); state - .connections - .write() - .await - .insert("gauss-stream".to_string(), crate::connection::PoolKind::Sqlite(pool.clone())); + .update_connection_pools(|connections| { + connections.insert("gauss-stream".to_string(), crate::connection::PoolKind::Sqlite(pool.clone())); + }) + .await; let path = temporary_sql_file( b"CREATE TABLE side_effects(value INTEGER);\nINSERT INTO missing_before_control VALUES (1);\nINSERT INTO side_effects VALUES (1);\n\\set ON_ERROR_STOP on\nINSERT INTO missing_after_control VALUES (1);\nINSERT INTO side_effects VALUES (2);", diff --git a/crates/dbx-core/src/table_export.rs b/crates/dbx-core/src/table_export.rs index 180e92d891..a75c963e27 100644 --- a/crates/dbx-core/src/table_export.rs +++ b/crates/dbx-core/src/table_export.rs @@ -519,8 +519,8 @@ enum TableExportCursorSession { } async fn table_export_cursor_kind(state: &AppState, pool_key: &str) -> Option { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Agent(_)) => Some(TableExportCursorKind::Agent), Some(PoolKind::ExternalDriver { .. }) => Some(TableExportCursorKind::ExternalDriver), _ => None, @@ -677,12 +677,11 @@ async fn fetch_table_export_batch( fetch_size: Some(active_batch_size), timeout_secs: (query_timeout > 0).then_some(query_timeout), }; - let connections = state.connections.read().await; - let Some(PoolKind::Agent(client)) = connections.get(pool_key) else { + let pool_handle = state.pool_handle(pool_key).await; + let Some(PoolKind::Agent(client)) = pool_handle.as_ref() else { return Err("Agent table read requires an agent connection".to_string()); }; let client = client.clone(); - drop(connections); let mut client = client.lock().await; match client.start_table_read::(params).await { Ok(result) => { @@ -731,12 +730,11 @@ async fn fetch_table_export_batch( if let Some(session) = cursor_session.clone() { return match session { TableExportCursorSession::Agent(session_id) => { - let connections = state.connections.read().await; - let Some(PoolKind::Agent(client)) = connections.get(pool_key) else { + let pool_handle = state.pool_handle(pool_key).await; + let Some(PoolKind::Agent(client)) = pool_handle.as_ref() else { return Err("Table read session requires an agent connection".to_string()); }; let client = client.clone(); - drop(connections); let mut client = client.lock().await; match client.fetch_table_read_page::(&session_id, active_batch_size).await { Ok(result) => { @@ -848,12 +846,11 @@ async fn close_table_export_cursor_if_open( }; match session { TableExportCursorSession::Agent(session_id) => { - let connections = state.connections.read().await; - let Some(PoolKind::Agent(client)) = connections.get(pool_key) else { + let pool_handle = state.pool_handle(pool_key).await; + let Some(PoolKind::Agent(client)) = pool_handle.as_ref() else { return; }; let client = client.clone(); - drop(connections); let mut client = client.lock().await; let _ = client.close_table_read_session::(&session_id).await; } @@ -893,12 +890,11 @@ async fn stream_native_table_rows( cancel_token: CancellationToken, on_row: impl FnMut(&[Value]) -> Result<(), String>, ) -> Result { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Mysql(pool, mode)) => { let pool = pool.clone(); let bare = *mode == MysqlMode::Bare; - drop(connections); crate::db::mysql::stream_query_rows( &pool, sql, @@ -913,13 +909,11 @@ async fn stream_native_table_rows( } Some(PoolKind::Postgres(pool)) => { let pool = pool.clone(); - drop(connections); crate::db::postgres::stream_query_rows(&pool, sql, row_limit, cancelled, on_row).await?; Ok(true) } Some(PoolKind::SqlServer(client)) => { let client = client.clone(); - drop(connections); let mut on_row = on_row; let mut client = client.lock().await; crate::db::sqlserver::stream_first_result_set(&mut client, sql, row_limit, Some(cancel_token), |item| { @@ -2188,10 +2182,14 @@ mod tests { let export_id = format!("export-{}", uuid::Uuid::new_v4()); let pool_key = format!("{}:session:{}", config.id, table_export_client_session_id(&export_id).replace(':', "_")); - state.connections.write().await.insert( - pool_key, - PoolKind::ExternalDriver { driver_id: "jdbc".to_string(), config: Arc::new(config), session }, - ); + state + .update_connection_pools(|connections| { + connections.insert( + pool_key, + PoolKind::ExternalDriver { driver_id: "jdbc".to_string(), config: Arc::new(config), session }, + ); + }) + .await; let output = dir.join("export.csv"); let request = TableExportRequest { @@ -2942,7 +2940,7 @@ mod tests { ); assert_eq!(progress.last().and_then(|event| event.total_rows), Some(3)); assert!(matches!(progress.last().map(|event| &event.status), Some(ExportStatus::Done))); - assert!(fixture.state.connections.read().await.is_empty()); + assert!(fixture.state.with_connection_pools(|pools| pools.is_empty()).await); cleanup_external_driver_export_fixture(fixture); } @@ -2999,7 +2997,7 @@ mod tests { run_external_driver_export(&fixture).await.expect("row-limited JDBC export should succeed"); assert_eq!(std::fs::read_to_string(&fixture.calls).unwrap(), "executeQueryPage\ncloseQuerySession\n"); - assert!(fixture.state.connections.read().await.is_empty()); + assert!(fixture.state.with_connection_pools(|pools| pools.is_empty()).await); cleanup_external_driver_export_fixture(fixture); } @@ -3052,7 +3050,7 @@ mod tests { std::fs::read_to_string(&fixture.calls).unwrap(), "executeQueryPage\nfetchQueryPage\ncloseQuerySession\n" ); - assert!(fixture.state.connections.read().await.is_empty()); + assert!(fixture.state.with_connection_pools(|pools| pools.is_empty()).await); cleanup_external_driver_export_fixture(fixture); } @@ -3087,7 +3085,7 @@ mod tests { assert!(cancel_requested_at.elapsed() < Duration::from_secs(2)); assert!(matches!(progress.last().map(|event| &event.status), Some(ExportStatus::Cancelled))); - assert!(fixture.state.connections.read().await.is_empty()); + assert!(fixture.state.with_connection_pools(|pools| pools.is_empty()).await); clear_export_cancelled(&fixture.request.export_id).await; cleanup_external_driver_export_fixture(fixture); } @@ -3126,7 +3124,7 @@ mod tests { assert!(cancel_requested_at.elapsed() < Duration::from_secs(2)); assert!(matches!(progress.last().map(|event| &event.status), Some(ExportStatus::Cancelled))); - assert!(fixture.state.connections.read().await.is_empty()); + assert!(fixture.state.with_connection_pools(|pools| pools.is_empty()).await); clear_export_cancelled(&fixture.request.export_id).await; cleanup_external_driver_export_fixture(fixture); } diff --git a/crates/dbx-core/src/table_import.rs b/crates/dbx-core/src/table_import.rs index 6a5b3f9934..1bfbc74dab 100644 --- a/crates/dbx-core/src/table_import.rs +++ b/crates/dbx-core/src/table_import.rs @@ -4541,8 +4541,8 @@ async fn execute_postgres_copy_batch( statement_count: &mut usize, ) -> Result<(), String> { let pool = { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Postgres(pool)) => pool.clone(), _ => return Err("PostgreSQL pool not found for COPY import".to_string()), } @@ -5333,8 +5333,8 @@ async fn kingbase_oracle_compatibility_mode(state: &AppState, pool_key: &str, db return false; } let client = { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Agent(client)) => client.clone(), _ => return false, } @@ -5350,8 +5350,8 @@ async fn kingbase_oracle_compatibility_mode(state: &AppState, pool_key: &str, db async fn mysql_import_sql_hard_limit(state: &AppState, pool_key: &str) -> Option { let pool = { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Mysql(pool, _)) => pool.clone(), _ => return None, } @@ -5480,8 +5480,8 @@ async fn sqlserver_bulk_import_plan_for_pool( } let import_plan = import_plan?; let client = { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::SqlServer(client)) => client.clone(), _ => return None, } @@ -5722,8 +5722,8 @@ async fn execute_sqlserver_bulk_rows_batch( .await .map_err(ImportRowsBatchError::before_write)?; let client = { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::SqlServer(client)) => client.clone(), _ => { return Err(ImportRowsBatchError::before_write( @@ -8735,7 +8735,11 @@ mod tests { ) .await .unwrap(); - state.connections.write().await.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + state + .update_connection_pools(|connections| { + connections.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + }) + .await; let config: ConnectionConfig = serde_json::from_value(serde_json::json!({ "id": connection_id, "name": "XLSX truncate tail test", @@ -8823,7 +8827,11 @@ mod tests { ) .await .unwrap(); - state.connections.write().await.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + state + .update_connection_pools(|connections| { + connections.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + }) + .await; let config: ConnectionConfig = serde_json::from_value(serde_json::json!({ "id": connection_id, "name": "Cancel XLSX validation test", @@ -8907,7 +8915,11 @@ mod tests { ) .await .unwrap(); - state.connections.write().await.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + state + .update_connection_pools(|connections| { + connections.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + }) + .await; let config: ConnectionConfig = serde_json::from_value(serde_json::json!({ "id": connection_id, "name": "Cancel truncate first batch test", @@ -10475,7 +10487,11 @@ mod tests { let pool_key = "sqlserver-cleanup-failure"; let database_path = dir.path().join("target.db"); let sqlite = crate::db::sqlite::connect_path_create_if_missing(database_path.to_str().unwrap()).await.unwrap(); - state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(sqlite)); + state + .update_connection_pools(|connections| { + connections.insert(pool_key.to_string(), PoolKind::Sqlite(sqlite)); + }) + .await; invalidate_sqlserver_pool_after_staging_cleanup_failure( &state, @@ -10486,7 +10502,7 @@ mod tests { ) .await; - assert!(!state.connections.read().await.contains_key(pool_key)); + assert!(!state.pool_handle(pool_key).await.is_some()); } #[test] @@ -10520,7 +10536,11 @@ mod tests { let database_path = dir.path().join("target.db"); let sqlite = crate::db::sqlite::connect_path_create_if_missing(database_path.to_str().unwrap()).await.unwrap(); crate::db::sqlite::execute_query(&sqlite, "CREATE TABLE items (payload TEXT)").await.unwrap(); - state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(sqlite.clone())); + state + .update_connection_pools(|connections| { + connections.insert(pool_key.to_string(), PoolKind::Sqlite(sqlite.clone())); + }) + .await; let rows = vec![vec![serde_json::json!("a".repeat(300 * 1024))], vec![serde_json::json!("b".repeat(300 * 1024))]]; @@ -10736,7 +10756,11 @@ mod tests { let sqlite = crate::db::sqlite::connect_path_create_if_missing(database_path.to_str().unwrap()).await.unwrap(); crate::db::sqlite::execute_query(&sqlite, "CREATE TABLE items (id INTEGER PRIMARY KEY)").await.unwrap(); - state.connections.write().await.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + state + .update_connection_pools(|connections| { + connections.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + }) + .await; Self { _dir: dir, state, @@ -10848,7 +10872,11 @@ mod tests { let database_path = dir.path().join("target.db"); let sqlite = crate::db::sqlite::connect_path_create_if_missing(database_path.to_str().unwrap()).await.unwrap(); crate::db::sqlite::execute_query(&sqlite, "CREATE TABLE items (id INTEGER, name TEXT)").await.unwrap(); - state.connections.write().await.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + state + .update_connection_pools(|connections| { + connections.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + }) + .await; let config: ConnectionConfig = serde_json::from_value(serde_json::json!({ "id": connection_id, "name": "SQLite delimited append test", diff --git a/crates/dbx-core/src/table_structure_sql/sqlite_rebuild.rs b/crates/dbx-core/src/table_structure_sql/sqlite_rebuild.rs index 598e18dade..022806f67a 100644 --- a/crates/dbx-core/src/table_structure_sql/sqlite_rebuild.rs +++ b/crates/dbx-core/src/table_structure_sql/sqlite_rebuild.rs @@ -85,8 +85,8 @@ async fn native_sqlite_pool( ) -> Result<(String, db::sqlite::SqliteHandle), String> { let database = (!database.trim().is_empty()).then_some(database); let pool_key = state.get_or_create_pool(connection_id, database).await?; - let connections = state.connections.read().await; - match connections.get(&pool_key) { + let pool_handle = state.pool_handle(&pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Sqlite(pool)) => Ok((pool_key, pool.clone())), Some(_) => Err("SQLite table rebuild is only available for native SQLite connections.".to_string()), None => Err("SQLite connection pool not found.".to_string()), diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index b186d75f29..9b0e48b980 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -5094,35 +5094,30 @@ async fn execute_on_pool_once( ) -> Result { // Read-only check: block transfer operations in readonly mode crate::query::check_read_only_for_connection(state, pool_key, sql).await?; - let connections = state.connections.read().await; - let pool = connections.get(pool_key).ok_or("Connection not found")?; + let pool_handle = state.pool_handle(pool_key).await; + let pool = pool_handle.as_ref().ok_or("Connection not found")?; match pool { PoolKind::Mysql(p, mode) => { let p = p.clone(); let bare = *mode == crate::connection::MysqlMode::Bare; - drop(connections); db::mysql::execute_query_with_max_rows(&p, sql, bare, max_rows, Default::default()).await } PoolKind::Postgres(p) => { let p = p.clone(); - drop(connections); db::postgres::execute_query_with_max_rows(&p, sql, max_rows).await } PoolKind::Sqlite(p) => { let p = p.clone(); - drop(connections); db::sqlite::execute_query_with_max_rows(&p, sql, max_rows).await } PoolKind::ClickHouse(client) => { let client = client.clone(); let database = database_from_pool_key(pool_key).unwrap_or("default").to_string(); - drop(connections); db::clickhouse_driver::execute_query_with_max_rows(&client, &database, sql, max_rows).await } PoolKind::SqlServer(client) => { let client = client.clone(); - drop(connections); let mut client = client.lock().await; let result = db::sqlserver::execute_query_with_max_rows(&mut client, sql, max_rows).await; drop(client); @@ -5132,7 +5127,6 @@ async fn execute_on_pool_once( let client = client.clone(); let database = database_from_pool_key(pool_key).map(str::to_string); let sql = sql.to_string(); - drop(connections); let mut client = client.lock().await; let params = agent_execute_query_params( &sql, @@ -5159,7 +5153,6 @@ async fn execute_on_pool_once( PoolKind::DuckDbWorker(client) => { let client = client.clone(); let sql = sql.to_string(); - drop(connections); client.execute(None, sql, max_rows, None, None).await } _ => Err("Unsupported database type for transfer".to_string()), @@ -5223,63 +5216,56 @@ pub async fn get_columns_for_transfer( table: &str, catalog: Option<&str>, ) -> Result, String> { - let connections = state.connections.read().await; + let pool_handle = state.pool_handle(pool_key).await; #[cfg(feature = "duckdb-sidecar")] - if let Some(PoolKind::DuckDbWorker(client)) = connections.get(pool_key) { + if let Some(PoolKind::DuckDbWorker(client)) = pool_handle.as_ref() { let client = client.clone(); let database = database.to_string(); let schema = schema.to_string(); let table = table.to_string(); - drop(connections); return client.list_columns(database, schema, table).await; } - if let Some(PoolKind::ClickHouse(client)) = connections.get(pool_key) { + if let Some(PoolKind::ClickHouse(client)) = pool_handle.as_ref() { let client = client.clone(); let database = database.to_string(); let table = table.to_string(); - drop(connections); return db::clickhouse_driver::get_columns(&client, &database, &table).await; } - if let Some(PoolKind::SqlServer(client)) = connections.get(pool_key) { + if let Some(PoolKind::SqlServer(client)) = pool_handle.as_ref() { let client = client.clone(); let schema = schema.to_string(); let table = table.to_string(); - drop(connections); let mut client = client.lock().await; return db::sqlserver::get_columns(&mut client, &schema, &table).await; } - if let Some(PoolKind::InfluxDb(client)) = connections.get(pool_key) { + if let Some(PoolKind::InfluxDb(client)) = pool_handle.as_ref() { let client = client.clone(); let database = database.to_string(); let table = table.to_string(); - drop(connections); return db::influxdb_driver::get_columns(&client, &database, &table).await; } - if let Some(PoolKind::InfluxDb3(client)) = connections.get(pool_key) { + if let Some(PoolKind::InfluxDb3(client)) = pool_handle.as_ref() { let client = client.clone(); let database = database.to_string(); let table = table.to_string(); - drop(connections); return db::influxdb3_driver::get_columns(&client, &database, &table).await; } - if let Some(PoolKind::Agent(client)) = connections.get(pool_key) { + if let Some(PoolKind::Agent(client)) = pool_handle.as_ref() { let client = client.clone(); let database = database.to_string(); let schema = schema.to_string(); let table = table.to_string(); - drop(connections); let mut client = client.lock().await; return client.get_columns(&database, &schema, &table, None).await; } - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(pool_key) { + if let Some(PoolKind::ExternalDriver { config, session, .. }) = pool_handle.as_ref() { let config = config.clone(); let session = session.clone(); let database = database.to_string(); let schema = schema.to_string(); let table = table.to_string(); - drop(connections); return session .invoke_with_timeout( "getColumns", @@ -5293,14 +5279,13 @@ pub async fn get_columns_for_transfer( ) .await; } - let pool = connections.get(pool_key).ok_or("Pool not found")?; + let pool = pool_handle.as_ref().ok_or("Pool not found")?; let schema = schema.to_string(); let table = table.to_string(); match pool { PoolKind::Mysql(p, _) => { let p = p.clone(); let catalog = normalize_external_catalog_name(catalog).map(str::to_string); - drop(connections); if let Some(catalog) = catalog { // Use 3-part qualified column lookup for Doris/StarRocks external catalogs db::doris::get_catalog_columns(&p, &catalog, &schema, &table).await @@ -5310,12 +5295,10 @@ pub async fn get_columns_for_transfer( } PoolKind::Postgres(p) => { let p = p.clone(); - drop(connections); db::postgres::get_columns(&p, &schema, &table).await } PoolKind::Sqlite(p) => { let p = p.clone(); - drop(connections); db::sqlite::get_columns(&p, &schema, &table).await } _ => Err("Unsupported database type".to_string()), @@ -5329,21 +5312,19 @@ async fn get_postgres_indexes_for_transfer( schema: &str, table: &str, ) -> Result, String> { - let connections = state.connections.read().await; - if let Some(PoolKind::Agent(client)) = connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + if let Some(PoolKind::Agent(client)) = pool_handle.as_ref() { let client = client.clone(); let database = database.to_string(); let schema = schema.to_string(); let table = table.to_string(); - drop(connections); let mut client = client.lock().await; return client.list_indexes(&database, &schema, &table, None).await; } - let Some(PoolKind::Postgres(pool)) = connections.get(pool_key) else { + let Some(PoolKind::Postgres(pool)) = pool_handle.as_ref() else { return Err("PostgreSQL pool not found".to_string()); }; let pool = pool.clone(); - drop(connections); db::postgres::list_indexes(&pool, schema, table).await } @@ -5354,21 +5335,19 @@ async fn get_postgres_foreign_keys_for_transfer( schema: &str, table: &str, ) -> Result, String> { - let connections = state.connections.read().await; - if let Some(PoolKind::Agent(client)) = connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + if let Some(PoolKind::Agent(client)) = pool_handle.as_ref() { let client = client.clone(); let database = database.to_string(); let schema = schema.to_string(); let table = table.to_string(); - drop(connections); let mut client = client.lock().await; return client.list_foreign_keys(&database, &schema, &table, None).await; } - let Some(PoolKind::Postgres(pool)) = connections.get(pool_key) else { + let Some(PoolKind::Postgres(pool)) = pool_handle.as_ref() else { return Err("PostgreSQL pool not found".to_string()); }; let pool = pool.clone(); - drop(connections); db::postgres::list_foreign_keys(&pool, schema, table).await } @@ -5383,8 +5362,8 @@ async fn get_postgres_owned_sequences_for_transfer( } let pool = { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Postgres(pool)) => pool.clone(), _ => return Ok(Vec::new()), } @@ -5430,8 +5409,8 @@ async fn get_postgres_sequence_snapshots_for_transfer( schema: &str, ) -> Result, String> { let pool = { - let connections = state.connections.read().await; - match connections.get(pool_key) { + let pool_handle = state.pool_handle(pool_key).await; + match pool_handle.as_ref() { Some(PoolKind::Postgres(pool)) => pool.clone(), _ => return Ok(Vec::new()), } @@ -6763,8 +6742,8 @@ pub async fn sort_tables_by_fk_dependency_with_foreign_keys( let postgres_pool = if db_type == DatabaseType::Postgres { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; { - let connections = state.connections.read().await; - native_postgres_dependency_pool(connections.get(&pool_key)) + let pool_handle = state.pool_handle(&pool_key).await; + native_postgres_dependency_pool(pool_handle.as_ref()) } } else { None @@ -7184,12 +7163,11 @@ async fn fetch_hive_server_transfer_batch( let configs = state.configs.read().await; configs.get(&request.source_connection_id).map(|config| config.query_timeout_secs).unwrap_or(0) }; - let connections = state.connections.read().await; - let Some(PoolKind::Agent(client)) = connections.get(pool_key) else { + let pool_handle = state.pool_handle(pool_key).await; + let Some(PoolKind::Agent(client)) = pool_handle.as_ref() else { return Err("Impala transfer requires an Agent connection".to_string()); }; let client = client.clone(); - drop(connections); let mut client = client.lock().await; let result = if cursor.started { @@ -7226,12 +7204,11 @@ async fn close_hive_server_transfer_cursor(state: &AppState, pool_key: &str, cur let Some(session_id) = cursor.session_id.take() else { return; }; - let connections = state.connections.read().await; - let Some(PoolKind::Agent(client)) = connections.get(pool_key) else { + let pool_handle = state.pool_handle(pool_key).await; + let Some(PoolKind::Agent(client)) = pool_handle.as_ref() else { return; }; let client = client.clone(); - drop(connections); let mut client = client.lock().await; if let Err(error) = client.close_table_read_session::(&session_id).await { log::warn!("[transfer] failed to close Impala transfer cursor: {error}"); @@ -7433,10 +7410,11 @@ where // SHOW CREATE TABLE catalog.database.table using the // existing source pool (bare MySQL — addresses any catalog). let pool = { - let connections = state.connections.read().await; - let pool = - connections.get(source_pool_key).ok_or_else(|| "Source pool not found".to_string())?; - let PoolKind::Mysql(p, _) = pool else { + let pool = state + .pool_handle(source_pool_key) + .await + .ok_or_else(|| "Source pool not found".to_string())?; + let PoolKind::Mysql(p, _) = &pool else { return Err("Source pool must be MySQL-family for catalog DDL".to_string()); }; p.clone() @@ -8637,10 +8615,14 @@ mod tests { #[tokio::test] async fn postgres_transfer_metadata_routes_agent_pools() { let (state, dir) = test_app_state().await; - state.connections.write().await.insert( - "source:source_db".to_string(), - PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub()), - ); + state + .update_connection_pools(|connections| { + connections.insert( + "source:source_db".to_string(), + PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub()), + ); + }) + .await; let index_error = get_postgres_indexes_for_transfer(&state, "source:source_db", "source_db", "source_schema", "items") diff --git a/crates/dbx-core/tests/live_mongodb_agent_tools.rs b/crates/dbx-core/tests/live_mongodb_agent_tools.rs index 2d8caf5f20..b30b2b7efe 100644 --- a/crates/dbx-core/tests/live_mongodb_agent_tools.rs +++ b/crates/dbx-core/tests/live_mongodb_agent_tools.rs @@ -1,10 +1,13 @@ use std::sync::{Arc, Mutex}; +use std::time::Duration; use dbx_core::agent_events::{ToolCall, ToolResult}; use dbx_core::agent_tools::{execute_tool, AgentSqlPermissions}; use dbx_core::connection::{AppState, PoolKind}; use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; +use dbx_core::mongo_ops::mongo_create_index_core; use dbx_core::storage::Storage; +use mongodb::bson::{doc, Bson}; use mongodb::event::{command::CommandEvent, EventHandler}; use mongodb::options::ClientOptions; use mongodb::Client; @@ -107,7 +110,11 @@ async fn mongodb_agent_enforces_limits_and_find_skips_total_count() { let client = Client::with_options(options).unwrap(); state.configs.write().await.insert(config.id.clone(), config.clone()); - state.connections.write().await.insert(config.id.clone(), PoolKind::MongoDb(client)); + state + .update_connection_pools(|connections| { + connections.insert(config.id.clone(), PoolKind::MongoDb(client)); + }) + .await; let zero_limit = call_agent_tool_with_limit(&state, &config.id, &database, "db.products.find({}).limit(0)", Some(0)).await; @@ -128,3 +135,126 @@ async fn mongodb_agent_enforces_limits_and_find_skips_total_count() { assert!(!distinct.is_error, "{}", distinct.content); assert!(distinct.content.contains("(7 rows,"), "{}", distinct.content); } + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MONGODB_42_URL pointing at MongoDB 4.2 with enableTestCommands=1"] +async fn stalled_mongodb_index_build_does_not_block_another_connection() { + let uri = std::env::var("DBX_LIVE_MONGODB_42_URL").expect("DBX_LIVE_MONGODB_42_URL"); + let directory = tempfile::tempdir().unwrap(); + let storage = Storage::open(&directory.path().join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let primary_id = "live-mongodb-42-index"; + let database = "dbx_issue_7720_primary"; + let collection = "records"; + let index_name = "issue_7720_pool_registry"; + + let client = Client::with_uri_str(&uri).await.expect("MongoDB 4.2 client"); + let primary_config: ConnectionConfig = serde_json::from_value(serde_json::json!({ + "id": primary_id, + "name": "MongoDB 4.2 stalled index", + "db_type": DatabaseType::MongoDb, + "host": "127.0.0.1", + "port": 17720, + "username": "", + "password": "", + "database": database, + "connection_string": uri, + "driver_profile": "mongodb-native", + "connect_timeout_secs": 5, + "query_timeout_secs": 30, + "idle_timeout_secs": 60, + "keepalive_interval_secs": 0 + })) + .unwrap(); + state.configs.write().await.insert(primary_id.to_string(), primary_config); + state + .update_connection_pools(|connections| { + connections.insert(primary_id.to_string(), PoolKind::MongoDb(client.clone())); + }) + .await; + + let database_handle = client.database(database); + let _ = database_handle.run_command(doc! { "dropIndexes": collection, "index": index_name }).await; + let enabled = client + .database("admin") + .run_command(doc! { "configureFailPoint": "hangAfterStartingIndexBuild", "mode": "alwaysOn" }) + .await + .expect("enable index-build failpoint"); + let entered = match enabled.get("count") { + Some(Bson::Int32(value)) => i64::from(*value), + Some(Bson::Int64(value)) => *value, + value => panic!("unexpected failpoint counter: {value:?}"), + }; + + let index_state = Arc::clone(&state); + let mut index_task = tokio::spawn(async move { + mongo_create_index_core( + &index_state, + primary_id, + database, + collection, + r#"{"pool_registry_probe":1}"#, + Some(&format!(r#"{{"name":"{index_name}"}}"#)), + ) + .await + }); + let admin_handle = client.database("admin"); + let wait_for_index = admin_handle.run_command(doc! { + "waitForFailPoint": "hangAfterStartingIndexBuild", + "timesEntered": entered + 1, + "maxTimeMS": 5_000_i64 + }); + let early_index_result = tokio::select! { + wait_result = wait_for_index => { + if let Err(error) = wait_result { + admin_handle + .run_command(doc! { "configureFailPoint": "hangAfterStartingIndexBuild", "mode": "off" }) + .await + .expect("disable index-build failpoint after wait failure"); + index_task.abort(); + let _ = (&mut index_task).await; + panic!("index build must reach the failpoint: {error}"); + } + None + }, + index_result = &mut index_task => Some(index_result), + }; + if let Some(index_result) = early_index_result { + admin_handle + .run_command(doc! { "configureFailPoint": "hangAfterStartingIndexBuild", "mode": "off" }) + .await + .expect("disable index-build failpoint after early completion"); + panic!("index build ended before reaching the failpoint: {index_result:?}"); + } + + let sqlite_path = directory.path().join("other-connection.db"); + std::fs::File::create(&sqlite_path).expect("create SQLite fixture"); + let other_config: ConnectionConfig = serde_json::from_value(serde_json::json!({ + "id": "other-connection", + "name": "Other connection", + "db_type": DatabaseType::Sqlite, + "host": sqlite_path.to_string_lossy(), + "port": 0, + "username": "", + "password": "", + "connect_timeout_secs": 5, + "query_timeout_secs": 30, + "idle_timeout_secs": 60, + "keepalive_interval_secs": 0 + })) + .unwrap(); + state.configs.write().await.insert(other_config.id.clone(), other_config); + let other_connection = + tokio::time::timeout(Duration::from_secs(2), state.get_or_create_pool("other-connection", None)).await; + + client + .database("admin") + .run_command(doc! { "configureFailPoint": "hangAfterStartingIndexBuild", "mode": "off" }) + .await + .expect("disable index-build failpoint"); + let index_result = index_task.await.expect("index task join"); + let _ = database_handle.run_command(doc! { "dropIndexes": collection, "index": index_name }).await; + + assert!(other_connection.expect("other connection must not wait on the registry lock").is_ok()); + assert_eq!(index_result.as_deref(), Ok(index_name)); +} diff --git a/crates/dbx-core/tests/live_postgres_transfer.rs b/crates/dbx-core/tests/live_postgres_transfer.rs index 169d78bacb..f6f24cb0ec 100644 --- a/crates/dbx-core/tests/live_postgres_transfer.rs +++ b/crates/dbx-core/tests/live_postgres_transfer.rs @@ -199,8 +199,12 @@ async fn live_postgres_transfer_upserts_generated_always_identity_values() { let target_connection_id = "live-always-target"; let source_pool_key = format!("{source_connection_id}:{source_database}"); let target_pool_key = format!("{target_connection_id}:{target_database}"); - state.connections.write().await.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); - state.connections.write().await.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + state + .update_connection_pools(|connections| { + connections.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + connections.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + }) + .await; state .configs .write() @@ -363,8 +367,12 @@ async fn live_postgres_structure_only_preserves_table_indexes() { let target_connection_id = "live-structure-only-target"; let source_pool_key = format!("{source_connection_id}:{source_database}"); let target_pool_key = format!("{target_connection_id}:{target_database}"); - state.connections.write().await.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); - state.connections.write().await.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + state + .update_connection_pools(|connections| { + connections.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + connections.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + }) + .await; state .configs .write() @@ -608,8 +616,12 @@ async fn live_postgres_transfer_preserves_data_and_schema_objects() { let source_pool_key = format!("{source_connection_id}:{source_database}"); let target_pool_key = format!("{target_connection_id}:{target_database}"); - state.connections.write().await.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); - state.connections.write().await.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + state + .update_connection_pools(|connections| { + connections.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + connections.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + }) + .await; state .configs .write() @@ -896,8 +908,12 @@ async fn live_postgres_transfer_skips_create_ddl_for_existing_target_table() { let source_pool_key = format!("{source_connection_id}:{source_database}"); let target_pool_key = format!("{target_connection_id}:{target_database}"); - state.connections.write().await.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); - state.connections.write().await.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + state + .update_connection_pools(|connections| { + connections.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + connections.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + }) + .await; state .configs .write() @@ -1013,8 +1029,12 @@ async fn live_postgres_transfer_creates_selected_sequence_before_referencing_tab let target_connection_id = "live-sequence-target"; let source_pool_key = format!("{source_connection_id}:{source_database}"); let target_pool_key = format!("{target_connection_id}:{target_database}"); - state.connections.write().await.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); - state.connections.write().await.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + state + .update_connection_pools(|connections| { + connections.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + connections.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + }) + .await; state .configs .write() diff --git a/crates/dbx-core/tests/live_redis_pool_registry.rs b/crates/dbx-core/tests/live_redis_pool_registry.rs new file mode 100644 index 0000000000..d83b10bc70 --- /dev/null +++ b/crates/dbx-core/tests/live_redis_pool_registry.rs @@ -0,0 +1,65 @@ +use std::sync::Arc; +use std::time::Duration; + +use dbx_core::connection::AppState; +use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; +use dbx_core::redis_ops::redis_execute_command_core; +use dbx_core::storage::Storage; + +#[tokio::test] +#[ignore = "requires DBX_LIVE_REDIS_HOST and DBX_LIVE_REDIS_PORT"] +async fn blocking_redis_command_does_not_block_another_connection() { + let host = std::env::var("DBX_LIVE_REDIS_HOST").expect("DBX_LIVE_REDIS_HOST"); + let port = std::env::var("DBX_LIVE_REDIS_PORT").expect("DBX_LIVE_REDIS_PORT").parse::().expect("Redis port"); + let directory = tempfile::tempdir().unwrap(); + let storage = Storage::open(&directory.path().join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let redis_config: ConnectionConfig = serde_json::from_value(serde_json::json!({ + "id": "live-redis-blocking", + "name": "Live blocking Redis", + "db_type": DatabaseType::Redis, + "host": host, + "port": port, + "username": "", + "password": "", + "database": null, + "connect_timeout_secs": 5, + "query_timeout_secs": 5, + "idle_timeout_secs": 60, + "keepalive_interval_secs": 0 + })) + .unwrap(); + state.configs.write().await.insert(redis_config.id.clone(), redis_config.clone()); + state.get_or_create_pool(&redis_config.id, None).await.expect("open Redis connection"); + + let redis_state = Arc::clone(&state); + let redis_id = redis_config.id.clone(); + let blocking = tokio::spawn(async move { + redis_execute_command_core(&redis_state, &redis_id, 0, "BLPOP dbx:issue:7720:missing 1", true).await + }); + tokio::time::sleep(Duration::from_millis(100)).await; + + let sqlite_path = directory.path().join("other-connection.db"); + std::fs::File::create(&sqlite_path).expect("create SQLite fixture"); + let sqlite_config: ConnectionConfig = serde_json::from_value(serde_json::json!({ + "id": "other-connection", + "name": "Other connection", + "db_type": DatabaseType::Sqlite, + "host": sqlite_path.to_string_lossy(), + "port": 0, + "username": "", + "password": "", + "database": null, + "connect_timeout_secs": 5, + "query_timeout_secs": 5, + "idle_timeout_secs": 60, + "keepalive_interval_secs": 0 + })) + .unwrap(); + state.configs.write().await.insert(sqlite_config.id.clone(), sqlite_config); + let other_connection = + tokio::time::timeout(Duration::from_millis(300), state.get_or_create_pool("other-connection", None)).await; + + assert!(other_connection.expect("other connection must not wait on Redis").is_ok()); + assert!(blocking.await.expect("Redis task join").is_ok()); +} diff --git a/crates/dbx-core/tests/live_sqlserver_completion.rs b/crates/dbx-core/tests/live_sqlserver_completion.rs index 052fb841b7..db45dcccf9 100644 --- a/crates/dbx-core/tests/live_sqlserver_completion.rs +++ b/crates/dbx-core/tests/live_sqlserver_completion.rs @@ -1167,10 +1167,11 @@ async fn live_sqlserver_query_result_export_streams_cte_query_to_csv() { let connection_id = "live-sqlserver-export"; let pool_key = format!("{connection_id}:{database}"); state - .connections - .write() - .await - .insert(pool_key, PoolKind::SqlServer(std::sync::Arc::new(tokio::sync::Mutex::new(export_client)))); + .update_connection_pools(|connections| { + connections + .insert(pool_key, PoolKind::SqlServer(std::sync::Arc::new(tokio::sync::Mutex::new(export_client)))); + }) + .await; let file_path = dir.join("result.csv"); let sql = format!( @@ -1258,10 +1259,14 @@ async fn live_sqlserver_sql_file_import_executes_go_batches() { config.username = user; config.password = password; state.configs.write().await.insert(connection_id.to_string(), config); - state.connections.write().await.insert( - format!("{connection_id}:{database}"), - PoolKind::SqlServer(std::sync::Arc::new(tokio::sync::Mutex::new(client))), - ); + state + .update_connection_pools(|connections| { + connections.insert( + format!("{connection_id}:{database}"), + PoolKind::SqlServer(std::sync::Arc::new(tokio::sync::Mutex::new(client))), + ); + }) + .await; let script = format!( "CREATE TABLE [dbo].[{table}] (id INT NOT NULL);\n\ @@ -1292,8 +1297,7 @@ async fn live_sqlserver_sql_file_import_executes_go_batches() { .expect("execute SQL Server file with GO batches"); let pool_key = format!("{connection_id}:{database}"); - let connections = state.connections.read().await; - let PoolKind::SqlServer(client) = connections.get(&pool_key).expect("SQL Server pool") else { + let PoolKind::SqlServer(client) = state.pool_handle(&pool_key).await.expect("SQL Server pool") else { panic!("expected SQL Server pool"); }; let mut client = client.lock().await; @@ -1301,7 +1305,6 @@ async fn live_sqlserver_sql_file_import_executes_go_batches() { let cleanup = format!("DROP PROCEDURE [dbo].[{procedure}]; DROP TABLE [dbo].[{table}];"); let _ = dbx_core::db::sqlserver::execute_batch(&mut client, &cleanup).await; drop(client); - drop(connections); let _ = std::fs::remove_dir_all(&dir); assert!(done_seen.load(Ordering::Relaxed)); diff --git a/crates/dbx-core/tests/live_sqlserver_query_result_export.rs b/crates/dbx-core/tests/live_sqlserver_query_result_export.rs index 456b9da7a9..0245c6962f 100644 --- a/crates/dbx-core/tests/live_sqlserver_query_result_export.rs +++ b/crates/dbx-core/tests/live_sqlserver_query_result_export.rs @@ -89,7 +89,11 @@ async fn live_sqlserver_xlsx_export_can_outlive_query_timeout_while_rows_keep_ar let connection_id = "live-sqlserver-xlsx-export"; let pool_key = format!("{connection_id}:{database}"); state.configs.write().await.insert(connection_id.to_string(), live_sqlserver_config(connection_id, &database)); - state.connections.write().await.insert(pool_key, PoolKind::SqlServer(Arc::new(tokio::sync::Mutex::new(client)))); + state + .update_connection_pools(|connections| { + connections.insert(pool_key, PoolKind::SqlServer(Arc::new(tokio::sync::Mutex::new(client)))); + }) + .await; let file_path = dir.join("result.xlsx"); let sql = "WITH numbers AS (\ diff --git a/crates/dbx-web/src/routes/connection.rs b/crates/dbx-web/src/routes/connection.rs index c9851dd642..c598431eaa 100644 --- a/crates/dbx-web/src/routes/connection.rs +++ b/crates/dbx-web/src/routes/connection.rs @@ -256,10 +256,7 @@ async fn sync_mongo_legacy_driver_fallback(state: &WebState, config: &Connection if config.db_type != DatabaseType::MongoDb { return Ok(()); } - let uses_legacy_agent = { - let connections = state.app.connections.read().await; - matches!(connections.get(&config.id), Some(PoolKind::Agent(_))) - }; + let uses_legacy_agent = matches!(state.app.pool_handle(&config.id).await, Some(PoolKind::Agent(_))); if !uses_legacy_agent { return Ok(()); } @@ -314,12 +311,9 @@ async fn run_temporary_connection_test( // runs before either a successful result or an error is returned. let result: Result = async { let success_message = if pool_result.is_ok() && config.db_type == DatabaseType::Consul { - let client = { - let connections = app.connections.read().await; - match connections.get(&temp_id) { - Some(PoolKind::Consul(client)) => Some(client.clone()), - _ => None, - } + let client = match app.pool_handle(&temp_id).await { + Some(PoolKind::Consul(client)) => Some(client), + _ => None, }; if let Some(client) = client { let configured_target = @@ -1040,7 +1034,7 @@ mod tests { assert_eq!(detailed.0.message, "Connection successful"); assert_eq!(detailed.0.database_info, None); assert!(state.app.configs.read().await.keys().all(|key| !key.starts_with("__test_"))); - assert!(state.app.connections.read().await.keys().all(|key| !key.starts_with("__test_"))); + assert!(state.app.with_connection_pools(|pools| pools.keys().all(|key| !key.starts_with("__test_"))).await); let _ = std::fs::remove_dir_all(dir); } @@ -1055,7 +1049,7 @@ mod tests { assert!(error.contains("CONSUL_AGENT_TARGET_MISMATCH"), "unexpected error: {error}"); assert!(state.app.configs.read().await.keys().all(|key| !key.starts_with("__test_"))); - assert!(state.app.connections.read().await.keys().all(|key| !key.starts_with("__test_"))); + assert!(state.app.with_connection_pools(|pools| pools.keys().all(|key| !key.starts_with("__test_"))).await); let _ = std::fs::remove_dir_all(dir); } @@ -1103,7 +1097,7 @@ mod tests { assert_eq!(result.0, "Connection successful"); assert!(state.app.configs.read().await.keys().all(|key| !key.starts_with("__test_"))); - assert!(state.app.connections.read().await.keys().all(|key| !key.starts_with("__test_"))); + assert!(state.app.with_connection_pools(|pools| pools.keys().all(|key| !key.starts_with("__test_"))).await); let cached = state.app.mq_registry.cached_connection_ids().await; assert!( cached.iter().all(|id| !id.starts_with("__test_")), @@ -1126,7 +1120,12 @@ mod tests { .await .unwrap(); state.app.configs.write().await.insert(initial.id.clone(), initial.clone()); - state.app.connections.write().await.insert(initial.id.clone(), PoolKind::Sqlite(pool.clone())); + state + .app + .update_connection_pools(|connections| { + connections.insert(initial.id.clone(), PoolKind::Sqlite(pool.clone())); + }) + .await; let mut invalid = initial.clone(); invalid.attached_databases.push(AttachedDatabaseConfig { @@ -1163,7 +1162,7 @@ mod tests { .unwrap_err(); assert!(proxy_error.message.contains("in-memory main database"), "{}", proxy_error.message); - assert!(state.app.connections.read().await.contains_key(&initial.id)); + assert!(state.app.pool_handle(&initial.id).await.is_some()); assert_eq!(state.app.configs.read().await.get(&initial.id), Some(&initial)); let retained = dbx_core::db::sqlite::execute_query(&pool, "SELECT value FROM retained;").await.unwrap(); assert_eq!(retained.rows[0][0], serde_json::json!("yes")); @@ -1359,7 +1358,12 @@ mod tests { let config = mq_config("mq-info", "http://127.0.0.1:8080"); state.app.storage.save_connections(std::slice::from_ref(&config)).await.unwrap(); state.app.configs.write().await.insert(config.id.clone(), config.clone()); - state.app.connections.write().await.insert(config.id.clone(), PoolKind::MessageQueue); + state + .app + .update_connection_pools(|connections| { + connections.insert(config.id.clone(), PoolKind::MessageQueue); + }) + .await; let database_info = DatabaseConnectionInfo { product_name: Some("Apache Pulsar".to_string()), product_version: Some("3.3.0".to_string()), @@ -1376,7 +1380,7 @@ mod tests { .await; assert!(result.is_ok()); - assert!(state.app.connections.read().await.contains_key(&config.id)); + assert!(state.app.pool_handle(&config.id).await.is_some()); assert_eq!(state.app.configs.read().await[&config.id].database_info, Some(database_info.clone())); assert_eq!(state.app.storage.load_connections().await.unwrap()[0].database_info, Some(database_info)); @@ -1389,7 +1393,12 @@ mod tests { let (state, dir) = test_web_state().await; let initial = mq_config("mq-conn", "http://127.0.0.1:8080"); state.app.configs.write().await.insert(initial.id.clone(), initial.clone()); - state.app.connections.write().await.insert(initial.id.clone(), PoolKind::MessageQueue); + state + .app + .update_connection_pools(|connections| { + connections.insert(initial.id.clone(), PoolKind::MessageQueue); + }) + .await; let first = state.app.mq_registry.get_or_build(&initial).await.unwrap().adapter; let updated = mq_config("mq-conn", "http://127.0.0.1:8081"); @@ -1415,7 +1424,7 @@ mod tests { let second = state.app.mq_registry.get_or_build(&updated).await.unwrap().adapter; assert!(!Arc::ptr_eq(&first, &second)); - assert!(!state.app.connections.read().await.contains_key(&initial.id)); + assert!(!state.app.pool_handle(&initial.id).await.is_some()); let _ = std::fs::remove_dir_all(dir); } @@ -1426,7 +1435,12 @@ mod tests { let (state, dir) = test_web_state().await; let initial = mq_config("mq-conn", "http://127.0.0.1:8080"); state.app.configs.write().await.insert(initial.id.clone(), initial.clone()); - state.app.connections.write().await.insert(initial.id.clone(), PoolKind::MessageQueue); + state + .app + .update_connection_pools(|connections| { + connections.insert(initial.id.clone(), PoolKind::MessageQueue); + }) + .await; let first = state.app.mq_registry.get_or_build(&initial).await.unwrap().adapter; let updated = mq_config("mq-conn", &spawn_pulsar_clusters_server().await); @@ -1648,7 +1662,12 @@ mod tests { let updated = mq_config("mq-conn", "http://127.0.0.1:8081"); state.app.storage.save_connections(std::slice::from_ref(&updated)).await.unwrap(); state.app.configs.write().await.insert(initial.id.clone(), initial.clone()); - state.app.connections.write().await.insert(initial.id.clone(), PoolKind::MessageQueue); + state + .app + .update_connection_pools(|connections| { + connections.insert(initial.id.clone(), PoolKind::MessageQueue); + }) + .await; let result = load_connections(State(state.clone()), HeaderMap::new()).await; assert!(result.is_ok()); @@ -1661,7 +1680,7 @@ mod tests { .and_then(serde_json::Value::as_str); assert_eq!(cached_admin_url, Some("http://127.0.0.1:8081")); drop(configs); - assert!(!state.app.connections.read().await.contains_key(&initial.id)); + assert!(!state.app.pool_handle(&initial.id).await.is_some()); let _ = std::fs::remove_dir_all(dir); } @@ -1838,7 +1857,12 @@ mod tests { configs.insert(kept.id.clone(), kept.clone()); configs.insert(removed.id.clone(), removed.clone()); } - state.app.connections.write().await.insert(removed.id.clone(), PoolKind::MessageQueue); + state + .app + .update_connection_pools(|connections| { + connections.insert(removed.id.clone(), PoolKind::MessageQueue); + }) + .await; let result = save_connections( State(state.clone()), @@ -1848,7 +1872,7 @@ mod tests { .await; assert!(result.is_ok()); - assert!(!state.app.connections.read().await.contains_key(&removed.id)); + assert!(!state.app.pool_handle(&removed.id).await.is_some()); let _ = std::fs::remove_dir_all(dir); } @@ -1863,11 +1887,13 @@ mod tests { let conn_pool = dbx_core::db::sqlite::connect_path(&conn_path.to_string_lossy()).await.unwrap(); let conn2_pool = dbx_core::db::sqlite::connect_path(&conn2_path.to_string_lossy()).await.unwrap(); - { - let mut connections = state.app.connections.write().await; - connections.insert("conn".to_string(), PoolKind::Sqlite(conn_pool)); - connections.insert("conn2".to_string(), PoolKind::Sqlite(conn2_pool)); - } + state + .app + .update_connection_pools(|connections| { + connections.insert("conn".to_string(), PoolKind::Sqlite(conn_pool)); + connections.insert("conn2".to_string(), PoolKind::Sqlite(conn2_pool)); + }) + .await; let result = disconnect_db( State(state.clone()), @@ -1876,9 +1902,10 @@ mod tests { .await; assert!(result.is_ok()); - let connections = state.app.connections.read().await; - assert!(!connections.contains_key("conn")); - assert!(connections.contains_key("conn2")); + let (has_conn, has_conn2) = + state.app.with_connection_pools(|pools| (pools.contains_key("conn"), pools.contains_key("conn2"))).await; + assert!(!has_conn); + assert!(has_conn2); let _ = std::fs::remove_dir_all(dir); } @@ -1891,7 +1918,12 @@ mod tests { let conn_pool = dbx_core::db::sqlite::connect_path(&conn_path.to_string_lossy()).await.unwrap(); state.app.begin_connection_attempt_with_client_attempt("conn", Some(1)).await; let current_attempt = state.app.begin_connection_attempt_with_client_attempt("conn", Some(2)).await; - state.app.connections.write().await.insert("conn".to_string(), PoolKind::Sqlite(conn_pool)); + state + .app + .update_connection_pools(|connections| { + connections.insert("conn".to_string(), PoolKind::Sqlite(conn_pool)); + }) + .await; let result = disconnect_db( State(state.clone()), @@ -1900,7 +1932,7 @@ mod tests { .await; assert!(result.is_ok()); - assert!(state.app.connections.read().await.contains_key("conn")); + assert!(state.app.pool_handle("conn").await.is_some()); assert!(state.app.ensure_current_connection_attempt("conn", Some(current_attempt)).await.is_ok()); let result = disconnect_db( @@ -1910,7 +1942,7 @@ mod tests { .await; assert!(result.is_ok()); - assert!(!state.app.connections.read().await.contains_key("conn")); + assert!(!state.app.pool_handle("conn").await.is_some()); assert!(state.app.ensure_current_connection_attempt("conn", Some(current_attempt)).await.is_err()); let _ = std::fs::remove_dir_all(dir); @@ -1923,10 +1955,12 @@ mod tests { std::fs::File::create(&conn_path).unwrap(); let conn_pool = dbx_core::db::sqlite::connect_path(&conn_path.to_string_lossy()).await.unwrap(); - { - let mut connections = state.app.connections.write().await; - connections.insert("conn".to_string(), PoolKind::Sqlite(conn_pool)); - } + state + .app + .update_connection_pools(|connections| { + connections.insert("conn".to_string(), PoolKind::Sqlite(conn_pool)); + }) + .await; { let mut configs = state.app.configs.write().await; configs.insert("conn".to_string(), sqlite_config("conn", &conn_path.to_string_lossy())); @@ -1951,7 +1985,12 @@ mod tests { let (state, dir) = test_web_state().await; let config = mq_config("mq-conn", "http://127.0.0.1:8080"); state.app.configs.write().await.insert(config.id.clone(), config.clone()); - state.app.connections.write().await.insert(config.id.clone(), PoolKind::MessageQueue); + state + .app + .update_connection_pools(|connections| { + connections.insert(config.id.clone(), PoolKind::MessageQueue); + }) + .await; let first = state.app.mq_registry.get_or_build(&config).await.unwrap().adapter; let result = disconnect_db( @@ -1961,7 +2000,7 @@ mod tests { .await; assert!(result.is_ok()); - assert!(!state.app.connections.read().await.contains_key(&config.id)); + assert!(!state.app.pool_handle(&config.id).await.is_some()); let second = state.app.mq_registry.get_or_build(&config).await.unwrap().adapter; assert!(!Arc::ptr_eq(&first, &second)); diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index cf54aa343b..87e0c46942 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -420,7 +420,11 @@ mod tests { .await .unwrap(); state.configs.write().await.insert(initial.id.clone(), initial.clone()); - state.connections.write().await.insert(initial.id.clone(), PoolKind::Sqlite(pool.clone())); + state + .update_connection_pools(|connections| { + connections.insert(initial.id.clone(), PoolKind::Sqlite(pool.clone())); + }) + .await; let mut invalid = initial.clone(); invalid.attached_databases.push(AttachedDatabaseConfig { @@ -430,7 +434,7 @@ mod tests { let error = save_connection_configs(&state, &[invalid]).await.unwrap_err(); assert!(error.contains("in-memory main database"), "{error}"); - assert!(state.connections.read().await.contains_key(&initial.id)); + assert!(state.pool_handle(&initial.id).await.is_some()); assert_eq!(state.configs.read().await.get(&initial.id), Some(&initial)); let retained = dbx_core::db::sqlite::execute_query(&pool, "SELECT value FROM retained;").await.unwrap(); assert_eq!(retained.rows[0][0], serde_json::json!("yes")); @@ -642,7 +646,11 @@ mod tests { let state = AppState::new_with_plugin_dir(storage, dir.join("plugins")); let initial = mq_config("mq-conn", "http://127.0.0.1:8080"); state.configs.write().await.insert(initial.id.clone(), initial.clone()); - state.connections.write().await.insert(initial.id.clone(), PoolKind::MessageQueue); + state + .update_connection_pools(|connections| { + connections.insert(initial.id.clone(), PoolKind::MessageQueue); + }) + .await; let first = state.mq_registry.get_or_build(&initial).await.unwrap().adapter; let updated = mq_config("mq-conn", "http://127.0.0.1:8081"); @@ -661,7 +669,7 @@ mod tests { let second = state.mq_registry.get_or_build(&updated).await.unwrap().adapter; assert!(!std::sync::Arc::ptr_eq(&first, &second)); - assert!(!state.connections.read().await.contains_key(&initial.id)); + assert!(state.pool_handle(&initial.id).await.is_none()); let _ = std::fs::remove_dir_all(dir); } @@ -677,7 +685,11 @@ mod tests { let updated = mq_config("mq-conn", "http://127.0.0.1:8081"); state.storage.save_connections(std::slice::from_ref(&updated)).await.unwrap(); state.configs.write().await.insert(initial.id.clone(), initial.clone()); - state.connections.write().await.insert(initial.id.clone(), PoolKind::MessageQueue); + state + .update_connection_pools(|connections| { + connections.insert(initial.id.clone(), PoolKind::MessageQueue); + }) + .await; let loaded = load_connection_configs(&state).await.unwrap(); @@ -692,7 +704,7 @@ mod tests { .and_then(serde_json::Value::as_str) .map(str::to_string); assert_eq!(cached_admin_url.as_deref(), Some("http://127.0.0.1:8081")); - assert!(!state.connections.read().await.contains_key(&initial.id)); + assert!(state.pool_handle(&initial.id).await.is_none()); let _ = std::fs::remove_dir_all(dir); } @@ -804,11 +816,15 @@ mod tests { configs.insert(kept.id.clone(), kept.clone()); configs.insert(removed.id.clone(), removed.clone()); } - state.connections.write().await.insert(removed.id.clone(), PoolKind::MessageQueue); + state + .update_connection_pools(|connections| { + connections.insert(removed.id.clone(), PoolKind::MessageQueue); + }) + .await; save_connection_configs(&state, std::slice::from_ref(&kept)).await.unwrap(); - assert!(!state.connections.read().await.contains_key(&removed.id)); + assert!(state.pool_handle(&removed.id).await.is_none()); let _ = std::fs::remove_dir_all(dir); } @@ -1696,19 +1712,17 @@ pub async fn connect_db( ), DatabaseType::Redis => { let con = if db_config.uses_redis_cluster() { - PoolKind::Redis(db::redis_driver::RedisConnection::Cluster( - state.connect_redis_cluster(&id, &db_config).await?, - )) + db::redis_driver::RedisConnection::Cluster(state.connect_redis_cluster(&id, &db_config).await?) } else if db_config.uses_redis_sentinel() { - PoolKind::Redis(db::redis_driver::RedisConnection::Direct(tokio::sync::Mutex::new( + db::redis_driver::RedisConnection::Direct(tokio::sync::Mutex::new( state.connect_redis_sentinel(&id, &db_config).await?, - ))) + )) } else { - PoolKind::Redis(db::redis_driver::RedisConnection::Direct(tokio::sync::Mutex::new( + db::redis_driver::RedisConnection::Direct(tokio::sync::Mutex::new( db::redis_driver::connect_standalone(&db_config, &host, port, connect_timeout).await?, - ))) + )) }; - con + PoolKind::Redis(Arc::new(con)) } #[cfg(feature = "duckdb-sidecar")] DatabaseType::DuckDb => state.create_duckdb_pool(&db_config).await?, diff --git a/src-tauri/src/commands/mqtt_cmd.rs b/src-tauri/src/commands/mqtt_cmd.rs index 8e54f0b975..ef91b23029 100644 --- a/src-tauri/src/commands/mqtt_cmd.rs +++ b/src-tauri/src/commands/mqtt_cmd.rs @@ -12,10 +12,9 @@ async fn get_mqtt_client( state: &AppState, connection_id: &str, ) -> Result, String> { - let connections = state.connections.read().await; - let pool = connections.get(connection_id).ok_or_else(|| format!("连接 {} 未建立", connection_id))?; + let pool = state.pool_handle(connection_id).await.ok_or_else(|| format!("连接 {} 未建立", connection_id))?; match pool { - PoolKind::Mqtt(client) => Ok(Arc::clone(client)), + PoolKind::Mqtt(client) => Ok(client), _ => Err(format!("连接 {} 不是 MQTT 类型", connection_id)), } } diff --git a/src-tauri/src/commands/sqlite_backup.rs b/src-tauri/src/commands/sqlite_backup.rs index 052ff7b1d2..bc892d258b 100644 --- a/src-tauri/src/commands/sqlite_backup.rs +++ b/src-tauri/src/commands/sqlite_backup.rs @@ -54,9 +54,8 @@ async fn sqlite_source_path(state: &Arc, connection_id: &str) -> Resul } async fn existing_sqlite_pool(state: &Arc, pool_key: &str) -> Result, String> { - let connections = state.connections.read().await; - match connections.get(pool_key) { - Some(PoolKind::Sqlite(pool)) => Ok(Some(pool.clone())), + match state.pool_handle(pool_key).await { + Some(PoolKind::Sqlite(pool)) => Ok(Some(pool)), Some(_) => Err("SQLite backup is only available for SQLite connections".to_string()), None => Ok(None), }