use std::{collections::HashMap, sync::Arc, time::Instant}; use rmcp::{ handler::server::{router::tool::ToolRouter, wrapper::Parameters}, model::{CallToolResult, ContentBlock, Implementation, ServerCapabilities, ServerInfo}, schemars, tool, tool_handler, tool_router, ServerHandler, }; use serde::Deserialize; use serde_json::json; use uuid::Uuid; use crate::backend::{format_query_result, new_connection_config, parse_database_type, ConnectionSummary, DbxBackend}; use crate::mongo::{self, MongoCommand, MongoSafetyError}; use crate::session::{McpSession, McpSessionStore}; use crate::transaction::{TransactionFailure, TransactionResult, TransactionStatus}; use dbx_core::{ agent_tools::{ format_query_result_as_text, normalize_sql_for_confirmation, QueryCellWindow, MAX_EXECUTE_QUERY_ROWS, }, database_manifest, db::redis_driver::{classify_command, parse_command_argv, RedisCommandResult, RedisCommandSafety}, history::HistoryEntry, models::connection::ConnectionConfig, models::connection::DatabaseType, production_safety::{ is_production_database, mongo_pipeline_targets_production_database, sql_references_disallowed_database, targets_production_database, }, query_execution_sql::is_write_sql_for_database, sql_risk::{ classify_sql_risk_for_database, is_dangerous_sql_for_database, mcp_sql_has_forbidden_database_switch, SqlRisk, }, storage::{McpDatabaseScope, McpGlobalPolicy}, }; #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ListConnectionsRequest {} #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ListDatabasesRequest { #[serde(flatten)] pub selector: ConnectionSelector, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ConnectionSelector { #[schemars(description = "Unique ID of the DBX connection")] #[schemars(extend("type" = "string"))] pub connection_id: Option, #[schemars(description = "Name of the DBX connection")] #[schemars(extend("type" = "string"))] pub connection_name: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ListTablesRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Database name")] #[schemars(extend("type" = "string"))] pub database: Option, #[schemars(description = "Schema name")] #[schemars(extend("type" = "string"))] pub schema: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct DescribeTableRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Table name")] pub table: String, #[schemars(description = "Database name")] #[schemars(extend("type" = "string"))] pub database: Option, #[schemars(description = "Schema name")] #[schemars(extend("type" = "string"))] pub schema: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ListRoutinesRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Database name")] #[schemars(extend("type" = "string"))] pub database: Option, #[schemars(description = "Schema name")] #[schemars(extend("type" = "string"))] pub schema: Option, #[schemars( description = "Optional routine type filter: PROCEDURE or FUNCTION. Omit to list both. Unsupported databases return an empty list." )] #[schemars(extend("type" = "string"))] pub routine_type: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct GetRoutineSourceRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Routine name (as returned by dbx_list_routines)")] pub name: String, #[schemars(description = "Routine type: PROCEDURE or FUNCTION")] pub object_type: String, #[schemars(description = "Optional routine signature for databases that overload routine names (e.g. Oracle)")] #[schemars(extend("type" = "string"))] pub signature: Option, #[schemars(description = "Database name")] #[schemars(extend("type" = "string"))] pub database: Option, #[schemars(description = "Schema name")] #[schemars(extend("type" = "string"))] pub schema: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ExecuteQueryRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Database name")] #[schemars(extend("type" = "string"))] pub database: Option, #[schemars(description = "SQL query to execute")] pub sql: String, #[schemars( description = "Session ID from dbx_open_session. When set, the query runs on the session's pinned connection, preserving USE/SET and other session state across calls." )] #[schemars(extend("type" = "string"))] pub session_id: Option, #[schemars( description = "Start character offset for every string cell (default 0, max 1000000). Use the next offset reported by a truncated result to slide through a long value; narrow the query to the target row and column first." )] #[schemars(extend("type" = "integer"))] pub cell_char_offset: Option, #[schemars( description = "Maximum characters returned per string cell (default 200, max 4000). Increase only for an explicit long-value expansion." )] #[schemars(extend("type" = "integer"))] pub cell_char_limit: Option, #[schemars( description = "Maximum rows returned for this call (default 100, max 1000; larger values are clamped). SQL connections only - MongoDB shell commands always return at most 100 rows, and multi-statement scripts routed to the batch executor return at most 100 rows per statement." )] #[schemars(extend("type" = "integer"))] pub max_rows: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct OpenSessionRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Database name")] #[schemars(extend("type" = "string"))] pub database: Option, #[serde(default)] #[schemars(description = "Eagerly reserve one native MySQL connection for explicit transaction tools")] pub enable_transactions: bool, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct CloseSessionRequest { #[schemars(description = "Session ID returned by dbx_open_session")] pub session_id: String, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct TransactionSessionRequest { #[schemars(description = "Transaction-enabled session ID returned by dbx_open_session")] pub session_id: String, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct AddConnectionRequest { pub name: String, pub db_type: String, pub host: String, #[schemars(extend("type" = "integer"))] pub port: Option, #[serde(default)] pub username: String, #[serde(default)] pub password: String, #[schemars(extend("type" = "string"))] pub database: Option, #[serde(default)] pub ssl: bool, #[schemars(extend("type" = "string"))] pub driver_profile: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct DuplicateConnectionRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Name for the copied connection")] pub new_name: String, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct RemoveConnectionRequest { #[schemars(description = "Name of the DBX connection when connection_id is not provided")] #[schemars(extend("type" = "string"))] pub connection_name: Option, #[schemars(extend("type" = "string"))] pub connection_id: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ExecuteRedisCommandRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Redis logical database number")] #[schemars(extend("type" = "integer"))] pub db: Option, #[schemars(description = "Redis command to execute, for example GET mykey or INFO")] pub command: String, } #[derive(Debug, Default, serde::Serialize, Deserialize, schemars::JsonSchema)] #[serde(rename_all = "lowercase")] pub enum PeekStartPosition { #[default] Latest, Earliest, Offset, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct PeekMessagesRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Kafka topic name")] pub topic: String, #[schemars(description = "Number of messages, 1 to 100 (default 20)", extend("type" = "integer"))] pub count: Option, #[serde(default)] #[schemars(description = "Read latest messages (default), earliest messages, or from an offset")] pub start_position: PeekStartPosition, #[schemars(description = "Non-negative partition; omit to read across partitions", extend("type" = "integer"))] pub partition: Option, #[schemars(description = "Non-negative offset, required only for start_position=offset", extend("type" = "integer"))] pub offset: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct SendMessageRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Message queue topic name")] pub topic: String, #[schemars(description = "Message payload encoded as standard base64")] pub payload_base64: String, #[serde(default)] #[schemars(description = "Optional message key used for partitioning")] #[schemars(extend("type" = "string"))] pub key: Option, #[serde(default)] #[schemars(description = "Optional UTF-8 text preview of the payload")] #[schemars(extend("type" = "string"))] pub payload_text: Option, #[serde(default)] #[schemars(description = "Optional message headers")] pub headers: std::collections::HashMap, #[serde(default)] #[schemars(description = "Optional target partition")] #[schemars(extend("type" = "integer"))] pub partition: Option, #[serde(default)] #[schemars(description = "RabbitMQ exchange name")] #[schemars(extend("type" = "string"))] pub exchange: Option, #[serde(default)] #[schemars(description = "RabbitMQ routing key")] #[schemars(extend("type" = "string"))] pub routing_key: Option, #[serde(default)] #[schemars(description = "RabbitMQ virtual-host namespace")] #[schemars(extend("type" = "string"))] pub namespace: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct SchemaContextRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(extend("type" = "string"))] pub database: Option, #[schemars(extend("type" = "string"))] pub schema: Option, #[schemars(description = "Specific table names to include")] #[schemars(extend("type" = "array"))] pub tables: Option>, #[schemars(description = "Maximum number of tables to include, from 1 to 20")] #[schemars(extend("type" = "integer"))] pub max_tables: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct OpenTableRequest { #[serde(flatten)] pub selector: ConnectionSelector, pub table: String, #[schemars(extend("type" = "string"))] pub database: Option, #[schemars(extend("type" = "string"))] pub schema: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ExecuteAndShowRequest { #[serde(flatten)] pub selector: ConnectionSelector, pub sql: String, #[schemars(extend("type" = "string"))] pub database: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ExecuteBatchQueryRequest { #[serde(flatten)] pub selector: ConnectionSelector, #[schemars(description = "Database name")] #[schemars(extend("type" = "string"))] pub database: Option, #[schemars( description = "SQL script containing one or more statements. Statements are split using the database-dialect-aware splitter, so semicolons inside strings, comments, and stored procedures are handled." )] pub sql: String, #[schemars( description = "Session ID from dbx_open_session. When set, every statement runs on the session's pinned connection, preserving USE/SET, temporary tables and other session state across statements in the script and across calls." )] #[schemars(extend("type" = "string"))] pub session_id: Option, #[schemars( description = "When true, execution continues after a statement error instead of stopping at the first failure. Connection-level failures always stop the batch. Default false (stop at first failure)." )] #[schemars(extend("type" = "boolean"))] pub continue_on_error: Option, #[schemars( description = "When true and the script has more than one statement, all statements run inside one transaction (BEGIN … COMMIT) so the whole batch succeeds or rolls back, and the call returns a single merged result. For a single-statement script the option is ignored and the statement runs as normal auto-commit. Backends that cannot provide a rollbackable transaction reject this option. Cannot be combined with session_id or continue_on_error. Rejected for MySQL-family DDL scripts because DDL implicitly commits and cannot be rolled back. Default is auto-commit for each statement." )] #[schemars(extend("type" = "boolean"))] pub use_transaction: Option, } #[derive(Clone)] pub struct DbxMcpServer { backend: Arc, scope: McpScope, sessions: Arc, tool_router: ToolRouter, } #[derive(Clone, Debug, Default)] pub struct McpScope { pub connection_ids: Vec, pub connection_name: Option, pub database: Option, pub schema: Option, } struct ResolvedConnection { connection: dbx_core::models::connection::ConnectionConfig, policy: McpGlobalPolicy, database_scope: DatabaseScope, group_ids: Vec, } struct SessionCleanupResult { status: Option, error: Option, } #[derive(Clone, Debug)] enum DatabaseScope { All, Selected(Vec), None, } impl McpScope { pub fn from_env() -> Self { let mut connection_ids = scoped_connection_ids(std::env::var("DBX_MCP_SCOPE_CONNECTION_IDS").ok().as_deref()); if connection_ids.is_empty() { if let Some(connection_id) = non_empty_env("DBX_MCP_SCOPE_CONNECTION_ID") { connection_ids.push(connection_id); } } Self { connection_ids, connection_name: non_empty_env("DBX_MCP_SCOPE_CONNECTION_NAME"), database: non_empty_env("DBX_MCP_SCOPE_DATABASE"), schema: non_empty_env("DBX_MCP_SCOPE_SCHEMA"), } } fn enabled(&self) -> bool { self.connection_scope_enabled() || self.database.is_some() || self.schema.is_some() } fn connection_scope_enabled(&self) -> bool { !self.connection_ids.is_empty() || self.connection_name.is_some() } fn matches(&self, connection: &dbx_core::models::connection::ConnectionConfig) -> bool { if !self.connection_ids.is_empty() { return self.connection_ids.iter().any(|id| id == &connection.id); } self.connection_name.as_deref() == Some(connection.name.as_str()) } } impl DbxMcpServer { pub fn new(backend: Arc) -> Self { Self::with_runtime_options(backend, McpScope::from_env(), std::env::var_os("DBX_WEB_URL").is_some()) } pub fn with_runtime_options(backend: Arc, scope: McpScope, web_mode: bool) -> Self { // The workspace enables more than one rustls crypto feature through // transitive dependencies. Native MCP runs outside the desktop/web // startup paths, so select the same provider before any TLS tool call. let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let mut tool_router = Self::tool_router(); if scope.enabled() { tool_router.disable_route("dbx_add_connection"); tool_router.disable_route("dbx_duplicate_connection"); tool_router.disable_route("dbx_remove_connection"); } // Desktop UI bridge operations are intentionally unavailable remotely and in scoped AI sessions. if web_mode || scope.enabled() { tool_router.disable_route("dbx_open_table"); tool_router.disable_route("dbx_execute_and_show"); } #[cfg(not(feature = "mq-admin"))] { tool_router.disable_route("dbx_send_message"); tool_router.disable_route("dbx_peek_messages"); } Self { backend, scope, sessions: McpSessionStore::new(), tool_router } } fn spawn_session_cleanup(&self, session: McpSession) -> tokio::sync::oneshot::Receiver { let backend = self.backend.clone(); let sessions = self.sessions.clone(); let (result_tx, result_rx) = tokio::sync::oneshot::channel(); tokio::spawn(async move { let status = if let Some(owner) = &session.transaction_owner { match owner.close().await { Ok(status) => Some(status), Err(_) => { owner.wait_closed().await; Some(owner.status()) } } } else { None }; let close_result = backend .close_client_session(&session.connection_id, &session.database, &session.client_session_id) .await; let error = match close_result { Ok(_) => { sessions.finish_close(&session.id).await; None } Err(error) if session.transaction_owner.is_none() => { sessions.restore_after_failed_close(session).await; Some(error) } Err(error) => { sessions.finish_close(&session.id).await; Some(error) } }; let _ = result_tx.send(SessionCleanupResult { status, error }); }); result_rx } async fn close_backend_sessions_best_effort(&self, sessions: Vec) { let cleanups = sessions.into_iter().map(|session| self.spawn_session_cleanup(session)).collect::>(); for cleanup in cleanups { let _ = cleanup.await; } } fn spawn_transaction_session_open( &self, connection: ConnectionConfig, database: String, session: McpSession, ) -> ( tokio::sync::oneshot::Receiver, String>>, tokio::sync::oneshot::Sender<()>, ) { let backend = self.backend.clone(); let sessions = self.sessions.clone(); let (prepared_tx, prepared_rx) = tokio::sync::oneshot::channel(); let (ack_tx, ack_rx) = tokio::sync::oneshot::channel(); tokio::spawn(async move { let owner = match backend.open_transaction_owner(&connection, &database, &session.client_session_id).await { Ok(owner) => owner, Err(error) => { sessions.cancel_opening(&session.id).await; let _ = backend .close_client_session(&session.connection_id, &session.database, &session.client_session_id) .await; let _ = prepared_tx.send(Err(error)); return; } }; if prepared_tx.send(Ok(owner.clone())).is_err() || ack_rx.await.is_err() { sessions.cancel_opening(&session.id).await; let _ = owner.close().await; owner.wait_closed().await; let _ = backend .close_client_session(&session.connection_id, &session.database, &session.client_session_id) .await; } }); (prepared_rx, ack_tx) } async fn save_mcp_sql_history( &self, connection: &ConnectionConfig, database: &str, sql: &str, started_at: Instant, success: bool, error: Option, affected_rows: Option, ) { let entry = HistoryEntry { id: Uuid::new_v4().to_string(), connection_id: connection.id.clone(), connection_name: connection.name.clone(), database: database.to_string(), sql: sql.to_string(), executed_at: chrono::Utc::now().to_rfc3339(), execution_time_ms: started_at.elapsed().as_millis(), success, error, activity_kind: mcp_sql_activity_kind(sql, connection.db_type).to_string(), operation: mcp_sql_operation(sql, connection.db_type), target: String::new(), affected_rows, rollback_sql: None, details_json: Some(r#"{"source":"mcp"}"#.to_string()), }; if let Err(error) = self.backend.save_history_entry(&entry).await { log::warn!("failed to save MCP SQL history for connection {}: {error}", connection.id); } } } #[tool_router] impl DbxMcpServer { #[tool( name = "dbx_list_connections", description = "List database connections configured in DBX. Returns connection IDs, names, group paths, database types, endpoints, and selected databases." )] async fn list_connections( &self, Parameters(ListConnectionsRequest {}): Parameters, ) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_list_connections").await { return error; } match self.load_scoped_connections().await { Ok(connections) if connections.is_empty() => text("No connections configured in DBX."), Ok(connections) => { let group_paths = self.backend.load_connection_group_details().await.unwrap_or_default(); let rows = connections .iter() .map(|connection| { let mut summary = ConnectionSummary::from(connection); summary.group_path = group_paths.get(&connection.id).map(|path| path.names.clone()).unwrap_or_default(); summary }) .collect::>(); text(format_connections(&rows)) } Err(error) => backend_tool_error("CONNECTION_LOAD_ERROR", error), } } #[tool( name = "dbx_list_databases", description = "List database names available through a DBX connection. When the connection has a database allowlist, returns only the databases allowed by DBX MCP settings; use this before listing tables if the connection has no default database." )] async fn list_databases(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_list_databases").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; self.list_databases_for_resolved(&resolved).await } #[tool(name = "dbx_list_tables", description = "List tables and views for a database connection")] async fn list_tables(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_list_tables").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; let schema = match self.resolve_schema(request.schema) { Ok(schema) => schema, Err(error) => return error, }; match self.backend.list_tables(&resolved.connection, &database, &schema).await { Ok(tables) if tables.is_empty() => text("No tables found."), Ok(tables) => text( tables .into_iter() .map(|table| { let comment = table .comment .filter(|comment| !comment.is_empty()) .map(|comment| format!(" -- {comment}")) .unwrap_or_default(); format!("- {} ({}){}", table.name, table.table_type, comment) }) .collect::>() .join("\n"), ), Err(error) => tool_error("TABLE_LIST_ERROR", error), } } #[tool(name = "dbx_describe_table", description = "Get column definitions for a table")] async fn describe_table(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_describe_table").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; let schema = match self.resolve_schema(request.schema) { Ok(schema) => schema, Err(error) => return error, }; match self.backend.get_columns(&resolved.connection, &database, &schema, &request.table).await { Ok(columns) if columns.is_empty() => text("No columns found."), Ok(columns) => text(format_columns(&columns)), Err(error) => tool_error("TABLE_DESCRIPTION_ERROR", error), } } #[tool(name = "dbx_list_routines", description = "List stored procedures and functions for a database schema")] async fn list_routines(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_list_routines").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; let schema = match self.resolve_schema(request.schema) { Ok(schema) => schema, Err(error) => return error, }; let routine_types = match request.routine_type.as_deref() { None => None, Some(kind) => match normalize_routine_type(kind) { Ok(kind) => Some(vec![kind]), Err(error) => return tool_error("ROUTINE_LIST_ERROR", error), }, }; match self.backend.list_routines(&resolved.connection, &database, &schema, routine_types.as_deref()).await { Ok(routines) if routines.is_empty() => text("No routines found."), Ok(routines) => text(format_routines(&routines)), Err(error) => tool_error("ROUTINE_LIST_ERROR", error), } } #[tool(name = "dbx_get_routine_source", description = "Get the full source text of a stored procedure or function")] async fn get_routine_source(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_get_routine_source").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; let schema = match self.resolve_schema(request.schema) { Ok(schema) => schema, Err(error) => return error, }; let object_type = match normalize_routine_type(&request.object_type) { Ok(kind) => kind, Err(error) => return tool_error("ROUTINE_SOURCE_ERROR", error), }; match self .backend .get_routine_source( &resolved.connection, &database, &schema, &request.name, &object_type, request.signature.as_deref(), ) .await { Ok(source) if source.source.trim().is_empty() => text("Routine source is empty."), Ok(source) => text(source.source), Err(error) => tool_error("ROUTINE_SOURCE_ERROR", error), } } #[tool( name = "dbx_execute_query", description = "Execute a SQL query on a database connection (default 100 rows; pass max_rows to return up to 1000 rows; larger values are clamped; MongoDB shell commands and multi-statement scripts routed through the batch executor stay at 100 rows). For backwards compatibility, multi-statement scripts and stored-procedure scripts are routed through the dialect-aware batch executor." )] async fn execute_query_tool( &self, cancellation: tokio_util::sync::CancellationToken, parameters: Parameters, ) -> CallToolResult { cancellable_tool_call(cancellation, self.execute_query(parameters)).await } async fn execute_query(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_execute_query").await { return error; } let explicit_cell_window = (request.cell_char_offset.is_some() || request.cell_char_limit.is_some()) .then(|| QueryCellWindow::from_options(request.cell_char_offset, request.cell_char_limit)); let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let connection = &resolved.connection; if connection.db_type == dbx_core::models::connection::DatabaseType::Redis { return tool_error( "REDIS_COMMAND_REQUIRED", "Redis connections do not accept SQL through dbx_execute_query. Use dbx_execute_redis_command.", ); } if sql_requires_batch_execution(&request.sql, connection.db_type) { return self .execute_batch_request(ExecuteBatchQueryRequest { selector: ConnectionSelector { connection_id: Some(connection.id.clone()), connection_name: None }, database: request.database.clone(), sql: request.sql.clone(), session_id: request.session_id.clone(), continue_on_error: None, use_transaction: None, }) .await; } // Database discovery does not require a default database. In a // narrowed MCP scope it must never reveal names outside the allowlist, // so serve it directly from the scoped discovery path instead of // passing SHOW DATABASES to the database driver. if is_database_discovery_sql(&request.sql) { return self.list_databases_for_resolved(&resolved).await; } // Resolve the session before the database so its connection/database // binding is enforced on every stateful query. let session = match request.session_id.as_deref().map(str::trim).filter(|id| !id.is_empty()) { Some(session_id) => { let (session, expired) = self.sessions.resolve(session_id).await.into_parts(); self.close_backend_sessions_best_effort(expired).await; match session { Some(session) if session.connection_id == connection.id => Some(session), Some(_) => { return tool_error( "SESSION_CONNECTION_MISMATCH", format!("Session \"{session_id}\" is bound to a different connection."), ); } None => { return tool_error( "SESSION_NOT_FOUND", format!( "Session \"{session_id}\" not found or expired. Open a new one with dbx_open_session." ), ); } } } None => None, }; let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; let policy = effective_policy_for_database_with_groups(&resolved.policy, &resolved.group_ids, connection, &database); if let Some(session) = &session { if session.database != database { return tool_error( "SESSION_DATABASE_MISMATCH", format!( "Session \"{}\" is bound to database \"{}\", not \"{database}\".", session.id, session.database ), ); } } if connection.db_type == DatabaseType::MongoDb { // max_rows is deliberately ignored here: MongoDB shell commands // always return at most 100 rows, and this branch returns before the // arguments below are built. let command = match validate_mongo_command_with_groups( connection, &resolved.policy, &resolved.group_ids, &resolved.database_scope, &database, &request.sql, ) { Ok(command) => command, Err(error) => return error, }; return match self.backend.execute_mongo_command(connection, &database, &command).await { Ok(result) => match explicit_cell_window { Some(window) => match format_query_result_as_text(&result, 100, window) { Ok(output) => text(output), Err(error) => backend_tool_error("QUERY_FORMAT_ERROR", error), }, None => text(format_query_result(&result, 100)), }, Err(error) => backend_tool_error("QUERY_ERROR", error), }; } // A pinned session makes USE/SET CATALOG meaningful, so database // switching is allowed — unless a hard database scope is configured, // which a USE statement could otherwise escape. let allow_database_switch = session.is_some() && self.scope.database.is_none(); if connection.db_type != DatabaseType::MongoDb { if let Err(error) = ensure_sql_database_scope(&resolved.database_scope, connection, &database, &request.sql) { return error; } if let Err(error) = ensure_sql_database_execution_scope(&resolved.policy, connection, &database, &request.sql) { return error; } } let permissions = match validate_sql_policy(connection, &policy, &database, &request.sql, allow_database_switch) { Ok(permissions) => permissions, Err(error) => return error, }; let max_rows = request.max_rows.map(|value| value.clamp(1, MAX_EXECUTE_QUERY_ROWS as u64)).unwrap_or(100); if let Some(session) = session.as_ref().filter(|session| session.transaction_owner.is_some()) { if let Err(error) = validate_transaction_sql_shape(&request.sql) { return tool_error("TRANSACTION_SQL_BLOCKED", error); } if let Some(reason) = confirmed_batch_sql_block_reason( &request.sql, connection.db_type, permissions.confirmed_write_sql.as_deref(), ) { return tool_error("SQL_BLOCKED", reason); } let owner = session.transaction_owner.as_ref().expect("filtered transaction owner").clone(); let lease = match owner.acquire().await { Ok(lease) => lease, Err(error) => return transaction_failure_result(&session.id, error), }; // The lease is the queue boundary. Re-read every mutable policy and // connection binding after waiting so queued work cannot outrun a // revocation or configuration refresh. if let Err(error) = self.ensure_tool_allowed("dbx_execute_query").await { return error; } let refreshed = match self .resolve_connection(&ConnectionSelector { connection_id: Some(session.connection_id.clone()), connection_name: None, }) .await { Ok(resolved) => resolved, Err(error) => return error, }; if !transaction_connection_supported(&refreshed.connection) || refreshed.connection.id != session.connection_id { return tool_error("TRANSACTION_UNSUPPORTED", "The bound native MySQL connection changed."); } if let Err(error) = ensure_sql_database_scope( &refreshed.database_scope, &refreshed.connection, &session.database, &request.sql, ) { return error; } if let Err(error) = ensure_sql_database_execution_scope( &refreshed.policy, &refreshed.connection, &session.database, &request.sql, ) { return error; } let refreshed_policy = effective_policy_for_database_with_groups( &refreshed.policy, &refreshed.group_ids, &refreshed.connection, &session.database, ); let refreshed_permissions = match validate_sql_policy( &refreshed.connection, &refreshed_policy, &session.database, &request.sql, false, ) { Ok(permissions) => permissions, Err(error) => return error, }; if let Err(error) = validate_transaction_sql_shape(&request.sql) { return tool_error("TRANSACTION_SQL_BLOCKED", error); } if let Some(reason) = confirmed_batch_sql_block_reason( &request.sql, refreshed.connection.db_type, refreshed_permissions.confirmed_write_sql.as_deref(), ) { return tool_error("SQL_BLOCKED", reason); } let history_sql = request.sql.clone(); let started_at = Instant::now(); return match lease.query(request.sql, Some(max_rows as usize)).await { Ok(execution) => { let rendered = explicit_cell_window .and_then(|window| { format_query_result_as_text(&execution.result, max_rows as usize, window).ok() }) .unwrap_or_else(|| format_query_result(&execution.result, max_rows as usize)); self.save_mcp_sql_history( &refreshed.connection, &session.database, &history_sql, started_at, true, None, Some(execution.result.affected_rows as i64), ) .await; transaction_success_result(&session.id, &rendered, execution) } Err(error) => { self.save_mcp_sql_history( &refreshed.connection, &session.database, &history_sql, started_at, false, Some(error.message.clone()), None, ) .await; transaction_failure_result(&session.id, error) } }; } let history_sql = request.sql.clone(); let mut arguments = json!({ "sql": request.sql, "limit": max_rows }); if let Some(schema) = self.scope.schema.as_deref() { arguments["schema"] = json!(schema); } if let Some(session) = &session { arguments["client_session_id"] = json!(session.client_session_id); } if let Some(offset) = request.cell_char_offset { arguments["cell_char_offset"] = json!(offset); } if let Some(limit) = request.cell_char_limit { arguments["cell_char_limit"] = json!(limit); } // Surface the global MCP query timeout as a per-query argument. It is // re-read on every tool call, so a settings-page change takes effect on // the next MCP query with no client-config regeneration. A future // per-call `timeout_secs` tool parameter can slot in above this by not // being clobbered when set by the model. if let Some(secs) = resolved.policy.query_timeout_secs { arguments["timeout_secs"] = json!(secs); } let started_at = Instant::now(); let result = self.backend.execute_agent_tool(connection, &database, "execute_query", arguments, permissions).await; let success = !result.is_error; let error = result.is_error.then(|| result.content.trim_start_matches("Error: ").to_string()); self.save_mcp_sql_history(connection, &database, &history_sql, started_at, success, error, None).await; agent_result(result) } #[tool( name = "dbx_execute_batch", description = "Execute a SQL script containing multiple statements in one call and return a result per statement. Statements are split with a database-dialect-aware parser, so semicolons inside strings, comments and stored procedures are handled. Stops at the first failing statement unless continue_on_error is true. Pass session_id (from dbx_open_session) when statements must share one connection (e.g. temporary tables, USE/SET). When use_transaction is true and the script has multiple statements, the whole script runs in one transaction and the call returns a single merged result instead of one result per statement (single-statement scripts run normally and return that one result); use_transaction cannot be combined with session_id or continue_on_error, and is rejected when the backend cannot provide a rollbackable transaction or for MySQL-family DDL scripts (DDL implicitly commits and cannot be rolled back)." )] async fn execute_batch_tool( &self, cancellation: tokio_util::sync::CancellationToken, parameters: Parameters, ) -> CallToolResult { cancellable_tool_call(cancellation, self.execute_batch(parameters)).await } async fn execute_batch(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_execute_batch").await { return error; } self.execute_batch_request(request).await } async fn execute_batch_request(&self, request: ExecuteBatchQueryRequest) -> CallToolResult { let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let connection = &resolved.connection; if matches!(connection.db_type, DatabaseType::Redis | DatabaseType::MongoDb) { return tool_error( "DBX_BATCH_UNSUPPORTED", format!("Batch SQL execution is not available for {connection:?} connections."), ); } let sql = request.sql.trim(); if sql.is_empty() { return tool_error("SQL_BATCH_EMPTY", "SQL script cannot be empty."); } let session = match request.session_id.as_deref().map(str::trim).filter(|id| !id.is_empty()) { Some(session_id) => { let (session, expired) = self.sessions.resolve(session_id).await.into_parts(); self.close_backend_sessions_best_effort(expired).await; match session { Some(session) if session.connection_id == connection.id => Some(session), Some(_) => { return tool_error( "SESSION_CONNECTION_MISMATCH", format!("Session \"{session_id}\" is bound to a different connection."), ); } None => { return tool_error( "SESSION_NOT_FOUND", format!( "Session \"{session_id}\" not found or expired. Open a new one with dbx_open_session." ), ); } } } None => None, }; let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; if let Some(session) = &session { if session.database != database { return tool_error( "SESSION_DATABASE_MISMATCH", format!( "Session \"{}\" is bound to database \"{}\", not \"{database}\".", session.id, session.database ), ); } } // A pinned session makes USE/SET CATALOG meaningful, so database // switching is allowed — unless a hard database scope is configured. let allow_database_switch = session.is_some() && self.scope.database.is_none(); let policy = effective_policy_for_database_with_groups(&resolved.policy, &resolved.group_ids, connection, &database); if let Err(error) = ensure_sql_database_scope(&resolved.database_scope, connection, &database, sql) { return error; } if let Err(error) = ensure_sql_database_execution_scope(&resolved.policy, connection, &database, sql) { return error; } // Split with the same dialect-aware splitter the core uses, early, so the // option-validity checks below only fire for scripts that actually enter // transaction mode (more than one statement). A single-statement script // ignores use_transaction and runs as normal auto-commit, so combining it // with session_id / continue_on_error must still be allowed. let execution_plan = self.backend.execution_plan(connection, &database, sql).await; let statement_count = execution_plan.statements.len(); let transactional = request.use_transaction == Some(true) && statement_count > 1; // A transactional batch runs on a pooled connection and returns one // merged result, so it cannot preserve session state or yield per- // statement results. The core transaction path takes no client session // id, so combining the two would silently drop the session pin. Only // relevant when the script actually enters transaction mode. if transactional && session.is_some() { return tool_error( "TRANSACTION_WITH_SESSION_UNSUPPORTED", "use_transaction cannot be combined with session_id: transactional batches run on a pooled connection, discard session state, and return a single merged result. Run the batch either without use_transaction (auto-commit, one result per statement) or without session_id.", ); } // The core transaction path rolls back and stops on the first failure // (query.rs:2970), so continue_on_error is silently ignored there. // Accepting both would promise per-statement continuation that never // happens — reject the combination instead. Only relevant when the // script actually enters transaction mode. if transactional && request.continue_on_error == Some(true) { return tool_error( "TRANSACTION_WITH_CONTINUE_UNSUPPORTED", "use_transaction cannot be combined with continue_on_error: a transactional batch stops and rolls back at the first failure, so continue_on_error would be ignored. Run the batch either without use_transaction (auto-commit, one result per statement, may continue on error) or without continue_on_error.", ); } // DDL implicitly commits on MySQL-family and Oracle engines (and any // backend without transactional DDL — the shared capability check), so // the transaction wrapper cannot roll back a DDL that already committed // (e.g. "CREATE TABLE a; INSERT INTO missing" leaves the table behind // after the failed INSERT rolls back). A transactional batch containing // DDL would over-promise atomicity — reject it instead. Gated on multi- // statement like the other rejects: a single DDL statement is plain // auto-commit and safe. if transactional && dbx_core::query::batch_transaction_ddl_is_unrollbackable( Some(connection.db_type), &execution_plan.statements, ) { return tool_error( "TRANSACTION_WITH_DDL_UNSUPPORTED", "use_transaction cannot be used with a batch whose DDL cannot be rolled back: DDL statements implicitly commit and cannot be rolled back, so the batch would not be atomic. Run the batch without use_transaction (auto-commit, one result per statement) or split the DDL and DML into separate calls.", ); } // Classify the whole script as a unit so a single write/DDL statement in // the batch fails closed under read-only / dangerous-SQL / production // policy, exactly as the Web /api/query/execute-multi route does. let permissions = match validate_sql_policy(connection, &policy, &database, sql, allow_database_switch) { Ok(permissions) => permissions, Err(error) => return error, }; // When the user confirmed a specific write SQL for this run, the batch // must not smuggle a different write/DDL statement past that binding. // Fail closed on the whole script when it contains a write and is not // exactly the confirmed SQL — matching execute_query's anti-replay rule. if let Some(message) = confirmed_batch_sql_block_reason(sql, connection.db_type, permissions.confirmed_write_sql.as_deref()) { return tool_error("SQL_BLOCKED", message); } if let Some(session) = session.as_ref().filter(|session| session.transaction_owner.is_some()) { if request.use_transaction == Some(true) { return tool_error( "TRANSACTION_WITH_SESSION_UNSUPPORTED", "use_transaction cannot be combined with a transaction-enabled session; use dbx_begin_transaction and dbx_commit_transaction.", ); } for statement in &execution_plan.statements { if let Err(error) = validate_transaction_sql_shape(statement) { return tool_error("TRANSACTION_SQL_BLOCKED", error); } } let owner = session.transaction_owner.as_ref().expect("filtered transaction owner").clone(); let lease = match owner.acquire().await { Ok(lease) => lease, Err(error) => return transaction_failure_result(&session.id, error), }; if let Err(error) = self.ensure_tool_allowed("dbx_execute_batch").await { return error; } let refreshed = match self .resolve_connection(&ConnectionSelector { connection_id: Some(session.connection_id.clone()), connection_name: None, }) .await { Ok(resolved) => resolved, Err(error) => return error, }; if !transaction_connection_supported(&refreshed.connection) || refreshed.connection.id != session.connection_id { return tool_error("TRANSACTION_UNSUPPORTED", "The bound native MySQL connection changed."); } if let Err(error) = ensure_sql_database_scope(&refreshed.database_scope, &refreshed.connection, &session.database, sql) { return error; } if let Err(error) = ensure_sql_database_execution_scope(&refreshed.policy, &refreshed.connection, &session.database, sql) { return error; } let refreshed_policy = effective_policy_for_database_with_groups( &refreshed.policy, &refreshed.group_ids, &refreshed.connection, &session.database, ); let refreshed_permissions = match validate_sql_policy(&refreshed.connection, &refreshed_policy, &session.database, sql, false) { Ok(permissions) => permissions, Err(error) => return error, }; if let Some(message) = confirmed_batch_sql_block_reason( sql, refreshed.connection.db_type, refreshed_permissions.confirmed_write_sql.as_deref(), ) { return tool_error("SQL_BLOCKED", message); } for statement in &execution_plan.statements { if let Err(error) = validate_transaction_sql_shape(statement) { return tool_error("TRANSACTION_SQL_BLOCKED", error); } } let started_at = Instant::now(); let executions = lease .batch( execution_plan.statements.clone(), Some(BATCH_MAX_ROWS), request.continue_on_error.unwrap_or(false), ) .await; let mut results = Vec::with_capacity(executions.len()); let mut failures = Vec::new(); for (index, execution) in executions.into_iter().enumerate() { match execution { Ok(execution) => results.push(crate::backend::BatchStatementResult { result: execution.result, execution_error: false, statement_index: Some(index), error_message: None, merged: false, transaction_state: Some(execution.state), transaction_outcome: execution.outcome, }), Err(error) => { failures.push(error.message.clone()); results.push(crate::backend::BatchStatementResult { result: empty_query_result(), execution_error: true, statement_index: Some(index), error_message: Some(error.message), merged: false, transaction_state: Some(error.state), transaction_outcome: error.outcome, }); } } } let status = owner.status(); self.save_mcp_sql_history( &refreshed.connection, &session.database, sql, started_at, failures.is_empty(), (!failures.is_empty()).then(|| failures.join("; ")), failures.is_empty().then(|| { results.iter().map(|result| result.result.affected_rows).sum::().min(i64::MAX as u64) as i64 }), ) .await; let mut tool_result = text(format_batch_results(&results)); tool_result.structured_content = Some(json!({ "session_id": session.id, "transaction_state": status.state, "transaction_outcome": status.outcome, "results": results, })); return tool_result; } // A transactional batch runs on the plain (non-session) pool: the core // transaction wrapper ignores client_session_id (query.rs:2970 → // execute_statements_in_transaction_typed re-resolves the default pool // with get_or_create_pool). Creating an ephemeral session pool here // would open physical connections that are closed unused, so only // auto-commit batches get an ephemeral pinned pool — it keeps temp // tables and SET state alive for every statement in the script. let ephemeral_client_session_id = (!transactional && session.is_none()).then(|| format!("mcp-batch-{}", Uuid::new_v4())); let client_session_id = session .as_ref() .map(|session| session.client_session_id.clone()) .or_else(|| ephemeral_client_session_id.clone()); let mut options = dbx_core::query::QueryExecutionOptions { max_rows: Some(BATCH_MAX_ROWS), continue_on_error: request.continue_on_error.unwrap_or(false), use_transaction: request.use_transaction, client_session_id: client_session_id.clone(), ..Default::default() }; // Surface the global MCP query timeout. It is re-read on every tool // call, so a settings change takes effect without client-config // regeneration. options.timeout_secs = resolved.policy.query_timeout_secs; let schema = self.scope.schema.as_deref(); let started_at = Instant::now(); let execution = self.backend.execute_batch(connection, &database, schema, sql, options).await; if let Some(client_session_id) = ephemeral_client_session_id.as_deref() { if let Err(error) = self.backend.close_client_session(&connection.id, &database, client_session_id).await { log::warn!("failed to close ephemeral MCP batch session for {}: {error}", connection.id); } } match execution { Ok(mut results) => { // The core transaction path collapses the whole script into one // merged QueryResult. Mark it so callers can tell a merged // outcome from a per-statement result, and render it distinctly // instead of implying per-statement outcomes. The len()==1 guard // keeps drivers that return per-statement results even with // use_transaction (e.g. SQL Server scripts with result sets) // from being mislabelled as transactional. if transactional && results.len() == 1 { results[0].merged = true; results[0].statement_index = None; } let failures = results .iter() .filter(|result| result.execution_error) .map(|result| { result.error_message.clone().unwrap_or_else(|| match result.statement_index { Some(index) => format!("Statement {} failed", index + 1), None => "Batch execution failed".to_string(), }) }) .collect::>(); let success = failures.is_empty(); let affected_rows = success.then(|| { results.iter().map(|result| result.result.affected_rows).sum::().min(i64::MAX as u64) as i64 }); self.save_mcp_sql_history( connection, &database, sql, started_at, success, (!failures.is_empty()).then(|| failures.join("; ")), affected_rows, ) .await; let markdown = format_batch_results(&results); // Issue #7548 requires structured per-statement results so callers do not // parse concatenated text. MCP requires structuredContent to be an object, // so the array lives under `results`. let mut tool_result = CallToolResult::success(vec![ContentBlock::text(markdown)]); tool_result.structured_content = Some(serde_json::json!({ "results": results })); tool_result } Err(error) => { self.save_mcp_sql_history(connection, &database, sql, started_at, false, Some(error.clone()), None) .await; backend_tool_error("DBX_BATCH_EXECUTION_ERROR", error) } } } #[tool( name = "dbx_open_session", description = "Open a stateful query session pinned to a single backend connection. Returns a session ID for dbx_execute_query: USE, SET CATALOG, session variables and temporary tables persist across calls within the session. Close with dbx_close_session when done; idle sessions expire after 30 minutes." )] async fn open_session_tool( &self, cancellation: tokio_util::sync::CancellationToken, parameters: Parameters, ) -> CallToolResult { cancellable_tool_call(cancellation, self.open_session(parameters)).await } async fn open_session(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_open_session").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let connection = &resolved.connection; if matches!(connection.db_type, DatabaseType::Redis | DatabaseType::MongoDb) { return tool_error( "SESSION_UNSUPPORTED", format!("Sessions are only supported for SQL connections; \"{}\" is not one.", connection.name), ); } let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; if request.enable_transactions && !transaction_connection_supported(connection) { return tool_error( "TRANSACTION_UNSUPPORTED", "Fixed-session transactions require a native local MySQL connection without an external driver profile.", ); } let (session, expired) = if request.enable_transactions { self.sessions.reserve_opening(&connection.id, &database).await.into_parts() } else { self.sessions.open(&connection.id, &database).await.into_parts() }; self.close_backend_sessions_best_effort(expired).await; let session = match session { Ok(session) => session, Err(error) => return tool_error("SESSION_LIMIT", error), }; if request.enable_transactions { let (prepared, acknowledge) = self.spawn_transaction_session_open(connection.clone(), database.clone(), session.clone()); let owner = match prepared.await { Ok(Ok(owner)) => owner, Ok(Err(error)) => return backend_tool_error("TRANSACTION_UNSUPPORTED", error), Err(_) => return tool_error("SESSION_OPEN_ERROR", "Transaction session open worker stopped."), }; let session = match self.sessions.promote_opening(&session.id, owner).await { Ok(session) => session, Err(error) => return tool_error("SESSION_OPEN_ERROR", error), }; let mut result = text(format!( "Session opened.\nsession_id: {}\nconnection: {} (id: {})\ndatabase: {}\ntransactions: enabled\n\nPass session_id to dbx_execute_query to run every query on the same pinned connection. Close with dbx_close_session when done.", session.id, connection.name, connection.id, database )); result.structured_content = Some(json!({ "session_id": session.id, "transaction_state": "idle", "transaction_outcome": null, })); let _ = acknowledge.send(()); return result; } text(format!( "Session opened.\nsession_id: {}\nconnection: {} (id: {})\ndatabase: {}\n\nPass session_id to dbx_execute_query to run every query on the same pinned connection. Close with dbx_close_session when done.", session.id, connection.name, connection.id, database )) } #[tool( name = "dbx_begin_transaction", description = "Begin a transaction on a transaction-enabled native MySQL session" )] async fn begin_transaction_tool( &self, cancellation: tokio_util::sync::CancellationToken, parameters: Parameters, ) -> CallToolResult { cancellable_tool_call(cancellation, self.begin_transaction(parameters)).await } async fn begin_transaction(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_begin_transaction").await { return error; } let (session, expired) = self.sessions.resolve(&request.session_id).await.into_parts(); self.close_backend_sessions_best_effort(expired).await; let Some(session) = session else { return tool_error( "SESSION_NOT_FOUND", format!("Session \"{}\" not found or expired.", request.session_id), ); }; let Some(owner) = session.transaction_owner.clone() else { return tool_error("TRANSACTION_UNSUPPORTED", "This session was not opened with enable_transactions=true."); }; let lease = match owner.acquire().await { Ok(lease) => lease, Err(error) => return transaction_failure_result(&session.id, error), }; if let Err(error) = self.ensure_tool_allowed("dbx_begin_transaction").await { return error; } if let Err(error) = self.validate_transaction_control(&session, "BEGIN").await { return error; } match lease.begin().await { Ok(result) => transaction_success_result(&session.id, "Transaction started.", result), Err(error) => transaction_failure_result(&session.id, error), } } #[tool( name = "dbx_commit_transaction", description = "Commit the active transaction; committed is reported only after MySQL acknowledges COMMIT" )] async fn commit_transaction_tool( &self, cancellation: tokio_util::sync::CancellationToken, parameters: Parameters, ) -> CallToolResult { cancellable_tool_call(cancellation, self.commit_transaction(parameters)).await } async fn commit_transaction(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_commit_transaction").await { return error; } let (session, expired) = self.sessions.resolve(&request.session_id).await.into_parts(); self.close_backend_sessions_best_effort(expired).await; let Some(session) = session else { return tool_error( "SESSION_NOT_FOUND", format!("Session \"{}\" not found or expired.", request.session_id), ); }; let Some(owner) = session.transaction_owner.clone() else { return tool_error("TRANSACTION_UNSUPPORTED", "This session was not opened with enable_transactions=true."); }; let lease = match owner.acquire().await { Ok(lease) => lease, Err(error) => return transaction_failure_result(&session.id, error), }; if let Err(error) = self.ensure_tool_allowed("dbx_commit_transaction").await { return error; } if let Err(error) = self.validate_transaction_control(&session, "COMMIT").await { return error; } match lease.commit().await { Ok(result) => transaction_success_result(&session.id, "Transaction committed.", result), Err(error) => transaction_failure_result(&session.id, error), } } #[tool( name = "dbx_rollback_transaction", description = "Roll back the active transaction. Cleanup remains available after write-policy revocation." )] async fn rollback_transaction_tool( &self, cancellation: tokio_util::sync::CancellationToken, parameters: Parameters, ) -> CallToolResult { cancellable_tool_call(cancellation, self.rollback_transaction(parameters)).await } async fn rollback_transaction(&self, Parameters(request): Parameters) -> CallToolResult { let (session, expired) = self.sessions.resolve(&request.session_id).await.into_parts(); self.close_backend_sessions_best_effort(expired).await; let Some(session) = session else { return tool_error( "SESSION_NOT_FOUND", format!("Session \"{}\" not found or expired.", request.session_id), ); }; let Some(owner) = session.transaction_owner.clone() else { return tool_error("TRANSACTION_UNSUPPORTED", "This session was not opened with enable_transactions=true."); }; let lease = match owner.acquire().await { Ok(lease) => lease, Err(error) => return transaction_failure_result(&session.id, error), }; match lease.rollback().await { Ok(result) => transaction_success_result(&session.id, "Transaction rolled back.", result), Err(error) => transaction_failure_result(&session.id, error), } } #[tool( name = "dbx_close_session", description = "Close a stateful query session and release its pinned backend connection" )] async fn close_session_tool( &self, cancellation: tokio_util::sync::CancellationToken, parameters: Parameters, ) -> CallToolResult { cancellable_tool_call(cancellation, self.close_session(parameters)).await } async fn close_session(&self, Parameters(request): Parameters) -> CallToolResult { let (session, expired) = self.sessions.begin_close(&request.session_id).await.into_parts(); self.close_backend_sessions_best_effort(expired).await; let Some(session) = session else { return tool_error( "SESSION_NOT_FOUND", format!("Session \"{}\" not found or already closed.", request.session_id), ); }; let session_id = session.id.clone(); let cleanup = self.spawn_session_cleanup(session); let cleanup = match cleanup.await { Ok(cleanup) => cleanup, Err(_) => return tool_error("SESSION_CLOSE_ERROR", "Session cleanup worker stopped."), }; if let Some(error) = cleanup.error { return backend_tool_error("SESSION_CLOSE_ERROR", error); } let mut result = text(format!("Session \"{}\" closed.", session_id)); if let Some(status) = cleanup.status { result.structured_content = Some(json!({ "session_id": session_id, "transaction_state": status.state, "transaction_outcome": status.outcome, })); } result } #[tool(name = "dbx_execute_redis_command", description = "Execute a Redis command on a Redis connection")] async fn execute_redis_command( &self, Parameters(request): Parameters, ) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_execute_redis_command").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let connection = &resolved.connection; if connection.db_type != DatabaseType::Redis { return tool_error("INVALID_CONNECTION_TYPE", format!("Connection \"{}\" is not Redis.", connection.name)); } let database = match self.resolve_redis_database(request.db, &resolved) { Ok(database) => database, Err(error) => return error, }; let argv = match parse_command_argv(&request.command) { Ok(argv) => argv, Err(error) => return tool_error("REDIS_COMMAND_BLOCKED", error), }; let safety = classify_command(&argv[0]); let policy = effective_policy_for_database_with_groups( &resolved.policy, &resolved.group_ids, connection, &database.to_string(), ); let permissions = mcp_permissions(connection, &policy); if safety != RedisCommandSafety::Allowed && policy.read_only { return tool_error( "MCP_READ_ONLY", format!("MCP execution permission for Redis database {database} is read-only. Redis command blocked."), ); } if safety != RedisCommandSafety::Allowed && connection.read_only { return tool_error( "CONNECTION_READ_ONLY", format!("Connection \"{}\" has read-only protection enabled. Redis command blocked.", connection.name), ); } if safety == RedisCommandSafety::Blocked && !permissions.allow_dangerous { return tool_error( "REDIS_COMMAND_BLOCKED", format!( "Dangerous Redis command \"{}\" is disabled in DBX MCP settings.", argv[0].to_ascii_uppercase() ), ); } if safety != RedisCommandSafety::Allowed && !permissions.allow_writes { return tool_error( "REDIS_COMMAND_BLOCKED", "MCP Redis command execution is read-only in DBX MCP settings.", ); } // Production protection is stricter than the opt-in write flags by design. if safety != RedisCommandSafety::Allowed && is_production_database(connection, &database.to_string()) { return tool_error( "PRODUCTION_WRITE_BLOCKED", "MCP cannot execute write or dangerous Redis commands against a production database.", ); } match self .backend .execute_redis_command( connection, database, &request.command, safety == RedisCommandSafety::Blocked && permissions.allow_dangerous, ) .await { Ok(result) => text(format_redis_result(&result)), Err(error) => backend_tool_error("REDIS_COMMAND_ERROR", error), } } #[tool( name = "dbx_peek_messages", description = "Read up to 100 Kafka messages without committing consumer offsets. Defaults to the latest 20 messages. Returns JSON with base64 payloads, UTF-8 previews, metadata, incomplete (partial broker read), and outputTruncated (256 KiB message output budget). Only Kafka is supported." )] async fn peek_messages(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_peek_messages").await { return error; } #[cfg(not(feature = "mq-admin"))] { let _ = request; tool_error("MQ_UNSUPPORTED", "Message queue support is not compiled into this DBX MCP build.") } #[cfg(feature = "mq-admin")] { use dbx_core::mq::{ config::MqAdminConfig, MqSystemKind, PeekMessagesOptions, PeekStartPosition as Start, TopicRef, }; let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; match MqAdminConfig::from_connection(&resolved.connection) { Ok(config) if config.system_kind == MqSystemKind::Kafka => {} _ => { return tool_error("MQ_UNSUPPORTED", "Message reading through MCP supports Kafka connections only.") } } let count = request.count.unwrap_or(20); let start_position = match request.start_position { PeekStartPosition::Latest => Start::Latest, PeekStartPosition::Earliest => Start::Earliest, PeekStartPosition::Offset => Start::Offset, }; if request.topic.trim().is_empty() { return tool_error("MESSAGE_TOPIC_REQUIRED", "Message topic must not be empty."); } if !(1..=100).contains(&count) || request.partition.is_some_and(|value| value < 0) || request.offset.is_some_and(|value| value < 0) || (start_position == Start::Offset) != request.offset.is_some() { return tool_error("INVALID_PEEK_OPTIONS", "count must be 1..100; partition and offset must be non-negative; offset is required only for start_position=offset."); } let topic = TopicRef { tenant: "_kafka".into(), namespace: "default".into(), topic: request.topic.trim().into(), ..Default::default() }; let options = PeekMessagesOptions { start_position: Some(start_position), partition: request.partition, offset: request.offset, }; match self.backend.peek_messages(&resolved.connection, topic, count, options).await { Ok(result) => text(format_peek_messages_result(result, count as usize)), Err(error) => backend_tool_error("MESSAGE_PEEK_ERROR", error), } } } #[tool( name = "dbx_send_message", description = "Send a base64-encoded message to a Kafka, RocketMQ, or RabbitMQ topic/queue" )] async fn send_message(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_send_message").await { return error; } #[cfg(not(feature = "mq-admin"))] { let _ = request; return tool_error("MQ_UNSUPPORTED", "Message queue support is not compiled into this DBX MCP build."); } #[cfg(feature = "mq-admin")] { let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let connection = &resolved.connection; let mq_config = match dbx_core::mq::config::MqAdminConfig::from_connection(connection) { Ok(config) => config, Err(error) => return tool_error("MQ_UNSUPPORTED", error), }; if request.topic.trim().is_empty() { return tool_error("MESSAGE_TOPIC_REQUIRED", "Message topic must not be empty."); } let policy = effective_policy_for_database_with_groups( &resolved.policy, &resolved.group_ids, connection, connection.database.as_deref().unwrap_or(""), ); if policy.read_only { return tool_error( "MCP_READ_ONLY", "The effective MCP execution permission is read-only. Message sending is blocked.", ); } if connection.read_only { return tool_error( "CONNECTION_READ_ONLY", format!( "Connection \"{}\" has read-only protection enabled. Message sending is blocked.", connection.name ), ); } if dbx_core::production_safety::is_production_database( connection, connection.database.as_deref().unwrap_or(""), ) { return tool_error("PRODUCTION_WRITE_BLOCKED", "MCP cannot send messages to a production database."); } let request = dbx_core::mq::SendMessageRequest { topic: request.topic.trim().to_string(), key: request.key, payload_base64: request.payload_base64, payload_text: request.payload_text, headers: request.headers, partition: request.partition, exchange: request.exchange, routing_key: request.routing_key, namespace: request.namespace, }; match self.backend.send_message(connection, request).await { Ok(result) => text(format_send_message_result(&mq_config.system_kind, &result)), Err(error) => backend_tool_error("MESSAGE_SEND_ERROR", error), } } } #[tool(name = "dbx_get_schema_context", description = "Get compact table and column context for writing SQL")] async fn get_schema_context(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_get_schema_context").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let connection = &resolved.connection; let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; let schema = match self.resolve_schema(request.schema) { Ok(schema) => schema, Err(error) => return error, }; let max_tables = request.max_tables.unwrap_or(8).clamp(1, 20); let available = match self.backend.list_tables(connection, &database, &schema).await { Ok(tables) => tables, Err(error) => return tool_error("SCHEMA_CONTEXT_ERROR", error), }; let requested = request .tables .unwrap_or_default() .into_iter() .map(|name| name.to_ascii_lowercase()) .collect::>(); let mut selected = if requested.is_empty() { available.iter().collect::>() } else { available.iter().filter(|table| requested.contains(&table.name.to_ascii_lowercase())).collect::>() }; let truncated = selected.len() > max_tables || (requested.is_empty() && available.len() > max_tables); selected.truncate(max_tables); if selected.is_empty() { return text("No matching tables found."); } let mut tables = Vec::with_capacity(selected.len()); for table in selected { // Keep metadata calls sequential because some embedded drivers expose a single physical connection. let columns = match self.backend.get_columns(connection, &database, &schema, &table.name).await { Ok(columns) => columns, Err(error) => return tool_error("SCHEMA_CONTEXT_ERROR", error), }; tables.push((table.clone(), columns)); } text(format_schema_context(&connection.name, &database, &schema, &tables, truncated)) } #[tool(name = "dbx_add_connection", description = "Add a new database connection to DBX")] async fn add_connection(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_add_connection").await { return error; } let policy = match self.load_policy().await { Ok(policy) => policy, Err(error) => return error, }; if policy.read_only { return tool_error( "MCP_READ_ONLY", "DBX global MCP read-only mode is enabled. Connection management is not allowed.", ); } let connections = match self.backend.load_connections().await { Ok(connections) => connections, Err(error) => return tool_error("CONNECTION_LOAD_ERROR", error), }; if connections.iter().any(|connection| connection.name.eq_ignore_ascii_case(&request.name)) { return text(format!("Connection \"{}\" already exists.", request.name)); } let db_type = match parse_database_type(&request.db_type) { Ok(db_type) => db_type, Err(error) => return tool_error("INVALID_CONNECTION_TYPE", error), }; let port = match request.port.or_else(|| database_manifest::default_port(&db_type)) { Some(port) => port, None => return text("Port is required for this database type."), }; let config = match new_connection_config( Uuid::new_v4().to_string(), request.name, db_type, request.host, port, request.username, request.password, request.database, request.ssl, request.driver_profile, ) { Ok(config) => config, Err(error) => return tool_error("INVALID_CONNECTION", error), }; match self.backend.add_connection_for_mcp(config).await { Ok(config) => text(format!("Connection \"{}\" added (id: {}).", config.name, config.id)), Err(error) => backend_tool_error("CONNECTION_SAVE_ERROR", error), } } #[tool( name = "dbx_duplicate_connection", description = "Duplicate a DBX connection with its complete settings, credentials, tunnels, and sidebar group" )] async fn duplicate_connection( &self, Parameters(request): Parameters, ) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_duplicate_connection").await { return error; } let policy = match self.load_policy().await { Ok(policy) => policy, Err(error) => return error, }; if policy.read_only { return tool_error( "MCP_READ_ONLY", "DBX global MCP read-only mode is enabled. Connection management is not allowed.", ); } let connections = match self.backend.load_connections().await { Ok(connections) => connections, Err(error) => return tool_error("CONNECTION_LOAD_ERROR", error), }; let group_paths = match self.load_group_paths_for_policy(&policy).await { Ok(paths) => paths, Err(error) => return backend_tool_error("MCP_POLICY_UNAVAILABLE", error), }; let allowed = connections .iter() .filter(|connection| policy_allows_connection(&policy, group_paths.get(&connection.id), connection)) .cloned() .collect::>(); let source = if let Some(id) = request.selector.connection_id.as_deref().map(str::trim).filter(|id| !id.is_empty()) { allowed.iter().find(|connection| connection.id == id).cloned() } else if let Some(name) = request.selector.connection_name.as_deref().map(str::trim).filter(|name| !name.is_empty()) { let matching = allowed .iter() .filter(|connection| connection.name.eq_ignore_ascii_case(name)) .cloned() .collect::>(); if matching.len() > 1 { return tool_error("AMBIGUOUS_CONNECTION", ambiguous_connections(name, &matching)); } matching.into_iter().next() } else { return tool_error("CONNECTION_NOT_FOUND", "Either connection_id or connection_name is required."); }; let Some(source) = source else { return tool_error("CONNECTION_NOT_FOUND", "The source connection was not found or is outside MCP scope."); }; let new_name = request.new_name.trim(); if new_name.is_empty() { return tool_error("INVALID_CONNECTION", "The copied connection name must not be empty."); } if connections.iter().any(|connection| connection.name.eq_ignore_ascii_case(new_name)) { return tool_error("CONNECTION_ALREADY_EXISTS", format!("Connection \"{new_name}\" already exists.")); } match self.backend.duplicate_connection_for_mcp(&source.id, &Uuid::new_v4().to_string(), new_name).await { Ok(copy) => text(format!("Connection \"{}\" duplicated (id: {}).", copy.name, copy.id)), Err(error) => backend_tool_error("CONNECTION_SAVE_ERROR", error), } } #[tool(name = "dbx_remove_connection", description = "Remove a database connection from DBX")] async fn remove_connection(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_remove_connection").await { return error; } let policy = match self.load_policy().await { Ok(policy) => policy, Err(error) => return error, }; if policy.read_only { return tool_error( "MCP_READ_ONLY", "DBX global MCP read-only mode is enabled. Connection management is not allowed.", ); } let connections = match self.backend.load_connections().await { Ok(connections) => connections, Err(error) => return tool_error("CONNECTION_LOAD_ERROR", error), }; let target = if let Some(id) = request.connection_id.as_deref().map(str::trim).filter(|id| !id.is_empty()) { connections.iter().find(|connection| connection.id == id).cloned() } else { let Some(name) = request.connection_name.as_deref().map(str::trim).filter(|name| !name.is_empty()) else { return tool_error("CONNECTION_NOT_FOUND", "Either connection_id or connection_name is required."); }; let matching = connections .iter() .filter(|connection| connection.name.eq_ignore_ascii_case(name)) .cloned() .collect::>(); if matching.len() > 1 { return tool_error("AMBIGUOUS_CONNECTION", ambiguous_connections(name, &matching)); } matching.into_iter().next() }; let Some(target) = target else { return tool_error( "CONNECTION_NOT_FOUND", request .connection_id .as_deref() .filter(|id| !id.trim().is_empty()) .map(|id| format!("Connection with id \"{id}\" not found.")) .unwrap_or_else(|| { format!("Connection \"{}\" not found.", request.connection_name.as_deref().unwrap_or_default()) }), ); }; let group_paths = match self.load_group_paths_for_policy(&policy).await { Ok(paths) => paths, Err(error) => return backend_tool_error("MCP_POLICY_UNAVAILABLE", error), }; if !policy_allows_connection(&policy, group_paths.get(&target.id), &target) { return tool_error( "CONNECTION_OUT_OF_SCOPE", format!("Connection \"{}\" is not allowed by DBX MCP settings.", target.id), ); } match self.backend.remove_connection_for_mcp(&target.id).await { Ok(true) => text(format!("Connection \"{}\" (id: {}) removed.", target.name, target.id)), Ok(false) => tool_error("CONNECTION_NOT_FOUND", format!("Connection \"{}\" not found.", target.name)), Err(error) => backend_tool_error("CONNECTION_SAVE_ERROR", error), } } #[tool(name = "dbx_open_table", description = "Open a table in DBX desktop app. Requires DBX to be running.")] async fn open_table(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_open_table").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let connection = &resolved.connection; let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; let schema = match self.resolve_schema(request.schema) { Ok(schema) => schema, Err(error) => return error, }; match self .backend .bridge_request( "/open-table", json!({ "connection_id": connection.id, "connection_name": connection.name, "table": request.table, "database": database, "schema": schema, }), ) .await { Ok(()) => text(format!("Opened {} in DBX", request.table)), Err(error) => backend_tool_error("DBX_NOT_RUNNING", error), } } #[tool( name = "dbx_execute_and_show", description = "Execute a SQL query in DBX desktop app UI and show results there. Requires DBX to be running." )] async fn execute_and_show(&self, Parameters(request): Parameters) -> CallToolResult { if let Err(error) = self.ensure_tool_allowed("dbx_execute_and_show").await { return error; } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, }; let connection = &resolved.connection; if connection.db_type == DatabaseType::Redis { return tool_error("REDIS_COMMAND_REQUIRED", "Use dbx_execute_redis_command for Redis connections."); } let database = match self.resolve_database(request.database, &resolved) { Ok(database) => database, Err(error) => return error, }; let policy = effective_policy_for_database_with_groups(&resolved.policy, &resolved.group_ids, connection, &database); if connection.db_type != DatabaseType::MongoDb { if let Err(error) = ensure_sql_database_scope(&resolved.database_scope, connection, &database, &request.sql) { return error; } if let Err(error) = ensure_sql_database_execution_scope(&resolved.policy, connection, &database, &request.sql) { return error; } } let permissions = if connection.db_type == DatabaseType::MongoDb { mcp_permissions(connection, &policy) } else { match validate_sql_policy(connection, &policy, &database, &request.sql, false) { Ok(permissions) => permissions, Err(error) => return error, } }; if connection.db_type == DatabaseType::MongoDb { if let Err(error) = validate_mongo_command_with_groups( connection, &resolved.policy, &resolved.group_ids, &resolved.database_scope, &database, &request.sql, ) { return error; } } match self .backend .bridge_request( "/execute-query", json!({ "connection_id": connection.id, "connection_name": connection.name, "sql": request.sql, "database": database, "allow_writes": permissions.allow_writes, "allow_dangerous": permissions.allow_dangerous, }), ) .await { Ok(()) => text("Query sent to DBX"), Err(error) => backend_tool_error("DBX_NOT_RUNNING", error), } } } impl DbxMcpServer { async fn list_databases_for_resolved(&self, resolved: &ResolvedConnection) -> CallToolResult { match &resolved.database_scope { DatabaseScope::None => tool_error( "DATABASE_OUT_OF_SCOPE", "This connection is configured with no database access in DBX MCP settings.", ), DatabaseScope::Selected(databases) => text(format_database_names(databases)), DatabaseScope::All => match self.backend.list_databases(&resolved.connection).await { Ok(databases) => text(format_database_names(&databases)), Err(error) => tool_error("DATABASE_LIST_ERROR", error), }, } } async fn load_scoped_connections(&self) -> Result, String> { let policy = self.backend.load_mcp_global_policy().await?; let connections = self.backend.load_connections().await?; let group_paths = self.load_group_paths_for_policy(&policy).await?; Ok(connections .into_iter() .filter(|connection| policy_allows_connection(&policy, group_paths.get(&connection.id), connection)) .filter(|connection| !self.scope.connection_scope_enabled() || self.scope.matches(connection)) .collect()) } async fn load_group_paths_for_policy( &self, policy: &McpGlobalPolicy, ) -> Result, String> { match self.backend.load_connection_group_details().await { Ok(paths) => Ok(paths), Err(error) if dbx_core::mcp_policy::policy_uses_connection_groups(policy) => { Err(format!("MCP_POLICY_UNAVAILABLE: {error}")) } Err(_) => Ok(HashMap::new()), } } // CallToolResult is the rmcp wire response type; keeping it unboxed avoids conversions at every tool boundary. #[allow(clippy::result_large_err)] async fn load_policy(&self) -> Result { self.backend.load_mcp_global_policy().await.map_err(|error| backend_tool_error("MCP_POLICY_UNAVAILABLE", error)) } // CallToolResult is the rmcp wire response type; keeping it unboxed avoids conversions at every tool boundary. #[allow(clippy::result_large_err)] async fn validate_transaction_control(&self, session: &McpSession, operation: &str) -> Result<(), CallToolResult> { let resolved = self .resolve_connection(&ConnectionSelector { connection_id: Some(session.connection_id.clone()), connection_name: None, }) .await?; if !transaction_connection_supported(&resolved.connection) { return Err(tool_error( "TRANSACTION_UNSUPPORTED", "The session connection is no longer a native MySQL connection.", )); } ensure_database_in_scope(&resolved.database_scope, &session.database)?; let policy = effective_policy_for_database_with_groups( &resolved.policy, &resolved.group_ids, &resolved.connection, &session.database, ); if policy.read_only { return Err(tool_error( "MCP_READ_ONLY", format!("MCP execution permission is read-only; {operation} is write-capable."), )); } if resolved.connection.read_only { return Err(tool_error( "CONNECTION_READ_ONLY", format!("Connection \"{}\" is read-only; {operation} is blocked.", resolved.connection.name), )); } if is_production_database(&resolved.connection, &session.database) { return Err(tool_error( "PRODUCTION_WRITE_BLOCKED", format!("{operation} is not allowed against a production database."), )); } Ok(()) } // CallToolResult is the rmcp wire response type; keeping it unboxed avoids conversions at every tool boundary. #[allow(clippy::result_large_err)] async fn ensure_tool_allowed(&self, tool_name: &str) -> Result<(), CallToolResult> { let policy = self.load_policy().await?; if policy_allows_tool(&policy, tool_name) { Ok(()) } else { Err(tool_error("TOOL_OUT_OF_SCOPE", format!("Tool \"{tool_name}\" is not allowed by DBX MCP settings."))) } } /// The advertised `tools/list` view: router-enabled tools narrowed to the /// global policy's tool allowlist. When the policy cannot be loaded every /// tool call already fails with `MCP_POLICY_UNAVAILABLE`, so fall back to /// the plain router view rather than adding a new failure mode here. async fn policy_filtered_tools(&self) -> Vec { match self.load_policy().await { Ok(policy) => self .tool_router .list_all() .into_iter() .filter(|tool| policy_allows_tool(&policy, tool.name.as_ref())) .collect(), Err(_) => self.tool_router.list_all(), } } // CallToolResult is the rmcp wire response type; keeping it unboxed avoids conversions at every tool boundary. #[allow(clippy::result_large_err)] fn resolve_database( &self, requested: Option, resolved: &ResolvedConnection, ) -> Result { let requested = requested.map(|database| database.trim().to_string()).filter(|database| !database.is_empty()); let database = if let Some(scoped) = self.scope.database.as_deref() { if let Some(requested) = requested.as_deref() { if requested != scoped { return Err(tool_error( "DATABASE_OUT_OF_SCOPE", format!("Database \"{requested}\" is outside the scoped database \"{scoped}\"."), )); } } scoped.to_string() } else { requested.or_else(|| resolved.connection.database.clone()).unwrap_or_default() }; ensure_database_in_scope(&resolved.database_scope, &database)?; Ok(database) } /// Resolve the schema for scoped CLI agents. A selected schema is a hard /// bound, matching the existing database scope behavior. #[allow(clippy::result_large_err)] fn resolve_schema(&self, requested: Option) -> Result { let requested = requested.map(|schema| schema.trim().to_string()).filter(|schema| !schema.is_empty()); if let Some(scoped) = self.scope.schema.as_deref() { if let Some(requested) = requested.as_deref() { if requested != scoped { return Err(tool_error( "SCHEMA_OUT_OF_SCOPE", format!("Schema \"{requested}\" is outside the scoped schema \"{scoped}\"."), )); } } return Ok(scoped.to_string()); } Ok(requested.unwrap_or_default()) } // CallToolResult is the rmcp wire response type; keeping it unboxed avoids conversions at every tool boundary. #[allow(clippy::result_large_err)] fn resolve_redis_database( &self, requested: Option, resolved: &ResolvedConnection, ) -> Result { if let Some(scoped) = self.scope.database.as_deref() { let scoped_database = parse_redis_database(scoped).ok_or_else(|| { tool_error( "INVALID_DATABASE_SCOPE", format!("Redis database scope \"{scoped}\" must be a non-negative integer."), ) })?; if let Some(requested) = requested { if requested != scoped_database { return Err(tool_error( "DATABASE_OUT_OF_SCOPE", format!("Redis database {requested} is outside the scoped database {scoped_database}."), )); } } ensure_database_in_scope(&resolved.database_scope, &scoped_database.to_string())?; return Ok(scoped_database); } let database = requested.or_else(|| redis_database(&resolved.connection)).unwrap_or(0); ensure_database_in_scope(&resolved.database_scope, &database.to_string())?; Ok(database) } // CallToolResult is the rmcp wire response type; keeping it unboxed avoids conversions at every tool boundary. #[allow(clippy::result_large_err)] async fn resolve_connection(&self, selector: &ConnectionSelector) -> Result { let policy = self.load_policy().await?; let group_paths = self .load_group_paths_for_policy(&policy) .await .map_err(|error| backend_tool_error("MCP_POLICY_UNAVAILABLE", error))?; let connections = self.backend.load_connections().await.map_err(|error| tool_error("CONNECTION_LOAD_ERROR", error))?; if let Some(id) = selector.connection_id.as_deref().map(str::trim).filter(|id| !id.is_empty()) { let connection = connections .into_iter() .find(|connection| connection.id == id) .ok_or_else(|| tool_error("CONNECTION_NOT_FOUND", format!("Connection with id \"{id}\" not found.")))?; if self.scope.connection_scope_enabled() && !self.scope.matches(&connection) { return Err(tool_error( "CONNECTION_OUT_OF_SCOPE", format!("Connection \"{id}\" is outside this DBX AI session scope."), )); } if !policy_allows_connection(&policy, group_paths.get(&connection.id), &connection) { return Err(tool_error( "CONNECTION_OUT_OF_SCOPE", format!("Connection \"{id}\" is not allowed by DBX MCP settings."), )); } return Ok(resolved_connection(policy, connection, group_paths.get(id))); } if self.scope.connection_scope_enabled() { let connection = connections .into_iter() .find(|connection| self.scope.matches(connection)) .ok_or_else(|| tool_error("CONNECTION_NOT_FOUND", "Scoped DBX connection was not found."))?; if let Some(name) = selector.connection_name.as_deref().map(str::trim).filter(|name| !name.is_empty()) { if name != connection.name && name != connection.id { return Err(tool_error( "CONNECTION_OUT_OF_SCOPE", format!("Connection \"{name}\" is outside this DBX AI session scope."), )); } } if !policy_allows_connection(&policy, group_paths.get(&connection.id), &connection) { return Err(tool_error( "CONNECTION_OUT_OF_SCOPE", "The DBX AI session scope is outside the global MCP connection allowlist.", )); } let path = group_paths.get(&connection.id); return Ok(resolved_connection(policy, connection, path)); } let Some(name) = selector.connection_name.as_deref().map(str::trim).filter(|name| !name.is_empty()) else { return Err(tool_error("CONNECTION_NOT_FOUND", "Either connection_id or connection_name is required.")); }; let matching = connections.into_iter().filter(|connection| connection.name.eq_ignore_ascii_case(name)).collect::>(); let allowed = matching .iter() .filter(|connection| policy_allows_connection(&policy, group_paths.get(&connection.id), connection)) .cloned() .collect::>(); match allowed.as_slice() { [] if matching.is_empty() => { Err(tool_error("CONNECTION_NOT_FOUND", format!("Connection \"{name}\" not found."))) } [] => Err(tool_error( "CONNECTION_OUT_OF_SCOPE", format!("Connection \"{name}\" is not allowed by DBX MCP settings."), )), [connection] => Ok(resolved_connection(policy, connection.clone(), group_paths.get(&connection.id))), _ => Err(tool_error("AMBIGUOUS_CONNECTION", ambiguous_connections(name, &allowed))), } } } #[tool_handler(router = self.tool_router)] impl ServerHandler for DbxMcpServer { fn get_info(&self) -> ServerInfo { ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) .with_server_info(Implementation::new("dbx", env!("CARGO_PKG_VERSION"))) .with_instructions("Use DBX connections to inspect schemas and query databases safely.") } /// Hide tools the global policy disallows from the advertised list, the /// same way scope-based route disabling does. Per-call enforcement stays /// in `ensure_tool_allowed`; this only stops clients (especially LLMs) /// from seeing — and repeatedly retrying — capabilities that would be /// rejected anyway. When the policy cannot be loaded, every call already /// fails with `MCP_POLICY_UNAVAILABLE`, so the listing falls back to the /// router view instead of adding a new failure mode here. async fn list_tools( &self, _request: Option, _context: rmcp::service::RequestContext, ) -> Result { Ok(rmcp::model::ListToolsResult { tools: self.policy_filtered_tools().await, meta: None, next_cursor: None }) } } fn text(value: impl Into) -> CallToolResult { CallToolResult::success(vec![ContentBlock::text(value)]) } async fn cancellable_tool_call( cancellation: tokio_util::sync::CancellationToken, operation: impl std::future::Future, ) -> CallToolResult { tokio::select! { result = operation => result, _ = cancellation.cancelled() => tool_error("REQUEST_CANCELLED", "The MCP request was cancelled."), } } fn transaction_success_result(session_id: &str, message: &str, execution: TransactionResult) -> CallToolResult { let mut result = text(format!( "{message}\nsession_id: {session_id}\ntransaction_state: {}{}", execution.state.as_str(), execution.outcome.map(|outcome| format!("\ntransaction_outcome: {}", outcome.as_str())).unwrap_or_default() )); result.structured_content = Some(json!({ "session_id": session_id, "transaction_state": execution.state, "transaction_outcome": execution.outcome, "result": execution.result, })); result } fn transaction_failure_result(session_id: &str, failure: TransactionFailure) -> CallToolResult { let mut result = tool_error(failure.code, &failure.message); result.structured_content = Some(json!({ "session_id": session_id, "transaction_state": failure.state, "transaction_outcome": failure.outcome, "mysql_code": failure.mysql_code, "error": failure.message, })); result } fn tool_error(code: &str, message: impl Into) -> CallToolResult { CallToolResult::error(vec![ContentBlock::text(format!("Error [{code}]: {}", message.into()))]) } fn backend_tool_error(default_code: &str, error: impl Into) -> CallToolResult { let error = error.into(); for code in [ "MCP_POLICY_UNAVAILABLE", "MCP_READ_ONLY", "CONNECTION_OUT_OF_SCOPE", "DATABASE_OUT_OF_SCOPE", "DATABASE_EXECUTION_POLICY_OUT_OF_SCOPE", "INVALID_DATABASE_SCOPE", "CONNECTION_READ_ONLY", "PRODUCTION_DATABASE_READ_ONLY", "PRODUCTION_WRITE_BLOCKED", "SQL_BLOCKED", "TRANSACTION_UNSUPPORTED", ] { let marker = format!("{code}:"); if let Some(index) = error.find(&marker) { return tool_error(code, error[index + marker.len()..].trim()); } } tool_error(default_code, error) } /// Keep older MCP clients compatible with the original single-query tool when /// they send a complete SQL script. The batch executor removes client-side /// commands such as MySQL `DELIMITER` and preserves semicolons inside routine /// bodies before dispatching statements to the database. fn sql_requires_batch_execution(sql: &str, database_type: DatabaseType) -> bool { if sql.lines().any(|line| line.trim_start().to_ascii_lowercase().starts_with("delimiter ")) { return true; } dbx_core::sql::sql_execution_plan_for_database(sql, database_type).statements.len() > 1 } /// Maximum rows returned per statement in a `dbx_execute_batch` call. The batch /// executor is deliberately tighter than dbx_execute_query's default: a script /// has many statements, so its row budget multiplies. const BATCH_MAX_ROWS: usize = 100; fn mcp_sql_activity_kind(sql: &str, database_type: DatabaseType) -> &'static str { let risks = dbx_core::sql::sql_execution_plan_for_database(sql, database_type) .statements .into_iter() .filter_map(|statement| classify_sql_risk_for_database(&statement, database_type).ok()) .collect::>(); if risks.contains(&SqlRisk::Ddl) { "schema_change" } else if risks.contains(&SqlRisk::Write) { "data_change" } else { "query" } } fn mcp_sql_operation(sql: &str, database_type: DatabaseType) -> String { dbx_core::sql::sql_execution_plan_for_database(sql, database_type) .statements .first() .map(|statement| dbx_core::query_execution_sql::strip_sql_comments(statement)) .and_then(|statement| statement.split_whitespace().next().map(str::to_string)) .map(|operation| operation.trim_matches(|character: char| !character.is_ascii_alphabetic()).to_string()) .filter(|operation| !operation.is_empty()) .map(|operation| operation.to_ascii_uppercase()) .unwrap_or_else(|| "SQL".to_string()) } /// Render a multi-statement batch result as one Markdown text block per /// statement. Each statement is labelled by its index so callers can see which /// statement failed, how many rows each affected, or what each returned. A /// `merged` entry (use_transaction mode) is labelled as the transaction outcome /// instead of a per-statement result. fn format_batch_results(results: &[crate::backend::BatchStatementResult]) -> String { let mut output = String::new(); for (index, result) in results.iter().enumerate() { if index > 0 { output.push('\n'); } if result.merged { output.push_str("### Transaction outcome"); output.push('\n'); if result.execution_error { let message = result.error_message.as_deref().unwrap_or("Transaction failed"); output.push_str(&format!("**Status:** failed\n\n{message}\n")); } else { let rendered = format_query_result_as_text(&result.result, BATCH_MAX_ROWS, QueryCellWindow::default()) .unwrap_or_else(|error| error); output.push_str(rendered.trim_start_matches("Query executed. ")); } continue; } let heading = match result.statement_index { Some(statement_index) => format!("### Statement {}", statement_index + 1), None => format!("### Statement {}", index + 1), }; output.push_str(&heading); output.push('\n'); if result.execution_error { let message = result.error_message.as_deref().unwrap_or("Statement failed"); output.push_str(&format!("**Status:** failed\n\n{message}\n")); } else { let rendered = format_query_result_as_text(&result.result, BATCH_MAX_ROWS, QueryCellWindow::default()) .unwrap_or_else(|error| error); output.push_str(rendered.trim_start_matches("Query executed. ")); } } output } fn empty_query_result() -> dbx_core::db::QueryResult { dbx_core::db::QueryResult { columns: Vec::new(), column_types: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: Vec::new(), affected_rows: 0, execution_time_ms: 0, server_execute_time_us: None, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), } } fn agent_result(result: dbx_core::agent_events::ToolResult) -> CallToolResult { if result.is_error { backend_tool_error("DBX_TOOL_ERROR", result.content.trim_start_matches("Error: ")) } else { text(result.content) } } fn policy_allows_connection( policy: &McpGlobalPolicy, group_path: Option<&dbx_core::mcp_policy::McpConnectionGroupPath>, connection: &dbx_core::models::connection::ConnectionConfig, ) -> bool { dbx_core::mcp_policy::policy_allows_connection(policy, group_path, &connection.id) } fn policy_allows_tool(policy: &McpGlobalPolicy, tool_name: &str) -> bool { policy.allowed_tool_names.as_ref().is_none_or(|allowed| allowed.iter().any(|name| name == tool_name)) } fn database_scope_for_connection( policy: &McpGlobalPolicy, connection: &dbx_core::models::connection::ConnectionConfig, ) -> DatabaseScope { let Some(rule) = policy.connection_policies.iter().find(|rule| rule.connection_id == connection.id) else { return DatabaseScope::All; }; match rule.database_scope { McpDatabaseScope::All => DatabaseScope::All, McpDatabaseScope::Selected => DatabaseScope::Selected(rule.allowed_databases.clone()), McpDatabaseScope::None => DatabaseScope::None, } } fn resolved_connection( policy: McpGlobalPolicy, connection: dbx_core::models::connection::ConnectionConfig, group_path: Option<&dbx_core::mcp_policy::McpConnectionGroupPath>, ) -> ResolvedConnection { let database_scope = database_scope_for_connection(&policy, &connection); let group_ids = group_path.map(|path| path.ids.clone()).unwrap_or_default(); ResolvedConnection { connection, policy, database_scope, group_ids } } #[allow(clippy::result_large_err)] fn ensure_database_in_scope(scope: &DatabaseScope, database: &str) -> Result<(), CallToolResult> { match scope { DatabaseScope::All => Ok(()), DatabaseScope::Selected(allowed) if allowed.iter().any(|allowed| allowed == database) => Ok(()), DatabaseScope::Selected(_) | DatabaseScope::None => Err(tool_error( "DATABASE_OUT_OF_SCOPE", format!("Database \"{database}\" is not allowed by DBX MCP settings for this connection."), )), } } /// A selected database scope applies to SQL object references as well as the /// request-level `database` field. Without this guard an MCP client could pick /// an allowed default and access another MySQL/SQL Server database through a /// qualified name such as `other_db.users`. #[allow(clippy::result_large_err)] fn ensure_sql_database_scope( scope: &DatabaseScope, connection: &dbx_core::models::connection::ConnectionConfig, active_database: &str, sql: &str, ) -> Result<(), CallToolResult> { let DatabaseScope::Selected(allowed_databases) = scope else { return Ok(()); }; if sql_references_disallowed_database(sql, &connection.db_type, active_database, allowed_databases) { return Err(tool_error( "DATABASE_OUT_OF_SCOPE", "SQL references a database that is not allowed by DBX MCP settings for this connection.", )); } Ok(()) } fn is_database_discovery_sql(sql: &str) -> bool { let normalized = sql.trim().trim_end_matches(';').trim().to_ascii_lowercase(); matches!(normalized.as_str(), "show databases" | "show schemas") } fn format_database_names(databases: &[String]) -> String { if databases.is_empty() { "No databases are available through this MCP connection.".to_string() } else { databases.iter().map(|database| format!("- {database}")).collect::>().join("\n") } } /// Execution modes are scoped defaults: a database rule overrides a configured /// connection default, which in turn overrides the global default. Connection /// read-only protection and production-database protections are enforced /// separately and cannot be bypassed by these defaults. #[cfg(test)] fn effective_policy_for_database( policy: &McpGlobalPolicy, connection: &dbx_core::models::connection::ConnectionConfig, database: &str, ) -> McpGlobalPolicy { effective_policy_for_database_with_groups(policy, &[], connection, database) } fn effective_policy_for_database_with_groups( policy: &McpGlobalPolicy, group_ids: &[String], connection: &dbx_core::models::connection::ConnectionConfig, database: &str, ) -> McpGlobalPolicy { let mut policy = policy.clone(); (policy.read_only, policy.allow_dangerous_sql) = dbx_core::mcp_policy::effective_database_execution_policy_with_groups( &policy, group_ids, &connection.id, database, ); // This returned value is carried through the request only; retaining a // complete policy document here could accidentally be reused for another // connection by a future caller. policy.connection_policies.clear(); policy } /// A database-level execution rule must not be bypassed by qualifying another /// allowed database in SQL. Until individual references are evaluated with /// their own policies, cross-database SQL is intentionally rejected whenever /// this connection has database-specific execution rules. #[allow(clippy::result_large_err)] fn ensure_sql_database_execution_scope( policy: &McpGlobalPolicy, connection: &dbx_core::models::connection::ConnectionConfig, active_database: &str, sql: &str, ) -> Result<(), CallToolResult> { dbx_core::mcp_policy::ensure_sql_database_execution_scope(policy, connection, active_database, sql) .map_err(|error| tool_error("DATABASE_EXECUTION_POLICY_OUT_OF_SCOPE", error)) } fn mcp_permissions( connection: &dbx_core::models::connection::ConnectionConfig, policy: &McpGlobalPolicy, ) -> dbx_core::agent_tools::AgentSqlPermissions { dbx_core::agent_tools::AgentSqlPermissions { allow_writes: !policy.read_only && !connection.read_only, allow_dangerous: !policy.read_only && !connection.read_only && policy.allow_dangerous_sql, confirmed_write_sql: mcp_confirmed_write_sql_from_env(), } } fn transaction_connection_supported(connection: &dbx_core::models::connection::ConnectionConfig) -> bool { crate::backend::native_mysql_transaction_connection(connection) } /// Read the DBX_MCP_CONFIRMED_WRITE_SQL env var (set by the CLI agent when the /// user confirmed a specific write SQL). Returns None when the var is unset or /// empty, so desktop-embedded MCP contexts (which don't set this var) continue /// to work without a confirmed-SQL binding. fn mcp_confirmed_write_sql_from_env() -> Option { normalize_confirmed_write_sql(std::env::var("DBX_MCP_CONFIRMED_WRITE_SQL").ok()) } fn normalize_confirmed_write_sql(value: Option) -> Option { let trimmed = value?.trim().to_string(); (!trimmed.is_empty()).then_some(trimmed) } /// When the user confirmed a specific write SQL for this run, a batch that /// contains a write/DDL statement must be exactly that confirmed SQL. Returns /// the block message when the script differs from the confirmation and is not /// read-only; `None` means the script is allowed (read-only, exact match, or /// no binding). Unparseable SQL is treated as a write and fails closed. fn confirmed_batch_sql_block_reason( sql: &str, db_type: dbx_core::models::connection::DatabaseType, confirmed: Option<&str>, ) -> Option { let confirmed = confirmed?; if normalize_sql_for_confirmation(sql) == normalize_sql_for_confirmation(confirmed) { return None; } let is_write = match classify_sql_risk_for_database(sql, db_type) { Ok(risk) => risk != dbx_core::sql_risk::SqlRisk::ReadOnly, Err(_) => true, }; if !is_write { return None; } Some(format!( "Blocked: the executed SQL does not match the user-confirmed SQL.\n\ Confirmed: {}\n\ Attempted: {}", confirmed, sql )) } fn validate_transaction_sql_shape(sql: &str) -> Result<(), String> { use sqlparser::{ ast::{Expr, Statement, Visit, Visitor}, dialect::MySqlDialect, parser::Parser, tokenizer::{Token, Tokenizer}, }; use std::ops::ControlFlow; const SAFE_FUNCTIONS: &[&str] = &[ "abs", "ascii", "avg", "bin", "ceiling", "char_length", "character_length", "coalesce", "concat", "concat_ws", "connection_id", "conv", "count", "crc32", "curdate", "curtime", "current_date", "current_time", "current_timestamp", "date_format", "datediff", "dayname", "dayofmonth", "dayofweek", "dayofyear", "exp", "floor", "format", "greatest", "hex", "hour", "if", "ifnull", "inet_aton", "inet_ntoa", "instr", "isnull", "json_extract", "json_length", "json_unquote", "json_valid", "last_day", "lcase", "least", "left", "length", "ln", "locate", "log", "log10", "log2", "lower", "lpad", "ltrim", "max", "md5", "microsecond", "min", "minute", "mod", "month", "monthname", "now", "nullif", "oct", "ord", "position", "pow", "power", "quarter", "rand", "repeat", "replace", "reverse", "right", "round", "rpad", "rtrim", "second", "sha", "sha1", "sha2", "sign", "sleep", "space", "sqrt", "str_to_date", "substring", "substr", "time_format", "timediff", "timestampadd", "timestampdiff", "truncate", "unhex", "unix_timestamp", "upper", "ucase", "utc_date", "utc_time", "utc_timestamp", "uuid", "version", "week", "weekday", "year", ]; struct TransactionFunctionVisitor { rejected: bool, } impl Visitor for TransactionFunctionVisitor { type Break = (); fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<()> { if let Expr::Function(function) = expr { let parts = &function.name.0; let allowed = parts.len() == 1 && parts.last().and_then(|part| part.as_ident()).is_some_and(|ident| { SAFE_FUNCTIONS.iter().any(|candidate| ident.value.eq_ignore_ascii_case(candidate)) }); if !allowed { self.rejected = true; return ControlFlow::Break(()); } } ControlFlow::Continue(()) } } if sql.contains("/*!") || sql.contains("/*+") { return Err("Executable comments and optimizer hints are not allowed in transaction sessions.".to_string()); } let dialect = MySqlDialect {}; let statements = Parser::parse_sql(&dialect, sql).map_err(|error| format!("SQL cannot be proven safe: {error}"))?; let [statement] = statements.as_slice() else { return Err("Transaction session calls must contain exactly one SQL statement.".to_string()); }; let tokens = Tokenizer::new(&dialect, sql).tokenize().map_err(|error| format!("SQL cannot be tokenized safely: {error}"))?; let has_word = |expected: &str| { tokens.iter().any(|token| matches!(token, Token::Word(word) if word.value.eq_ignore_ascii_case(expected))) }; if tokens.iter().any(|token| matches!(token, Token::Assignment)) { return Err("User-variable assignment is not allowed in transaction sessions.".to_string()); } match statement { Statement::Query(_) => { for forbidden in [ "INTO", "OUTFILE", "DUMPFILE", "LOAD_FILE", "GET_LOCK", "RELEASE_LOCK", "IS_FREE_LOCK", "IS_USED_LOCK", "MASTER_POS_WAIT", "SQL_CALC_FOUND_ROWS", ] { if has_word(forbidden) { return Err(format!("{forbidden} is not allowed in transaction session queries.")); } } Ok(()) } Statement::Insert(insert) if insert.output.is_none() && insert.returning.is_none() && insert.settings.is_none() && insert.format_clause.is_none() && insert.multi_table_insert_type.is_none() && insert.multi_table_into_clauses.is_empty() && insert.multi_table_when_clauses.is_empty() && insert.multi_table_else_clause.is_none() => { Ok(()) } Statement::Update(update) if update.output.is_none() && update.returning.is_none() => Ok(()), Statement::Delete(delete) if delete.output.is_none() && delete.returning.is_none() => Ok(()), _ => Err( "Transaction sessions allow only SELECT (including locking reads), INSERT, UPDATE, DELETE, and REPLACE." .to_string(), ), }?; let mut visitor = TransactionFunctionVisitor { rejected: false }; let _ = statement.visit(&mut visitor); if visitor.rejected { return Err( "Only unqualified, known-safe MySQL built-in functions are allowed in transaction sessions.".to_string() ); } Ok(()) } // CallToolResult is the transport-native error payload; boxing it would complicate every MCP call site. #[allow(clippy::result_large_err)] fn validate_sql_policy( connection: &dbx_core::models::connection::ConnectionConfig, policy: &McpGlobalPolicy, database: &str, sql: &str, allow_database_switch: bool, ) -> Result { if !allow_database_switch && mcp_sql_has_forbidden_database_switch(sql, connection.db_type) { return Err(tool_error( "SQL_BLOCKED", "MCP does not allow USE or persistent database switching outside a session. Open one with dbx_open_session to run stateful queries.", )); } let risk = classify_sql_risk_for_database(sql, connection.db_type).map_err(|error| tool_error("SQL_BLOCKED", error))?; if risk == SqlRisk::Transaction { return Err(tool_error("SQL_BLOCKED", "Transaction statements are not supported by MCP.")); } // The keyword scan alone misses write-capable SQL that the risk classifier // does recognize (locking reads, side-effect functions, writable CTEs), and // those statements would otherwise reach the database whenever high-risk SQL // is permitted. Fail closed on either signal so read-only stays read-only. let is_write = risk != SqlRisk::ReadOnly || is_write_sql_for_database(sql, connection.db_type); if policy.read_only && is_write { return Err(tool_error( "MCP_READ_ONLY", format!("MCP execution permission for database \"{database}\" is read-only. SQL write blocked."), )); } if connection.read_only && is_write { return Err(tool_error( "CONNECTION_READ_ONLY", format!("Connection \"{}\" has read-only protection enabled. SQL write blocked.", connection.name), )); } let high_risk = risk == SqlRisk::Ddl || is_dangerous_sql_for_database(sql, connection.db_type); if high_risk && !policy.allow_dangerous_sql { return Err(tool_error("SQL_BLOCKED", "High-risk SQL is disabled in DBX MCP settings.")); } if is_write && targets_production_database(connection, database, sql) { return Err(tool_error("PRODUCTION_WRITE_BLOCKED", "MCP cannot execute writes against a production database.")); } Ok(mcp_permissions(connection, policy)) } // CallToolResult is the transport-native error payload; boxing it would complicate every MCP call site. #[cfg(test)] #[allow(clippy::result_large_err)] #[cfg(test)] fn validate_mongo_command( connection: &dbx_core::models::connection::ConnectionConfig, policy: &McpGlobalPolicy, database_scope: &DatabaseScope, database: &str, source: &str, ) -> Result { validate_mongo_command_with_groups(connection, policy, &[], database_scope, database, source) } #[allow(clippy::result_large_err)] fn validate_mongo_command_with_groups( connection: &dbx_core::models::connection::ConnectionConfig, policy: &McpGlobalPolicy, group_ids: &[String], database_scope: &DatabaseScope, database: &str, source: &str, ) -> Result { let command = mongo::parse(source).map_err(|error| { tool_error( "QUERY_ERROR", format!( "{error} Use MongoDB shell-style commands such as db.collection.find({{}}), db.collection.aggregate([]), or db.collection.countDocuments({{}})." ), ) })?; let (target_database, command_to_validate) = match &command { MongoCommand::InDatabase { database, command } => (database.as_str(), command.as_ref()), command => (database, command), }; ensure_database_in_scope(database_scope, target_database)?; if matches!(command_to_validate, MongoCommand::RunCommand { .. }) { return Err(tool_error( "SQL_BLOCKED", "MongoDB runCommand is not available through MCP; review and execute it manually in DBX.", )); } if let MongoCommand::Aggregate { pipeline, .. } = command_to_validate { dbx_core::mcp_policy::ensure_mongo_database_execution_scope(policy, &connection.id, target_database, pipeline) .map_err(|error| { let (code, message) = error .strip_prefix("QUERY_ERROR: ") .map_or(("DATABASE_EXECUTION_POLICY_OUT_OF_SCOPE", error.as_str()), |message| { ("QUERY_ERROR", message) }); tool_error(code, message) })?; for output_database in mongo_aggregate_target_databases(pipeline, target_database)? { ensure_database_in_scope(database_scope, &output_database)?; } } let effective_policy = effective_policy_for_database_with_groups(policy, group_ids, connection, target_database); let permissions = mcp_permissions(connection, &effective_policy); let production_database = match command_to_validate { MongoCommand::Aggregate { pipeline, .. } => { mongo_pipeline_targets_production_database(connection, target_database, pipeline) } _ => is_production_database(connection, target_database), }; if let Err(error) = mongo::validate_safety( command_to_validate, permissions.allow_writes, permissions.allow_dangerous, production_database, ) { return Err(match error { MongoSafetyError::WritesDisabled => tool_error( if effective_policy.read_only { "MCP_READ_ONLY" } else { "CONNECTION_READ_ONLY" }, "MCP MongoDB execution is read-only for this database in DBX MCP settings.", ), MongoSafetyError::EmptyFilter => tool_error( "SQL_BLOCKED", "MongoDB update/delete commands must include a non-empty filter unless high-risk operations are enabled in DBX MCP settings.", ), MongoSafetyError::Dangerous => { tool_error("SQL_BLOCKED", "Dangerous MongoDB command is disabled in DBX MCP settings.") } MongoSafetyError::ProductionWrite => { tool_error("PRODUCTION_WRITE_BLOCKED", "MCP cannot execute writes against a production database.") } }); } Ok(command) } #[allow(clippy::result_large_err)] fn mongo_aggregate_target_databases(pipeline: &str, active_database: &str) -> Result, CallToolResult> { dbx_core::mcp_policy::mongo_pipeline_output_databases(pipeline, active_database) .map_err(|error| tool_error("QUERY_ERROR", error.strip_prefix("QUERY_ERROR: ").unwrap_or(&error))) } fn non_empty_env(name: &str) -> Option { std::env::var(name).ok().map(|value| value.trim().to_string()).filter(|value| !value.is_empty()) } #[cfg(feature = "mq-admin")] fn format_peek_messages_result(result: dbx_core::mq::PeekMessagesResult, count: usize) -> String { let mut messages = Vec::new(); let mut bytes = 0; let mut truncated = result.messages.len() > count; for message in result.messages.into_iter().take(count) { let value = serde_json::to_value(message).expect("message is JSON serializable"); bytes += value.to_string().len(); if bytes > 256 * 1024 { truncated = true; break; } messages.push(value); } json!({ "messages": messages, "incomplete": result.incomplete, "outputTruncated": truncated }).to_string() } #[cfg(feature = "mq-admin")] fn format_send_message_result( system_kind: &dbx_core::mq::MqSystemKind, result: &dbx_core::mq::SendMessageResponse, ) -> String { let mut output = format!( "Message sent to {}.\ntopic: {}\npartition: {}\noffset: {}", system_kind.as_str(), result.topic, result.partition, result.offset ); if let Some(timestamp) = result.timestamp.as_deref() { output.push_str(&format!("\ntimestamp: {timestamp}")); } output } fn scoped_connection_ids(value: Option<&str>) -> Vec { let mut ids = Vec::new(); for id in value.unwrap_or_default().split(',').map(str::trim).filter(|id| !id.is_empty()) { if !ids.iter().any(|existing| existing == id) { ids.push(id.to_string()); } } ids } fn ambiguous_connections(name: &str, connections: &[dbx_core::models::connection::ConnectionConfig]) -> String { let lines = connections .iter() .map(|connection| { format!("- {}: {:?} @ {}:{}", connection.id, connection.db_type, connection.host, connection.port) }) .collect::>() .join("\n"); format!("Multiple connections found with name \"{name}\". Please specify connection_id:\n{lines}") } fn format_connections(connections: &[ConnectionSummary]) -> String { let mut output = String::from( "| ID | Name | Group Path | Type | Host | Port | Database |\n| --- | --- | --- | --- | --- | --- | --- |", ); for connection in connections { output.push_str(&format!( "\n| {} | {} | {} | {} | {} | {} | {} |", escape_cell(&connection.id), escape_cell(&connection.name), escape_cell(&connection.group_path.join(" / ")), escape_cell(&connection.db_type), escape_cell(&connection.host), connection.port, escape_cell(&connection.database), )); } output } /// Normalize a user-supplied routine type ("procedure", "PROCEDURE", ...) to /// the canonical SCREAMING form used by the object listing protocol. fn normalize_routine_type(value: &str) -> Result { match value.trim().to_ascii_uppercase().as_str() { "PROCEDURE" => Ok("PROCEDURE".to_string()), "FUNCTION" => Ok("FUNCTION".to_string()), other => Err(format!("Unsupported routine type \"{other}\"; use PROCEDURE or FUNCTION.")), } } fn format_routines(routines: &[dbx_core::db::ObjectInfo]) -> String { let rows = routines .iter() .map(|routine| { vec![ routine.name.clone(), routine.object_type.clone(), routine.schema.clone().unwrap_or_default(), routine.signature.clone().unwrap_or_default(), routine.comment.clone().unwrap_or_default(), ] }) .collect::>(); markdown_table(&["Routine", "Type", "Schema", "Signature", "Comment"], &rows) } fn format_columns(columns: &[dbx_core::db::ColumnInfo]) -> String { let rows = columns .iter() .map(|column| { vec![ if column.is_primary_key { format!("{} (PK)", column.name) } else { column.name.clone() }, column.data_type.clone(), if column.is_nullable { "YES".to_string() } else { "NO".to_string() }, column.column_default.clone().unwrap_or_default(), column.comment.clone().unwrap_or_default(), ] }) .collect::>(); markdown_table(&["Column", "Type", "Nullable", "Default", "Comment"], &rows) } fn markdown_table(headers: &[&str], rows: &[Vec]) -> String { let mut output = format!("| {} |\n| {} |", headers.join(" | "), vec!["---"; headers.len()].join(" | ")); for row in rows { output .push_str(&format!("\n| {} |", row.iter().map(|value| escape_cell(value)).collect::>().join(" | "))); } output } fn escape_cell(value: &str) -> String { value.replace('|', "\\|").replace('\n', " ") } fn redis_database(connection: &dbx_core::models::connection::ConnectionConfig) -> Option { connection.database.as_deref().and_then(parse_redis_database) } fn parse_redis_database(value: &str) -> Option { value.trim().parse().ok() } fn format_redis_result(result: &RedisCommandResult) -> String { let value = result.value.as_str().map(ToOwned::to_owned).unwrap_or_else(|| { serde_json::to_string_pretty(&result.value).unwrap_or_else(|_| result.value.to_string()) }); let safety = serde_json::to_value(&result.safety) .ok() .and_then(|value| value.as_str().map(ToOwned::to_owned)) .unwrap_or_else(|| format!("{:?}", result.safety).to_ascii_lowercase()); format!("Command: {}\nSafety: {}\n\n{}", result.command, safety, value) } fn format_schema_context( connection: &str, database: &str, schema: &str, tables: &[(dbx_core::db::TableInfo, Vec)], truncated: bool, ) -> String { let mut output = format!("Connection: {connection}"); if !database.is_empty() { output.push_str(&format!("\nDatabase: {database}")); } if !schema.is_empty() { output.push_str(&format!("\nSchema: {schema}")); } for (table, columns) in tables { output.push_str(&format!("\n\n## {}\nType: {}", table.name, table.table_type)); for column in columns { output.push_str(&format!( "\n- {} {} {}{}{}", column.name, column.data_type, if column.is_nullable { "NULL" } else { "NOT NULL" }, if column.is_primary_key { " PK" } else { "" }, column.comment.as_ref().map(|comment| format!(" -- {comment}")).unwrap_or_default(), )); } } if truncated { output.push_str("\n\nNote: table list was truncated; request specific table names for more context."); } output } #[cfg(test)] mod tests { use super::*; use async_trait::async_trait; use dbx_core::models::connection::ConnectionConfig; use rmcp::{ model::{CallToolRequest, CallToolRequestParams, ClientRequest}, service::PeerRequestOptions, ServiceExt, }; use std::{collections::HashSet, sync::atomic::AtomicUsize, time::Duration}; struct FakeTransactionIo { in_transaction: bool, executed_sql: Arc>>, block_next_sql: Arc>>>, disconnects: Arc, } #[async_trait] impl crate::transaction::TransactionIo for FakeTransactionIo { async fn execute( &mut self, sql: &str, _max_rows: Option, ) -> Result { self.executed_sql.lock().unwrap().push(sql.to_string()); let block = self.block_next_sql.lock().unwrap().take(); if let Some(block) = block { block.notified().await; } match sql { "START TRANSACTION" => self.in_transaction = true, "COMMIT" | "ROLLBACK" => self.in_transaction = false, _ => {} } Ok(crate::transaction::TransactionIoSuccess { result: dbx_core::db::QueryResult { columns: Vec::new(), column_types: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: Vec::new(), affected_rows: 0, execution_time_ms: 0, server_execute_time_us: None, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }, in_transaction: self.in_transaction, }) } async fn ping_in_transaction(&mut self) -> Result { Ok(self.in_transaction) } async fn disconnect(&mut self) { self.disconnects.fetch_add(1, std::sync::atomic::Ordering::SeqCst); } } #[test] fn transaction_sql_allowlist_accepts_only_rollback_safe_mysql_dml_and_queries() { for sql in [ "SELECT id FROM accounts WHERE id = 1", "SELECT id FROM accounts WHERE id = 1 FOR UPDATE", "SELECT CONNECTION_ID(), VERSION(), SLEEP(0)", "SELECT COUNT(*), LOWER(name), COALESCE(name, '') FROM accounts", "INSERT INTO accounts(id) VALUES (1)", "UPDATE accounts SET name = 'safe' WHERE id = 1", "DELETE FROM accounts WHERE id = 1", "REPLACE INTO accounts(id) VALUES (1)", ] { assert!(validate_transaction_sql_shape(sql).is_ok(), "expected allowed: {sql}"); } for sql in [ "CREATE TABLE t(id INT)", "CREATE TEMPORARY TABLE t(id INT)", "ALTER TABLE t ADD COLUMN name TEXT", "DROP TABLE t", "CALL mutate_data()", "LOAD DATA INFILE '/tmp/input' INTO TABLE t", "LOCK TABLES t WRITE", "UNLOCK TABLES", "SET autocommit = 0", "USE another_database", "XA START 'x'", "BEGIN", "COMMIT", "ROLLBACK", "SELECT 1 INTO OUTFILE '/tmp/output'", "SELECT @value := 1", "SELECT evil_udf(id) FROM accounts", "SELECT app.evil_udf(id) FROM accounts", "SELECT (SELECT evil_udf())", "INSERT INTO accounts(id) VALUES (evil_udf(1))", "UPDATE accounts SET name = evil_udf(name) WHERE id = 1", "DELETE FROM accounts WHERE evil_udf(id) = 1", "/*!40101 SET autocommit = 0 */", "SELECT 1; DELETE FROM accounts", ] { assert!(validate_transaction_sql_shape(sql).is_err(), "expected rejected: {sql}"); } } struct FakeBackend { connections: Vec, policy: McpGlobalPolicy, recorded_arguments: std::sync::Mutex>, history: std::sync::Mutex>, closed_sessions: std::sync::Mutex>, pinned_sessions: std::sync::Mutex>, close_failures_remaining: std::sync::Mutex, transaction_owners_opened: AtomicUsize, transaction_owner_sql: Arc>>, transaction_block_next_sql: Arc>>>, transaction_owner_disconnects: Arc, transaction_owner_budget: Arc, transaction_open_before: Option<(Arc, Arc)>, transaction_open_after_owner: Option<(Arc, Arc)>, transaction_open_error: Option, policy_override: std::sync::Mutex>, } impl Default for FakeBackend { fn default() -> Self { Self { connections: Vec::new(), policy: McpGlobalPolicy::default(), recorded_arguments: std::sync::Mutex::new(Vec::new()), history: std::sync::Mutex::new(Vec::new()), closed_sessions: std::sync::Mutex::new(Vec::new()), pinned_sessions: std::sync::Mutex::new(HashSet::new()), close_failures_remaining: std::sync::Mutex::new(0), transaction_owners_opened: AtomicUsize::new(0), transaction_owner_sql: Arc::new(std::sync::Mutex::new(Vec::new())), transaction_block_next_sql: Arc::new(std::sync::Mutex::new(None)), transaction_owner_disconnects: Arc::new(AtomicUsize::new(0)), transaction_owner_budget: Arc::new(tokio::sync::Semaphore::new(32)), transaction_open_before: None, transaction_open_after_owner: None, transaction_open_error: None, policy_override: std::sync::Mutex::new(None), } } } fn connection(id: &str, name: &str, db_type: &str, database: &str) -> ConnectionConfig { serde_json::from_value(serde_json::json!({ "id": id, "name": name, "db_type": db_type, "host": "", "port": 0, "username": "", "password": "", "database": database, "ssl": false })) .unwrap() } fn resolved_connection_for_test(connection: ConnectionConfig) -> ResolvedConnection { ResolvedConnection { connection, policy: McpGlobalPolicy::default(), database_scope: DatabaseScope::All, group_ids: Vec::new(), } } fn result_text(result: &CallToolResult) -> &str { result.content[0].as_text().expect("text tool result").text.as_str() } fn opened_session_id(result: &CallToolResult) -> String { result_text(result) .lines() .find_map(|line| line.strip_prefix("session_id: ")) .expect("open_session returns a session_id") .to_string() } #[async_trait] impl DbxBackend for FakeBackend { async fn save_history_entry(&self, entry: &dbx_core::history::HistoryEntry) -> Result<(), String> { self.history.lock().unwrap().push(entry.clone()); Ok(()) } async fn load_mcp_global_policy(&self) -> Result { Ok(self.policy_override.lock().unwrap().clone().unwrap_or_else(|| self.policy.clone())) } async fn load_connections(&self) -> Result, String> { Ok(self.connections.clone()) } async fn open_transaction_owner( &self, _connection: &ConnectionConfig, _database: &str, _client_session_id: &str, ) -> Result, String> { if let Some(error) = &self.transaction_open_error { return Err(error.clone()); } if let Some((entered, release)) = &self.transaction_open_before { entered.notify_one(); release.notified().await; } self.transaction_owners_opened.fetch_add(1, std::sync::atomic::Ordering::SeqCst); let permit = self .transaction_owner_budget .clone() .try_acquire_owned() .map_err(|_| "Too many open transaction-enabled MCP sessions (max 32).".to_string())?; let owner = crate::transaction::TransactionOwner::spawn_with_resource_permit( FakeTransactionIo { in_transaction: false, executed_sql: self.transaction_owner_sql.clone(), block_next_sql: self.transaction_block_next_sql.clone(), disconnects: self.transaction_owner_disconnects.clone(), }, crate::transaction::TransactionOwnerConfig::default(), permit, ); if let Some((entered, release)) = &self.transaction_open_after_owner { entered.notify_one(); release.notified().await; } Ok(owner) } async fn execute_agent_tool( &self, _connection: &ConnectionConfig, _database: &str, tool_name: &str, arguments: serde_json::Value, _permissions: dbx_core::agent_tools::AgentSqlPermissions, ) -> dbx_core::agent_events::ToolResult { let is_error = tool_name == "execute_query" && arguments["sql"].as_str().is_some_and(|sql| sql.contains("FAIL_MCP_TEST")); if let Some(client_session_id) = arguments.get("client_session_id").and_then(serde_json::Value::as_str).filter(|id| !id.is_empty()) { self.pinned_sessions.lock().unwrap().insert(client_session_id.to_string()); } self.recorded_arguments.lock().unwrap().push((tool_name.to_string(), arguments)); dbx_core::agent_events::ToolResult { tool_call_id: "test".to_string(), tool_name: tool_name.to_string(), content: if is_error { "Error: query failed" } else { "ok" }.to_string(), is_error, explain_data: None, } } #[cfg(feature = "mq-admin")] async fn peek_messages( &self, connection: &ConnectionConfig, topic: dbx_core::mq::TopicRef, count: u32, options: dbx_core::mq::PeekMessagesOptions, ) -> Result { self.recorded_arguments.lock().unwrap().push(( "peek_messages".into(), json!({"connectionId": connection.id, "topic": topic, "count": count, "options": options}), )); Ok(dbx_core::mq::PeekMessagesResult { messages: vec![dbx_core::mq::PeekedMessage { payload_base64: "aGk=".into(), payload_text: Some("hi".into()), ..Default::default() }], incomplete: true, }) } async fn close_client_session( &self, _connection_id: &str, _database: &str, client_session_id: &str, ) -> Result { let mut failures = self.close_failures_remaining.lock().unwrap(); if *failures > 0 { *failures -= 1; return Err("temporary close failure".to_string()); } drop(failures); self.closed_sessions.lock().unwrap().push(client_session_id.to_string()); self.pinned_sessions.lock().unwrap().remove(client_session_id); Ok(true) } async fn add_connection_for_mcp(&self, config: ConnectionConfig) -> Result { Ok(config) } async fn duplicate_connection_for_mcp( &self, source_id: &str, copy_id: &str, copy_name: &str, ) -> Result { let mut copy = self .connections .iter() .find(|connection| connection.id == source_id) .cloned() .ok_or_else(|| "source not found".to_string())?; copy.id = copy_id.to_string(); copy.name = copy_name.to_string(); Ok(copy) } async fn remove_connection_for_mcp(&self, _connection_id: &str) -> Result { Ok(true) } async fn execute_batch( &self, _connection: &ConnectionConfig, _database: &str, schema: Option<&str>, sql: &str, options: dbx_core::query::QueryExecutionOptions, ) -> Result, String> { let client_session_id = options.client_session_id.as_deref().filter(|id| !id.trim().is_empty()).map(str::to_owned); if let Some(client_session_id) = &client_session_id { self.pinned_sessions.lock().unwrap().insert(client_session_id.clone()); } self.recorded_arguments.lock().unwrap().push(( "execute_batch".to_string(), json!({ "sql": sql, "schema": schema, "continue_on_error": options.continue_on_error, "use_transaction": options.use_transaction, "client_session_id": client_session_id, }), )); let execution_error = sql.contains("FAIL_MCP_TEST") || sql.contains("FAIL_MCP_NO_MESSAGE_TEST"); Ok(vec![crate::backend::BatchStatementResult { result: dbx_core::db::QueryResult { columns: vec![], column_types: vec![], column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: 0, execution_time_ms: 1, server_execute_time_us: None, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: vec![], }, execution_error, statement_index: Some(0), error_message: sql.contains("FAIL_MCP_TEST").then(|| "statement failed".to_string()), merged: false, transaction_state: None, transaction_outcome: None, }]) } } #[cfg(feature = "mq-admin")] fn kafka_connection() -> ConnectionConfig { let mut conn = connection("kafka", "Kafka", "mq", ""); conn.read_only = true; conn.is_production = true; conn.external_config = Some( json!({"systemKind":"kafka", "adminUrl":"", "auth":{"kind":"none"}, "extra":{"bootstrapServers":"localhost:9092"}}), ); conn } #[cfg(feature = "mq-admin")] #[tokio::test] async fn peek_messages_reads_with_read_only_policy_and_forwards_positions() { let backend = Arc::new(FakeBackend { connections: vec![kafka_connection()], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); for (extra, expected) in [ (json!({}), json!({"startPosition":"latest"})), (json!({"start_position":"earliest", "partition":0}), json!({"startPosition":"earliest", "partition":0})), ( json!({"start_position":"offset", "partition":2, "offset":17}), json!({"startPosition":"offset", "partition":2, "offset":17}), ), (json!({"start_position":"offset", "offset":0}), json!({"startPosition":"offset", "offset":0})), ] { let mut request = json!({"connection_id":"kafka", "topic":" events "}); request.as_object_mut().unwrap().extend(extra.as_object().unwrap().clone()); let result = server.peek_messages(Parameters(serde_json::from_value(request).unwrap())).await; assert_ne!(result.is_error, Some(true), "{}", result_text(&result)); let response: serde_json::Value = serde_json::from_str(result_text(&result)).unwrap(); assert_eq!(response["messages"][0]["payloadBase64"], "aGk="); assert_eq!(response["incomplete"], true); assert_eq!(response["outputTruncated"], false); let calls = backend.recorded_arguments.lock().unwrap(); let args = &calls.last().unwrap().1; assert_eq!(args["options"], expected); assert_eq!(args["count"], 20); assert_eq!(args["topic"]["topic"], "events"); } } #[cfg(feature = "mq-admin")] #[tokio::test] async fn peek_messages_rejects_invalid_options_before_backend_call() { let backend = Arc::new(FakeBackend { connections: vec![kafka_connection()], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); for extra in [ json!({"count":0}), json!({"count":101}), json!({"partition":-1}), json!({"start_position":"offset"}), json!({"offset":1}), json!({"start_position":"earliest", "offset":0}), json!({"start_position":"offset", "offset":-1}), json!({"topic":" "}), ] { let mut request = json!({"connection_id":"kafka", "topic":"events"}); request.as_object_mut().unwrap().extend(extra.as_object().unwrap().clone()); let result = server.peek_messages(Parameters(serde_json::from_value(request).unwrap())).await; assert_eq!(result.is_error, Some(true), "{}", result_text(&result)); } assert!(backend.recorded_arguments.lock().unwrap().is_empty()); } #[cfg(feature = "mq-admin")] #[tokio::test] async fn peek_messages_enforces_tool_connection_scope_and_kafka_only() { for (policy, conn, expected) in [ ( McpGlobalPolicy { allowed_tool_names: Some(vec![]), ..Default::default() }, kafka_connection(), "TOOL_OUT_OF_SCOPE", ), ( McpGlobalPolicy { allowed_connection_ids: Some(vec![]), ..Default::default() }, kafka_connection(), "CONNECTION_OUT_OF_SCOPE", ), (McpGlobalPolicy::default(), connection("kafka", "redis", "redis", "0"), "MQ_UNSUPPORTED"), ] { let backend = Arc::new(FakeBackend { policy, connections: vec![conn], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let result = server .peek_messages(Parameters( serde_json::from_value(json!({"connection_id":"kafka", "topic":"events"})).unwrap(), )) .await; assert_eq!(result.is_error, Some(true)); assert!(result_text(&result).contains(expected), "{}", result_text(&result)); assert!(backend.recorded_arguments.lock().unwrap().is_empty()); } } #[cfg(feature = "mq-admin")] #[tokio::test] async fn peek_messages_respects_session_scope_and_rejects_other_queues() { let mut rabbit = kafka_connection(); rabbit.external_config.as_mut().unwrap()["systemKind"] = json!("rabbitmq"); for (conn, scope, expected) in [ ( kafka_connection(), McpScope { connection_ids: vec!["other".into()], ..Default::default() }, "CONNECTION_OUT_OF_SCOPE", ), (rabbit, McpScope::default(), "MQ_UNSUPPORTED"), ] { let backend = Arc::new(FakeBackend { connections: vec![conn], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), scope, false); let result = server .peek_messages(Parameters( serde_json::from_value(json!({"connection_id":"kafka", "topic":"events"})).unwrap(), )) .await; assert_eq!(result.is_error, Some(true)); assert!(result_text(&result).contains(expected), "{}", result_text(&result)); assert!(backend.recorded_arguments.lock().unwrap().is_empty()); } } #[cfg(feature = "mq-admin")] #[test] fn peek_messages_output_preserves_metadata_and_bounds_whole_messages() { use dbx_core::mq::{PeekMessagesResult, PeekedMessage}; let message = PeekedMessage { payload_base64: "/w==".into(), message_id: Some("2:17".into()), headers: [("kind".into(), "event".into())].into(), ..Default::default() }; let output = format_peek_messages_result(PeekMessagesResult::complete(vec![message.clone(), message]), 1); let result: serde_json::Value = serde_json::from_str(&output).unwrap(); assert_eq!(result["messages"].as_array().unwrap().len(), 1); assert_eq!(result["messages"][0]["messageId"], "2:17"); assert_eq!(result["messages"][0]["headers"]["kind"], "event"); assert_eq!(result["outputTruncated"], true); assert_eq!(result["incomplete"], false); let output = format_peek_messages_result( PeekMessagesResult::complete(vec![PeekedMessage { payload_base64: "a".repeat(300_000), ..Default::default() }]), 20, ); let result: serde_json::Value = serde_json::from_str(&output).unwrap(); assert!(result["messages"].as_array().unwrap().is_empty()); assert_eq!(result["outputTruncated"], true); assert!(output.len() < 256 * 1024); } #[test] fn connection_table_escapes_markdown_cells() { let output = format_connections(&[ConnectionSummary { id: "id|1".to_string(), name: "local\npg".to_string(), db_type: "postgres".to_string(), host: "127.0.0.1".to_string(), port: 5432, database: "app".to_string(), group_path: vec!["Project|A".to_string(), "Staging\nWest".to_string()], }]); assert!(output.contains("id\\|1")); assert!(output.contains("local pg")); assert!(output.contains("Project\\|A / Staging West")); } #[test] fn server_registers_list_connections_tool() { let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); let tools = server.tool_router.list_all(); let names = tools.iter().map(|tool| tool.name.as_ref()).collect::>(); #[cfg(feature = "mq-admin")] assert_eq!(tools.len(), 22); #[cfg(not(feature = "mq-admin"))] assert_eq!(tools.len(), 20); #[cfg(feature = "mq-admin")] assert!(names.contains(&"dbx_peek_messages")); #[cfg(not(feature = "mq-admin"))] assert!(!names.contains(&"dbx_peek_messages")); assert!(names.contains(&"dbx_list_connections")); assert!(names.contains(&"dbx_list_databases")); assert!(names.contains(&"dbx_list_tables")); assert!(names.contains(&"dbx_describe_table")); assert!(names.contains(&"dbx_list_routines")); assert!(names.contains(&"dbx_get_routine_source")); assert!(names.contains(&"dbx_execute_query")); assert!(names.contains(&"dbx_execute_batch")); assert!(names.contains(&"dbx_begin_transaction")); assert!(names.contains(&"dbx_commit_transaction")); assert!(names.contains(&"dbx_rollback_transaction")); assert!(names.contains(&"dbx_add_connection")); assert!(names.contains(&"dbx_duplicate_connection")); assert!(names.contains(&"dbx_remove_connection")); assert!(names.contains(&"dbx_execute_redis_command")); assert!(names.contains(&"dbx_get_schema_context")); assert!(names.contains(&"dbx_open_table")); assert!(names.contains(&"dbx_execute_and_show")); assert!(names.contains(&"dbx_open_session")); assert!(names.contains(&"dbx_close_session")); #[cfg(feature = "mq-admin")] assert!(names.contains(&"dbx_send_message")); } #[tokio::test] async fn tools_list_hides_tools_the_policy_disallows() { let policy = McpGlobalPolicy { allowed_tool_names: Some(vec!["dbx_list_connections".to_string(), "dbx_execute_query".to_string()]), ..Default::default() }; let server = DbxMcpServer::with_runtime_options( Arc::new(FakeBackend { policy, ..Default::default() }), McpScope::default(), false, ); let tools = server.policy_filtered_tools().await; let names = tools.iter().map(|tool| tool.name.as_ref()).collect::>(); assert_eq!(names, vec!["dbx_execute_query", "dbx_list_connections"]); } #[tokio::test] async fn tools_list_keeps_the_full_view_without_an_allowlist() { let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); let tools = server.policy_filtered_tools().await; assert_eq!(tools.len(), server.tool_router.list_all().len()); } #[test] fn schema_context_tables_schema_is_gemini_compatible() { let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); let tool = server .tool_router .list_all() .into_iter() .find(|tool| tool.name == "dbx_get_schema_context") .expect("schema context tool should be registered"); let tables = tool .input_schema .get("properties") .and_then(serde_json::Value::as_object) .and_then(|properties| properties.get("tables")) .expect("tables property should be published"); assert_eq!(tables.get("type"), Some(&serde_json::json!("array"))); assert_eq!(tables.pointer("/items/type"), Some(&serde_json::json!("string"))); assert!(!tool .input_schema .get("required") .and_then(serde_json::Value::as_array) .is_some_and(|required| required.iter().any(|field| field == "tables"))); } #[test] fn connection_selector_schema_uses_optional_strings() { let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); let tools = server.tool_router.list_all(); for tool_name in ["dbx_execute_query", "dbx_list_tables", "dbx_open_session"] { let tool = tools.iter().find(|tool| tool.name == tool_name).expect("selector tool should be registered"); let properties = tool .input_schema .get("properties") .and_then(serde_json::Value::as_object) .expect("selector tool should publish object properties"); let required = tool .input_schema .get("required") .and_then(serde_json::Value::as_array) .into_iter() .flatten() .collect::>(); for field in ["connection_id", "connection_name"] { let selector = properties.get(field).expect("selector field should be published"); assert_eq!(selector.get("type"), Some(&serde_json::json!("string")), "{tool_name}.{field}"); assert!(!required.iter().any(|required| *required == field), "{tool_name}.{field} must stay optional"); } } } #[test] fn optional_fields_never_publish_nullable_union_types() { // Some MCP clients (e.g. OpenCode, see #6344) cannot resolve a JSON Schema // `"type": ["string", "null"]` union and fall back to wrapping the argument in a // nested error object instead of passing the value through. b521d0377 fixed this for // `ConnectionSelector`'s connection_id/connection_name but left every other optional // field on these request structs emitting the same union shape. Every optional field // must instead publish a single concrete `type`, relying on omission from `required` // (not a `"null"` union member) to signal optionality. let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); let tools = server.tool_router.list_all(); #[allow(unused_mut)] let mut checks: Vec<(&str, &[&str])> = vec![ ("dbx_list_tables", &["database", "schema"]), ("dbx_describe_table", &["database", "schema"]), ("dbx_list_routines", &["database", "schema", "routine_type"]), ("dbx_get_routine_source", &["database", "schema", "signature"]), ("dbx_execute_query", &["database", "session_id", "cell_char_offset", "cell_char_limit", "max_rows"]), ("dbx_execute_batch", &["database", "session_id", "continue_on_error", "use_transaction"]), ("dbx_open_session", &["database", "enable_transactions"]), ("dbx_open_table", &["database", "schema"]), ("dbx_execute_and_show", &["database"]), ("dbx_add_connection", &["port", "database", "driver_profile"]), ("dbx_remove_connection", &["connection_id", "connection_name"]), ("dbx_execute_redis_command", &["db"]), ("dbx_get_schema_context", &["database", "schema", "max_tables"]), ]; #[cfg(feature = "mq-admin")] checks .push(("dbx_send_message", &["key", "payload_text", "partition", "exchange", "routing_key", "namespace"])); #[cfg(feature = "mq-admin")] checks.push(("dbx_peek_messages", &["count", "partition", "offset"])); for (tool_name, fields) in checks { let tool = tools.iter().find(|tool| tool.name == *tool_name).expect("tool should be registered"); let properties = tool .input_schema .get("properties") .and_then(serde_json::Value::as_object) .expect("tool should publish object properties"); let required = tool .input_schema .get("required") .and_then(serde_json::Value::as_array) .into_iter() .flatten() .collect::>(); for &field in fields { let schema = properties.get(field).unwrap_or_else(|| panic!("{tool_name}.{field} should be published")); let type_value = schema.get("type").unwrap_or_else(|| panic!("{tool_name}.{field} should publish a type")); assert!( type_value.is_string(), "{tool_name}.{field} must publish a single concrete type, not a union: {type_value:?}" ); assert!(!required.iter().any(|required| *required == field), "{tool_name}.{field} must stay optional"); } } } #[test] fn execute_query_selector_preserves_serde_inputs() { let omitted: ExecuteQueryRequest = serde_json::from_str(r#"{"sql":"SELECT 1"}"#).unwrap(); let explicit_nulls: ExecuteQueryRequest = serde_json::from_str(r#"{"connection_id":null,"connection_name":null,"sql":"SELECT 1"}"#).unwrap(); let by_name: ExecuteQueryRequest = serde_json::from_str(r#"{"connection_name":"test_conn","sql":"SELECT 1"}"#).unwrap(); let by_id: ExecuteQueryRequest = serde_json::from_str(r#"{"connection_id":"123e4567-e89b-12d3-a456-426614174000","sql":"SELECT 1"}"#) .unwrap(); assert!(omitted.selector.connection_id.is_none()); assert!(omitted.selector.connection_name.is_none()); assert!(omitted.max_rows.is_none()); assert!(explicit_nulls.selector.connection_id.is_none()); assert!(explicit_nulls.selector.connection_name.is_none()); assert_eq!(by_name.selector.connection_name.as_deref(), Some("test_conn")); assert_eq!(by_id.selector.connection_id.as_deref(), Some("123e4567-e89b-12d3-a456-426614174000")); let with_max_rows: ExecuteQueryRequest = serde_json::from_str(r#"{"sql":"SELECT 1","max_rows":500}"#).unwrap(); assert_eq!(with_max_rows.max_rows, Some(500)); let nested = serde_json::from_str::( r#"{"connection_name":{"tool":"dbx_dbx_execute_query","error":"Invalid input"},"sql":"SELECT 1"}"#, ) .unwrap_err(); assert!(nested.to_string().contains("invalid type: map, expected a string")); } #[cfg(feature = "mq-admin")] #[test] fn send_message_request_preserves_binary_payload_and_optional_fields() { let request: SendMessageRequest = serde_json::from_str( r#"{"connection_name":"events","topic":"orders","payload_base64":"AP8=","key":"order-1","partition":2,"headers":{"trace":"abc"}}"#, ) .expect("send message request"); assert_eq!(request.selector.connection_name.as_deref(), Some("events")); assert_eq!(request.topic, "orders"); assert_eq!(request.payload_base64, "AP8="); assert_eq!(request.key.as_deref(), Some("order-1")); assert_eq!(request.partition, Some(2)); assert_eq!(request.headers.get("trace").map(String::as_str), Some("abc")); } #[cfg(feature = "mq-admin")] #[test] fn format_send_message_result_is_concise_and_includes_delivery_metadata() { let result = dbx_core::mq::SendMessageResponse { topic: "orders".to_string(), partition: 2, offset: 41, timestamp: Some("1700000000000".to_string()), }; let output = format_send_message_result(&dbx_core::mq::MqSystemKind::Kafka, &result); assert!(output.contains("Message sent to kafka.")); assert!(output.contains("topic: orders")); assert!(output.contains("partition: 2")); assert!(output.contains("offset: 41")); assert!(output.contains("timestamp: 1700000000000")); assert!(!output.contains("AP8=")); } #[test] fn schema_context_tables_preserve_optional_inputs() { let omitted: SchemaContextRequest = serde_json::from_str("{}").unwrap(); let explicit_null: SchemaContextRequest = serde_json::from_str(r#"{"tables":null}"#).unwrap(); let empty: SchemaContextRequest = serde_json::from_str(r#"{"tables":[]}"#).unwrap(); let populated: SchemaContextRequest = serde_json::from_str(r#"{"tables":["users","orders"]}"#).unwrap(); assert_eq!(omitted.tables, None); assert_eq!(explicit_null.tables, None); assert_eq!(empty.tables, Some(Vec::new())); assert_eq!(populated.tables, Some(vec!["users".to_string(), "orders".to_string()])); } #[test] fn scoped_server_hides_mutating_and_desktop_tools() { let server = DbxMcpServer::with_runtime_options( Arc::new(FakeBackend::default()), McpScope { connection_ids: vec!["scoped".to_string()], ..Default::default() }, false, ); let names = server.tool_router.list_all().into_iter().map(|tool| tool.name).collect::>(); #[cfg(feature = "mq-admin")] assert_eq!(names.len(), 17); #[cfg(not(feature = "mq-admin"))] assert_eq!(names.len(), 15); assert!(!names.iter().any(|name| name == "dbx_add_connection")); assert!(!names.iter().any(|name| name == "dbx_duplicate_connection")); assert!(!names.iter().any(|name| name == "dbx_remove_connection")); assert!(!names.iter().any(|name| name == "dbx_open_table")); assert!(!names.iter().any(|name| name == "dbx_execute_and_show")); assert!(names.iter().any(|name| name == "dbx_execute_batch")); assert!(names.iter().any(|name| name == "dbx_list_routines")); assert!(names.iter().any(|name| name == "dbx_get_routine_source")); assert!(names.iter().any(|name| name == "dbx_open_session")); assert!(names.iter().any(|name| name == "dbx_close_session")); assert!(names.iter().any(|name| name == "dbx_begin_transaction")); assert!(names.iter().any(|name| name == "dbx_commit_transaction")); assert!(names.iter().any(|name| name == "dbx_rollback_transaction")); #[cfg(feature = "mq-admin")] assert!(names.iter().any(|name| name == "dbx_send_message")); } #[test] fn scoped_connection_ids_are_deduplicated_and_take_precedence_over_name() { assert_eq!(scoped_connection_ids(Some(" first, second,first ,, ")), vec!["first", "second"]); let first = connection("first", "other", "sqlite", ":memory:"); let named = ConnectionConfig { id: "named".to_string(), name: "scope-name".to_string(), ..first.clone() }; let scope = McpScope { connection_ids: vec!["first".to_string()], connection_name: Some("scope-name".to_string()), database: None, schema: None, }; assert!(scope.matches(&first)); assert!(!scope.matches(&named)); } #[tokio::test] async fn database_scope_is_a_hard_bound_without_filtering_connections() { let scoped = connection("scoped", "scoped", "postgres", "configured"); let server = DbxMcpServer::with_runtime_options( Arc::new(FakeBackend { connections: vec![scoped.clone()], ..Default::default() }), McpScope { database: Some("analytics".to_string()), ..Default::default() }, false, ); assert_eq!(server.load_scoped_connections().await.unwrap().len(), 1); let resolved = resolved_connection_for_test(scoped.clone()); assert_eq!(server.resolve_database(None, &resolved).unwrap(), "analytics"); assert_eq!(server.resolve_database(Some("analytics".to_string()), &resolved).unwrap(), "analytics"); let error = server.resolve_database(Some("production".to_string()), &resolved).unwrap_err(); assert!(result_text(&error).contains("DATABASE_OUT_OF_SCOPE")); let names = server.tool_router.list_all().into_iter().map(|tool| tool.name).collect::>(); assert!(!names.iter().any(|name| name == "dbx_add_connection")); assert!(!names.iter().any(|name| name == "dbx_execute_and_show")); } #[test] fn configured_database_allowlist_is_enforced_for_sql_and_redis_targets() { let sql = connection("sql", "sql", "mysql", "reporting"); let selected = ResolvedConnection { connection: sql, policy: McpGlobalPolicy::default(), database_scope: DatabaseScope::Selected(vec!["reporting".to_string()]), group_ids: Vec::new(), }; let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); assert_eq!(server.resolve_database(Some("reporting".to_string()), &selected).unwrap(), "reporting"); let error = server.resolve_database(Some("production".to_string()), &selected).unwrap_err(); assert!(result_text(&error).contains("DATABASE_OUT_OF_SCOPE")); let redis = ResolvedConnection { connection: connection("redis", "redis", "redis", "0"), policy: McpGlobalPolicy::default(), database_scope: DatabaseScope::Selected(vec!["2".to_string()]), group_ids: Vec::new(), }; assert_eq!(server.resolve_redis_database(Some(2), &redis).unwrap(), 2); let error = server.resolve_redis_database(Some(3), &redis).unwrap_err(); assert!(result_text(&error).contains("DATABASE_OUT_OF_SCOPE")); } #[test] fn configured_database_allowlist_blocks_qualified_sql_references() { let connection = connection("sql", "sql", "mysql", "reporting"); let scope = DatabaseScope::Selected(vec!["reporting".to_string()]); assert!(ensure_sql_database_scope(&scope, &connection, "reporting", "SELECT * FROM reporting.users").is_ok()); let error = ensure_sql_database_scope(&scope, &connection, "reporting", "SELECT * FROM production.users").unwrap_err(); assert!(result_text(&error).contains("DATABASE_OUT_OF_SCOPE")); } #[test] fn database_execution_policy_overrides_connection_and_global_defaults_and_blocks_cross_database_sql() { let connection = connection("sql", "sql", "mysql", "operations"); let policy = McpGlobalPolicy { connection_policies: vec![dbx_core::storage::McpConnectionPolicy { connection_id: "sql".to_string(), read_only: false, allow_dangerous_sql: false, execution_mode_configured: true, execution_mode_policy_version: Some(dbx_core::mcp_policy::MCP_EXECUTION_POLICY_VERSION), database_scope: McpDatabaseScope::Selected, allowed_databases: vec!["operations".to_string(), "reporting".to_string()], database_policies: vec![ dbx_core::storage::McpDatabasePolicy { database_name: "operations".to_string(), read_only: false, allow_dangerous_sql: true, }, dbx_core::storage::McpDatabasePolicy { database_name: "reporting".to_string(), read_only: true, allow_dangerous_sql: false, }, ], }], ..Default::default() }; let operations = effective_policy_for_database(&policy, &connection, "operations"); assert!(!operations.read_only); assert!(operations.allow_dangerous_sql); let reporting = effective_policy_for_database(&policy, &connection, "reporting"); assert!(reporting.read_only); assert!(!reporting.allow_dangerous_sql); let error = ensure_sql_database_execution_scope( &policy, &connection, "operations", "UPDATE reporting.jobs SET done = 1", ) .unwrap_err(); assert!(result_text(&error).contains("DATABASE_EXECUTION_POLICY_OUT_OF_SCOPE")); } #[tokio::test] async fn selected_database_scope_can_be_discovered_without_a_connection_default_database() { let connection = connection("sql", "sql", "mysql", ""); let backend = Arc::new(FakeBackend { connections: vec![connection], policy: McpGlobalPolicy { allowed_connection_ids: Some(vec!["sql".to_string()]), connection_policies: vec![dbx_core::storage::McpConnectionPolicy { connection_id: "sql".to_string(), read_only: false, allow_dangerous_sql: false, database_scope: McpDatabaseScope::Selected, allowed_databases: vec!["aa".to_string(), "aaa".to_string(), "abc".to_string()], database_policies: Vec::new(), execution_mode_configured: false, execution_mode_policy_version: None, }], ..Default::default() }, ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); let result = server.list_databases(Parameters(ListDatabasesRequest { selector: selector("sql") })).await; assert_eq!(result_text(&result), "- aa\n- aaa\n- abc"); let show_databases = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("sql"), database: None, sql: "SHOW DATABASES".to_string(), session_id: None, cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await; assert_eq!(result_text(&show_databases), "- aa\n- aaa\n- abc"); assert!(is_database_discovery_sql(" SHOW DATABASES; ")); assert!(is_database_discovery_sql("show schemas")); assert!(!is_database_discovery_sql("show tables")); } #[test] fn schema_scope_is_a_hard_bound() { let dameng = connection("dameng-1", "Dameng", "dameng", "APPDB"); let server = DbxMcpServer::with_runtime_options( Arc::new(FakeBackend::default()), McpScope { database: Some("APPDB".to_string()), schema: Some("REPORTING".to_string()), ..Default::default() }, false, ); let resolved = resolved_connection_for_test(dameng.clone()); assert_eq!(server.resolve_database(None, &resolved).unwrap(), "APPDB"); assert_eq!(server.resolve_schema(None).unwrap(), "REPORTING"); assert_eq!(server.resolve_schema(Some("REPORTING".to_string())).unwrap(), "REPORTING"); let error = server.resolve_schema(Some("APP_USER".to_string())).unwrap_err(); assert!(result_text(&error).contains("SCHEMA_OUT_OF_SCOPE")); } #[test] fn redis_database_scope_fails_closed_and_cannot_be_overridden() { let redis = connection("redis", "redis", "redis", "1"); let scoped = DbxMcpServer::with_runtime_options( Arc::new(FakeBackend { connections: vec![redis.clone()], ..Default::default() }), McpScope { database: Some("2".to_string()), ..Default::default() }, false, ); let resolved = resolved_connection_for_test(redis.clone()); assert_eq!(scoped.resolve_redis_database(None, &resolved).unwrap(), 2); let error = scoped.resolve_redis_database(Some(3), &resolved).unwrap_err(); assert!(result_text(&error).contains("DATABASE_OUT_OF_SCOPE")); let invalid = DbxMcpServer::with_runtime_options( Arc::new(FakeBackend { connections: vec![redis.clone()], ..Default::default() }), McpScope { database: Some("analytics".to_string()), ..Default::default() }, false, ); let error = invalid.resolve_redis_database(None, &resolved).unwrap_err(); assert!(result_text(&error).contains("INVALID_DATABASE_SCOPE")); } #[test] fn local_mongo_aggregate_cannot_write_to_a_production_database() { let mut mongo = connection("mongo", "mongo", "mongodb", "staging"); mongo.production_databases = vec!["production".to_string()]; let policy = McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None, ..Default::default() }; let error = validate_mongo_command( &mongo, &policy, &DatabaseScope::All, "staging", r#"db.items.aggregate([{"$out":{"db":"production","coll":"archive"}}])"#, ) .unwrap_err(); assert!(result_text(&error).contains("PRODUCTION_WRITE_BLOCKED")); assert!(validate_mongo_command( &mongo, &policy, &DatabaseScope::All, "staging", r#"db.items.aggregate([{"$out":"archive"}])"#, ) .is_ok()); } #[test] fn mongo_run_command_is_never_exposed_through_mcp() { let mongo = connection("mongo", "mongo", "mongodb", "staging"); let source = r#"db.runCommand({ping: 1})"#; for policy in [ McpGlobalPolicy { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: None, ..Default::default() }, McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None, ..Default::default() }, McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None, ..Default::default() }, ] { let error = validate_mongo_command(&mongo, &policy, &DatabaseScope::All, "staging", source).unwrap_err(); assert!(result_text(&error).contains("SQL_BLOCKED")); assert!(result_text(&error).contains("runCommand")); } } #[test] fn mongo_sibling_database_run_command_is_never_exposed_through_mcp() { let mongo = connection("mongo", "mongo", "mongodb", "operations"); let policy = McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None, ..Default::default() }; let error = validate_mongo_command( &mongo, &policy, &DatabaseScope::All, "operations", r#"db.getSiblingDB("staging").runCommand({ping: 1})"#, ) .unwrap_err(); assert!(result_text(&error).contains("SQL_BLOCKED")); assert!(result_text(&error).contains("runCommand")); } #[test] fn mongo_sibling_database_honors_target_scope_and_read_only_policy() { let mongo = connection("mongo", "mongo", "mongodb", "operations"); let policy = McpGlobalPolicy { connection_policies: vec![dbx_core::storage::McpConnectionPolicy { connection_id: "mongo".to_string(), read_only: false, allow_dangerous_sql: true, execution_mode_configured: true, execution_mode_policy_version: Some(dbx_core::mcp_policy::MCP_EXECUTION_POLICY_VERSION), database_scope: McpDatabaseScope::Selected, allowed_databases: vec!["operations".to_string(), "reporting".to_string()], database_policies: vec![ dbx_core::storage::McpDatabasePolicy { database_name: "operations".to_string(), read_only: false, allow_dangerous_sql: true, }, dbx_core::storage::McpDatabasePolicy { database_name: "reporting".to_string(), read_only: true, allow_dangerous_sql: false, }, ], }], ..Default::default() }; let read_only_error = validate_mongo_command( &mongo, &policy, &DatabaseScope::Selected(vec!["operations".to_string(), "reporting".to_string()]), "operations", r#"db.getSiblingDB("reporting").users.deleteMany({_id: 1})"#, ) .unwrap_err(); assert!(result_text(&read_only_error).contains("MCP_READ_ONLY")); let scope_error = validate_mongo_command( &mongo, &policy, &DatabaseScope::Selected(vec!["operations".to_string()]), "operations", r#"db.getSiblingDB("reporting").users.find({})"#, ) .unwrap_err(); assert!(result_text(&scope_error).contains("DATABASE_OUT_OF_SCOPE")); } #[test] fn mongo_sibling_database_aggregate_honors_output_scope_and_production_protection() { let mut mongo = connection("mongo", "mongo", "mongodb", "operations"); mongo.production_databases = vec!["production".to_string()]; let policy = McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None, ..Default::default() }; let production_error = validate_mongo_command( &mongo, &policy, &DatabaseScope::All, "operations", r#"db.getSiblingDB("staging").items.aggregate([{"$out":{"db":"production","coll":"archive"}}])"#, ) .unwrap_err(); assert!(result_text(&production_error).contains("PRODUCTION_WRITE_BLOCKED")); let scope_error = validate_mongo_command( &mongo, &policy, &DatabaseScope::Selected(vec!["operations".to_string(), "staging".to_string()]), "operations", r#"db.getSiblingDB("staging").items.aggregate([{"$merge":{"into":{"db":"archive","coll":"items"}}}])"#, ) .unwrap_err(); assert!(result_text(&scope_error).contains("DATABASE_OUT_OF_SCOPE")); } #[test] fn agent_results_preserve_stable_backend_policy_errors() { let result = agent_result(dbx_core::agent_events::ToolResult { tool_call_id: "test".to_string(), tool_name: "execute_query".to_string(), content: "Error: API request failed: MCP_READ_ONLY: policy changed".to_string(), is_error: true, explain_data: None, }); assert!(result_text(&result).contains("Error [MCP_READ_ONLY]: policy changed")); } #[test] fn mcp_confirmation_and_policy_guards_fail_closed() { assert_eq!( normalize_confirmed_write_sql(Some(" DELETE FROM sessions WHERE id = 7 ".to_string())), Some("DELETE FROM sessions WHERE id = 7".to_string()) ); assert_eq!(normalize_confirmed_write_sql(Some(" \n ".to_string())), None); let read_only = ConnectionConfig { read_only: true, ..connection("readonly", "readonly", "postgres", "app") }; let writable_policy = McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None, ..Default::default() }; let read_only_error = validate_sql_policy(&read_only, &writable_policy, "app", "DELETE FROM sessions", false).unwrap_err(); assert!(result_text(&read_only_error).contains("CONNECTION_READ_ONLY")); let mut production = connection("production", "production", "postgres", "app"); production.production_databases = vec!["app".to_string()]; let production_error = validate_sql_policy(&production, &writable_policy, "app", "DROP TABLE sessions", false).unwrap_err(); assert!(result_text(&production_error).contains("PRODUCTION_WRITE_BLOCKED")); } /// RAII guard that sets an env var and restores the original value (or /// removes the var) on drop. Panic-safe — cleanup runs even when an /// assertion fails. struct EnvGuard { key: &'static str, original: Option, } impl EnvGuard { fn set(key: &'static str, value: &str) -> Self { let original = std::env::var(key).ok(); std::env::set_var(key, value); Self { key, original } } fn remove(key: &'static str) -> Self { let original = std::env::var(key).ok(); std::env::remove_var(key); Self { key, original } } } impl Drop for EnvGuard { fn drop(&mut self) { match &self.original { Some(original) => std::env::set_var(self.key, original), None => std::env::remove_var(self.key), } } } #[test] fn confirmed_sql_binding_cannot_elevate_central_policy() { let connection = connection("dev", "dev", "postgres", "app"); // An exact confirmation remains available as a narrowing constraint, // but cannot turn the central safe-write policy into full access. let _guard = EnvGuard::set("DBX_MCP_CONFIRMED_WRITE_SQL", "CREATE TABLE metrics (id INT)"); let safe_write = McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None, ..Default::default() }; let permissions = mcp_permissions(&connection, &safe_write); assert!(!permissions.allow_dangerous, "confirmed SQL must NOT elevate allow_dangerous for Redis/Mongo paths"); assert!(permissions.allow_writes); assert_eq!(permissions.confirmed_write_sql.as_deref(), Some("CREATE TABLE metrics (id INT)")); let error = validate_sql_policy(&connection, &safe_write, "app", "CREATE TABLE metrics (id INT)", false).unwrap_err(); assert!(result_text(&error).contains("SQL_BLOCKED")); assert!(result_text(&error).contains("High-risk SQL is disabled")); drop(_guard); // A missing binding cannot change the same safe-write boundary. let _guard = EnvGuard::remove("DBX_MCP_CONFIRMED_WRITE_SQL"); let error = validate_sql_policy(&connection, &safe_write, "app", "DROP TABLE sessions", false).unwrap_err(); assert!(result_text(&error).contains("SQL_BLOCKED")); assert!(result_text(&error).contains("High-risk SQL is disabled")); drop(_guard); // Full access still permits an exact confirmed DDL statement and keeps // the binding for the execution-time anti-replay check. let _guard = EnvGuard::set("DBX_MCP_CONFIRMED_WRITE_SQL", "CREATE TABLE metrics (id INT)"); let full_access = McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None, ..Default::default() }; let permissions = validate_sql_policy(&connection, &full_access, "app", "CREATE TABLE metrics (id INT)", false).unwrap(); assert!(permissions.allow_writes); assert!(permissions.allow_dangerous); assert_eq!(permissions.confirmed_write_sql.as_deref(), Some("CREATE TABLE metrics (id INT)")); // A global read-only policy is still authoritative even for the exact // confirmed statement, while read queries remain available. let read_only_policy = McpGlobalPolicy { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: None, ..Default::default() }; let error = validate_sql_policy(&connection, &read_only_policy, "app", "CREATE TABLE metrics (id INT)", false) .unwrap_err(); assert!(result_text(&error).contains("MCP_READ_ONLY"), "confirmed SQL must not bypass global read_only"); assert!(validate_sql_policy(&connection, &read_only_policy, "app", "SELECT 1", false).is_ok()); drop(_guard); // Safe-write DML remains allowed, and an exact binding never grants the // DDL/high-risk bit implicitly. let _guard = EnvGuard::set("DBX_MCP_CONFIRMED_WRITE_SQL", "INSERT INTO metrics (id) VALUES (1)"); let permissions = validate_sql_policy(&connection, &safe_write, "app", "INSERT INTO metrics (id) VALUES (1)", false).unwrap(); assert!(permissions.allow_writes); assert!(!permissions.allow_dangerous); assert_eq!(permissions.confirmed_write_sql.as_deref(), Some("INSERT INTO metrics (id) VALUES (1)")); } #[test] fn use_statements_require_a_session() { let starrocks = connection("sr", "sr", "starrocks", "default_catalog"); let policy = McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None, ..Default::default() }; let blocked = validate_sql_policy(&starrocks, &policy, "default_catalog", "USE analytics", false).unwrap_err(); assert!(result_text(&blocked).contains("SQL_BLOCKED")); assert!(result_text(&blocked).contains("dbx_open_session")); // Inside a pinned session, USE is meaningful and passes policy checks. assert!(validate_sql_policy(&starrocks, &policy, "default_catalog", "USE analytics", true).is_ok()); } fn selector(id: &str) -> ConnectionSelector { ConnectionSelector { connection_id: Some(id.to_string()), connection_name: None } } #[tokio::test] async fn opted_in_session_exposes_transaction_tools_and_structured_state() { let policy = McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, ..Default::default() }; let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], policy, ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; let session_id = opened_session_id(&opened); assert_eq!(backend.transaction_owners_opened.load(std::sync::atomic::Ordering::SeqCst), 1); assert_eq!(opened.structured_content.as_ref().unwrap()["transaction_state"], "idle"); let begun = server.begin_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; assert!(result_text(&begun).contains("transaction_state: active")); assert_eq!(begun.structured_content.as_ref().unwrap()["transaction_state"], "active"); let queried = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("mysql"), database: None, sql: "SELECT id FROM accounts WHERE id = 1 FOR UPDATE".to_string(), session_id: Some(session_id.clone()), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await; assert!(queried.structured_content.is_some(), "{}", result_text(&queried)); assert_eq!(queried.structured_content.as_ref().unwrap()["transaction_state"], "active"); let committed = server.commit_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; let structured = committed.structured_content.as_ref().unwrap(); assert!(result_text(&committed).contains("transaction_outcome: committed")); assert_eq!(structured["transaction_state"], "idle"); assert_eq!(structured["transaction_outcome"], "committed"); let closed = server.close_session(Parameters(CloseSessionRequest { session_id })).await; assert!(!closed.is_error.unwrap_or(false)); } #[tokio::test] async fn legacy_session_open_text_is_unchanged_when_transactions_are_disabled() { let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: false, })) .await; assert!(!result_text(&opened).contains("transactions:")); } #[tokio::test] async fn transaction_opt_in_rejects_non_mysql_external_profiles_and_unsupported_backends() { let postgres_backend = Arc::new(FakeBackend { connections: vec![connection("pg", "pg", "postgres", "app")], ..Default::default() }); let postgres_server = DbxMcpServer::with_runtime_options(postgres_backend.clone(), McpScope::default(), false); let rejected = postgres_server .open_session(Parameters(OpenSessionRequest { selector: selector("pg"), database: None, enable_transactions: true, })) .await; assert!(result_text(&rejected).contains("TRANSACTION_UNSUPPORTED")); assert_eq!(postgres_backend.transaction_owners_opened.load(std::sync::atomic::Ordering::SeqCst), 0); let mut builtin_mysql = connection("builtin", "builtin", "mysql", "app"); builtin_mysql.driver_profile = Some("mysql".to_string()); let builtin_backend = Arc::new(FakeBackend { connections: vec![builtin_mysql], ..Default::default() }); let builtin_server = DbxMcpServer::with_runtime_options(builtin_backend.clone(), McpScope::default(), false); let opened = builtin_server .open_session(Parameters(OpenSessionRequest { selector: selector("builtin"), database: None, enable_transactions: true, })) .await; assert!(!opened.is_error.unwrap_or(false)); assert_eq!(builtin_backend.transaction_owners_opened.load(std::sync::atomic::Ordering::SeqCst), 1); let mut external_mysql = connection("external", "external", "mysql", "app"); external_mysql.driver_profile = Some("custom-profile".to_string()); let external_backend = Arc::new(FakeBackend { connections: vec![external_mysql], ..Default::default() }); let external_server = DbxMcpServer::with_runtime_options(external_backend.clone(), McpScope::default(), false); let rejected = external_server .open_session(Parameters(OpenSessionRequest { selector: selector("external"), database: None, enable_transactions: true, })) .await; assert!(result_text(&rejected).contains("TRANSACTION_UNSUPPORTED")); assert_eq!(external_backend.transaction_owners_opened.load(std::sync::atomic::Ordering::SeqCst), 0); let unsupported_backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], transaction_open_error: Some("TRANSACTION_UNSUPPORTED: backend has no native fixed session".to_string()), ..Default::default() }); let unsupported_server = DbxMcpServer::with_runtime_options(unsupported_backend, McpScope::default(), false); let rejected = unsupported_server .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; assert!(result_text(&rejected).contains("TRANSACTION_UNSUPPORTED")); } #[tokio::test] async fn invalid_transaction_transitions_do_not_send_sql() { let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; let session_id = opened_session_id(&opened); let no_active_commit = server.commit_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; assert!(result_text(&no_active_commit).contains("TRANSACTION_NOT_ACTIVE")); server.begin_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; let nested = server.begin_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; assert!(result_text(&nested).contains("TRANSACTION_ALREADY_ACTIVE")); server.rollback_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; let no_active_rollback = server.rollback_transaction(Parameters(TransactionSessionRequest { session_id })).await; assert!(result_text(&no_active_rollback).contains("TRANSACTION_NOT_ACTIVE")); assert_eq!(backend.transaction_owner_sql.lock().unwrap().as_slice(), ["START TRANSACTION", "ROLLBACK"]); } #[tokio::test] async fn queued_transaction_query_rechecks_revoked_policy_before_io() { let policy = McpGlobalPolicy { allow_dangerous_sql: true, ..Default::default() }; let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], policy: policy.clone(), ..Default::default() }); let server = Arc::new(DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false)); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; let session_id = opened_session_id(&opened); server.begin_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; let release_first = Arc::new(tokio::sync::Notify::new()); *backend.transaction_block_next_sql.lock().unwrap() = Some(release_first.clone()); let first = { let server = server.clone(); let session_id = session_id.clone(); tokio::spawn(async move { server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("mysql"), database: None, sql: "SELECT id FROM accounts WHERE id = 1 FOR UPDATE".to_string(), session_id: Some(session_id), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await }) }; while backend.transaction_owner_sql.lock().unwrap().len() < 2 { tokio::task::yield_now().await; } let queued = { let server = server.clone(); let session_id = session_id.clone(); tokio::spawn(async move { server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("mysql"), database: None, sql: "UPDATE accounts SET name = 'revoked' WHERE id = 1".to_string(), session_id: Some(session_id), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await }) }; tokio::task::yield_now().await; let mut revoked = policy; revoked.read_only = true; *backend.policy_override.lock().unwrap() = Some(revoked); release_first.notify_one(); assert!(!first.await.unwrap().is_error.unwrap_or(false)); let denied = queued.await.unwrap(); assert!(result_text(&denied).contains("MCP_READ_ONLY")); assert_eq!( backend.transaction_owner_sql.lock().unwrap().as_slice(), ["START TRANSACTION", "SELECT id FROM accounts WHERE id = 1 FOR UPDATE"] ); let rolled_back = server.rollback_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; assert!(!rolled_back.is_error.unwrap_or(false)); let closed = server.close_session(Parameters(CloseSessionRequest { session_id })).await; assert!(!closed.is_error.unwrap_or(false)); } #[tokio::test] async fn transaction_batch_keeps_one_lease_and_reports_each_resulting_state() { let policy = McpGlobalPolicy { allow_dangerous_sql: true, ..Default::default() }; let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], policy, ..Default::default() }); let server = Arc::new(DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false)); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; let session_id = opened_session_id(&opened); server.begin_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; let release_first = Arc::new(tokio::sync::Notify::new()); *backend.transaction_block_next_sql.lock().unwrap() = Some(release_first.clone()); let batch = { let server = server.clone(); let session_id = session_id.clone(); tokio::spawn(async move { server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("mysql"), database: None, sql: "UPDATE accounts SET name = 'first' WHERE id = 1; SELECT id FROM accounts WHERE id = 1" .to_string(), session_id: Some(session_id), continue_on_error: Some(true), use_transaction: None, })) .await }) }; while backend.transaction_owner_sql.lock().unwrap().len() < 2 { tokio::task::yield_now().await; } let queued_query = { let server = server.clone(); let session_id = session_id.clone(); tokio::spawn(async move { server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("mysql"), database: None, sql: "SELECT id FROM accounts WHERE id = 2".to_string(), session_id: Some(session_id), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await }) }; release_first.notify_one(); let batch = batch.await.unwrap(); assert!(!batch.is_error.unwrap_or(false)); let structured = batch.structured_content.as_ref().unwrap(); assert_eq!(structured["transaction_state"], "active"); assert_eq!(structured["results"][0]["transaction_state"], "active"); assert_eq!(structured["results"][1]["transaction_state"], "active"); assert!(!queued_query.await.unwrap().is_error.unwrap_or(false)); assert_eq!( backend.transaction_owner_sql.lock().unwrap().as_slice(), [ "START TRANSACTION", "UPDATE accounts SET name = 'first' WHERE id = 1", "SELECT id FROM accounts WHERE id = 1", "SELECT id FROM accounts WHERE id = 2", ] ); } #[tokio::test] async fn dropping_protocol_server_rolls_back_its_active_transaction_sessions() { let policy = McpGlobalPolicy { allow_dangerous_sql: true, ..Default::default() }; let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], policy, ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; let session_id = opened_session_id(&opened); server.begin_transaction(Parameters(TransactionSessionRequest { session_id })).await; drop(server); tokio::time::timeout(std::time::Duration::from_millis(250), async { loop { if backend.transaction_owner_sql.lock().unwrap().as_slice() == ["START TRANSACTION", "ROLLBACK"] { break; } tokio::task::yield_now().await; } }) .await .expect("protocol session teardown must trigger bounded rollback"); } #[tokio::test] async fn protocol_cancellation_drops_a_busy_transaction_tool_request() { let policy = McpGlobalPolicy { allow_dangerous_sql: true, ..Default::default() }; let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], policy, ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let (server_transport, client_transport) = tokio::io::duplex(64 * 1024); let server_task = tokio::spawn(async move { server.serve(server_transport).await }); let mut client = ().serve(client_transport).await.unwrap(); let mut service = server_task.await.unwrap().unwrap(); let arguments = |value: serde_json::Value| value.as_object().cloned().unwrap(); let opened = client .peer() .call_tool(CallToolRequestParams::new("dbx_open_session").with_arguments(arguments(json!({ "connection_id": "mysql", "database": "app", "enable_transactions": true })))) .await .unwrap(); let session_id = opened.structured_content.unwrap()["session_id"].as_str().unwrap().to_string(); client .peer() .call_tool( CallToolRequestParams::new("dbx_begin_transaction") .with_arguments(arguments(json!({"session_id": session_id}))), ) .await .unwrap(); *backend.transaction_block_next_sql.lock().unwrap() = Some(Arc::new(tokio::sync::Notify::new())); let request = ClientRequest::CallToolRequest(CallToolRequest::new( CallToolRequestParams::new("dbx_execute_query").with_arguments(arguments(json!({ "connection_id": "mysql", "database": "app", "session_id": session_id, "sql": "SELECT id FROM accounts WHERE id = 1 FOR UPDATE" }))), )); let request = client.peer().send_cancellable_request(request, PeerRequestOptions::no_options()).await.unwrap(); while backend.transaction_owner_sql.lock().unwrap().len() < 2 { tokio::task::yield_now().await; } request.cancel(Some("test cancellation".to_string())).await.unwrap(); let follow_up = tokio::time::timeout( std::time::Duration::from_millis(250), client.peer().call_tool(CallToolRequestParams::new("dbx_execute_query").with_arguments(arguments(json!({ "connection_id": "mysql", "database": "app", "session_id": session_id, "sql": "SELECT 2" })))), ) .await .expect("protocol cancellation must release the transaction operation slot") .unwrap(); assert_eq!(follow_up.is_error, Some(true)); let _ = client.close_with_timeout(std::time::Duration::from_millis(250)).await; let _ = service.close_with_timeout(std::time::Duration::from_millis(250)).await; } #[tokio::test] async fn session_queries_pin_client_session_and_close_releases_pool() { let starrocks = connection("sr", "sr", "starrocks", "default_catalog"); let backend = Arc::new(FakeBackend { connections: vec![starrocks], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("sr"), database: None, enable_transactions: false, })) .await; let session_id = opened_session_id(&opened); // A USE statement is allowed inside the session and runs with the // session's pinned client_session_id. let result = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("sr"), database: None, sql: "USE analytics".to_string(), session_id: Some(session_id.clone()), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await; assert_eq!(result_text(&result), "ok"); let pinned_client_session = { let recorded = backend.recorded_arguments.lock().unwrap(); let (_, arguments) = recorded.iter().find(|(name, _)| name == "execute_query").unwrap(); arguments["client_session_id"].as_str().unwrap().to_string() }; assert_eq!(pinned_client_session, format!("mcp:{session_id}")); // Queries bound to another database are rejected. let mismatch = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("sr"), database: Some("other".to_string()), sql: "SELECT 1".to_string(), session_id: Some(session_id.clone()), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await; assert!(result_text(&mismatch).contains("SESSION_DATABASE_MISMATCH")); let closed = server.close_session(Parameters(CloseSessionRequest { session_id: session_id.clone() })).await; assert!(result_text(&closed).contains("closed")); assert_eq!(backend.closed_sessions.lock().unwrap().as_slice(), [format!("mcp:{session_id}")]); // The session is gone: further queries fail instead of silently // falling back to an unpinned connection. let missing = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("sr"), database: None, sql: "SELECT 1".to_string(), session_id: Some(session_id.clone()), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await; assert!(result_text(&missing).contains("SESSION_NOT_FOUND")); let second_close = server.close_session(Parameters(CloseSessionRequest { session_id })).await; assert!(result_text(&second_close).contains("SESSION_NOT_FOUND")); } #[tokio::test] async fn execute_query_forwards_character_window_options() { let elasticsearch = connection("es", "es", "elasticsearch", ""); let backend = Arc::new(FakeBackend { connections: vec![elasticsearch], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let result = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("es"), database: None, sql: "GET /logs/_search".to_string(), session_id: None, cell_char_offset: Some(200), cell_char_limit: Some(800), max_rows: None, })) .await; assert_eq!(result_text(&result), "ok"); let recorded = backend.recorded_arguments.lock().unwrap(); let (_, arguments) = recorded.iter().find(|(name, _)| name == "execute_query").unwrap(); assert_eq!(arguments["cell_char_offset"], 200); assert_eq!(arguments["cell_char_limit"], 800); } #[tokio::test] async fn execute_query_forwards_and_clamps_max_rows() { let elasticsearch = connection("esm", "esm", "elasticsearch", ""); let backend = Arc::new(FakeBackend { connections: vec![elasticsearch], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); for max_rows in [None, Some(500), Some(100_000), Some(0)] { let result = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("esm"), database: None, sql: "GET /logs/_search".to_string(), session_id: None, cell_char_offset: None, cell_char_limit: None, max_rows, })) .await; assert_eq!(result_text(&result), "ok"); } let recorded = backend.recorded_arguments.lock().unwrap(); let limits = recorded .iter() .filter(|(name, _)| name == "execute_query") .map(|(_, arguments)| arguments["limit"].as_u64().expect("limit should be published")) .collect::>(); // Omitted keeps the historical 100-row default; explicit values are clamped // into 1..=MAX_EXECUTE_QUERY_ROWS rather than rejected. assert_eq!(limits, vec![100, 500, MAX_EXECUTE_QUERY_ROWS as u64, 1]); } #[tokio::test] async fn expired_sessions_close_backend_pools_without_accumulating() { let starrocks = connection("sr", "sr", "starrocks", "default_catalog"); let backend = Arc::new(FakeBackend { connections: vec![starrocks], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let mut expired_client_session_ids = Vec::new(); for _ in 0..3 { let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("sr"), database: None, enable_transactions: false, })) .await; let session_id = opened_session_id(&opened); let result = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("sr"), database: None, sql: "SELECT 1".to_string(), session_id: Some(session_id.clone()), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await; assert_eq!(result_text(&result), "ok"); assert_eq!(backend.pinned_sessions.lock().unwrap().len(), 1); server.sessions.expire_for_test(&session_id).await; expired_client_session_ids.push(format!("mcp:{session_id}")); } let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("sr"), database: None, enable_transactions: false, })) .await; let final_session_id = opened_session_id(&opened); assert!(backend.pinned_sessions.lock().unwrap().is_empty()); assert_eq!(backend.closed_sessions.lock().unwrap().as_slice(), expired_client_session_ids.as_slice()); let closed = server.close_session(Parameters(CloseSessionRequest { session_id: final_session_id })).await; assert!(result_text(&closed).contains("closed")); } #[tokio::test] async fn cancelled_eager_open_while_backend_acquisition_is_blocked_does_not_leak_capacity() { let entered = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Notify::new()); let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], transaction_open_before: Some((entered.clone(), release.clone())), ..Default::default() }); let server = Arc::new(DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false)); let opening = { let server = server.clone(); tokio::spawn(async move { server .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await }) }; entered.notified().await; opening.abort(); assert!(opening.await.unwrap_err().is_cancelled()); release.notify_waiters(); tokio::time::timeout(Duration::from_millis(500), async { while backend.closed_sessions.lock().unwrap().is_empty() { tokio::task::yield_now().await; } }) .await .expect("cancelled provisional open must close its backend session"); assert_eq!(backend.transaction_owner_disconnects.load(std::sync::atomic::Ordering::SeqCst), 1); } #[tokio::test] async fn cancelled_eager_open_after_owner_acquisition_disposes_owner_and_backend_session() { let entered = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Notify::new()); let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], transaction_open_after_owner: Some((entered.clone(), release.clone())), ..Default::default() }); let server = Arc::new(DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false)); let opening = { let server = server.clone(); tokio::spawn(async move { server .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await }) }; entered.notified().await; opening.abort(); assert!(opening.await.unwrap_err().is_cancelled()); release.notify_waiters(); tokio::time::timeout(Duration::from_millis(500), async { while backend.closed_sessions.lock().unwrap().is_empty() { tokio::task::yield_now().await; } }) .await .expect("cancelled prepared open must close its backend session"); assert_eq!(backend.transaction_owner_disconnects.load(std::sync::atomic::Ordering::SeqCst), 1); } #[tokio::test] async fn two_protocol_servers_share_physical_owner_budget_until_cleanup_finishes() { let backend = Arc::new(FakeBackend { connections: vec![connection("mysql", "mysql", "mysql", "app")], ..Default::default() }); let server_a = Arc::new(DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false)); let server_b = Arc::new(DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false)); let mut sessions_a = Vec::new(); for _ in 0..16 { let opened = server_a .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; assert!(!opened.is_error.unwrap_or(false)); sessions_a.push(opened_session_id(&opened)); } for _ in 0..16 { let opened = server_b .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; assert!(!opened.is_error.unwrap_or(false)); } let rejected = server_b .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; assert!(result_text(&rejected).contains("max 32")); let session_id = sessions_a.pop().unwrap(); let hidden_from_other_protocol = server_b.begin_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; assert!(result_text(&hidden_from_other_protocol).contains("SESSION_NOT_FOUND")); server_a.begin_transaction(Parameters(TransactionSessionRequest { session_id: session_id.clone() })).await; let rollback_release = Arc::new(tokio::sync::Notify::new()); *backend.transaction_block_next_sql.lock().unwrap() = Some(rollback_release.clone()); let closing = { let server = server_a.clone(); let session_id = session_id.clone(); tokio::spawn(async move { server.close_session(Parameters(CloseSessionRequest { session_id })).await }) }; while !backend.transaction_owner_sql.lock().unwrap().iter().any(|sql| sql == "ROLLBACK") { tokio::task::yield_now().await; } closing.abort(); assert!(closing.await.unwrap_err().is_cancelled()); let still_rejected = server_b .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; assert!(result_text(&still_rejected).contains("max 32")); rollback_release.notify_one(); tokio::time::timeout(Duration::from_millis(500), async { while backend.transaction_owner_budget.available_permits() == 0 { tokio::task::yield_now().await; } }) .await .expect("physical owner permit must return after disposal"); let opened = server_b .open_session(Parameters(OpenSessionRequest { selector: selector("mysql"), database: None, enable_transactions: true, })) .await; assert!(!opened.is_error.unwrap_or(false)); } #[tokio::test] async fn failed_session_close_can_be_retried() { let starrocks = connection("sr", "sr", "starrocks", "default_catalog"); let backend = Arc::new(FakeBackend { connections: vec![starrocks], ..Default::default() }); *backend.close_failures_remaining.lock().unwrap() = 1; let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("sr"), database: None, enable_transactions: false, })) .await; let session_id = opened_session_id(&opened); let query = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("sr"), database: None, sql: "SELECT 1".to_string(), session_id: Some(session_id.clone()), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await; assert_eq!(result_text(&query), "ok"); let failed = server.close_session(Parameters(CloseSessionRequest { session_id: session_id.clone() })).await; assert!(result_text(&failed).contains("SESSION_CLOSE_ERROR")); assert_eq!(backend.pinned_sessions.lock().unwrap().len(), 1); let retry_query = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("sr"), database: None, sql: "SELECT 1".to_string(), session_id: Some(session_id.clone()), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await; assert_eq!(result_text(&retry_query), "ok"); let closed = server.close_session(Parameters(CloseSessionRequest { session_id })).await; assert!(result_text(&closed).contains("closed")); assert!(backend.pinned_sessions.lock().unwrap().is_empty()); assert_eq!(backend.closed_sessions.lock().unwrap().len(), 1); } #[tokio::test] async fn open_session_rejects_non_sql_connections_and_unknown_sessions_fail_closed() { let redis = connection("redis", "redis", "redis", "0"); let pg = connection("pg", "pg", "postgres", "app"); let server = DbxMcpServer::with_runtime_options( Arc::new(FakeBackend { connections: vec![redis, pg], ..Default::default() }), McpScope::default(), false, ); let rejected = server .open_session(Parameters(OpenSessionRequest { selector: selector("redis"), database: None, enable_transactions: false, })) .await; assert!(result_text(&rejected).contains("SESSION_UNSUPPORTED")); let missing = server .execute_query(Parameters(ExecuteQueryRequest { selector: selector("pg"), database: None, sql: "SELECT 1".to_string(), session_id: Some("mcp-session-nope".to_string()), cell_char_offset: None, cell_char_limit: None, max_rows: None, })) .await; assert!(result_text(&missing).contains("SESSION_NOT_FOUND")); } #[tokio::test] async fn execute_batch_forwards_options_and_surfaces_results() { let postgres = connection("pg", "pg", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: Some("app".to_string()), sql: "SELECT 1; SELECT 2".to_string(), session_id: None, continue_on_error: Some(true), use_transaction: None, })) .await; // Auto-commit mode returns per-statement results; continue_on_error is // forwarded to the core. assert!(result_text(&result).contains("Statement 1")); let ephemeral_session_id = { let recorded = backend.recorded_arguments.lock().unwrap(); let (_, arguments) = recorded.iter().find(|(name, _)| name == "execute_batch").unwrap(); assert_eq!(arguments["continue_on_error"], true); assert!( arguments["use_transaction"].is_null(), "use_transaction must be forwarded as null, got: {}", arguments["use_transaction"] ); arguments["client_session_id"] .as_str() .expect("auto-commit batch must use an ephemeral session") .to_string() }; assert!(ephemeral_session_id.starts_with("mcp-batch-")); assert_eq!(backend.closed_sessions.lock().unwrap().as_slice(), [ephemeral_session_id]); } #[tokio::test] async fn execute_batch_pins_session_and_forwards_client_session_id() { let postgres = connection("pg", "pg", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("pg"), database: None, enable_transactions: false, })) .await; let session_id = opened_session_id(&opened); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "SELECT 1; SELECT 2".to_string(), session_id: Some(session_id.clone()), continue_on_error: Some(false), use_transaction: None, })) .await; assert!(result_text(&result).contains("Statement 1"), "got: {}", result_text(&result)); let recorded = backend.recorded_arguments.lock().unwrap(); let (_, arguments) = recorded.iter().find(|(name, _)| name == "execute_batch").unwrap(); assert_eq!(arguments["client_session_id"], format!("mcp:{session_id}")); assert!(backend.pinned_sessions.lock().unwrap().contains(&format!("mcp:{session_id}"))); } #[tokio::test] async fn execute_batch_blocks_redis_and_mongo_connections() { for (id, db_type) in [("redis", "redis"), ("mongo", "mongodb")] { let backend = Arc::new(FakeBackend { connections: vec![connection(id, id, db_type, "db")], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector(id), database: None, sql: "SELECT 1".to_string(), session_id: None, continue_on_error: None, use_transaction: None, })) .await; assert!(result_text(&result).contains("DBX_BATCH_UNSUPPORTED")); } } #[tokio::test] async fn execute_batch_forwards_scoped_schema() { let postgres = connection("pg", "pg", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options( backend.clone(), McpScope { schema: Some("REPORTING".to_string()), ..Default::default() }, false, ); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "SELECT 1; SELECT 2".to_string(), session_id: None, continue_on_error: None, use_transaction: None, })) .await; assert!(result_text(&result).contains("Statement 1")); let recorded = backend.recorded_arguments.lock().unwrap(); let (_, arguments) = recorded.iter().find(|(name, _)| name == "execute_batch").unwrap(); assert_eq!(arguments["schema"], "REPORTING"); } #[tokio::test] async fn execute_batch_rejects_empty_script() { let postgres = connection("pg", "pg", "postgres", "app"); let server = DbxMcpServer::with_runtime_options( Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }), McpScope::default(), false, ); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: " ".to_string(), session_id: None, continue_on_error: None, use_transaction: None, })) .await; assert!(result_text(&result).contains("SQL_BATCH_EMPTY")); } #[tokio::test] async fn executed_mcp_query_is_saved_with_source_and_failure_status() { let postgres = connection("pg", "Reporting", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let query = |sql: &str| ExecuteQueryRequest { selector: selector("pg"), database: None, sql: sql.to_string(), session_id: None, max_rows: None, cell_char_offset: None, cell_char_limit: None, }; let success = server.execute_query(Parameters(query("SELECT 1"))).await; let failure = server.execute_query(Parameters(query("SELECT 1 /* FAIL_MCP_TEST */"))).await; assert_ne!(success.is_error, Some(true)); assert_eq!(failure.is_error, Some(true)); let history = backend.history.lock().unwrap(); assert_eq!(history.len(), 2); assert_eq!(history[0].connection_id, "pg"); assert_eq!(history[0].connection_name, "Reporting"); assert_eq!(history[0].database, "app"); assert_eq!(history[0].sql, "SELECT 1"); assert!(history[0].success); assert_eq!(history[0].details_json.as_deref(), Some(r#"{"source":"mcp"}"#)); assert!(!history[1].success); assert!(history[1].error.as_deref().unwrap_or_default().contains("query failed")); } #[tokio::test] async fn mcp_batch_failure_is_recorded_once_without_claiming_success() { let postgres = connection("pg", "Reporting", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let sql = "SELECT 1 /* FAIL_MCP_TEST */; SELECT 2"; let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: sql.to_string(), session_id: None, continue_on_error: Some(true), use_transaction: None, })) .await; assert!(result_text(&result).contains("failed")); let history = backend.history.lock().unwrap(); assert_eq!(history.len(), 1); assert_eq!(history[0].sql, sql); assert!(!history[0].success); assert!(history[0].error.as_deref().unwrap_or_default().contains("statement failed")); } #[tokio::test] async fn mcp_batch_error_without_message_is_still_recorded_as_failed() { let postgres = connection("pg", "Reporting", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "SELECT 1 /* FAIL_MCP_NO_MESSAGE_TEST */".to_string(), session_id: None, continue_on_error: None, use_transaction: None, })) .await; assert!(result_text(&result).contains("failed")); let history = backend.history.lock().unwrap(); assert_eq!(history.len(), 1); assert!(!history[0].success); assert_eq!(history[0].error.as_deref(), Some("Statement 1 failed")); } #[tokio::test] async fn preflight_rejection_does_not_enter_execution_history() { let postgres = connection("pg", "Reporting", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: " ".to_string(), session_id: None, continue_on_error: None, use_transaction: None, })) .await; assert_eq!(result.is_error, Some(true)); assert!(backend.history.lock().unwrap().is_empty()); } #[test] fn mcp_history_metadata_handles_comments_and_mixed_batches() { assert_eq!(mcp_sql_operation("-- context\nSELECT 1", DatabaseType::Postgres), "SELECT"); assert_eq!( mcp_sql_activity_kind( "UPDATE accounts SET active = true; CREATE INDEX active_idx ON accounts(active)", DatabaseType::Postgres ), "schema_change" ); } #[tokio::test] async fn execute_batch_rejects_transaction_with_session() { let postgres = connection("pg", "pg", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("pg"), database: None, enable_transactions: false, })) .await; let session_id = opened_session_id(&opened); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "SELECT 1; SELECT 2".to_string(), session_id: Some(session_id), continue_on_error: None, use_transaction: Some(true), })) .await; assert!(result_text(&result).contains("TRANSACTION_WITH_SESSION_UNSUPPORTED")); } #[tokio::test] async fn execute_batch_rejects_transaction_with_continue_on_error() { // The core transaction path rolls back and stops at the first failure // (query.rs:2970), so continue_on_error would be silently ignored. // Reject the combination instead of promising continuation that never // happens, and make sure neither option reaches the backend. let postgres = connection("pg", "pg", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)".to_string(), session_id: None, continue_on_error: Some(true), use_transaction: Some(true), })) .await; assert!(result_text(&result).contains("TRANSACTION_WITH_CONTINUE_UNSUPPORTED")); assert!(backend.recorded_arguments.lock().unwrap().iter().all(|(name, _)| name != "execute_batch")); } #[tokio::test] async fn execute_batch_rejects_mysql_family_ddl_with_transaction() { // MySQL-family engines implicitly commit before DDL, so a DDL batch // with use_transaction cannot roll back a DDL that already committed. // Reject the combination instead of over-promising atomicity. Plain // DML batches and transactional DDL engines (PostgreSQL) still accept // use_transaction. let mysql = connection("my", "my", "mysql", "app"); let backend = Arc::new(FakeBackend { connections: vec![mysql], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let ddl = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("my"), database: None, sql: "CREATE TABLE a (id INT); INSERT INTO missing VALUES (1)".to_string(), session_id: None, continue_on_error: None, use_transaction: Some(true), })) .await; assert!(result_text(&ddl).contains("TRANSACTION_WITH_DDL_UNSUPPORTED")); assert!(backend.recorded_arguments.lock().unwrap().iter().all(|(name, _)| name != "execute_batch")); // GoldenDB is a MySQL-compatible Agent driver, so its DDL has the // same implicit-commit boundary and must not be advertised as atomic. let goldendb = connection("goldendb", "goldendb", "goldendb", "app"); let goldendb_backend = Arc::new(FakeBackend { connections: vec![goldendb], ..Default::default() }); let goldendb_server = DbxMcpServer::with_runtime_options(goldendb_backend.clone(), McpScope::default(), false); let goldendb_ddl = goldendb_server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("goldendb"), database: None, sql: "CREATE TABLE a (id INT); INSERT INTO missing VALUES (1)".to_string(), session_id: None, continue_on_error: None, use_transaction: Some(true), })) .await; assert!(result_text(&goldendb_ddl).contains("TRANSACTION_WITH_DDL_UNSUPPORTED")); assert!(goldendb_backend.recorded_arguments.lock().unwrap().iter().all(|(name, _)| name != "execute_batch")); let dml = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("my"), database: None, sql: "INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)".to_string(), session_id: None, continue_on_error: None, use_transaction: Some(true), })) .await; assert!(result_text(&dml).contains("Transaction outcome")); assert!(backend.recorded_arguments.lock().unwrap().iter().any(|(name, _)| name == "execute_batch")); let postgres = connection("pg", "pg", "postgres", "app"); let postgres_backend = Arc::new(FakeBackend { connections: vec![postgres], policy: McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, ..Default::default() }, ..Default::default() }); let postgres_server = DbxMcpServer::with_runtime_options(postgres_backend.clone(), McpScope::default(), false); let pg_ddl = postgres_server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "CREATE TABLE a (id INT); INSERT INTO t VALUES (1)".to_string(), session_id: None, continue_on_error: None, use_transaction: Some(true), })) .await; assert!(result_text(&pg_ddl).contains("Transaction outcome")); } #[tokio::test] async fn execute_batch_rejects_oracle_ddl_with_transaction() { // Oracle DDL implicitly commits just like MySQL-family engines, so a DDL // batch with use_transaction over-promises atomicity. The rejection must // come from the shared capability check, not a MySQL-only engine list — // otherwise Oracle would silently run a non-rollbackable "transaction". let oracle = connection("ora", "ora", "oracle", "app"); let backend = Arc::new(FakeBackend { connections: vec![oracle], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let ddl = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("ora"), database: None, sql: "CREATE TABLE a (id NUMBER); INSERT INTO missing VALUES (1)".to_string(), session_id: None, continue_on_error: None, use_transaction: Some(true), })) .await; assert!(result_text(&ddl).contains("TRANSACTION_WITH_DDL_UNSUPPORTED")); assert!(backend.recorded_arguments.lock().unwrap().iter().all(|(name, _)| name != "execute_batch")); } #[tokio::test] async fn execute_batch_single_statement_allows_transaction_option_combinations() { // A single-statement script ignores use_transaction and runs as normal // auto-commit, so it never enters transaction mode. Combining it with // session_id or continue_on_error is therefore harmless and must not be // rejected — the option-validity rejects apply only when the script // actually enters transaction mode (statement_count > 1). Contrast with // execute_batch_rejects_transaction_with_session / _continue_on_error, // which use multi-statement scripts and are still rejected. let postgres = connection("pg", "pg", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); let with_continue = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "INSERT INTO t VALUES (1)".to_string(), session_id: None, continue_on_error: Some(true), use_transaction: Some(true), })) .await; assert!( !result_text(&with_continue).contains("TRANSACTION_WITH_CONTINUE_UNSUPPORTED"), "single-statement use_transaction + continue_on_error must not be rejected" ); assert!(backend.recorded_arguments.lock().unwrap().iter().any(|(name, _)| name == "execute_batch")); // A real session bound to the same connection must likewise reach // execution rather than be rejected as a transaction/session conflict. let opened = server .open_session(Parameters(OpenSessionRequest { selector: selector("pg"), database: None, enable_transactions: false, })) .await; let session_id = opened_session_id(&opened); let with_session = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "INSERT INTO t VALUES (1)".to_string(), session_id: Some(session_id), continue_on_error: None, use_transaction: Some(true), })) .await; assert!( !result_text(&with_session).contains("TRANSACTION_WITH_SESSION_UNSUPPORTED"), "single-statement use_transaction + session_id must not be rejected as a transaction/session conflict" ); } #[tokio::test] async fn execute_batch_returns_structured_content() { let postgres = connection("pg", "pg", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "SELECT 1; SELECT 2".to_string(), session_id: None, continue_on_error: None, use_transaction: None, })) .await; // The human-readable block stays in content… assert!(result_text(&result).contains("Statement 1")); // …and structuredContent is a protocol-valid object carrying one result per statement. let structured = result.structured_content.as_ref().expect("structured content must be populated"); assert!(structured.is_object(), "MCP structuredContent must be a JSON object"); let results = structured["results"].as_array().expect("results must be an array"); assert_eq!(results.len(), 1); assert_eq!(results[0]["statement_index"], 0); } #[tokio::test] async fn execute_batch_enforces_confirmed_write_sql_binding() { // `confirmed_batch_sql_block_reason` is exercised directly (not through // the handler with DBX_MCP_CONFIRMED_WRITE_SQL) because the env var is // process-global and would race with `confirmed_sql_binding_cannot_elevate_central_policy`. let postgres_db_type = dbx_core::models::connection::DatabaseType::Postgres; let confirmed = Some("INSERT INTO t VALUES (1)"); // An exact match (modulo surrounding whitespace) passes. assert!(confirmed_batch_sql_block_reason(" INSERT INTO t VALUES (1) ", postgres_db_type, confirmed).is_none()); // A differing write script fails closed. let blocked = confirmed_batch_sql_block_reason( "INSERT INTO t VALUES (1); INSERT INTO u VALUES (2)", postgres_db_type, confirmed, ) .expect("differing write script must be blocked"); assert!(blocked.contains("does not match the user-confirmed SQL")); // A read-only script is never subject to the write confirmation. assert!(confirmed_batch_sql_block_reason("SELECT 1; SELECT 2", postgres_db_type, confirmed).is_none()); // No binding never blocks. assert!(confirmed_batch_sql_block_reason("INSERT INTO t VALUES (1)", postgres_db_type, None).is_none()); // Unparseable SQL fails closed (treated as a write). assert!(confirmed_batch_sql_block_reason("NOT VALID SQL ;;;", postgres_db_type, confirmed).is_some()); } #[test] fn execute_query_routes_scripts_to_the_dialect_aware_batch_path() { assert!(sql_requires_batch_execution("SELECT 1; SELECT 2", DatabaseType::Postgres)); assert!(sql_requires_batch_execution( "DELIMITER $$\nCREATE PROCEDURE p() BEGIN SELECT 1; END$$\nDELIMITER ;", DatabaseType::Mysql )); assert!(!sql_requires_batch_execution("SELECT 1;", DatabaseType::Mysql)); assert!(!sql_requires_batch_execution("SELECT 1", DatabaseType::Mysql)); } #[test] fn format_batch_results_renders_failed_statement() { let results = vec![crate::backend::BatchStatementResult { result: dbx_core::db::QueryResult { columns: vec![], column_types: vec![], column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: 0, execution_time_ms: 0, server_execute_time_us: None, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: vec![], }, execution_error: true, statement_index: Some(2), error_message: Some("syntax error near SELECT".to_string()), merged: false, transaction_state: None, transaction_outcome: None, }]; let output = format_batch_results(&results); assert!(output.contains("### Statement 3")); assert!(output.contains("**Status:** failed")); assert!(output.contains("syntax error near SELECT")); } #[test] fn format_batch_results_renders_transaction_outcome() { let results = vec![crate::backend::BatchStatementResult { result: dbx_core::db::QueryResult { columns: vec![], column_types: vec![], column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: 2, execution_time_ms: 0, server_execute_time_us: None, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: vec![], }, execution_error: false, statement_index: None, error_message: None, merged: true, transaction_state: None, transaction_outcome: None, }]; let output = format_batch_results(&results); assert!(output.contains("### Transaction outcome")); assert!(output.contains("2 row(s) affected")); assert!(!output.contains("Statement 1")); } #[tokio::test] async fn execute_batch_transaction_mode_marks_merged_outcome() { let postgres = connection("pg", "pg", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)".to_string(), session_id: None, continue_on_error: None, use_transaction: Some(true), })) .await; // The merged transaction outcome is labelled distinctly and does not // pretend to be per-statement. assert!(result_text(&result).contains("Transaction outcome")); assert!(!result_text(&result).contains("Statement 1")); let structured = result.structured_content.as_ref().expect("structured content must be populated"); assert_eq!(structured["results"][0]["merged"], true); assert!(structured["results"][0]["statement_index"].is_null()); } #[tokio::test] async fn execute_batch_single_statement_with_transaction_is_not_merged() { // The core transaction wrapper only applies to scripts with more than // one statement (query.rs:2970). A single statement with // use_transaction=true runs as a normal auto-commit statement and must // not be labelled as a transaction outcome. let postgres = connection("pg", "pg", "postgres", "app"); let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("pg"), database: None, sql: "INSERT INTO t VALUES (1)".to_string(), session_id: None, continue_on_error: None, use_transaction: Some(true), })) .await; assert!(result_text(&result).contains("Statement 1")); assert!(!result_text(&result).contains("Transaction outcome")); let structured = result.structured_content.as_ref().expect("structured content must be populated"); // merged=false is skipped by serde, so the single statement must not // carry a merged marker (null/absent), unlike the transaction outcome. assert!(structured["results"][0]["merged"].is_null()); } #[tokio::test] async fn execute_batch_sqlserver_use_transaction_marks_merged_outcome() { // SQL Server enters the core transaction path before its normal batch // fast path, so a successful multi-statement request is merged. let sqlserver = connection("mssql", "mssql", "sqlserver", "app"); let backend = Arc::new(FakeBackend { connections: vec![sqlserver], ..Default::default() }); let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); let result = server .execute_batch(Parameters(ExecuteBatchQueryRequest { selector: selector("mssql"), database: None, sql: "INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)".to_string(), session_id: None, continue_on_error: None, use_transaction: Some(true), })) .await; assert!(result_text(&result).contains("Transaction outcome")); assert!(!result_text(&result).contains("Statement 1")); let structured = result.structured_content.as_ref().expect("structured content must be populated"); assert_eq!(structured["results"][0]["merged"], true); } }