From 929dc351bc5398886bb815f10344e44d5bee8fa7 Mon Sep 17 00:00:00 2001 From: blankll Date: Fri, 14 Aug 2026 16:41:46 +0800 Subject: [PATCH] feat(sql): add session/slow-query/privilege DBA capabilities MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds 5 new sqlkit capabilities — list_sessions, kill_session, get_slow_queries, grant_privilege, revoke_privilege — via a new capabilities/dba.rs module wired into the registry: - DatabaseAdapter trait gains 5 methods with unsupported default stubs - Implemented for PostgreSQL (pg_stat_activity / pg_terminate_backend / GRANT/REVOKE), MySQL (processlist / KILL / GRANT/REVOKE) and SQL Server (sys.dm_exec_sessions / KILL / sys.dm_exec_query_stats / GRANT/REVOKE) - Identifier whitelist + privilege allowlist guard against SQL injection in GRANT/REVOKE object/grantee arguments - agent-tag regression test asserts all 25 sqlkit capabilities are agent-visible Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src-tauri/src/capabilities/dba.rs | 380 ++++++++++++++++++++++++++++ src-tauri/src/capabilities/mod.rs | 1 + src-tauri/src/capabilities/sql.rs | 315 +++++++++++++++++++++++ src-tauri/src/database/adapter.rs | 39 +++ src-tauri/src/database/mysql.rs | 234 +++++++++++++++++ src-tauri/src/database/postgres.rs | 267 +++++++++++++++++++ src-tauri/src/database/sqlserver.rs | 232 +++++++++++++++++ src-tauri/src/lib.rs | 1 + src-tauri/src/mcp_bridge.rs | 30 +++ 9 files changed, 1499 insertions(+) create mode 100644 src-tauri/src/capabilities/dba.rs diff --git a/src-tauri/src/capabilities/dba.rs b/src-tauri/src/capabilities/dba.rs new file mode 100644 index 00000000..8d248ff9 --- /dev/null +++ b/src-tauri/src/capabilities/dba.rs @@ -0,0 +1,380 @@ +//! DBA SQL tools for the MCP bridge. +//! +//! Dedicated capabilities for session inspection/termination, slow-query +//! discovery, and privilege management. These call dedicated +//! `DatabaseAdapter` methods that bypass `classify_sql`, so +//! `Statement::Grant/Revoke/Kill` never reaches the write/DDL gate. + +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::{json, Value}; + +use data_studio_agent::capabilities::registry::CapabilityRegistry; +use data_studio_agent::capabilities::types::{ + Capability, CapabilityHandler, RiskLevel, SourceKind, +}; + +use super::sql::{ + get_connection_id, get_slow_queries_on_adapter, grant_privilege_on_adapter, + kill_session_on_adapter, list_sessions_on_adapter, resolve_adapter, + revoke_privilege_on_adapter, +}; + +// --------------------------------------------------------------------------- +// Handler structs +// --------------------------------------------------------------------------- + +struct ListSessionsHandler; +struct KillSessionHandler; +struct GetSlowQueriesHandler; +struct GrantPrivilegeHandler; +struct RevokePrivilegeHandler; + +// --------------------------------------------------------------------------- +// Handler implementations +// --------------------------------------------------------------------------- + +#[async_trait] +impl CapabilityHandler for ListSessionsHandler { + async fn handle( + &self, + args: &Value, + connection_config: Option<&Value>, + ) -> Result { + let conn_id = get_connection_id(connection_config)?; + let database = args.get("database").and_then(|v| v.as_str()); + let adapter = resolve_adapter(&conn_id).await?; + let sessions = list_sessions_on_adapter(&adapter, database).await?; + serde_json::to_string(&sessions).map_err(|e| e.to_string()) + } +} + +#[async_trait] +impl CapabilityHandler for KillSessionHandler { + async fn handle( + &self, + args: &Value, + connection_config: Option<&Value>, + ) -> Result { + let conn_id = get_connection_id(connection_config)?; + let session_id = args + .get("session_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| "Missing 'session_id' argument".to_string())?; + let adapter = resolve_adapter(&conn_id).await?; + kill_session_on_adapter(&adapter, session_id).await?; + serde_json::to_string(&json!({ "status": "ok", "session_id": session_id })) + .map_err(|e| e.to_string()) + } +} + +#[async_trait] +impl CapabilityHandler for GetSlowQueriesHandler { + async fn handle( + &self, + args: &Value, + connection_config: Option<&Value>, + ) -> Result { + let conn_id = get_connection_id(connection_config)?; + let database = args.get("database").and_then(|v| v.as_str()); + let limit = args.get("limit").and_then(|v| v.as_u64()).map(|n| n as u32); + let adapter = resolve_adapter(&conn_id).await?; + let slow = get_slow_queries_on_adapter(&adapter, database, limit).await?; + serde_json::to_string(&slow).map_err(|e| e.to_string()) + } +} + +#[async_trait] +impl CapabilityHandler for GrantPrivilegeHandler { + async fn handle( + &self, + args: &Value, + connection_config: Option<&Value>, + ) -> Result { + let conn_id = get_connection_id(connection_config)?; + let privilege = args + .get("privilege") + .and_then(|v| v.as_str()) + .ok_or_else(|| "Missing 'privilege' argument".to_string())?; + let object = args + .get("object") + .and_then(|v| v.as_str()) + .ok_or_else(|| "Missing 'object' argument".to_string())?; + let grantee = args + .get("grantee") + .and_then(|v| v.as_str()) + .ok_or_else(|| "Missing 'grantee' argument".to_string())?; + let adapter = resolve_adapter(&conn_id).await?; + grant_privilege_on_adapter(&adapter, privilege, object, grantee).await?; + serde_json::to_string(&json!({ "status": "ok" })).map_err(|e| e.to_string()) + } +} + +#[async_trait] +impl CapabilityHandler for RevokePrivilegeHandler { + async fn handle( + &self, + args: &Value, + connection_config: Option<&Value>, + ) -> Result { + let conn_id = get_connection_id(connection_config)?; + let privilege = args + .get("privilege") + .and_then(|v| v.as_str()) + .ok_or_else(|| "Missing 'privilege' argument".to_string())?; + let object = args + .get("object") + .and_then(|v| v.as_str()) + .ok_or_else(|| "Missing 'object' argument".to_string())?; + let grantee = args + .get("grantee") + .and_then(|v| v.as_str()) + .ok_or_else(|| "Missing 'grantee' argument".to_string())?; + let adapter = resolve_adapter(&conn_id).await?; + revoke_privilege_on_adapter(&adapter, privilege, object, grantee).await?; + serde_json::to_string(&json!({ "status": "ok" })).map_err(|e| e.to_string()) + } +} + +// --------------------------------------------------------------------------- +// 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_dba_tools(reg: &mut CapabilityRegistry) { + reg.register(Capability { + name: "sqlkit__list_sessions", + description: "List active database sessions/connections: id, user, database, state, and running query. Use to see who is connected and what they are running.", + handler: Arc::new(ListSessionsHandler), + input_schema: json!({"type": "object", "properties": { + "connection_id": connection_id_schema(), + "database": {"type": "string", "description": "Database name (optional)"} + }, "required": ["connection_id"]}), + risk_level: RiskLevel::Safe, + required_permission: "read", + source_kind: SourceKind::SqlDatabase, + tags: &["agent"], + parallel_ok: true, + }); + + reg.register(Capability { + name: "sqlkit__kill_session", + description: "Terminate a database session by its session id (PostgreSQL PID for postgres). DANGEROUS: kills an in-flight query/connection. Requires Full Access in Settings → MCP Bridge.", + handler: Arc::new(KillSessionHandler), + input_schema: json!({"type": "object", "properties": { + "connection_id": connection_id_schema(), + "session_id": {"type": "string", "description": "Session id to terminate (numeric; for PostgreSQL this is the backend PID)"} + }, "required": ["connection_id", "session_id"]}), + risk_level: RiskLevel::Elevated, + required_permission: "create", + source_kind: SourceKind::SqlDatabase, + tags: &["agent"], + parallel_ok: false, + }); + + reg.register(Capability { + name: "sqlkit__get_slow_queries", + description: "List currently slow-running queries (or cached slow query statistics) on the server: duration, user, and query text. Use when investigating performance issues.", + handler: Arc::new(GetSlowQueriesHandler), + input_schema: json!({"type": "object", "properties": { + "connection_id": connection_id_schema(), + "database": {"type": "string", "description": "Database name (optional)"}, + "limit": {"type": "integer", "description": "Maximum number of queries to return (optional, default 20)"} + }, "required": ["connection_id"]}), + risk_level: RiskLevel::Safe, + required_permission: "read", + source_kind: SourceKind::SqlDatabase, + tags: &["agent"], + parallel_ok: true, + }); + + reg.register(Capability { + name: "sqlkit__grant_privilege", + description: "Grant a privilege (e.g. SELECT, INSERT) on an object (e.g. a table) to a user/role. Requires Full Access in Settings → MCP Bridge.", + handler: Arc::new(GrantPrivilegeHandler), + input_schema: json!({"type": "object", "properties": { + "connection_id": connection_id_schema(), + "privilege": {"type": "string", "description": "Privilege(s) to grant, e.g. SELECT, INSERT, or a comma-separated list"}, + "object": {"type": "string", "description": "Object to grant on, e.g. public.users or db.table"}, + "grantee": {"type": "string", "description": "User or role to grant to"} + }, "required": ["connection_id", "privilege", "object", "grantee"]}), + risk_level: RiskLevel::Elevated, + required_permission: "create", + source_kind: SourceKind::SqlDatabase, + tags: &["agent"], + parallel_ok: false, + }); + + reg.register(Capability { + name: "sqlkit__revoke_privilege", + description: "Revoke a privilege (e.g. SELECT, INSERT) on an object from a user/role. Requires Full Access in Settings → MCP Bridge.", + handler: Arc::new(RevokePrivilegeHandler), + input_schema: json!({"type": "object", "properties": { + "connection_id": connection_id_schema(), + "privilege": {"type": "string", "description": "Privilege(s) to revoke, e.g. SELECT, INSERT, or a comma-separated list"}, + "object": {"type": "string", "description": "Object to revoke on, e.g. public.users or db.table"}, + "grantee": {"type": "string", "description": "User or role to revoke from"} + }, "required": ["connection_id", "privilege", "object", "grantee"]}), + risk_level: RiskLevel::Elevated, + required_permission: "create", + source_kind: SourceKind::SqlDatabase, + tags: &["agent"], + parallel_ok: false, + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn list_sessions_missing_config() { + let err = ListSessionsHandler + .handle(&json!({}), None) + .await + .unwrap_err(); + assert!(err.contains("connection_id"), "got: {}", err); + assert!(err.contains("Settings → MCP Bridge"), "got: {}", err); + } + + #[tokio::test] + async fn kill_session_missing_config() { + let err = KillSessionHandler + .handle(&json!({ "session_id": "123" }), None) + .await + .unwrap_err(); + assert!(err.contains("connection_id"), "got: {}", err); + assert!(err.contains("Settings → MCP Bridge"), "got: {}", err); + } + + #[tokio::test] + async fn kill_session_rejects_missing_session_id() { + let config = json!({ "connectionId": "conn-1" }); + let err = KillSessionHandler + .handle(&json!({}), Some(&config)) + .await + .unwrap_err(); + assert!(err.contains("Missing 'session_id'"), "got: {}", err); + } + + #[tokio::test] + async fn get_slow_queries_missing_config() { + let err = GetSlowQueriesHandler + .handle(&json!({}), None) + .await + .unwrap_err(); + assert!(err.contains("connection_id"), "got: {}", err); + assert!(err.contains("Settings → MCP Bridge"), "got: {}", err); + } + + #[tokio::test] + async fn grant_privilege_missing_config() { + let err = GrantPrivilegeHandler + .handle( + &json!({ "privilege": "SELECT", "object": "users", "grantee": "app" }), + None, + ) + .await + .unwrap_err(); + assert!(err.contains("connection_id"), "got: {}", err); + assert!(err.contains("Settings → MCP Bridge"), "got: {}", err); + } + + #[tokio::test] + async fn grant_privilege_rejects_missing_privilege() { + let config = json!({ "connectionId": "conn-1" }); + let err = GrantPrivilegeHandler + .handle( + &json!({ "object": "users", "grantee": "app" }), + Some(&config), + ) + .await + .unwrap_err(); + assert!(err.contains("Missing 'privilege'"), "got: {}", err); + } + + #[tokio::test] + async fn grant_privilege_rejects_missing_object() { + let config = json!({ "connectionId": "conn-1" }); + let err = GrantPrivilegeHandler + .handle( + &json!({ "privilege": "SELECT", "grantee": "app" }), + Some(&config), + ) + .await + .unwrap_err(); + assert!(err.contains("Missing 'object'"), "got: {}", err); + } + + #[tokio::test] + async fn grant_privilege_rejects_missing_grantee() { + let config = json!({ "connectionId": "conn-1" }); + let err = GrantPrivilegeHandler + .handle( + &json!({ "privilege": "SELECT", "object": "users" }), + Some(&config), + ) + .await + .unwrap_err(); + assert!(err.contains("Missing 'grantee'"), "got: {}", err); + } + + #[tokio::test] + async fn revoke_privilege_missing_config() { + let err = RevokePrivilegeHandler + .handle( + &json!({ "privilege": "SELECT", "object": "users", "grantee": "app" }), + None, + ) + .await + .unwrap_err(); + assert!(err.contains("connection_id"), "got: {}", err); + assert!(err.contains("Settings → MCP Bridge"), "got: {}", err); + } + + #[tokio::test] + async fn revoke_privilege_rejects_missing_privilege() { + let config = json!({ "connectionId": "conn-1" }); + let err = RevokePrivilegeHandler + .handle( + &json!({ "object": "users", "grantee": "app" }), + Some(&config), + ) + .await + .unwrap_err(); + assert!(err.contains("Missing 'privilege'"), "got: {}", err); + } + + #[tokio::test] + async fn revoke_privilege_rejects_missing_object() { + let config = json!({ "connectionId": "conn-1" }); + let err = RevokePrivilegeHandler + .handle( + &json!({ "privilege": "SELECT", "grantee": "app" }), + Some(&config), + ) + .await + .unwrap_err(); + assert!(err.contains("Missing 'object'"), "got: {}", err); + } + + #[tokio::test] + async fn revoke_privilege_rejects_missing_grantee() { + let config = json!({ "connectionId": "conn-1" }); + let err = RevokePrivilegeHandler + .handle( + &json!({ "privilege": "SELECT", "object": "users" }), + Some(&config), + ) + .await + .unwrap_err(); + assert!(err.contains("Missing 'grantee'"), "got: {}", err); + } +} diff --git a/src-tauri/src/capabilities/mod.rs b/src-tauri/src/capabilities/mod.rs index 56d38819..d22f7f9b 100644 --- a/src-tauri/src/capabilities/mod.rs +++ b/src-tauri/src/capabilities/mod.rs @@ -1,4 +1,5 @@ pub mod commands; +pub mod dba; pub mod mysql; pub mod postgres; pub mod sql; diff --git a/src-tauri/src/capabilities/sql.rs b/src-tauri/src/capabilities/sql.rs index 97029f12..68462f2a 100644 --- a/src-tauri/src/capabilities/sql.rs +++ b/src-tauri/src/capabilities/sql.rs @@ -142,6 +142,321 @@ pub(crate) async fn execute_on_adapter( } } +pub(crate) async fn list_sessions_on_adapter( + adapter: &ActiveConnection, + database: Option<&str>, +) -> Result, String> { + match adapter { + ActiveConnection::Postgres(a) => a + .lock() + .await + .list_sessions(database) + .await + .map_err(|e| e.to_string()), + ActiveConnection::MySQL(a) => a + .lock() + .await + .list_sessions(database) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLite(a) => a + .lock() + .await + .list_sessions(None) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLServer(a) => a + .lock() + .await + .list_sessions(database) + .await + .map_err(|e| e.to_string()), + ActiveConnection::ClickHouse(a) => a + .lock() + .await + .list_sessions(database) + .await + .map_err(|e| e.to_string()), + ActiveConnection::JdbcBridge(a) => a + .lock() + .await + .list_sessions(database) + .await + .map_err(|e| e.to_string()), + ActiveConnection::HttpSql(a) => a + .lock() + .await + .list_sessions(database) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Rqlite(a) => a + .lock() + .await + .list_sessions(database) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Turso(a) => a + .lock() + .await + .list_sessions(database) + .await + .map_err(|e| e.to_string()), + } +} + +pub(crate) async fn kill_session_on_adapter( + adapter: &ActiveConnection, + session_id: &str, +) -> Result<(), String> { + match adapter { + ActiveConnection::Postgres(a) => a + .lock() + .await + .kill_session(session_id) + .await + .map_err(|e| e.to_string()), + ActiveConnection::MySQL(a) => a + .lock() + .await + .kill_session(session_id) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLite(a) => a + .lock() + .await + .kill_session(session_id) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLServer(a) => a + .lock() + .await + .kill_session(session_id) + .await + .map_err(|e| e.to_string()), + ActiveConnection::ClickHouse(a) => a + .lock() + .await + .kill_session(session_id) + .await + .map_err(|e| e.to_string()), + ActiveConnection::JdbcBridge(a) => a + .lock() + .await + .kill_session(session_id) + .await + .map_err(|e| e.to_string()), + ActiveConnection::HttpSql(a) => a + .lock() + .await + .kill_session(session_id) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Rqlite(a) => a + .lock() + .await + .kill_session(session_id) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Turso(a) => a + .lock() + .await + .kill_session(session_id) + .await + .map_err(|e| e.to_string()), + } +} + +pub(crate) async fn get_slow_queries_on_adapter( + adapter: &ActiveConnection, + database: Option<&str>, + limit: Option, +) -> Result, String> { + match adapter { + ActiveConnection::Postgres(a) => a + .lock() + .await + .get_slow_queries(database, limit) + .await + .map_err(|e| e.to_string()), + ActiveConnection::MySQL(a) => a + .lock() + .await + .get_slow_queries(database, limit) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLite(a) => a + .lock() + .await + .get_slow_queries(None, limit) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLServer(a) => a + .lock() + .await + .get_slow_queries(database, limit) + .await + .map_err(|e| e.to_string()), + ActiveConnection::ClickHouse(a) => a + .lock() + .await + .get_slow_queries(database, limit) + .await + .map_err(|e| e.to_string()), + ActiveConnection::JdbcBridge(a) => a + .lock() + .await + .get_slow_queries(database, limit) + .await + .map_err(|e| e.to_string()), + ActiveConnection::HttpSql(a) => a + .lock() + .await + .get_slow_queries(database, limit) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Rqlite(a) => a + .lock() + .await + .get_slow_queries(database, limit) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Turso(a) => a + .lock() + .await + .get_slow_queries(database, limit) + .await + .map_err(|e| e.to_string()), + } +} + +pub(crate) async fn grant_privilege_on_adapter( + adapter: &ActiveConnection, + privilege: &str, + object: &str, + grantee: &str, +) -> Result<(), String> { + match adapter { + ActiveConnection::Postgres(a) => a + .lock() + .await + .grant_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::MySQL(a) => a + .lock() + .await + .grant_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLite(a) => a + .lock() + .await + .grant_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLServer(a) => a + .lock() + .await + .grant_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::ClickHouse(a) => a + .lock() + .await + .grant_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::JdbcBridge(a) => a + .lock() + .await + .grant_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::HttpSql(a) => a + .lock() + .await + .grant_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Rqlite(a) => a + .lock() + .await + .grant_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Turso(a) => a + .lock() + .await + .grant_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + } +} + +pub(crate) async fn revoke_privilege_on_adapter( + adapter: &ActiveConnection, + privilege: &str, + object: &str, + grantee: &str, +) -> Result<(), String> { + match adapter { + ActiveConnection::Postgres(a) => a + .lock() + .await + .revoke_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::MySQL(a) => a + .lock() + .await + .revoke_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLite(a) => a + .lock() + .await + .revoke_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::SQLServer(a) => a + .lock() + .await + .revoke_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::ClickHouse(a) => a + .lock() + .await + .revoke_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::JdbcBridge(a) => a + .lock() + .await + .revoke_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::HttpSql(a) => a + .lock() + .await + .revoke_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Rqlite(a) => a + .lock() + .await + .revoke_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + ActiveConnection::Turso(a) => a + .lock() + .await + .revoke_privilege(privilege, object, grantee) + .await + .map_err(|e| e.to_string()), + } +} + pub(crate) fn get_connection_id(config: Option<&Value>) -> Result { match config { None => Err( diff --git a/src-tauri/src/database/adapter.rs b/src-tauri/src/database/adapter.rs index 6be846ac..a90a7a5f 100644 --- a/src-tauri/src/database/adapter.rs +++ b/src-tauri/src/database/adapter.rs @@ -251,6 +251,45 @@ pub trait DatabaseAdapter: Send + Sync { Err(DbError::unsupported("rename_object")) } + /// List active database sessions/connections. + async fn list_sessions(&self, _database: Option<&str>) -> DbResult> { + Err(DbError::unsupported("list_sessions")) + } + + /// Terminate a database session/connection by id (or PID for PostgreSQL). + async fn kill_session(&self, _session_id: &str) -> DbResult<()> { + Err(DbError::unsupported("kill_session")) + } + + /// List slow-running queries currently executing on the server. + async fn get_slow_queries( + &self, + _database: Option<&str>, + _limit: Option, + ) -> DbResult> { + Err(DbError::unsupported("get_slow_queries")) + } + + /// Grant a privilege on an object to a grantee (user/role). + async fn grant_privilege( + &self, + _privilege: &str, + _object: &str, + _grantee: &str, + ) -> DbResult<()> { + Err(DbError::unsupported("grant_privilege")) + } + + /// Revoke a privilege on an object from a grantee (user/role). + async fn revoke_privilege( + &self, + _privilege: &str, + _object: &str, + _grantee: &str, + ) -> DbResult<()> { + Err(DbError::unsupported("revoke_privilege")) + } + /// Get the connection pool. /// /// This method returns a reference to the connection pool used by this adapter. diff --git a/src-tauri/src/database/mysql.rs b/src-tauri/src/database/mysql.rs index 8579d277..c2d3672f 100644 --- a/src-tauri/src/database/mysql.rs +++ b/src-tauri/src/database/mysql.rs @@ -353,6 +353,109 @@ impl MySQLAdapter { .await .map_err(|e| DbError::Connection(format!("Failed to get connection: {}", e))) } + + fn quote_object(object: &str) -> DbResult { + if object.is_empty() || object.len() > 256 { + return Err(DbError::InvalidQuery(format!( + "Invalid object identifier: '{}'", + object + ))); + } + for part in object.split('.') { + if part == "*" { + continue; + } + let trimmed = part.trim_matches('`'); + let valid = !trimmed.is_empty() + && trimmed + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_') + && !trimmed.chars().next().unwrap().is_ascii_digit(); + if !valid { + return Err(DbError::InvalidQuery(format!( + "Invalid object identifier: '{}'", + object + ))); + } + } + Ok(object + .split('.') + .map(|part| { + if part == "*" { + "*".to_string() + } else { + format!("`{}`", part.trim_matches('`')) + } + }) + .collect::>() + .join(".")) + } + + fn validate_user(grantee: &str) -> DbResult { + let valid = !grantee.is_empty() + && grantee.len() <= 64 + && grantee + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '$') + && !grantee.chars().next().unwrap().is_ascii_digit(); + if !valid { + return Err(DbError::InvalidQuery(format!( + "Invalid user name: '{}'", + grantee + ))); + } + Ok(grantee.to_string()) + } + + fn validate_privilege(privilege: &str) -> DbResult { + const PRIVILEGES: &[&str] = &[ + "SELECT", + "INSERT", + "UPDATE", + "DELETE", + "CREATE", + "DROP", + "RELOAD", + "SHUTDOWN", + "PROCESS", + "FILE", + "REFERENCES", + "INDEX", + "ALTER", + "SHOW DATABASES", + "SUPER", + "CREATE TEMPORARY TABLES", + "LOCK TABLES", + "EXECUTE", + "REPLICATION SLAVE", + "REPLICATION CLIENT", + "CREATE VIEW", + "SHOW VIEW", + "CREATE ROUTINE", + "ALTER ROUTINE", + "CREATE USER", + "EVENT", + "TRIGGER", + "CREATE TABLESPACE", + "CREATE ROLE", + "DROP ROLE", + "ALL", + "ALL PRIVILEGES", + ]; + let mut validated = Vec::new(); + for part in privilege.split(',') { + let trimmed = part.trim().to_uppercase(); + if trimmed.is_empty() || !PRIVILEGES.contains(&trimmed.as_str()) { + return Err(DbError::InvalidQuery(format!( + "Unsupported privilege: '{}'. Supported: {}", + part.trim(), + PRIVILEGES.join(", ") + ))); + } + validated.push(trimmed); + } + Ok(validated.join(", ")) + } } /// Check if a mysql_async error is SSL/TLS related. @@ -1241,6 +1344,137 @@ impl DatabaseAdapter for MySQLAdapter { Ok(()) } + async fn list_sessions(&self, _database: Option<&str>) -> DbResult> { + let mut conn = self.get_conn().await?; + + let query = "SELECT id, user, host, db, command, time, state, info \ + FROM information_schema.processlist"; + + let rows: Vec = conn + .query(query) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + let sessions = rows + .into_iter() + .map(|row| { + let id: Option = row.get_opt(0).and_then(|r| r.ok()).flatten(); + let user: String = get_str(&row, 1); + let host: String = get_str(&row, 2); + let db: Option = get_opt_str(&row, 3); + let command: String = get_str(&row, 4); + let time: Option = row.get_opt(5).and_then(|r| r.ok()).flatten(); + let state: Option = get_opt_str(&row, 6); + let info: Option = get_opt_str(&row, 7); + serde_json::json!({ + "id": id, + "user": user, + "host": host, + "db": db, + "command": command, + "time_seconds": time, + "state": state, + "info": info, + }) + }) + .collect(); + + Ok(sessions) + } + + async fn kill_session(&self, session_id: &str) -> DbResult<()> { + let id: u64 = session_id.parse().map_err(|_| { + DbError::InvalidQuery(format!( + "Invalid session id '{}': MySQL session ids are numeric", + session_id + )) + })?; + + let mut conn = self.get_conn().await?; + + conn.query_drop(format!("KILL {}", id)) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + Ok(()) + } + + async fn get_slow_queries( + &self, + _database: Option<&str>, + limit: Option, + ) -> DbResult> { + let limit = limit.unwrap_or(20); + + let mut conn = self.get_conn().await?; + + let query = format!( + "SELECT id, user, host, db, time, state, info \ + FROM information_schema.processlist \ + WHERE time > 1 ORDER BY time DESC LIMIT {}", + limit + ); + + let rows: Vec = conn + .query(&query) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + let slow = rows + .into_iter() + .map(|row| { + let id: Option = row.get_opt(0).and_then(|r| r.ok()).flatten(); + let user: String = get_str(&row, 1); + let host: String = get_str(&row, 2); + let db: Option = get_opt_str(&row, 3); + let time: Option = row.get_opt(4).and_then(|r| r.ok()).flatten(); + let state: Option = get_opt_str(&row, 5); + let info: Option = get_opt_str(&row, 6); + serde_json::json!({ + "id": id, + "user": user, + "host": host, + "db": db, + "time_seconds": time, + "state": state, + "info": info, + }) + }) + .collect(); + + Ok(slow) + } + + async fn grant_privilege(&self, privilege: &str, object: &str, grantee: &str) -> DbResult<()> { + let privs = Self::validate_privilege(privilege)?; + let obj = Self::quote_object(object)?; + let user = Self::validate_user(grantee)?; + let sql = format!("GRANT {} ON {} TO '{}'@'%'", privs, obj, user); + + let mut conn = self.get_conn().await?; + + conn.query_drop(&sql) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + Ok(()) + } + + async fn revoke_privilege(&self, privilege: &str, object: &str, grantee: &str) -> DbResult<()> { + let privs = Self::validate_privilege(privilege)?; + let obj = Self::quote_object(object)?; + let user = Self::validate_user(grantee)?; + let sql = format!("REVOKE {} ON {} FROM '{}'@'%'", privs, obj, user); + + let mut conn = self.get_conn().await?; + + conn.query_drop(&sql) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + Ok(()) + } + fn get_pool(&self) -> Option> { self.pool.clone() } diff --git a/src-tauri/src/database/postgres.rs b/src-tauri/src/database/postgres.rs index 47524f25..33c9c1ac 100644 --- a/src-tauri/src/database/postgres.rs +++ b/src-tauri/src/database/postgres.rs @@ -243,6 +243,61 @@ impl PostgresAdapter { Self { config, pool: None } } + fn quote_ident(identifier: &str) -> DbResult { + if identifier.is_empty() || identifier.len() > 128 { + return Err(DbError::InvalidQuery(format!( + "Invalid identifier: '{}'", + identifier + ))); + } + for part in identifier.split('.') { + let trimmed = part.trim_matches('"'); + let valid = !trimmed.is_empty() + && trimmed + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_') + && !trimmed.chars().next().unwrap().is_ascii_digit(); + if !valid { + return Err(DbError::InvalidQuery(format!( + "Invalid identifier: '{}'", + identifier + ))); + } + } + Ok(identifier + .split('.') + .map(|part| format!("\"{}\"", part.trim_matches('"'))) + .collect::>() + .join(".")) + } + + fn validate_privilege(privilege: &str) -> DbResult { + const PRIVILEGES: &[&str] = &[ + "SELECT", + "INSERT", + "UPDATE", + "DELETE", + "TRUNCATE", + "REFERENCES", + "TRIGGER", + "ALL", + "ALL PRIVILEGES", + ]; + let mut validated = Vec::new(); + for part in privilege.split(',') { + let trimmed = part.trim().to_uppercase(); + if trimmed.is_empty() || !PRIVILEGES.contains(&trimmed.as_str()) { + return Err(DbError::InvalidQuery(format!( + "Unsupported privilege: '{}'. Supported: {}", + part.trim(), + PRIVILEGES.join(", ") + ))); + } + validated.push(trimmed); + } + Ok(validated.join(", ")) + } + /// Build the PostgreSQL connection string. fn build_connection_string(&self) -> String { let mut parts = Vec::new(); @@ -1942,6 +1997,218 @@ impl DatabaseAdapter for PostgresAdapter { Ok(indexes) } + async fn list_sessions(&self, database: Option<&str>) -> DbResult> { + if database.is_some() && database != self.config.database.as_deref() { + return Err(DbError::UnsupportedOperation( + "Cannot list sessions from a different database without reconnecting".to_string(), + )); + } + + let pool = self + .pool + .as_ref() + .ok_or_else(|| DbError::Connection("Not connected".to_string()))?; + + let client = pool + .pool + .get() + .await + .map_err(|e| DbError::Connection(format!("Failed to get connection: {}", e)))?; + + let query = r#" + SELECT + pid, + usename, + datname, + state, + query_start, + EXTRACT(EPOCH FROM (NOW() - query_start))::float8 AS duration_seconds, + query + FROM pg_stat_activity + WHERE state IS NOT NULL + AND pid <> pg_backend_pid() + "#; + + let rows = client + .query(query, &[]) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + let sessions = rows + .iter() + .map(|row| { + let pid: i32 = row.get(0); + let usename: Option = row.get(1); + let datname: Option = row.get(2); + let state: Option = row.get(3); + let query_start: Option> = row.get(4); + let duration_seconds: Option = row.get(5); + let query: Option = row.get(6); + serde_json::json!({ + "pid": pid, + "usename": usename, + "datname": datname, + "state": state, + "query_start": query_start.map(|d| d.to_string()), + "duration_seconds": duration_seconds, + "query": query, + }) + }) + .collect(); + + Ok(sessions) + } + + async fn kill_session(&self, session_id: &str) -> DbResult<()> { + let pid: i32 = session_id.parse().map_err(|_| { + DbError::InvalidQuery(format!( + "Invalid session id '{}': PostgreSQL session ids are numeric PIDs", + session_id + )) + })?; + + let pool = self + .pool + .as_ref() + .ok_or_else(|| DbError::Connection("Not connected".to_string()))?; + + let client = pool + .pool + .get() + .await + .map_err(|e| DbError::Connection(format!("Failed to get connection: {}", e)))?; + + let rows = client + .query("SELECT pg_terminate_backend($1::int)", &[&pid]) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + let terminated: bool = rows.first().map(|row| row.get(0)).unwrap_or(false); + if !terminated { + return Err(DbError::QueryExecution(format!( + "Failed to terminate session {} (no such backend or insufficient privilege)", + session_id + ))); + } + + Ok(()) + } + + async fn get_slow_queries( + &self, + database: Option<&str>, + limit: Option, + ) -> DbResult> { + if database.is_some() && database != self.config.database.as_deref() { + return Err(DbError::UnsupportedOperation( + "Cannot list slow queries from a different database without reconnecting" + .to_string(), + )); + } + + let limit = limit.unwrap_or(20) as i64; + + let pool = self + .pool + .as_ref() + .ok_or_else(|| DbError::Connection("Not connected".to_string()))?; + + let client = pool + .pool + .get() + .await + .map_err(|e| DbError::Connection(format!("Failed to get connection: {}", e)))?; + + let query = r#" + SELECT + pid, + usename, + datname, + EXTRACT(EPOCH FROM (NOW() - query_start))::float8 AS duration_seconds, + query + FROM pg_stat_activity + WHERE state = 'active' + AND NOW() - query_start > interval '1 second' + ORDER BY duration_seconds DESC + LIMIT $1 + "#; + + let rows = client + .query(query, &[&limit]) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + let slow = rows + .iter() + .map(|row| { + let pid: i32 = row.get(0); + let usename: String = row.get(1); + let datname: Option = row.get(2); + let duration_seconds: f64 = row.get(3); + let query: String = row.get(4); + serde_json::json!({ + "pid": pid, + "usename": usename, + "datname": datname, + "duration_seconds": duration_seconds, + "query": query, + }) + }) + .collect(); + + Ok(slow) + } + + async fn grant_privilege(&self, privilege: &str, object: &str, grantee: &str) -> DbResult<()> { + let privs = Self::validate_privilege(privilege)?; + let obj = Self::quote_ident(object)?; + let grantee = Self::quote_ident(grantee)?; + let sql = format!("GRANT {} ON {} TO {}", privs, obj, grantee); + + let pool = self + .pool + .as_ref() + .ok_or_else(|| DbError::Connection("Not connected".to_string()))?; + + let client = pool + .pool + .get() + .await + .map_err(|e| DbError::Connection(format!("Failed to get connection: {}", e)))?; + + client + .execute(&sql, &[]) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + Ok(()) + } + + async fn revoke_privilege(&self, privilege: &str, object: &str, grantee: &str) -> DbResult<()> { + let privs = Self::validate_privilege(privilege)?; + let obj = Self::quote_ident(object)?; + let grantee = Self::quote_ident(grantee)?; + let sql = format!("REVOKE {} ON {} FROM {}", privs, obj, grantee); + + let pool = self + .pool + .as_ref() + .ok_or_else(|| DbError::Connection("Not connected".to_string()))?; + + let client = pool + .pool + .get() + .await + .map_err(|e| DbError::Connection(format!("Failed to get connection: {}", e)))?; + + client + .execute(&sql, &[]) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + Ok(()) + } + async fn list_foreign_keys( &self, database: Option<&str>, diff --git a/src-tauri/src/database/sqlserver.rs b/src-tauri/src/database/sqlserver.rs index db68b4ce..a6cc0667 100644 --- a/src-tauri/src/database/sqlserver.rs +++ b/src-tauri/src/database/sqlserver.rs @@ -234,6 +234,64 @@ impl SqlServerAdapter { } } + fn quote_ident(identifier: &str) -> DbResult { + if identifier.is_empty() || identifier.len() > 128 { + return Err(DbError::InvalidQuery(format!( + "Invalid identifier: '{}'", + identifier + ))); + } + for part in identifier.split('.') { + let trimmed = part.trim_matches(|c| c == '[' || c == ']'); + let valid = !trimmed.is_empty() + && trimmed + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_') + && !trimmed.chars().next().unwrap().is_ascii_digit(); + if !valid { + return Err(DbError::InvalidQuery(format!( + "Invalid identifier: '{}'", + identifier + ))); + } + } + Ok(identifier + .split('.') + .map(|part| format!("[{}]", part.trim_matches(|c| c == '[' || c == ']'))) + .collect::>() + .join(".")) + } + + fn validate_privilege(privilege: &str) -> DbResult { + const PRIVILEGES: &[&str] = &[ + "SELECT", + "INSERT", + "UPDATE", + "DELETE", + "REFERENCES", + "EXECUTE", + "ALTER", + "CONTROL", + "TAKE OWNERSHIP", + "VIEW DEFINITION", + "ALL", + "ALL PRIVILEGES", + ]; + let mut validated = Vec::new(); + for part in privilege.split(',') { + let trimmed = part.trim().to_uppercase(); + if trimmed.is_empty() || !PRIVILEGES.contains(&trimmed.as_str()) { + return Err(DbError::InvalidQuery(format!( + "Unsupported privilege: '{}'. Supported: {}", + part.trim(), + PRIVILEGES.join(", ") + ))); + } + validated.push(trimmed); + } + Ok(validated.join(", ")) + } + /// Convert a tiberius Row to QueryRow. fn row_to_query_row(row: &Row) -> DbResult { let mut query_row = HashMap::new(); @@ -1661,6 +1719,180 @@ impl DatabaseAdapter for SqlServerAdapter { Ok(()) } + async fn list_sessions(&self, _database: Option<&str>) -> DbResult> { + let client = self.get_client().await?; + let mut client = client.lock().await; + + let stream = client + .simple_query( + "SELECT session_id, login_name, database_id, status, login_time \ + FROM sys.dm_exec_sessions ORDER BY session_id", + ) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + let rows = stream + .into_first_result() + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + let sessions = rows + .iter() + .map(|row| { + let session_id = row.try_get::(0).ok().flatten(); + let login_name = row + .try_get::<&str, _>(1) + .ok() + .flatten() + .map(|s| s.to_string()); + let database_id = row.try_get::(2).ok().flatten(); + let status = row + .try_get::<&str, _>(3) + .ok() + .flatten() + .map(|s| s.to_string()); + let login_time = row + .try_get::(4) + .ok() + .flatten() + .map(|d| d.to_string()); + serde_json::json!({ + "session_id": session_id, + "login_name": login_name, + "database_id": database_id, + "status": status, + "login_time": login_time, + }) + }) + .collect(); + + Ok(sessions) + } + + async fn kill_session(&self, session_id: &str) -> DbResult<()> { + let id: u16 = session_id.parse().map_err(|_| { + DbError::InvalidQuery(format!( + "Invalid session id '{}': SQL Server session ids are numeric", + session_id + )) + })?; + + let client = self.get_client().await?; + let mut client = client.lock().await; + + client + .simple_query(&format!("KILL {}", id)) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))? + .into_results() + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + Ok(()) + } + + async fn get_slow_queries( + &self, + _database: Option<&str>, + limit: Option, + ) -> DbResult> { + let limit = limit.unwrap_or(20); + + let client = self.get_client().await?; + let mut client = client.lock().await; + + let query = format!( + "SELECT TOP ({}) \ + qs.total_elapsed_time / 1000000.0 AS elapsed_seconds, \ + qs.execution_count, \ + qs.last_execution_time, \ + SUBSTRING(st.text, (qs.statement_start_offset / 2) + 1, \ + ((CASE qs.statement_end_offset \ + WHEN -1 THEN DATALENGTH(st.text) \ + ELSE qs.statement_end_offset END - qs.statement_start_offset) / 2) + 1) AS query_text \ + FROM sys.dm_exec_query_stats qs \ + CROSS APPLY sys.dm_exec_sql_text(qs.sql_handle) st \ + ORDER BY qs.total_elapsed_time DESC", + limit + ); + + let stream = client + .simple_query(&query) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + let rows = stream + .into_first_result() + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + let slow = rows + .iter() + .map(|row| { + let elapsed_seconds = row.try_get::(0).ok().flatten(); + let execution_count = row.try_get::(1).ok().flatten(); + let last_execution_time = row + .try_get::(2) + .ok() + .flatten() + .map(|d| d.to_string()); + let query_text = row + .try_get::<&str, _>(3) + .ok() + .flatten() + .map(|s| s.to_string()); + serde_json::json!({ + "elapsed_seconds": elapsed_seconds, + "execution_count": execution_count, + "last_execution_time": last_execution_time, + "query_text": query_text, + }) + }) + .collect(); + + Ok(slow) + } + + async fn grant_privilege(&self, privilege: &str, object: &str, grantee: &str) -> DbResult<()> { + let privs = Self::validate_privilege(privilege)?; + let obj = Self::quote_ident(object)?; + let grantee = Self::quote_ident(grantee)?; + let sql = format!("GRANT {} ON {} TO {}", privs, obj, grantee); + + let client = self.get_client().await?; + let mut client = client.lock().await; + + client + .simple_query(&sql) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))? + .into_results() + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + Ok(()) + } + + async fn revoke_privilege(&self, privilege: &str, object: &str, grantee: &str) -> DbResult<()> { + let privs = Self::validate_privilege(privilege)?; + let obj = Self::quote_ident(object)?; + let grantee = Self::quote_ident(grantee)?; + let sql = format!("REVOKE {} ON {} FROM {}", privs, obj, grantee); + + let client = self.get_client().await?; + let mut client = client.lock().await; + + client + .simple_query(&sql) + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))? + .into_results() + .await + .map_err(|e| DbError::QueryExecution(e.to_string()))?; + + Ok(()) + } + fn get_pool(&self) -> Option> { self.pool.clone() } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index dbaac2e1..f111122e 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -117,6 +117,7 @@ pub fn run() { crate::capabilities::sqlkit::register_all, crate::capabilities::sql::register_sql_tools, crate::capabilities::sql_write::register_write_tools, + crate::capabilities::dba::register_dba_tools, ]); // Initialize agent SQLite database diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs index 7db755ba..f13bb411 100644 --- a/src-tauri/src/mcp_bridge.rs +++ b/src-tauri/src/mcp_bridge.rs @@ -633,6 +633,7 @@ mod tests { crate::capabilities::sqlkit::register_all, crate::capabilities::sql::register_sql_tools, crate::capabilities::sql_write::register_write_tools, + crate::capabilities::dba::register_dba_tools, ]); } @@ -999,4 +1000,33 @@ mod tests { // proving execution reached the invoke path. assert_eq!(resp.status, 400); } + + #[test] + fn test_all_sql_capabilities_are_agent_tagged() { + init_registry_for_tests(); + let reg = data_studio_agent::capabilities::registry::registry(); + let all_tools = reg.agent_tools(); + let sql_tools: Vec<_> = all_tools + .iter() + .filter(|c| c.name.starts_with("sqlkit__")) + .collect(); + assert!( + sql_tools.len() >= 25, + "expected at least 25 sqlkit capabilities tagged for agent, got {}", + sql_tools.len() + ); + for name in &[ + "sqlkit__list_sessions", + "sqlkit__kill_session", + "sqlkit__get_slow_queries", + "sqlkit__grant_privilege", + "sqlkit__revoke_privilege", + ] { + assert!( + sql_tools.iter().any(|c| c.name == *name), + "missing expected DBA capability: {}", + name + ); + } + } }