From 8d590e4682d7d65159813fec804dfc43478837f5 Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 7 Aug 2026 20:11:03 +0800 Subject: [PATCH 1/4] fix(mcp): resolve saved-but-inactive connections in bridge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MCP bridge resolve_connection only checked in-memory connections, so connections saved in .store.dat but not yet activated in the UI session failed with misleading errors. Reuse capabilities::sql::resolve_adapter, which reads the store in-process and connects — credentials never leave the app. McpPolicy authorization still runs upstream in invoke_with_policy. Also make get_connection_id errors actionable: distinguish 'no connection provided' (point to sqlkit__list_connections / Settings → MCP Bridge) from a malformed internal config. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src-tauri/src/capabilities/sql.rs | 43 ++++++++++++++++++++++++++----- src-tauri/src/mcp_bridge.rs | 9 +++---- 2 files changed, 40 insertions(+), 12 deletions(-) diff --git a/src-tauri/src/capabilities/sql.rs b/src-tauri/src/capabilities/sql.rs index 6ce0ce0..b531696 100644 --- a/src-tauri/src/capabilities/sql.rs +++ b/src-tauri/src/capabilities/sql.rs @@ -20,7 +20,7 @@ fn app_handle() -> AppHandle { .clone() } -async fn resolve_adapter(connection_id: &str) -> Result { +pub(crate) async fn resolve_adapter(connection_id: &str) -> Result { let app = app_handle(); // Check if already connected @@ -132,11 +132,18 @@ async fn execute_on_adapter(adapter: &ActiveConnection, sql: &str) -> Result) -> Result { - config - .and_then(|c| c.get("connectionId")) - .and_then(|v| v.as_str()) - .map(|s| s.to_string()) - .ok_or_else(|| "Missing connectionId in connection config".to_string()) + match config { + None => Err( + "No connection was provided for this tool call. Supply a connection_id \ + (list them with sqlkit__list_connections) or enable it in Settings → MCP Bridge" + .to_string(), + ), + Some(c) => c + .get("connectionId") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .ok_or_else(|| "Connection config is missing the 'connectionId' field".to_string()), + } } // --------------------------------------------------------------------------- @@ -777,3 +784,27 @@ pub fn register_sql_tools(reg: &mut CapabilityRegistry) { parallel_ok: true, }); } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn get_connection_id_returns_id_when_present() { + let config = json!({ "connectionId": "conn-1" }); + assert_eq!(get_connection_id(Some(&config)), Ok("conn-1".to_string())); + } + + #[test] + fn get_connection_id_explains_missing_config() { + let err = get_connection_id(None).unwrap_err(); + assert!(err.contains("connection_id"), "got: {}", err); + assert!(err.contains("Settings → MCP Bridge"), "got: {}", err); + } + + #[test] + fn get_connection_id_rejects_config_without_field() { + let err = get_connection_id(Some(&json!({ "host": "x" }))).unwrap_err(); + assert!(err.contains("connectionId"), "got: {}", err); + } +} diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs index 72b6c07..d35dddc 100644 --- a/src-tauri/src/mcp_bridge.rs +++ b/src-tauri/src/mcp_bridge.rs @@ -506,12 +506,9 @@ fn to_metadata(cap: &Capability) -> Value { } async fn resolve_connection(connection_id: &str) -> Result { - let handle = crate::APP_HANDLE - .get() - .ok_or_else(|| "AppHandle not initialized".to_string())?; - use tauri::State; - let state: State<'_, crate::state::AppState> = handle.state(); - state.ensure_connection(connection_id).await?; + // Shared resolver: connects saved-but-inactive connections in-process; + // credentials never leave the app. McpPolicy auth already ran upstream. + crate::capabilities::sql::resolve_adapter(connection_id).await?; Ok(json!({ "connectionId": connection_id })) } From ab2594ac428ef6ef0d4c4ab829ac6240e0869480 Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 7 Aug 2026 20:11:11 +0800 Subject: [PATCH 2/4] chore: sync Cargo.lock sqlkit version to 0.8.4 Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src-tauri/Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 702c9e1..04b74b4 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -6509,7 +6509,7 @@ dependencies = [ [[package]] name = "sqlkit" -version = "0.8.1" +version = "0.8.4" dependencies = [ "async-trait", "axum", From 367f0c42d1499922719a431509845eda2a7c9815 Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 7 Aug 2026 20:47:57 +0800 Subject: [PATCH 3/4] fix(mcp): remove get_store_value stub capability MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The handler always returned {"value": null} without reading the store — dead code that misleads agents. MCP should not expose internal app store data, so remove the capability and its registration. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src-tauri/src/capabilities/sqlkit.rs | 42 ---------------------------- 1 file changed, 42 deletions(-) diff --git a/src-tauri/src/capabilities/sqlkit.rs b/src-tauri/src/capabilities/sqlkit.rs index 4a42115..bc9f4aa 100644 --- a/src-tauri/src/capabilities/sqlkit.rs +++ b/src-tauri/src/capabilities/sqlkit.rs @@ -88,27 +88,6 @@ impl CapabilityHandler for ListConnections { } } -// --------------------------------------------------------------------------- -// GetStoreValue handler (kept from original) -// --------------------------------------------------------------------------- - -struct GetStoreValueHandler; - -#[async_trait::async_trait] -impl CapabilityHandler for GetStoreValueHandler { - async fn handle( - &self, - args: &Value, - _connection_config: Option<&Value>, - ) -> Result { - let key = args - .get("key") - .and_then(|v| v.as_str()) - .ok_or_else(|| "Missing 'key' argument".to_string())?; - Ok(serde_json::json!({ "key": key, "value": null }).to_string()) - } -} - // --------------------------------------------------------------------------- // Registration // --------------------------------------------------------------------------- @@ -130,27 +109,6 @@ pub(crate) fn register_all(reg: &mut CapabilityRegistry) { tags: &["agent"], parallel_ok: true, }); - - reg.register(Capability { - name: "sqlkit__get_store_value", - description: "Get a value from the persistent key-value store.", - handler: Arc::new(GetStoreValueHandler), - input_schema: serde_json::json!({ - "type": "object", - "properties": { - "key": { - "type": "string", - "description": "The key to look up" - } - }, - "required": ["key"] - }), - risk_level: RiskLevel::Safe, - required_permission: "read", - source_kind: SourceKind::AppLocal, - tags: &["agent"], - parallel_ok: true, - }); } // --------------------------------------------------------------------------- From 47747a0830b64cbfb47b44726286629890e5fee5 Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 7 Aug 2026 20:48:05 +0800 Subject: [PATCH 4/4] feat(mcp): split SQL write tools by risk level MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sqlkit__execute_query was a single Elevated capability carrying all SQL — INSERT/UPDATE/DELETE/DDL ran with no Destructive gate, and parallel_ok blocked concurrent reads. Split by statement class so McpPolicy can gate each risk level: - sqlkit__execute_query: read-only (SELECT/SHOW/EXPLAIN) → Safe, parallel - sqlkit__execute_write: INSERT/UPDATE/MERGE → Elevated - sqlkit__execute_delete: DELETE/TRUNCATE → Destructive - sqlkit__execute_ddl: CREATE/ALTER/DROP → Destructive classify_sql parses with sqlparser (dialect-aware). execute_query now rejects write/delete/ddl statements with actionable guidance pointing to the split tools. New module sql_write.rs shares resolve_adapter and execute_on_adapter from sql.rs. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src-tauri/src/capabilities/mod.rs | 1 + src-tauri/src/capabilities/sql.rs | 17 +- src-tauri/src/capabilities/sql_write.rs | 437 ++++++++++++++++++++++++ src-tauri/src/lib.rs | 1 + src-tauri/src/mcp_bridge.rs | 1 + 5 files changed, 451 insertions(+), 6 deletions(-) create mode 100644 src-tauri/src/capabilities/sql_write.rs diff --git a/src-tauri/src/capabilities/mod.rs b/src-tauri/src/capabilities/mod.rs index dfcb172..56d3881 100644 --- a/src-tauri/src/capabilities/mod.rs +++ b/src-tauri/src/capabilities/mod.rs @@ -2,6 +2,7 @@ pub mod commands; pub mod mysql; pub mod postgres; pub mod sql; +pub mod sql_write; pub mod sqlite; pub mod sqlkit; pub mod sqlserver; diff --git a/src-tauri/src/capabilities/sql.rs b/src-tauri/src/capabilities/sql.rs index b531696..832ac2b 100644 --- a/src-tauri/src/capabilities/sql.rs +++ b/src-tauri/src/capabilities/sql.rs @@ -72,7 +72,7 @@ pub(crate) async fn resolve_adapter(connection_id: &str) -> Result Result { +pub(crate) async fn execute_on_adapter(adapter: &ActiveConnection, sql: &str) -> Result { match adapter { ActiveConnection::Postgres(a) => a .lock() @@ -131,7 +131,7 @@ async fn execute_on_adapter(adapter: &ActiveConnection, sql: &str) -> Result) -> Result { +pub(crate) fn get_connection_id(config: Option<&Value>) -> Result { match config { None => Err( "No connection was provided for this tool call. Supply a connection_id \ @@ -176,6 +176,11 @@ impl CapabilityHandler for ExecuteQueryHandler { .ok_or_else(|| "Missing 'sql' argument".to_string())?; let adapter = resolve_adapter(&conn_id).await?; + // Read-only guard: reject write/delete/ddl statements with actionable + // guidance so agents migrate to the split write tools. + let db_type = crate::capabilities::sql_write::adapter_db_type(&adapter); + crate::capabilities::sql_write::ensure_read_only(&db_type, sql)?; + // Check connection quality and warn the AI agent about flaky connections let mut guardian_warning: Option = None; if let Some(guardian) = crate::GUARDIAN.get() { @@ -683,17 +688,17 @@ fn connection_id_schema() -> Value { pub fn register_sql_tools(reg: &mut CapabilityRegistry) { reg.register(Capability { name: "sqlkit__execute_query", - description: "Execute an arbitrary SQL query and return the result set. Supports SELECT, INSERT, UPDATE, DELETE, DDL, and any other SQL statement.", + description: "Execute a read-only SQL query (SELECT, SHOW, EXPLAIN) and return the result set. Write statements (INSERT/UPDATE/MERGE), deletes (DELETE/TRUNCATE), and DDL (CREATE/ALTER/DROP) are rejected — use sqlkit__execute_write, sqlkit__execute_delete, or sqlkit__execute_ddl respectively.", handler: Arc::new(ExecuteQueryHandler), input_schema: json!({"type": "object", "properties": { "connection_id": connection_id_schema(), - "sql": {"type": "string", "description": "The SQL query to execute"} + "sql": {"type": "string", "description": "The read-only SQL query (SELECT/SHOW/EXPLAIN)"} }, "required": ["connection_id", "sql"]}), - risk_level: RiskLevel::Elevated, + risk_level: RiskLevel::Safe, required_permission: "read", source_kind: SourceKind::SqlDatabase, tags: &["agent"], - parallel_ok: false, + parallel_ok: true, }); reg.register(Capability { diff --git a/src-tauri/src/capabilities/sql_write.rs b/src-tauri/src/capabilities/sql_write.rs new file mode 100644 index 0000000..2c323ce --- /dev/null +++ b/src-tauri/src/capabilities/sql_write.rs @@ -0,0 +1,437 @@ +//! SQL write-capability tools for the MCP bridge. +//! +//! Splits write operations out of `sqlkit__execute_query` so the policy gate +//! (McpPolicy) can enforce them by risk level: writes are Elevated, deletes +//! and DDL are Destructive (gated by `confirm_destructive`). + +use std::sync::Arc; + +use serde_json::{json, Value}; +use sqlparser::ast::Statement; +use sqlparser::dialect::{GenericDialect, MsSqlDialect, MySqlDialect, PostgreSqlDialect}; +use sqlparser::parser::Parser; + +use data_studio_agent::capabilities::registry::CapabilityRegistry; +use data_studio_agent::capabilities::types::{Capability, CapabilityHandler, RiskLevel, SourceKind}; + +use super::sql::{execute_on_adapter, resolve_adapter}; + +/// Statement category for risk-gated dispatch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SqlKind { + Read, + Write, + Delete, + Ddl, + Other, +} + +/// Classify a SQL statement by its top-level AST variant, using the dialect +/// that matches the connection's database type. +pub(crate) fn classify_sql(db_type: &str, sql: &str) -> Result { + let dialect: Box = match db_type.to_lowercase().as_str() { + "postgres" | "postgresql" | "duckdb" | "cockroachdb" | "gbase8c" | "kingbasees" + | "yashandb" | "xugudb" | "oceanbase" | "dameng" => Box::new(PostgreSqlDialect {}), + "mysql" | "clickhouse" | "gbase8a" => Box::new(MySqlDialect {}), + "sqlserver" => Box::new(MsSqlDialect {}), + _ => Box::new(GenericDialect {}), + }; + + let stmts = Parser::parse_sql(&*dialect, sql) + .map_err(|e| format!("parse error: {}", e))?; + if stmts.len() > 1 { + return Err("multiple statements are not supported".to_string()); + } + let Some(stmt) = stmts.into_iter().next() else { + return Err("empty query".to_string()); + }; + Ok(classify_statement(&stmt)) +} + +fn classify_statement(stmt: &Statement) -> SqlKind { + match stmt { + Statement::Query(_) + | Statement::Explain { .. } + | Statement::ExplainTable { .. } + | Statement::ShowVariable { .. } + | Statement::ShowVariables { .. } + | Statement::ShowTables { .. } + | Statement::ShowColumns { .. } + | Statement::ShowViews { .. } + | Statement::ShowSchemas { .. } + | Statement::ShowDatabases { .. } + | Statement::ShowFunctions { .. } + | Statement::ShowStatus { .. } + | Statement::ShowCollation { .. } + | Statement::ShowCreate { .. } + | Statement::ShowObjects(_) => SqlKind::Read, + Statement::Insert(_) | Statement::Update { .. } | Statement::Merge { .. } => SqlKind::Write, + Statement::Delete(_) | Statement::Truncate { .. } => SqlKind::Delete, + Statement::CreateTable(_) + | Statement::CreateIndex(_) + | Statement::CreateDatabase { .. } + | Statement::CreateSchema { .. } + | Statement::CreateView { .. } + | Statement::CreateFunction(_) + | Statement::CreateProcedure { .. } + | Statement::CreateTrigger { .. } + | Statement::CreateSequence { .. } + | Statement::CreateType { .. } + | Statement::CreateExtension { .. } + | Statement::CreateRole { .. } + | Statement::CreatePolicy { .. } + | Statement::CreateMacro { .. } + | Statement::CreateStage { .. } + | Statement::CreateSecret { .. } + | Statement::CreateVirtualTable { .. } + | Statement::CreateConnector(_) + | Statement::AlterTable { .. } + | Statement::AlterIndex { .. } + | Statement::AlterView { .. } + | Statement::AlterRole { .. } + | Statement::AlterType(_) + | Statement::AlterSession { .. } + | Statement::AlterPolicy { .. } + | Statement::AlterConnector { .. } + | Statement::Drop { .. } + | Statement::DropFunction { .. } + | Statement::DropProcedure { .. } + | Statement::DropTrigger { .. } + | Statement::DropExtension { .. } + | Statement::DropPolicy { .. } + | Statement::DropSecret { .. } + | Statement::DropConnector { .. } + | Statement::RenameTable(_) + | Statement::AttachDatabase { .. } + | Statement::AttachDuckDBDatabase { .. } + | Statement::DetachDuckDBDatabase { .. } + | Statement::Unload { .. } + | Statement::Load { .. } + | Statement::LoadData { .. } + | Statement::Cache { .. } + | Statement::UNCache { .. } + | Statement::OptimizeTable { .. } + | Statement::Analyze { .. } + | Statement::Msck { .. } + | Statement::Grant { .. } + | Statement::Revoke { .. } + | Statement::Comment { .. } + | Statement::LockTables { .. } + | Statement::UnlockTables { .. } + | Statement::SetVariable { .. } + | Statement::SetNames { .. } + | Statement::SetNamesDefault { .. } + | Statement::SetRole { .. } + | Statement::SetSessionParam(_) + | Statement::SetTimeZone { .. } + | Statement::SetTransaction { .. } + | Statement::Commit { .. } + | Statement::Rollback { .. } + | Statement::Savepoint { .. } + | Statement::ReleaseSavepoint { .. } + | Statement::StartTransaction { .. } + | Statement::Declare { .. } + | Statement::Prepare { .. } + | Statement::Execute { .. } + | Statement::Deallocate { .. } + | Statement::Call(_) + | Statement::Copy { .. } + | Statement::CopyIntoSnowflake { .. } + | Statement::Kill { .. } + | Statement::Flush { .. } + | Statement::Pragma { .. } + | Statement::Use(_) + | Statement::Install { .. } + | Statement::RaisError { .. } + | Statement::Fetch { .. } + | Statement::Close { .. } + | Statement::Discard { .. } + | Statement::Assert { .. } + | Statement::Directory { .. } + | Statement::Remove(_) + | Statement::List(_) + | Statement::LISTEN { .. } + | Statement::NOTIFY { .. } + | Statement::UNLISTEN { .. } => SqlKind::Ddl, + // Future sqlparser variants fall here rather than failing to compile. + #[allow(unreachable_patterns)] + _ => SqlKind::Other, + } +} + +/// Guard a capability against the statement category it is allowed to run. +/// `execute_query` is read-only; anything else returns an actionable error. +pub(crate) fn ensure_read_only(db_type: &str, sql: &str) -> Result<(), String> { + match classify_sql(db_type, sql)? { + SqlKind::Read => Ok(()), + SqlKind::Write => Err( + "Only SELECT/SHOW/EXPLAIN statements are allowed in sqlkit__execute_query. \ + Use sqlkit__execute_write for INSERT/UPDATE/MERGE." + .to_string(), + ), + SqlKind::Delete => Err( + "DELETE/TRUNCATE must use sqlkit__execute_delete (destructive, requires \ + Full Access with Confirm Destructive in Settings → MCP Bridge)." + .to_string(), + ), + SqlKind::Ddl => Err( + "DDL statements (CREATE/ALTER/DROP) must use sqlkit__execute_ddl (destructive, \ + requires Full Access with Confirm Destructive in Settings → MCP Bridge)." + .to_string(), + ), + SqlKind::Other => Err( + "Statement type is not recognized as read-only. Use sqlkit__execute_write, \ + sqlkit__execute_delete, or sqlkit__execute_ddl as appropriate." + .to_string(), + ), + } +} + +fn connection_id_from(config: Option<&Value>) -> Result { + super::sql::get_connection_id(config) +} + +fn sql_from(args: &Value) -> Result { + args.get("sql") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .ok_or_else(|| "Missing 'sql' argument".to_string()) +} + +/// Guard that the statement belongs to the allowed categories, then execute. +async fn run_classified( + allowed: &[SqlKind], + args: &Value, + connection_config: Option<&Value>, +) -> Result { + let conn_id = connection_id_from(connection_config)?; + let sql = sql_from(args)?; + let adapter = resolve_adapter(&conn_id).await?; + + let db_type = adapter_db_type(&adapter); + let kind = classify_sql(&db_type, &sql)?; + if !allowed.contains(&kind) { + return Err(format!( + "sqlkit__execute_{} does not accept {:?} statements. {:?} statements go to {}.", + match allowed[0] { + SqlKind::Write => "write", + SqlKind::Delete => "delete", + _ => "ddl", + }, + kind, + kind, + match kind { + SqlKind::Read => "sqlkit__execute_query", + SqlKind::Write => "sqlkit__execute_write", + SqlKind::Delete => "sqlkit__execute_delete", + SqlKind::Ddl => "sqlkit__execute_ddl", + SqlKind::Other => "an explicit tool", + }, + )); + } + + let result = execute_on_adapter(&adapter, &sql).await?; + serde_json::to_string(&result).map_err(|e| e.to_string()) +} + +pub(crate) fn adapter_db_type(adapter: &crate::state::ActiveConnection) -> String { + match adapter { + crate::state::ActiveConnection::Postgres(_) => "postgres".to_string(), + crate::state::ActiveConnection::MySQL(_) => "mysql".to_string(), + crate::state::ActiveConnection::SQLite(_) => "sqlite".to_string(), + crate::state::ActiveConnection::SQLServer(_) => "sqlserver".to_string(), + crate::state::ActiveConnection::ClickHouse(_) => "clickhouse".to_string(), + crate::state::ActiveConnection::JdbcBridge(_) => "generic".to_string(), + crate::state::ActiveConnection::HttpSql(_) => "generic".to_string(), + crate::state::ActiveConnection::Rqlite(_) => "generic".to_string(), + crate::state::ActiveConnection::Turso(_) => "generic".to_string(), + } +} + +// --------------------------------------------------------------------------- +// Handlers +// --------------------------------------------------------------------------- + +struct ExecuteWriteHandler; +struct ExecuteDeleteHandler; +struct ExecuteDdlHandler; + +#[async_trait::async_trait] +#[async_trait::async_trait] +impl CapabilityHandler for ExecuteWriteHandler { + async fn handle( + &self, + args: &Value, + connection_config: Option<&Value>, + ) -> Result { + run_classified(&[SqlKind::Write], args, connection_config).await + } +} + +#[async_trait::async_trait] +#[async_trait::async_trait] +impl CapabilityHandler for ExecuteDeleteHandler { + async fn handle( + &self, + args: &Value, + connection_config: Option<&Value>, + ) -> Result { + run_classified(&[SqlKind::Delete], args, connection_config).await + } +} + +#[async_trait::async_trait] +#[async_trait::async_trait] +impl CapabilityHandler for ExecuteDdlHandler { + async fn handle( + &self, + args: &Value, + connection_config: Option<&Value>, + ) -> Result { + run_classified(&[SqlKind::Ddl], args, connection_config).await + } +} + +// --------------------------------------------------------------------------- +// Registration +// --------------------------------------------------------------------------- + +fn connection_id_schema() -> Value { + json!({ + "type": "string", + "description": "The connection alias to use (e.g. 'mac-postgresql'). Use sqlkit__list_connections to see available connections." + }) +} + +pub(crate) fn register_write_tools(reg: &mut CapabilityRegistry) { + reg.register(Capability { + name: "sqlkit__execute_write", + description: "Execute a data-modifying SQL statement: INSERT, UPDATE, or MERGE. Use sqlkit__execute_query for reads and sqlkit__execute_delete / sqlkit__execute_ddl for destructive statements.", + handler: Arc::new(ExecuteWriteHandler), + input_schema: json!({"type": "object", "properties": { + "connection_id": connection_id_schema(), + "sql": {"type": "string", "description": "The INSERT/UPDATE/MERGE SQL statement"} + }, "required": ["connection_id", "sql"]}), + risk_level: RiskLevel::Elevated, + required_permission: "create", + source_kind: SourceKind::SqlDatabase, + tags: &["agent"], + parallel_ok: false, + }); + + reg.register(Capability { + name: "sqlkit__execute_delete", + description: "Execute a destructive DELETE or TRUNCATE statement. DESTRUCTIVE: permanently removes data. Requires Full Access with Confirm Destructive in Settings → MCP Bridge.", + handler: Arc::new(ExecuteDeleteHandler), + input_schema: json!({"type": "object", "properties": { + "connection_id": connection_id_schema(), + "sql": {"type": "string", "description": "The DELETE or TRUNCATE SQL statement"} + }, "required": ["connection_id", "sql"]}), + risk_level: RiskLevel::Destructive, + required_permission: "delete", + source_kind: SourceKind::SqlDatabase, + tags: &["agent"], + parallel_ok: false, + }); + + reg.register(Capability { + name: "sqlkit__execute_ddl", + description: "Execute a DDL statement (CREATE, ALTER, DROP, and other schema changes). DESTRUCTIVE for DROP/ALTER: permanently changes schema. Requires Full Access with Confirm Destructive in Settings → MCP Bridge.", + handler: Arc::new(ExecuteDdlHandler), + input_schema: json!({"type": "object", "properties": { + "connection_id": connection_id_schema(), + "sql": {"type": "string", "description": "The DDL SQL statement"} + }, "required": ["connection_id", "sql"]}), + risk_level: RiskLevel::Destructive, + required_permission: "delete", + source_kind: SourceKind::SqlDatabase, + tags: &["agent"], + parallel_ok: false, + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn classifies_select_as_read() { + assert_eq!(classify_sql("postgres", "SELECT * FROM users").unwrap(), SqlKind::Read); + assert_eq!(classify_sql("postgres", "WITH x AS (SELECT 1) SELECT * FROM x").unwrap(), SqlKind::Read); + } + + #[test] + fn classifies_show_as_read() { + assert_eq!(classify_sql("mysql", "SHOW TABLES").unwrap(), SqlKind::Read); + assert_eq!(classify_sql("sqlserver", "SHOW VARIABLES").unwrap(), SqlKind::Read); + } + + #[test] + fn classifies_insert_update_merge_as_write() { + assert_eq!( + classify_sql("postgres", "INSERT INTO users (id) VALUES (1)").unwrap(), + SqlKind::Write + ); + assert_eq!( + classify_sql("postgres", "UPDATE users SET name = 'x' WHERE id = 1").unwrap(), + SqlKind::Write + ); + assert_eq!( + classify_sql("sqlserver", "MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN UPDATE SET x = s.x").unwrap(), + SqlKind::Write + ); + } + + #[test] + fn classifies_delete_truncate_as_delete() { + assert_eq!( + classify_sql("postgres", "DELETE FROM users WHERE id = 1").unwrap(), + SqlKind::Delete + ); + assert_eq!( + classify_sql("postgres", "TRUNCATE TABLE users").unwrap(), + SqlKind::Delete + ); + } + + #[test] + fn classifies_ddl_as_ddl() { + assert_eq!( + classify_sql("postgres", "CREATE TABLE t (id int)").unwrap(), + SqlKind::Ddl + ); + assert_eq!( + classify_sql("postgres", "ALTER TABLE t ADD COLUMN c int").unwrap(), + SqlKind::Ddl + ); + assert_eq!(classify_sql("postgres", "DROP TABLE t").unwrap(), SqlKind::Ddl); + } + + #[test] + fn rejects_multiple_statements() { + assert!(classify_sql("postgres", "SELECT 1; SELECT 2").is_err()); + } + + #[test] + fn read_guard_rejects_write() { + let err = ensure_read_only("postgres", "INSERT INTO t VALUES (1)").unwrap_err(); + assert!(err.contains("execute_write"), "got: {}", err); + } + + #[test] + fn read_guard_rejects_delete() { + let err = ensure_read_only("postgres", "DELETE FROM t").unwrap_err(); + assert!(err.contains("execute_delete"), "got: {}", err); + } + + #[test] + fn read_guard_rejects_ddl() { + let err = ensure_read_only("postgres", "DROP TABLE t").unwrap_err(); + assert!(err.contains("execute_ddl"), "got: {}", err); + } + + #[test] + fn read_guard_allows_select() { + assert!(ensure_read_only("postgres", "SELECT 1").is_ok()); + } +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 858d04e..b2f2aaa 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -116,6 +116,7 @@ pub fn run() { data_studio_agent::capabilities::registry::init_registry(&[ crate::capabilities::sqlkit::register_all, crate::capabilities::sql::register_sql_tools, + crate::capabilities::sql_write::register_write_tools, ]); // Initialize agent SQLite database diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs index d35dddc..7450224 100644 --- a/src-tauri/src/mcp_bridge.rs +++ b/src-tauri/src/mcp_bridge.rs @@ -616,6 +616,7 @@ mod tests { data_studio_agent::capabilities::registry::init_registry(&[ crate::capabilities::sqlkit::register_all, crate::capabilities::sql::register_sql_tools, + crate::capabilities::sql_write::register_write_tools, ]); }