From 1294708eea83efd90620e8a188a2b00ed0034e78 Mon Sep 17 00:00:00 2001 From: xcstatus <40375067+xcstatus@users.noreply.github.com> Date: Sun, 20 Sep 2026 12:23:12 +0800 Subject: [PATCH] feat(mcp): add native MySQL session transactions --- Cargo.lock | 3 + apps/desktop/src/i18n/locales/az.ts | 3 + apps/desktop/src/i18n/locales/en.ts | 3 + apps/desktop/src/i18n/locales/es.ts | 3 + apps/desktop/src/i18n/locales/it.ts | 3 + apps/desktop/src/i18n/locales/ja.ts | 3 + apps/desktop/src/i18n/locales/ko.ts | 3 + apps/desktop/src/i18n/locales/pt-BR.ts | 3 + apps/desktop/src/i18n/locales/tr.ts | 3 + apps/desktop/src/i18n/locales/zh-CN.ts | 3 + apps/desktop/src/i18n/locales/zh-TW.ts | 3 + .../desktop/src/lib/mcp/mcpPolicySelection.ts | 3 + crates/dbx-core/src/connection/mod.rs | 151 ++ crates/dbx-drivers/src/db/mysql.rs | 134 +- crates/dbx-mcp/Cargo.toml | 4 + crates/dbx-mcp/src/backend.rs | 301 ++- crates/dbx-mcp/src/http.rs | 407 +++- crates/dbx-mcp/src/lib.rs | 1 + crates/dbx-mcp/src/server.rs | 1720 ++++++++++++++++- crates/dbx-mcp/src/session.rs | 136 +- crates/dbx-mcp/src/transaction.rs | 1214 ++++++++++++ crates/dbx-mcp/tests/protocol.rs | 7 +- .../tests/session_transactions_mysql.rs | 1188 ++++++++++++ crates/dbx-mcp/tests/support/mod.rs | 1 + .../dbx-mcp/tests/support/mysql_wire_proxy.rs | 386 ++++ docs/content/docs/mcp.cn.mdx | 13 + docs/content/docs/mcp.mdx | 13 + 27 files changed, 5621 insertions(+), 91 deletions(-) create mode 100644 crates/dbx-mcp/src/transaction.rs create mode 100644 crates/dbx-mcp/tests/session_transactions_mysql.rs create mode 100644 crates/dbx-mcp/tests/support/mod.rs create mode 100644 crates/dbx-mcp/tests/support/mysql_wire_proxy.rs diff --git a/Cargo.lock b/Cargo.lock index 22d5bdb7ba..55e29ff161 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2259,6 +2259,7 @@ dependencies = [ "dbx-core", "dirs", "log", + "mysql_async", "reqwest 0.12.28", "rmcp", "rustls 0.23.43", @@ -2266,6 +2267,7 @@ dependencies = [ "schemars 1.2.2", "serde", "serde_json", + "sqlparser", "tempfile", "tokio", "tokio-rustls 0.26.4", @@ -7343,6 +7345,7 @@ dependencies = [ "pastey", "pin-project-lite", "rand 0.10.2", + "reqwest 0.13.4", "rmcp-macros", "schemars 1.2.2", "serde", diff --git a/apps/desktop/src/i18n/locales/az.ts b/apps/desktop/src/i18n/locales/az.ts index 99948b68d7..470a48c419 100644 --- a/apps/desktop/src/i18n/locales/az.ts +++ b/apps/desktop/src/i18n/locales/az.ts @@ -7974,6 +7974,9 @@ export default withEnglishFallback({ mcpToolExecuteQuery: "SQL / Mongo əmrlərini icra et", mcpToolExecuteBatch: "SQL əmrlərini toplu icra et", mcpToolOpenSession: "Sorğu sessiyasını aç", + mcpToolBeginTransaction: "Tranzaksiyanı başlat", + mcpToolCommitTransaction: "Tranzaksiyanı təsdiqlə", + mcpToolRollbackTransaction: "Tranzaksiyanı geri qaytar", mcpToolCloseSession: "Sorğu sessiyasını bağla", mcpToolExecuteRedisCommand: "Redis əmrini icra et", mcpToolPeekMessages: "Kafka mesajlarını oxu", diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 0ba32b0a5c..147db36ec1 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -8240,6 +8240,9 @@ export default { mcpToolExecuteQuery: "Execute SQL / Mongo commands", mcpToolExecuteBatch: "Execute SQL batch", mcpToolOpenSession: "Open query session", + mcpToolBeginTransaction: "Begin transaction", + mcpToolCommitTransaction: "Commit transaction", + mcpToolRollbackTransaction: "Roll back transaction", mcpToolCloseSession: "Close query session", mcpToolExecuteRedisCommand: "Execute Redis command", mcpToolPeekMessages: "Read Kafka messages", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index 2304a4aba4..83aa3d8003 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -7734,6 +7734,9 @@ export default withEnglishFallback({ mcpToolExecuteQuery: "Ejecutar SQL / comandos Mongo", mcpToolExecuteBatch: "Ejecutar lote SQL", mcpToolOpenSession: "Abrir sesión de consulta", + mcpToolBeginTransaction: "Iniciar transacción", + mcpToolCommitTransaction: "Confirmar transacción", + mcpToolRollbackTransaction: "Revertir transacción", mcpToolCloseSession: "Cerrar sesión de consulta", mcpToolExecuteRedisCommand: "Ejecutar comando Redis", mcpToolPeekMessages: "Leer mensajes de Kafka", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index 150e1a72bf..b1513f4fb0 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -7733,6 +7733,9 @@ export default withEnglishFallback({ mcpToolExecuteQuery: "Esegui SQL / comandi Mongo", mcpToolExecuteBatch: "Esegui batch SQL", mcpToolOpenSession: "Apri sessione di query", + mcpToolBeginTransaction: "Avvia transazione", + mcpToolCommitTransaction: "Conferma transazione", + mcpToolRollbackTransaction: "Annulla transazione", mcpToolCloseSession: "Chiudi sessione di query", mcpToolExecuteRedisCommand: "Esegui comando Redis", mcpToolPeekMessages: "Leggi messaggi Kafka", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index 10ea900031..a54270dddf 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -7738,6 +7738,9 @@ export default withEnglishFallback({ mcpToolExecuteQuery: "SQL / Mongo コマンドを実行", mcpToolExecuteBatch: "SQL バッチを実行", mcpToolOpenSession: "クエリセッションを開く", + mcpToolBeginTransaction: "トランザクションを開始", + mcpToolCommitTransaction: "トランザクションをコミット", + mcpToolRollbackTransaction: "トランザクションをロールバック", mcpToolCloseSession: "クエリセッションを閉じる", mcpToolExecuteRedisCommand: "Redis コマンドを実行", mcpToolPeekMessages: "Kafka メッセージを読み取る", diff --git a/apps/desktop/src/i18n/locales/ko.ts b/apps/desktop/src/i18n/locales/ko.ts index 67504fed51..59f3fc168c 100644 --- a/apps/desktop/src/i18n/locales/ko.ts +++ b/apps/desktop/src/i18n/locales/ko.ts @@ -7503,6 +7503,9 @@ export default withEnglishFallback({ mcpToolExecuteQuery: "SQL / Mongo 명령 실행", mcpToolExecuteBatch: "SQL 배치 실행", mcpToolOpenSession: "쿼리 세션 열기", + mcpToolBeginTransaction: "트랜잭션 시작", + mcpToolCommitTransaction: "트랜잭션 커밋", + mcpToolRollbackTransaction: "트랜잭션 롤백", mcpToolCloseSession: "쿼리 세션 닫기", mcpToolExecuteRedisCommand: "Redis 명령 실행", mcpToolPeekMessages: "Kafka 메시지 읽기", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index e58b50aff6..3a27815804 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -7735,6 +7735,9 @@ export default withEnglishFallback({ mcpToolExecuteQuery: "Executar SQL / comandos Mongo", mcpToolExecuteBatch: "Executar lote SQL", mcpToolOpenSession: "Abrir sessão de consulta", + mcpToolBeginTransaction: "Iniciar transação", + mcpToolCommitTransaction: "Confirmar transação", + mcpToolRollbackTransaction: "Reverter transação", mcpToolCloseSession: "Fechar sessão de consulta", mcpToolExecuteRedisCommand: "Executar comando Redis", mcpToolPeekMessages: "Ler mensagens Kafka", diff --git a/apps/desktop/src/i18n/locales/tr.ts b/apps/desktop/src/i18n/locales/tr.ts index 1ae17b7c50..f4e0c684b0 100644 --- a/apps/desktop/src/i18n/locales/tr.ts +++ b/apps/desktop/src/i18n/locales/tr.ts @@ -7872,6 +7872,9 @@ export default withEnglishFallback({ mcpToolExecuteQuery: "SQL / Mongo komutlarını çalıştır", mcpToolExecuteBatch: "SQL toplu işlemini çalıştır", mcpToolOpenSession: "Sorgu oturumu aç", + mcpToolBeginTransaction: "İşlemi başlat", + mcpToolCommitTransaction: "İşlemi onayla", + mcpToolRollbackTransaction: "İşlemi geri al", mcpToolCloseSession: "Sorgu oturumunu kapat", mcpToolExecuteRedisCommand: "Redis komutunu çalıştır", mcpToolPeekMessages: "Kafka mesajlarını oku", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index 094b2af83b..26457ff851 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -8215,6 +8215,9 @@ export default withEnglishFallback({ mcpToolExecuteQuery: "执行 SQL / Mongo 命令", mcpToolExecuteBatch: "批量执行 SQL", mcpToolOpenSession: "打开查询会话", + mcpToolBeginTransaction: "开始事务", + mcpToolCommitTransaction: "提交事务", + mcpToolRollbackTransaction: "回滚事务", mcpToolCloseSession: "关闭查询会话", mcpToolExecuteRedisCommand: "执行 Redis 命令", mcpToolPeekMessages: "读取 Kafka 消息", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index 2d21976393..85740b8610 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -7104,6 +7104,9 @@ export default withEnglishFallback({ mcpToolExecuteQuery: "執行 SQL / Mongo 命令", mcpToolExecuteBatch: "批次執行 SQL", mcpToolOpenSession: "開啟查詢工作階段", + mcpToolBeginTransaction: "開始交易", + mcpToolCommitTransaction: "提交交易", + mcpToolRollbackTransaction: "回滾交易", mcpToolCloseSession: "關閉查詢工作階段", mcpToolExecuteRedisCommand: "執行 Redis 命令", mcpToolPeekMessages: "讀取 Kafka 訊息", diff --git a/apps/desktop/src/lib/mcp/mcpPolicySelection.ts b/apps/desktop/src/lib/mcp/mcpPolicySelection.ts index b2cb830f20..e6dcc75c89 100644 --- a/apps/desktop/src/lib/mcp/mcpPolicySelection.ts +++ b/apps/desktop/src/lib/mcp/mcpPolicySelection.ts @@ -32,6 +32,9 @@ export const MCP_TOOL_OPTIONS = [ { name: "dbx_execute_query", labelKey: "settings.mcpToolExecuteQuery" }, { name: "dbx_execute_batch", labelKey: "settings.mcpToolExecuteBatch" }, { name: "dbx_open_session", labelKey: "settings.mcpToolOpenSession" }, + { name: "dbx_begin_transaction", labelKey: "settings.mcpToolBeginTransaction" }, + { name: "dbx_commit_transaction", labelKey: "settings.mcpToolCommitTransaction" }, + { name: "dbx_rollback_transaction", labelKey: "settings.mcpToolRollbackTransaction" }, { name: "dbx_close_session", labelKey: "settings.mcpToolCloseSession" }, { name: "dbx_execute_redis_command", labelKey: "settings.mcpToolExecuteRedisCommand" }, { name: "dbx_peek_messages", labelKey: "settings.mcpToolPeekMessages" }, diff --git a/crates/dbx-core/src/connection/mod.rs b/crates/dbx-core/src/connection/mod.rs index 658cf8310a..1c111fa133 100644 --- a/crates/dbx-core/src/connection/mod.rs +++ b/crates/dbx-core/src/connection/mod.rs @@ -11,6 +11,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::{watch, Mutex, OwnedSemaphorePermit, RwLock, Semaphore}; +use tokio_util::sync::CancellationToken; use mysql_async::prelude::Queryable; use mysql_async::Row as MysqlRow; @@ -337,12 +338,36 @@ macro_rules! agent_connection_pool_database_type { }; } +#[derive(Clone)] +pub struct ConnectionLifecycleSnapshot { + generation: u64, + cancellation: CancellationToken, +} + +impl ConnectionLifecycleSnapshot { + pub fn cancellation(&self) -> &CancellationToken { + &self.cancellation + } +} + +struct ConnectionLifecycle { + generation: u64, + cancellation: CancellationToken, +} + +struct SharedResourceBudget { + capacity: usize, + semaphore: Arc, +} + pub struct AppState { connections: Arc>, task_supervisor: TaskSupervisor, pool_activity: Arc>>, draining_pools: Arc>>>, connection_attempts: RwLock>, + connection_lifecycles: std::sync::Mutex>, + shared_resource_budgets: std::sync::Mutex>, pub configs: RwLock>, pub running_queries: RunningQueries, pub tunnels: TunnelManager, @@ -1244,6 +1269,68 @@ fn mysql_metadata_fallback_url( } impl AppState { + pub fn shared_resource_budget(&self, name: &str, capacity: usize) -> Result, String> { + if capacity == 0 { + return Err("Shared resource budget capacity must be greater than zero".to_string()); + } + let mut budgets = self.shared_resource_budgets.lock().unwrap_or_else(|error| error.into_inner()); + if let Some(budget) = budgets.get(name) { + if budget.capacity != capacity { + return Err(format!( + "Shared resource budget {name:?} already has capacity {}, not {capacity}", + budget.capacity + )); + } + return Ok(budget.semaphore.clone()); + } + let semaphore = Arc::new(Semaphore::new(capacity)); + budgets.insert(name.to_string(), SharedResourceBudget { capacity, semaphore: semaphore.clone() }); + Ok(semaphore) + } + + pub fn connection_lifecycle_snapshot(&self, connection_id: &str) -> ConnectionLifecycleSnapshot { + let mut lifecycles = self.connection_lifecycles.lock().unwrap_or_else(|error| error.into_inner()); + let lifecycle = lifecycles + .entry(connection_id.to_string()) + .or_insert_with(|| ConnectionLifecycle { generation: 0, cancellation: CancellationToken::new() }); + ConnectionLifecycleSnapshot { generation: lifecycle.generation, cancellation: lifecycle.cancellation.clone() } + } + + pub fn connection_lifecycle_is_current(&self, connection_id: &str, snapshot: &ConnectionLifecycleSnapshot) -> bool { + self.connection_lifecycles.lock().unwrap_or_else(|error| error.into_inner()).get(connection_id).is_some_and( + |lifecycle| lifecycle.generation == snapshot.generation && !snapshot.cancellation.is_cancelled(), + ) + } + + pub fn invalidate_connection_lifecycle(&self, connection_id: &str) { + let previous = { + let mut lifecycles = self.connection_lifecycles.lock().unwrap_or_else(|error| error.into_inner()); + let lifecycle = lifecycles + .entry(connection_id.to_string()) + .or_insert_with(|| ConnectionLifecycle { generation: 0, cancellation: CancellationToken::new() }); + let previous = std::mem::replace(&mut lifecycle.cancellation, CancellationToken::new()); + lifecycle.generation = lifecycle.generation.wrapping_add(1); + previous + }; + previous.cancel(); + } + + fn invalidate_all_connection_lifecycles(&self) { + let previous = { + let mut lifecycles = self.connection_lifecycles.lock().unwrap_or_else(|error| error.into_inner()); + lifecycles + .values_mut() + .map(|lifecycle| { + lifecycle.generation = lifecycle.generation.wrapping_add(1); + std::mem::replace(&mut lifecycle.cancellation, CancellationToken::new()) + }) + .collect::>() + }; + for cancellation in previous { + cancellation.cancel(); + } + } + /// Return an owned pool handle. The registry read lock is released before /// the caller can perform any asynchronous database operation. pub async fn pool_handle(&self, pool_key: &str) -> Option { @@ -1398,6 +1485,8 @@ impl AppState { pool_activity: Arc::new(RwLock::new(HashMap::new())), draining_pools: Arc::new(std::sync::Mutex::new(HashMap::new())), connection_attempts: RwLock::new(HashMap::new()), + connection_lifecycles: std::sync::Mutex::new(HashMap::new()), + shared_resource_budgets: std::sync::Mutex::new(HashMap::new()), configs: RwLock::new(HashMap::new()), running_queries: RunningQueries::default(), tunnels: TunnelManager::new(data_dir), @@ -2064,6 +2153,7 @@ impl AppState { } pub async fn shutdown(&self, deadline: Duration) { + self.invalidate_all_connection_lifecycles(); self.running_queries.cancel_all(); let removed_pools = self.drain_all_connection_pools().await; self.transaction_sessions.write().await.clear(); @@ -5251,6 +5341,7 @@ impl AppState { } pub async fn remove_connection_pools(&self, connection_id: &str) { + self.invalidate_connection_lifecycle(connection_id); self.rollback_manual_transaction_sessions(connection_id).await; let removed = self.drain_connection_pools(connection_id).await; self.clear_metadata_gates_for_connection(connection_id).await; @@ -5266,6 +5357,7 @@ impl AppState { /// user disconnect still goes through remove_connection_pools* and does /// send connection/disconnect. pub async fn drop_connection_pools_without_close(&self, connection_id: &str) { + self.invalidate_connection_lifecycle(connection_id); self.rollback_manual_transaction_sessions(connection_id).await; let removed = self.drain_connection_pools(connection_id).await; self.clear_metadata_gates_for_connection(connection_id).await; @@ -5273,6 +5365,7 @@ impl AppState { } pub async fn remove_connection_pools_detached(&self, connection_id: &str) { + self.invalidate_connection_lifecycle(connection_id); self.rollback_manual_transaction_sessions(connection_id).await; let removed = self.drain_connection_pools(connection_id).await; self.clear_metadata_gates_for_connection(connection_id).await; @@ -7387,6 +7480,64 @@ mod tests { (AppState::new(storage), dir) } + #[tokio::test] + async fn connection_lifecycle_invalidation_rotates_generation_and_cancels_previous_snapshot() { + let (state, dir) = test_app_state().await; + let first = state.connection_lifecycle_snapshot("conn"); + assert!(state.connection_lifecycle_is_current("conn", &first)); + + state.invalidate_connection_lifecycle("conn"); + + tokio::time::timeout(Duration::from_millis(100), first.cancellation().cancelled()) + .await + .expect("previous lifecycle must be cancelled"); + assert!(!state.connection_lifecycle_is_current("conn", &first)); + let second = state.connection_lifecycle_snapshot("conn"); + assert!(state.connection_lifecycle_is_current("conn", &second)); + assert!(!second.cancellation().is_cancelled()); + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn every_connection_pool_removal_boundary_invalidates_lifecycle_snapshots() { + let (state, dir) = test_app_state().await; + + let removed = state.connection_lifecycle_snapshot("removed"); + state.remove_connection_pools("removed").await; + assert!(removed.cancellation().is_cancelled()); + + let dropped = state.connection_lifecycle_snapshot("dropped"); + state.drop_connection_pools_without_close("dropped").await; + assert!(dropped.cancellation().is_cancelled()); + + let detached = state.connection_lifecycle_snapshot("detached"); + state.remove_connection_pools_detached("detached").await; + assert!(detached.cancellation().is_cancelled()); + + let shutdown = state.connection_lifecycle_snapshot("shutdown"); + state.shutdown(Duration::from_millis(100)).await; + assert!(shutdown.cancellation().is_cancelled()); + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn named_resource_budget_is_shared_across_app_state_views() { + let (state, dir) = test_app_state().await; + let first = state.shared_resource_budget("fixed-owner", 2).unwrap(); + let second = state.shared_resource_budget("fixed-owner", 2).unwrap(); + assert!(Arc::ptr_eq(&first, &second)); + + let first_permit = first.clone().try_acquire_owned().unwrap(); + let second_permit = second.clone().try_acquire_owned().unwrap(); + assert!(first.clone().try_acquire_owned().is_err()); + drop(first_permit); + assert!(second.clone().try_acquire_owned().is_ok()); + drop(second_permit); + + assert!(state.shared_resource_budget("fixed-owner", 3).is_err()); + let _ = std::fs::remove_dir_all(dir); + } + fn agent_pool_stub() -> PoolKind { PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub()) } diff --git a/crates/dbx-drivers/src/db/mysql.rs b/crates/dbx-drivers/src/db/mysql.rs index 463b3b574e..63a07a759e 100644 --- a/crates/dbx-drivers/src/db/mysql.rs +++ b/crates/dbx-drivers/src/db/mysql.rs @@ -1,6 +1,6 @@ use chrono::{NaiveDate, NaiveDateTime, NaiveTime}; use futures::StreamExt; -use mysql_async::consts::{ColumnFlags, ColumnType}; +use mysql_async::consts::{ColumnFlags, ColumnType, StatusFlags}; use mysql_async::prelude::*; use percent_encoding::percent_decode_str; use rust_decimal::Decimal; @@ -191,6 +191,47 @@ pub struct MySqlQueryResult { pub large_value_cells: Vec, } +/// Result of one MCP fixed-session statement executed through MySQL's text +/// protocol. The transaction bit comes from the final server OK/EOF packet, +/// never from SQL text or client-side state inference. +#[derive(Debug)] +pub struct MySqlTransactionExecution { + pub result: QueryResult, + pub in_transaction: bool, +} + +/// Error classification retained by the fixed-session owner. Server errors +/// are recoverable protocol responses; every other driver error means the +/// connection state cannot be trusted or replayed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum MySqlTransactionError { + Server { code: u16, message: String, state: String }, + Transport(String), +} + +fn transaction_error_from_mysql_error(error: mysql_async::Error) -> MySqlTransactionError { + match error { + mysql_async::Error::Server(error) => { + MySqlTransactionError::Server { code: error.code, message: error.message, state: error.state } + } + error => MySqlTransactionError::Transport(error.to_string()), + } +} + +fn transaction_status_from_last_ok(conn: &mysql_async::Conn) -> Result { + conn.last_ok_packet() + .map(|packet| packet.status_flags().contains(StatusFlags::SERVER_STATUS_IN_TRANS)) + .ok_or_else(|| MySqlTransactionError::Transport("MySQL did not return a final status packet".to_string())) +} + +/// Read transaction state after a recoverable server ERR packet. mysql_async +/// intentionally clears the cached OK packet for ERR, so COM_PING is required +/// to obtain fresh status from the same physical connection. +pub async fn ping_transaction_status_on_conn(conn: &mut mysql_async::Conn) -> Result { + conn.ping().await.map_err(transaction_error_from_mysql_error)?; + transaction_status_from_last_ok(conn) +} + impl MySqlQueryResult { fn exact(result: QueryResult) -> Self { Self { result, large_value_cells: Vec::new() } @@ -5274,6 +5315,79 @@ pub async fn execute_query_on_conn_with_max_rows( .map(|result| result.result) } +/// Execute exactly one policy-approved MCP transaction statement using +/// COM_QUERY/text protocol. This path never retries, rewrites or replays user +/// SQL. It fully drains the result before exposing the final server status. +pub async fn execute_transaction_statement_on_conn( + conn: &mut mysql_async::Conn, + sql: &str, + max_rows: Option, +) -> Result { + let started_at = Instant::now(); + let row_limit = query_result_row_limit(max_rows); + let mut query = conn.query_iter(sql).await.map_err(transaction_error_from_mysql_error)?; + let columns: Vec = query.columns_ref().iter().map(|column| column.name_str().to_string()).collect(); + let column_types: Vec = query.columns_ref().iter().map(mysql_column_type_name).collect(); + + let result = if columns.is_empty() { + let affected_rows = query.affected_rows(); + let warnings = query.warnings(); + let info = query.info().into_owned(); + query.drop_result().await.map_err(transaction_error_from_mysql_error)?; + let messages = collect_mysql_server_messages(conn, warnings, &info).await; + QueryResult { + columns, + column_types, + column_sortables: Vec::new(), + spatial_columns: Vec::new(), + spatial_values: Vec::new(), + rows: Vec::new(), + affected_rows, + execution_time_ms: started_at.elapsed().as_millis(), + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages, + } + } else { + let mut spatial_columns = mysql_spatial_column_builder(query.columns_ref()); + let mut rows = Vec::with_capacity(row_limit.min(128)); + let mut spatial_values = Vec::new(); + let mut truncated = false; + while let Some(row) = query.next().await.map_err(transaction_error_from_mysql_error)? { + if rows.len() < row_limit { + let (values, srids) = mysql_row_to_json_with_srids(&row, &mut spatial_columns); + spatial_values.push(srids); + rows.push(values); + } else { + truncated = true; + } + } + // `next` stops at the first result-set boundary. Drain every remaining + // result set before reading the final status or allowing another + // operation on this physical connection. + query.drop_result().await.map_err(transaction_error_from_mysql_error)?; + QueryResult { + columns, + column_types, + column_sortables: Vec::new(), + spatial_columns: spatial_columns.finish(), + spatial_values, + rows, + affected_rows: 0, + execution_time_ms: started_at.elapsed().as_millis(), + truncated, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages: Vec::new(), + } + }; + let in_transaction = transaction_status_from_last_ok(conn)?; + Ok(MySqlTransactionExecution { result, in_transaction }) +} + pub async fn execute_query_on_conn_with_limits( conn: &mut mysql_async::Conn, sql: &str, @@ -6060,6 +6174,24 @@ mod tests { use mysql_async::{consts::ColumnType, Column, Value}; use mysql_common::row::new_row; + #[test] + fn transaction_error_preserves_server_code_without_treating_it_as_transport_failure() { + let error = mysql_async::Error::Server(mysql_async::ServerError { + code: 1062, + message: "Duplicate entry".to_string(), + state: "23000".to_string(), + }); + + assert_eq!( + transaction_error_from_mysql_error(error), + MySqlTransactionError::Server { + code: 1062, + message: "Duplicate entry".to_string(), + state: "23000".to_string(), + } + ); + } + fn mysql_test_row(values: Vec) -> mysql_async::Row { let columns = values .iter() diff --git a/crates/dbx-mcp/Cargo.toml b/crates/dbx-mcp/Cargo.toml index 26a8b44c08..dc83a7d981 100644 --- a/crates/dbx-mcp/Cargo.toml +++ b/crates/dbx-mcp/Cargo.toml @@ -9,6 +9,7 @@ default = ["duckdb-sidecar", "mq-admin", "sqlite-bundled"] duckdb-sidecar = ["dbx-core/duckdb-sidecar"] mq-admin = ["dbx-core/mq-admin"] sqlite-bundled = ["dbx-core/sqlite-bundled"] +transaction-test-hooks = [] [[bin]] name = "dbx-mcp" @@ -19,6 +20,7 @@ async-trait = "0.1" dbx-core = { path = "../dbx-core", default-features = false } dirs = "6" log = "0.4" +mysql_async = { version = "0.37", default-features = false, features = ["default-rustls", "client_ed25519", "chrono", "rust_decimal"] } axum = "0.8" chrono = "0.4" rmcp = { version = "2.2.0", features = ["client", "transport-io", "transport-streamable-http-server"] } @@ -27,6 +29,7 @@ rustls = { version = "0.23", features = ["aws-lc-rs"] } schemars = "1" serde = { version = "1", features = ["derive"] } serde_json = "1" +sqlparser = { version = "0.62", features = ["visitor"] } tokio = { version = "1", features = ["fs", "macros", "net", "rt-multi-thread", "signal"] } tokio-util = "0.7" tower-http = { version = "0.6", features = ["cors"] } @@ -34,6 +37,7 @@ uuid = { version = "1", features = ["v4"] } url = "2" [dev-dependencies] +rmcp = { version = "2.2.0", features = ["transport-streamable-http-client-reqwest"] } rustls-pemfile = "2.2" tempfile = "3" tokio = { version = "1", features = ["net", "io-util"] } diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index b6ff7b79d2..a02d0785b1 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -8,7 +8,7 @@ use async_trait::async_trait; use dbx_core::{ agent_events::{ToolCall, ToolResult}, agent_tools::{self, format_query_result_as_text, AgentSqlPermissions, QueryCellWindow}, - connection::{connection_configs_pool_equivalent, AppState}, + connection::{connection_configs_pool_equivalent, AppState, ConnectionLifecycleSnapshot}, db::{mongo_driver::MongoIndexSpec, redis_driver::RedisCommandResult, ColumnInfo, TableInfo}, history::HistoryEntry, mcp_policy::{connection_group_paths, McpConnectionGroupPath}, @@ -21,7 +21,10 @@ use tokio::sync::Mutex; use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; -use crate::mongo::MongoCommand; +use crate::{ + mongo::MongoCommand, + transaction::{MysqlTransactionIo, TransactionOwner, TransactionOwnerConfig, TransactionOwnerRegistry}, +}; #[derive(Debug, Clone, Serialize, PartialEq, Eq)] pub struct ConnectionSummary { @@ -119,12 +122,34 @@ pub struct BatchStatementResult { /// false. #[serde(skip_serializing_if = "is_false")] pub merged: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub transaction_state: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub transaction_outcome: Option, } fn is_false(value: &bool) -> bool { !*value } +pub(crate) fn native_mysql_transaction_connection(connection: &ConnectionConfig) -> bool { + connection.db_type == DatabaseType::Mysql + && connection.driver_profile.as_deref().is_none_or(|profile| { + let profile = profile.trim(); + profile.is_empty() || profile.eq_ignore_ascii_case("mysql") + }) +} + +fn transaction_operation_timeout( + connection: &ConnectionConfig, + configured: std::time::Duration, +) -> std::time::Duration { + match connection.effective_query_timeout_secs() { + 0 => configured, + seconds => std::time::Duration::from_secs(seconds), + } +} + impl From for BatchStatementResult { fn from(result: dbx_core::query::ExecuteMultiResult) -> Self { let error_message = result.error.as_ref().and_then(|error| error.detail().map(str::to_owned)).or_else(|| { @@ -136,6 +161,8 @@ impl From for BatchStatementResult { statement_index: result.statement_index, error_message, merged: false, + transaction_state: None, + transaction_outcome: None, } } } @@ -160,7 +187,15 @@ fn batch_statement_result_from_json(value: &Value) -> Result Result, String> { + let _ = (connection, database, client_session_id); + Err("TRANSACTION_UNSUPPORTED: fixed-session transactions require a native local MySQL connection.".to_string()) + } /// Execute a multi-statement SQL script, returning one result per statement. /// /// The script text is split using the database-dialect-aware splitter so @@ -374,6 +418,8 @@ pub trait DbxBackend: Send + Sync { pub struct LocalBackend { state: Arc, data_dir: std::path::PathBuf, + transaction_owners: Arc, + transaction_owner_config: TransactionOwnerConfig, } #[derive(Debug, Default)] @@ -391,6 +437,12 @@ pub struct WebBackend { connected: Mutex>, } +impl Drop for LocalBackend { + fn drop(&mut self) { + self.transaction_owners.invalidate_all(); + } +} + // Manual impl: the derived one would print `password` (and the session cookie // inside `auth`) in plaintext. Redact credentials and skip lockable state. impl std::fmt::Debug for WebBackend { @@ -580,10 +632,51 @@ impl WebBackend { } impl LocalBackend { + fn spawn_connection_lifecycle_watcher( + &self, + lifecycle: ConnectionLifecycleSnapshot, + owner: &Arc, + ) -> tokio::task::JoinHandle<()> { + let cancellation = lifecycle.cancellation().clone(); + let owner_closed = owner.closed_token(); + let watched_owner = Arc::downgrade(owner); + tokio::spawn(async move { + tokio::select! { + _ = cancellation.cancelled() => { + if let Some(owner) = watched_owner.upgrade() { + owner.invalidate(); + } + } + _ = owner_closed.cancelled() => {} + } + }) + } + + async fn finalize_transaction_owner( + &self, + connection_id: &str, + lifecycle: ConnectionLifecycleSnapshot, + owner: Arc, + ) -> Result, String> { + self.spawn_connection_lifecycle_watcher(lifecycle.clone(), &owner); + self.transaction_owners.register(connection_id, &owner); + if !self.state.connection_lifecycle_is_current(connection_id, &lifecycle) { + owner.invalidate(); + let _ = owner.close().await; + return Err("The connection was removed while its fixed transaction session was opening.".to_string()); + } + Ok(owner) + } + /// Reuse an already initialized DBX application state, for hosts that /// embed MCP alongside their own HTTP server. pub fn from_app_state(state: Arc, data_dir: PathBuf) -> Self { - Self { state, data_dir } + Self { + state, + data_dir, + transaction_owners: Arc::new(TransactionOwnerRegistry::default()), + transaction_owner_config: TransactionOwnerConfig::default(), + } } pub async fn open(path: &Path) -> Result { @@ -613,7 +706,18 @@ impl LocalBackend { let config_map: HashMap = configs.into_iter().map(|config| (config.id.clone(), config)).collect(); *state.configs.write().await = config_map; - Ok(Self { state, data_dir }) + Ok(Self { + state, + data_dir, + transaction_owners: Arc::new(TransactionOwnerRegistry::default()), + transaction_owner_config: TransactionOwnerConfig::default(), + }) + } + + #[cfg(feature = "transaction-test-hooks")] + pub fn with_transaction_test_config(mut self, config: TransactionOwnerConfig) -> Self { + self.transaction_owner_config = config; + self } pub fn state(&self) -> &Arc { @@ -711,6 +815,7 @@ impl LocalBackend { }; for id in pool_ids_to_drop { + self.transaction_owners.invalidate_connection(&id); self.state.remove_connection_pools_detached(&id).await; } } @@ -875,6 +980,57 @@ impl DbxBackend for LocalBackend { .await } + async fn open_transaction_owner( + &self, + connection: &ConnectionConfig, + database: &str, + client_session_id: &str, + ) -> Result, String> { + if !native_mysql_transaction_connection(connection) { + return Err("TRANSACTION_UNSUPPORTED: fixed-session transactions require a native local MySQL connection." + .to_string()); + } + let lifecycle = self.state.connection_lifecycle_snapshot(&connection.id); + let resource_permit = + self.state.shared_resource_budget("dbx-mcp-transaction-owners", 32)?.try_acquire_owned().map_err(|_| { + "Too many open transaction-enabled MCP sessions (max 32). Close an existing transaction session first." + .to_string() + })?; + let pool_database = (!database.trim().is_empty()).then_some(database); + let pool_key = + self.state.get_or_create_pool_for_session(&connection.id, pool_database, Some(client_session_id)).await?; + let pool = match self.state.pool_handle(&pool_key).await { + Some(dbx_core::connection::PoolKind::Mysql(pool, dbx_core::connection::MysqlMode::Normal)) => pool, + _ => { + let _ = self.state.close_client_session_pool(&connection.id, pool_database, client_session_id).await; + return Err( + "TRANSACTION_UNSUPPORTED: fixed-session transactions require the native MySQL driver.".to_string() + ); + } + }; + let conn = match dbx_core::db::mysql::get_conn_with_health_check(&pool).await { + Ok(conn) => conn, + Err(error) => { + let _ = self.state.close_client_session_pool(&connection.id, pool_database, client_session_id).await; + return Err(format!("Failed to acquire the fixed MySQL session connection: {error}")); + } + }; + let mut owner_config = self.transaction_owner_config; + owner_config.operation_timeout = transaction_operation_timeout(connection, owner_config.operation_timeout); + let owner = TransactionOwner::spawn_with_resource_permit( + MysqlTransactionIo::new( + conn, + self.state.clone(), + connection.id.clone(), + database.to_string(), + client_session_id.to_string(), + ), + owner_config, + resource_permit, + ); + self.finalize_transaction_owner(&connection.id, lifecycle, owner).await + } + async fn execute_query( &self, connection: &ConnectionConfig, @@ -952,6 +1108,7 @@ impl DbxBackend for LocalBackend { let removed = self.state.storage.remove_connection_for_mcp(connection_id).await?; if removed { self.state.configs.write().await.remove(connection_id); + self.transaction_owners.invalidate_connection(connection_id); self.state.remove_connection_pools_detached(connection_id).await; } Ok(removed) @@ -2464,7 +2621,10 @@ mod tests { use std::{ io::{Read, Write}, net::TcpListener, - sync::mpsc, + sync::{ + atomic::{AtomicUsize, Ordering}, + mpsc, + }, time::Duration, }; @@ -2505,6 +2665,50 @@ mod tests { assert!(parse_database_type("unknown").is_err()); } + #[test] + fn native_mysql_transaction_gate_accepts_the_builtin_mysql_profile_only() { + let mut connection: ConnectionConfig = serde_json::from_value(json!({ + "id": "mysql", + "name": "mysql", + "db_type": "mysql", + "host": "", + "port": 3306, + "username": "", + "password": "", + "database": "test", + "ssl": false + })) + .unwrap(); + assert!(native_mysql_transaction_connection(&connection)); + connection.driver_profile = Some("mysql".to_string()); + assert!(native_mysql_transaction_connection(&connection)); + connection.driver_profile = Some("MySQL".to_string()); + assert!(native_mysql_transaction_connection(&connection)); + for profile in ["dolt", "tidb", "oceanbase", "custom-profile"] { + connection.driver_profile = Some(profile.to_string()); + assert!(!native_mysql_transaction_connection(&connection), "unexpectedly accepted {profile}"); + } + } + + #[test] + fn unlimited_connection_timeout_preserves_bounded_transaction_default() { + let mut connection: ConnectionConfig = serde_json::from_value(json!({ + "id": "mysql", + "name": "mysql", + "db_type": "mysql", + "host": "", + "port": 3306, + "username": "", + "password": "", + "database": "test", + "ssl": false + })) + .unwrap(); + connection.query_timeout_secs = 0; + + assert_eq!(transaction_operation_timeout(&connection, Duration::from_secs(300)), Duration::from_secs(300)); + } + #[test] fn parses_nested_current_and_legacy_connection_group_paths() { let layout = json!({ @@ -3833,6 +4037,91 @@ mod tests { assert_eq!(backend.state().agent_manager.base_dir(), &agent_dir); } + struct LifecycleTestIo { + disconnects: Arc, + } + + #[async_trait] + impl crate::transaction::TransactionIo for LifecycleTestIo { + async fn execute( + &mut self, + _sql: &str, + _max_rows: Option, + ) -> Result { + panic!("lifecycle test owner must stay idle") + } + + async fn ping_in_transaction(&mut self) -> Result { + panic!("lifecycle test owner must not probe") + } + + async fn disconnect(&mut self) { + self.disconnects.fetch_add(1, Ordering::SeqCst); + } + } + + #[tokio::test] + async fn local_transaction_owner_rejects_stale_connection_lifecycle_before_registration() { + let data_dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&data_dir.path().join("dbx.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let backend = LocalBackend::from_app_state(state.clone(), data_dir.path().to_path_buf()); + let lifecycle = state.connection_lifecycle_snapshot("conn"); + state.invalidate_connection_lifecycle("conn"); + let disconnects = Arc::new(AtomicUsize::new(0)); + let owner = TransactionOwner::spawn( + LifecycleTestIo { disconnects: disconnects.clone() }, + TransactionOwnerConfig::default(), + ); + + let error = backend.finalize_transaction_owner("conn", lifecycle, owner.clone()).await.unwrap_err(); + + assert!(error.contains("removed"), "{error}"); + owner.wait_closed().await; + assert_eq!(disconnects.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn shared_app_state_removal_cancels_a_registered_local_transaction_owner() { + let data_dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&data_dir.path().join("dbx.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let backend = LocalBackend::from_app_state(state.clone(), data_dir.path().to_path_buf()); + let lifecycle = state.connection_lifecycle_snapshot("conn"); + let disconnects = Arc::new(AtomicUsize::new(0)); + let owner = TransactionOwner::spawn( + LifecycleTestIo { disconnects: disconnects.clone() }, + TransactionOwnerConfig::default(), + ); + backend.finalize_transaction_owner("conn", lifecycle, owner.clone()).await.unwrap(); + + state.remove_connection_pools_detached("conn").await; + owner.wait_closed().await; + + assert_eq!(disconnects.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn connection_lifecycle_watcher_exits_after_normal_owner_disposal() { + let data_dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&data_dir.path().join("dbx.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let backend = LocalBackend::from_app_state(state.clone(), data_dir.path().to_path_buf()); + let lifecycle = state.connection_lifecycle_snapshot("conn"); + let owner = TransactionOwner::spawn( + LifecycleTestIo { disconnects: Arc::new(AtomicUsize::new(0)) }, + TransactionOwnerConfig::default(), + ); + let watcher = backend.spawn_connection_lifecycle_watcher(lifecycle, &owner); + + owner.close().await.unwrap(); + + tokio::time::timeout(Duration::from_millis(100), watcher) + .await + .expect("lifecycle watcher must not outlive a normally disposed owner") + .unwrap(); + } + #[tokio::test] async fn local_backend_standalone_open_skips_plugin_dbx_engine_gate() { let data_dir = tempfile::tempdir().unwrap(); diff --git a/crates/dbx-mcp/src/http.rs b/crates/dbx-mcp/src/http.rs index 5304d7e9b5..ff4eeb6319 100644 --- a/crates/dbx-mcp/src/http.rs +++ b/crates/dbx-mcp/src/http.rs @@ -2,7 +2,10 @@ use std::{io, sync::Arc}; use axum::{http::Method, middleware, routing::get, Router}; use rmcp::transport::{ - streamable_http_server::{session::local::LocalSessionManager, tower::StreamableHttpService}, + streamable_http_server::{ + session::{local::LocalSessionManager, SessionManager}, + tower::StreamableHttpService, + }, StreamableHttpServerConfig, }; use tokio_util::sync::CancellationToken; @@ -25,7 +28,7 @@ pub fn streamable_http_router( allowed_hosts: Vec, web_mode: bool, ) -> Router { - build_streamable_http_router(backend, path, auth, allowed_hosts, web_mode, None) + build_streamable_http_router(backend, path, auth, allowed_hosts, web_mode, None, Default::default()) } fn build_streamable_http_router( @@ -35,6 +38,7 @@ fn build_streamable_http_router( allowed_hosts: Vec, web_mode: bool, cancellation: Option, + session_manager: Arc, ) -> Router { let mut rmcp_config = StreamableHttpServerConfig::default().with_allowed_hosts(allowed_hosts); if let Some(cancellation) = cancellation { @@ -44,7 +48,7 @@ fn build_streamable_http_router( let scope = McpScope::from_env(); let service: StreamableHttpService = StreamableHttpService::new( move || Ok(DbxMcpServer::with_runtime_options(server_backend.clone(), scope.clone(), web_mode)), - Default::default(), + session_manager, rmcp_config, ); @@ -98,6 +102,7 @@ pub async fn serve_streamable_http_on_listener( cancellation: CancellationToken, listener: tokio::net::TcpListener, ) -> io::Result<()> { + let session_manager = Arc::new(LocalSessionManager::default()); let mcp_router = build_streamable_http_router( backend, &config.path, @@ -105,14 +110,408 @@ pub async fn serve_streamable_http_on_listener( config.allowed_hosts, false, Some(cancellation.child_token()), + session_manager.clone(), ); let router = Router::new().route("/healthz", get(health)).route("/readyz", get(health)).merge(mcp_router); eprintln!("DBX MCP Streamable HTTP listening on http://{}{}", config.bind_addr, config.path); - axum::serve(listener, router) + let result = axum::serve(listener, router) .with_graceful_shutdown(async move { cancellation.cancelled().await; }) + .await; + close_local_sessions_bounded(&session_manager).await; + result +} + +async fn close_local_sessions_bounded(session_manager: &Arc) { + let session_ids = session_manager.sessions.read().await.keys().cloned().collect::>(); + let cleanup = async { + for session_id in session_ids { + let _ = session_manager.close_session(&session_id).await; + } + }; + if tokio::time::timeout(std::time::Duration::from_secs(10), cleanup).await.is_err() { + log::warn!("Timed out draining MCP HTTP protocol sessions during shutdown"); + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use async_trait::async_trait; + use dbx_core::{ + agent_events::ToolResult, agent_tools::AgentSqlPermissions, models::connection::ConnectionConfig, + storage::McpGlobalPolicy, + }; + use rmcp::{ + model::CallToolRequestParams, + service::ServiceExt, + transport::{ + streamable_http_client::StreamableHttpClientTransportConfig, + streamable_http_server::session::local::SessionConfig, StreamableHttpClientTransport, + }, + }; + use serde_json::{json, Value}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + use super::*; + use crate::{ + backend::DbxBackend, + transaction::{ + TransactionIo, TransactionIoError, TransactionIoSuccess, TransactionOwner, TransactionOwnerConfig, + }, + }; + + struct HttpTestIo { + sql: Arc>>, + disconnects: Arc, + in_transaction: bool, + } + + #[async_trait] + impl TransactionIo for HttpTestIo { + async fn execute( + &mut self, + sql: &str, + _max_rows: Option, + ) -> Result { + self.sql.lock().unwrap().push(sql.to_string()); + match sql { + "START TRANSACTION" => self.in_transaction = true, + "COMMIT" | "ROLLBACK" => self.in_transaction = false, + _ => {} + } + Ok(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, + 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, Ordering::SeqCst); + } + } + + struct HttpTestBackend { + connection: ConnectionConfig, + sql: Arc>>, + disconnects: Arc, + } + + impl HttpTestBackend { + fn new() -> Self { + Self { + connection: serde_json::from_value(json!({ + "id": "mysql", + "name": "mysql", + "db_type": "mysql", + "host": "", + "port": 3306, + "username": "", + "password": "", + "database": "app", + "ssl": false + })) + .unwrap(), + sql: Arc::new(std::sync::Mutex::new(Vec::new())), + disconnects: Arc::new(AtomicUsize::new(0)), + } + } + } + + #[async_trait] + impl DbxBackend for HttpTestBackend { + async fn load_mcp_global_policy(&self) -> Result { + Ok(McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, ..Default::default() }) + } + + async fn load_connections(&self) -> Result, String> { + Ok(vec![self.connection.clone()]) + } + + async fn execute_agent_tool( + &self, + _connection: &ConnectionConfig, + _database: &str, + tool_name: &str, + _arguments: Value, + _permissions: AgentSqlPermissions, + ) -> ToolResult { + ToolResult { + tool_call_id: "http-test".to_string(), + tool_name: tool_name.to_string(), + content: "unused".to_string(), + is_error: false, + explain_data: None, + } + } + + async fn open_transaction_owner( + &self, + _connection: &ConnectionConfig, + _database: &str, + _client_session_id: &str, + ) -> Result, String> { + Ok(TransactionOwner::spawn( + HttpTestIo { sql: self.sql.clone(), disconnects: self.disconnects.clone(), in_transaction: false }, + TransactionOwnerConfig { cleanup_timeout: std::time::Duration::from_millis(50), ..Default::default() }, + )) + } + + 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 { + Err("unused".to_string()) + } + async fn remove_connection_for_mcp(&self, _connection_id: &str) -> Result { + Ok(false) + } + } + + async fn start_http_test_server( + backend: Arc, + keep_alive: std::time::Duration, + ) -> (String, Arc, CancellationToken, tokio::task::JoinHandle<()>) { + // A single default provider avoids the "No rustls crypto provider is + // configured" panic when tests build reqwest clients in workspace + // builds where multiple rustls crypto features are present; the + // install is idempotent, so subsequent calls are no-ops. + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); + let address = listener.local_addr().unwrap(); + let mut session_config = SessionConfig::default(); + session_config.keep_alive = Some(keep_alive); + let mut local_manager = LocalSessionManager::default(); + local_manager.session_config = session_config; + let manager = Arc::new(local_manager); + let cancellation = CancellationToken::new(); + let router = build_streamable_http_router( + backend, + "/mcp", + HttpAuth::new("http-test-token".to_string(), Vec::::new(), true).unwrap(), + vec![address.to_string()], + false, + Some(cancellation.child_token()), + manager.clone(), + ); + let shutdown = cancellation.clone(); + let shutdown_manager = manager.clone(); + let task = tokio::spawn(async move { + axum::serve(listener, router) + .with_graceful_shutdown(async move { shutdown.cancelled().await }) + .await + .unwrap(); + close_local_sessions_bounded(&shutdown_manager).await; + }); + (format!("http://{address}/mcp"), manager, cancellation, task) + } + + async fn open_active_transaction(url: &str) -> (rmcp::service::RunningService, String) { + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(url.to_string()).auth_header("http-test-token"), + ); + let client = ().serve(transport).await.unwrap(); + let opened = client + .call_tool( + CallToolRequestParams::new("dbx_open_session").with_arguments( + serde_json::from_value(json!({ + "connection_id": "mysql", + "database": "app", + "enable_transactions": true + })) + .unwrap(), + ), + ) + .await + .unwrap(); + let session_id = opened.structured_content.unwrap()["session_id"].as_str().unwrap().to_string(); + let begun = client + .call_tool( + CallToolRequestParams::new("dbx_begin_transaction") + .with_arguments(serde_json::from_value(json!({"session_id": session_id.clone()})).unwrap()), + ) + .await + .unwrap(); + assert_ne!(begun.is_error, Some(true)); + (client, session_id) + } + + async fn wait_for_disposal(backend: &HttpTestBackend) { + tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if backend.disconnects.load(Ordering::SeqCst) == 1 + && backend.sql.lock().unwrap().iter().any(|sql| sql == "ROLLBACK") + { + break; + } + tokio::task::yield_now().await; + } + }) .await + .expect("HTTP session cleanup must roll back and disconnect the owner"); + } + + async fn open_and_drop_raw_sse(url: &str, outer_session_id: &str) { + let parsed = url::Url::parse(url).unwrap(); + let host = parsed.host_str().unwrap(); + let port = parsed.port_or_known_default().unwrap(); + let mut stream = tokio::net::TcpStream::connect((host, port)).await.unwrap(); + let path = match parsed.query() { + Some(query) => format!("{}?{query}", parsed.path()), + None => parsed.path().to_string(), + }; + let request = format!( + "GET {path} HTTP/1.1\r\nHost: {host}:{port}\r\nAuthorization: Bearer http-test-token\r\nAccept: text/event-stream\r\nMcp-Session-Id: {outer_session_id}\r\nMcp-Protocol-Version: 2025-06-18\r\nConnection: keep-alive\r\n\r\n" + ); + stream.write_all(request.as_bytes()).await.unwrap(); + + let headers = tokio::time::timeout(std::time::Duration::from_secs(1), async { + let mut response = Vec::new(); + let mut buffer = [0_u8; 1024]; + loop { + let read = stream.read(&mut buffer).await.unwrap(); + assert!(read > 0, "raw SSE connection closed before HTTP headers"); + response.extend_from_slice(&buffer[..read]); + if response.windows(4).any(|window| window == b"\r\n\r\n") { + break response; + } + assert!(response.len() <= 16 * 1024, "raw SSE response headers exceeded 16 KiB"); + } + }) + .await + .expect("raw SSE GET must return HTTP headers"); + let headers = std::str::from_utf8(&headers).unwrap(); + assert!(headers.starts_with("HTTP/1.1 200 "), "raw SSE GET did not return HTTP 200: {headers}"); + assert!(headers.to_ascii_lowercase().contains("content-type: text/event-stream")); + + drop(stream); + } + + #[tokio::test] + async fn authenticated_http_delete_rolls_back_and_disconnects_inner_owner() { + let backend = Arc::new(HttpTestBackend::new()); + let (url, manager, cancellation, server_task) = + start_http_test_server(backend.clone(), std::time::Duration::from_secs(30)).await; + let (client, _) = open_active_transaction(&url).await; + let outer_session_id = manager.sessions.read().await.keys().next().unwrap().to_string(); + + let response = reqwest::Client::new() + .delete(&url) + .bearer_auth("http-test-token") + .header("mcp-session-id", outer_session_id) + .send() + .await + .unwrap(); + assert!(response.status().is_success(), "DELETE returned {}", response.status()); + wait_for_disposal(&backend).await; + + drop(client); + cancellation.cancel(); + server_task.await.unwrap(); + } + + #[tokio::test] + async fn http_inactivity_expiry_rolls_back_and_disconnects_inner_owner() { + let backend = Arc::new(HttpTestBackend::new()); + let (url, manager, cancellation, server_task) = + start_http_test_server(backend.clone(), std::time::Duration::from_millis(50)).await; + let (client, _) = open_active_transaction(&url).await; + + wait_for_disposal(&backend).await; + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while !manager.sessions.read().await.is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .expect("expired outer HTTP session must be removed"); + + drop(client); + cancellation.cancel(); + server_task.await.unwrap(); + } + + #[tokio::test] + async fn transient_http_connections_preserve_one_outer_session() { + let backend = Arc::new(HttpTestBackend::new()); + let (url, manager, cancellation, server_task) = + start_http_test_server(backend.clone(), std::time::Duration::from_secs(30)).await; + let (client, inner_session_id) = open_active_transaction(&url).await; + let outer_session_id = manager.sessions.read().await.keys().next().unwrap().to_string(); + + open_and_drop_raw_sse(&url, &outer_session_id).await; + + let query = client + .call_tool( + CallToolRequestParams::new("dbx_execute_query").with_arguments( + serde_json::from_value(json!({ + "connection_id": "mysql", + "database": "app", + "session_id": inner_session_id, + "sql": "SELECT 1" + })) + .unwrap(), + ), + ) + .await + .unwrap(); + assert_ne!(query.is_error, Some(true)); + assert_eq!(query.structured_content.as_ref().unwrap()["transaction_state"], "active"); + assert_eq!(manager.sessions.read().await.len(), 1); + assert_eq!(backend.disconnects.load(Ordering::SeqCst), 0); + assert!(!backend.sql.lock().unwrap().iter().any(|sql| sql == "ROLLBACK")); + + let response = reqwest::Client::new() + .delete(&url) + .bearer_auth("http-test-token") + .header("mcp-session-id", outer_session_id) + .send() + .await + .unwrap(); + assert!(response.status().is_success()); + wait_for_disposal(&backend).await; + drop(client); + cancellation.cancel(); + server_task.await.unwrap(); + } + + #[tokio::test] + async fn http_service_shutdown_rolls_back_and_disconnects_inner_owner() { + let backend = Arc::new(HttpTestBackend::new()); + let (url, _manager, cancellation, server_task) = + start_http_test_server(backend.clone(), std::time::Duration::from_secs(30)).await; + let (client, _) = open_active_transaction(&url).await; + + cancellation.cancel(); + server_task.await.unwrap(); + wait_for_disposal(&backend).await; + drop(client); + } } diff --git a/crates/dbx-mcp/src/lib.rs b/crates/dbx-mcp/src/lib.rs index dc4c359b7b..5d2717ee84 100644 --- a/crates/dbx-mcp/src/lib.rs +++ b/crates/dbx-mcp/src/lib.rs @@ -8,6 +8,7 @@ pub mod paths; pub mod runtime; pub mod server; pub mod session; +pub mod transaction; pub mod transport; pub use backend::{ConnectionSummary, DbxBackend, LocalBackend, WebBackend}; diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 0e50d1abe3..c292e0cf4e 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -12,6 +12,7 @@ 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, @@ -151,6 +152,9 @@ pub struct OpenSessionRequest { #[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)] @@ -159,6 +163,12 @@ pub struct CloseSessionRequest { 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, @@ -356,6 +366,11 @@ struct ResolvedConnection { group_ids: Vec, } +struct SessionCleanupResult { + status: Option, + error: Option, +} + #[derive(Clone, Debug)] enum DatabaseScope { All, @@ -424,15 +439,87 @@ impl DbxMcpServer { Self { backend, scope, sessions: McpSessionStore::new(), tool_router } } - async fn close_backend_sessions_best_effort(&self, sessions: Vec) { - for session in sessions { - let _ = self - .backend + 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, @@ -648,6 +735,14 @@ impl DbxMcpServer { 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; @@ -774,6 +869,123 @@ impl DbxMcpServer { 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() { @@ -809,6 +1021,14 @@ impl DbxMcpServer { 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; @@ -947,6 +1167,132 @@ impl DbxMcpServer { { 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 @@ -1036,6 +1382,14 @@ impl DbxMcpServer { 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; @@ -1055,14 +1409,169 @@ impl DbxMcpServer { Ok(database) => database, Err(error) => return error, }; - let (session, expired) = self.sessions.open(&connection.id, &database).await.into_parts(); + 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; - match session { - Ok(session) => 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.", + 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 - )), - Err(error) => tool_error("SESSION_LIMIT", error), + )); + 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), } } @@ -1070,10 +1579,15 @@ impl DbxMcpServer { 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 { - if let Err(error) = self.ensure_tool_allowed("dbx_close_session").await { - return error; - } 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 { @@ -1082,20 +1596,24 @@ impl DbxMcpServer { format!("Session \"{}\" not found or already closed.", request.session_id), ); }; - match self - .backend - .close_client_session(&session.connection_id, &session.database, &session.client_session_id) - .await - { - Ok(_) => { - self.sessions.finish_close(&session.id).await; - text(format!("Session \"{}\" closed.", session.id)) - } - Err(error) => { - self.sessions.restore_after_failed_close(session).await; - backend_tool_error("SESSION_CLOSE_ERROR", error) - } + 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")] @@ -1699,10 +2217,57 @@ impl DbxMcpServer { } } + // 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) { @@ -1802,6 +2367,8 @@ impl DbxMcpServer { 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 @@ -1903,6 +2470,43 @@ 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()))]) } @@ -1920,6 +2524,7 @@ fn backend_tool_error(default_code: &str, error: impl Into) -> CallToolR "PRODUCTION_DATABASE_READ_ONLY", "PRODUCTION_WRITE_BLOCKED", "SQL_BLOCKED", + "TRANSACTION_UNSUPPORTED", ] { let marker = format!("{code}:"); if let Some(index) = error.find(&marker) { @@ -2014,6 +2619,24 @@ fn format_batch_results(results: &[crate::backend::BatchStatementResult]) -> Str 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, + 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: ")) @@ -2166,6 +2789,10 @@ fn mcp_permissions( } } +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 @@ -2208,6 +2835,214 @@ fn confirmed_batch_sql_block_reason( )) } +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( @@ -2543,7 +3378,110 @@ mod tests { use super::*; use async_trait::async_trait; use dbx_core::models::connection::ConnectionConfig; - use std::collections::HashSet; + 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, + 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, @@ -2553,6 +3491,15 @@ mod tests { 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 { @@ -2565,6 +3512,15 @@ mod tests { 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), } } } @@ -2612,13 +3568,49 @@ mod tests { Ok(()) } async fn load_mcp_global_policy(&self) -> Result { - Ok(self.policy.clone()) + 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, @@ -2753,6 +3745,8 @@ mod tests { statement_index: Some(0), error_message: sql.contains("FAIL_MCP_TEST").then(|| "statement failed".to_string()), merged: false, + transaction_state: None, + transaction_outcome: None, }]) } } @@ -2928,9 +3922,9 @@ mod tests { 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(), 19); + assert_eq!(tools.len(), 22); #[cfg(not(feature = "mq-admin"))] - assert_eq!(tools.len(), 17); + assert_eq!(tools.len(), 20); #[cfg(feature = "mq-admin")] assert!(names.contains(&"dbx_peek_messages")); #[cfg(not(feature = "mq-admin"))] @@ -2943,6 +3937,9 @@ mod tests { 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")); @@ -3056,7 +4053,7 @@ mod tests { ("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"]), + ("dbx_open_session", &["database", "enable_transactions"]), ("dbx_open_table", &["database", "schema"]), ("dbx_execute_and_show", &["database"]), ("dbx_add_connection", &["port", "database", "driver_profile"]), @@ -3184,9 +4181,9 @@ mod tests { ); let names = server.tool_router.list_all().into_iter().map(|tool| tool.name).collect::>(); #[cfg(feature = "mq-admin")] - assert_eq!(names.len(), 14); + assert_eq!(names.len(), 17); #[cfg(not(feature = "mq-admin"))] - assert_eq!(names.len(), 12); + 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")); @@ -3197,6 +4194,9 @@ mod tests { 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")); } @@ -3731,14 +4731,435 @@ mod tests { 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 })).await; + 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 @@ -3864,8 +5285,13 @@ mod tests { let mut expired_client_session_ids = Vec::new(); for _ in 0..3 { - let opened = - server.open_session(Parameters(OpenSessionRequest { selector: selector("sr"), database: None })).await; + 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 { @@ -3885,8 +5311,13 @@ mod tests { expired_client_session_ids.push(format!("mcp:{session_id}")); } - let opened = - server.open_session(Parameters(OpenSessionRequest { selector: selector("sr"), database: None })).await; + 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()); @@ -3895,6 +5326,168 @@ mod tests { 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"); @@ -3902,8 +5495,13 @@ mod tests { *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 })).await; + 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 { @@ -3950,8 +5548,13 @@ mod tests { McpScope::default(), false, ); - let rejected = - server.open_session(Parameters(OpenSessionRequest { selector: selector("redis"), database: None })).await; + 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 @@ -4012,8 +5615,13 @@ mod tests { 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 })).await; + 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 @@ -4216,8 +5824,13 @@ mod tests { 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 })).await; + 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 @@ -4387,8 +6000,13 @@ mod tests { // 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 })).await; + 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 { @@ -4491,6 +6109,8 @@ mod tests { 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")); @@ -4520,6 +6140,8 @@ mod tests { 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")); diff --git a/crates/dbx-mcp/src/session.rs b/crates/dbx-mcp/src/session.rs index 65c3cbac89..89d7acae5e 100644 --- a/crates/dbx-mcp/src/session.rs +++ b/crates/dbx-mcp/src/session.rs @@ -5,6 +5,8 @@ use std::time::{Duration, Instant}; use tokio::sync::Mutex; use uuid::Uuid; +use crate::transaction::TransactionOwner; + /// Idle time after which an MCP session is considered expired and its pinned /// backend connection pool may be reclaimed. const SESSION_IDLE_TTL: Duration = Duration::from_secs(30 * 60); @@ -22,6 +24,7 @@ pub struct McpSession { /// Value forwarded to the backend as `client_session_id`, pinning all /// queries in this session to the same connection pool. pub client_session_id: String, + pub transaction_owner: Option>, last_used: Instant, } @@ -43,6 +46,7 @@ impl McpSessionStoreResult { #[derive(Default)] struct McpSessionState { + opening: HashMap, active: HashMap, closing: HashMap, } @@ -63,8 +67,8 @@ impl McpSessionStore { /// reached; expired sessions are swept before the cap is enforced. pub async fn open(&self, connection_id: &str, database: &str) -> McpSessionStoreResult> { let mut state = self.state.lock().await; - let expired = sweep_expired(&mut state.active); - if state.active.len() + state.closing.len() >= MAX_SESSIONS { + let expired = sweep_expired(&mut state); + if session_count(&state) >= MAX_SESSIONS { return McpSessionStoreResult::new( Err(format!( "Too many open MCP sessions (max {MAX_SESSIONS}). Close unused sessions with dbx_close_session." @@ -72,25 +76,55 @@ impl McpSessionStore { expired, ); } - let id = format!("mcp-session-{}", Uuid::new_v4()); - let session = McpSession { - id: id.clone(), - connection_id: connection_id.to_string(), - database: database.to_string(), - // Prefixed so these pools are easy to distinguish from desktop UI - // sessions in backend diagnostics. `:` is normalized away by the - // backend pool key sanitizer. - client_session_id: format!("mcp:{id}"), - last_used: Instant::now(), - }; - state.active.insert(id, session.clone()); + let session = new_session(connection_id, database); + state.active.insert(session.id.clone(), session.clone()); McpSessionStoreResult::new(Ok(session), expired) } + pub async fn reserve_opening( + &self, + connection_id: &str, + database: &str, + ) -> McpSessionStoreResult> { + let mut state = self.state.lock().await; + let expired = sweep_expired(&mut state); + if session_count(&state) >= MAX_SESSIONS { + return McpSessionStoreResult::new( + Err(format!( + "Too many open MCP sessions (max {MAX_SESSIONS}). Close unused sessions with dbx_close_session." + )), + expired, + ); + } + let session = new_session(connection_id, database); + state.opening.insert(session.id.clone(), session.clone()); + McpSessionStoreResult::new(Ok(session), expired) + } + + pub async fn promote_opening(&self, session_id: &str, owner: Arc) -> Result { + let mut state = self.state.lock().await; + let mut session = + state.opening.remove(session_id).ok_or_else(|| "MCP session opening was cancelled".to_string())?; + session.transaction_owner = Some(owner); + state.active.insert(session.id.clone(), session.clone()); + Ok(session) + } + + pub async fn cancel_opening(&self, session_id: &str) -> Option { + self.state.lock().await.opening.remove(session_id) + } + + pub async fn attach_transaction_owner(&self, session_id: &str, owner: Arc) -> Result<(), String> { + let mut state = self.state.lock().await; + let session = state.active.get_mut(session_id).ok_or_else(|| "MCP session is no longer active".to_string())?; + session.transaction_owner = Some(owner); + Ok(()) + } + /// Resolve a session id and refresh its idle timer. pub async fn resolve(&self, session_id: &str) -> McpSessionStoreResult> { let mut state = self.state.lock().await; - let expired = sweep_expired(&mut state.active); + let expired = sweep_expired(&mut state); let session = state.active.get_mut(session_id).map(|session| { session.last_used = Instant::now(); session.clone() @@ -102,10 +136,10 @@ impl McpSessionStore { /// cap until the backend pool is closed successfully. pub async fn begin_close(&self, session_id: &str) -> McpSessionStoreResult> { let mut state = self.state.lock().await; - let expired = sweep_expired(&mut state.active); + let expired = sweep_expired(&mut state); let session = state.active.remove(session_id); if let Some(session) = &session { - state.closing.insert(session.id.clone(), session.clone()); + state.closing.entry(session.id.clone()).or_insert_with(|| session.clone()); } McpSessionStoreResult::new(session, expired) } @@ -129,16 +163,40 @@ impl McpSessionStore { } } -fn sweep_expired(sessions: &mut HashMap) -> Vec { +fn new_session(connection_id: &str, database: &str) -> McpSession { + let id = format!("mcp-session-{}", Uuid::new_v4()); + McpSession { + id: id.clone(), + connection_id: connection_id.to_string(), + database: database.to_string(), + // Prefixed so these pools are easy to distinguish from desktop UI + // sessions in backend diagnostics. `:` is normalized away by the + // backend pool key sanitizer. + client_session_id: format!("mcp:{id}"), + transaction_owner: None, + last_used: Instant::now(), + } +} + +fn session_count(state: &McpSessionState) -> usize { + state.opening.len() + state.active.len() + state.closing.len() +} + +fn sweep_expired(state: &mut McpSessionState) -> Vec { let now = Instant::now(); - let mut expired = Vec::new(); - sessions.retain(|_, session| { - let active = now.duration_since(session.last_used) < SESSION_IDLE_TTL; - if !active { - expired.push(session.clone()); + let expired_ids = state + .active + .iter() + .filter(|(_, session)| now.duration_since(session.last_used) >= SESSION_IDLE_TTL) + .map(|(id, _)| id.clone()) + .collect::>(); + let mut expired = Vec::with_capacity(expired_ids.len()); + for id in expired_ids { + if let Some(session) = state.active.remove(&id) { + state.closing.insert(id, session.clone()); + expired.push(session); } - active - }); + } expired } @@ -165,6 +223,7 @@ mod tests { assert!(expired.is_empty()); assert_eq!(closing.unwrap().id, session.id); assert!(store.resolve(&session.id).await.into_parts().0.is_none()); + assert!(store.begin_close(&session.id).await.into_parts().0.is_none()); store.finish_close(&session.id).await; assert!(store.begin_close(&session.id).await.into_parts().0.is_none()); } @@ -180,7 +239,7 @@ mod tests { } #[tokio::test] - async fn expired_sessions_are_swept_and_free_capacity() { + async fn expired_sessions_remain_counted_until_cleanup_finishes() { let store = McpSessionStore::new(); let mut session_ids = Vec::new(); for _ in 0..MAX_SESSIONS { @@ -190,11 +249,15 @@ mod tests { for session_id in &session_ids { store.expire_for_test(session_id).await; } - // All sessions are expired: the sweep must reclaim them before the cap - // check, return them for backend cleanup, and allow a new session. + // Expiry transfers ownership to closing. Capacity is not reusable until + // the detached cleanup worker confirms physical disposal. let (opened, expired) = store.open("conn-1", "").await.into_parts(); - opened.unwrap(); + assert!(opened.unwrap_err().contains("Too many open MCP sessions")); assert_eq!(expired.len(), MAX_SESSIONS); + for session in expired { + store.finish_close(&session.id).await; + } + store.open("conn-1", "").await.into_parts().0.unwrap(); // Resolving an expired session must fail instead of silently pinning a // fresh backend connection. @@ -205,4 +268,19 @@ mod tests { assert_eq!(expired.len(), 1); assert_eq!(expired[0].id, session.id); } + + #[tokio::test] + async fn provisional_openings_count_toward_cap_until_promoted_or_cancelled() { + let store = McpSessionStore::new(); + let mut openings = Vec::new(); + for _ in 0..MAX_SESSIONS { + openings.push(store.reserve_opening("conn-1", "").await.into_parts().0.unwrap()); + } + assert!(store.reserve_opening("conn-1", "").await.into_parts().0.is_err()); + assert!(store.open("conn-1", "").await.into_parts().0.is_err()); + + let cancelled = openings.pop().unwrap(); + assert!(store.cancel_opening(&cancelled.id).await.is_some()); + store.open("conn-1", "").await.into_parts().0.unwrap(); + } } diff --git a/crates/dbx-mcp/src/transaction.rs b/crates/dbx-mcp/src/transaction.rs new file mode 100644 index 0000000000..1274734a1b --- /dev/null +++ b/crates/dbx-mcp/src/transaction.rs @@ -0,0 +1,1214 @@ +use std::{ + collections::HashMap, + fmt, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, Weak, + }, + time::Duration, +}; + +use async_trait::async_trait; +use dbx_core::db::QueryResult; +use serde::{Deserialize, Serialize}; +use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit, Semaphore}; +use tokio_util::sync::CancellationToken; + +use dbx_core::connection::AppState; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TransactionState { + Idle, + Active, + Unknown, +} + +impl TransactionState { + pub fn as_str(self) -> &'static str { + match self { + Self::Idle => "idle", + Self::Active => "active", + Self::Unknown => "unknown", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TransactionOutcome { + Committed, + RolledBack, + Unknown, +} + +impl TransactionOutcome { + pub fn as_str(self) -> &'static str { + match self { + Self::Committed => "committed", + Self::RolledBack => "rolled_back", + Self::Unknown => "unknown", + } + } +} + +#[derive(Debug)] +pub struct TransactionIoSuccess { + pub result: QueryResult, + pub in_transaction: bool, +} + +#[derive(Debug)] +pub enum TransactionIoError { + Server { code: u16, message: String }, + Transport(String), +} + +#[async_trait] +pub trait TransactionIo: Send + 'static { + async fn execute(&mut self, sql: &str, max_rows: Option) + -> Result; + + async fn ping_in_transaction(&mut self) -> Result; + + async fn disconnect(&mut self); +} + +pub struct MysqlTransactionIo { + conn: Option, + state: Arc, + connection_id: String, + database: String, + client_session_id: String, +} + +impl MysqlTransactionIo { + pub fn new( + conn: mysql_async::Conn, + state: Arc, + connection_id: String, + database: String, + client_session_id: String, + ) -> Self { + Self { conn: Some(conn), state, connection_id, database, client_session_id } + } + + fn conn(&mut self) -> Result<&mut mysql_async::Conn, TransactionIoError> { + self.conn.as_mut().ok_or_else(|| TransactionIoError::Transport("MySQL connection is closed".to_string())) + } +} + +#[async_trait] +impl TransactionIo for MysqlTransactionIo { + async fn execute( + &mut self, + sql: &str, + max_rows: Option, + ) -> Result { + let execution = dbx_core::db::mysql::execute_transaction_statement_on_conn(self.conn()?, sql, max_rows) + .await + .map_err(map_mysql_transaction_error)?; + Ok(TransactionIoSuccess { result: execution.result, in_transaction: execution.in_transaction }) + } + + async fn ping_in_transaction(&mut self) -> Result { + dbx_core::db::mysql::ping_transaction_status_on_conn(self.conn()?).await.map_err(map_mysql_transaction_error) + } + + async fn disconnect(&mut self) { + if let Some(conn) = self.conn.take() { + let _ = conn.disconnect().await; + } + let _ = self + .state + .close_client_session_pool(&self.connection_id, Some(&self.database), &self.client_session_id) + .await; + } +} + +fn map_mysql_transaction_error(error: dbx_core::db::mysql::MySqlTransactionError) -> TransactionIoError { + match error { + dbx_core::db::mysql::MySqlTransactionError::Server { code, message, .. } => { + TransactionIoError::Server { code, message } + } + dbx_core::db::mysql::MySqlTransactionError::Transport(message) => TransactionIoError::Transport(message), + } +} + +#[derive(Default)] +pub struct TransactionOwnerRegistry { + owners: Mutex>>>, +} + +impl TransactionOwnerRegistry { + pub fn register(&self, connection_id: &str, owner: &Arc) { + let mut owners = self.owners.lock().expect("transaction owner registry lock"); + let entries = owners.entry(connection_id.to_string()).or_default(); + entries.retain(|entry| entry.strong_count() > 0); + entries.push(Arc::downgrade(owner)); + } + + pub fn invalidate_connection(&self, connection_id: &str) { + let owners = self + .owners + .lock() + .expect("transaction owner registry lock") + .remove(connection_id) + .unwrap_or_default() + .into_iter() + .filter_map(|owner| owner.upgrade()) + .collect::>(); + for owner in owners { + owner.invalidate(); + } + } + + pub fn invalidate_all(&self) { + let owners = self + .owners + .lock() + .expect("transaction owner registry lock") + .drain() + .flat_map(|(_, owners)| owners) + .filter_map(|owner| owner.upgrade()) + .collect::>(); + for owner in owners { + owner.invalidate(); + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct TransactionOwnerConfig { + pub idle_ttl: Duration, + pub operation_timeout: Duration, + pub cleanup_timeout: Duration, +} + +impl Default for TransactionOwnerConfig { + fn default() -> Self { + Self { + idle_ttl: Duration::from_secs(30 * 60), + operation_timeout: Duration::from_secs(300), + cleanup_timeout: Duration::from_secs(5), + } + } +} + +#[derive(Debug)] +pub struct TransactionResult { + pub result: QueryResult, + pub state: TransactionState, + pub outcome: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TransactionFailure { + pub code: &'static str, + pub message: String, + pub mysql_code: Option, + pub state: TransactionState, + pub outcome: Option, +} + +#[derive(Debug, Clone, Copy)] +pub struct TransactionStatus { + pub state: TransactionState, + pub outcome: Option, +} + +impl Default for TransactionStatus { + fn default() -> Self { + Self { state: TransactionState::Idle, outcome: None } + } +} + +enum Operation { + Begin, + Query { sql: String, max_rows: Option }, + Commit, + Rollback, +} + +struct Command { + operation: Operation, + request_cancellation: CancellationToken, + response: oneshot::Sender>, +} + +struct TransactionOwnerInner { + sender: mpsc::Sender, + operation_gate: Arc, + cancellation: CancellationToken, + status: Mutex, + closed: AtomicBool, + closed_notify: tokio::sync::Notify, + closed_token: CancellationToken, + cleanup_timeout: Duration, +} + +pub struct TransactionOwner { + inner: Arc, +} + +impl fmt::Debug for TransactionOwner { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.debug_struct("TransactionOwner").field("status", &self.status()).finish_non_exhaustive() + } +} + +impl TransactionOwner { + pub fn spawn(io: impl TransactionIo, config: TransactionOwnerConfig) -> Arc { + Self::spawn_inner(io, config, None) + } + + pub fn spawn_with_resource_permit( + io: impl TransactionIo, + config: TransactionOwnerConfig, + resource_permit: OwnedSemaphorePermit, + ) -> Arc { + Self::spawn_inner(io, config, Some(resource_permit)) + } + + fn spawn_inner( + io: impl TransactionIo, + config: TransactionOwnerConfig, + resource_permit: Option, + ) -> Arc { + let (sender, receiver) = mpsc::channel(1); + let inner = Arc::new(TransactionOwnerInner { + sender, + operation_gate: Arc::new(Semaphore::new(1)), + cancellation: CancellationToken::new(), + status: Mutex::new(TransactionStatus::default()), + closed: AtomicBool::new(false), + closed_notify: tokio::sync::Notify::new(), + closed_token: CancellationToken::new(), + cleanup_timeout: config.cleanup_timeout, + }); + tokio::spawn(run_owner(Box::new(io), receiver, inner.clone(), config, resource_permit)); + Arc::new(Self { inner }) + } + + pub fn status(&self) -> TransactionStatus { + *self.inner.status.lock().expect("transaction status lock") + } + + pub async fn acquire(self: &Arc) -> Result { + let permit = tokio::select! { + permit = self.inner.operation_gate.clone().acquire_owned() => { + permit.map_err(|_| self.terminal_failure("TRANSACTION_CLOSED", "Transaction session is closed."))? + } + _ = self.inner.cancellation.cancelled() => { + return Err(self.terminal_failure("TRANSACTION_CLOSED", "Transaction session is closed.")); + } + }; + if self.inner.closed.load(Ordering::Acquire) || self.inner.cancellation.is_cancelled() { + return Err(self.terminal_failure("TRANSACTION_CLOSED", "Transaction session is closed.")); + } + if self.status().state == TransactionState::Unknown { + return Err(self.terminal_failure( + "TRANSACTION_OUTCOME_UNKNOWN", + "Transaction outcome is unknown; this session cannot execute more SQL.", + )); + } + Ok(TransactionLease { owner: self.clone(), _permit: permit }) + } + + pub async fn close(&self) -> Result { + self.inner.cancellation.cancel(); + if !self.inner.closed.load(Ordering::Acquire) { + let close_timeout = self.inner.cleanup_timeout.saturating_mul(2).saturating_add(Duration::from_millis(100)); + if tokio::time::timeout(close_timeout, self.wait_closed()).await.is_err() { + return Err(self.status()); + } + } + Ok(self.status()) + } + + pub async fn wait_closed(&self) { + while !self.inner.closed.load(Ordering::Acquire) { + let notified = self.inner.closed_notify.notified(); + if self.inner.closed.load(Ordering::Acquire) { + break; + } + notified.await; + } + } + + pub fn invalidate(&self) { + self.inner.cancellation.cancel(); + } + + pub fn closed_token(&self) -> CancellationToken { + self.inner.closed_token.clone() + } + + fn terminal_failure(&self, code: &'static str, message: impl Into) -> TransactionFailure { + let status = self.status(); + TransactionFailure { + code, + message: message.into(), + mysql_code: None, + state: status.state, + outcome: status.outcome, + } + } +} + +impl Drop for TransactionOwner { + fn drop(&mut self) { + // The worker owns only `TransactionOwnerInner`, not another public + // handle. Dropping the last Arc therefore means its + // protocol/session owner disappeared (HTTP DELETE, stdio shutdown, or + // server teardown) and must start bounded cleanup out-of-band. + self.inner.cancellation.cancel(); + } +} + +pub struct TransactionLease { + owner: Arc, + _permit: OwnedSemaphorePermit, +} + +impl fmt::Debug for TransactionLease { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.debug_struct("TransactionLease").finish_non_exhaustive() + } +} + +impl TransactionLease { + pub async fn begin(self) -> Result { + self.run(Operation::Begin).await + } + + pub async fn query( + self, + sql: impl Into, + max_rows: Option, + ) -> Result { + self.run(Operation::Query { sql: sql.into(), max_rows }).await + } + + pub async fn commit(self) -> Result { + self.run(Operation::Commit).await + } + + pub async fn rollback(self) -> Result { + self.run(Operation::Rollback).await + } + + pub async fn batch( + mut self, + statements: Vec, + max_rows: Option, + continue_on_error: bool, + ) -> Vec> { + let mut results = Vec::with_capacity(statements.len()); + for statement in statements { + let result = self.run_ref(Operation::Query { sql: statement, max_rows }).await; + let stop = match &result { + Ok(_) => false, + Err(error) => error.state == TransactionState::Unknown || !continue_on_error, + }; + results.push(result); + if stop { + break; + } + } + results + } + + async fn run(self, operation: Operation) -> Result { + let mut this = self; + this.run_ref(operation).await + } + + async fn run_ref(&mut self, operation: Operation) -> Result { + let request_cancellation = CancellationToken::new(); + let mut drop_guard = RequestDropGuard { cancellation: request_cancellation.clone(), armed: true }; + let (response, result) = oneshot::channel(); + self.owner + .inner + .sender + .send(Command { operation, request_cancellation, response }) + .await + .map_err(|_| self.owner.terminal_failure("TRANSACTION_CLOSED", "Transaction session is closed."))?; + let result = result.await.map_err(|_| { + self.owner.terminal_failure("TRANSACTION_CLOSED", "Transaction session closed during the operation.") + })?; + drop_guard.armed = false; + result + } +} + +struct RequestDropGuard { + cancellation: CancellationToken, + armed: bool, +} + +impl Drop for RequestDropGuard { + fn drop(&mut self) { + if self.armed { + self.cancellation.cancel(); + } + } +} + +async fn run_owner( + mut io: Box, + mut receiver: mpsc::Receiver, + inner: Arc, + config: TransactionOwnerConfig, + resource_permit: Option, +) { + loop { + let command = tokio::select! { + _ = inner.cancellation.cancelled() => None, + _ = tokio::time::sleep(config.idle_ttl) => None, + command = receiver.recv() => command, + }; + let Some(command) = command else { break }; + let terminal = handle_command(io.as_mut(), &inner, command, config.operation_timeout).await; + if terminal { + break; + } + } + + let status = *inner.status.lock().expect("transaction status lock"); + if status.state == TransactionState::Active { + let cleanup = tokio::time::timeout(config.cleanup_timeout, io.execute("ROLLBACK", None)).await; + match cleanup { + Ok(Ok(success)) if !success.in_transaction => { + *inner.status.lock().expect("transaction status lock") = + TransactionStatus { state: TransactionState::Idle, outcome: Some(TransactionOutcome::RolledBack) }; + } + _ => set_unknown(&inner), + } + } + tokio::time::timeout(config.cleanup_timeout, io.disconnect()).await.ok(); + inner.closed.store(true, Ordering::Release); + inner.operation_gate.close(); + inner.closed_notify.notify_waiters(); + inner.closed_token.cancel(); + drop(resource_permit); +} + +async fn handle_command( + io: &mut dyn TransactionIo, + inner: &Arc, + command: Command, + timeout: Duration, +) -> bool { + let request_cancellation = command.request_cancellation.clone(); + let current = *inner.status.lock().expect("transaction status lock"); + let invalid = match &command.operation { + Operation::Begin if current.state == TransactionState::Active => { + Some(("TRANSACTION_ALREADY_ACTIVE", "A transaction is already active for this session.")) + } + Operation::Commit | Operation::Rollback if current.state != TransactionState::Active => { + Some(("TRANSACTION_NOT_ACTIVE", "No active transaction exists for this session.")) + } + _ => None, + }; + if let Some((code, message)) = invalid { + let _ = command.response.send(Err(failure(code, message, None, current))); + return false; + } + + let is_begin = matches!(&command.operation, Operation::Begin); + let is_commit = matches!(&command.operation, Operation::Commit); + let (sql, max_rows, success_outcome) = match command.operation { + Operation::Begin => ("START TRANSACTION".to_string(), None, None), + Operation::Query { sql, max_rows } => (sql, max_rows, None), + Operation::Commit => ("COMMIT".to_string(), None, Some(TransactionOutcome::Committed)), + Operation::Rollback => ("ROLLBACK".to_string(), None, Some(TransactionOutcome::RolledBack)), + }; + + let execution = tokio::select! { + _ = inner.cancellation.cancelled() => { + set_unknown(inner); + let status = *inner.status.lock().expect("transaction status lock"); + let _ = command.response.send(Err(failure("TRANSACTION_CANCELLED", "Transaction operation was cancelled.", None, status))); + return true; + } + _ = request_cancellation.cancelled() => { + set_unknown(inner); + let status = *inner.status.lock().expect("transaction status lock"); + let _ = command.response.send(Err(failure("TRANSACTION_CANCELLED", "Transaction request was dropped before completion.", None, status))); + return true; + } + result = tokio::time::timeout(timeout, io.execute(&sql, max_rows)) => result, + }; + + match execution { + Ok(Ok(success)) => { + let expected_in_transaction = match success_outcome { + Some(_) => false, + None if is_begin => true, + None => success.in_transaction, + }; + if success.in_transaction != expected_in_transaction { + set_unknown(inner); + let status = *inner.status.lock().expect("transaction status lock"); + let _ = command.response.send(Err(failure( + "TRANSACTION_STATE_UNKNOWN", + "MySQL acknowledged the operation with an unexpected transaction state.", + None, + status, + ))); + return true; + } + let status = TransactionStatus { + state: if success.in_transaction { TransactionState::Active } else { TransactionState::Idle }, + outcome: success_outcome, + }; + *inner.status.lock().expect("transaction status lock") = status; + let _ = command.response.send(Ok(TransactionResult { + result: success.result, + state: status.state, + outcome: status.outcome, + })); + false + } + Ok(Err(TransactionIoError::Server { code, message })) => { + if is_commit { + set_unknown(inner); + let status = *inner.status.lock().expect("transaction status lock"); + let _ = command.response.send(Err(failure("TRANSACTION_STATE_UNKNOWN", message, Some(code), status))); + return true; + } + let probe = tokio::select! { + _ = inner.cancellation.cancelled() => { + set_unknown(inner); + let status = *inner.status.lock().expect("transaction status lock"); + let _ = command.response.send(Err(failure( + "TRANSACTION_CANCELLED", + "Transaction operation was cancelled during status verification.", + Some(code), + status, + ))); + return true; + } + _ = request_cancellation.cancelled() => { + set_unknown(inner); + let status = *inner.status.lock().expect("transaction status lock"); + let _ = command.response.send(Err(failure( + "TRANSACTION_CANCELLED", + "Transaction request was dropped during status verification.", + Some(code), + status, + ))); + return true; + } + result = tokio::time::timeout(timeout, io.ping_in_transaction()) => result, + }; + match probe { + Ok(Ok(in_transaction)) => { + let status = TransactionStatus { + state: if in_transaction { TransactionState::Active } else { TransactionState::Idle }, + outcome: if current.state == TransactionState::Active && !in_transaction { + Some(TransactionOutcome::RolledBack) + } else { + None + }, + }; + *inner.status.lock().expect("transaction status lock") = status; + let _ = command.response.send(Err(failure("MYSQL_SERVER_ERROR", message, Some(code), status))); + false + } + _ => { + set_unknown(inner); + let status = *inner.status.lock().expect("transaction status lock"); + let _ = + command.response.send(Err(failure("TRANSACTION_STATE_UNKNOWN", message, Some(code), status))); + true + } + } + } + Ok(Err(TransactionIoError::Transport(message))) => { + set_unknown(inner); + let status = *inner.status.lock().expect("transaction status lock"); + let _ = command.response.send(Err(failure("TRANSACTION_STATE_UNKNOWN", message, None, status))); + true + } + Err(_) => { + set_unknown(inner); + let status = *inner.status.lock().expect("transaction status lock"); + let _ = command.response.send(Err(failure( + "TRANSACTION_TIMEOUT", + format!("Transaction operation timed out after {} seconds.", timeout.as_secs()), + None, + status, + ))); + true + } + } +} + +fn set_unknown(inner: &TransactionOwnerInner) { + *inner.status.lock().expect("transaction status lock") = + TransactionStatus { state: TransactionState::Unknown, outcome: Some(TransactionOutcome::Unknown) }; +} + +fn failure( + code: &'static str, + message: impl Into, + mysql_code: Option, + status: TransactionStatus, +) -> TransactionFailure { + TransactionFailure { code, message: message.into(), mysql_code, state: status.state, outcome: status.outcome } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::VecDeque, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, + time::Duration, + }; + + use async_trait::async_trait; + use dbx_core::db::QueryResult; + use tokio::sync::{Notify, Semaphore}; + + use super::{ + TransactionIo, TransactionIoError, TransactionIoSuccess, TransactionOutcome, TransactionOwner, + TransactionOwnerConfig, TransactionOwnerRegistry, TransactionState, + }; + + #[derive(Clone)] + enum Step { + Success { in_transaction: bool }, + ServerError { code: u16, message: &'static str }, + TransportError(&'static str), + Ping { in_transaction: bool }, + WaitPing { entered: Arc, release: Arc, in_transaction: bool }, + Wait(Arc), + } + + struct ScriptedIo { + steps: VecDeque, + active_calls: Arc, + max_active_calls: Arc, + disconnects: Arc, + ping_calls: Arc, + executed_sql: Arc>>, + } + + struct IoProbe { + active_calls: Arc, + max_active_calls: Arc, + disconnects: Arc, + ping_calls: Arc, + executed_sql: Arc>>, + } + + impl ScriptedIo { + fn new(steps: impl IntoIterator) -> (Self, IoProbe) { + let active_calls = Arc::new(AtomicUsize::new(0)); + let max_active_calls = Arc::new(AtomicUsize::new(0)); + let disconnects = Arc::new(AtomicUsize::new(0)); + let ping_calls = Arc::new(AtomicUsize::new(0)); + let executed_sql = Arc::new(Mutex::new(Vec::new())); + ( + Self { + steps: steps.into_iter().collect(), + active_calls: active_calls.clone(), + max_active_calls: max_active_calls.clone(), + disconnects: disconnects.clone(), + ping_calls: ping_calls.clone(), + executed_sql: executed_sql.clone(), + }, + IoProbe { active_calls, max_active_calls, disconnects, ping_calls, executed_sql }, + ) + } + + async fn run_step(&mut self, step: Step) -> Result { + let active = self.active_calls.fetch_add(1, Ordering::SeqCst) + 1; + self.max_active_calls.fetch_max(active, Ordering::SeqCst); + let result = match step { + Step::Success { in_transaction } => Ok(TransactionIoSuccess { result: empty_result(), in_transaction }), + Step::ServerError { code, message } => { + Err(TransactionIoError::Server { code, message: message.to_string() }) + } + Step::TransportError(message) => Err(TransactionIoError::Transport(message.to_string())), + Step::Ping { .. } => panic!("ping step used as SQL execution"), + Step::WaitPing { .. } => panic!("waiting ping step used as SQL execution"), + Step::Wait(notify) => { + notify.notified().await; + Ok(TransactionIoSuccess { result: empty_result(), in_transaction: true }) + } + }; + self.active_calls.fetch_sub(1, Ordering::SeqCst); + result + } + } + + #[async_trait] + impl TransactionIo for ScriptedIo { + async fn execute( + &mut self, + sql: &str, + _max_rows: Option, + ) -> Result { + self.executed_sql.lock().unwrap().push(sql.to_string()); + let step = self.steps.pop_front().expect("scripted SQL step"); + self.run_step(step).await + } + + async fn ping_in_transaction(&mut self) -> Result { + self.ping_calls.fetch_add(1, Ordering::SeqCst); + match self.steps.pop_front().expect("scripted ping step") { + Step::Ping { in_transaction } => Ok(in_transaction), + Step::WaitPing { entered, release, in_transaction } => { + entered.notify_one(); + release.notified().await; + Ok(in_transaction) + } + other => panic!("expected ping step, got {}", step_name(&other)), + } + } + + async fn disconnect(&mut self) { + self.disconnects.fetch_add(1, Ordering::SeqCst); + } + } + + fn step_name(step: &Step) -> &'static str { + match step { + Step::Success { .. } => "success", + Step::ServerError { .. } => "server error", + Step::TransportError(_) => "transport error", + Step::Ping { .. } => "ping", + Step::WaitPing { .. } => "waiting ping", + Step::Wait(_) => "wait", + } + } + + fn empty_result() -> QueryResult { + 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, + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages: Vec::new(), + } + } + + fn config() -> TransactionOwnerConfig { + TransactionOwnerConfig { + idle_ttl: Duration::from_secs(30 * 60), + operation_timeout: Duration::from_secs(30), + cleanup_timeout: Duration::from_secs(1), + } + } + + #[tokio::test] + async fn begin_ack_and_commit_ack_drive_visible_state_and_outcome() { + let (io, _) = + ScriptedIo::new([Step::Success { in_transaction: true }, Step::Success { in_transaction: false }]); + let owner = TransactionOwner::spawn(io, config()); + + let begin = owner.acquire().await.unwrap().begin().await.unwrap(); + assert_eq!(begin.state, TransactionState::Active); + assert_eq!(begin.outcome, None); + + let commit = owner.acquire().await.unwrap().commit().await.unwrap(); + assert_eq!(commit.state, TransactionState::Idle); + assert_eq!(commit.outcome, Some(TransactionOutcome::Committed)); + } + + #[tokio::test] + async fn recoverable_server_error_probes_fresh_status_and_keeps_transaction_active() { + let (io, probe) = ScriptedIo::new([ + Step::Success { in_transaction: true }, + Step::ServerError { code: 1062, message: "Duplicate entry" }, + Step::Ping { in_transaction: true }, + Step::Success { in_transaction: true }, + Step::Success { in_transaction: false }, + ]); + let owner = TransactionOwner::spawn(io, config()); + owner.acquire().await.unwrap().begin().await.unwrap(); + + let duplicate = owner.acquire().await.unwrap().query("INSERT INTO t VALUES (1)", None).await.unwrap_err(); + assert_eq!(duplicate.mysql_code, Some(1062)); + assert_eq!(duplicate.state, TransactionState::Active); + + let read = owner.acquire().await.unwrap().query("SELECT * FROM t FOR UPDATE", None).await.unwrap(); + assert_eq!(read.state, TransactionState::Active); + let rollback = owner.acquire().await.unwrap().rollback().await.unwrap(); + assert_eq!(rollback.outcome, Some(TransactionOutcome::RolledBack)); + assert_eq!(probe.executed_sql.lock().unwrap().len(), 4); + } + + #[tokio::test] + async fn lost_commit_ack_is_permanently_unknown_even_after_cleanup() { + let (io, probe) = ScriptedIo::new([ + Step::Success { in_transaction: true }, + Step::TransportError("connection reset after COMMIT"), + ]); + let owner = TransactionOwner::spawn(io, config()); + owner.acquire().await.unwrap().begin().await.unwrap(); + + let failure = owner.acquire().await.unwrap().commit().await.unwrap_err(); + assert_eq!(failure.state, TransactionState::Unknown); + assert_eq!(failure.outcome, Some(TransactionOutcome::Unknown)); + + let close = owner.close().await.unwrap(); + assert_eq!(close.state, TransactionState::Unknown); + assert_eq!(close.outcome, Some(TransactionOutcome::Unknown)); + assert_eq!(probe.disconnects.load(Ordering::SeqCst), 1); + let later = owner.acquire().await.unwrap_err(); + assert_eq!(later.state, TransactionState::Unknown); + } + + #[tokio::test] + async fn commit_server_error_is_terminal_unknown_without_status_probe() { + let (io, probe) = ScriptedIo::new([ + Step::Success { in_transaction: true }, + Step::ServerError { code: 1180, message: "Got error during COMMIT" }, + Step::Ping { in_transaction: false }, + ]); + let owner = TransactionOwner::spawn(io, config()); + owner.acquire().await.unwrap().begin().await.unwrap(); + + let failure = owner.acquire().await.unwrap().commit().await.unwrap_err(); + + assert_eq!(failure.code, "TRANSACTION_STATE_UNKNOWN"); + assert_eq!(failure.mysql_code, Some(1180)); + assert_eq!(failure.state, TransactionState::Unknown); + assert_eq!(failure.outcome, Some(TransactionOutcome::Unknown)); + assert_eq!(probe.ping_calls.load(Ordering::SeqCst), 0); + + let close = owner.close().await.unwrap(); + assert_eq!(close.state, TransactionState::Unknown); + assert_eq!(close.outcome, Some(TransactionOutcome::Unknown)); + assert_eq!(probe.ping_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn one_owner_serializes_operations_while_distinct_owners_run_independently() { + let release_a = Arc::new(Notify::new()); + let release_b = Arc::new(Notify::new()); + let (io_a, probe_a) = ScriptedIo::new([Step::Wait(release_a.clone()), Step::Success { in_transaction: true }]); + let (io_b, probe_b) = ScriptedIo::new([Step::Wait(release_b.clone())]); + let owner_a = TransactionOwner::spawn(io_a, config()); + let owner_b = TransactionOwner::spawn(io_b, config()); + + let a1 = { + let owner = owner_a.clone(); + tokio::spawn(async move { owner.acquire().await.unwrap().query("SELECT 1", None).await }) + }; + while probe_a.active_calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + let a2 = { + let owner = owner_a.clone(); + tokio::spawn(async move { owner.acquire().await.unwrap().query("SELECT 2", None).await }) + }; + let b1 = { + let owner = owner_b.clone(); + tokio::spawn(async move { owner.acquire().await.unwrap().query("SELECT 3", None).await }) + }; + while probe_b.active_calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + + assert_eq!(probe_a.max_active_calls.load(Ordering::SeqCst), 1); + assert_eq!(probe_b.max_active_calls.load(Ordering::SeqCst), 1); + assert_eq!(probe_a.executed_sql.lock().unwrap().as_slice(), ["SELECT 1"]); + + release_b.notify_one(); + b1.await.unwrap().unwrap(); + release_a.notify_one(); + a1.await.unwrap().unwrap(); + a2.await.unwrap().unwrap(); + assert_eq!(probe_a.max_active_calls.load(Ordering::SeqCst), 1); + assert_eq!(probe_a.executed_sql.lock().unwrap().as_slice(), ["SELECT 1", "SELECT 2"]); + } + + #[tokio::test] + async fn dropping_an_in_flight_request_terminates_the_owner_and_disposes_the_connection() { + let never_release = Arc::new(Notify::new()); + let (io, probe) = ScriptedIo::new([Step::Wait(never_release)]); + let owner = TransactionOwner::spawn(io, config()); + let request = { + let owner = owner.clone(); + tokio::spawn(async move { owner.acquire().await.unwrap().query("SELECT SLEEP(30)", None).await }) + }; + while probe.active_calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + + request.abort(); + assert!(request.await.unwrap_err().is_cancelled()); + let later = tokio::time::timeout(Duration::from_millis(250), async { + match owner.acquire().await { + Ok(lease) => lease.query("SELECT 2", None).await.unwrap_err(), + Err(failure) => failure, + } + }) + .await + .expect("dropped caller must terminate the owner"); + + assert_eq!(later.state, TransactionState::Unknown); + assert_eq!(later.outcome, Some(TransactionOutcome::Unknown)); + assert_eq!(probe.disconnects.load(Ordering::SeqCst), 1); + assert_eq!(probe.executed_sql.lock().unwrap().as_slice(), ["SELECT SLEEP(30)"]); + } + + #[tokio::test] + async fn dropping_a_request_during_server_error_probe_disposes_the_connection() { + let ping_entered = Arc::new(Notify::new()); + let never_release = Arc::new(Notify::new()); + let (io, probe) = ScriptedIo::new([ + Step::ServerError { code: 1062, message: "Duplicate entry" }, + Step::WaitPing { entered: ping_entered.clone(), release: never_release, in_transaction: true }, + ]); + let owner = TransactionOwner::spawn(io, config()); + let request = { + let owner = owner.clone(); + tokio::spawn(async move { owner.acquire().await.unwrap().query("INSERT duplicate", None).await }) + }; + ping_entered.notified().await; + + request.abort(); + assert!(request.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_millis(250), owner.wait_closed()) + .await + .expect("dropped caller must cancel a busy status probe"); + + assert_eq!(owner.status().state, TransactionState::Unknown); + assert_eq!(probe.disconnects.load(Ordering::SeqCst), 1); + assert_eq!(probe.executed_sql.lock().unwrap().as_slice(), ["INSERT duplicate"]); + } + + #[tokio::test] + async fn active_idle_expiry_rolls_back_and_disposes_the_connection() { + let (io, probe) = + ScriptedIo::new([Step::Success { in_transaction: true }, Step::Success { in_transaction: false }]); + let owner = TransactionOwner::spawn( + io, + TransactionOwnerConfig { + idle_ttl: Duration::from_millis(20), + operation_timeout: Duration::from_secs(1), + cleanup_timeout: Duration::from_millis(20), + }, + ); + owner.acquire().await.unwrap().begin().await.unwrap(); + + tokio::time::timeout(Duration::from_millis(250), async { + while probe.disconnects.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("idle owner must be disposed without another request"); + + assert_eq!(probe.executed_sql.lock().unwrap().as_slice(), ["START TRANSACTION", "ROLLBACK"]); + let status = owner.status(); + assert_eq!(status.state, TransactionState::Idle); + assert_eq!(status.outcome, Some(TransactionOutcome::RolledBack)); + } + + #[tokio::test] + async fn close_waits_until_bounded_active_cleanup_releases_the_owner() { + let rollback_release = Arc::new(Notify::new()); + let (io, probe) = ScriptedIo::new([Step::Success { in_transaction: true }, Step::Wait(rollback_release)]); + let owner = TransactionOwner::spawn( + io, + TransactionOwnerConfig { + idle_ttl: Duration::from_secs(60), + operation_timeout: Duration::from_secs(1), + cleanup_timeout: Duration::from_millis(20), + }, + ); + owner.acquire().await.unwrap().begin().await.unwrap(); + + let status = owner.close().await.expect("owner cleanup must reach its bounded disposal path"); + + assert_eq!(status.state, TransactionState::Unknown); + assert_eq!(probe.disconnects.load(Ordering::SeqCst), 1); + assert!(owner.inner.closed.load(Ordering::Acquire)); + } + + #[tokio::test] + async fn resource_permit_is_held_until_pending_owner_cleanup_finishes() { + let rollback_release = Arc::new(Notify::new()); + let budget = Arc::new(Semaphore::new(1)); + let permit = budget.clone().acquire_owned().await.unwrap(); + let (io, probe) = + ScriptedIo::new([Step::Success { in_transaction: true }, Step::Wait(rollback_release.clone())]); + let owner = TransactionOwner::spawn_with_resource_permit(io, config(), permit); + owner.acquire().await.unwrap().begin().await.unwrap(); + + let closing = { + let owner = owner.clone(); + tokio::spawn(async move { owner.close().await }) + }; + while probe.active_calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + assert!(budget.clone().try_acquire_owned().is_err()); + + rollback_release.notify_one(); + closing.await.unwrap().unwrap(); + assert!(budget.clone().try_acquire_owned().is_ok()); + } + + #[tokio::test] + async fn operation_timeout_is_terminal_unknown_and_disposes_the_connection() { + let never_release = Arc::new(Notify::new()); + let (io, probe) = ScriptedIo::new([Step::Wait(never_release)]); + let owner = TransactionOwner::spawn( + io, + TransactionOwnerConfig { + idle_ttl: Duration::from_secs(60), + operation_timeout: Duration::from_millis(20), + cleanup_timeout: Duration::from_millis(20), + }, + ); + + let failure = owner.acquire().await.unwrap().query("SELECT SLEEP(30)", None).await.unwrap_err(); + + assert_eq!(failure.code, "TRANSACTION_TIMEOUT"); + assert_eq!(failure.state, TransactionState::Unknown); + owner.wait_closed().await; + assert_eq!(probe.disconnects.load(Ordering::SeqCst), 1); + assert_eq!(probe.executed_sql.lock().unwrap().as_slice(), ["SELECT SLEEP(30)"]); + } + + #[tokio::test] + async fn cancelling_a_queued_operation_never_sends_its_sql() { + let release_first = Arc::new(Notify::new()); + let (io, probe) = ScriptedIo::new([ + Step::Wait(release_first.clone()), + Step::Success { in_transaction: true }, + Step::Success { in_transaction: false }, + ]); + let owner = TransactionOwner::spawn(io, config()); + let first = { + let owner = owner.clone(); + tokio::spawn(async move { owner.acquire().await.unwrap().query("SELECT 1", None).await }) + }; + while probe.active_calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + let queued = { + let owner = owner.clone(); + tokio::spawn(async move { owner.acquire().await.unwrap().query("SELECT 2", None).await }) + }; + tokio::task::yield_now().await; + queued.abort(); + assert!(queued.await.unwrap_err().is_cancelled()); + + release_first.notify_one(); + first.await.unwrap().unwrap(); + owner.acquire().await.unwrap().query("SELECT 3", None).await.unwrap(); + owner.close().await.unwrap(); + + assert_eq!(probe.executed_sql.lock().unwrap().as_slice(), ["SELECT 1", "SELECT 3", "ROLLBACK"]); + } + + #[tokio::test] + async fn server_error_probe_can_report_an_implicit_rollback() { + let (io, probe) = ScriptedIo::new([ + Step::Success { in_transaction: true }, + Step::ServerError { code: 1213, message: "Deadlock found" }, + Step::Ping { in_transaction: false }, + ]); + let owner = TransactionOwner::spawn(io, config()); + owner.acquire().await.unwrap().begin().await.unwrap(); + + let failure = owner.acquire().await.unwrap().query("UPDATE t SET value = 1", None).await.unwrap_err(); + + assert_eq!(failure.mysql_code, Some(1213)); + assert_eq!(failure.state, TransactionState::Idle); + assert_eq!(failure.outcome, Some(TransactionOutcome::RolledBack)); + let commit = owner.acquire().await.unwrap().commit().await.unwrap_err(); + assert_eq!(commit.code, "TRANSACTION_NOT_ACTIVE"); + assert_eq!(probe.executed_sql.lock().unwrap().as_slice(), ["START TRANSACTION", "UPDATE t SET value = 1"]); + } + + #[tokio::test] + async fn registry_invalidation_rolls_back_active_owners_without_holding_the_registry_lock() { + let (io, probe) = + ScriptedIo::new([Step::Success { in_transaction: true }, Step::Success { in_transaction: false }]); + let owner = TransactionOwner::spawn(io, config()); + let registry = TransactionOwnerRegistry::default(); + registry.register("connection-1", &owner); + owner.acquire().await.unwrap().begin().await.unwrap(); + + registry.invalidate_connection("connection-1"); + owner.wait_closed().await; + + assert_eq!(probe.executed_sql.lock().unwrap().as_slice(), ["START TRANSACTION", "ROLLBACK"]); + assert_eq!(probe.disconnects.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn dropping_the_last_public_handle_disposes_an_idle_owner() { + let (io, probe) = ScriptedIo::new([]); + let owner = TransactionOwner::spawn(io, config()); + + drop(owner); + + tokio::time::timeout(Duration::from_millis(250), async { + while probe.disconnects.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping the protocol-owned handle must dispose its connection"); + } + + #[tokio::test] + async fn batch_continues_after_recoverable_error_but_stops_on_unknown() { + let (recoverable_io, recoverable_probe) = ScriptedIo::new([ + Step::ServerError { code: 1062, message: "Duplicate entry" }, + Step::Ping { in_transaction: true }, + Step::Success { in_transaction: true }, + Step::Success { in_transaction: false }, + ]); + let recoverable = TransactionOwner::spawn(recoverable_io, config()); + let results = recoverable + .acquire() + .await + .unwrap() + .batch(vec!["INSERT duplicate".to_string(), "SELECT winner".to_string()], None, true) + .await; + assert_eq!(results.len(), 2); + assert!(results[0].is_err()); + assert!(results[1].is_ok()); + recoverable.close().await.unwrap(); + assert_eq!( + recoverable_probe.executed_sql.lock().unwrap().as_slice(), + ["INSERT duplicate", "SELECT winner", "ROLLBACK"] + ); + + let (unknown_io, unknown_probe) = ScriptedIo::new([Step::TransportError("connection reset")]); + let unknown = TransactionOwner::spawn(unknown_io, config()); + let results = unknown + .acquire() + .await + .unwrap() + .batch(vec!["UPDATE first".to_string(), "UPDATE forbidden".to_string()], None, true) + .await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].as_ref().unwrap_err().state, TransactionState::Unknown); + unknown.wait_closed().await; + assert_eq!(unknown_probe.executed_sql.lock().unwrap().as_slice(), ["UPDATE first"]); + } +} diff --git a/crates/dbx-mcp/tests/protocol.rs b/crates/dbx-mcp/tests/protocol.rs index c0ceb98dde..83b03d905f 100644 --- a/crates/dbx-mcp/tests/protocol.rs +++ b/crates/dbx-mcp/tests/protocol.rs @@ -328,9 +328,9 @@ async fn initializes_lists_tools_and_calls_a_tool() { let tools = client.peer().list_tools(None).await.expect("list tools"); let names = tools.tools.iter().map(|tool| tool.name.as_ref()).collect::>(); #[cfg(feature = "mq-admin")] - assert_eq!(names.len(), 19); + assert_eq!(names.len(), 22); #[cfg(not(feature = "mq-admin"))] - assert_eq!(names.len(), 17); + assert_eq!(names.len(), 20); #[cfg(feature = "mq-admin")] assert!(names.contains(&"dbx_peek_messages")); #[cfg(not(feature = "mq-admin"))] @@ -344,6 +344,9 @@ async fn initializes_lists_tools_and_calls_a_tool() { assert!(names.contains(&"dbx_list_routines")); assert!(names.contains(&"dbx_get_routine_source")); assert!(names.contains(&"dbx_open_session")); + assert!(names.contains(&"dbx_begin_transaction")); + assert!(names.contains(&"dbx_commit_transaction")); + assert!(names.contains(&"dbx_rollback_transaction")); assert!(names.contains(&"dbx_close_session")); #[cfg(feature = "mq-admin")] assert!(names.contains(&"dbx_send_message")); diff --git a/crates/dbx-mcp/tests/session_transactions_mysql.rs b/crates/dbx-mcp/tests/session_transactions_mysql.rs new file mode 100644 index 0000000000..35c69f578f --- /dev/null +++ b/crates/dbx-mcp/tests/session_transactions_mysql.rs @@ -0,0 +1,1188 @@ +//! Native MySQL 5.7 contention evidence for MCP fixed-session transactions. +//! +//! This ignored test drives the public MCP tools through the production +//! `LocalBackend`. It creates an isolated temporary DBX storage file containing +//! only the candidate connection supplied by environment variables; it never +//! opens or copies the installed application's `dbx.db`. +//! +//! Preferred private environment variable: +//! +//! ```text +//! DBX_TXN_MYSQL_CONNECTION_JSON +//! ``` +//! +//! Its value is one already-authorized native MySQL `ConnectionConfig` JSON, +//! including the existing TLS/tunnel settings. An optional +//! `DBX_TXN_MYSQL_FIXTURE_CONNECTION_JSON` may supply a separately authorized +//! fixture configuration; otherwise the candidate configuration is cloned into +//! a distinct logical fixture channel. No account creation or permission +//! widening is performed. The candidate needs SELECT/INSERT/UPDATE/DELETE/ +//! REPLACE; the fixture channel must already be allowed to CREATE/DROP only the +//! exact uniquely prefixed table used by this run. Secrets, endpoint details, +//! and the database name are never printed. +//! A runner that has already verified the saved connection is direct plaintext +//! may instead set `DBX_TXN_MYSQL_DIRECT_PLAINTEXT=1` plus the HOST/PORT/ +//! DATABASE/USER/PASSWORD fields and optional `DBX_TXN_MYSQL_DRIVER_PROFILE`. + +use std::{sync::Arc, time::Duration}; + +use dbx_core::{models::connection::ConnectionConfig, storage::McpGlobalPolicy, storage::Storage}; +use dbx_mcp::{DbxMcpServer, LocalBackend, McpScope}; +use rmcp::{ + model::{CallToolRequest, CallToolRequestParams, ClientRequest}, + service::PeerRequestOptions, + ServiceExt, +}; +use serde_json::{json, Map, Value}; + +mod support; +use support::mysql_wire_proxy::MysqlWireProxy; + +fn required_env(name: &str) -> Result { + std::env::var(name).map_err(|_| format!("set {name} through the private live-test environment")) +} + +fn candidate_connection_from_env() -> Result { + if let Ok(value) = std::env::var("DBX_TXN_MYSQL_CONNECTION_JSON") { + return serde_json::from_str(&value) + .map_err(|error| format!("DBX_TXN_MYSQL_CONNECTION_JSON is not a valid ConnectionConfig JSON: {error}")); + } + if std::env::var("DBX_TXN_MYSQL_DIRECT_PLAINTEXT").as_deref() != Ok("1") { + return Err( + "field-based credentials require DBX_TXN_MYSQL_DIRECT_PLAINTEXT=1 after the runner verifies that the authorized saved connection uses no TLS, certificate, tunnel, plugin, JDBC, or URL parameters" + .to_string(), + ); + } + let profile = std::env::var("DBX_TXN_MYSQL_DRIVER_PROFILE") + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + serde_json::from_value(json!({ + "id": "mcp-session-transaction-live", + "name": "mcp-session-transaction-live", + "db_type": "mysql", + "driver_profile": profile, + "host": required_env("DBX_TXN_MYSQL_HOST")?, + "port": required_env("DBX_TXN_MYSQL_PORT")?.parse::().map_err(|_| "DBX_TXN_MYSQL_PORT must be numeric".to_string())?, + "username": required_env("DBX_TXN_MYSQL_USER")?, + "password": required_env("DBX_TXN_MYSQL_PASSWORD")?, + "database": required_env("DBX_TXN_MYSQL_DATABASE")?, + "ssl": false + })) + .map_err(|error| format!("build direct-plaintext candidate connection: {error}")) +} + +fn is_builtin_native_mysql(connection: &ConnectionConfig) -> bool { + connection.db_type == dbx_core::models::connection::DatabaseType::Mysql + && connection.driver_profile.as_deref().is_none_or(|profile| { + let profile = profile.trim(); + profile.is_empty() || profile.eq_ignore_ascii_case("mysql") + }) +} + +fn arguments(value: Value) -> Map { + value.as_object().cloned().unwrap_or_default() +} + +fn result_error_code(result: &rmcp::model::CallToolResult) -> Option<&str> { + result.content.iter().find_map(|content| { + let text = content.as_text()?.text.strip_prefix("Error [")?; + text.split_once(']').map(|(code, _)| code) + }) +} + +fn redacted_result_diagnostics(result: &rmcp::model::CallToolResult) -> String { + let structured_keys = result.structured_content.as_ref().and_then(Value::as_object).map(|object| { + let mut keys = object.keys().map(String::as_str).collect::>(); + keys.sort_unstable(); + keys + }); + let error_code = result_error_code(result); + let transaction_state = + result.structured_content.as_ref().and_then(|value| value.get("transaction_state")).and_then(Value::as_str); + let mysql_code = + result.structured_content.as_ref().and_then(|value| value.get("mysql_code")).and_then(Value::as_u64); + format!( + "is_error={:?}, structured_keys={structured_keys:?}, error_code={error_code:?}, transaction_state={transaction_state:?}, mysql_code={mysql_code:?}", + result.is_error + ) +} + +fn structured<'a>(result: &'a rmcp::model::CallToolResult, operation: &str) -> Result<&'a Value, String> { + result.structured_content.as_ref().ok_or_else(|| { + format!("{operation} response omitted structured content ({})", redacted_result_diagnostics(result)) + }) +} + +fn expect_success(result: &rmcp::model::CallToolResult, operation: &str) -> Result<(), String> { + if result.is_error == Some(true) { + Err(format!("{operation} failed ({})", redacted_result_diagnostics(result))) + } else { + Ok(()) + } +} + +fn expect_state(result: &rmcp::model::CallToolResult, operation: &str, state: &str) -> Result<(), String> { + let actual = structured(result, operation)?.get("transaction_state").and_then(Value::as_str); + if actual == Some(state) { + Ok(()) + } else { + Err(format!( + "{operation} expected transaction_state={state}, got {actual:?} ({})", + redacted_result_diagnostics(result) + )) + } +} + +fn result_cell<'a>( + result: &'a rmcp::model::CallToolResult, + operation: &str, + row: usize, + column: usize, +) -> Result<&'a Value, String> { + structured(result, operation)?.pointer(&format!("/result/rows/{row}/{column}")).ok_or_else(|| { + format!("{operation} missing result cell {row}/{column} ({})", redacted_result_diagnostics(result)) + }) +} + +#[test] +fn result_diagnostics_expose_only_redacted_envelope_metadata() { + let mut result = rmcp::model::CallToolResult::error(vec![rmcp::model::ContentBlock::text( + "Error [TRANSACTION_CLOSED]: sensitive endpoint and database details", + )]); + result.structured_content = Some(json!({ + "session_id": "sensitive-session-id", + "transaction_state": "idle", + "mysql_code": 2013, + "error": "sensitive backend message" + })); + + let diagnostic = redacted_result_diagnostics(&result); + + assert_eq!( + diagnostic, + "is_error=Some(true), structured_keys=Some([\"error\", \"mysql_code\", \"session_id\", \"transaction_state\"]), error_code=Some(\"TRANSACTION_CLOSED\"), transaction_state=Some(\"idle\"), mysql_code=Some(2013)" + ); + assert!(!diagnostic.contains("sensitive")); +} + +macro_rules! call_tool { + ($peer:expr, $name:literal, $arguments:expr) => { + $peer + .call_tool(CallToolRequestParams::new($name).with_arguments(arguments($arguments))) + .await + .map_err(|error| format!("{} transport failure: {error}", $name)) + }; +} + +async fn read_fixture_cell( + fixture: &rmcp::service::RunningService, + database: &str, + sql: String, + operation: &str, +) -> Result { + let opened = call_tool!( + fixture.peer(), + "dbx_open_session", + json!({ + "connection_id": "mcp-session-transaction-fixture", + "database": database, + "enable_transactions": true + }) + )?; + expect_success(&opened, &format!("open fixture session for {operation}"))?; + let session = structured(&opened, &format!("open fixture session for {operation}"))?["session_id"] + .as_str() + .ok_or_else(|| format!("fixture session for {operation} omitted session_id"))? + .to_string(); + + let read = async { + let result = call_tool!( + fixture.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-fixture", + "database": database, + "session_id": session.clone(), + "sql": sql + }) + )?; + expect_success(&result, operation)?; + Ok::(result_cell(&result, operation, 0, 0)?.clone()) + } + .await; + + let closed = call_tool!(fixture.peer(), "dbx_close_session", json!({"session_id": session}))?; + expect_success(&closed, &format!("close fixture session for {operation}"))?; + read +} + +async fn run_contention_trace( + client: &rmcp::service::RunningService, + database: &str, + table: &str, + wire_proxy: Option<&MysqlWireProxy>, +) -> Result<(), String> { + let peer = client.peer(); + let connection = "mcp-session-transaction-live"; + let open_a = call_tool!( + peer, + "dbx_open_session", + json!({"connection_id": connection, "database": database, "enable_transactions": true}) + )?; + let open_b = call_tool!( + peer, + "dbx_open_session", + json!({"connection_id": connection, "database": database, "enable_transactions": true}) + )?; + expect_success(&open_a, "open A")?; + expect_success(&open_b, "open B")?; + expect_state(&open_a, "open A", "idle")?; + expect_state(&open_b, "open B", "idle")?; + let session_a = structured(&open_a, "open A")?["session_id"] + .as_str() + .ok_or_else(|| "open A omitted session_id".to_string())? + .to_string(); + let session_b = structured(&open_b, "open B")?["session_id"] + .as_str() + .ok_or_else(|| "open B omitted session_id".to_string())? + .to_string(); + + let id_a_before = call_tool!( + peer, + "dbx_execute_query", + json!({"connection_id": connection, "database": database, "session_id": session_a, "sql": "SELECT CONNECTION_ID() AS connection_id"}) + )?; + let id_b_before = call_tool!( + peer, + "dbx_execute_query", + json!({"connection_id": connection, "database": database, "session_id": session_b, "sql": "SELECT CONNECTION_ID() AS connection_id"}) + )?; + expect_success(&id_a_before, "read A physical connection ID before contention")?; + expect_success(&id_b_before, "read B physical connection ID before contention")?; + let id_a_before = result_cell(&id_a_before, "read A physical connection ID before contention", 0, 0)?.clone(); + let id_b_before = result_cell(&id_b_before, "read B physical connection ID before contention", 0, 0)?.clone(); + if id_a_before == id_b_before { + return Err("A and B unexpectedly share one physical MySQL connection".to_string()); + } + + let begin_b = call_tool!(peer, "dbx_begin_transaction", json!({"session_id": session_b}))?; + expect_success(&begin_b, "B.begin")?; + expect_state(&begin_b, "B.begin", "active")?; + let b_absent = call_tool!( + peer, + "dbx_execute_query", + json!({ + "connection_id": connection, + "database": database, + "session_id": session_b, + "sql": format!("SELECT payload FROM {table} WHERE id = 1") + }) + )?; + expect_success(&b_absent, "B initial snapshot read")?; + expect_state(&b_absent, "B initial snapshot read", "active")?; + let rows = structured(&b_absent, "B initial snapshot read")? + .pointer("/result/rows") + .and_then(Value::as_array) + .ok_or_else(|| "B initial read omitted rows".to_string())?; + if !rows.is_empty() { + return Err("B initial snapshot unexpectedly found the fixture key".to_string()); + } + + let begin_a = call_tool!(peer, "dbx_begin_transaction", json!({"session_id": session_a}))?; + expect_success(&begin_a, "A.begin")?; + expect_state(&begin_a, "A.begin", "active")?; + let a_insert = call_tool!( + peer, + "dbx_execute_query", + json!({ + "connection_id": connection, + "database": database, + "session_id": session_a, + "sql": format!("INSERT INTO {table}(id, payload) VALUES (1, 'exact_winner')") + }) + )?; + expect_success(&a_insert, "A.insert")?; + expect_state(&a_insert, "A.insert", "active")?; + + let pending_peer = peer.clone(); + let pending_database = database.to_string(); + let pending_table = table.to_string(); + let pending_session = session_b.clone(); + if let Some(proxy) = wire_proxy { + proxy.set_contention_marker("VALUES (1, 'loser')").await; + } + let mut b_insert = tokio::spawn(async move { + pending_peer + .call_tool(CallToolRequestParams::new("dbx_execute_query").with_arguments(arguments(json!({ + "connection_id": connection, + "database": pending_database, + "session_id": pending_session, + "sql": format!("INSERT INTO {pending_table}(id, payload) VALUES (1, 'loser')") + })))) + .await + }); + if tokio::time::timeout(Duration::from_millis(300), &mut b_insert).await.is_ok() { + return Err("B.insert completed before A.commit; contention was not established".to_string()); + } + if let Some(proxy) = wire_proxy { + if !proxy.contention_forwarded_and_pending() { + return Err("wire proxy did not prove that B.insert was forwarded with no server response".to_string()); + } + } + + let commit_a = call_tool!(peer, "dbx_commit_transaction", json!({"session_id": session_a}))?; + expect_success(&commit_a, "A.commit")?; + expect_state(&commit_a, "A.commit", "idle")?; + if structured(&commit_a, "A.commit")?["transaction_outcome"] != "committed" { + return Err("A.commit did not report committed".to_string()); + } + let b_duplicate = tokio::time::timeout(Duration::from_secs(5), b_insert) + .await + .map_err(|_| "B.insert did not finish after A.commit".to_string())? + .map_err(|error| format!("B.insert task failed: {error}"))? + .map_err(|error| format!("B.insert transport failed: {error}"))?; + if b_duplicate.is_error != Some(true) { + return Err("B.insert unexpectedly succeeded".to_string()); + } + expect_state(&b_duplicate, "B.insert duplicate response", "active")?; + if structured(&b_duplicate, "B.insert duplicate response")?["mysql_code"] != 1062 { + return Err("B.insert did not preserve MySQL 1062".to_string()); + } + + let b_winner = call_tool!( + peer, + "dbx_execute_query", + json!({ + "connection_id": connection, + "database": database, + "session_id": session_b, + "sql": format!("SELECT payload FROM {table} WHERE id = 1 FOR UPDATE") + }) + )?; + expect_success(&b_winner, "B locking read")?; + expect_state(&b_winner, "B locking read", "active")?; + if result_cell(&b_winner, "B locking read", 0, 0)?.as_str() != Some("exact_winner") { + return Err("B locking read did not return A's exact winner".to_string()); + } + + let id_a_after = call_tool!( + peer, + "dbx_execute_query", + json!({"connection_id": connection, "database": database, "session_id": session_a, "sql": "SELECT CONNECTION_ID() AS connection_id"}) + )?; + let id_b_during = call_tool!( + peer, + "dbx_execute_query", + json!({"connection_id": connection, "database": database, "session_id": session_b, "sql": "SELECT CONNECTION_ID() AS connection_id"}) + )?; + expect_success(&id_a_after, "read A physical connection ID after commit")?; + expect_success(&id_b_during, "read B physical connection ID during transaction")?; + if result_cell(&id_a_after, "read A physical connection ID after commit", 0, 0)? != &id_a_before { + return Err("A physical connection identity changed".to_string()); + } + if result_cell(&id_b_during, "read B physical connection ID during transaction", 0, 0)? != &id_b_before { + return Err("B physical connection identity changed".to_string()); + } + + let rollback_b = call_tool!(peer, "dbx_rollback_transaction", json!({"session_id": session_b}))?; + expect_success(&rollback_b, "B.rollback")?; + expect_state(&rollback_b, "B.rollback", "idle")?; + if structured(&rollback_b, "B.rollback")?["transaction_outcome"] != "rolled_back" { + return Err("B.rollback did not report rolled_back".to_string()); + } + let close_a = call_tool!(peer, "dbx_close_session", json!({"session_id": session_a}))?; + let close_b = call_tool!(peer, "dbx_close_session", json!({"session_id": session_b}))?; + expect_success(&close_a, "close A")?; + expect_success(&close_b, "close B")?; + Ok(()) +} + +async fn run_native_cleanup_traces( + candidate: &rmcp::service::RunningService, + fixture: &rmcp::service::RunningService, + database: &str, + table: &str, +) -> Result<(), String> { + let connection = "mcp-session-transaction-live"; + let fixture_connection = "mcp-session-transaction-fixture"; + + let seed_close = call_tool!( + fixture.peer(), + "dbx_execute_query", + json!({ + "connection_id": fixture_connection, + "database": database, + "sql": format!("INSERT INTO {table}(id, payload) VALUES (2, 'close_original')") + }) + )?; + expect_success(&seed_close, "seed close fixture row")?; + let opened = call_tool!( + candidate.peer(), + "dbx_open_session", + json!({"connection_id": connection, "database": database, "enable_transactions": true}) + )?; + expect_success(&opened, "open close-test session")?; + let session = structured(&opened, "open close-test session")?["session_id"] + .as_str() + .ok_or_else(|| "close-test open omitted session_id".to_string())? + .to_string(); + expect_success( + &call_tool!(candidate.peer(), "dbx_begin_transaction", json!({"session_id": session}))?, + "begin close-test transaction", + )?; + expect_success( + &call_tool!( + candidate.peer(), + "dbx_execute_query", + json!({ + "connection_id": connection, + "database": database, + "session_id": session, + "sql": format!("UPDATE {table} SET payload = 'uncommitted_close' WHERE id = 2") + }) + )?, + "hold close-test row lock", + )?; + let fixture_peer = fixture.peer().clone(); + let fixture_database = database.to_string(); + let fixture_table = table.to_string(); + let mut blocked_update = tokio::spawn(async move { + fixture_peer + .call_tool(CallToolRequestParams::new("dbx_execute_query").with_arguments(arguments(json!({ + "connection_id": fixture_connection, + "database": fixture_database, + "sql": format!("UPDATE {fixture_table} SET payload = 'close_released' WHERE id = 2") + })))) + .await + }); + if tokio::time::timeout(Duration::from_millis(300), &mut blocked_update).await.is_ok() { + return Err("fixture update did not block on the close-test transaction".to_string()); + } + let close = call_tool!(candidate.peer(), "dbx_close_session", json!({"session_id": session}))?; + expect_success(&close, "close active transaction session")?; + let released = tokio::time::timeout(Duration::from_secs(10), blocked_update) + .await + .map_err(|_| "row lock was not released after dbx_close_session".to_string())? + .map_err(|error| format!("close-test fixture task failed: {error}"))? + .map_err(|error| format!("close-test fixture transport failed: {error}"))?; + expect_success(&released, "fixture update after close")?; + + let seed_cancel = call_tool!( + fixture.peer(), + "dbx_execute_query", + json!({ + "connection_id": fixture_connection, + "database": database, + "sql": format!("INSERT INTO {table}(id, payload) VALUES (3, 'cancel_original')") + }) + )?; + expect_success(&seed_cancel, "seed cancellation fixture row")?; + let opened = call_tool!( + candidate.peer(), + "dbx_open_session", + json!({"connection_id": connection, "database": database, "enable_transactions": true}) + )?; + expect_success(&opened, "open cancellation-test session")?; + let session = structured(&opened, "open cancellation-test session")?["session_id"] + .as_str() + .ok_or_else(|| "cancellation-test open omitted session_id".to_string())? + .to_string(); + expect_success( + &call_tool!(candidate.peer(), "dbx_begin_transaction", json!({"session_id": session}))?, + "begin cancellation-test transaction", + )?; + let request = ClientRequest::CallToolRequest(CallToolRequest::new( + CallToolRequestParams::new("dbx_execute_query").with_arguments(arguments(json!({ + "connection_id": connection, + "database": database, + "session_id": session, + "sql": format!("SELECT SLEEP(30), payload FROM {table} WHERE id = 3 FOR UPDATE") + }))), + )); + let request = candidate + .peer() + .send_cancellable_request(request, PeerRequestOptions::no_options()) + .await + .map_err(|error| format!("send cancellable locking query: {error}"))?; + tokio::time::sleep(Duration::from_millis(300)).await; + + let fixture_peer = fixture.peer().clone(); + let fixture_database = database.to_string(); + let fixture_table = table.to_string(); + let mut blocked_update = tokio::spawn(async move { + fixture_peer + .call_tool(CallToolRequestParams::new("dbx_execute_query").with_arguments(arguments(json!({ + "connection_id": fixture_connection, + "database": fixture_database, + "sql": format!("UPDATE {fixture_table} SET payload = 'cancel_released' WHERE id = 3") + })))) + .await + }); + if tokio::time::timeout(Duration::from_millis(300), &mut blocked_update).await.is_ok() { + return Err("fixture update did not block on the cancellable locking query".to_string()); + } + request + .cancel(Some("native cancellation lock-release probe".to_string())) + .await + .map_err(|error| format!("cancel locking query: {error}"))?; + let released = tokio::time::timeout(Duration::from_secs(10), blocked_update) + .await + .map_err(|_| "row lock was not released after MCP cancellation".to_string())? + .map_err(|error| format!("cancellation fixture task failed: {error}"))? + .map_err(|error| format!("cancellation fixture transport failed: {error}"))?; + expect_success(&released, "fixture update after cancellation")?; + let close = call_tool!(candidate.peer(), "dbx_close_session", json!({"session_id": session}))?; + expect_success(&close, "close terminal cancellation-test session")?; + Ok(()) +} + +async fn run_fault_proxy_traces( + candidate: &rmcp::service::RunningService, + fixture: &rmcp::service::RunningService, + database: &str, + table: &str, + proxy: &MysqlWireProxy, +) -> Result<(), String> { + let connection = "mcp-session-transaction-live"; + let open = call_tool!( + candidate.peer(), + "dbx_open_session", + json!({"connection_id": connection, "database": database, "enable_transactions": true}) + )?; + expect_success(&open, "open lost-commit-ack session")?; + let session = structured(&open, "open lost-commit-ack session")?["session_id"] + .as_str() + .ok_or_else(|| "lost-commit-ack open omitted session_id".to_string())? + .to_string(); + expect_success( + &call_tool!(candidate.peer(), "dbx_begin_transaction", json!({"session_id": session}))?, + "begin lost-commit-ack transaction", + )?; + expect_success( + &call_tool!( + candidate.peer(), + "dbx_execute_query", + json!({ + "connection_id": connection, + "database": database, + "session_id": session, + "sql": format!("INSERT INTO {table}(id, payload) VALUES (10, 'commit_ack_dropped')") + }) + )?, + "insert before lost COMMIT ACK", + )?; + proxy.arm_drop_commit_ack().await; + let commit = call_tool!(candidate.peer(), "dbx_commit_transaction", json!({"session_id": session}))?; + if commit.is_error != Some(true) { + return Err("COMMIT with a dropped server ACK unexpectedly reported success".to_string()); + } + expect_state(&commit, "COMMIT with dropped server ACK", "unknown")?; + if proxy.commit_forward_count() != 1 || !proxy.commit_ok_observed() || !proxy.commit_ack_dropped() { + return Err( + "wire proxy did not prove exactly one forwarded COMMIT + observed server OK + dropped ACK".to_string() + ); + } + let committed = read_fixture_cell( + fixture, + database, + format!("SELECT payload FROM {table} WHERE id = 10"), + "verify commit after dropped ACK", + ) + .await?; + if committed.as_str() != Some("commit_ack_dropped") { + return Err("fixture channel did not observe the COMMIT whose ACK was dropped".to_string()); + } + + let rejected_retry = + call_tool!(candidate.peer(), "dbx_commit_transaction", json!({"session_id": session.clone()}))?; + if rejected_retry.is_error != Some(true) { + return Err("retrying COMMIT after an unknown outcome unexpectedly succeeded".to_string()); + } + expect_state(&rejected_retry, "retry COMMIT after unknown outcome", "unknown")?; + let close = call_tool!(candidate.peer(), "dbx_close_session", json!({"session_id": session}))?; + expect_success(&close, "close unknown lost-commit-ack session")?; + expect_state(&close, "close unknown lost-commit-ack session", "unknown")?; + if proxy.commit_forward_count() != 1 { + return Err("COMMIT was forwarded more than once across retry refusal or cleanup".to_string()); + } + + let open = call_tool!( + candidate.peer(), + "dbx_open_session", + json!({"connection_id": connection, "database": database, "enable_transactions": true}) + )?; + expect_success(&open, "open non-COMMIT transport-loss session")?; + let session = structured(&open, "open non-COMMIT transport-loss session")?["session_id"] + .as_str() + .ok_or_else(|| "non-COMMIT transport-loss open omitted session_id".to_string())? + .to_string(); + expect_success( + &call_tool!(candidate.peer(), "dbx_begin_transaction", json!({"session_id": session}))?, + "begin non-COMMIT transport-loss transaction", + )?; + proxy.arm_drop_query_response("SELECT SLEEP(2)").await; + let lost = call_tool!( + candidate.peer(), + "dbx_execute_query", + json!({ + "connection_id": connection, + "database": database, + "session_id": session, + "sql": "SELECT SLEEP(2)" + }) + )?; + if lost.is_error != Some(true) { + return Err("non-COMMIT transport loss unexpectedly reported success".to_string()); + } + expect_state(&lost, "non-COMMIT transport loss", "unknown")?; + if !proxy.dropped_query_forwarded() { + return Err("wire proxy did not prove the non-COMMIT command was forwarded".to_string()); + } + let close = call_tool!(candidate.peer(), "dbx_close_session", json!({"session_id": session}))?; + expect_success(&close, "close unknown non-COMMIT transport-loss session")?; + expect_state(&close, "close unknown non-COMMIT transport-loss session", "unknown")?; + Ok(()) +} + +async fn run_permission_refusal_trace( + candidate: &rmcp::service::RunningService, + fixture: &rmcp::service::RunningService, + policy_storage: &Storage, + database: &str, + table: &str, +) -> Result<(), String> { + let writable_policy = || McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: true, + allowed_connection_ids: Some(vec![ + "mcp-session-transaction-live".to_string(), + "mcp-session-transaction-fixture".to_string(), + ]), + ..Default::default() + }; + let seed = call_tool!( + fixture.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-fixture", + "database": database, + "sql": format!("INSERT INTO {table}(id, payload) VALUES (30, 'policy_seed')") + }) + )?; + expect_success(&seed, "seed permission-refusal row")?; + let open = call_tool!( + candidate.peer(), + "dbx_open_session", + json!({ + "connection_id": "mcp-session-transaction-live", + "database": database, + "enable_transactions": true + }) + )?; + expect_success(&open, "open permission-refusal transaction")?; + let session = structured(&open, "open permission-refusal transaction")?["session_id"] + .as_str() + .ok_or_else(|| "permission-refusal open omitted session_id".to_string())? + .to_string(); + expect_success( + &call_tool!(candidate.peer(), "dbx_begin_transaction", json!({"session_id": session}))?, + "begin permission-refusal transaction", + )?; + + let mut read_only_policy = writable_policy(); + read_only_policy.read_only = true; + policy_storage + .save_mcp_global_policy(&read_only_policy) + .await + .map_err(|error| format!("set isolated read-only policy: {error}"))?; + let read_only_check = async { + let rejected = call_tool!( + candidate.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-live", + "database": database, + "session_id": session, + "sql": format!("UPDATE {table} SET payload = 'read_only_bypass' WHERE id = 30") + }) + )?; + if rejected.is_error != Some(true) { + return Err("isolated read-only policy unexpectedly allowed a transaction write".to_string()); + } + let state = call_tool!( + candidate.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-live", + "database": database, + "session_id": session, + "sql": format!("SELECT payload FROM {table} WHERE id = 30") + }) + )?; + expect_success(&state, "read transaction state after read-only refusal")?; + expect_state(&state, "read transaction state after read-only refusal", "active")?; + if result_cell(&state, "read transaction state after read-only refusal", 0, 0)?.as_str() != Some("policy_seed") + { + return Err("read-only refusal changed the row inside the active transaction".to_string()); + } + Result::<(), String>::Ok(()) + } + .await; + policy_storage + .save_mcp_global_policy(&writable_policy()) + .await + .map_err(|error| format!("restore isolated writable policy: {error}"))?; + read_only_check?; + + let previous_confirmation = std::env::var_os("DBX_MCP_CONFIRMED_WRITE_SQL"); + unsafe { + std::env::set_var( + "DBX_MCP_CONFIRMED_WRITE_SQL", + format!("UPDATE {table} SET payload = 'only_exact_sql' WHERE id = 30"), + ); + } + let confirmation_check = async { + let rejected = call_tool!( + candidate.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-live", + "database": database, + "session_id": session, + "sql": format!("UPDATE {table} SET payload = 'different_sql' WHERE id = 30") + }) + )?; + if rejected.is_error != Some(true) { + return Err("mismatched confirmed SQL unexpectedly executed in the transaction".to_string()); + } + let state = call_tool!( + candidate.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-live", + "database": database, + "session_id": session, + "sql": format!("SELECT payload FROM {table} WHERE id = 30") + }) + )?; + expect_success(&state, "read transaction state after confirmed-SQL refusal")?; + expect_state(&state, "read transaction state after confirmed-SQL refusal", "active")?; + if result_cell(&state, "read transaction state after confirmed-SQL refusal", 0, 0)?.as_str() + != Some("policy_seed") + { + return Err("confirmed-SQL refusal changed the row inside the active transaction".to_string()); + } + Result::<(), String>::Ok(()) + } + .await; + unsafe { + match previous_confirmation { + Some(value) => std::env::set_var("DBX_MCP_CONFIRMED_WRITE_SQL", value), + None => std::env::remove_var("DBX_MCP_CONFIRMED_WRITE_SQL"), + } + } + confirmation_check?; + + expect_success( + &call_tool!(candidate.peer(), "dbx_rollback_transaction", json!({"session_id": session}))?, + "rollback permission-refusal transaction", + )?; + expect_success( + &call_tool!(candidate.peer(), "dbx_close_session", json!({"session_id": session}))?, + "close permission-refusal transaction", + )?; + + let unchanged = read_fixture_cell( + fixture, + database, + format!("SELECT payload FROM {table} WHERE id = 30"), + "verify permission refusals left target data unchanged", + ) + .await?; + if unchanged.as_str() != Some("policy_seed") { + return Err("permission refusal trace changed the target fixture row".to_string()); + } + Ok(()) +} + +#[cfg(feature = "transaction-test-hooks")] +async fn run_idle_ttl_trace( + candidate: &rmcp::service::RunningService, + fixture: &rmcp::service::RunningService, + database: &str, + table: &str, +) -> Result<(), String> { + let seeded = call_tool!( + fixture.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-fixture", + "database": database, + "sql": format!("INSERT INTO {table}(id, payload) VALUES (20, 'ttl_seed')") + }) + )?; + expect_success(&seeded, "seed idle-TTL fixture row")?; + + let open = call_tool!( + candidate.peer(), + "dbx_open_session", + json!({ + "connection_id": "mcp-session-transaction-live", + "database": database, + "enable_transactions": true + }) + )?; + expect_success(&open, "open idle-TTL transaction session")?; + let session = structured(&open, "open idle-TTL transaction session")?["session_id"] + .as_str() + .ok_or_else(|| "idle-TTL open omitted session_id".to_string())? + .to_string(); + expect_success( + &call_tool!(candidate.peer(), "dbx_begin_transaction", json!({"session_id": session}))?, + "begin idle-TTL transaction", + )?; + let locked = call_tool!( + candidate.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-live", + "database": database, + "session_id": session, + "sql": format!("UPDATE {table} SET payload = 'ttl_uncommitted' WHERE id = 20") + }) + )?; + expect_success(&locked, "hold idle-TTL row lock")?; + + let fixture_peer = fixture.peer().clone(); + let fixture_database = database.to_string(); + let fixture_table = table.to_string(); + let mut conditional_update = tokio::spawn(async move { + fixture_peer + .call_tool(CallToolRequestParams::new("dbx_execute_query").with_arguments(arguments(json!({ + "connection_id": "mcp-session-transaction-fixture", + "database": fixture_database, + "sql": format!( + "UPDATE {fixture_table} SET payload = 'ttl_released' WHERE id = 20 AND payload = 'ttl_seed'" + ) + })))) + .await + }); + if tokio::time::timeout(Duration::from_millis(100), &mut conditional_update).await.is_ok() { + return Err("idle-TTL fixture update did not remain pending for 100ms".to_string()); + } + let released = tokio::time::timeout(Duration::from_secs(3), conditional_update) + .await + .map_err(|_| "idle-TTL rollback did not release the row lock".to_string())? + .map_err(|error| format!("idle-TTL fixture task failed: {error}"))? + .map_err(|error| format!("idle-TTL fixture transport failed: {error}"))?; + expect_success(&released, "conditional update after idle-TTL rollback")?; + + let expired = call_tool!( + candidate.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-live", + "database": database, + "session_id": session.clone(), + "sql": "SELECT 1" + }) + )?; + if expired.is_error != Some(true) || result_error_code(&expired) != Some("TRANSACTION_CLOSED") { + return Err(format!( + "idle-TTL session accepted SQL after terminal cleanup ({})", + redacted_result_diagnostics(&expired) + )); + } + expect_state(&expired, "query expired idle-TTL session", "idle")?; + + let released_payload = read_fixture_cell( + fixture, + database, + format!("SELECT payload FROM {table} WHERE id = 20"), + "verify idle-TTL rollback released exactly the seeded fixture row", + ) + .await?; + if released_payload.as_str() != Some("ttl_released") { + return Err("idle-TTL conditional fixture update did not change the seeded row exactly once".to_string()); + } + + let closed = call_tool!(candidate.peer(), "dbx_close_session", json!({"session_id": session}))?; + expect_success(&closed, "close idle-TTL session")?; + Ok(()) +} + +#[tokio::test] +#[ignore = "requires an authorized native MySQL ConnectionConfig in DBX_TXN_MYSQL_CONNECTION_JSON"] +async fn native_mysql_fixed_sessions_preserve_identity_and_real_contention() { + async fn run() -> Result<(), String> { + let mut candidate = candidate_connection_from_env()?; + let mut fixture = match std::env::var("DBX_TXN_MYSQL_FIXTURE_CONNECTION_JSON") { + Ok(value) => serde_json::from_str::(&value) + .map_err(|error| format!("DBX_TXN_MYSQL_FIXTURE_CONNECTION_JSON is invalid: {error}"))?, + Err(std::env::VarError::NotPresent) => candidate.clone(), + Err(error) => return Err(format!("cannot read fixture connection environment: {error}")), + }; + if !is_builtin_native_mysql(&candidate) { + return Err( + "candidate must use the native MySQL route (empty profile or the built-in mysql profile)".to_string() + ); + } + if !is_builtin_native_mysql(&fixture) { + return Err("fixture channel must use the native MySQL route".to_string()); + } + let database = candidate + .database + .as_deref() + .map(str::trim) + .filter(|database| !database.is_empty()) + .ok_or_else(|| "candidate must name one non-empty authorized test database".to_string())? + .to_string(); + if fixture.database.as_deref().map(str::trim) != Some(database.as_str()) { + return Err( + "candidate and fixture channels must name the same non-empty authorized test database".to_string() + ); + } + let mut wire_proxy = if std::env::var("DBX_TXN_MYSQL_FAULT_PROXY").as_deref() == Ok("1") { + if std::env::var("DBX_TXN_MYSQL_DIRECT_PLAINTEXT").as_deref() != Ok("1") { + return Err( + "DBX_TXN_MYSQL_FAULT_PROXY=1 also requires the coordinator's DBX_TXN_MYSQL_DIRECT_PLAINTEXT=1 attestation" + .to_string(), + ); + } + if candidate.ssl + || !candidate.transport_layers.is_empty() + || candidate.connection_string.as_deref().is_some_and(|value| !value.trim().is_empty()) + { + return Err("fault proxy requires a direct plaintext candidate without TLS, tunnels, or a custom connection string".to_string()); + } + let proxy = MysqlWireProxy::start(candidate.host.clone(), candidate.port) + .await + .map_err(|error| format!("start loopback MySQL fault proxy: {error}"))?; + candidate.host = proxy.listen_addr().ip().to_string(); + candidate.port = proxy.listen_addr().port(); + Some(proxy) + } else { + None + }; + candidate.id = "mcp-session-transaction-live".to_string(); + candidate.name = candidate.id.clone(); + fixture.id = "mcp-session-transaction-fixture".to_string(); + fixture.name = fixture.id.clone(); + + let directory = tempfile::tempdir().map_err(|error| format!("create isolated storage directory: {error}"))?; + let db_path = directory.path().join("dbx.db"); + let storage = Storage::open(&db_path).await.map_err(|error| format!("open isolated storage: {error}"))?; + storage + .add_connection_for_mcp(candidate) + .await + .map_err(|_| "seed candidate connection failed (details redacted)".to_string())?; + storage + .add_connection_for_mcp(fixture) + .await + .map_err(|_| "seed fixture connection failed (details redacted)".to_string())?; + storage + .save_mcp_global_policy(&McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: true, + allowed_connection_ids: Some(vec![ + "mcp-session-transaction-live".to_string(), + "mcp-session-transaction-fixture".to_string(), + ]), + ..Default::default() + }) + .await + .map_err(|error| format!("seed isolated policy: {error}"))?; + drop(storage); + + let backend = LocalBackend::open(&db_path) + .await + .map_err(|_| "open candidate LocalBackend failed (details redacted)".to_string())?; + #[cfg(feature = "transaction-test-hooks")] + let idle_ttl_backend = LocalBackend::from_app_state(backend.state().clone(), directory.path().to_path_buf()) + .with_transaction_test_config(dbx_mcp::transaction::TransactionOwnerConfig { + idle_ttl: Duration::from_millis(750), + operation_timeout: Duration::from_secs(30), + cleanup_timeout: Duration::from_secs(5), + }); + let backend = Arc::new(backend); + let policy_storage = + Storage::open(&db_path).await.map_err(|error| format!("reopen isolated policy storage: {error}"))?; + let candidate_server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); + let fixture_server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + #[cfg(feature = "transaction-test-hooks")] + let idle_ttl_server = + DbxMcpServer::with_runtime_options(Arc::new(idle_ttl_backend), McpScope::default(), false); + let (candidate_server_transport, candidate_client_transport) = tokio::io::duplex(64 * 1024); + let (fixture_server_transport, fixture_client_transport) = tokio::io::duplex(64 * 1024); + #[cfg(feature = "transaction-test-hooks")] + let (idle_ttl_server_transport, idle_ttl_client_transport) = tokio::io::duplex(64 * 1024); + let candidate_server_task = + tokio::spawn(async move { candidate_server.serve(candidate_server_transport).await }); + let fixture_server_task = tokio::spawn(async move { fixture_server.serve(fixture_server_transport).await }); + #[cfg(feature = "transaction-test-hooks")] + let idle_ttl_server_task = tokio::spawn(async move { idle_ttl_server.serve(idle_ttl_server_transport).await }); + let mut candidate_client = + ().serve(candidate_client_transport) + .await + .map_err(|error| format!("initialize candidate MCP client: {error}"))?; + let mut fixture_client = + ().serve(fixture_client_transport) + .await + .map_err(|error| format!("initialize fixture MCP client: {error}"))?; + #[cfg(feature = "transaction-test-hooks")] + let mut idle_ttl_client = + ().serve(idle_ttl_client_transport) + .await + .map_err(|error| format!("initialize idle-TTL MCP client: {error}"))?; + let mut candidate_service = candidate_server_task + .await + .map_err(|error| format!("candidate server task failed: {error}"))? + .map_err(|error| format!("initialize candidate MCP server: {error}"))?; + let mut fixture_service = fixture_server_task + .await + .map_err(|error| format!("fixture server task failed: {error}"))? + .map_err(|error| format!("initialize fixture MCP server: {error}"))?; + #[cfg(feature = "transaction-test-hooks")] + let mut idle_ttl_service = idle_ttl_server_task + .await + .map_err(|error| format!("idle-TTL server task failed: {error}"))? + .map_err(|error| format!("initialize idle-TTL MCP server: {error}"))?; + + let fixture_open = tokio::time::timeout( + Duration::from_secs(10), + fixture_client.peer().call_tool(CallToolRequestParams::new("dbx_open_session").with_arguments(arguments( + json!({ + "connection_id": "mcp-session-transaction-fixture", + "database": database, + "enable_transactions": true + }), + ))), + ) + .await + .map_err(|_| "fixture capability probe timed out".to_string())? + .map_err(|error| format!("fixture capability probe transport failed: {error}"))?; + expect_success(&fixture_open, "fixture capability session open")?; + let fixture_session = structured(&fixture_open, "fixture capability session open")?["session_id"] + .as_str() + .ok_or_else(|| "fixture session omitted session_id".to_string())? + .to_string(); + let capability = call_tool!( + fixture_client.peer(), + "dbx_execute_query", + json!({ + "connection_id": "mcp-session-transaction-fixture", + "database": database, + "session_id": fixture_session, + "sql": "SELECT VERSION(), @@tx_isolation" + }) + )?; + expect_success(&capability, "read MySQL version/isolation")?; + let version = result_cell(&capability, "read MySQL version/isolation", 0, 0)?.as_str().unwrap_or_default(); + let isolation = result_cell(&capability, "read MySQL version/isolation", 0, 1)?.as_str().unwrap_or_default(); + if !version.starts_with("5.7.") || !isolation.eq_ignore_ascii_case("REPEATABLE-READ") { + return Err("live gate requires MySQL 5.7 with REPEATABLE READ".to_string()); + } + let fixture_close = + call_tool!(fixture_client.peer(), "dbx_close_session", json!({"session_id": fixture_session}))?; + expect_success(&fixture_close, "close fixture capability session")?; + + let table = format!("dbx_mcp_txn_{}", uuid::Uuid::new_v4().simple()); + // Record the exact attributable name before CREATE is dispatched. A + // lost response can still mean the server committed the DDL, so every + // path from this point must attempt the exact-name bounded cleanup. + eprintln!("DBX transaction fixture token: {table}"); + let create = async { + let create = tokio::time::timeout( + Duration::from_secs(10), + fixture_client.peer().call_tool(CallToolRequestParams::new("dbx_execute_query").with_arguments( + arguments(json!({ + "connection_id": "mcp-session-transaction-fixture", + "database": database, + "sql": format!( + "CREATE TABLE {table} (id BIGINT PRIMARY KEY, payload VARCHAR(255) NOT NULL) ENGINE=InnoDB" + ) + })), + )), + ) + .await + .map_err(|_| "fixture CREATE timed out".to_string())? + .map_err(|error| format!("fixture CREATE transport failed: {error}"))?; + expect_success(&create, "create unique InnoDB fixture") + } + .await; + + let scenario = match create { + Ok(()) => match tokio::time::timeout(Duration::from_secs(90), async { + run_contention_trace(&candidate_client, &database, &table, wire_proxy.as_ref()).await?; + run_native_cleanup_traces(&candidate_client, &fixture_client, &database, &table).await?; + #[cfg(feature = "transaction-test-hooks")] + run_idle_ttl_trace(&idle_ttl_client, &fixture_client, &database, &table).await?; + run_permission_refusal_trace(&candidate_client, &fixture_client, &policy_storage, &database, &table) + .await?; + if let Some(proxy) = wire_proxy.as_ref() { + run_fault_proxy_traces(&candidate_client, &fixture_client, &database, &table, proxy).await?; + } + Ok(()) + }) + .await + { + Ok(result) => result, + Err(_) => Err("native transaction scenarios exceeded 90 seconds".to_string()), + }, + Err(error) => Err(error), + }; + + let cleanup = tokio::time::timeout(Duration::from_secs(60), async { + let close_candidate = async { + let _ = candidate_client.close_with_timeout(Duration::from_secs(12)).await; + let _ = candidate_service.close_with_timeout(Duration::from_secs(12)).await; + drop(candidate_client); + drop(candidate_service); + }; + + #[cfg(feature = "transaction-test-hooks")] + let close_idle_ttl = async { + let _ = idle_ttl_client.close_with_timeout(Duration::from_secs(12)).await; + let _ = idle_ttl_service.close_with_timeout(Duration::from_secs(12)).await; + drop(idle_ttl_client); + drop(idle_ttl_service); + }; + + #[cfg(feature = "transaction-test-hooks")] + tokio::join!(close_candidate, close_idle_ttl); + #[cfg(not(feature = "transaction-test-hooks"))] + close_candidate.await; + + let drop_result = tokio::time::timeout( + Duration::from_secs(15), + fixture_client.peer().call_tool(CallToolRequestParams::new("dbx_execute_query").with_arguments( + arguments(json!({ + "connection_id": "mcp-session-transaction-fixture", + "database": database, + "sql": format!("DROP TABLE IF EXISTS {table}") + })), + )), + ) + .await + .map_err(|_| "fixture DROP timed out".to_string())? + .map_err(|error| format!("fixture DROP transport failed: {error}"))?; + expect_success(&drop_result, "drop exact fixture table")?; + let _ = fixture_client.close_with_timeout(Duration::from_secs(5)).await; + let _ = fixture_service.close_with_timeout(Duration::from_secs(5)).await; + Result::<(), String>::Ok(()) + }) + .await + .map_err(|_| format!("bounded cleanup timed out; attributable fixture token: {table}"))?; + + if let Err(error) = cleanup { + return Err(format!("{error}; attributable fixture token: {table}")); + } + if let Some(proxy) = wire_proxy.take() { + proxy.shutdown().await; + } + scenario + } + + run().await.expect("native fixed-session contention gate"); +} diff --git a/crates/dbx-mcp/tests/support/mod.rs b/crates/dbx-mcp/tests/support/mod.rs new file mode 100644 index 0000000000..77b44fe360 --- /dev/null +++ b/crates/dbx-mcp/tests/support/mod.rs @@ -0,0 +1 @@ +pub mod mysql_wire_proxy; diff --git a/crates/dbx-mcp/tests/support/mysql_wire_proxy.rs b/crates/dbx-mcp/tests/support/mysql_wire_proxy.rs new file mode 100644 index 0000000000..f5b7d45512 --- /dev/null +++ b/crates/dbx-mcp/tests/support/mysql_wire_proxy.rs @@ -0,0 +1,386 @@ +use std::{ + io, + net::SocketAddr, + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, + }, +}; + +use tokio::{ + io::{AsyncRead, AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::Mutex, +}; +use tokio_util::sync::CancellationToken; + +enum FaultDirective { + None, + DropCommitAck, + DropQueryResponse { marker: String }, +} + +enum PendingResponse { + Contention, + DropCommitAck, +} + +struct ProxyState { + fault: Mutex, + contention_marker: Mutex>, + contention_forwarded: AtomicBool, + contention_response_pending: AtomicBool, + commit_tracking: AtomicBool, + commit_forward_count: AtomicUsize, + commit_ok_observed: AtomicBool, + commit_ack_dropped: AtomicBool, + dropped_query_forwarded: AtomicBool, +} + +impl Default for ProxyState { + fn default() -> Self { + Self { + fault: Mutex::new(FaultDirective::None), + contention_marker: Mutex::new(None), + contention_forwarded: AtomicBool::new(false), + contention_response_pending: AtomicBool::new(false), + commit_tracking: AtomicBool::new(false), + commit_forward_count: AtomicUsize::new(0), + commit_ok_observed: AtomicBool::new(false), + commit_ack_dropped: AtomicBool::new(false), + dropped_query_forwarded: AtomicBool::new(false), + } + } +} + +pub struct MysqlWireProxy { + listen_addr: SocketAddr, + state: Arc, + cancellation: CancellationToken, + task: Option>, +} + +impl MysqlWireProxy { + pub async fn start(target_host: String, target_port: u16) -> io::Result { + let listener = TcpListener::bind(("127.0.0.1", 0)).await?; + let listen_addr = listener.local_addr()?; + let state = Arc::new(ProxyState::default()); + let cancellation = CancellationToken::new(); + let task_state = state.clone(); + let task_cancellation = cancellation.clone(); + let task = tokio::spawn(async move { + loop { + let accepted = tokio::select! { + _ = task_cancellation.cancelled() => break, + accepted = listener.accept() => accepted, + }; + let Ok((client, _)) = accepted else { break }; + let target_host = target_host.clone(); + let state = task_state.clone(); + let cancellation = task_cancellation.child_token(); + tokio::spawn(async move { + let Ok(server) = TcpStream::connect((target_host.as_str(), target_port)).await else { + return; + }; + let _ = proxy_connection(client, server, state, cancellation).await; + }); + } + }); + Ok(Self { listen_addr, state, cancellation, task: Some(task) }) + } + + pub fn listen_addr(&self) -> SocketAddr { + self.listen_addr + } + + pub async fn set_contention_marker(&self, marker: impl Into) { + *self.state.contention_marker.lock().await = Some(marker.into()); + self.state.contention_forwarded.store(false, Ordering::SeqCst); + self.state.contention_response_pending.store(false, Ordering::SeqCst); + } + + pub fn contention_forwarded_and_pending(&self) -> bool { + self.state.contention_forwarded.load(Ordering::SeqCst) + && self.state.contention_response_pending.load(Ordering::SeqCst) + } + + pub async fn arm_drop_commit_ack(&self) { + *self.state.fault.lock().await = FaultDirective::DropCommitAck; + self.state.commit_tracking.store(true, Ordering::SeqCst); + self.state.commit_forward_count.store(0, Ordering::SeqCst); + self.state.commit_ok_observed.store(false, Ordering::SeqCst); + self.state.commit_ack_dropped.store(false, Ordering::SeqCst); + } + + pub fn commit_forward_count(&self) -> usize { + self.state.commit_forward_count.load(Ordering::SeqCst) + } + + pub fn commit_ok_observed(&self) -> bool { + self.state.commit_ok_observed.load(Ordering::SeqCst) + } + + pub fn commit_ack_dropped(&self) -> bool { + self.state.commit_ack_dropped.load(Ordering::SeqCst) + } + + pub async fn arm_drop_query_response(&self, marker: impl Into) { + *self.state.fault.lock().await = FaultDirective::DropQueryResponse { marker: marker.into() }; + self.state.dropped_query_forwarded.store(false, Ordering::SeqCst); + } + + pub fn dropped_query_forwarded(&self) -> bool { + self.state.dropped_query_forwarded.load(Ordering::SeqCst) + } + + pub async fn shutdown(mut self) { + self.cancellation.cancel(); + if let Some(task) = self.task.take() { + let _ = task.await; + } + } +} + +impl Drop for MysqlWireProxy { + fn drop(&mut self) { + self.cancellation.cancel(); + if let Some(task) = self.task.take() { + task.abort(); + } + } +} + +async fn proxy_connection( + client: TcpStream, + server: TcpStream, + state: Arc, + cancellation: CancellationToken, +) -> io::Result<()> { + let (mut client_reader, mut client_writer) = client.into_split(); + let (mut server_reader, mut server_writer) = server.into_split(); + let pending = Arc::new(Mutex::new(None::)); + let client_pending = pending.clone(); + let client_state = state.clone(); + let client_cancellation = cancellation.clone(); + let upstream = async move { + loop { + let packet = tokio::select! { + _ = client_cancellation.cancelled() => break, + packet = read_mysql_packet(&mut client_reader) => packet?, + }; + let Some(packet) = packet else { break }; + let sql = mysql_query_sql(&packet); + let contention = if let Some(sql) = sql { + client_state.contention_marker.lock().await.as_deref().is_some_and(|marker| sql.contains(marker)) + } else { + false + }; + if contention { + client_state.contention_forwarded.store(true, Ordering::SeqCst); + client_state.contention_response_pending.store(true, Ordering::SeqCst); + *client_pending.lock().await = Some(PendingResponse::Contention); + } + + let mut drop_after_forward = false; + if let Some(sql) = sql { + if sql.trim().eq_ignore_ascii_case("COMMIT") && client_state.commit_tracking.load(Ordering::SeqCst) { + client_state.commit_forward_count.fetch_add(1, Ordering::SeqCst); + } + let mut fault = client_state.fault.lock().await; + match &*fault { + FaultDirective::DropCommitAck if sql.trim().eq_ignore_ascii_case("COMMIT") => { + *client_pending.lock().await = Some(PendingResponse::DropCommitAck); + *fault = FaultDirective::None; + } + FaultDirective::DropQueryResponse { marker } if sql.contains(marker) => { + client_state.dropped_query_forwarded.store(true, Ordering::SeqCst); + drop_after_forward = true; + *fault = FaultDirective::None; + } + _ => {} + } + } + server_writer.write_all(&packet).await?; + if drop_after_forward { + client_cancellation.cancel(); + break; + } + } + io::Result::Ok(()) + }; + + let downstream_state = state; + let downstream_cancellation = cancellation.clone(); + let downstream = async move { + loop { + let packet = tokio::select! { + _ = downstream_cancellation.cancelled() => break, + packet = read_mysql_packet(&mut server_reader) => packet?, + }; + let Some(packet) = packet else { break }; + match pending.lock().await.take() { + Some(PendingResponse::Contention) => { + downstream_state.contention_response_pending.store(false, Ordering::SeqCst); + } + Some(PendingResponse::DropCommitAck) if packet.get(4) == Some(&0x00) => { + downstream_state.commit_ok_observed.store(true, Ordering::SeqCst); + downstream_state.commit_ack_dropped.store(true, Ordering::SeqCst); + downstream_cancellation.cancel(); + break; + } + Some(PendingResponse::DropCommitAck) | None => {} + } + client_writer.write_all(&packet).await?; + } + io::Result::Ok(()) + }; + + let (upstream, downstream) = tokio::join!(upstream, downstream); + upstream?; + downstream?; + Ok(()) +} + +pub async fn read_mysql_packet(reader: &mut (impl AsyncRead + Unpin)) -> io::Result>> { + let mut header = [0_u8; 4]; + match reader.read(&mut header[..1]).await? { + 0 => return Ok(None), + 1 => { + reader.read_exact(&mut header[1..]).await?; + } + _ => unreachable!("one-byte read returned more than one byte"), + } + let payload_len = usize::from(header[0]) | (usize::from(header[1]) << 8) | (usize::from(header[2]) << 16); + let mut packet = Vec::with_capacity(4 + payload_len); + packet.extend_from_slice(&header); + packet.resize(4 + payload_len, 0); + reader.read_exact(&mut packet[4..]).await?; + Ok(Some(packet)) +} + +pub fn mysql_query_sql(packet: &[u8]) -> Option<&str> { + let header: [u8; 4] = packet.get(..4)?.try_into().ok()?; + let payload_len = usize::from(header[0]) | (usize::from(header[1]) << 8) | (usize::from(header[2]) << 16); + if packet.len() != 4 + payload_len || packet.get(4) != Some(&0x03) { + return None; + } + std::str::from_utf8(packet.get(5..)?).ok() +} + +#[cfg(test)] +mod tests { + use super::{mysql_query_sql, read_mysql_packet, MysqlWireProxy}; + + fn packet(sequence: u8, payload: &[u8]) -> Vec { + let length = payload.len(); + let mut encoded = + vec![(length & 0xff) as u8, ((length >> 8) & 0xff) as u8, ((length >> 16) & 0xff) as u8, sequence]; + encoded.extend_from_slice(payload); + encoded + } + + #[tokio::test] + async fn packet_reader_handles_fragmented_headers_and_coalesced_packets() { + let first = packet(0, &[0x03, b'S', b'E', b'L', b'E', b'C', b'T', b' ', b'1']); + let second = packet(1, &[0x00, 0x01]); + let bytes = [first.clone(), second.clone()].concat(); + let (mut writer, mut reader) = tokio::io::duplex(64); + let sender = tokio::spawn(async move { + use tokio::io::AsyncWriteExt; + for byte in bytes { + writer.write_all(&[byte]).await.unwrap(); + } + }); + + let first_read = read_mysql_packet(&mut reader).await.unwrap().unwrap(); + let second_read = read_mysql_packet(&mut reader).await.unwrap().unwrap(); + assert_eq!(first_read, first); + assert_eq!(second_read, second); + assert!(read_mysql_packet(&mut reader).await.unwrap().is_none()); + sender.await.unwrap(); + } + + #[test] + fn query_decoder_accepts_only_complete_com_query_payloads() { + assert_eq!(mysql_query_sql(&packet(0, b"\x03COMMIT")), Some("COMMIT")); + assert_eq!(mysql_query_sql(&packet(0, b"\x03SELECT 1")), Some("SELECT 1")); + assert_eq!(mysql_query_sql(&packet(0, b"\x01COMMIT")), None); + assert_eq!(mysql_query_sql(&[7, 0, 0, 0, 0x03, b'C']), None); + } + + #[tokio::test] + async fn proxy_forwards_commit_observes_ok_and_drops_only_the_downstream_ack() { + use tokio::io::AsyncWriteExt; + + let upstream_listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); + let upstream_address = upstream_listener.local_addr().unwrap(); + let proxy = MysqlWireProxy::start(upstream_address.ip().to_string(), upstream_address.port()).await.unwrap(); + proxy.arm_drop_commit_ack().await; + let upstream = tokio::spawn(async move { + let (mut socket, _) = upstream_listener.accept().await.unwrap(); + assert_eq!(read_mysql_packet(&mut socket).await.unwrap().unwrap(), packet(0, b"\x03COMMIT")); + socket.write_all(&packet(1, &[0x00, 0x00, 0x00])).await.unwrap(); + }); + + let mut client = tokio::net::TcpStream::connect(proxy.listen_addr()).await.unwrap(); + client.write_all(&packet(0, b"\x03COMMIT")).await.unwrap(); + assert!(read_mysql_packet(&mut client).await.unwrap().is_none()); + upstream.await.unwrap(); + assert_eq!(proxy.commit_forward_count(), 1); + assert!(proxy.commit_ok_observed()); + assert!(proxy.commit_ack_dropped()); + proxy.shutdown().await; + } + + #[tokio::test] + async fn commit_counter_tracks_every_forward_after_the_fault_is_consumed() { + use tokio::io::AsyncWriteExt; + + let upstream_listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); + let upstream_address = upstream_listener.local_addr().unwrap(); + let proxy = MysqlWireProxy::start(upstream_address.ip().to_string(), upstream_address.port()).await.unwrap(); + proxy.arm_drop_commit_ack().await; + let upstream = tokio::spawn(async move { + for _ in 0..2 { + let (mut socket, _) = upstream_listener.accept().await.unwrap(); + assert_eq!(read_mysql_packet(&mut socket).await.unwrap().unwrap(), packet(0, b"\x03COMMIT")); + socket.write_all(&packet(1, &[0x00, 0x00, 0x00])).await.unwrap(); + } + }); + + let mut first = tokio::net::TcpStream::connect(proxy.listen_addr()).await.unwrap(); + first.write_all(&packet(0, b"\x03COMMIT")).await.unwrap(); + assert!(read_mysql_packet(&mut first).await.unwrap().is_none()); + + let mut second = tokio::net::TcpStream::connect(proxy.listen_addr()).await.unwrap(); + second.write_all(&packet(0, b"\x03COMMIT")).await.unwrap(); + assert!(read_mysql_packet(&mut second).await.unwrap().is_some()); + + upstream.await.unwrap(); + assert_eq!(proxy.commit_forward_count(), 2); + proxy.shutdown().await; + } + + #[tokio::test] + async fn dropping_proxy_releases_its_listener() { + let upstream_listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); + let upstream_address = upstream_listener.local_addr().unwrap(); + let proxy = MysqlWireProxy::start(upstream_address.ip().to_string(), upstream_address.port()).await.unwrap(); + let proxy_address = proxy.listen_addr(); + + drop(proxy); + + let rebound = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + match tokio::net::TcpListener::bind(proxy_address).await { + Ok(listener) => break listener, + Err(error) if error.kind() == std::io::ErrorKind::AddrInUse => tokio::task::yield_now().await, + Err(error) => panic!("unexpected listener rebind error: {error}"), + } + } + }) + .await; + + assert!(rebound.is_ok(), "dropping the proxy must release its listener"); + } +} diff --git a/docs/content/docs/mcp.cn.mdx b/docs/content/docs/mcp.cn.mdx index c431dd4f5c..17fb8745e0 100644 --- a/docs/content/docs/mcp.cn.mdx +++ b/docs/content/docs/mcp.cn.mdx @@ -104,6 +104,9 @@ DBX MCP 当前提供以下工具（具体可用工具取决于构建功能及权 | `dbx_execute_query` | 执行 SQL 或支持的 MongoDB shell 命令，默认返回 100 行，可用 `max_rows` 参数提高到最多 1000 行 | | `dbx_execute_batch` | 一次执行包含多条语句的 SQL 脚本，按语句返回结果（多语句脚本搭配 `use_transaction` 时返回单个合并结果） | | `dbx_open_session` | 为 SQL 连接打开固定后端连接的有状态查询会话 | +| `dbx_begin_transaction` | 在已启用事务的原生 MySQL 会话上开始事务 | +| `dbx_commit_transaction` | 在 MySQL 确认 `COMMIT` 后提交当前会话事务 | +| `dbx_rollback_transaction` | 回滚当前会话事务 | | `dbx_close_session` | 关闭会话并释放固定连接资源 | | `dbx_execute_redis_command` | 执行 Redis 命令 | | `dbx_peek_messages` | 读取 Kafka 消息，不提交消费位点 | @@ -163,6 +166,16 @@ DBX MCP 当前提供以下工具（具体可用工具取决于构建功能及权 有状态会话不会放宽 SQL 策略。`USE` 只有在会话中才有意义，但写入、DDL、生产保护、连接只读和数据库权限仍按每次请求重新检查。 +### 原生 MySQL 显式事务 + +调用 `dbx_open_session` 时传入 `enable_transactions: true`，DBX 会立即保留一条物理原生 MySQL 连接。随后依次调用 `dbx_begin_transaction`、带该 `session_id` 的一个或多个 `dbx_execute_query`/`dbx_execute_batch`，最后调用 `dbx_commit_transaction` 或 `dbx_rollback_transaction`。事务前后的查询仍使用同一条受管连接。`use_transaction` 仍不能与 `session_id` 同时使用；它是前文所述的独立原子批处理功能。 + +该 opt-in 模式仅支持本地原生 MySQL 连接。DBX Web/`DBX_WEB_URL`、外部驱动 profile、兼容引擎及其他数据库类型会返回 `TRANSACTION_UNSUPPORTED`，不会静默退化为自动提交。事务会话每次操作只接受一条经过解析的 `SELECT`（含锁定读）、`INSERT`、`UPDATE`、`DELETE` 或 `REPLACE`；原始事务控制、DDL/临时 DDL、`CALL`、`LOAD`、表锁、`SET`、`USE`、XA、文件/输出子句、可执行注释、优化器 hint 及多语句均会被拒绝。完整批次持有同一个操作槽位，并为每条语句返回执行后的事务状态。 + +响应包含 `transaction_state`（`idle`、`active` 或 `unknown`），并在适用时包含 `transaction_outcome`（`committed`、`rolled_back` 或 `unknown`）。普通 MySQL 语句错误可能仍让事务保持 active；DBX 会在同一连接上重新探测服务器状态后再返回。超时、取消、传输故障或丢失 `COMMIT` 确认会把结果永久标为 `unknown`，禁止该会话继续执行 SQL，并在不重放语句或 `COMMIT` 的前提下销毁连接。显式正值的连接查询超时同样用于事务操作；当连接配置为“无限制”（`0`）时，事务 owner 仍使用默认 300 秒操作上限，确保 busy 操作保持有界。关闭、过期、协议会话删除和服务退出会对状态已知的 active 事务执行有界回滚，然后销毁连接。 + +回滚保证依赖 InnoDB 等事务型存储。触发器、存储函数、非事务表及外部副作用可能产生 MySQL 无法回滚的效果。DBX 会在排队完成后重新检查当前连接/数据库范围、工具 allowlist、只读与生产保护以及已确认 SQL。即使写权限被撤销，仍允许原会话执行回滚和关闭以完成清理。 + ## 数据库访问 ### DBX Desktop 原生 Streamable HTTP diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index bad0eabf32..84c8706e80 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -104,6 +104,9 @@ DBX MCP provides the following tools (availability depends on build features and | `dbx_execute_query` | Execute SQL or a supported MongoDB shell command, returning 100 rows by default (up to 1000 with the `max_rows` parameter) | | `dbx_execute_batch` | Execute a SQL script containing multiple statements in one call, returning a result per statement (or a single merged result with `use_transaction` on a multi-statement script) | | `dbx_open_session` | Open a stateful SQL query session pinned to one backend connection | +| `dbx_begin_transaction` | Begin a transaction on a transaction-enabled native MySQL session | +| `dbx_commit_transaction` | Commit the active session transaction after MySQL acknowledges `COMMIT` | +| `dbx_rollback_transaction` | Roll back the active session transaction | | `dbx_close_session` | Close a session and release its pinned connection resources | | `dbx_execute_redis_command` | Execute a Redis command | | `dbx_peek_messages` | Read Kafka messages without committing consumer offsets | @@ -163,6 +166,16 @@ Sessions support SQL connections only and are bound to one connection and databa A stateful session does not widen SQL policy. `USE` is meaningful only inside a session, while writes, DDL, production protection, connection read-only state, and database privileges are rechecked for every request. +### Explicit native MySQL transactions + +Set `enable_transactions: true` when calling `dbx_open_session` to reserve one physical native MySQL connection eagerly. Then use `dbx_begin_transaction`, one or more `dbx_execute_query` or `dbx_execute_batch` calls with the returned `session_id`, and finally `dbx_commit_transaction` or `dbx_rollback_transaction`. Queries before and after a transaction still use the same owned connection. `use_transaction` remains incompatible with `session_id`; it is the separate atomic-batch feature described above. + +This opt-in mode is available only for local native MySQL connections. DBX Web/`DBX_WEB_URL`, external driver profiles, compatibility engines, and other database types return `TRANSACTION_UNSUPPORTED` instead of falling back to auto-commit. Transaction-session SQL accepts one parsed `SELECT` (including locking reads), `INSERT`, `UPDATE`, `DELETE`, or `REPLACE` statement per operation. It rejects raw transaction control, DDL and temporary DDL, `CALL`, `LOAD`, table locks, `SET`, `USE`, XA, file/output clauses, executable comments, optimizer hints, and multi-statements. Complete batches hold one operation slot and report the resulting transaction state for every statement. + +Responses include `transaction_state` (`idle`, `active`, or `unknown`) and, when known, `transaction_outcome` (`committed`, `rolled_back`, or `unknown`). A normal MySQL statement error can leave the transaction active; DBX probes fresh state on the same connection before reporting it. A timeout, cancellation, transport failure, or lost `COMMIT` acknowledgement makes the outcome `unknown`, permanently blocks more SQL on that session, and disposes the connection without replaying the statement or `COMMIT`. A positive connection query timeout applies to transaction operations; when the connection is configured as unlimited (`0`), transaction owners still apply their default 300-second operation cap so busy work remains bounded. Closing, expiry, protocol-session deletion, and server shutdown attempt bounded rollback for a known active transaction and then dispose the connection. + +Rollback guarantees require transactional storage such as InnoDB. Triggers, stored functions, nontransactional tables, and external side effects can have effects that MySQL cannot roll back. DBX rechecks the current connection/database scope, tool allowlist, read-only and production protections, and confirmed SQL after queue waiting. Rollback and close remain available for cleanup after write permission is revoked. + ## Database Access ### Native Streamable HTTP for DBX Desktop