From d326cb0cedfcfc98ee6feba13da509c81cc2a3df Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 25 May 2026 17:23:36 +0800 Subject: [PATCH] fix: scope query connections by client session --- apps/desktop/src/lib/api.ts | 1 + apps/desktop/src/lib/http.ts | 33 +++++++-- apps/desktop/src/lib/tauri.ts | 33 +++++++-- apps/desktop/src/stores/queryStore.ts | 27 ++++++-- crates/dbx-core/src/connection.rs | 67 ++++++++++++++++++- crates/dbx-core/src/db/clickhouse_driver.rs | 2 +- .../dbx-core/src/db/elasticsearch_driver.rs | 2 +- crates/dbx-core/src/db/mongo_driver.rs | 2 +- crates/dbx-core/src/db/mysql.rs | 20 +++--- crates/dbx-core/src/db/postgres.rs | 4 +- crates/dbx-core/src/query.rs | 33 ++++++--- crates/dbx-web/src/main.rs | 1 + crates/dbx-web/src/routes/query.rs | 34 +++++++++- .../app-tests/connectionSessionModel.test.ts | 39 +++++++++++ packages/app-tests/paginationPageSize.test.ts | 2 +- src-tauri/src/commands/query.rs | 33 ++++++++- src-tauri/src/lib.rs | 1 + 17 files changed, 290 insertions(+), 44 deletions(-) create mode 100644 packages/app-tests/connectionSessionModel.test.ts diff --git a/apps/desktop/src/lib/api.ts b/apps/desktop/src/lib/api.ts index 6bf83e664..1efeec9da 100644 --- a/apps/desktop/src/lib/api.ts +++ b/apps/desktop/src/lib/api.ts @@ -89,6 +89,7 @@ export const executeScript = forward("executeScript"); export const executeInTransaction = forward("executeInTransaction"); export const cancelQuery = forward("cancelQuery"); export const closeQuerySession = forward("closeQuerySession"); +export const closeClientConnectionSession = forward("closeClientConnectionSession"); export const analyzeSqlReferences = forward("analyzeSqlReferences"); export const findStatementAtCursor = forward("findStatementAtCursor"); export const prepareQueryPaginationExecutionPlan = forward("prepareQueryPaginationExecutionPlan"); diff --git a/apps/desktop/src/lib/http.ts b/apps/desktop/src/lib/http.ts index b7c83a763..9d786ea91 100644 --- a/apps/desktop/src/lib/http.ts +++ b/apps/desktop/src/lib/http.ts @@ -407,7 +407,13 @@ export async function executeQuery( sql: string, schema?: string, executionId?: string, - options?: { maxRows?: number; fetchSize?: number; pageSize?: number; resultSessionId?: string }, + options?: { + maxRows?: number; + fetchSize?: number; + pageSize?: number; + resultSessionId?: string; + clientSessionId?: string; + }, ): Promise { return post("/api/query/execute", { connectionId, database, sql, schema, executionId, ...options }); } @@ -418,13 +424,32 @@ export async function executeMulti( sql: string, schema?: string, executionId?: string, - options?: { maxRows?: number; fetchSize?: number; pageSize?: number; resultSessionId?: string }, + options?: { + maxRows?: number; + fetchSize?: number; + pageSize?: number; + resultSessionId?: string; + clientSessionId?: string; + }, ): Promise { return post("/api/query/execute-multi", { connectionId, database, sql, schema, executionId, ...options }); } -export async function closeQuerySession(connectionId: string, database: string, sessionId: string): Promise { - return post("/api/query/close-session", { connectionId, database, sessionId }); +export async function closeQuerySession( + connectionId: string, + database: string, + sessionId: string, + clientSessionId?: string, +): Promise { + return post("/api/query/close-session", { connectionId, database, sessionId, clientSessionId }); +} + +export async function closeClientConnectionSession( + connectionId: string, + database: string, + clientSessionId: string, +): Promise { + return post("/api/query/close-client-session", { connectionId, database, clientSessionId }); } export async function executeBatch( diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index 2c627d744..20814e370 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -361,7 +361,13 @@ export async function executeQuery( sql: string, schema?: string, executionId?: string, - options?: { maxRows?: number; fetchSize?: number; pageSize?: number; resultSessionId?: string }, + options?: { + maxRows?: number; + fetchSize?: number; + pageSize?: number; + resultSessionId?: string; + clientSessionId?: string; + }, ): Promise { return invoke("execute_query", { connectionId, database, sql, schema, executionId, ...options }); } @@ -372,7 +378,13 @@ export async function executeMulti( sql: string, schema?: string, executionId?: string, - options?: { maxRows?: number; fetchSize?: number; pageSize?: number; resultSessionId?: string }, + options?: { + maxRows?: number; + fetchSize?: number; + pageSize?: number; + resultSessionId?: string; + clientSessionId?: string; + }, ): Promise { return invoke("execute_multi", { connectionId, database, sql, schema, executionId, ...options }); } @@ -381,8 +393,21 @@ export async function cancelQuery(executionId: string): Promise { return invoke("cancel_query", { executionId }); } -export async function closeQuerySession(connectionId: string, database: string, sessionId: string): Promise { - return invoke("close_query_session", { connectionId, database, sessionId }); +export async function closeQuerySession( + connectionId: string, + database: string, + sessionId: string, + clientSessionId?: string, +): Promise { + return invoke("close_query_session", { connectionId, database, sessionId, clientSessionId }); +} + +export async function closeClientConnectionSession( + connectionId: string, + database: string, + clientSessionId: string, +): Promise { + return invoke("close_client_connection_session", { connectionId, database, clientSessionId }); } export async function executeBatch( diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 26d7d2880..6e0625bb1 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -62,7 +62,7 @@ export const useQueryStore = defineStore("query", () => { const sessionId = tab?.resultSessionId ?? tab?.result?.session_id; if (!tab || !sessionId || sessionId === preserveSessionId) return; try { - await api.closeQuerySession(tab.connectionId, tab.database, sessionId); + await api.closeQuerySession(tab.connectionId, tab.database, sessionId, tab.id); } catch (error) { console.warn("[DBX][query-session:close:error]", { tabId: tab.id, sessionId, error }); } finally { @@ -71,6 +71,15 @@ export const useQueryStore = defineStore("query", () => { } } + async function closeClientConnectionSession(tab: QueryTab | undefined) { + if (!tab?.connectionId) return; + try { + await api.closeClientConnectionSession(tab.connectionId, tab.database, tab.id); + } catch (error) { + console.warn("[DBX][client-session:close:error]", { tabId: tab.id, error }); + } + } + function clearResultPayload(tab: QueryTab, options: { evicted?: boolean } = {}) { tab.result = undefined; tab.results = undefined; @@ -245,6 +254,7 @@ export const useQueryStore = defineStore("query", () => { if (tabs.value[idx].isExecuting) void cancelTabExecution(id); if (tabs.value[idx].isExplaining) void cancelTabExplain(id); void closeResultSession(tabs.value[idx]); + void closeClientConnectionSession(tabs.value[idx]); clearResultPayload(tabs.value[idx]); tabs.value.splice(idx, 1); if (activeTabId.value === id) { @@ -256,6 +266,7 @@ export const useQueryStore = defineStore("query", () => { tabs.value.filter((tab) => tab.id !== id && tab.isExecuting).forEach((tab) => void cancelTabExecution(tab.id)); tabs.value.filter((tab) => tab.id !== id && tab.isExplaining).forEach((tab) => void cancelTabExplain(tab.id)); tabs.value.filter((tab) => tab.id !== id).forEach((tab) => void closeResultSession(tab)); + tabs.value.filter((tab) => tab.id !== id).forEach((tab) => void closeClientConnectionSession(tab)); const next = closeOtherTabsState(tabs.value, activeTabId.value, id); tabs.value = next.tabs; activeTabId.value = next.activeTabId; @@ -265,6 +276,7 @@ export const useQueryStore = defineStore("query", () => { tabs.value.filter((tab) => tab.isExecuting).forEach((tab) => void cancelTabExecution(tab.id)); tabs.value.filter((tab) => tab.isExplaining).forEach((tab) => void cancelTabExplain(tab.id)); tabs.value.forEach((tab) => void closeResultSession(tab)); + tabs.value.forEach((tab) => void closeClientConnectionSession(tab)); const next = closeAllTabsState(tabs.value, activeTabId.value); tabs.value = next.tabs; activeTabId.value = next.activeTabId; @@ -325,10 +337,11 @@ export const useQueryStore = defineStore("query", () => { function updateDatabase(id: string, database: string) { const tab = tabs.value.find((t) => t.id === id); if (!tab || tab.database === database) return; + void closeResultSession(tab); + void closeClientConnectionSession(tab); tab.database = database; tab.schema = undefined; tab.objectBrowser = undefined; - void closeResultSession(tab); clearResultPayload(tab); tab.lastExecutedSql = undefined; tab.resultBaseSql = undefined; @@ -347,10 +360,11 @@ export const useQueryStore = defineStore("query", () => { function updateConnection(id: string, connectionId: string, database = "") { const tab = tabs.value.find((t) => t.id === id); if (!tab || tab.connectionId === connectionId) return; + void closeResultSession(tab); + void closeClientConnectionSession(tab); tab.connectionId = connectionId; tab.database = database; tab.schema = undefined; - void closeResultSession(tab); clearResultPayload(tab); tab.lastExecutedSql = undefined; tab.resultBaseSql = undefined; @@ -637,8 +651,9 @@ export const useQueryStore = defineStore("query", () => { fetchSize: pageLimit, pageSize: pageLimit, resultSessionId: options?.pagination?.sessionId, + clientSessionId: tab.id, } - : { maxRows: pageLimit, fetchSize: pageLimit } + : { maxRows: pageLimit, fetchSize: pageLimit, clientSessionId: tab.id } : undefined; const results = await api.executeMulti( tab.connectionId, @@ -738,7 +753,9 @@ export const useQueryStore = defineStore("query", () => { tab.explainSql = built.sql; tab.lastExplainedSql = sql; try { - const result = await api.executeQuery(tab.connectionId, tab.database, built.sql, tab.schema, executionId); + const result = await api.executeQuery(tab.connectionId, tab.database, built.sql, tab.schema, executionId, { + clientSessionId: tab.id, + }); const current = tabs.value.find((t) => t.id === id); if (current?.explainExecutionId === executionId) { current.explainPlan = parseExplainResult(databaseType as "mysql" | "postgres", result); diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 5ae22cd32..cedfce964 100644 --- a/crates/dbx-core/src/connection.rs +++ b/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 { + 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 { 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 { + 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 { 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 { + 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 { @@ -515,6 +562,16 @@ impl AppState { } } +fn normalize_client_session_id(client_session_id: Option<&str>) -> Option { + 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); diff --git a/crates/dbx-core/src/db/clickhouse_driver.rs b/crates/dbx-core/src/db/clickhouse_driver.rs index 6ef09ce3c..c3aa3cf59 100644 --- a/crates/dbx-core/src/db/clickhouse_driver.rs +++ b/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?; diff --git a/crates/dbx-core/src/db/elasticsearch_driver.rs b/crates/dbx-core/src/db/elasticsearch_driver.rs index 4406e1f73..448bbc338 100644 --- a/crates/dbx-core/src/db/elasticsearch_driver.rs +++ b/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?; diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index 5ad4d4531..19cb7f0a6 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -13,7 +13,7 @@ pub struct MongoDocumentResult { } pub async fn connect(url: &str) -> Result { - 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 diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 6460b6541..029b1fdaa 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -187,7 +187,7 @@ pub async fn connect(url: &str) -> Result { fn create_pool(url: &str) -> Result { 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); } } diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 5eb677e49..742cc6367 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -125,7 +125,7 @@ pub async fn connect(url: &str) -> Result { 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 { mgr_config, ); let pool = Pool::builder(mgr) - .max_size(10) + .max_size(1) .runtime(Runtime::Tokio1) .wait_timeout(Some(super::connection_timeout())) .build() diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 9fec9f2a9..bd89d4702 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -26,6 +26,7 @@ pub struct QueryExecutionOptions { pub fetch_size: Option, pub page_size: Option, pub result_session_id: Option, + pub client_session_id: Option, } fn query_result_row_limit(max_rows: Option) -> usize { @@ -602,9 +603,11 @@ pub async fn execute_sql_statement_with_options( options: QueryExecutionOptions, ) -> Result { 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 { 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, 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 { diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index f0e6ab39e..b04552ec2 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -185,6 +185,7 @@ async fn main() { .route("/data-compare/prepare-from-tables", post(routes::data_compare::prepare_data_compare_from_tables)) .route("/query/cancel", post(routes::query::cancel_query)) .route("/query/close-session", post(routes::query::close_query_session)) + .route("/query/close-client-session", post(routes::query::close_client_connection_session)) .route("/export/query-result-json", post(routes::text_export::export_query_result_json)) .route("/export/query-result-markdown", post(routes::text_export::export_query_result_markdown)) // Redis diff --git a/crates/dbx-web/src/routes/query.rs b/crates/dbx-web/src/routes/query.rs index 457d1bc00..dfad94408 100644 --- a/crates/dbx-web/src/routes/query.rs +++ b/crates/dbx-web/src/routes/query.rs @@ -19,6 +19,7 @@ pub struct ExecuteQueryRequest { pub fetch_size: Option, pub page_size: Option, pub result_session_id: Option, + pub client_session_id: Option, } #[derive(Deserialize)] @@ -33,6 +34,15 @@ pub struct CloseSessionRequest { pub connection_id: String, pub database: String, pub session_id: String, + pub client_session_id: Option, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CloseClientConnectionSessionRequest { + pub connection_id: String, + pub database: String, + pub client_session_id: String, } #[derive(Deserialize)] @@ -260,6 +270,7 @@ pub async fn execute_query( fetch_size: req.fetch_size, page_size: req.page_size, result_session_id: req.result_session_id, + client_session_id: req.client_session_id, }, ) .await @@ -290,6 +301,7 @@ pub async fn execute_multi( fetch_size: req.fetch_size, page_size: req.page_size, result_session_id: req.result_session_id, + client_session_id: req.client_session_id, }, ) .await @@ -328,7 +340,27 @@ pub async fn close_query_session( State(state): State>, Json(req): Json, ) -> Result, AppError> { - let closed = dbx_core::query::close_query_session(&state.app, &req.connection_id, &req.database, &req.session_id) + let closed = dbx_core::query::close_query_session( + &state.app, + &req.connection_id, + &req.database, + &req.session_id, + req.client_session_id.as_deref(), + ) + .await + .map_err(AppError)?; + + Ok(Json(serde_json::json!(closed))) +} + +pub async fn close_client_connection_session( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + let database = if req.database.trim().is_empty() { None } else { Some(req.database.as_str()) }; + let closed = state + .app + .close_client_session_pool(&req.connection_id, database, &req.client_session_id) .await .map_err(AppError)?; diff --git a/packages/app-tests/connectionSessionModel.test.ts b/packages/app-tests/connectionSessionModel.test.ts new file mode 100644 index 000000000..4daeb8693 --- /dev/null +++ b/packages/app-tests/connectionSessionModel.test.ts @@ -0,0 +1,39 @@ +import { readFileSync } from "node:fs"; +import { strict as assert } from "node:assert"; +import test from "node:test"; + +const connectionSource = readFileSync("crates/dbx-core/src/connection.rs", "utf8"); +const querySource = readFileSync("crates/dbx-core/src/query.rs", "utf8"); +const queryStoreSource = readFileSync("apps/desktop/src/stores/queryStore.ts", "utf8"); +const apiSource = readFileSync("apps/desktop/src/lib/api.ts", "utf8"); +const tauriApiSource = readFileSync("apps/desktop/src/lib/tauri.ts", "utf8"); +const httpApiSource = readFileSync("apps/desktop/src/lib/http.ts", "utf8"); +const tauriCommandSource = readFileSync("src-tauri/src/commands/query.rs", "utf8"); +const webRouteSource = readFileSync("crates/dbx-web/src/routes/query.rs", "utf8"); +const mysqlSource = readFileSync("crates/dbx-core/src/db/mysql.rs", "utf8"); +const postgresSource = readFileSync("crates/dbx-core/src/db/postgres.rs", "utf8"); + +test("query tabs pass a stable client session id to backend execution", () => { + assert.match(queryStoreSource, /clientSessionId: tab\.id/); + assert.match(queryStoreSource, /closeQuerySession\(tab\.connectionId, tab\.database, sessionId, tab\.id\)/); + assert.match(apiSource, /closeClientConnectionSession = forward\("closeClientConnectionSession"\)/); + assert.match(tauriApiSource, /clientSessionId\?: string/); + assert.match(httpApiSource, /clientSessionId\?: string/); + assert.match(tauriCommandSource, /client_session_id: Option/); + assert.match(webRouteSource, /pub client_session_id: Option/); +}); + +test("backend scopes query pools by client session and can close them", () => { + assert.match(connectionSource, /get_or_create_pool_for_session/); + assert.match(connectionSource, /session_scoped_pool_key/); + assert.match(connectionSource, /close_client_session_pool/); + assert.match(querySource, /client_session_id: Option/); + assert.match(querySource, /get_or_create_pool_for_session\(connection_id, Some\(database\), options\.client_session_id\.as_deref\(\)\)/); + assert.match(querySource, /close_query_session\([\s\S]*client_session_id: Option<&str>/); + assert.match(querySource, /execute_sql_statement_with_options\([\s\S]*options\.clone\(\)/); +}); + +test("native SQL query sessions use a single physical connection", () => { + assert.match(mysqlSource, /PoolConstraints::new\(1, 1\)/); + assert.match(postgresSource, /\.max_size\(1\)/); +}); diff --git a/packages/app-tests/paginationPageSize.test.ts b/packages/app-tests/paginationPageSize.test.ts index d5ab8aec2..8df23c692 100644 --- a/packages/app-tests/paginationPageSize.test.ts +++ b/packages/app-tests/paginationPageSize.test.ts @@ -55,7 +55,7 @@ test("query execution sends the selected page size to agent drivers", () => { assert.match(source, /if \(tab\.mode === "data"\) \{/); assert.match(source, /pageLimit = settingsStore\.editorSettings\.pageSize/); - assert.match(source, /maxRows: pageLimit,\s*fetchSize: pageLimit,\s*pageSize: pageLimit/s); + assert.match(source, /maxRows: pageLimit,\s*fetchSize: pageLimit,\s*pageSize: pageLimit,\s*resultSessionId: options\?\.pagination\?\.sessionId,\s*clientSessionId: tab\.id/s); assert.doesNotMatch(source, /maxRows: 10000,\s*fetchSize: pageLimit,\s*pageSize: pageLimit/s); }); diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index e9cf97f25..15ba96270 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -21,6 +21,7 @@ pub async fn execute_query( fetch_size: Option, page_size: Option, result_session_id: Option, + client_session_id: Option, ) -> Result { let registered_query = execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); @@ -33,7 +34,13 @@ pub async fn execute_query( &sql, schema.as_deref(), cancel_token, - dbx_core::query::QueryExecutionOptions { max_rows, fetch_size, page_size, result_session_id }, + dbx_core::query::QueryExecutionOptions { + max_rows, + fetch_size, + page_size, + result_session_id, + client_session_id, + }, ) .await } @@ -50,6 +57,7 @@ pub async fn execute_multi( fetch_size: Option, page_size: Option, result_session_id: Option, + client_session_id: Option, ) -> Result, String> { let registered_query = execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); @@ -71,7 +79,13 @@ pub async fn execute_multi( &sql, schema.as_deref(), cancel_token, - dbx_core::query::QueryExecutionOptions { max_rows, fetch_size, page_size, result_session_id }, + dbx_core::query::QueryExecutionOptions { + max_rows, + fetch_size, + page_size, + result_session_id, + client_session_id, + }, ) .await; match &result { @@ -97,8 +111,21 @@ pub async fn close_query_session( connection_id: String, database: String, session_id: String, + client_session_id: Option, ) -> Result { - dbx_core::query::close_query_session(&state, &connection_id, &database, &session_id).await + dbx_core::query::close_query_session(&state, &connection_id, &database, &session_id, client_session_id.as_deref()) + .await +} + +#[tauri::command] +pub async fn close_client_connection_session( + state: State<'_, Arc>, + connection_id: String, + database: String, + client_session_id: String, +) -> Result { + let database = if database.trim().is_empty() { None } else { Some(database.as_str()) }; + state.close_client_session_pool(&connection_id, database, &client_session_id).await } #[tauri::command] diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index c92a99dc2..3944506ab 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -281,6 +281,7 @@ pub fn run() { commands::query::execute_multi, commands::query::cancel_query, commands::query::close_query_session, + commands::query::close_client_connection_session, commands::query::execute_batch, commands::query::execute_script, commands::query::execute_in_transaction,