diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index b69fb0af7..7ceba0d79 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -1211,6 +1211,7 @@ export const useQueryStore = defineStore("query", () => { } const countSql = options.countSql; const clientSessionId = tabClientSessionId({ id: options.tabId }, "count"); + const countExecutionId = `${options.executionId}:count`; if (typeof options.pageLimit === "number" && resultRowCount < options.pageLimit) { setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, (options.pageOffset ?? 0) + resultRowCount); @@ -1220,7 +1221,7 @@ export const useQueryStore = defineStore("query", () => { void (async () => { try { console.info("[DBX][executeTabSql:count:start]", { traceId: options.traceId, elapsed: options.elapsed() }); - const countResult = await api.executeQuery(options.connectionId, options.database, countSql, options.schema, undefined, { + const countResult = await api.executeQuery(options.connectionId, options.database, countSql, options.schema, countExecutionId, { clientSessionId, timeoutSecs: options.timeoutSecs, }); @@ -2040,6 +2041,7 @@ export const useQueryStore = defineStore("query", () => { let executionTimeMs = 0; let offset = 0; const clientSessionId = tabClientSessionId(tab, "export"); + const exportExecutionId = uuid(); try { while (true) { @@ -2055,7 +2057,7 @@ export const useQueryStore = defineStore("query", () => { limit: pageLimit, offset, }); - const results = await api.executeMulti(tab.connectionId, tab.database, sql, undefined, undefined, { + const results = await api.executeMulti(tab.connectionId, tab.database, sql, undefined, exportExecutionId, { maxRows: pageLimit, fetchSize: pageLimit, clientSessionId, @@ -2106,6 +2108,7 @@ export const useQueryStore = defineStore("query", () => { let offset = 0; let sessionId: string | undefined; const clientSessionId = tabClientSessionId(tab, "export"); + const exportExecutionId = uuid(); try { while (true) { @@ -2126,7 +2129,7 @@ export const useQueryStore = defineStore("query", () => { timeoutSecs: queryTimeoutSecs, } : { maxRows: plan.pageLimit, fetchSize: plan.pageLimit, clientSessionId, timeoutSecs: queryTimeoutSecs }; - const results = await api.executeMulti(tab.connectionId, tab.database, plan.sqlToExecute, tab.schema, undefined, executionOptions); + const results = await api.executeMulti(tab.connectionId, tab.database, plan.sqlToExecute, tab.schema, exportExecutionId, executionOptions); const result = results[0]; if (!result) break; if (columns.length === 0) columns = result.columns; diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 1401e39bb..89c58c831 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -1,7 +1,7 @@ use std::collections::{HashMap, HashSet}; use std::path::PathBuf; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use tokio::sync::RwLock; use tokio::task::JoinHandle; @@ -84,6 +84,7 @@ pub enum PoolKind { pub struct AppState { pub connections: Arc>>, keepalive_tasks: Arc>>>, + pool_activity: Arc>>, pub configs: RwLock>, pub running_queries: RunningQueries, pub tunnels: TunnelManager, @@ -95,6 +96,40 @@ pub struct AppState { pub mq_registry: crate::mq::MqAdminRegistry, } +#[derive(Clone, Copy)] +struct PoolActivity { + last_used_at: Instant, +} + +impl PoolActivity { + fn now() -> Self { + Self { last_used_at: Instant::now() } + } +} + +pub struct PoolActivityTouch { + pool_key: String, + connections: Arc>>, + pool_activity: Arc>>, +} + +impl Drop for PoolActivityTouch { + fn drop(&mut self) { + let pool_key = self.pool_key.clone(); + let connections = self.connections.clone(); + let pool_activity = self.pool_activity.clone(); + let Ok(handle) = tokio::runtime::Handle::try_current() else { + return; + }; + handle.spawn(async move { + if !connections.read().await.contains_key(&pool_key) { + return; + } + pool_activity.write().await.insert(pool_key, PoolActivity::now()); + }); + } +} + pub fn metadata_connection_config(config: &ConnectionConfig) -> ConnectionConfig { let mut db_config = config.canonicalized(); if database_capabilities::is_metadata_connection_scoped(&db_config.db_type) { @@ -263,6 +298,7 @@ impl AppState { Self { connections: Arc::new(RwLock::new(HashMap::new())), keepalive_tasks: Arc::new(RwLock::new(HashMap::new())), + pool_activity: Arc::new(RwLock::new(HashMap::new())), configs: RwLock::new(HashMap::new()), running_queries: RunningQueries::default(), tunnels: TunnelManager::new(), @@ -326,6 +362,7 @@ impl AppState { pub async fn insert_connection_pool(&self, pool_key: String, pool: PoolKind, config: &ConnectionConfig) { self.stop_keepalive_task(&pool_key).await; + self.pool_activity.write().await.insert(pool_key.clone(), PoolActivity::now()); self.start_keepalive_task(&pool_key, &pool, config).await; let previous = self.connections.write().await.insert(pool_key, pool); if let Some(pool) = previous { @@ -335,37 +372,77 @@ impl AppState { async fn start_keepalive_task(&self, pool_key: &str, pool: &PoolKind, config: &ConnectionConfig) { let interval_secs = config.keepalive_interval_secs; - if interval_secs == 0 { + let idle_timeout_secs = config.idle_timeout_secs; + let idle_cleanup_enabled = is_session_scoped_pool_key(pool_key) && idle_timeout_secs > 0; + let mut target = keepalive_target_from_pool(pool, config); + if interval_secs == 0 && !idle_cleanup_enabled { return; } - let Some(mut target) = keepalive_target_from_pool(pool, config) else { + if interval_secs > 0 && target.is_none() { log::debug!( "Connection keepalive requested for '{pool_key}', but this database driver does not keep a pingable client handle." ); - return; + if !idle_cleanup_enabled { + return; + } }; let key = pool_key.to_string(); - let interval = Duration::from_secs(interval_secs.max(1)); + let interval = if interval_secs > 0 { + Duration::from_secs(interval_secs.max(1)) + } else { + Duration::from_secs(idle_timeout_secs.min(60).max(1)) + }; let timeout = Duration::from_secs(config.effective_connect_timeout_secs().max(1)); let connections = self.connections.clone(); let keepalive_tasks = self.keepalive_tasks.clone(); + let pool_activity = self.pool_activity.clone(); + let running_queries = self.running_queries.clone(); + let idle_timeout = Duration::from_secs(idle_timeout_secs.max(1)); let handle = tokio::spawn(async move { loop { tokio::time::sleep(interval).await; - let result = tokio::time::timeout(timeout, ping_keepalive_target(&mut target, timeout)).await; - match result { - Ok(Ok(())) => {} - Ok(Err(err)) => { - log::warn!("Connection keepalive failed for '{key}': {err}; invalidating pool"); + + if running_queries.is_pool_active(&key) { + continue; + } + + if idle_cleanup_enabled { + let idle_for = { + let activity = pool_activity.read().await; + activity.get(&key).map(|activity| activity.last_used_at.elapsed()) + }; + if idle_for.is_some_and(|elapsed| elapsed >= idle_timeout) { + log::info!( + "Closing idle session-scoped connection pool '{key}' after {}s", + idle_timeout.as_secs() + ); keepalive_tasks.write().await.remove(&key); + pool_activity.write().await.remove(&key); let removed = connections.write().await.remove(&key); if let Some(pool) = removed { close_pool_kind_with_timeout(key, pool).await; } break; } - Err(_) => log::warn!("Connection keepalive timed out for '{key}' after {}s", timeout.as_secs()), + } + + if let Some(target) = target.as_mut() { + let result = tokio::time::timeout(timeout, ping_keepalive_target(target, timeout)).await; + match result { + Ok(Ok(())) => {} + Ok(Err(err)) => { + log::warn!("Connection keepalive failed for '{key}': {err}; invalidating pool"); + keepalive_tasks.write().await.remove(&key); + pool_activity.write().await.remove(&key); + let removed = connections.write().await.remove(&key); + if let Some(pool) = removed { + close_pool_kind_with_timeout(key, pool).await; + } + break; + } + Err(_) => log::warn!("Connection keepalive timed out for '{key}' after {}s", timeout.as_secs()), + } } } }); @@ -391,6 +468,18 @@ impl AppState { } } + pub async fn touch_pool_activity(&self, pool_key: &str) { + self.pool_activity.write().await.insert(pool_key.to_string(), PoolActivity::now()); + } + + pub fn pool_activity_touch(&self, pool_key: &str) -> PoolActivityTouch { + PoolActivityTouch { + pool_key: pool_key.to_string(), + connections: self.connections.clone(), + pool_activity: self.pool_activity.clone(), + } + } + 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 } @@ -413,6 +502,7 @@ impl AppState { if conns.contains_key(&pool_key) { drop(conns); if !self.remove_stale_connection_pool(&pool_key).await { + self.touch_pool_activity(&pool_key).await; return Ok(pool_key); } } else { @@ -904,6 +994,10 @@ impl AppState { } async fn remove_stale_connection_pool(&self, pool_key: &str) -> bool { + if self.running_queries.is_pool_active(pool_key) { + return false; + } + let stale = { let connections = self.connections.read().await; let Some(pool) = connections.get(pool_key) else { @@ -1085,6 +1179,7 @@ impl AppState { } self.stop_keepalive_task(pool_key).await; + self.pool_activity.write().await.remove(pool_key); let removed = self.connections.write().await.remove(pool_key); if let Some(pool) = removed { close_pool_kind(pool).await; @@ -1115,6 +1210,7 @@ impl AppState { self.reset_connection_transport(connection_id).await; } else { self.stop_keepalive_task(&pool_key).await; + self.pool_activity.write().await.remove(&pool_key); let removed = self.connections.write().await.remove(&pool_key); if let Some(pool) = removed { close_pool_kind(pool).await; @@ -1143,6 +1239,7 @@ impl AppState { return Ok(false); } self.stop_keepalive_task(&pool_key).await; + self.pool_activity.write().await.remove(&pool_key); let removed = self.connections.write().await.remove(&pool_key); if let Some(pool) = removed { close_pool_kind(pool).await; @@ -1154,6 +1251,7 @@ impl AppState { pub async fn remove_pool_by_key(&self, pool_key: &str) -> bool { self.stop_keepalive_task(pool_key).await; + self.pool_activity.write().await.remove(pool_key); let removed = self.connections.write().await.remove(pool_key); if let Some(pool) = removed { close_pool_kind(pool).await; @@ -1182,6 +1280,12 @@ impl AppState { .cloned() .collect(); self.stop_keepalive_tasks(&keys_to_remove).await; + { + let mut activity = self.pool_activity.write().await; + for key in &keys_to_remove { + activity.remove(key); + } + } let mut conns = self.connections.write().await; let mut removed = Vec::with_capacity(keys_to_remove.len()); for key in keys_to_remove { @@ -1472,6 +1576,12 @@ impl AppState { // Remove dead pools if !dead_keys.is_empty() { self.stop_keepalive_tasks(&dead_keys).await; + { + let mut activity = self.pool_activity.write().await; + for key in &dead_keys { + activity.remove(key); + } + } let mut conns = self.connections.write().await; for key in &dead_keys { if let Some(pool) = conns.remove(key) { @@ -1517,6 +1627,12 @@ impl AppState { .cloned() .collect(); self.stop_keepalive_tasks(&keys_to_remove).await; + { + let mut activity = self.pool_activity.write().await; + for key in &keys_to_remove { + activity.remove(key); + } + } let mut conns = self.connections.write().await; let mut removed = Vec::with_capacity(keys_to_remove.len()); for key in keys_to_remove { @@ -1545,6 +1661,12 @@ impl AppState { }) .collect(); self.stop_keepalive_tasks(&keys_to_remove).await; + { + let mut activity = self.pool_activity.write().await; + for key in &keys_to_remove { + activity.remove(key); + } + } let mut conns = self.connections.write().await; let mut removed = Vec::with_capacity(keys_to_remove.len()); for key in keys_to_remove { @@ -1697,6 +1819,10 @@ fn session_scoped_pool_key(base_pool_key: String, client_session_id: Option<&str .unwrap_or(base_pool_key) } +fn is_session_scoped_pool_key(pool_key: &str) -> bool { + pool_key.contains(":session:") +} + pub(crate) fn config_for_pool_key<'a>( pool_key: &str, configs: &'a HashMap, @@ -2763,6 +2889,23 @@ mod tests { ); } + #[test] + fn session_scoped_pool_keys_are_sanitized_and_detected() { + let key = super::session_scoped_pool_key_for( + Some(DatabaseType::Mysql), + "mysql-conn:analytics".to_string(), + Some("tab-1:count"), + ); + + assert_eq!(key, "mysql-conn:analytics:session:tab-1_count"); + assert!(super::is_session_scoped_pool_key(&key)); + assert!(!super::is_session_scoped_pool_key("mysql-conn:analytics")); + assert_eq!( + super::session_scoped_pool_key_for(Some(DatabaseType::DuckDb), "duckdb-conn".to_string(), Some("tab-1")), + "duckdb-conn" + ); + } + #[test] fn mysql_direct_connections_skip_tcp_probe() { let mut config = mysql_config(Some("app")); @@ -2884,6 +3027,36 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + #[tokio::test] + async fn pool_activity_touch_updates_existing_pool_only() { + let (state, dir) = test_app_state().await; + let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap(); + let pool_key = "conn:session:tab-1"; + + state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool)); + state.pool_activity.write().await.insert( + pool_key.to_string(), + super::PoolActivity { last_used_at: std::time::Instant::now() - std::time::Duration::from_secs(10) }, + ); + + { + let _touch = state.pool_activity_touch(pool_key); + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + let elapsed = state.pool_activity.read().await.get(pool_key).unwrap().last_used_at.elapsed(); + assert!(elapsed < std::time::Duration::from_secs(10)); + + { + let _touch = state.pool_activity_touch(pool_key); + state.connections.write().await.remove(pool_key); + state.pool_activity.write().await.remove(pool_key); + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + assert!(!state.pool_activity.read().await.contains_key(pool_key)); + + let _ = std::fs::remove_dir_all(dir); + } + #[cfg(feature = "duckdb-bundled")] #[tokio::test] async fn duckdb_client_session_reuses_base_pool_to_avoid_file_locks() { diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 6fb32027b..741081dee 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -805,6 +805,12 @@ pub async fn do_execute( cancel_token: Option, options: QueryExecutionOptions, ) -> Result { + if let Some(execution_id) = options.execution_id.as_deref() { + state.running_queries.set_pool_key(execution_id, pool_key.to_string()); + } + state.touch_pool_activity(pool_key).await; + let _activity_touch = state.pool_activity_touch(pool_key); + let query_timeout = resolve_query_timeout(options.timeout_secs); let (_duckdb_attached_names, conn_name_if_readonly) = { let configs = state.configs.read().await; @@ -1257,6 +1263,11 @@ async fn execute_postgres_drop_database( let pool_key = state .get_or_create_pool_for_session(connection_id, Some(admin_database), options.client_session_id.as_deref()) .await?; + if let Some(execution_id) = options.execution_id.as_deref() { + state.running_queries.set_pool_key(execution_id, pool_key.clone()); + } + state.touch_pool_activity(&pool_key).await; + let _activity_touch = state.pool_activity_touch(pool_key.as_str()); if is_canceled(&cancel_token) { return Err(canceled_error()); @@ -1390,6 +1401,11 @@ pub async fn execute_multi_core_with_options( .get_or_create_pool_for_session(connection_id, Some(database), options.client_session_id.as_deref()) .await? }; + if let Some(execution_id) = options.execution_id.as_deref() { + state.running_queries.set_pool_key(execution_id, pool_key.clone()); + } + 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; diff --git a/crates/dbx-core/src/query_cancel.rs b/crates/dbx-core/src/query_cancel.rs index 6193a6f42..e985b4162 100644 --- a/crates/dbx-core/src/query_cancel.rs +++ b/crates/dbx-core/src/query_cancel.rs @@ -4,16 +4,64 @@ use tokio_util::sync::CancellationToken; type InterruptFn = Box; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RunningTaskKind { + Query, + Count, + Explain, + Export, + Unknown, +} + +impl Default for RunningTaskKind { + fn default() -> Self { + Self::Unknown + } +} + +#[derive(Clone, Debug, Default)] +pub struct RunningTaskMetadata { + pub kind: RunningTaskKind, + pub connection_id: Option, + pub database: Option, + pub client_session_id: Option, +} + +impl RunningTaskMetadata { + pub fn query( + connection_id: impl Into, + database: impl Into, + client_session_id: Option, + ) -> Self { + let kind = task_kind_from_client_session_id(client_session_id.as_deref()); + Self { kind, connection_id: Some(connection_id.into()), database: Some(database.into()), client_session_id } + } +} + +#[derive(Clone)] +struct RunningTask { + token: CancellationToken, + metadata: RunningTaskMetadata, + pool_key: Option, +} + #[derive(Clone, Default)] pub struct RunningQueries { - inner: Arc>>, + inner: Arc>>, interrupts: Arc>>, } impl RunningQueries { pub fn register(&self, execution_id: String) -> RegisteredQuery { + self.register_task(execution_id, RunningTaskMetadata::default()) + } + + pub fn register_task(&self, execution_id: String, metadata: RunningTaskMetadata) -> RegisteredQuery { let token = CancellationToken::new(); - self.inner.lock().unwrap_or_else(|e| e.into_inner()).insert(execution_id.clone(), token.clone()); + self.inner + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(execution_id.clone(), RunningTask { token: token.clone(), metadata, pool_key: None }); RegisteredQuery { execution_id, token, running_queries: self.clone() } } @@ -23,7 +71,8 @@ impl RunningQueries { } pub fn cancel(&self, execution_id: &str) -> bool { - let token = self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).cloned(); + let token = + self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).map(|task| task.token.clone()); let interrupt = self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id); if let Some(interrupt) = interrupt { @@ -37,17 +86,51 @@ impl RunningQueries { } } + pub fn set_pool_key(&self, execution_id: &str, pool_key: impl Into) { + if let Some(task) = self.inner.lock().unwrap_or_else(|e| e.into_inner()).get_mut(execution_id) { + task.pool_key = Some(pool_key.into()); + } + } + + pub fn is_pool_active(&self, pool_key: &str) -> bool { + self.inner.lock().unwrap_or_else(|e| e.into_inner()).values().any(|task| { + let _kind = task.metadata.kind; + task.pool_key.as_deref() == Some(pool_key) + }) + } + #[cfg(test)] pub fn has(&self, execution_id: &str) -> bool { self.inner.lock().unwrap_or_else(|e| e.into_inner()).contains_key(execution_id) } + #[cfg(test)] + pub fn task_kind(&self, execution_id: &str) -> Option { + self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).map(|task| task.metadata.kind) + } + fn remove(&self, execution_id: &str) { self.inner.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id); self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id); } } +fn task_kind_from_client_session_id(client_session_id: Option<&str>) -> RunningTaskKind { + let Some(session_id) = client_session_id else { + return RunningTaskKind::Query; + }; + let session_id = session_id.trim().to_ascii_lowercase(); + if session_id.ends_with(":count") { + RunningTaskKind::Count + } else if session_id.ends_with(":explain") { + RunningTaskKind::Explain + } else if session_id.ends_with(":export") { + RunningTaskKind::Export + } else { + RunningTaskKind::Query + } +} + pub struct RegisteredQuery { execution_id: String, token: CancellationToken, @@ -68,7 +151,7 @@ impl Drop for RegisteredQuery { #[cfg(test)] mod tests { - use super::RunningQueries; + use super::{RunningQueries, RunningTaskKind, RunningTaskMetadata}; #[test] fn cancel_marks_registered_query_as_cancelled() { @@ -103,4 +186,20 @@ mod tests { assert!(!running.has("exec-1")); } + + #[test] + fn register_task_tracks_kind_and_pool_activity() { + let running = RunningQueries::default(); + let registered = running.register_task( + "exec-1".to_string(), + RunningTaskMetadata::query("conn-1", "main", Some("tab-1:export".to_string())), + ); + + running.set_pool_key("exec-1", "conn-1:main:session:tab-1_export"); + + assert_eq!(running.task_kind("exec-1"), Some(RunningTaskKind::Export)); + assert!(running.is_pool_active("conn-1:main:session:tab-1_export")); + drop(registered); + assert!(!running.is_pool_active("conn-1:main:session:tab-1_export")); + } } diff --git a/crates/dbx-web/src/routes/query.rs b/crates/dbx-web/src/routes/query.rs index 4aabf8cc9..fbde51303 100644 --- a/crates/dbx-web/src/routes/query.rs +++ b/crates/dbx-web/src/routes/query.rs @@ -6,6 +6,7 @@ use serde::Deserialize; use crate::error::AppError; use crate::state::WebState; +use dbx_core::query_cancel::RunningTaskMetadata; #[derive(Deserialize)] #[serde(rename_all = "camelCase")] @@ -269,7 +270,10 @@ pub async fn execute_query( ) -> Result, AppError> { let execution_id = req.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let registered = state.app.running_queries.register(execution_id.clone()); + let registered = state.app.running_queries.register_task( + execution_id.clone(), + RunningTaskMetadata::query(req.connection_id.clone(), req.database.clone(), req.client_session_id.clone()), + ); let cancel_token = registered.token(); let result = dbx_core::query::execute_sql_statement_with_options( @@ -302,7 +306,10 @@ pub async fn execute_multi( ) -> Result, AppError> { let execution_id = req.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let registered = state.app.running_queries.register(execution_id.clone()); + let registered = state.app.running_queries.register_task( + execution_id.clone(), + RunningTaskMetadata::query(req.connection_id.clone(), req.database.clone(), req.client_session_id.clone()), + ); let cancel_token = registered.token(); let result = dbx_core::query::execute_multi_core_with_options( diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index 4cd283c6d..dd366f78b 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -5,6 +5,7 @@ use tauri::State; use crate::commands::connection::AppState; use dbx_core::db; use dbx_core::models::connection::DatabaseType; +use dbx_core::query_cancel::RunningTaskMetadata; use dbx_core::sql::split_sql_statements; #[tauri::command] @@ -23,8 +24,13 @@ pub async fn execute_query( client_session_id: Option, timeout_secs: Option, ) -> Result { - let registered_query = - execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); + let execution_id = execution_id.filter(|id| !id.trim().is_empty()); + let registered_query = execution_id.as_ref().map(|id| { + state.running_queries.register_task( + id.clone(), + RunningTaskMetadata::query(connection_id.clone(), database.clone(), client_session_id.clone()), + ) + }); let cancel_token = registered_query.as_ref().map(|query| query.token()); dbx_core::query::execute_sql_statement_with_options( @@ -41,7 +47,7 @@ pub async fn execute_query( result_session_id, client_session_id, timeout_secs, - execution_id: execution_id.filter(|id| !id.trim().is_empty()), + execution_id, }, ) .await @@ -63,10 +69,15 @@ pub async fn execute_multi( client_session_id: Option, timeout_secs: Option, ) -> Result, String> { - let registered_query = - execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); + let execution_id = execution_id.filter(|id| !id.trim().is_empty()); + let registered_query = execution_id.as_ref().map(|id| { + state.running_queries.register_task( + id.clone(), + RunningTaskMetadata::query(connection_id.clone(), database.clone(), client_session_id.clone()), + ) + }); let cancel_token = registered_query.as_ref().map(|query| query.token()); - let trace_id = execution_id.clone().as_deref().unwrap_or("no-execution-id").to_string(); + let trace_id = execution_id.as_deref().unwrap_or("no-execution-id").to_string(); let started_at = Instant::now(); log::info!( "[query][execute_multi:start] trace_id={} connection_id={} database={} schema={:?} sql={}", @@ -91,7 +102,7 @@ pub async fn execute_multi( result_session_id, client_session_id, timeout_secs, - execution_id: execution_id.filter(|id| !id.trim().is_empty()), + execution_id, }, ) .await;