From ae71272d02164b0d3c9334bec6b1386a49c65fdf Mon Sep 17 00:00:00 2001 From: ZhuchkaTriplesix Date: Tue, 29 Sep 2026 09:58:23 +0300 Subject: [PATCH] feat(rpc): support Command Palette actions via commands.execute and manifest contributions. Closes #48 Querya Desktop's Command Palette invokes commands.execute for actions declared under contributions.commands in manifest.json, but the driver neither declared any commands nor implemented the method (returning -32601 Method not found). - manifest.json now declares four ClickHouse maintenance commands: clickhouse.optimizeFinal, clickhouse.deduplicate, clickhouse.serverStats, clickhouse.dropPartition. - New src/rpc/handlers/commands.rs implements handle_execute, dispatching on commandId. clickhouse.serverStats delegates straight to handle_get_server_stats. The other three build the same SQL as their sdui.contextActions equivalents (OPTIMIZE ... FINAL / DEDUPLICATE, ALTER ... DROP PARTITION) and run it through handle_query, reusing existing Safe Mode and injection-safe quoting/escaping. - The Command Palette doesn't yet forward workspace selection (nodeId) to drivers, so database/table/partition are accepted as optional params; a command that needs a target and doesn't get one returns a clear -32602 error instead of guessing or operating on the wrong object. - Registered "commands.execute" in router.rs. --- manifest.json | 22 +++ src/rpc/handlers/commands.rs | 284 +++++++++++++++++++++++++++++++++++ src/rpc/handlers/mod.rs | 1 + src/rpc/router.rs | 25 ++- 4 files changed, 331 insertions(+), 1 deletion(-) create mode 100644 src/rpc/handlers/commands.rs diff --git a/manifest.json b/manifest.json index 6216fd1..1bc1c3f 100644 --- a/manifest.json +++ b/manifest.json @@ -41,6 +41,28 @@ "defaultPort": 8123, "connectionFormSchema": "assets/connection_form.json" } + ], + "commands": [ + { + "id": "clickhouse.optimizeFinal", + "title": "ClickHouse: Optimize Table (FINAL)", + "category": "ClickHouse" + }, + { + "id": "clickhouse.deduplicate", + "title": "ClickHouse: Deduplicate Table", + "category": "ClickHouse" + }, + { + "id": "clickhouse.serverStats", + "title": "ClickHouse: Show Server Statistics", + "category": "ClickHouse" + }, + { + "id": "clickhouse.dropPartition", + "title": "ClickHouse: Drop Selected Partition", + "category": "ClickHouse" + } ] } } diff --git a/src/rpc/handlers/commands.rs b/src/rpc/handlers/commands.rs new file mode 100644 index 0000000..34f65ce --- /dev/null +++ b/src/rpc/handlers/commands.rs @@ -0,0 +1,284 @@ +use crate::error::DriverError; +use crate::rpc::handlers::{query, schema}; +use crate::utils::sql_escape::{escape_sql_string_literal, quote_identifier}; +use serde::Deserialize; +use serde_json::{Value, json}; +use tracing::info; + +/// Parameters for `commands.execute`, sent by Querya Desktop's Command Palette +/// when the user runs one of this driver's `contributions.commands` entries. +/// +/// The Command Palette does not yet forward workspace selection (selected +/// table/partition) to the driver, so `database`/`table`/`partition` are +/// optional here and validated per-command: a command that needs a target +/// returns a clear `-32602` error explaining what's missing instead of +/// silently operating on the wrong object. +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExecuteCommandParams { + pub command_id: String, + pub connection_id: u64, + #[serde(default)] + pub database: Option, + #[serde(default)] + pub table: Option, + #[serde(default)] + pub partition: Option, +} + +/// Handler for `commands.execute`. Dispatches a Command Palette invocation +/// (`commandId`) to the matching ClickHouse maintenance operation. +pub async fn handle_execute(params: Option) -> Result { + let params_val = params.ok_or_else(|| DriverError::Rpc { + code: -32602, + message: "Invalid params: commands.execute requires commandId and connectionId".to_string(), + data: None, + })?; + + let p: ExecuteCommandParams = + serde_json::from_value(params_val).map_err(|e| DriverError::Rpc { + code: -32602, + message: format!("Malformed commands.execute parameters: {}", e), + data: None, + })?; + + info!( + "Executing command '{}' on connectionId={}", + p.command_id, p.connection_id + ); + + match p.command_id.as_str() { + "clickhouse.serverStats" => { + schema::handle_get_server_stats(Some(json!({ "connectionId": p.connection_id }))).await + } + "clickhouse.optimizeFinal" => { + let (db, tbl) = require_table_target(&p)?; + run_sql( + p.connection_id, + format!( + "OPTIMIZE TABLE {}.{} FINAL", + quote_identifier(db), + quote_identifier(tbl) + ), + ) + .await + } + "clickhouse.deduplicate" => { + let (db, tbl) = require_table_target(&p)?; + run_sql( + p.connection_id, + format!( + "OPTIMIZE TABLE {}.{} DEDUPLICATE", + quote_identifier(db), + quote_identifier(tbl) + ), + ) + .await + } + "clickhouse.dropPartition" => { + let (db, tbl) = require_table_target(&p)?; + let partition = p + .partition + .as_deref() + .filter(|s| !s.is_empty()) + .ok_or_else(|| DriverError::Rpc { + code: -32602, + message: + "commandId 'clickhouse.dropPartition' also requires a 'partition' parameter" + .to_string(), + data: None, + })?; + run_sql( + p.connection_id, + format!( + "ALTER TABLE {}.{} DROP PARTITION '{}'", + quote_identifier(db), + quote_identifier(tbl), + escape_sql_string_literal(partition) + ), + ) + .await + } + other => Err(DriverError::Rpc { + code: -32602, + message: format!("Unknown commandId: '{}'", other), + data: None, + }), + } +} + +fn require_table_target(p: &ExecuteCommandParams) -> Result<(&str, &str), DriverError> { + match ( + p.database.as_deref().filter(|s| !s.is_empty()), + p.table.as_deref().filter(|s| !s.is_empty()), + ) { + (Some(db), Some(tbl)) => Ok((db, tbl)), + _ => Err(DriverError::Rpc { + code: -32602, + message: format!( + "commandId '{}' requires 'database' and 'table' parameters (select a table first)", + p.command_id + ), + data: None, + }), + } +} + +async fn run_sql(connection_id: u64, sql: String) -> Result { + query::handle_query(Some(json!({ + "connectionId": connection_id, + "sql": sql + }))) + .await +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::driver::client::{ClickHouseClient, ConnectParams}; + use crate::driver::pool::ConnectionPool; + + #[tokio::test] + async fn test_handle_execute_server_stats() { + let _guard = crate::utils::test_lock::GLOBAL_TEST_LOCK.lock().await; + let client = ClickHouseClient::from_params(ConnectParams { + connection_id: 501, + connection_string: Some("mock://localhost:8123/default".to_string()), + ..Default::default() + }) + .unwrap(); + ConnectionPool::global().insert(client); + + let res = handle_execute(Some(json!({ + "commandId": "clickhouse.serverStats", + "connectionId": 501 + }))) + .await + .unwrap(); + assert_eq!(res["serverVersion"], "ClickHouse 24.3 (Mock)"); + + ConnectionPool::global().remove(501); + } + + #[tokio::test] + async fn test_handle_execute_optimize_final_and_deduplicate() { + let _guard = crate::utils::test_lock::GLOBAL_TEST_LOCK.lock().await; + let client = ClickHouseClient::from_params(ConnectParams { + connection_id: 502, + connection_string: Some("mock://localhost:8123/default".to_string()), + ..Default::default() + }) + .unwrap(); + ConnectionPool::global().insert(client); + + let res = handle_execute(Some(json!({ + "commandId": "clickhouse.optimizeFinal", + "connectionId": 502, + "database": "analytics", + "table": "events" + }))) + .await + .unwrap(); + assert_eq!(res["operation"], "optimize"); + + let res = handle_execute(Some(json!({ + "commandId": "clickhouse.deduplicate", + "connectionId": 502, + "database": "analytics", + "table": "events" + }))) + .await + .unwrap(); + assert_eq!(res["operation"], "optimize"); + + ConnectionPool::global().remove(502); + } + + #[tokio::test] + async fn test_handle_execute_drop_partition() { + let _guard = crate::utils::test_lock::GLOBAL_TEST_LOCK.lock().await; + let client = ClickHouseClient::from_params(ConnectParams { + connection_id: 503, + connection_string: Some("mock://localhost:8123/default".to_string()), + readonly: Some(false), + ..Default::default() + }) + .unwrap(); + ConnectionPool::global().insert(client); + + let res = handle_execute(Some(json!({ + "commandId": "clickhouse.dropPartition", + "connectionId": 503, + "database": "analytics", + "table": "events", + "partition": "202607" + }))) + .await + .unwrap(); + assert_eq!(res["operation"], "alter"); + + ConnectionPool::global().remove(503); + } + + #[tokio::test] + async fn test_handle_execute_requires_table_target() { + let _guard = crate::utils::test_lock::GLOBAL_TEST_LOCK.lock().await; + let client = ClickHouseClient::from_params(ConnectParams { + connection_id: 504, + connection_string: Some("mock://localhost:8123/default".to_string()), + ..Default::default() + }) + .unwrap(); + ConnectionPool::global().insert(client); + + let err = handle_execute(Some(json!({ + "commandId": "clickhouse.optimizeFinal", + "connectionId": 504 + }))) + .await + .unwrap_err(); + assert_eq!(err.to_rpc_code(), -32602); + assert!(err.to_string().contains("requires 'database' and 'table'")); + + let err = handle_execute(Some(json!({ + "commandId": "clickhouse.dropPartition", + "connectionId": 504, + "database": "analytics", + "table": "events" + }))) + .await + .unwrap_err(); + assert!(err.to_string().contains("also requires a 'partition'")); + + ConnectionPool::global().remove(504); + } + + #[tokio::test] + async fn test_handle_execute_unknown_command_id() { + let _guard = crate::utils::test_lock::GLOBAL_TEST_LOCK.lock().await; + let client = ClickHouseClient::from_params(ConnectParams { + connection_id: 505, + connection_string: Some("mock://localhost:8123/default".to_string()), + ..Default::default() + }) + .unwrap(); + ConnectionPool::global().insert(client); + + let err = handle_execute(Some(json!({ + "commandId": "clickhouse.doesNotExist", + "connectionId": 505 + }))) + .await + .unwrap_err(); + assert_eq!(err.to_rpc_code(), -32602); + assert!(err.to_string().contains("Unknown commandId")); + + ConnectionPool::global().remove(505); + } + + #[tokio::test] + async fn test_handle_execute_missing_params() { + let err = handle_execute(None).await.unwrap_err(); + assert_eq!(err.to_rpc_code(), -32602); + } +} diff --git a/src/rpc/handlers/mod.rs b/src/rpc/handlers/mod.rs index c5e083f..eeaf771 100644 --- a/src/rpc/handlers/mod.rs +++ b/src/rpc/handlers/mod.rs @@ -1,3 +1,4 @@ +pub mod commands; pub mod connection; pub mod query; pub mod schema; diff --git a/src/rpc/router.rs b/src/rpc/router.rs index 51097c8..9422a9f 100644 --- a/src/rpc/router.rs +++ b/src/rpc/router.rs @@ -1,5 +1,5 @@ use crate::error::DriverError; -use crate::rpc::handlers::{connection, query, schema, system}; +use crate::rpc::handlers::{commands, connection, query, schema, system}; use serde_json::Value; pub async fn dispatch(method: &str, params: Option) -> Result { @@ -22,6 +22,7 @@ pub async fn dispatch(method: &str, params: Option) -> Result { schema::handle_get_object_metadata(params).await } + "commands.execute" => commands::handle_execute(params).await, _ => Err(DriverError::Rpc { code: -32601, message: format!("Method not found: {}", method), @@ -91,6 +92,28 @@ mod tests { ConnectionPool::global().remove(889); } + #[tokio::test] + async fn test_dispatch_commands_execute() { + let _guard = crate::utils::test_lock::GLOBAL_TEST_LOCK.lock().await; + let client = ClickHouseClient::from_params(ConnectParams { + connection_id: 890, + connection_string: Some("mock://localhost:8123/default".to_string()), + ..Default::default() + }) + .unwrap(); + ConnectionPool::global().insert(client); + + let params = json!({ + "commandId": "clickhouse.serverStats", + "connectionId": 890 + }); + + let res = dispatch("commands.execute", Some(params)).await.unwrap(); + assert_eq!(res["serverVersion"], "ClickHouse 24.3 (Mock)"); + + ConnectionPool::global().remove(890); + } + #[tokio::test] async fn test_dispatch_system_handshake_and_ping() { let handshake_res = dispatch("system.handshake", None).await.unwrap();