fix: scope query connections by client session

crates/dbx-core/src/connection.rs

@@ -137,13 +137,22 @@ impl AppState {
     }
 
     pub async fn get_or_create_pool(&self, connection_id: &str, database: Option<&str>) -> Result<String, String> {
+        self.get_or_create_pool_for_session(connection_id, database, None).await
+    }
+
+    pub async fn get_or_create_pool_for_session(
+        &self,
+        connection_id: &str,
+        database: Option<&str>,
+        client_session_id: Option<&str>,
+    ) -> Result<String, String> {
         let db_type = {
             let configs = self.configs.read().await;
             configs.get(connection_id).map(|c| c.db_type.clone())
         };
 
         let is_single_conn = db_type.as_ref().is_some_and(database_capabilities::is_single_connection_pool);
-        let pool_key = if is_single_conn {
+        let base_pool_key = if is_single_conn {
             connection_id.to_string()
         } else {
             match database {
@@ -151,6 +160,7 @@ impl AppState {
                 None => connection_id.to_string(),
             }
         };
+        let pool_key = session_scoped_pool_key(base_pool_key, client_session_id);
 
         let conns = self.connections.read().await;
         if conns.contains_key(&pool_key) {
@@ -430,6 +440,15 @@ impl AppState {
     }
 
     pub async fn reconnect_pool(&self, connection_id: &str, database: Option<&str>) -> Result<String, String> {
+        self.reconnect_pool_for_session(connection_id, database, None).await
+    }
+
+    pub async fn reconnect_pool_for_session(
+        &self,
+        connection_id: &str,
+        database: Option<&str>,
+        client_session_id: Option<&str>,
+    ) -> Result<String, String> {
         let is_single_conn = {
             let configs = self.configs.read().await;
             configs
@@ -440,7 +459,7 @@ impl AppState {
                 })
                 .unwrap_or(false)
         };
-        let pool_key = if is_single_conn {
+        let base_pool_key = if is_single_conn {
             connection_id.to_string()
         } else {
             match database {
@@ -448,13 +467,41 @@ impl AppState {
                 None => connection_id.to_string(),
             }
         };
+        let pool_key = session_scoped_pool_key(base_pool_key, client_session_id);
         if self.uses_forwarded_transport(connection_id).await {
             self.remove_connection_pools(connection_id).await;
             self.reset_connection_transport(connection_id).await;
         } else {
             self.connections.write().await.remove(&pool_key);
         }
-        self.get_or_create_pool(connection_id, database).await
+        self.get_or_create_pool_for_session(connection_id, database, client_session_id).await
+    }
+
+    pub async fn close_client_session_pool(
+        &self,
+        connection_id: &str,
+        database: Option<&str>,
+        client_session_id: &str,
+    ) -> Result<bool, String> {
+        let session = normalize_client_session_id(Some(client_session_id));
+        let Some(session) = session else {
+            return Ok(false);
+        };
+        let db_type = {
+            let configs = self.configs.read().await;
+            configs.get(connection_id).map(|c| c.db_type.clone())
+        };
+        let is_single_conn = db_type.as_ref().is_some_and(database_capabilities::is_single_connection_pool);
+        let base_pool_key = if is_single_conn {
+            connection_id.to_string()
+        } else {
+            match database {
+                Some(db) => format!("{connection_id}:{db}"),
+                None => connection_id.to_string(),
+            }
+        };
+        let pool_key = session_scoped_pool_key(base_pool_key, Some(&session));
+        Ok(self.connections.write().await.remove(&pool_key).is_some())
     }
 
     pub async fn duckdb_existing_pool_is_usable_for_config(&self, config: &ConnectionConfig) -> Result<bool, String> {
@@ -515,6 +562,16 @@ impl AppState {
     }
 }
 
+fn normalize_client_session_id(client_session_id: Option<&str>) -> Option<String> {
+    client_session_id.map(str::trim).filter(|session| !session.is_empty()).map(|session| session.replace(':', "_"))
+}
+
+fn session_scoped_pool_key(base_pool_key: String, client_session_id: Option<&str>) -> String {
+    normalize_client_session_id(client_session_id)
+        .map(|session| format!("{base_pool_key}:session:{session}"))
+        .unwrap_or(base_pool_key)
+}
+
 fn default_plugin_dir() -> PathBuf {
     default_dbx_dir().join("plugins")
 }
@@ -1104,6 +1161,8 @@ mod tests {
             let mut conns = state.connections.write().await;
             conns.insert("conn".to_string(), PoolKind::Sqlite(pool.clone()));
             conns.insert("conn:analytics".to_string(), PoolKind::Sqlite(pool.clone()));
+            conns.insert("conn:session:tab-1".to_string(), PoolKind::Sqlite(pool.clone()));
+            conns.insert("conn:analytics:session:tab-1".to_string(), PoolKind::Sqlite(pool.clone()));
             conns.insert("other".to_string(), PoolKind::Sqlite(pool));
         }
 
@@ -1112,6 +1171,8 @@ mod tests {
         let conns = state.connections.read().await;
         assert!(!conns.contains_key("conn"));
         assert!(!conns.contains_key("conn:analytics"));
+        assert!(!conns.contains_key("conn:session:tab-1"));
+        assert!(!conns.contains_key("conn:analytics:session:tab-1"));
         assert!(conns.contains_key("other"));
 
         let _ = std::fs::remove_dir_all(dir);

crates/dbx-core/src/db/clickhouse_driver.rs

@@ -114,7 +114,7 @@ fn limited_query_result(result: ChJsonResult, execution_time_ms: u128, max_rows:
 pub async fn test_connection(client: &ChClient) -> Result<(), String> {
     let url = format!("{}/?query=SELECT%201", client.base_url);
     let req = build_request(client, client.http.get(&url));
-    let resp = with_connection_timeout("ClickHouse", async {
+    let resp = with_connection_timeout("ClickHouse", connection_timeout(), async {
         req.send().await.map_err(|e| format!("ClickHouse connection failed: {e}"))
     })
     .await?;

crates/dbx-core/src/db/elasticsearch_driver.rs

@@ -60,7 +60,7 @@ impl Clone for EsClient {
 }
 
 pub async fn test_connection(client: &EsClient) -> Result<(), String> {
-    let resp = with_connection_timeout("Elasticsearch", async {
+    let resp = with_connection_timeout("Elasticsearch", connection_timeout(), async {
         client.get("/").send().await.map_err(|e| format!("Elasticsearch connection failed: {e}"))
     })
     .await?;

crates/dbx-core/src/db/mongo_driver.rs

@@ -13,7 +13,7 @@ pub struct MongoDocumentResult {
 }
 
 pub async fn connect(url: &str) -> Result<Client, String> {
-    with_connection_timeout("MongoDB", async {
+    with_connection_timeout("MongoDB", connection_timeout(), async {
         Client::with_uri_str(url).await.map_err(|e| format!("MongoDB connection failed: {e}"))
     })
     .await

crates/dbx-core/src/db/mysql.rs

@@ -187,7 +187,7 @@ pub async fn connect(url: &str) -> Result<MySqlPool, String> {
 fn create_pool(url: &str) -> Result<MySqlPool, String> {
     let opts = mysql_async::Opts::from_url(&mysql_async_url(url)).map_err(|e| format!("Invalid MySQL URL: {e}"))?;
     let pool_opts = mysql_async::PoolOpts::new()
-        .with_constraints(mysql_async::PoolConstraints::new(1, 5).unwrap())
+        .with_constraints(mysql_async::PoolConstraints::new(1, 1).unwrap())
         .with_inactive_connection_ttl(Duration::from_secs(300));
     let builder =
         mysql_async::OptsBuilder::from_opts(opts).stmt_cache_size(0).prefer_socket(false).pool_opts(Some(pool_opts));
@@ -797,35 +797,35 @@ mod tests {
     #[test]
     fn parse_connect_timeout_extracts_underscore_form() {
         let url = "mysql://host:3306/db?connect_timeout=30";
-        assert_eq!(super::parse_connect_timeout(url), Duration::from_secs(30));
+        assert_eq!(crate::db::parse_connect_timeout(url), Duration::from_secs(30));
     }
 
     #[test]
     fn parse_connect_timeout_extracts_camelcase_form() {
         let url = "mysql://host:3306/db?connectTimeout=60";
-        assert_eq!(super::parse_connect_timeout(url), Duration::from_secs(60));
+        assert_eq!(crate::db::parse_connect_timeout(url), Duration::from_secs(60));
     }
 
     #[test]
     fn parse_connect_timeout_ignores_out_of_range() {
-        let default = super::connection_timeout();
+        let default = crate::db::connection_timeout();
         let url = "mysql://host:3306/db?connect_timeout=999";
-        assert_eq!(super::parse_connect_timeout(url), default);
+        assert_eq!(crate::db::parse_connect_timeout(url), default);
         let url2 = "mysql://host:3306/db?connect_timeout=0";
-        assert_eq!(super::parse_connect_timeout(url2), default);
+        assert_eq!(crate::db::parse_connect_timeout(url2), default);
     }
 
     #[test]
     fn parse_connect_timeout_returns_default_when_missing() {
-        let default = super::connection_timeout();
+        let default = crate::db::connection_timeout();
         let url = "mysql://host:3306/db?ssl-mode=preferred&charset=utf8mb4";
-        assert_eq!(super::parse_connect_timeout(url), default);
+        assert_eq!(crate::db::parse_connect_timeout(url), default);
     }
 
     #[test]
     fn parse_connect_timeout_returns_default_when_no_query() {
-        let default = super::connection_timeout();
+        let default = crate::db::connection_timeout();
         let url = "mysql://host:3306/db";
-        assert_eq!(super::parse_connect_timeout(url), default);
+        assert_eq!(crate::db::parse_connect_timeout(url), default);
     }
 }

crates/dbx-core/src/db/postgres.rs

@@ -125,7 +125,7 @@ pub async fn connect(url: &str) -> Result<Pool, String> {
 
     let tz = iana_time_zone::get_timezone().unwrap_or_else(|_| "UTC".to_string());
 
-    super::with_connection_timeout("PostgreSQL", async {
+    super::with_connection_timeout("PostgreSQL", super::connection_timeout(), async {
         let pg_config =
             tokio_postgres::Config::from_str(url).map_err(|e| format!("Invalid PostgreSQL connection URL: {e}"))?;
 
@@ -139,7 +139,7 @@ pub async fn connect(url: &str) -> Result<Pool, String> {
             mgr_config,
         );
         let pool = Pool::builder(mgr)
-            .max_size(10)
+            .max_size(1)
             .runtime(Runtime::Tokio1)
             .wait_timeout(Some(super::connection_timeout()))
             .build()

crates/dbx-core/src/query.rs

@@ -26,6 +26,7 @@ pub struct QueryExecutionOptions {
     pub fetch_size: Option<usize>,
     pub page_size: Option<usize>,
     pub result_session_id: Option<String>,
+    pub client_session_id: Option<String>,
 }
 
 fn query_result_row_limit(max_rows: Option<usize>) -> usize {
@@ -602,9 +603,11 @@ pub async fn execute_sql_statement_with_options(
     options: QueryExecutionOptions,
 ) -> Result<db::QueryResult, String> {
     let pool_key = if database.is_empty() {
-        connection_id.to_string()
+        state.get_or_create_pool_for_session(connection_id, None, options.client_session_id.as_deref()).await?
     } else {
-        state.get_or_create_pool(connection_id, Some(database)).await?
+        state
+            .get_or_create_pool_for_session(connection_id, Some(database), options.client_session_id.as_deref())
+            .await?
     };
 
     if is_canceled(&cancel_token) {
@@ -616,7 +619,8 @@ pub async fn execute_sql_statement_with_options(
     match &result {
         Err(e) if is_connection_error(e) && !is_canceled(&cancel_token) => {
             let db_opt = if database.is_empty() { None } else { Some(database) };
-            let new_key = state.reconnect_pool(connection_id, db_opt).await?;
+            let new_key =
+                state.reconnect_pool_for_session(connection_id, db_opt, options.client_session_id.as_deref()).await?;
             do_execute(state, &new_key, Some(database), sql, schema, cancel_token, options).await
         }
         _ => result,
@@ -628,11 +632,12 @@ pub async fn close_query_session(
     connection_id: &str,
     database: &str,
     session_id: &str,
+    client_session_id: Option<&str>,
 ) -> Result<bool, String> {
     let pool_key = if database.is_empty() {
-        connection_id.to_string()
+        state.get_or_create_pool_for_session(connection_id, None, client_session_id).await?
     } else {
-        state.get_or_create_pool(connection_id, Some(database)).await?
+        state.get_or_create_pool_for_session(connection_id, Some(database), client_session_id).await?
     };
 
     let connections = state.connections.read().await;
@@ -678,9 +683,11 @@ pub async fn execute_multi_core_with_options(
     options: QueryExecutionOptions,
 ) -> Result<Vec<db::QueryResult>, String> {
     let pool_key = if database.is_empty() {
-        connection_id.to_string()
+        state.get_or_create_pool_for_session(connection_id, None, options.client_session_id.as_deref()).await?
     } else {
-        state.get_or_create_pool(connection_id, Some(database)).await?
+        state
+            .get_or_create_pool_for_session(connection_id, Some(database), options.client_session_id.as_deref())
+            .await?
     };
 
     let is_sqlserver = {
@@ -726,7 +733,17 @@ pub async fn execute_multi_core_with_options(
             });
             break;
         }
-        match execute_sql_statement(state, connection_id, database, stmt, schema, cancel_token.clone()).await {
+        match execute_sql_statement_with_options(
+            state,
+            connection_id,
+            database,
+            stmt,
+            schema,
+            cancel_token.clone(),
+            options.clone(),
+        )
+        .await
+        {
             Ok(r) => results.push(r),
             Err(e) => {
                 results.push(db::QueryResult {