diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 138ef4ecb8..f942909ce7 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use tauri::{AppHandle, Manager, State}; use tokio::sync::Mutex; +use crate::commands::query_cancel::RunningQueries; use crate::db; use crate::db::ssh_tunnel::TunnelManager; use crate::models::connection::{ConnectionConfig, DatabaseType}; @@ -23,6 +24,7 @@ pub enum PoolKind { pub struct AppState { pub connections: Mutex>, pub configs: Mutex>, + pub running_queries: RunningQueries, pub tunnels: TunnelManager, } @@ -31,6 +33,7 @@ impl AppState { Self { connections: Mutex::new(HashMap::new()), configs: Mutex::new(HashMap::new()), + running_queries: RunningQueries::default(), tunnels: TunnelManager::new(), } } diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index 091d063d1f..f0bc01bf5a 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -3,6 +3,7 @@ pub mod connection; pub mod history; pub mod mongo_cmd; pub mod query; +pub mod query_cancel; pub mod redis_cmd; pub mod schema; pub mod transfer; diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index d4fcc2046c..596c6bfd86 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -1,13 +1,16 @@ +use std::future::Future; use std::sync::Arc; use std::time::Duration; use tauri::State; use tokio::time::timeout; +use tokio_util::sync::CancellationToken; use crate::commands::connection::{AppState, PoolKind}; use crate::db; const QUERY_TIMEOUT: Duration = Duration::from_secs(30); const MAX_ROWS: usize = 10000; +const QUERY_CANCELED: &str = "Query canceled"; fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result { let start = std::time::Instant::now(); @@ -68,10 +71,46 @@ fn is_connection_error(err: &str) -> bool { || lower.contains("eof") } +fn timeout_error() -> String { + format!("Query timed out after {} seconds", QUERY_TIMEOUT.as_secs()) +} + +fn canceled_error() -> String { + QUERY_CANCELED.to_string() +} + +fn is_canceled(cancel_token: &Option) -> bool { + cancel_token + .as_ref() + .map(|token| token.is_cancelled()) + .unwrap_or(false) +} + +async fn wait_for_query( + cancel_token: Option, + future: F, +) -> Result +where + F: Future>, +{ + if let Some(token) = cancel_token { + tokio::select! { + biased; + _ = token.cancelled() => Err(canceled_error()), + result = timeout(QUERY_TIMEOUT, future) => result.map_err(|_| timeout_error())?, + } + } else { + timeout(QUERY_TIMEOUT, future) + .await + .map_err(|_| timeout_error())? + } +} + async fn do_execute( state: &AppState, pool_key: &str, sql: &str, + cancel_token: Option, ) -> Result { let connections = state.connections.lock().await; let pool = connections.get(pool_key).ok_or("Connection not found")?; @@ -81,64 +120,72 @@ async fn do_execute( let con = con.clone(); let sql = sql.to_string(); drop(connections); - let task = tokio::task::spawn_blocking(move || { - let con = con.lock().map_err(|e| e.to_string())?; - duckdb_execute(&con, &sql) - }); - timeout(QUERY_TIMEOUT, task) - .await - .map_err(|_| format!("Query timed out after {} seconds", QUERY_TIMEOUT.as_secs()))? - .map_err(|e| e.to_string())? + wait_for_query(cancel_token, async move { + let task = tokio::task::spawn_blocking(move || { + let con = con.lock().map_err(|e| e.to_string())?; + duckdb_execute(&con, &sql) + }); + task.await.map_err(|e| e.to_string())? + }) + .await } PoolKind::Mysql(p) => { let p = p.clone(); drop(connections); - timeout(QUERY_TIMEOUT, db::mysql::execute_query(&p, sql)) + wait_for_query(cancel_token, db::mysql::execute_query(&p, sql)) .await - .map_err(|_| format!("Query timed out after {} seconds", QUERY_TIMEOUT.as_secs()))? .map(truncate_result) } PoolKind::Postgres(p) => { let p = p.clone(); drop(connections); - timeout(QUERY_TIMEOUT, db::postgres::execute_query(&p, sql)) + wait_for_query(cancel_token, db::postgres::execute_query(&p, sql)) .await - .map_err(|_| format!("Query timed out after {} seconds", QUERY_TIMEOUT.as_secs()))? .map(truncate_result) } PoolKind::Sqlite(p) => { let p = p.clone(); drop(connections); - timeout(QUERY_TIMEOUT, db::sqlite::execute_query(&p, sql)) + wait_for_query(cancel_token, db::sqlite::execute_query(&p, sql)) .await - .map_err(|_| format!("Query timed out after {} seconds", QUERY_TIMEOUT.as_secs()))? .map(truncate_result) } PoolKind::ClickHouse(client) => { let client = client.clone(); let database = pool_key.split(':').nth(1).unwrap_or("default").to_string(); drop(connections); - timeout(QUERY_TIMEOUT, db::clickhouse_driver::execute_query(&client, &database, sql)) + wait_for_query(cancel_token, db::clickhouse_driver::execute_query(&client, &database, sql)) .await - .map_err(|_| format!("Query timed out after {} seconds", QUERY_TIMEOUT.as_secs()))? .map(truncate_result) } PoolKind::SqlServer(client) => { let client = client.clone(); drop(connections); - let mut client = client.lock().await; - timeout(QUERY_TIMEOUT, db::sqlserver::execute_query(&mut client, sql)) + let mut client = match cancel_token.as_ref() { + Some(token) => tokio::select! { + biased; + _ = token.cancelled() => return Err(canceled_error()), + guard = client.lock() => guard, + }, + None => client.lock().await, + }; + wait_for_query(cancel_token, db::sqlserver::execute_query(&mut client, sql)) .await - .map_err(|_| format!("Query timed out after {} seconds", QUERY_TIMEOUT.as_secs()))? .map(truncate_result) } PoolKind::Oracle(client) => { let client = client.clone(); drop(connections); - let client = client.lock().await; - timeout(QUERY_TIMEOUT, db::oracle_driver::execute_query(&*client, sql)) + let client = match cancel_token.as_ref() { + Some(token) => tokio::select! { + biased; + _ = token.cancelled() => return Err(canceled_error()), + guard = client.lock() => guard, + }, + None => client.lock().await, + }; + wait_for_query(cancel_token, db::oracle_driver::execute_query(&*client, sql)) .await - .map_err(|_| format!("Query timed out after {} seconds", QUERY_TIMEOUT.as_secs()))? .map(truncate_result) } PoolKind::Elasticsearch(_) => Err("Use document browser for Elasticsearch".to_string()), @@ -153,25 +200,44 @@ pub async fn execute_query( connection_id: String, database: String, sql: String, + execution_id: Option, ) -> Result { + let registered_query = execution_id + .as_ref() + .filter(|id| !id.trim().is_empty()) + .map(|id| state.running_queries.register(id.clone())); + let cancel_token = registered_query.as_ref().map(|query| query.token()); + let pool_key = if database.is_empty() { connection_id.clone() } else { state.get_or_create_pool(&connection_id, Some(&database)).await? }; - let result = do_execute(&state, &pool_key, &sql).await; + if is_canceled(&cancel_token) { + return Err(canceled_error()); + } + + let result = do_execute(&state, &pool_key, &sql, cancel_token.clone()).await; match &result { - Err(e) if is_connection_error(e) => { + Err(e) if is_connection_error(e) && !is_canceled(&cancel_token) => { let db_opt = if database.is_empty() { None } else { Some(database.as_str()) }; let new_key = state.reconnect_pool(&connection_id, db_opt).await?; - do_execute(&state, &new_key, &sql).await + do_execute(&state, &new_key, &sql, cancel_token).await } _ => result, } } +#[tauri::command] +pub async fn cancel_query( + state: State<'_, Arc>, + execution_id: String, +) -> Result { + Ok(state.running_queries.cancel(&execution_id)) +} + #[tauri::command] pub async fn execute_batch( state: State<'_, Arc>, @@ -189,7 +255,7 @@ pub async fn execute_batch( let start = std::time::Instant::now(); for (i, sql) in statements.iter().enumerate() { - match do_execute(&state, &pool_key, sql).await { + match do_execute(&state, &pool_key, sql, None).await { Ok(result) => { total_affected += result.affected_rows; } @@ -214,3 +280,28 @@ pub async fn execute_batch( truncated: false, }) } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn wait_for_query_returns_cancelled_when_token_is_cancelled() { + let token = CancellationToken::new(); + token.cancel(); + + let result = wait_for_query(Some(token), async { + tokio::time::sleep(Duration::from_secs(30)).await; + Ok(db::QueryResult { + columns: vec![], + rows: vec![], + affected_rows: 0, + execution_time_ms: 0, + truncated: false, + }) + }) + .await; + + assert_eq!(result.unwrap_err(), QUERY_CANCELED); + } +} diff --git a/src-tauri/src/commands/query_cancel.rs b/src-tauri/src/commands/query_cancel.rs new file mode 100644 index 0000000000..bdb81ae2a6 --- /dev/null +++ b/src-tauri/src/commands/query_cancel.rs @@ -0,0 +1,97 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use tokio_util::sync::CancellationToken; + +#[derive(Clone, Default)] +pub struct RunningQueries { + inner: Arc>>, +} + +impl RunningQueries { + pub fn register(&self, execution_id: String) -> RegisteredQuery { + let token = CancellationToken::new(); + self.inner + .lock() + .expect("running query registry poisoned") + .insert(execution_id.clone(), token.clone()); + + RegisteredQuery { + execution_id, + token, + running_queries: self.clone(), + } + } + + pub fn cancel(&self, execution_id: &str) -> bool { + let token = self + .inner + .lock() + .expect("running query registry poisoned") + .get(execution_id) + .cloned(); + + if let Some(token) = token { + token.cancel(); + true + } else { + false + } + } + + pub fn has(&self, execution_id: &str) -> bool { + self.inner + .lock() + .expect("running query registry poisoned") + .contains_key(execution_id) + } + + fn remove(&self, execution_id: &str) { + self.inner + .lock() + .expect("running query registry poisoned") + .remove(execution_id); + } +} + +pub struct RegisteredQuery { + execution_id: String, + token: CancellationToken, + running_queries: RunningQueries, +} + +impl RegisteredQuery { + pub fn token(&self) -> CancellationToken { + self.token.clone() + } +} + +impl Drop for RegisteredQuery { + fn drop(&mut self) { + self.running_queries.remove(&self.execution_id); + } +} + +#[cfg(test)] +mod tests { + use super::RunningQueries; + + #[test] + fn cancel_marks_registered_query_as_cancelled() { + let running = RunningQueries::default(); + let registered = running.register("exec-1".to_string()); + + assert!(running.cancel("exec-1")); + assert!(registered.token().is_cancelled()); + } + + #[test] + fn dropping_registration_removes_running_query() { + let running = RunningQueries::default(); + let registered = running.register("exec-1".to_string()); + + assert!(running.has("exec-1")); + drop(registered); + + assert!(!running.has("exec-1")); + } +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 2094e812f3..ffee5939da 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -54,6 +54,7 @@ pub fn run() { commands::schema::list_triggers, commands::schema::get_table_ddl, commands::query::execute_query, + commands::query::cancel_query, commands::query::execute_batch, commands::redis_cmd::redis_list_databases, commands::redis_cmd::redis_scan_keys, diff --git a/src/App.vue b/src/App.vue index d6a9709a19..9d3050e2ff 100644 --- a/src/App.vue +++ b/src/App.vue @@ -1,7 +1,7 @@