diff --git a/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java b/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java index 3f1450701e..9e7f655aff 100644 --- a/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java +++ b/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java @@ -56,8 +56,9 @@ public final class TransactionExecutor { return unchecked(() -> { long start = System.currentTimeMillis(); if (!supportsTransactions(conn)) { - long totalAffected = executeAll(conn, statements, schema, setSchemaSql, resetSchemaSql, runner); - return result(totalAffected, start); + // An explicit transaction request must not degrade to auto-commit: + // callers rely on it to avoid partial writes when a later statement fails. + throw new UnsupportedOperationException("Transactions are not supported by this JDBC driver"); } boolean savedAutoCommit = conn.getAutoCommit(); diff --git a/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java b/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java index 8ca214d3f6..ae787bcd9e 100644 --- a/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java +++ b/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java @@ -626,17 +626,19 @@ class CommonJavaCompatibilityTest { } @Test - void executesTransactionsOneByOneWhenJdbcDriverDoesNotSupportTransactions() { + void rejectsTransactionsWhenJdbcDriverDoesNotSupportTransactions() { List calls = new ArrayList<>(); DatabaseAgent agent = new TransactionAgent(nonTransactionalConnection(calls)); - QueryResult result = agent.executeTransaction(Arrays.asList("UPDATE A SET ID = 1", "UPDATE B SET ID = 2"), "APP"); - - assertEquals(2L, result.getAffected_rows()); - assertEquals( - Arrays.asList("supportsTransactions", "execute:SET SCHEMA \"APP\"", "executeUpdate:UPDATE A SET ID = 1", "executeUpdate:UPDATE B SET ID = 2"), - calls + UnsupportedOperationException error = assertThrows( + UnsupportedOperationException.class, + () -> agent.executeTransaction(Arrays.asList("UPDATE A SET ID = 1", "UPDATE B SET ID = 2"), "APP") ); + + assertEquals("Transactions are not supported by this JDBC driver", error.getMessage()); + // Capability is checked before schema switching or statements, so a + // failed transaction request cannot leave a partially applied batch. + assertEquals(Collections.singletonList("supportsTransactions"), calls); } @Test diff --git a/agents/drivers/hive-go/query.go b/agents/drivers/hive-go/query.go index 734d4fa967..088a803976 100644 --- a/agents/drivers/hive-go/query.go +++ b/agents/drivers/hive-go/query.go @@ -333,9 +333,12 @@ func (server *server) executeStatements(params map[string]json.RawMessage, trans } return emptyQueryResult(affected, started), nil } - if !transactionUnsupported(beginErr) { - return queryResult{}, beginErr + if transactionUnsupported(beginErr) { + // Do not turn an explicit transaction request into auto-commit + // execution: a later error would leave prior statements applied. + return queryResult{}, errors.New("Hive does not support rollbackable transactions") } + return queryResult{}, beginErr } for _, statement := range statements { trimmed := trimStatementSQL(statement) diff --git a/agents/drivers/hive-go/query_test.go b/agents/drivers/hive-go/query_test.go index e49cdd5e46..2af52a96bf 100644 --- a/agents/drivers/hive-go/query_test.go +++ b/agents/drivers/hive-go/query_test.go @@ -536,31 +536,24 @@ func TestRuntimeSessionsRemainIsolated(t *testing.T) { } } -func TestTransactionFallbackDoesNotReplayFailedStatements(t *testing.T) { +func TestTransactionRejectsUnsupportedDriverBeforeStatementsRun(t *testing.T) { behavior := &scriptedBehavior{} behavior.exec = func(_ context.Context, query string) (driver.Result, error) { - if strings.Contains(query, "second") { - return nil, io.EOF - } - return driver.RowsAffected(1), nil + return nil, fmt.Errorf("must not execute %q", query) } server := newScriptedServer(t, behavior) _, err := server.executeStatements(rawParams(map[string]any{ "statements": []string{"INSERT first", "INSERT second", "INSERT third"}, }), true) - if !errors.Is(err, io.EOF) { - t.Fatalf("expected connection failure, got %v", err) + if err == nil || !strings.Contains(err.Error(), "does not support rollbackable transactions") { + t.Fatalf("expected unsupported transaction error, got %v", err) } _, executions, beginCalls, _ := behavior.snapshot() - if !reflect.DeepEqual(executions, []string{"INSERT first", "INSERT second"}) { - t.Fatalf("failed transaction was replayed or continued: %v", executions) + if len(executions) != 0 { + t.Fatalf("unsupported transaction executed statements: %v", executions) } if beginCalls == 0 { - t.Fatal("transaction capability was not attempted before fallback") - } - rpcErr := classifyRPCError("execute_transaction", "session-a", err) - if rpcErr.Data.Category != "connection" || rpcErr.Data.SessionDisposition != "quarantine" { - t.Fatalf("unexpected connection failure classification: %#v", rpcErr) + t.Fatal("transaction capability was not checked") } } diff --git a/agents/drivers/iotdb/main_test.go b/agents/drivers/iotdb/main_test.go index e7ebe91ba8..544e37c896 100644 --- a/agents/drivers/iotdb/main_test.go +++ b/agents/drivers/iotdb/main_test.go @@ -26,6 +26,16 @@ func TestHandshakeAdvertisesNativeCapabilities(t *testing.T) { } } +func TestExecuteTransactionIsRejectedBeforeConnecting(t *testing.T) { + server := &server{} + _, err := server.executeStatements(map[string]json.RawMessage{ + "statements": json.RawMessage(`["INSERT INTO metrics(device, time, value) VALUES ('d1', 1, 1)"]`), + }, true) + if err == nil || err.Error() != "IoTDB does not support transactions" { + t.Fatalf("expected unsupported transaction error, got %v", err) + } +} + func TestHandleLineClassifiesMissingSession(t *testing.T) { response, _ := newRuntimeServer().handleLine( `{"jsonrpc":"2.0","id":7,"method":"validate_session","params":{"agentSessionId":"missing"}}`, diff --git a/agents/drivers/iotdb/query.go b/agents/drivers/iotdb/query.go index 3770e849f5..65877d5d9a 100644 --- a/agents/drivers/iotdb/query.go +++ b/agents/drivers/iotdb/query.go @@ -306,6 +306,12 @@ func (s *server) closeAllQuerySessions() error { } func (s *server) executeStatements(params map[string]json.RawMessage, transaction bool) (queryResult, error) { + if transaction { + // IoTDB executes statements independently and exposes no rollbackable + // transaction boundary. Reject the RPC before obtaining a client so a + // requested atomic batch can never leave partial writes behind. + return queryResult{}, errors.New("IoTDB does not support transactions") + } started := time.Now() statements := stringSliceParam(params, "statements") if len(statements) == 0 { diff --git a/agents/scripts/driver_release_packages_test.py b/agents/scripts/driver_release_packages_test.py index 31b2162526..da3da6e02b 100644 --- a/agents/scripts/driver_release_packages_test.py +++ b/agents/scripts/driver_release_packages_test.py @@ -54,6 +54,7 @@ class DriverReleasePackagesTest(unittest.TestCase): "tdengine": "0.1.0", "etcd": "0.1.40", "etcd2": "0.1.0", + "argo": "0.1.0", "sqlite-worker": "0.1.0", } diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 1368a8265c..ba0517b08c 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -2248,7 +2248,15 @@ impl AppState { | DatabaseType::Kwdb | DatabaseType::Questdb | DatabaseType::OpenGauss => { - let pg_pool = db::postgres::connect(&url, connect_timeout).await?; + // A session pool must have one physical client: sequential + // statements can otherwise lose temp tables/SET state when a + // pool checkout selects another PostgreSQL connection. + let pg_pool = db::postgres::connect_with_max_connections( + &url, + connect_timeout, + postgres_pool_max_connections_for_session(client_session_id), + ) + .await?; // Build TLS cancel context for reconstructing TLS connection during cancel if let Some(ctx) = db::postgres::build_postgres_cancel_context(&url) { self.postgres_cancel_contexts.write().await.insert(pool_key.clone(), ctx); @@ -5397,6 +5405,14 @@ fn mysql_pool_max_connections_for_session(client_session_id: Option<&str>) -> us } } +fn postgres_pool_max_connections_for_session(client_session_id: Option<&str>) -> usize { + if normalize_client_session_id(client_session_id).is_some() { + 1 + } else { + 10 + } +} + fn redis_cluster_transport_prefix(connection_id: &str) -> String { format!("{connection_id}:redis-cluster:") } @@ -7729,6 +7745,22 @@ mod tests { assert_eq!(super::mysql_pool_max_connections_for_session(Some("tab-1")), 1); } + #[test] + fn postgres_pool_size_keeps_session_pools_single_connection() { + assert_eq!(super::postgres_pool_max_connections_for_session(None), 10); + assert_eq!(super::postgres_pool_max_connections_for_session(Some("")), 10); + assert_eq!(super::postgres_pool_max_connections_for_session(Some("mcp:batch-1")), 1); + } + + #[test] + fn stale_mysql_pool_observation_does_not_remove_replacement_generation() { + let checked = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10); + let checked_clone = checked.clone(); + let replacement = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10); + assert!(checked.is_same_pool(&checked_clone)); + assert!(!checked.is_same_pool(&replacement)); + } + #[tokio::test] async fn stale_postgres_cleanup_preserves_concurrent_replacement_publication() { let (state, dir) = test_app_state().await; diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index 6fed9b44e8..462833fb63 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -2558,9 +2558,15 @@ impl AgentDriverClient { database: Option<&str>, statements: &[String], schema: Option<&str>, + timeout_duration: Option, ) -> Result { self.invalidate_cached_query(); - self.call_method(AgentMethod::ExecuteTransaction, agent_transaction_params(database, statements, schema)).await + self.call_method_with_timeout( + AgentMethod::ExecuteTransaction, + agent_transaction_params(database, statements, schema), + timeout_duration, + ) + .await } pub async fn execute_transaction_typed( @@ -2568,10 +2574,15 @@ impl AgentDriverClient { database: Option<&str>, statements: &[String], schema: Option<&str>, + timeout_duration: Option, ) -> Result { self.invalidate_cached_query(); - self.call_method_typed(AgentMethod::ExecuteTransaction, agent_transaction_params(database, statements, schema)) - .await + self.call_method_typed_with_timeout( + AgentMethod::ExecuteTransaction, + agent_transaction_params(database, statements, schema), + timeout_duration, + ) + .await } pub async fn begin_manual_transaction( @@ -4161,7 +4172,7 @@ mod tests { let script_path = std::env::temp_dir().join(format!("dbx-agent-{prefix}-{}.py", uuid::Uuid::new_v4())); std::fs::write( &script_path, - r#"import json, sys, threading + r#"import json, sys, threading, time print(json.dumps({'ready': True}), flush=True) state_lock = threading.Lock() output_lock = threading.Lock() @@ -4232,6 +4243,12 @@ def respond(req): return write_response(req, {'sql': sql, 'count': current_count}) return + if method == 'execute_transaction': + statements = params.get('statements', []) + if statements == ['slow']: + time.sleep(1.2) + write_response(req, {'ok': True}) + return write_response(req, {'ok': True}) for line in sys.stdin: @@ -4486,7 +4503,10 @@ for line in sys.stdin: ) .await .unwrap(); - client.execute_transaction::(None, &["UPDATE t SET a = 2".to_string()], None).await.unwrap(); + client + .execute_transaction::(None, &["UPDATE t SET a = 2".to_string()], None, None) + .await + .unwrap(); client .execute_query_cached_with_timeout::( "transaction-invalidation".to_string(), @@ -4515,6 +4535,42 @@ for line in sys.stdin: let _ = std::fs::remove_file(script_path); } + #[tokio::test] + async fn execute_transaction_typed_applies_timeout_when_agent_stalls() { + let (runtime, script_path) = spawn_stateful_test_runtime("transaction-timeout-test").await; + let mut client = AgentDriverClient::shared_session(runtime.clone(), "transaction-timeout-session".to_string()); + + // A single slow statement with a short timeout must surface a typed timeout + // error instead of waiting for the agent indefinitely. + let error = client + .execute_transaction_typed::( + None, + &["slow".to_string()], + None, + Some(Duration::from_millis(75)), + ) + .await + .unwrap_err(); + assert!(matches!( + error, + AgentCallError::Timeout { + stage: AgentErrorStage::Execute, + operation_outcome: AgentOperationOutcome::Unknown, + } + )); + + // A fast transaction with a comfortable timeout still succeeds and + // invalidates the cached query, proving the timeout path did not break it. + let result: serde_json::Value = client + .execute_transaction_typed(None, &["UPDATE t SET a = 2".to_string()], None, Some(Duration::from_secs(2))) + .await + .unwrap(); + assert_eq!(result, serde_json::json!({"ok": true})); + + runtime.kill(); + let _ = std::fs::remove_file(script_path); + } + #[tokio::test] async fn shared_runtime_tolerates_post_startup_driver_stdout_noise() { let script_path = diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 29431a0e1e..d72e5132bf 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -2043,20 +2043,45 @@ async fn stream_query_rows_text_on_client( } pub async fn connect(url: &str, fallback_timeout: Duration) -> Result { + connect_with_max_connections(url, fallback_timeout, 10).await +} + +/// Creates a PostgreSQL pool with an explicit checkout bound. +/// +/// Session-scoped DBX pools use a single connection so temporary tables and +/// other connection-local state cannot migrate between physical clients. +pub async fn connect_with_max_connections( + url: &str, + fallback_timeout: Duration, + max_connections: usize, +) -> Result { #[cfg(all(windows, target_vendor = "win7"))] { - connect_with_optional_local_timezone(url, fallback_timeout, None).await + connect_with_optional_local_timezone_with_max_connections(url, fallback_timeout, None, max_connections).await } #[cfg(not(all(windows, target_vendor = "win7")))] { let timezone = iana_time_zone::get_timezone().unwrap_or_else(|_| "UTC".to_string()); - connect_with_local_timezone(url, fallback_timeout, &timezone).await + connect_with_local_timezone_with_max_connections(url, fallback_timeout, &timezone, max_connections).await } } +/// Test-only thin wrappers (the library now connects through +/// `connect_with_max_connections`, which carries the pool bound). +#[cfg(test)] async fn connect_with_local_timezone(url: &str, fallback_timeout: Duration, timezone: &str) -> Result { - connect_with_optional_local_timezone(url, fallback_timeout, Some(timezone)).await + connect_with_local_timezone_with_max_connections(url, fallback_timeout, timezone, 10).await +} + +async fn connect_with_local_timezone_with_max_connections( + url: &str, + fallback_timeout: Duration, + timezone: &str, + max_connections: usize, +) -> Result { + connect_with_optional_local_timezone_with_max_connections(url, fallback_timeout, Some(timezone), max_connections) + .await } /// Identity of one physical backend connection for notice attribution: @@ -2247,23 +2272,33 @@ async fn drain_postgres_notices(client: &deadpool_postgres::Client) -> Vec, +) -> Result { + connect_with_optional_local_timezone_with_max_connections(url, fallback_timeout, timezone, 10).await +} + +async fn connect_with_optional_local_timezone_with_max_connections( + url: &str, + fallback_timeout: Duration, + timezone: Option<&str>, + max_connections: usize, ) -> Result { let url_with_keepalive = inject_postgres_keepalive_params(url); let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); - let first_attempt = connect_postgres_pool_attempt(&url_with_keepalive, timeout).await; + let first_attempt = connect_postgres_pool_attempt(&url_with_keepalive, timeout, max_connections).await; let (pool, client) = match first_attempt { Err(error) if postgres_error_should_retry_without_tls(&error) => { let Some(fallback_url) = postgres_ssl_fallback_url(&url_with_keepalive) else { return Err(error); }; log::info!("PostgreSQL TLS handshake failed in sslmode=prefer; retrying without TLS"); - connect_postgres_pool_attempt(&fallback_url, timeout).await? + connect_postgres_pool_attempt(&fallback_url, timeout, max_connections).await? } result => result?, }; @@ -2285,6 +2320,7 @@ async fn connect_with_optional_local_timezone( async fn connect_postgres_pool_attempt( url: &str, timeout: Duration, + max_connections: usize, ) -> Result<(Pool, deadpool_postgres::Client), String> { let postgres_url = postgres_connection_url(url)?; super::with_connection_timeout("PostgreSQL", timeout, async { @@ -2320,7 +2356,7 @@ async fn connect_postgres_pool_attempt( ) }; let pool = Pool::builder(mgr) - .max_size(10) + .max_size(max_connections.max(1)) .runtime(Runtime::Tokio1) .wait_timeout(Some(timeout)) .create_timeout(Some(timeout)) diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 4b679a6f1b..2eb0c07172 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -11,7 +11,7 @@ use std::future::Future; use std::ops::ControlFlow; use std::sync::{ atomic::{AtomicU64, Ordering}, - Arc, + Arc, Condvar, Mutex, }; use std::time::Duration; use tokio::time::timeout; @@ -1718,6 +1718,15 @@ async fn configured_operation_budget_for_pool_key(state: &AppState, pool_key: &s .unwrap_or_else(DbOperationBudget::with_defaults) } +/// Override a transaction budget's query timeout from a per-call override (e.g. the MCP +/// global query-timeout policy). `None` leaves the budget unchanged; `Some(secs)` follows +/// `resolve_query_timeout` semantics (`Some(0)` clears the limit, meaning unlimited). +fn apply_query_timeout_override(budget: &mut DbOperationBudget, timeout_secs: Option) { + if let Some(secs) = timeout_secs { + budget.query_timeout = resolve_query_timeout(Some(secs)); + } +} + fn oceanbase_mysql_session_timeout_sql(config: Option<&ConnectionConfig>, timeout_secs: Option) -> Option { let config = config?; let timeout_secs = timeout_secs.unwrap_or(config.query_timeout_secs); @@ -2937,6 +2946,42 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( } }; + let execution_plan = query_execution_plan(sql, db_type, is_sqlserver_agent); + let continue_on_error = options.continue_on_error && !execution_plan.stop_on_error; + let statements = execution_plan.statements; + if statements.is_empty() { + return Ok(vec![empty_query_result(0).into()]); + } + + // Check the transaction request before database-specific fast paths. Otherwise + // a backend such as SQL Server or HTTP SQLite can return successful + // auto-commit results while the API has promised a rollbackable batch. + // The DDL cap applies only to this opt-in use_transaction contract: a script + // that the user asked to run atomically must not silently produce partial + // effects. Other callers that share the transaction kernel (schema-diff + // deploy, imports) document a mixed-outcome-on-failure behaviour instead, so + // they are deliberately not capped here. + if options.use_transaction == Some(true) && statements.len() > 1 { + if batch_transaction_ddl_is_unrollbackable(db_type, &statements) { + return Err( + "use_transaction cannot be used with a batch whose DDL cannot be rolled back: DDL statements implicitly commit and cannot be undone. Run the batch without use_transaction (auto-commit, one result per statement) or split the DDL and DML into separate calls." + .to_string() + .into(), + ); + } + let result = execute_statements_in_transaction_typed( + state, + connection_id, + database, + &statements, + schema, + options.catalog.as_deref(), + options.timeout_secs, + ) + .await?; + return Ok(vec![result.into()]); + } + if is_sqlserver { return execute_multi_sqlserver(state, &pool_key, sql, cancel_token, options).await.map_err(Into::into); } @@ -2949,7 +2994,7 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( }; // HTTP SQLite providers send all statements in one request so the provider - // can preserve batch ordering and atomicity. + // can preserve batch ordering and atomicity in the default batch mode. if is_http_sqlite { let table_data_preview = options.table_data_preview; return single_statement_multi_result( @@ -2967,28 +3012,6 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( ); } - let execution_plan = query_execution_plan(sql, db_type, is_sqlserver_agent); - let continue_on_error = options.continue_on_error && !execution_plan.stop_on_error; - let statements = execution_plan.statements; - if statements.is_empty() { - return Ok(vec![empty_query_result(0).into()]); - } - - // When use_transaction is explicitly true and we have multiple statements, - // route through the transaction wrapper instead of the sequential auto-commit loop. - if options.use_transaction == Some(true) && statements.len() > 1 { - let result = execute_statements_in_transaction_typed( - state, - connection_id, - database, - &statements, - schema, - options.catalog.as_deref(), - ) - .await?; - return Ok(vec![result.into()]); - } - let mysql_pool = { let pool_handle = state.pool_handle(&pool_key).await; match pool_handle.as_ref() { @@ -3106,7 +3129,12 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( Ok(results) } -fn query_execution_plan( +/// Split a batch script into the statements the core will execute, using the +/// same dialect-aware splitter for every caller. SQL Server agent pools split on +/// `GO` batch separators (`split_sql_batches`); everything else uses the +/// database-dialect statement splitter. Public so the MCP pre-check can align its +/// transaction-entry decision with the core's actual script split. +pub fn query_execution_plan( sql: &str, db_type: Option, preserve_sqlserver_batches: bool, @@ -3121,6 +3149,38 @@ fn query_execution_plan( ) } +/// Whether a connection's pool is the SQL Server agent driver, mirroring the +/// predicate the core uses to select the batch splitter in [`query_execution_plan`]. +/// Reads the existing pool without creating one so callers (the MCP pre-check) +/// can align their splitter without opening a connection. `false` for non-SQL +/// Server connections and when no pool is open yet. A SQL Server connection +/// backed by the legacy agent driver connects as `PoolKind::Agent` and is +/// detected here; the native driver connects as `PoolKind::SqlServer` and +/// returns `false`. +pub async fn connection_pool_is_sqlserver_agent(state: &AppState, connection_id: &str, database: &str) -> bool { + let db_type = connection_database_type(state, connection_id).await; + if db_type != Some(DatabaseType::SqlServer) { + return false; + } + + // SQL Server base pool keys are either `connection_id` (no database) or + // `connection_id:database`. Peek both without creating a pool. + let candidates = match database.trim() { + "" => vec![connection_id.to_string()], + db => vec![format!("{connection_id}:{db}"), connection_id.to_string()], + }; + for key in candidates { + if let Some(pool) = state.pool_handle(&key).await { + return match pool { + PoolKind::Agent(_) => true, + PoolKind::SqlServer(_) => false, + _ => false, + }; + } + } + false +} + fn single_statement_multi_result( result: Result, table_data_preview: bool, @@ -3945,7 +4005,7 @@ fn pool_kind_has_transactional_path(pool: &PoolKind) -> bool { /// - Returns a structured result (never re-executes statements to probe status). /// - Comment-only / empty scripts succeed as `committed` with zero statements. /// - When the target path cannot guarantee DDL atomicity (MySQL/Oracle DDL -/// auto-commit, `TxPath::None`, etc.), a failure reports `mixed` and +/// auto-commit, unsupported batch transaction paths, etc.), a failure reports `mixed` and /// `executed_count` reflects the statements that were issued before the /// error, so the caller can warn the user that partial effects may persist. pub async fn execute_schema_diff_deploy( @@ -4087,11 +4147,10 @@ pub async fn execute_schema_diff_deploy( /// Execute multiple SQL statements within a single transaction. /// For pooled drivers (Postgres/MySQL), uses the driver transaction API. -/// For SQLite and already-single-connection drivers (ClickHouse/SqlServer/Agent), -/// uses explicit BEGIN/COMMIT/ROLLBACK on the shared connection. -/// For databases that don't support explicit transactions (Redis, MongoDB, Oracle), -/// executes statements sequentially without transaction. -/// If BEGIN fails, returns an error instead of silently falling back to auto-commit. +/// For SQLite and SQL Server, uses a transaction on the driver's shared connection. +/// Agent drivers must provide the same rollbackable transaction contract. +/// Backends without a verified rollbackable path are rejected instead of being +/// silently executed one statement at a time in auto-commit mode. pub async fn execute_statements_in_transaction( state: &AppState, connection_id: &str, @@ -4100,7 +4159,7 @@ pub async fn execute_statements_in_transaction( schema: Option<&str>, catalog: Option<&str>, ) -> Result { - execute_statements_in_transaction_typed(state, connection_id, database, statements, schema, catalog) + execute_statements_in_transaction_typed(state, connection_id, database, statements, schema, catalog, None) .await .map_err(QueryExecutionError::into_legacy_string) } @@ -4113,6 +4172,7 @@ pub async fn execute_statements_in_transaction_typed( statements: &[String], schema: Option<&str>, catalog: Option<&str>, + timeout_secs: Option, ) -> Result { let sql_ctx = statements.first().map(|s| s.as_str()).unwrap_or(""); let pool_database = query_pool_database(database, catalog); @@ -4129,6 +4189,7 @@ pub async fn execute_statements_in_transaction_typed( statements, schema, catalog, + timeout_secs, ) .await } @@ -4152,6 +4213,7 @@ pub async fn execute_statements_in_transaction_on_pool( statements, schema, catalog, + None, ) .await .map_err(QueryExecutionError::into_legacy_string) @@ -4166,50 +4228,32 @@ pub async fn execute_statements_in_transaction_on_pool_typed( statements: &[String], schema: Option<&str>, catalog: Option<&str>, + timeout_secs: Option, ) -> Result { + let db_type = connection_database_type(state, connection_id).await; + // Read-only check: intercept all transaction paths before dispatching check_read_only_for_connection_multi(state, pool_key, statements).await?; let start = std::time::Instant::now(); - let db_type = connection_database_type(state, connection_id).await; let mysql_catalog_dialect = connection_mysql_catalog_dialect(state, connection_id).await; - let operation_budget = configured_operation_budget_for_pool_key(state, pool_key).await; + let mut operation_budget = configured_operation_budget_for_pool_key(state, pool_key).await; + apply_query_timeout_override(&mut operation_budget, timeout_secs); + // This is the single capability model for explicit batch transactions. + // Do not add a backend here unless its execution path keeps every statement + // on one transaction-capable connection and can roll back on failure. + // Agent drivers are delegated because their transaction RPC has the same + // contract and rejects drivers that cannot provide it. // Clone the pool handle within the lock, then drop it before any async work. - let path = { - state.pool_handle(pool_key).await.as_ref().map(|p| match p { - PoolKind::Postgres(pg) => TxPath::Pg(pg.clone()), - PoolKind::Mysql(mp, _mode) => TxPath::Mysql(mp.clone(), false), - PoolKind::Sqlite(sq) => TxPath::Sqlite(sq.clone()), - PoolKind::CloudflareD1(client) => TxPath::CloudflareD1(client.clone()), - PoolKind::ClickHouse(_) | PoolKind::Rqlite(_) | PoolKind::Turso(_) | PoolKind::SqlServer(_) => { - TxPath::Explicit - } - PoolKind::Agent(client) => TxPath::Agent(client.clone()), - PoolKind::MessageQueue | PoolKind::Nacos | PoolKind::Consul(_) | PoolKind::HBase(_) => TxPath::None, - #[cfg(feature = "mq-admin")] - PoolKind::Mqtt(_) => TxPath::None, - PoolKind::DuckDbWorker(_) - | PoolKind::Redis(_) - | PoolKind::MongoDb(_) - | PoolKind::DynamoDb(_) - | PoolKind::Elasticsearch(_) - | PoolKind::Easysearch(_) - | PoolKind::Meilisearch(_) - | PoolKind::VectorDb(_) - | PoolKind::InfluxDb(_) - | PoolKind::InfluxDb3(_) - | PoolKind::VictoriaMetrics(_) - | PoolKind::ExternalDriver { .. } => TxPath::None, - }) - }; + let path = { state.pool_handle(pool_key).await.as_ref().map(batch_transaction_path) }; let result = match path { - Some(TxPath::Pg(pool)) => { + Some(BatchTransactionPath::Pg(pool)) => { let cancel_context = state.get_postgres_cancel_context(pool_key).await; exec_tx_pg_inner(pool, statements, schema, start, operation_budget.clone(), cancel_context).await } - Some(TxPath::Mysql(pool, _bare)) => exec_tx_mysql_inner( + Some(BatchTransactionPath::Mysql(pool)) => exec_tx_mysql_inner( state, pool_key, pool, @@ -4222,37 +4266,37 @@ pub async fn execute_statements_in_transaction_on_pool_typed( ) .await .map_err(Into::into), - Some(TxPath::Sqlite(pool)) => exec_tx_sqlite_inner(pool, statements, start).await.map_err(Into::into), - Some(TxPath::CloudflareD1(client)) => { - let sql = statements.join(";\n"); - wait_for_query_opt( - None, - operation_budget.query_timeout, - db::cloudflare_d1_driver::execute_query_with_max_rows(&client, &sql, None), - ) - .await - .map_err(Into::into) + Some(BatchTransactionPath::Sqlite(pool)) => { + exec_tx_sqlite_inner(pool, statements, start, &operation_budget).await.map_err(Into::into) } - Some(TxPath::Agent(client)) => { - let result = exec_tx_agent_inner(client.clone(), db_type, Some(database), statements, schema, start).await; + Some(BatchTransactionPath::Agent(client)) => { + let result = exec_tx_agent_inner( + client.clone(), + db_type, + Some(database), + statements, + schema, + start, + &operation_budget, + ) + .await; if let Err(error) = result.as_ref() { discard_agent_pool_after_typed_error(state, pool_key, &client, error, RecoveryScope::UserOperation) .await; } return result.map_err(QueryExecutionError::Agent); } - Some(TxPath::Explicit) => { + Some(BatchTransactionPath::Explicit) => { let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; exec_tx_explicit_inner(state, pool_key, mysql_dialect, Some(database), statements, schema, start) .await .map_err(Into::into) } - Some(TxPath::None) => { - let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; - exec_tx_none_inner(state, pool_key, mysql_dialect, Some(database), statements, schema, start) - .await - .map_err(Into::into) - } + Some(BatchTransactionPath::Unsupported) => Err( + "The active backend cannot provide a rollbackable transaction for a batch; run without use_transaction." + .to_string() + .into(), + ), None => Err("Connection not found for transaction".to_string().into()), }; @@ -4263,15 +4307,64 @@ pub async fn execute_statements_in_transaction_on_pool_typed( result } -/// Owned pool variants for safe dispatch across async boundaries. -enum TxPath { +/// Whether an opt-in explicit batch transaction (`use_transaction`) must be +/// rejected because the backend's DDL statements implicitly commit and cannot be +/// rolled back. Used by the opt-in batch-transaction entry point (the +/// `use_transaction` branch of [`execute_multi_core_with_options_for_client_and_progress_typed`]) +/// and by the MCP layer's identical pre-check. It covers every backend whose DDL +/// is not rollbackable (MySQL-family and Oracle) via the single capability +/// predicate (`database_supports_transactional_ddl`) without re-listing engines +/// here. Paths that document a mixed-outcome-on-failure behaviour (schema-diff +/// deploy, imports) do not call this and keep running-and-reporting. +pub fn batch_transaction_ddl_is_unrollbackable(db_type: Option, statements: &[String]) -> bool { + let Some(db_type) = db_type else { + return false; + }; + if database_supports_transactional_ddl(db_type) { + return false; + } + statements.iter().any(|statement| matches!(classify_sql_risk_for_database(statement, db_type), Ok(SqlRisk::Ddl))) +} + +/// Owned transaction-capable pool variants for safe dispatch across async boundaries. +enum BatchTransactionPath { Pg(deadpool_postgres::Pool), - Mysql(db::mysql::MySqlPool, bool), + Mysql(db::mysql::MySqlPool), Sqlite(db::sqlite::SqliteHandle), - CloudflareD1(db::cloudflare_d1_driver::CloudflareD1Client), Agent(Arc), Explicit, - None, + Unsupported, +} + +fn batch_transaction_path(pool: &PoolKind) -> BatchTransactionPath { + match pool { + PoolKind::Postgres(pg) => BatchTransactionPath::Pg(pg.clone()), + PoolKind::Mysql(pool, _mode) => BatchTransactionPath::Mysql(pool.clone()), + PoolKind::Sqlite(pool) => BatchTransactionPath::Sqlite(pool.clone()), + PoolKind::SqlServer(_) => BatchTransactionPath::Explicit, + PoolKind::Agent(client) => BatchTransactionPath::Agent(client.clone()), + PoolKind::MessageQueue | PoolKind::Nacos | PoolKind::Consul(_) | PoolKind::HBase(_) => { + BatchTransactionPath::Unsupported + } + #[cfg(feature = "mq-admin")] + PoolKind::Mqtt(_) => BatchTransactionPath::Unsupported, + PoolKind::DuckDbWorker(_) + | PoolKind::Redis(_) + | PoolKind::MongoDb(_) + | PoolKind::DynamoDb(_) + | PoolKind::ClickHouse(_) + | PoolKind::Rqlite(_) + | PoolKind::Turso(_) + | PoolKind::CloudflareD1(_) + | PoolKind::Elasticsearch(_) + | PoolKind::Easysearch(_) + | PoolKind::Meilisearch(_) + | PoolKind::VectorDb(_) + | PoolKind::InfluxDb(_) + | PoolKind::InfluxDb3(_) + | PoolKind::VictoriaMetrics(_) + | PoolKind::ExternalDriver { .. } => BatchTransactionPath::Unsupported, + } } // Each of these acquires a dedicated connection and runs all statements within @@ -4457,40 +4550,173 @@ async fn exec_tx_sqlite_inner( pool: db::sqlite::SqliteHandle, statements: &[String], start: std::time::Instant, + budget: &DbOperationBudget, ) -> Result { let statements = statements.to_vec(); + let query_timeout = budget.query_timeout; tokio::task::spawn_blocking(move || { pool.with_connection(|conn| { conn.execute_batch("BEGIN").map_err(|e| format!("Failed to begin transaction: {}", e))?; + + // rusqlite's synchronous API blocks for the full duration of a single + // statement, so a per-statement-boundary elapsed check alone cannot + // interrupt a statement that itself runs past the query budget. Use a + // cross-thread watchdog: it sleeps on a condvar for the budget, then + // fires `InterruptHandle::interrupt()` (SQLITE_INTERRUPT) which aborts + // the currently executing statement mid-flight. The watchdog is disarmed + // and joined before COMMIT so a stale interrupt can never land on a + // future, unrelated query of this pooled connection. + let interrupt = conn.get_interrupt_handle(); + let armed = Arc::new((Mutex::new(true), Condvar::new())); + let watchdog = match query_timeout { + Some(timeout) => { + let armed = armed.clone(); + Some(std::thread::spawn(move || { + let (lock, cvar) = &*armed; + let mut guard = lock.lock().unwrap(); + let mut should_interrupt = false; + if *guard { + // Armed at wait start: wait for the budget or until the + // main thread disarms (notifies) after finishing. A + // spurious wakeup re-waits for the remaining budget, so + // interrupt() only fires when the budget genuinely + // elapsed while still armed. + let wait_start = std::time::Instant::now(); + let mut remaining = timeout; + loop { + let (guard2, wait_result) = cvar + .wait_timeout(guard, remaining) + .expect("sqlite tx watchdog condvar wait poisoned"); + guard = guard2; + if !*guard { + // Disarmed: the main thread finished first. + break; + } + if wait_result.timed_out() || wait_start.elapsed() >= timeout { + // Budget elapsed while still armed: interrupt. + should_interrupt = true; + break; + } + remaining = timeout.saturating_sub(wait_start.elapsed()); + } + } + if should_interrupt { + interrupt.interrupt(); + } + })) + } + None => None, + }; + + let mut timeout_error: Option = None; + let mut statement_error: Option = None; let mut total_affected: u64 = 0; - for (i, sql) in statements.iter().enumerate() { - match conn.execute_batch(sql) { - Ok(_) => total_affected += conn.changes(), - Err(e) => { - let _ = conn.execute_batch("ROLLBACK"); - return Err(query_error_with_omitted_sql_context( - &format!("Statement {} failed: {}", i + 1, e), - sql, - )); + let result = (|| { + for (i, sql) in statements.iter().enumerate() { + // Defense in depth: the boundary check still runs (it also + // guarantees we never COMMIT past the budget even if the + // watchdog was not armed because the timeout is None). + if let Some(timeout) = query_timeout { + if start.elapsed() >= timeout { + timeout_error = Some(format!("Query timed out after {} seconds", timeout.as_secs())); + return Err(()); + } + } + match conn.execute_batch(sql) { + Ok(_) => total_affected += conn.changes(), + Err(e) => { + // The watchdog interrupt aborts the statement with + // SQLITE_INTERRUPT; surface it as a query timeout + // (matching `mysql_query_iter_with_timeout` wording so + // `is_dbx_query_timeout_error` recognizes it). The + // interrupt is detected by the SQLITE_INTERRUPT error + // code, never by matching "interrupt" in the message + // text (user/trigger/constraint text could otherwise + // be misclassified as a timeout). + // + // The SQLITE_INTERRUPT error-code match is only + // consulted when a query budget is set: the watchdog + // (this function's only in-process source of + // SQLITE_INTERRUPT intended to be a timeout) is only + // armed when query_timeout is Some. An EXTERNAL + // interrupt (e.g. the query_cancel mechanism) with no + // budget must surface as a normal statement error, not + // a timeout — misclassifying it panics on + // `query_timeout.unwrap()` below (None) and skips the + // ROLLBACK/disarm below, leaking the open transaction. + let timed_out = query_timeout.is_some_and(|timeout| { + start.elapsed() >= timeout + || matches!( + e.sqlite_error_code(), + Some(rusqlite::ffi::ErrorCode::OperationInterrupted) + ) + }); + if timed_out { + timeout_error = Some(format!( + "Query timed out after {} seconds", + query_timeout.unwrap_or_default().as_secs() + )); + return Err(()); + } + statement_error = Some(query_error_with_omitted_sql_context( + &format!("Statement {} failed: {}", i + 1, e), + sql, + )); + return Err(()); + } } } + Ok(()) + })(); + + // Disarm and join the watchdog before issuing COMMIT (or ROLLBACK) so + // no interrupt can be delivered after the transaction ends. + { + let (lock, cvar) = &*armed; + *lock.lock().unwrap() = false; + cvar.notify_all(); + } + if let Some(watchdog) = watchdog { + watchdog.join().expect("sqlite tx watchdog thread joined"); + } + + match result { + Ok(()) => { + // The boundary check above runs at the top of each loop + // iteration only. After the last statement the loop returns + // Ok and COMMIT would run without any elapsed re-check, so a + // statement that started under budget and finished after the + // budget elapsed could still COMMIT. Guard once more here, + // before COMMIT, so we never COMMIT past the budget. + if let Some(timeout) = query_timeout { + if start.elapsed() >= timeout { + let _ = conn.execute_batch("ROLLBACK"); + return Err(format!("Query timed out after {} seconds", timeout.as_secs())); + } + } + conn.execute_batch("COMMIT").map_err(|e| format!("COMMIT failed: {}", e)).map(|_| db::QueryResult { + columns: vec![], + column_types: Vec::new(), + column_sortables: vec![], + spatial_columns: vec![], + spatial_values: vec![], + rows: vec![], + affected_rows: total_affected, + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages: Vec::new(), + }) + } + Err(()) => { + let _ = conn.execute_batch("ROLLBACK"); + Err(timeout_error + .or(statement_error) + .unwrap_or_else(|| "Statement execution failed inside the transaction".to_string())) + } } - conn.execute_batch("COMMIT").map_err(|e| format!("COMMIT failed: {}", e))?; - Ok(db::QueryResult { - columns: vec![], - column_types: Vec::new(), - column_sortables: vec![], - spatial_columns: vec![], - spatial_values: vec![], - rows: vec![], - affected_rows: total_affected, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - session_id: None, - has_more: false, - elasticsearch_raw_body: None, - messages: Vec::new(), - }) }) }) .await @@ -4575,6 +4801,7 @@ async fn exec_tx_agent_inner( statements: &[String], schema: Option<&str>, start: std::time::Instant, + budget: &DbOperationBudget, ) -> Result { let execution_schema = schema_for_execution_context(db_type, schema); let rewritten_statements; @@ -4588,56 +4815,11 @@ async fn exec_tx_agent_inner( statements }; let mut client = client.lock().await; - let result: db::QueryResult = client.execute_transaction_typed(database, statements, execution_schema).await?; + let result: db::QueryResult = + client.execute_transaction_typed(database, statements, execution_schema, budget.query_timeout).await?; Ok(db::QueryResult { execution_time_ms: start.elapsed().as_millis(), ..result }) } -async fn exec_tx_none_inner( - state: &AppState, - pool_key: &str, - mysql_dialect: db::mysql::MySqlQueryDialect, - database: Option<&str>, - statements: &[String], - schema: Option<&str>, - start: std::time::Instant, -) -> Result { - let mut total_affected: u64 = 0; - for (i, sql) in statements.iter().enumerate() { - log::info!("[query][tx-none:statement:start] index={}", i + 1); - match do_execute(state, pool_key, mysql_dialect, database, sql, schema, None, QueryExecutionOptions::default()) - .await - { - Ok(result) => { - total_affected += result.affected_rows; - log::info!("[query][tx-none:statement:done] index={} affected_rows={}", i + 1, result.affected_rows); - } - Err(e) => { - log::warn!("Statement {} failed (no transaction support): {}", i + 1, e); - return Err(query_error_with_omitted_sql_context( - &format!("Statement {} failed: {}. No transaction support for this database type.", i + 1, e), - sql, - )); - } - } - } - - Ok(db::QueryResult { - columns: vec![], - column_types: Vec::new(), - column_sortables: vec![], - spatial_columns: vec![], - spatial_values: vec![], - rows: vec![], - affected_rows: total_affected, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - session_id: None, - has_more: false, - elasticsearch_raw_body: None, - messages: Vec::new(), - }) -} - /// Start a manual transaction session, holding a connection from the pool. /// Returns a transaction session ID that must be passed to subsequent calls. pub async fn begin_manual_transaction( @@ -5801,6 +5983,285 @@ mod tests { assert!(!postgres_prefers_text_protocol(None)); } + #[test] + fn apply_query_timeout_override_respects_resolve_semantics() { + // None leaves the budget unchanged. + let mut budget = DbOperationBudget::with_defaults(); + let original = budget.query_timeout; + apply_query_timeout_override(&mut budget, None); + assert_eq!(budget.query_timeout, original); + + // Some(5) sets query_timeout to 5s. + let mut budget = DbOperationBudget::with_defaults(); + apply_query_timeout_override(&mut budget, Some(5)); + assert_eq!(budget.query_timeout, Some(Duration::from_secs(5))); + + // Some(0) clears the limit (unlimited), matching resolve_query_timeout. + let mut budget = DbOperationBudget::with_defaults(); + apply_query_timeout_override(&mut budget, Some(0)); + assert_eq!(budget.query_timeout, None); + } + + #[tokio::test] + async fn sqlite_transaction_timeout_rolls_back_and_commits_nothing() { + use std::sync::mpsc; + + let pool = db::sqlite::connect_path(":memory:").await.expect("connect in-memory SQLite"); + db::sqlite::execute_query(&pool, "CREATE TABLE t (id INTEGER PRIMARY KEY, val TEXT)") + .await + .expect("create table"); + + // Deterministic timeout: hold the connection lock from a helper thread so the + // transaction's first statement-boundary check is guaranteed to observe + // elapsed time >= the 1ms query budget, regardless of machine speed. + let mut budget = DbOperationBudget::with_defaults(); + budget.query_timeout = Some(Duration::from_millis(1)); + let (lock_held_tx, lock_held_rx) = mpsc::channel(); + let holder = { + let pool = pool.clone(); + std::thread::spawn(move || { + pool.with_connection(|_conn| { + let _ = lock_held_tx.send(()); + std::thread::sleep(Duration::from_millis(200)); + Ok(()) + }) + .expect("helper holds sqlite connection lock"); + }) + }; + lock_held_rx.recv().expect("helper acquired sqlite connection lock"); + + let error = exec_tx_sqlite_inner( + pool.clone(), + &["INSERT INTO t (val) VALUES ('one')".to_string(), "INSERT INTO t (val) VALUES ('two')".to_string()], + std::time::Instant::now(), + &budget, + ) + .await + .expect_err("sqlite transaction must time out"); + + assert!(error.contains("Query timed out after"), "unexpected error: {error}"); + + // The transaction was rolled back (or never got past the first statement): + // no partial rows may survive. + holder.join().expect("helper thread joined"); + let result = db::sqlite::execute_query(&pool, "SELECT COUNT(*) AS n FROM t").await.expect("count rows"); + assert_eq!(result.rows[0][0], serde_json::json!(0)); + } + + #[tokio::test] + async fn sqlite_slow_statement_is_interrupted_and_rolled_back() { + // A statement that genuinely BLOCKS past the query budget must be aborted + // mid-flight by the watchdog (SQLITE_INTERRUPT) rather than running to + // completion and then committing. This is the gap the lock-holder test + // above does not cover: that one only makes the first statement-boundary + // check observe an already-elapsed budget. + // + // Lock-holder approach verified empirically and REJECTED: holding a write + // lock from a second connection and letting the main handle busy-wait is + // NOT interrupted by sqlite3_interrupt on the bundled SQLite 3.45.3 + // (rusqlite 0.32). `pager_wait_on_lock` loops on the busy handler without + // re-checking `db->u1.isInterrupted`, so the wait runs out the full busy + // timeout (a 60s wait with a 50ms budget confirmed it) instead of failing + // fast. So the deterministic proof uses a slow-but-bounded statement whose + // VDBE loop re-checks the interrupt flag every iteration (WITH RECURSIVE + // row generator) — the watchdog interrupts it mid-flight. + let dir = tempfile::tempdir().expect("temp dir"); + let db_path = dir.path().join("slow.db"); + let db_path = db_path.to_str().expect("utf8 temp path"); + + let pool = db::sqlite::connect_path_create_if_missing(db_path).await.expect("connect sqlite"); + db::sqlite::execute_query(&pool, "CREATE TABLE t (id INTEGER PRIMARY KEY, val TEXT)") + .await + .expect("create table"); + + let mut budget = DbOperationBudget::with_defaults(); + budget.query_timeout = Some(Duration::from_millis(50)); + + // Generating 50M rows takes seconds, so this single statement is provably + // still running when the 50ms watchdog fires. A busy-wait would not prove + // the point (see above); this statement is interrupted mid-flight. It is + // the LAST statement of a 2-statement batch so it also exercises the + // loop-exit-then-final-guard path: a fast first statement, then a slow + // second one that is still running when the budget elapses. + let slow_sql = "INSERT INTO t (val) SELECT 'slow' FROM (WITH RECURSIVE cnt(x) AS \ + (SELECT 1 UNION ALL SELECT x + 1 FROM cnt WHERE x < 50000000) SELECT x FROM cnt)"; + let error = exec_tx_sqlite_inner( + pool.clone(), + &["INSERT INTO t (val) VALUES ('fast')".to_string(), slow_sql.to_string()], + std::time::Instant::now(), + &budget, + ) + .await + .expect_err("sqlite transaction with a slow statement must time out"); + + assert!(error.contains("Query timed out after"), "unexpected error: {error}"); + let result = db::sqlite::execute_query(&pool, "SELECT COUNT(*) AS n FROM t").await.expect("count rows"); + assert_eq!(result.rows[0][0], serde_json::json!(0)); + } + + #[tokio::test] + async fn sqlite_legitimate_interrupt_text_error_is_not_reported_as_timeout() { + // A genuine statement error whose MESSAGE contains "interrupt" (here a + // missing column named `interrupted_at`) must NOT be misclassified as a + // watchdog timeout. The interrupt is detected by the SQLITE_INTERRUPT + // error code only, never by matching the message text. A large 60s budget + // guarantees elapsed can never trigger the timeout path, so only the + // error-code match could classify it. + let pool = db::sqlite::connect_path(":memory:").await.expect("connect in-memory SQLite"); + db::sqlite::execute_query(&pool, "CREATE TABLE t (id INTEGER PRIMARY KEY, val TEXT)") + .await + .expect("create table"); + + let mut budget = DbOperationBudget::with_defaults(); + budget.query_timeout = Some(Duration::from_secs(60)); + + let error = exec_tx_sqlite_inner( + pool.clone(), + &["INSERT INTO t (interrupted_at) VALUES (1)".to_string()], + std::time::Instant::now(), + &budget, + ) + .await + .expect_err("sqlite transaction must fail with the statement error"); + + assert!( + !error.contains("Query timed out after"), + "legitimate 'interrupt'-text error must not be masked as a timeout: {error}" + ); + assert!(error.contains("Statement 1 failed") && error.contains("interrupted_at"), "unexpected error: {error}"); + } + + #[tokio::test] + async fn sqlite_external_interrupt_without_budget_does_not_panic() { + // Regression: with query_timeout = None (MCP policy can set it to 0 = + // unlimited; resolve_query_timeout(Some(0)) -> None), an SQLITE_INTERRUPT + // arriving from an EXTERNAL source (e.g. the query_cancel mechanism, not + // this function's watchdog, which is only armed when query_timeout is + // Some) used to be misclassified as a timeout, then `query_timeout.unwrap()` + // panicked on None. The panic fired before the watchdog disarm/join and + // before ROLLBACK, leaking the watchdog thread and leaving the BEGIN + // transaction open on the pooled connection. + // + // Determinism: the bundled SQLite 3.45.3 clears the interrupt flag at + // VDBE step start whenever `nVdbeActive == 0` (sqlite3Step), so a + // pre-set interrupt fires only if it lands mid-statement. This was + // verified empirically below. So the interrupt is issued from a timer + // thread DURING a slow-but-bounded statement whose VDBE loop re-checks + // the interrupt flag every iteration (WITH RECURSIVE row generator — + // same technique the watchdog test uses; a lock busy-wait would not be + // aborted because pager_wait_on_lock does not recheck isInterrupted). + let pool = db::sqlite::connect_path(":memory:").await.expect("connect in-memory SQLite"); + db::sqlite::execute_query(&pool, "CREATE TABLE t (id INTEGER PRIMARY KEY, val TEXT)") + .await + .expect("create table"); + + let interrupt = + pool.with_connection(|conn| Ok(conn.get_interrupt_handle())).expect("get sqlite interrupt handle"); + let mut budget = DbOperationBudget::with_defaults(); + budget.query_timeout = None; // external interrupt: no in-process watchdog is armed + + // Arm the timer BEFORE the statement starts and send the go signal first, + // so the interrupt can land mid-statement. A single pre-set interrupt + // would be cleared (nVdbeActive==0 at step start), so the timer fires in + // a short retry loop: as soon as the slow statement is running, the flag + // sticks and the VDBE aborts on its next interrupt check. The `done` flag + // stops the loop as soon as exec returns so no stray interrupt can hit a + // later query, and we join the timer before any further query anyway. + let done = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let (timer_tx, timer_rx) = std::sync::mpsc::channel(); + let timer = { + let done = done.clone(); + std::thread::spawn(move || { + timer_rx.recv().expect("go signal"); + let deadline = std::time::Instant::now() + Duration::from_secs(5); + while !done.load(Ordering::Relaxed) && std::time::Instant::now() < deadline { + interrupt.interrupt(); // external interrupt mid-statement + std::thread::sleep(Duration::from_millis(5)); + } + }) + }; + timer_tx.send(()).expect("send go signal"); + + // Single slow statement that takes seconds to run, so an interrupt that + // lands within milliseconds of its start always hits mid-execution. + let slow_sql = "INSERT INTO t (val) SELECT 'slow' FROM (WITH RECURSIVE cnt(x) AS \ + (SELECT 1 UNION ALL SELECT x + 1 FROM cnt WHERE x < 10000000) SELECT x FROM cnt)"; + let error = exec_tx_sqlite_inner(pool.clone(), &[slow_sql.to_string()], std::time::Instant::now(), &budget) + .await + .expect_err("sqlite transaction must fail with the external interrupt"); + + done.store(true, Ordering::Relaxed); + timer.join().expect("timer thread joined"); + + // Must NOT panic, must NOT be masked as a timeout (no budget), and must + // surface as a plain statement error like any other failure. + assert!( + !error.contains("Query timed out after"), + "external interrupt with no budget must not be reported as a timeout: {error}" + ); + assert!( + error.contains("Statement 1 failed") && error.contains("interrupted"), + "expected a plain interrupted-statement error, got: {error}" + ); + + // The external interrupt must still roll back the transaction (0 rows), + // just like any other statement failure. + let result = db::sqlite::execute_query(&pool, "SELECT COUNT(*) AS n FROM t").await.expect("count rows"); + assert_eq!(result.rows[0][0], serde_json::json!(0)); + } + + #[tokio::test] + async fn sqlite_pre_set_interrupt_is_cleared_before_next_statement() { + // Empirical probe for the pre-set-interrupt question: `sqlite3_interrupt` + // unconditionally sets the interrupt flag, but bundled SQLite 3.45.3 + // clears it at VDBE step start whenever `nVdbeActive == 0` (sqlite3Step, + // "prevents a call to sqlite3_interrupt from interrupting a statement + // that has not yet started"). Verified here: a pre-set interrupt (with + // nothing running) is cleared before the NEXT statement begins, so that + // statement runs to completion. A regression test therefore cannot rely + // on a pre-set interrupt — it must interrupt DURING a running statement + // (see sqlite_external_interrupt_without_budget_does_not_panic). + let pool = db::sqlite::connect_path(":memory:").await.expect("connect in-memory SQLite"); + db::sqlite::execute_query(&pool, "CREATE TABLE t (id INTEGER PRIMARY KEY, val TEXT)") + .await + .expect("create table"); + + let interrupt = + pool.with_connection(|conn| Ok(conn.get_interrupt_handle())).expect("get sqlite interrupt handle"); + + // Pre-set the interrupt, then run a fast statement on the SAME + // connection: the flag is cleared at step start, so the statement + // succeeds (it does NOT fail with SQLITE_INTERRUPT). + pool.with_connection(|conn| { + interrupt.interrupt(); // pre-set, nothing running yet + conn.execute_batch("INSERT INTO t (val) VALUES ('one')").map_err(|e| e.to_string()) + }) + .expect("pre-set interrupt must be cleared before the statement runs"); + + // Same for a transaction: pre-set the interrupt, then the transaction's + // first statement must still run and commit. + pool.with_connection(|_conn| { + interrupt.interrupt(); // pre-set, nothing running yet + Ok(()) + }) + .expect("pre-set interrupt on idle connection"); + + let mut budget = DbOperationBudget::with_defaults(); + budget.query_timeout = None; + let result = exec_tx_sqlite_inner( + pool.clone(), + &["INSERT INTO t (val) VALUES ('two')".to_string()], + std::time::Instant::now(), + &budget, + ) + .await + .expect("transaction with a pre-set interrupt must still succeed (flag cleared at step start)"); + assert_eq!(result.affected_rows, 1); + + let count = db::sqlite::execute_query(&pool, "SELECT COUNT(*) AS n FROM t").await.expect("count rows"); + assert_eq!(count.rows[0][0], serde_json::json!(2)); + } + #[test] fn execute_multi_result_manual_transaction_markers_serialize_conditionally() { let plain = ExecuteMultiResult::success_with_optional_server_large_values(empty_query_result(0), false); @@ -6649,6 +7110,79 @@ for line in sys.stdin: assert_eq!(!table_check.rows.is_empty(), continue_on_error); } + #[tokio::test] + async fn transactional_sqlite_batch_rolls_back_when_a_later_statement_fails() { + let dir = std::env::temp_dir().join(format!("dbx-query-transaction-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "sqlite-transaction"; + let sqlite = db::sqlite::connect_path_create_if_missing(dir.join("query.db").to_str().unwrap()).await.unwrap(); + state + .update_connection_pools(|pools| { + pools.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); + }) + .await; + state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Sqlite)); + + let error = execute_multi_core_with_options( + &state, + connection_id, + "", + "CREATE TABLE rolled_back_table (id INTEGER); INSERT INTO missing_table VALUES (1);", + None, + None, + QueryExecutionOptions { use_transaction: Some(true), ..Default::default() }, + ) + .await + .unwrap_err(); + assert!(error.contains("missing_table"), "unexpected transaction error: {error}"); + + let table_check = execute_sql_statement( + &state, + connection_id, + "", + "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'rolled_back_table'", + None, + None, + ) + .await + .unwrap(); + assert!(table_check.rows.is_empty(), "the failed batch must roll back its preceding DDL"); + + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn transactional_batch_rejects_an_unsupported_backend_before_execution() { + let dir = std::env::temp_dir().join(format!("dbx-query-unsupported-transaction-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "message-queue-transaction"; + state + .update_connection_pools(|pools| { + pools.insert(connection_id.to_string(), PoolKind::MessageQueue); + }) + .await; + state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Redis)); + + let error = execute_multi_core_with_options( + &state, + connection_id, + "", + "INSERT INTO first_statement VALUES (1); INSERT INTO second_statement VALUES (2);", + None, + None, + QueryExecutionOptions { use_transaction: Some(true), ..Default::default() }, + ) + .await + .unwrap_err(); + assert!(error.contains("cannot provide a rollbackable transaction"), "unexpected transaction error: {error}"); + + let _ = std::fs::remove_dir_all(dir); + } + #[test] fn agent_execute_batch_unsupported_detects_case_insensitive_method_errors() { assert!(is_agent_execute_batch_unsupported("Agent RPC error (-1): unknown method: execute_batch")); @@ -6662,6 +7196,98 @@ for line in sys.stdin: assert!(!is_agent_execute_batch_unsupported("Agent RPC error (-1): unknown method: execute_query")); } + #[test] + fn batch_transaction_ddl_is_unrollbackable_covers_oracle_and_mysql_family() { + let ddl = vec!["CREATE TABLE test_table (id INT)".to_string()]; + let dml = vec!["INSERT INTO test_table VALUES (1)".to_string()]; + + // MySQL-family and Oracle DDL implicitly commit; explicit transactions + // over such DDL cannot be rolled back, so use_transaction is rejected. + for db in [DatabaseType::Mysql, DatabaseType::Goldendb, DatabaseType::Oracle] { + assert!(batch_transaction_ddl_is_unrollbackable(Some(db), &ddl), "expected {db:?} to reject DDL"); + } + assert!(!batch_transaction_ddl_is_unrollbackable(Some(DatabaseType::Mysql), &dml)); + // Postgres and SQLite have transactional DDL and must both be allowed. + assert!(!batch_transaction_ddl_is_unrollbackable(Some(DatabaseType::Postgres), &ddl)); + assert!(!batch_transaction_ddl_is_unrollbackable(Some(DatabaseType::Sqlite), &ddl)); + // An unknown db type cannot be verified — do not risk rejecting valid DDL batches. + assert!(!batch_transaction_ddl_is_unrollbackable(None, &ddl)); + } + + #[tokio::test] + async fn connection_pool_is_sqlserver_agent_detects_agent_and_native_pools() { + let dir = std::env::temp_dir().join(format!("dbx-query-sqlserver-agent-flag-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + + // Non-SQL-Server connections never use the SQL Server agent splitter. + state.configs.write().await.insert("pg".to_string(), test_connection_config(DatabaseType::Postgres)); + assert!(!connection_pool_is_sqlserver_agent(&state, "pg", "").await); + + // SQL Server with no pool yet defaults to the native (non-agent) splitter. + state.configs.write().await.insert("mssql".to_string(), test_connection_config(DatabaseType::SqlServer)); + assert!(!connection_pool_is_sqlserver_agent(&state, "mssql", "").await); + + // A SQL Server connection backed by the agent driver is detected, so the + // MCP pre-check uses the same GO-batch splitter as the core. + state + .update_connection_pools(|pools| { + pools.insert( + "mssql".to_string(), + PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub()), + ); + }) + .await; + assert!(connection_pool_is_sqlserver_agent(&state, "mssql", "").await); + + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn transaction_kernel_dispatches_by_backend_capability_not_ddl_content() { + // Regression for the maintainer review: the unrollbackable-DDL cap now lives + // only in the opt-in use_transaction entry point. The shared transaction + // kernel dispatches on the backend's transaction capability, so the + // use_transaction error never leaks to callers that share the kernel but do + // not set it (schema-diff deploy, imports), which document a mixed outcome + // on failure for DDL that cannot be rolled back. + let dir = std::env::temp_dir().join(format!("dbx-query-tx-predispatch-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "message-queue-tx-ddl"; + state + .update_connection_pools(|pools| { + pools.insert(connection_id.to_string(), PoolKind::MessageQueue); + }) + .await; + state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Redis)); + + let error = execute_statements_in_transaction_on_pool( + &state, + connection_id, + connection_id, + "", + &["CREATE TABLE t (id INT)".to_string()], + None, + None, + ) + .await + .unwrap_err(); + + assert!( + error.contains("cannot provide a rollbackable transaction"), + "expected a backend-capability reject, got: {error}" + ); + assert!( + !error.contains("whose DDL cannot be rolled back"), + "the opt-in use_transaction DDL cap leaked into the shared transaction kernel: {error}" + ); + + let _ = std::fs::remove_dir_all(dir); + } + #[test] fn query_pool_error_policy_retries_reads_but_not_writes() { assert_eq!( @@ -7527,7 +8153,9 @@ for line in sys.stdin: ); assert_eq!( backend_error.detail(), - Some("Query timed out after 1 seconds\nSQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement.") + Some( + "Query timed out after 1 seconds\nSQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement." + ) ); } @@ -7678,7 +8306,9 @@ for line in sys.stdin: ); assert_eq!( external_driver_preview_fallback_sql(sql).as_deref(), - Some("SELECT \"id\", \"description\" AS \"description\", \"metadata\"::text AS \"metadata\", 'left(\"literal\", 1)' AS \"note\", 'T:139' AS \"__DBX_LARGE_VALUE_BYTES_T_1\" FROM \"job_details\" WHERE left(note, 1) = 'x' LIMIT 100") + Some( + "SELECT \"id\", \"description\" AS \"description\", \"metadata\"::text AS \"metadata\", 'left(\"literal\", 1)' AS \"note\", 'T:139' AS \"__DBX_LARGE_VALUE_BYTES_T_1\" FROM \"job_details\" WHERE left(note, 1) = 'x' LIMIT 100" + ) ); } @@ -7691,7 +8321,9 @@ for line in sys.stdin: ); assert_eq!( external_driver_preview_fallback_sql(sql).as_deref(), - Some("SELECT coalesce(\"description\", concat('a,b', \"fallback\")) AS \"description\", 'FROM left(\"literal\", 1)' AS \"note\", 'T:140' AS \"__DBX_LARGE_VALUE_BYTES_T_0\" FROM \"job_details\"") + Some( + "SELECT coalesce(\"description\", concat('a,b', \"fallback\")) AS \"description\", 'FROM left(\"literal\", 1)' AS \"note\", 'T:140' AS \"__DBX_LARGE_VALUE_BYTES_T_0\" FROM \"job_details\"" + ) ); } @@ -8248,7 +8880,9 @@ for line in sys.stdin: assert_eq!( error.detail(), - Some("Server error: `ERROR 1064 (42000): syntax error` SQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement.") + Some( + "Server error: `ERROR 1064 (42000): syntax error` SQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement." + ) ); } diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index a9a2c949b1..fc37769680 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -145,6 +145,74 @@ fn effective_mcp_policy_with_legacy_allow_writes( policy } +/// Wire-level result for one statement within a `dbx_execute_batch` call. +/// +/// Mirrors the per-statement metadata the Web `/api/query/execute-multi` route +/// returns. `dbx_core::query::ExecuteMultiResult` only derives `Serialize`, so +/// a separate type that also derives `Deserialize` lets the Web backend decode +/// the JSON response while LocalBackend adapts from the core type. The error is +/// kept as an optional message string so `Deserialize` never has to handle the +/// private-field `BackendError` envelope. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BatchStatementResult { + #[serde(flatten)] + pub result: dbx_core::db::QueryResult, + #[serde(skip_serializing_if = "is_false")] + pub execution_error: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub statement_index: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error_message: Option, + /// True when this entry is the single merged outcome of a `use_transaction` + /// batch rather than an auto-commit per-statement result. Set by the MCP + /// server, which knows the requested mode; the wire decoders default to + /// false. + #[serde(skip_serializing_if = "is_false")] + pub merged: bool, +} + +fn is_false(value: &bool) -> bool { + !*value +} + +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(|| { + result.result.rows.first().and_then(|row| row.first()).and_then(Value::as_str).map(str::to_owned) + }); + Self { + result: result.result, + execution_error: result.execution_error, + statement_index: result.statement_index, + error_message, + merged: false, + } + } +} + +/// Optionally decode a loose per-statement JSON object from the Web +/// `/api/query/execute-multi` response into a `BatchStatementResult`. +/// +/// The Web route serializes `dbx_core::query::ExecuteMultiResult`, whose fields +/// are flattened onto the `QueryResult` and whose `error` is a `BackendError` +/// envelope with private fields and no `Deserialize` impl. This decodes only +/// the fields the MCP renderer needs instead of reconstructing the envelope. +/// Returns an explicit error when the response shape is not what this client +/// expects, so a protocol drift fails loudly instead of silently dropping +/// statements from the batch. +fn batch_statement_result_from_json(value: &Value) -> Result { + let result: dbx_core::db::QueryResult = serde_json::from_value(value.clone()) + .map_err(|error| format!("Invalid execute-multi statement envelope: {error} (element: {value})"))?; + let execution_error = value.get("execution_error").and_then(Value::as_bool).unwrap_or(false); + let statement_index = value.get("statement_index").and_then(Value::as_u64).map(|index| index as usize); + let error_message = value + .pointer("/error/detail") + .and_then(Value::as_str) + .map(str::to_owned) + .or_else(|| value.pointer("/rows/0/0").and_then(Value::as_str).map(str::to_owned)); + Ok(BatchStatementResult { result, execution_error, statement_index, error_message, merged: false }) +} + /// Wire-level options for a documentation snapshot. Mirrors /// `dbx_core::docs::CollectOptions` minus the fields the backend fills in. #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -201,6 +269,38 @@ pub trait DbxBackend: Send + Sync { let _ = (connection, database, sql, max_rows, timeout_secs); Err("SQL queries are not supported by this backend.".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 + /// semicolons inside strings, comments, and stored procedures are handled. + /// `schema` carries the configured scope schema when present. `continue_on_error` + /// / `use_transaction` / `client_session_id` are carried through `options`. + /// Mirrors the Web `/api/query/execute-multi` route. + async fn execute_batch( + &self, + connection: &ConnectionConfig, + database: &str, + schema: Option<&str>, + sql: &str, + options: dbx_core::query::QueryExecutionOptions, + ) -> Result, String> { + let _ = (connection, database, schema, sql, options); + Err("SQL batch execution is not supported by this backend.".to_string()) + } + /// Dialect-aware execution plan for a batch script, computed the way the core + /// will split it (including the SQL Server agent `GO`-batch splitter), so the + /// MCP pre-check's transaction-entry decision never diverges from the core. + /// The in-process backend overrides this to resolve the pool's agent-ness; + /// the default is the database-dialect splitter. + async fn execution_plan( + &self, + connection: &ConnectionConfig, + database: &str, + sql: &str, + ) -> dbx_core::sql::SqlExecutionPlan { + let _ = database; + dbx_core::sql::sql_execution_plan_for_database(sql, connection.db_type) + } async fn add_connection_for_mcp(&self, config: ConnectionConfig) -> Result; async fn duplicate_connection_for_mcp( &self, @@ -687,6 +787,42 @@ impl DbxBackend for LocalBackend { .await } + async fn execute_batch( + &self, + connection: &ConnectionConfig, + database: &str, + schema: Option<&str>, + sql: &str, + options: dbx_core::query::QueryExecutionOptions, + ) -> Result, String> { + // The multi-statement executor dispatches across every database engine + // and produces a very large future. Boxing it keeps the async-trait + // implementation below Rust's type-layout recursion limit (query depth + // overflow reported at this async block) without changing behaviour. + let results = Box::pin(dbx_core::query::execute_multi_core_with_options_for_client( + &self.state, + &connection.id, + database, + sql, + schema, + None, + options, + )) + .await?; + Ok(results.into_iter().map(BatchStatementResult::from).collect()) + } + + async fn execution_plan( + &self, + connection: &ConnectionConfig, + database: &str, + sql: &str, + ) -> dbx_core::sql::SqlExecutionPlan { + let is_sqlserver_agent = + dbx_core::query::connection_pool_is_sqlserver_agent(self.state.as_ref(), &connection.id, database).await; + dbx_core::query::query_execution_plan(sql, Some(connection.db_type), is_sqlserver_agent) + } + async fn add_connection_for_mcp(&self, config: ConnectionConfig) -> Result { let config = self.state.storage.add_connection_for_mcp(config).await?; self.state.configs.write().await.insert(config.id.clone(), config.clone()); @@ -1052,6 +1188,55 @@ impl DbxBackend for WebBackend { .map_err(|error| format!("Invalid query response: {error}")) } + async fn execute_batch( + &self, + connection: &ConnectionConfig, + database: &str, + schema: Option<&str>, + sql: &str, + options: dbx_core::query::QueryExecutionOptions, + ) -> Result, String> { + if connection.db_type == DatabaseType::MongoDb { + return Err("MongoDB batch execution in DBX Web mode is not implemented by the Rust CLI yet.".to_string()); + } + self.ensure_connected(connection).await?; + let mut body = json!({ + "connectionId": connection.id, + "database": database, + "sql": sql, + // Keep Web-mode structured results within the MCP 100-row contract. + "maxRows": options.max_rows, + "timeoutSecs": agent_tools::agent_query_timeout_secs(options.timeout_secs, Some(connection)), + }); + if let Some(schema) = schema.map(str::trim).filter(|schema| !schema.is_empty()) { + body["schema"] = json!(schema); + } + if let Some(client_session_id) = options.client_session_id.as_deref().filter(|id| !id.trim().is_empty()) { + body["clientSessionId"] = json!(client_session_id); + } + if options.continue_on_error { + body["continueOnError"] = json!(true); + } + if options.use_transaction == Some(true) { + body["useTransaction"] = json!(true); + } + let value: Value = self + .request(reqwest::Method::POST, "/api/query/execute-multi", Some(body)) + .await? + .json() + .await + .map_err(|error| format!("Invalid execute-multi response: {error}"))?; + // Fail loudly on any element that does not match the expected envelope + // instead of silently dropping statements from the batch. + let results = value + .as_array() + .ok_or_else(|| "Invalid execute-multi response: expected an array of per-statement results".to_string())? + .iter() + .map(batch_statement_result_from_json) + .collect::, String>>()?; + Ok(results) + } + async fn close_client_session( &self, connection_id: &str, @@ -2341,6 +2526,109 @@ mod tests { assert_eq!(second_request["timeoutSecs"], 300); } + #[tokio::test] + async fn web_execute_batch_hits_execute_multi_and_decodes_statement_results() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let (request_sender, request_receiver) = mpsc::channel(); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + stream.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let count = stream.read(&mut buffer).unwrap(); + request.extend_from_slice(&buffer[..count]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length").then_some(value.trim()) + }) + .unwrap() + .parse::() + .unwrap(); + while request.len() < header_end + content_length { + let count = stream.read(&mut buffer).unwrap(); + request.extend_from_slice(&buffer[..count]); + } + let request = String::from_utf8(request).unwrap(); + let body = request[header_end..header_end + content_length].to_string(); + request_sender.send((request.lines().next().unwrap().to_string(), body)).unwrap(); + + let response_body = r#"[{"columns":["id"],"column_types":[],"column_sortables":[],"rows":[["1"],["2"]],"affected_rows":0,"execution_time_ms":1,"truncated":false,"has_more":false,"statement_index":0},{"columns":[],"column_types":[],"column_sortables":[],"rows":[],"affected_rows":2,"execution_time_ms":1,"truncated":false,"has_more":false,"statement_index":1}]"#; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + response_body.len(), + response_body + ) + .unwrap(); + }); + + // The mock is a loopback server, so it must not inherit a contributor's + // outbound proxy configuration. + let backend = + WebBackend::new_with_config(format!("http://{address}"), String::new(), None, None, None, false, None) + .unwrap(); + backend.auth.lock().await.checked = true; + let connection = new_connection_config( + "web-batch".to_string(), + "web-batch".to_string(), + DatabaseType::Postgres, + "localhost".to_string(), + 5432, + String::new(), + String::new(), + None, + false, + None, + ) + .unwrap(); + // Pre-seed so the mock only needs to answer /api/query/execute-multi. + backend.connected.lock().await.insert(connection.id.clone(), connection.clone()); + + let results = backend + .execute_batch( + &connection, + "postgres", + None, + "SELECT 1; INSERT INTO t VALUES (1)", + dbx_core::query::QueryExecutionOptions { + max_rows: Some(100), + timeout_secs: Some(0), + continue_on_error: true, + use_transaction: Some(true), + ..Default::default() + }, + ) + .await + .unwrap(); + + assert_eq!(results.len(), 2); + // Two metadata fields are transferred from the JSON envelope. + assert_eq!(results[0].result.columns, vec!["id".to_string()]); + assert_eq!(results[0].result.rows.len(), 2); + assert_eq!(results[1].statement_index, Some(1)); + assert_eq!(results[1].result.affected_rows, 2); + + let (request_line, body) = request_receiver.recv().unwrap(); + assert_eq!(request_line, "POST /api/query/execute-multi HTTP/1.1"); + let request: Value = serde_json::from_str(&body).unwrap(); + assert_eq!(request["sql"], "SELECT 1; INSERT INTO t VALUES (1)"); + assert_eq!(request["continueOnError"], true); + assert_eq!(request["useTransaction"], true); + assert_eq!(request["maxRows"], 100); + assert_eq!(request["timeoutSecs"], 0); + + server.join().unwrap(); + } + #[test] fn format_query_result_appends_server_messages() { let mut result = query_result(Vec::new(), Vec::new(), 3); diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 9f58b17ffe..78c21f0eba 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -13,7 +13,7 @@ use crate::backend::{format_query_result, new_connection_config, parse_database_ use crate::mongo::{self, MongoCommand, MongoSafetyError}; use crate::session::{McpSession, McpSessionStore}; use dbx_core::{ - agent_tools::{format_query_result_as_text, QueryCellWindow}, + agent_tools::{format_query_result_as_text, normalize_sql_for_confirmation, QueryCellWindow}, database_manifest, db::redis_driver::{classify_command, parse_command_argv, RedisCommandResult, RedisCommandSafety}, models::connection::DatabaseType, @@ -270,6 +270,34 @@ pub struct ExecuteAndShowRequest { pub database: Option, } +#[derive(Debug, Deserialize, schemars::JsonSchema)] +pub struct ExecuteBatchQueryRequest { + #[serde(flatten)] + pub selector: ConnectionSelector, + #[schemars(description = "Database name")] + #[schemars(extend("type" = "string"))] + pub database: Option, + #[schemars( + description = "SQL script containing one or more statements. Statements are split using the database-dialect-aware splitter, so semicolons inside strings, comments, and stored procedures are handled." + )] + pub sql: String, + #[schemars( + description = "Session ID from dbx_open_session. When set, every statement runs on the session's pinned connection, preserving USE/SET, temporary tables and other session state across statements in the script and across calls." + )] + #[schemars(extend("type" = "string"))] + pub session_id: Option, + #[schemars( + description = "When true, execution continues after a statement error instead of stopping at the first failure. Connection-level failures always stop the batch. Default false (stop at first failure)." + )] + #[schemars(extend("type" = "boolean"))] + pub continue_on_error: Option, + #[schemars( + description = "When true and the script has more than one statement, all statements run inside one transaction (BEGIN … COMMIT) so the whole batch succeeds or rolls back, and the call returns a single merged result. For a single-statement script the option is ignored and the statement runs as normal auto-commit. Backends that cannot provide a rollbackable transaction reject this option. Cannot be combined with session_id or continue_on_error. Rejected for MySQL-family DDL scripts because DDL implicitly commits and cannot be rolled back. Default is auto-commit for each statement." + )] + #[schemars(extend("type" = "boolean"))] + pub use_transaction: Option, +} + #[derive(Clone)] pub struct DbxMcpServer { backend: Arc, @@ -584,7 +612,7 @@ impl DbxMcpServer { return tool_error( "SESSION_CONNECTION_MISMATCH", format!("Session \"{session_id}\" is bound to a different connection."), - ) + ); } None => { return tool_error( @@ -592,7 +620,7 @@ impl DbxMcpServer { format!( "Session \"{session_id}\" not found or expired. Open a new one with dbx_open_session." ), - ) + ); } } } @@ -682,6 +710,194 @@ impl DbxMcpServer { agent_result(result) } + #[tool( + name = "dbx_execute_batch", + description = "Execute a SQL script containing multiple statements in one call and return a result per statement. Statements are split with a database-dialect-aware parser, so semicolons inside strings, comments and stored procedures are handled. Stops at the first failing statement unless continue_on_error is true. Pass session_id (from dbx_open_session) when statements must share one connection (e.g. temporary tables, USE/SET). When use_transaction is true and the script has multiple statements, the whole script runs in one transaction and the call returns a single merged result instead of one result per statement (single-statement scripts run normally and return that one result); use_transaction cannot be combined with session_id or continue_on_error, and is rejected when the backend cannot provide a rollbackable transaction or for MySQL-family DDL scripts (DDL implicitly commits and cannot be rolled back)." + )] + async fn execute_batch(&self, Parameters(request): Parameters) -> CallToolResult { + let resolved = match self.resolve_connection(&request.selector).await { + Ok(resolved) => resolved, + Err(error) => return error, + }; + let connection = &resolved.connection; + if matches!(connection.db_type, DatabaseType::Redis | DatabaseType::MongoDb) { + return tool_error( + "DBX_BATCH_UNSUPPORTED", + format!("Batch SQL execution is not available for {connection:?} connections."), + ); + } + let sql = request.sql.trim(); + if sql.is_empty() { + return tool_error("SQL_BATCH_EMPTY", "SQL script cannot be empty."); + } + let session = match request.session_id.as_deref().map(str::trim).filter(|id| !id.is_empty()) { + Some(session_id) => { + let (session, expired) = self.sessions.resolve(session_id).await.into_parts(); + self.close_backend_sessions_best_effort(expired).await; + match session { + Some(session) if session.connection_id == connection.id => Some(session), + Some(_) => { + return tool_error( + "SESSION_CONNECTION_MISMATCH", + format!("Session \"{session_id}\" is bound to a different connection."), + ); + } + None => { + return tool_error( + "SESSION_NOT_FOUND", + format!( + "Session \"{session_id}\" not found or expired. Open a new one with dbx_open_session." + ), + ); + } + } + } + None => None, + }; + let database = match self.resolve_database(request.database, &resolved) { + Ok(database) => database, + Err(error) => return error, + }; + if let Some(session) = &session { + if session.database != database { + return tool_error( + "SESSION_DATABASE_MISMATCH", + format!( + "Session \"{}\" is bound to database \"{}\", not \"{database}\".", + session.id, session.database + ), + ); + } + } + // A pinned session makes USE/SET CATALOG meaningful, so database + // switching is allowed — unless a hard database scope is configured. + let allow_database_switch = session.is_some() && self.scope.database.is_none(); + // Split with the same dialect-aware splitter the core uses, early, so the + // option-validity checks below only fire for scripts that actually enter + // transaction mode (more than one statement). A single-statement script + // ignores use_transaction and runs as normal auto-commit, so combining it + // with session_id / continue_on_error must still be allowed. + let execution_plan = self.backend.execution_plan(connection, &database, sql).await; + let statement_count = execution_plan.statements.len(); + let transactional = request.use_transaction == Some(true) && statement_count > 1; + // A transactional batch runs on a pooled connection and returns one + // merged result, so it cannot preserve session state or yield per- + // statement results. The core transaction path takes no client session + // id, so combining the two would silently drop the session pin. Only + // relevant when the script actually enters transaction mode. + if transactional && session.is_some() { + return tool_error( + "TRANSACTION_WITH_SESSION_UNSUPPORTED", + "use_transaction cannot be combined with session_id: transactional batches run on a pooled connection, discard session state, and return a single merged result. Run the batch either without use_transaction (auto-commit, one result per statement) or without session_id.", + ); + } + // The core transaction path rolls back and stops on the first failure + // (query.rs:2970), so continue_on_error is silently ignored there. + // Accepting both would promise per-statement continuation that never + // happens — reject the combination instead. Only relevant when the + // script actually enters transaction mode. + if transactional && request.continue_on_error == Some(true) { + return tool_error( + "TRANSACTION_WITH_CONTINUE_UNSUPPORTED", + "use_transaction cannot be combined with continue_on_error: a transactional batch stops and rolls back at the first failure, so continue_on_error would be ignored. Run the batch either without use_transaction (auto-commit, one result per statement, may continue on error) or without continue_on_error.", + ); + } + // DDL implicitly commits on MySQL-family and Oracle engines (and any + // backend without transactional DDL — the shared capability check), so + // the transaction wrapper cannot roll back a DDL that already committed + // (e.g. "CREATE TABLE a; INSERT INTO missing" leaves the table behind + // after the failed INSERT rolls back). A transactional batch containing + // DDL would over-promise atomicity — reject it instead. Gated on multi- + // statement like the other rejects: a single DDL statement is plain + // auto-commit and safe. + if transactional + && dbx_core::query::batch_transaction_ddl_is_unrollbackable( + Some(connection.db_type), + &execution_plan.statements, + ) + { + return tool_error( + "TRANSACTION_WITH_DDL_UNSUPPORTED", + "use_transaction cannot be used with a batch whose DDL cannot be rolled back: DDL statements implicitly commit and cannot be rolled back, so the batch would not be atomic. Run the batch without use_transaction (auto-commit, one result per statement) or split the DDL and DML into separate calls.", + ); + } + // Classify the whole script as a unit so a single write/DDL statement in + // the batch fails closed under read-only / dangerous-SQL / production + // policy, exactly as the Web /api/query/execute-multi route does. + let permissions = match validate_sql_policy(connection, &resolved.policy, &database, sql, allow_database_switch) + { + Ok(permissions) => permissions, + Err(error) => return error, + }; + // When the user confirmed a specific write SQL for this run, the batch + // must not smuggle a different write/DDL statement past that binding. + // Fail closed on the whole script when it contains a write and is not + // exactly the confirmed SQL — matching execute_query's anti-replay rule. + if let Some(message) = + confirmed_batch_sql_block_reason(sql, connection.db_type, permissions.confirmed_write_sql.as_deref()) + { + return tool_error("SQL_BLOCKED", message); + } + // A transactional batch runs on the plain (non-session) pool: the core + // transaction wrapper ignores client_session_id (query.rs:2970 → + // execute_statements_in_transaction_typed re-resolves the default pool + // with get_or_create_pool). Creating an ephemeral session pool here + // would open physical connections that are closed unused, so only + // auto-commit batches get an ephemeral pinned pool — it keeps temp + // tables and SET state alive for every statement in the script. + let ephemeral_client_session_id = + (!transactional && session.is_none()).then(|| format!("mcp-batch-{}", Uuid::new_v4())); + let client_session_id = session + .as_ref() + .map(|session| session.client_session_id.clone()) + .or_else(|| ephemeral_client_session_id.clone()); + let mut options = dbx_core::query::QueryExecutionOptions { + max_rows: Some(BATCH_MAX_ROWS), + continue_on_error: request.continue_on_error.unwrap_or(false), + use_transaction: request.use_transaction, + client_session_id: client_session_id.clone(), + ..Default::default() + }; + // Surface the global MCP query timeout. It is re-read on every tool + // call, so a settings change takes effect without client-config + // regeneration. + options.timeout_secs = resolved.policy.query_timeout_secs; + let schema = self.scope.schema.as_deref(); + let execution = self.backend.execute_batch(connection, &database, schema, sql, options).await; + if let Some(client_session_id) = ephemeral_client_session_id.as_deref() { + if let Err(error) = self.backend.close_client_session(&connection.id, &database, client_session_id).await { + log::warn!("failed to close ephemeral MCP batch session for {}: {error}", connection.id); + } + } + match execution { + Ok(mut results) => { + // The core transaction path collapses the whole script into one + // merged QueryResult. Mark it so callers can tell a merged + // outcome from a per-statement result, and render it distinctly + // instead of implying per-statement outcomes. The len()==1 guard + // keeps drivers that return per-statement results even with + // use_transaction (e.g. SQL Server scripts with result sets) + // from being mislabelled as transactional. + if transactional && results.len() == 1 { + results[0].merged = true; + results[0].statement_index = None; + } + let markdown = format_batch_results(&results); + // Issue #7548 requires a structured array so callers do not + // parse concatenated text. The Markdown block stays as a human- + // readable summary; structuredContent carries one object per + // statement (or the single merged outcome in transaction mode). + let mut tool_result = CallToolResult::success(vec![ContentBlock::text(markdown)]); + tool_result.structured_content = Some( + serde_json::to_value(&results) + .unwrap_or_else(|error| serde_json::json!({ "error": error.to_string() })), + ); + tool_result + } + Err(error) => backend_tool_error("DBX_BATCH_EXECUTION_ERROR", error), + } + } + #[tool( name = "dbx_open_session", description = "Open a stateful query session pinned to a single backend connection. Returns a session ID for dbx_execute_query: USE, SET CATALOG, session variables and temporary tables persist across calls within the session. Close with dbx_close_session when done; idle sessions expire after 30 minutes." @@ -1478,6 +1694,53 @@ fn backend_tool_error(default_code: &str, error: impl Into) -> CallToolR tool_error(default_code, error) } +/// Maximum rows returned per statement in a `dbx_execute_batch` call, matching +/// the `dbx_execute_query` "at most 100 rows" contract so batch responses stay +/// bounded. +const BATCH_MAX_ROWS: usize = 100; + +/// Render a multi-statement batch result as one Markdown text block per +/// statement. Each statement is labelled by its index so callers can see which +/// statement failed, how many rows each affected, or what each returned. A +/// `merged` entry (use_transaction mode) is labelled as the transaction outcome +/// instead of a per-statement result. +fn format_batch_results(results: &[crate::backend::BatchStatementResult]) -> String { + let mut output = String::new(); + for (index, result) in results.iter().enumerate() { + if index > 0 { + output.push('\n'); + } + if result.merged { + output.push_str("### Transaction outcome"); + output.push('\n'); + if result.execution_error { + let message = result.error_message.as_deref().unwrap_or("Transaction failed"); + output.push_str(&format!("**Status:** failed\n\n{message}\n")); + } else { + let rendered = format_query_result_as_text(&result.result, BATCH_MAX_ROWS, QueryCellWindow::default()) + .unwrap_or_else(|error| error); + output.push_str(rendered.trim_start_matches("Query executed. ")); + } + continue; + } + let heading = match result.statement_index { + Some(statement_index) => format!("### Statement {}", statement_index + 1), + None => format!("### Statement {}", index + 1), + }; + output.push_str(&heading); + output.push('\n'); + if result.execution_error { + let message = result.error_message.as_deref().unwrap_or("Statement failed"); + output.push_str(&format!("**Status:** failed\n\n{message}\n")); + } else { + let rendered = format_query_result_as_text(&result.result, BATCH_MAX_ROWS, QueryCellWindow::default()) + .unwrap_or_else(|error| error); + output.push_str(rendered.trim_start_matches("Query executed. ")); + } + } + output +} + fn 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: ")) @@ -1625,6 +1888,35 @@ fn normalize_confirmed_write_sql(value: Option) -> Option { (!trimmed.is_empty()).then_some(trimmed) } +/// When the user confirmed a specific write SQL for this run, a batch that +/// contains a write/DDL statement must be exactly that confirmed SQL. Returns +/// the block message when the script differs from the confirmation and is not +/// read-only; `None` means the script is allowed (read-only, exact match, or +/// no binding). Unparseable SQL is treated as a write and fails closed. +fn confirmed_batch_sql_block_reason( + sql: &str, + db_type: dbx_core::models::connection::DatabaseType, + confirmed: Option<&str>, +) -> Option { + let confirmed = confirmed?; + if normalize_sql_for_confirmation(sql) == normalize_sql_for_confirmation(confirmed) { + return None; + } + let is_write = match classify_sql_risk_for_database(sql, db_type) { + Ok(risk) => risk != dbx_core::sql_risk::SqlRisk::ReadOnly, + Err(_) => true, + }; + if !is_write { + return None; + } + Some(format!( + "Blocked: the executed SQL does not match the user-confirmed SQL.\n\ + Confirmed: {}\n\ + Attempted: {}", + confirmed, sql + )) +} + // CallToolResult is the transport-native error payload; boxing it would complicate every MCP call site. #[allow(clippy::result_large_err)] fn validate_sql_policy( @@ -1729,10 +2021,9 @@ fn validate_mongo_command( "SQL_BLOCKED", "MongoDB update/delete commands must include a non-empty filter unless high-risk operations are enabled in DBX MCP settings.", ), - MongoSafetyError::Dangerous => tool_error( - "SQL_BLOCKED", - "Dangerous MongoDB command is disabled in DBX MCP settings.", - ), + MongoSafetyError::Dangerous => { + tool_error("SQL_BLOCKED", "Dangerous MongoDB command is disabled in DBX MCP settings.") + } MongoSafetyError::ProductionWrite => { tool_error("PRODUCTION_WRITE_BLOCKED", "MCP cannot execute writes against a production database.") } @@ -2051,6 +2342,52 @@ mod tests { async fn remove_connection_for_mcp(&self, _connection_id: &str) -> Result { Ok(true) } + + async fn execute_batch( + &self, + _connection: &ConnectionConfig, + _database: &str, + schema: Option<&str>, + sql: &str, + options: dbx_core::query::QueryExecutionOptions, + ) -> Result, String> { + let client_session_id = + options.client_session_id.as_deref().filter(|id| !id.trim().is_empty()).map(str::to_owned); + if let Some(client_session_id) = &client_session_id { + self.pinned_sessions.lock().unwrap().insert(client_session_id.clone()); + } + self.recorded_arguments.lock().unwrap().push(( + "execute_batch".to_string(), + json!({ + "sql": sql, + "schema": schema, + "continue_on_error": options.continue_on_error, + "use_transaction": options.use_transaction, + "client_session_id": client_session_id, + }), + )); + Ok(vec![crate::backend::BatchStatementResult { + result: dbx_core::db::QueryResult { + columns: vec![], + column_types: vec![], + column_sortables: vec![], + spatial_columns: vec![], + spatial_values: vec![], + rows: vec![], + affected_rows: 0, + execution_time_ms: 1, + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages: vec![], + }, + execution_error: false, + statement_index: Some(0), + error_message: None, + merged: false, + }]) + } } #[test] @@ -2075,9 +2412,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(), 17); + assert_eq!(tools.len(), 18); #[cfg(not(feature = "mq-admin"))] - assert_eq!(tools.len(), 16); + assert_eq!(tools.len(), 17); assert!(names.contains(&"dbx_list_connections")); assert!(names.contains(&"dbx_list_databases")); assert!(names.contains(&"dbx_list_tables")); @@ -2085,6 +2422,7 @@ mod tests { assert!(names.contains(&"dbx_list_routines")); assert!(names.contains(&"dbx_get_routine_source")); assert!(names.contains(&"dbx_execute_query")); + assert!(names.contains(&"dbx_execute_batch")); assert!(names.contains(&"dbx_add_connection")); assert!(names.contains(&"dbx_duplicate_connection")); assert!(names.contains(&"dbx_remove_connection")); @@ -2197,6 +2535,7 @@ mod tests { ("dbx_list_routines", &["database", "schema", "routine_type"]), ("dbx_get_routine_source", &["database", "schema", "signature"]), ("dbx_execute_query", &["database", "session_id", "cell_char_offset", "cell_char_limit"]), + ("dbx_execute_batch", &["database", "session_id", "continue_on_error", "use_transaction"]), ("dbx_open_session", &["database"]), ("dbx_open_table", &["database", "schema"]), ("dbx_execute_and_show", &["database"]), @@ -2319,14 +2658,15 @@ mod tests { ); let names = server.tool_router.list_all().into_iter().map(|tool| tool.name).collect::>(); #[cfg(feature = "mq-admin")] - assert_eq!(names.len(), 12); + assert_eq!(names.len(), 13); #[cfg(not(feature = "mq-admin"))] - assert_eq!(names.len(), 11); + assert_eq!(names.len(), 12); assert!(!names.iter().any(|name| name == "dbx_add_connection")); assert!(!names.iter().any(|name| name == "dbx_duplicate_connection")); assert!(!names.iter().any(|name| name == "dbx_remove_connection")); assert!(!names.iter().any(|name| name == "dbx_open_table")); assert!(!names.iter().any(|name| name == "dbx_execute_and_show")); + assert!(names.iter().any(|name| name == "dbx_execute_batch")); assert!(names.iter().any(|name| name == "dbx_list_routines")); assert!(names.iter().any(|name| name == "dbx_get_routine_source")); assert!(names.iter().any(|name| name == "dbx_open_session")); @@ -2955,4 +3295,519 @@ mod tests { .await; assert!(result_text(&missing).contains("SESSION_NOT_FOUND")); } + + #[tokio::test] + async fn execute_batch_forwards_options_and_surfaces_results() { + let postgres = connection("pg", "pg", "postgres", "app"); + let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); + + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: Some("app".to_string()), + sql: "SELECT 1; SELECT 2".to_string(), + session_id: None, + continue_on_error: Some(true), + use_transaction: None, + })) + .await; + + // Auto-commit mode returns per-statement results; continue_on_error is + // forwarded to the core. + assert!(result_text(&result).contains("Statement 1")); + let ephemeral_session_id = { + let recorded = backend.recorded_arguments.lock().unwrap(); + let (_, arguments) = recorded.iter().find(|(name, _)| name == "execute_batch").unwrap(); + assert_eq!(arguments["continue_on_error"], true); + assert!( + arguments["use_transaction"].is_null(), + "use_transaction must be forwarded as null, got: {}", + arguments["use_transaction"] + ); + arguments["client_session_id"] + .as_str() + .expect("auto-commit batch must use an ephemeral session") + .to_string() + }; + assert!(ephemeral_session_id.starts_with("mcp-batch-")); + assert_eq!(backend.closed_sessions.lock().unwrap().as_slice(), [ephemeral_session_id]); + } + + #[tokio::test] + async fn execute_batch_pins_session_and_forwards_client_session_id() { + let postgres = connection("pg", "pg", "postgres", "app"); + let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); + + let opened = + server.open_session(Parameters(OpenSessionRequest { selector: selector("pg"), database: None })).await; + let session_id = opened_session_id(&opened); + + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "SELECT 1; SELECT 2".to_string(), + session_id: Some(session_id.clone()), + continue_on_error: Some(false), + use_transaction: None, + })) + .await; + assert!(result_text(&result).contains("Statement 1"), "got: {}", result_text(&result)); + let recorded = backend.recorded_arguments.lock().unwrap(); + let (_, arguments) = recorded.iter().find(|(name, _)| name == "execute_batch").unwrap(); + assert_eq!(arguments["client_session_id"], format!("mcp:{session_id}")); + assert!(backend.pinned_sessions.lock().unwrap().contains(&format!("mcp:{session_id}"))); + } + + #[tokio::test] + async fn execute_batch_blocks_redis_and_mongo_connections() { + for (id, db_type) in [("redis", "redis"), ("mongo", "mongodb")] { + let backend = + Arc::new(FakeBackend { connections: vec![connection(id, id, db_type, "db")], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector(id), + database: None, + sql: "SELECT 1".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: None, + })) + .await; + assert!(result_text(&result).contains("DBX_BATCH_UNSUPPORTED")); + } + } + + #[tokio::test] + async fn execute_batch_forwards_scoped_schema() { + let postgres = connection("pg", "pg", "postgres", "app"); + let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options( + backend.clone(), + McpScope { schema: Some("REPORTING".to_string()), ..Default::default() }, + false, + ); + + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "SELECT 1; SELECT 2".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: None, + })) + .await; + assert!(result_text(&result).contains("Statement 1")); + + let recorded = backend.recorded_arguments.lock().unwrap(); + let (_, arguments) = recorded.iter().find(|(name, _)| name == "execute_batch").unwrap(); + assert_eq!(arguments["schema"], "REPORTING"); + } + + #[tokio::test] + async fn execute_batch_rejects_empty_script() { + let postgres = connection("pg", "pg", "postgres", "app"); + let server = DbxMcpServer::with_runtime_options( + Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }), + McpScope::default(), + false, + ); + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: " ".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: None, + })) + .await; + assert!(result_text(&result).contains("SQL_BATCH_EMPTY")); + } + + #[tokio::test] + async fn execute_batch_rejects_transaction_with_session() { + let postgres = connection("pg", "pg", "postgres", "app"); + let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + + let opened = + server.open_session(Parameters(OpenSessionRequest { selector: selector("pg"), database: None })).await; + let session_id = opened_session_id(&opened); + + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "SELECT 1; SELECT 2".to_string(), + session_id: Some(session_id), + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + assert!(result_text(&result).contains("TRANSACTION_WITH_SESSION_UNSUPPORTED")); + } + + #[tokio::test] + async fn execute_batch_rejects_transaction_with_continue_on_error() { + // The core transaction path rolls back and stops at the first failure + // (query.rs:2970), so continue_on_error would be silently ignored. + // Reject the combination instead of promising continuation that never + // happens, and make sure neither option reaches the backend. + let postgres = connection("pg", "pg", "postgres", "app"); + let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); + + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)".to_string(), + session_id: None, + continue_on_error: Some(true), + use_transaction: Some(true), + })) + .await; + assert!(result_text(&result).contains("TRANSACTION_WITH_CONTINUE_UNSUPPORTED")); + assert!(backend.recorded_arguments.lock().unwrap().iter().all(|(name, _)| name != "execute_batch")); + } + + #[tokio::test] + async fn execute_batch_rejects_mysql_family_ddl_with_transaction() { + // MySQL-family engines implicitly commit before DDL, so a DDL batch + // with use_transaction cannot roll back a DDL that already committed. + // Reject the combination instead of over-promising atomicity. Plain + // DML batches and transactional DDL engines (PostgreSQL) still accept + // use_transaction. + let mysql = connection("my", "my", "mysql", "app"); + let backend = Arc::new(FakeBackend { connections: vec![mysql], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); + + let ddl = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("my"), + database: None, + sql: "CREATE TABLE a (id INT); INSERT INTO missing VALUES (1)".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + assert!(result_text(&ddl).contains("TRANSACTION_WITH_DDL_UNSUPPORTED")); + assert!(backend.recorded_arguments.lock().unwrap().iter().all(|(name, _)| name != "execute_batch")); + + // GoldenDB is a MySQL-compatible Agent driver, so its DDL has the + // same implicit-commit boundary and must not be advertised as atomic. + let goldendb = connection("goldendb", "goldendb", "goldendb", "app"); + let goldendb_backend = Arc::new(FakeBackend { connections: vec![goldendb], ..Default::default() }); + let goldendb_server = DbxMcpServer::with_runtime_options(goldendb_backend.clone(), McpScope::default(), false); + let goldendb_ddl = goldendb_server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("goldendb"), + database: None, + sql: "CREATE TABLE a (id INT); INSERT INTO missing VALUES (1)".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + assert!(result_text(&goldendb_ddl).contains("TRANSACTION_WITH_DDL_UNSUPPORTED")); + assert!(goldendb_backend.recorded_arguments.lock().unwrap().iter().all(|(name, _)| name != "execute_batch")); + + let dml = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("my"), + database: None, + sql: "INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + assert!(result_text(&dml).contains("Transaction outcome")); + assert!(backend.recorded_arguments.lock().unwrap().iter().any(|(name, _)| name == "execute_batch")); + + let postgres = connection("pg", "pg", "postgres", "app"); + let postgres_backend = Arc::new(FakeBackend { + connections: vec![postgres], + policy: McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, ..Default::default() }, + ..Default::default() + }); + let postgres_server = DbxMcpServer::with_runtime_options(postgres_backend.clone(), McpScope::default(), false); + let pg_ddl = postgres_server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "CREATE TABLE a (id INT); INSERT INTO t VALUES (1)".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + assert!(result_text(&pg_ddl).contains("Transaction outcome")); + } + + #[tokio::test] + async fn execute_batch_rejects_oracle_ddl_with_transaction() { + // Oracle DDL implicitly commits just like MySQL-family engines, so a DDL + // batch with use_transaction over-promises atomicity. The rejection must + // come from the shared capability check, not a MySQL-only engine list — + // otherwise Oracle would silently run a non-rollbackable "transaction". + let oracle = connection("ora", "ora", "oracle", "app"); + let backend = Arc::new(FakeBackend { connections: vec![oracle], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); + + let ddl = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("ora"), + database: None, + sql: "CREATE TABLE a (id NUMBER); INSERT INTO missing VALUES (1)".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + assert!(result_text(&ddl).contains("TRANSACTION_WITH_DDL_UNSUPPORTED")); + assert!(backend.recorded_arguments.lock().unwrap().iter().all(|(name, _)| name != "execute_batch")); + } + + #[tokio::test] + async fn execute_batch_single_statement_allows_transaction_option_combinations() { + // A single-statement script ignores use_transaction and runs as normal + // auto-commit, so it never enters transaction mode. Combining it with + // session_id or continue_on_error is therefore harmless and must not be + // rejected — the option-validity rejects apply only when the script + // actually enters transaction mode (statement_count > 1). Contrast with + // execute_batch_rejects_transaction_with_session / _continue_on_error, + // which use multi-statement scripts and are still rejected. + let postgres = connection("pg", "pg", "postgres", "app"); + let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false); + + let with_continue = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "INSERT INTO t VALUES (1)".to_string(), + session_id: None, + continue_on_error: Some(true), + use_transaction: Some(true), + })) + .await; + assert!( + !result_text(&with_continue).contains("TRANSACTION_WITH_CONTINUE_UNSUPPORTED"), + "single-statement use_transaction + continue_on_error must not be rejected" + ); + assert!(backend.recorded_arguments.lock().unwrap().iter().any(|(name, _)| name == "execute_batch")); + + // A real session bound to the same connection must likewise reach + // execution rather than be rejected as a transaction/session conflict. + let opened = + server.open_session(Parameters(OpenSessionRequest { selector: selector("pg"), database: None })).await; + let session_id = opened_session_id(&opened); + let with_session = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "INSERT INTO t VALUES (1)".to_string(), + session_id: Some(session_id), + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + assert!( + !result_text(&with_session).contains("TRANSACTION_WITH_SESSION_UNSUPPORTED"), + "single-statement use_transaction + session_id must not be rejected as a transaction/session conflict" + ); + } + + #[tokio::test] + async fn execute_batch_returns_structured_content() { + let postgres = connection("pg", "pg", "postgres", "app"); + let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "SELECT 1; SELECT 2".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: None, + })) + .await; + // The human-readable block stays in content… + assert!(result_text(&result).contains("Statement 1")); + // …and structuredContent carries one object per statement. + let structured = result.structured_content.as_ref().expect("structured content must be populated"); + let statements = structured.as_array().expect("structured content must be an array"); + assert_eq!(statements.len(), 1); + assert_eq!(statements[0]["statement_index"], 0); + } + + #[tokio::test] + async fn execute_batch_enforces_confirmed_write_sql_binding() { + // `confirmed_batch_sql_block_reason` is exercised directly (not through + // the handler with DBX_MCP_CONFIRMED_WRITE_SQL) because the env var is + // process-global and would race with `confirmed_sql_binding_cannot_elevate_central_policy`. + let postgres_db_type = dbx_core::models::connection::DatabaseType::Postgres; + let confirmed = Some("INSERT INTO t VALUES (1)"); + + // An exact match (modulo surrounding whitespace) passes. + assert!(confirmed_batch_sql_block_reason(" INSERT INTO t VALUES (1) ", postgres_db_type, confirmed).is_none()); + // A differing write script fails closed. + let blocked = confirmed_batch_sql_block_reason( + "INSERT INTO t VALUES (1); INSERT INTO u VALUES (2)", + postgres_db_type, + confirmed, + ) + .expect("differing write script must be blocked"); + assert!(blocked.contains("does not match the user-confirmed SQL")); + // A read-only script is never subject to the write confirmation. + assert!(confirmed_batch_sql_block_reason("SELECT 1; SELECT 2", postgres_db_type, confirmed).is_none()); + // No binding never blocks. + assert!(confirmed_batch_sql_block_reason("INSERT INTO t VALUES (1)", postgres_db_type, None).is_none()); + // Unparseable SQL fails closed (treated as a write). + assert!(confirmed_batch_sql_block_reason("NOT VALID SQL ;;;", postgres_db_type, confirmed).is_some()); + } + + #[test] + fn format_batch_results_renders_failed_statement() { + let results = vec![crate::backend::BatchStatementResult { + result: dbx_core::db::QueryResult { + columns: vec![], + column_types: vec![], + column_sortables: vec![], + spatial_columns: vec![], + spatial_values: vec![], + rows: vec![], + affected_rows: 0, + execution_time_ms: 0, + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages: vec![], + }, + execution_error: true, + statement_index: Some(2), + error_message: Some("syntax error near SELECT".to_string()), + merged: false, + }]; + let output = format_batch_results(&results); + assert!(output.contains("### Statement 3")); + assert!(output.contains("**Status:** failed")); + assert!(output.contains("syntax error near SELECT")); + } + + #[test] + fn format_batch_results_renders_transaction_outcome() { + let results = vec![crate::backend::BatchStatementResult { + result: dbx_core::db::QueryResult { + columns: vec![], + column_types: vec![], + column_sortables: vec![], + spatial_columns: vec![], + spatial_values: vec![], + rows: vec![], + affected_rows: 2, + execution_time_ms: 0, + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages: vec![], + }, + execution_error: false, + statement_index: None, + error_message: None, + merged: true, + }]; + let output = format_batch_results(&results); + assert!(output.contains("### Transaction outcome")); + assert!(output.contains("2 row(s) affected")); + assert!(!output.contains("Statement 1")); + } + + #[tokio::test] + async fn execute_batch_transaction_mode_marks_merged_outcome() { + let postgres = connection("pg", "pg", "postgres", "app"); + let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + // The merged transaction outcome is labelled distinctly and does not + // pretend to be per-statement. + assert!(result_text(&result).contains("Transaction outcome")); + assert!(!result_text(&result).contains("Statement 1")); + let structured = result.structured_content.as_ref().expect("structured content must be populated"); + assert_eq!(structured[0]["merged"], true); + assert!(structured[0]["statement_index"].is_null()); + } + + #[tokio::test] + async fn execute_batch_single_statement_with_transaction_is_not_merged() { + // The core transaction wrapper only applies to scripts with more than + // one statement (query.rs:2970). A single statement with + // use_transaction=true runs as a normal auto-commit statement and must + // not be labelled as a transaction outcome. + let postgres = connection("pg", "pg", "postgres", "app"); + let backend = Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("pg"), + database: None, + sql: "INSERT INTO t VALUES (1)".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + assert!(result_text(&result).contains("Statement 1")); + assert!(!result_text(&result).contains("Transaction outcome")); + let structured = result.structured_content.as_ref().expect("structured content must be populated"); + // merged=false is skipped by serde, so the single statement must not + // carry a merged marker (null/absent), unlike the transaction outcome. + assert!(structured[0]["merged"].is_null()); + } + + #[tokio::test] + async fn execute_batch_sqlserver_use_transaction_marks_merged_outcome() { + // SQL Server enters the core transaction path before its normal batch + // fast path, so a successful multi-statement request is merged. + let sqlserver = connection("mssql", "mssql", "sqlserver", "app"); + let backend = Arc::new(FakeBackend { connections: vec![sqlserver], ..Default::default() }); + let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + + let result = server + .execute_batch(Parameters(ExecuteBatchQueryRequest { + selector: selector("mssql"), + database: None, + sql: "INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)".to_string(), + session_id: None, + continue_on_error: None, + use_transaction: Some(true), + })) + .await; + assert!(result_text(&result).contains("Transaction outcome")); + assert!(!result_text(&result).contains("Statement 1")); + let structured = result.structured_content.as_ref().expect("structured content must be populated"); + assert_eq!(structured[0]["merged"], true); + } } diff --git a/crates/dbx-mcp/tests/live_batch.rs b/crates/dbx-mcp/tests/live_batch.rs new file mode 100644 index 0000000000..12e6717ce3 --- /dev/null +++ b/crates/dbx-mcp/tests/live_batch.rs @@ -0,0 +1,331 @@ +//! Live regression for `dbx_execute_batch` transaction semantics against real +//! MySQL and PostgreSQL servers (reviewer evidence for issue #7548). +//! +//! Credentials come ONLY from environment variables — never committed: +//! +//! ```text +//! DBX_LIVE_MYSQL_HOST / DBX_LIVE_MYSQL_PORT / DBX_LIVE_MYSQL_USER +//! DBX_LIVE_MYSQL_PASSWORD / DBX_LIVE_MYSQL_DATABASE +//! DBX_LIVE_PG_HOST / DBX_LIVE_PG_PORT / DBX_LIVE_PG_USER +//! DBX_LIVE_PG_PASSWORD / DBX_LIVE_PG_DATABASE +//! ``` +//! +//! Run with: +//! ```text +//! DBX_LIVE_MYSQL_*=... DBX_LIVE_PG_*=... cargo test -p dbx-mcp --test live_batch -- --ignored --nocapture +//! ``` +//! +//! What this proves: +//! * MySQL-family DDL scripts with `use_transaction=true` are rejected +//! (`TRANSACTION_WITH_DDL_UNSUPPORTED`) before touching the database, because +//! MySQL DDL implicitly commits and cannot be rolled back. Nothing is left +//! behind. +//! * MySQL pure-DML `use_transaction` batches actually run in a transaction: a +//! mid-batch failure rolls the whole batch back (0 rows persist). +//! * PostgreSQL DDL is transactional: a `CREATE TABLE` followed by a failing +//! statement rolls back the DDL too (the table does not exist afterward), +//! which is exactly the contrast that motivated the MySQL guard. + +use std::sync::Arc; + +use dbx_core::{models::connection::ConnectionConfig, storage::McpGlobalPolicy, storage::Storage}; +use dbx_mcp::{DbxMcpServer, LocalBackend, McpScope}; +use rmcp::{model::CallToolRequestParams, ServiceExt}; +use serde_json::{json, Map}; + +fn required_env(name: &str) -> String { + std::env::var(name).unwrap_or_else(|_| panic!("set {name} to run the live batch tests")) +} + +/// Build the MySQL connection config and database name from `DBX_LIVE_MYSQL_*`. +/// Each engine is independent so a test can run with only its own env set. +fn mysql_env() -> (ConnectionConfig, String) { + build_env("DBX_LIVE_MYSQL", "live-mysql", "mysql") +} + +/// Build the PostgreSQL connection config and database name from `DBX_LIVE_PG_*`. +fn postgres_env() -> (ConnectionConfig, String) { + build_env("DBX_LIVE_PG", "live-postgres", "postgres") +} + +fn build_env(prefix: &str, id: &str, db_type: &str) -> (ConnectionConfig, String) { + let config = serde_json::from_value(json!({ + "id": id, + "name": id, + "db_type": db_type, + "host": required_env(&format!("{prefix}_HOST")), + "port": required_env(&format!("{prefix}_PORT")).parse::().expect("numeric port"), + "username": required_env(&format!("{prefix}_USER")), + "password": required_env(&format!("{prefix}_PASSWORD")), + "database": required_env(&format!("{prefix}_DATABASE")), + "ssl": false, + })) + .expect("build live connection config"); + (config, required_env(&format!("{prefix}_DATABASE"))) +} + +fn full_access_policy() -> McpGlobalPolicy { + McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: true, + allowed_connection_ids: None, + allowed_tool_names: None, + connection_policies: vec![], + query_timeout_secs: None, + } +} + +/// Like `full_access_policy()` but with a 1s global MCP query timeout, used to +/// prove that `use_transaction` batches honor the MCP query-timeout policy. +fn short_timeout_policy() -> McpGlobalPolicy { + McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: true, + allowed_connection_ids: None, + allowed_tool_names: None, + connection_policies: vec![], + query_timeout_secs: Some(1), + } +} + +/// Boot a real `DbxMcpServer` over a live LocalBackend seeded from a temp +/// storage db. Returns `(client, server_task, temp_dir)`; `temp_dir` must stay +/// alive for the test so the sqlite file is not deleted mid-run. +macro_rules! boot_live_server { + ($connections:expr, $policy:expr) => {{ + let dir = tempfile::tempdir().expect("temp dir"); + let storage = Storage::open(&dir.path().join("dbx.db")).await.expect("open storage"); + for config in $connections { + storage.add_connection_for_mcp(config.clone()).await.expect("add connection"); + } + storage.save_mcp_global_policy(&$policy).await.expect("save policy"); + drop(storage); + let backend = Arc::new(LocalBackend::open(&dir.path().join("dbx.db")).await.expect("open LocalBackend")); + let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + let server_task = tokio::spawn(async move { server.serve(server_transport).await }); + let client = ().serve(client_transport).await.expect("initialize MCP client"); + (client, server_task, dir) + }}; +} + +/// Unique table names so repeated runs (and the two engines) never collide. +/// The tests intentionally leave tables behind on the live server, so a UUID +/// (rather than a per-process counter) also protects against PID reuse after a +/// crashed earlier run left a same-named table behind. +fn unique_table(prefix: &str) -> String { + format!("mcp_live_{prefix}_{}", uuid::Uuid::new_v4().simple()) +} + +fn arguments(object: serde_json::Value) -> Map { + object.as_object().cloned().unwrap_or_else(Map::new) +} + +/// Call `dbx_execute_query` on the live server; returns `(is_error, text)`. +macro_rules! query { + ($client:expr, $connection_id:expr, $database:expr, $sql:expr) => {{ + let result = $client + .peer() + .call_tool( + CallToolRequestParams::new("dbx_execute_query").with_arguments(arguments(json!({ + "connection_id": $connection_id, + "database": $database, + "sql": $sql, + }))), + ) + .await + .expect("call dbx_execute_query"); + let text = result + .content + .iter() + .filter_map(|block| block.as_text().map(|t| t.text.clone())) + .collect::>() + .join("\n"); + (result.is_error == Some(true), text) + }}; +} + +/// Call `dbx_execute_batch` on the live server; returns `(is_error, text)`. +macro_rules! batch { + ($client:expr, $connection_id:expr, $database:expr, $sql:expr, $use_transaction:expr) => {{ + let result = $client + .peer() + .call_tool( + CallToolRequestParams::new("dbx_execute_batch").with_arguments(arguments(json!({ + "connection_id": $connection_id, + "database": $database, + "sql": $sql, + "use_transaction": $use_transaction, + }))), + ) + .await + .expect("call dbx_execute_batch"); + let text = result + .content + .iter() + .filter_map(|block| block.as_text().map(|t| t.text.clone())) + .collect::>() + .join("\n"); + (result.is_error == Some(true), text) + }}; +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MYSQL_* and DBX_LIVE_PG_* env pointing at reachable MySQL 5.7 / PostgreSQL"] +async fn mysql_batch_rejects_ddl_transaction_and_rolls_back_failed_dml() { + let (mysql, mysql_database) = mysql_env(); + let (client, server_task, _dir) = boot_live_server!([mysql], full_access_policy()); + let db = mysql_database.as_str(); + + // 1) DDL + use_transaction is rejected BEFORE execution, and no partial DDL + // is left behind (MySQL DDL implicitly commits; a table must not exist). + let ddl_table = unique_table("mysql_ddl"); + let (is_error, text) = batch!( + &client, + "live-mysql", + db, + format!("CREATE TABLE {ddl_table} (id INT); INSERT INTO {ddl_table} VALUES (1)"), + true + ); + assert!(is_error, "DDL+transaction must be rejected, got: {text}"); + assert!( + text.contains("TRANSACTION_WITH_DDL_UNSUPPORTED"), + "expected TRANSACTION_WITH_DDL_UNSUPPORTED, got: {text}" + ); + let (select_error, select_text) = query!(&client, "live-mysql", db, format!("SELECT * FROM {ddl_table}")); + assert!(select_error, "DDL table must not exist after the rejected batch, got: {select_text}"); + + // 2) Pure-DML transaction batch succeeds and returns a merged outcome. + let ok_table = unique_table("mysql_dml_ok"); + let (create_error, create_text) = query!(&client, "live-mysql", db, format!("CREATE TABLE {ok_table} (n INT)")); + assert!(!create_error, "setup CREATE TABLE failed: {create_text}"); + let (is_error, text) = batch!( + &client, + "live-mysql", + db, + format!("INSERT INTO {ok_table} (n) VALUES (1); INSERT INTO {ok_table} (n) VALUES (2)"), + true + ); + assert!(!is_error, "DML transaction batch failed: {text}"); + assert!(text.contains("Transaction outcome"), "expected a merged transaction outcome, got: {text}"); + let (_, count_text) = query!(&client, "live-mysql", db, format!("SELECT COUNT(*) AS c FROM {ok_table}")); + assert!(count_text.contains("c") && count_text.contains("2"), "expected 2 persisted rows, got: {count_text}"); + + // 3) A mid-batch DML failure rolls the whole transaction back (0 rows). + let fail_table = unique_table("mysql_dml_fail"); + let (create_error, create_text) = query!(&client, "live-mysql", db, format!("CREATE TABLE {fail_table} (n INT)")); + assert!(!create_error, "setup CREATE TABLE failed: {create_text}"); + let (is_error, text) = batch!( + &client, + "live-mysql", + db, + format!( + "INSERT INTO {fail_table} (n) VALUES (1); INSERT INTO {fail_table} (n) VALUES (2); INSERT INTO {fail_table}_missing (n) VALUES (3)" + ), + true + ); + assert!(is_error, "expected the failing DML transaction batch to error, got: {text}"); + assert!(text.contains("DBX_BATCH_EXECUTION_ERROR"), "expected execution error, got: {text}"); + let (_, count_text) = query!(&client, "live-mysql", db, format!("SELECT COUNT(*) AS c FROM {fail_table}")); + assert!( + count_text.contains("c") && count_text.contains("0"), + "expected the whole batch to roll back (0 rows), got: {count_text}" + ); + + client.cancel().await.expect("close MCP client"); + server_task.abort(); +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MYSQL_* and DBX_LIVE_PG_* env pointing at reachable MySQL 5.7 / PostgreSQL"] +async fn postgres_transaction_batch_respects_mcp_query_timeout() { + let (postgres, postgres_database) = postgres_env(); + let (client, server_task, _dir) = boot_live_server!([postgres], short_timeout_policy()); + let db = postgres_database.as_str(); + + // A use_transaction batch must honor the MCP global query-timeout policy: + // pg_sleep(3) exceeds the 1s policy, so the whole transactional batch errors + // instead of returning a successful merged outcome. Regression for the P1 + // where the transaction path ignored options.timeout_secs and ran unbounded. + let (is_error, text) = batch!(&client, "live-postgres", db, "SELECT pg_sleep(3); SELECT 1", true); + assert!(is_error, "expected the PostgreSQL transaction batch to time out, got: {text}"); + assert!( + text.contains("DBX_BATCH_EXECUTION_ERROR") || text.contains("Query timed out"), + "expected a timeout surfaced as an execution error, got: {text}" + ); + assert!(!text.contains("Transaction outcome"), "the batch must NOT be a successful merged outcome, got: {text}"); + + client.cancel().await.expect("close MCP client"); + server_task.abort(); +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MYSQL_* and DBX_LIVE_PG_* env pointing at reachable MySQL 5.7 / PostgreSQL"] +async fn mysql_transaction_batch_respects_mcp_query_timeout() { + let (mysql, mysql_database) = mysql_env(); + let (client, server_task, _dir) = boot_live_server!([mysql], short_timeout_policy()); + let db = mysql_database.as_str(); + + // Same shape as the PostgreSQL timeout test: SLEEP(3) exceeds the 1s MCP + // policy, so the transactional batch errors instead of a merged outcome. + // (Whether the driver also aborts the server-side SLEEP is an observation; + // the requirement here is that the error surfaces to the caller.) + let (is_error, text) = batch!(&client, "live-mysql", db, "SELECT SLEEP(3); SELECT 1", true); + assert!(is_error, "expected the MySQL transaction batch to time out, got: {text}"); + assert!( + text.contains("DBX_BATCH_EXECUTION_ERROR") || text.contains("Query timed out"), + "expected a timeout surfaced as an execution error, got: {text}" + ); + assert!(!text.contains("Transaction outcome"), "the batch must NOT be a successful merged outcome, got: {text}"); + + client.cancel().await.expect("close MCP client"); + server_task.abort(); +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MYSQL_* and DBX_LIVE_PG_* env pointing at reachable MySQL 5.7 / PostgreSQL"] +async fn postgres_batch_runs_ddl_inside_a_rollbackable_transaction() { + let (postgres, postgres_database) = postgres_env(); + let (client, server_task, _dir) = boot_live_server!([postgres], full_access_policy()); + let db = postgres_database.as_str(); + + // PostgreSQL DDL is transactional: a CREATE TABLE + INSERT in one + // use_transaction batch succeeds and the table persists. + let ok_table = unique_table("pg_ddl_ok"); + let (is_error, text) = batch!( + &client, + "live-postgres", + db, + format!("CREATE TABLE {ok_table} (id INT); INSERT INTO {ok_table} VALUES (1)"), + true + ); + assert!(!is_error, "PostgreSQL DDL transaction batch failed: {text}"); + assert!(text.contains("Transaction outcome"), "expected a merged transaction outcome, got: {text}"); + let (_, count_text) = query!(&client, "live-postgres", db, format!("SELECT COUNT(*) AS c FROM {ok_table}")); + assert!( + count_text.contains("c") && count_text.contains("1"), + "expected the CREATE TABLE to persist (1 row), got: {count_text}" + ); + + // PostgreSQL DDL rolls back: CREATE TABLE + failing statement leaves no + // table behind — the exact contrast to MySQL's implicit-commit DDL. + let fail_table = unique_table("pg_ddl_fail"); + let missing = unique_table("pg_missing"); + let (is_error, text) = batch!( + &client, + "live-postgres", + db, + format!("CREATE TABLE {fail_table} (id INT); INSERT INTO {missing} VALUES (1)"), + true + ); + assert!(is_error, "expected the failing PostgreSQL transaction batch to error, got: {text}"); + assert!(text.contains("DBX_BATCH_EXECUTION_ERROR"), "expected execution error, got: {text}"); + let (select_error, select_text) = query!(&client, "live-postgres", db, format!("SELECT * FROM {fail_table}")); + assert!( + select_error, + "PostgreSQL CREATE TABLE must roll back with the failed batch, table still exists: {select_text}" + ); + + client.cancel().await.expect("close MCP client"); + server_task.abort(); +} diff --git a/crates/dbx-mcp/tests/protocol.rs b/crates/dbx-mcp/tests/protocol.rs index 0226022094..3ca7b550e0 100644 --- a/crates/dbx-mcp/tests/protocol.rs +++ b/crates/dbx-mcp/tests/protocol.rs @@ -285,14 +285,17 @@ 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(), 17); + assert_eq!(names.len(), 18); #[cfg(not(feature = "mq-admin"))] - assert_eq!(names.len(), 16); + assert_eq!(names.len(), 17); assert!(names.contains(&"dbx_list_connections")); assert!(names.contains(&"dbx_list_databases")); assert!(names.contains(&"dbx_duplicate_connection")); assert!(names.contains(&"dbx_execute_redis_command")); assert!(names.contains(&"dbx_execute_and_show")); + assert!(names.contains(&"dbx_execute_batch")); + assert!(names.contains(&"dbx_list_routines")); + assert!(names.contains(&"dbx_get_routine_source")); assert!(names.contains(&"dbx_open_session")); assert!(names.contains(&"dbx_close_session")); #[cfg(feature = "mq-admin")] diff --git a/docs/content/docs/mcp.cn.mdx b/docs/content/docs/mcp.cn.mdx index b0ace67460..ea70d84991 100644 --- a/docs/content/docs/mcp.cn.mdx +++ b/docs/content/docs/mcp.cn.mdx @@ -87,7 +87,7 @@ DeepSeek Harness 通过 Cordis 插件条目加载 MCP Server，不读取 `mcpSer ## 工具列表 -DBX MCP 当前提供 17 个工具： +DBX MCP 当前提供 18 个工具： | 工具 | 说明 | | --- | --- | @@ -102,6 +102,7 @@ DBX MCP 当前提供 17 个工具： | `dbx_get_routine_source` | 按名称返回存储过程或函数的源码，重载场景可用可选的 `signature` 区分 | | `dbx_get_schema_context` | 返回适合 AI 使用的紧凑 Schema 上下文 | | `dbx_execute_query` | 执行 SQL 或支持的 MongoDB shell 命令，最多返回 100 行 | +| `dbx_execute_batch` | 一次执行包含多条语句的 SQL 脚本，按语句返回结果（多语句脚本搭配 `use_transaction` 时返回单个合并结果） | | `dbx_open_session` | 为 SQL 连接打开固定后端连接的有状态查询会话 | | `dbx_close_session` | 关闭会话并释放固定连接资源 | | `dbx_execute_redis_command` | 执行 Redis 命令 | @@ -111,9 +112,20 @@ DBX MCP 当前提供 17 个工具： 启用连接作用域后，修改连接和桌面 UI 工具会被隐藏。 +## 批量 SQL 执行 + +`dbx_execute_batch` 在单次调用内执行包含多条语句的 SQL 脚本。默认（自动提交）模式下按语句返回结果，因此调用方可以清楚看到哪条失败、每条影响多少行、各自返回什么；使用 `use_transaction` 且脚本包含多条语句时，只返回一个合并的事务结果。 + +- 脚本通过方言感知的解析器拆分成独立语句，字符串、注释、存储过程中出现分号也不会破坏批次。 +- 默认遇到第一条失败语句即停止；传入 `continue_on_error` 可让批内在单条语句出错后继续执行（连接级错误仍会终止）。 +- 传入 `use_transaction` 可把多条语句放入同一个 `BEGIN … COMMIT`，使整批要么全部成功要么回滚；这要求脚本包含多条语句（单语句脚本会忽略该选项，按普通自动提交执行）。此模式下调用只返回一个合并结果（而非逐语句结果），且 `use_transaction` 不能与 `session_id` 同时使用——事务批处理运行在池连接上，会丢弃会话状态；也不能与 `continue_on_error` 同时使用——事务批处理在第一条失败时即回滚并停止。当当前后端无法提供可回滚事务时，DBX 会拒绝该选项。MySQL 系连接对含 DDL 的脚本会拒绝 `use_transaction`，因为 MySQL 的 DDL 会隐式提交、无法回滚。 +- SQL 策略（只读、危险 SQL、生产保护、切换数据库）作为整体在匹配前重新检查，与 `dbx_execute_query` 一致；只要批内含一条写入或 DDL，在只读策略下整次调用都会被拒绝。 +- 每个批次在其完整执行期间都会使用独占连接，因此临时表和 `SET` 状态可在批内语句之间保留。若要跨多次 MCP 调用保留这些状态，先用 `dbx_open_session` 打开会话，再传入其 `session_id`。 +- Redis 和 MongoDB 连接不支持 `dbx_execute_batch`。 + ## 有状态查询会话 -普通 `dbx_execute_query` 调用彼此独立。需要保持数据库 Session 状态时，先调用 `dbx_open_session`，再把返回的 `sessionId` 传给后续 `dbx_execute_query`： +普通 `dbx_execute_query` 调用彼此独立。需要保持数据库 Session 状态时，先调用 `dbx_open_session`，再把返回的 `sessionId` 传给后续 `dbx_execute_query` 或 `dbx_execute_batch`： - `USE` 或数据库上下文切换 - 临时表 diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index 0e3fed0c1d..ad8ec8213d 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -87,7 +87,7 @@ Run `dsh web --dump-config` to verify the composed configuration, then restart D ## Tools -DBX MCP currently provides 17 tools: +DBX MCP currently provides 18 tools: | Tool | Description | | --- | --- | @@ -102,6 +102,7 @@ DBX MCP currently provides 17 tools: | `dbx_get_routine_source` | Return the source of a stored procedure or function by name, with an optional `signature` for overloaded names | | `dbx_get_schema_context` | Return compact schema context for an AI model | | `dbx_execute_query` | Execute SQL or a supported MongoDB shell command, returning at most 100 rows | +| `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_close_session` | Close a session and release its pinned connection resources | | `dbx_execute_redis_command` | Execute a Redis command | @@ -111,9 +112,20 @@ DBX MCP currently provides 17 tools: Connection-scoped sessions hide connection-mutating and desktop UI tools. +## Batch SQL Execution + +`dbx_execute_batch` runs a SQL script containing multiple statements in a single call. In auto-commit mode (the default) it returns one result per statement, so callers see exactly which statement failed, how many rows each affected, and what each returned. With `use_transaction` and a script containing multiple statements it returns a single merged transaction outcome instead. + +- The script text is split with a database-dialect-aware parser, so semicolons inside string literals, comments, and stored procedures do not break the batch. +- By default the batch stops at the first failing statement. Pass `continue_on_error` to keep going after a per-statement error (connection-level failures always stop the batch). +- Pass `use_transaction` to run all statements inside one `BEGIN … COMMIT` so the whole batch either succeeds or rolls back; this requires a script with more than one statement (a single-statement script ignores the option and runs as normal auto-commit). In transaction mode the call returns a single merged result (not one result per statement), and `use_transaction` cannot be combined with `session_id` (a transactional batch runs on a pooled connection and would discard the session's state) or with `continue_on_error` (a transactional batch stops and rolls back at the first failure). DBX rejects the option when the active backend cannot provide a rollbackable transaction. MySQL-family connections reject `use_transaction` for DDL scripts because MySQL DDL implicitly commits and cannot be rolled back. +- SQL policy (read-only, dangerous-SQL, production protection, database switching) is rechecked against the whole script as a unit before anything runs, exactly as `dbx_execute_query` does. A single write or DDL statement in the batch fails the whole call under a read-only policy. +- Each batch gets a dedicated connection for its full lifetime, so temporary tables and `SET` state are preserved between its statements. To preserve that state across multiple MCP calls, open a session with `dbx_open_session` and pass its `session_id`. +- `dbx_execute_batch` is not available for Redis or MongoDB connections. + ## Stateful Query Sessions -Regular `dbx_execute_query` calls are independent. When a workflow must preserve database session state, call `dbx_open_session` first and pass its returned `sessionId` to later `dbx_execute_query` calls. Typical uses include: +Regular `dbx_execute_query` calls are independent. When a workflow must preserve database session state, call `dbx_open_session` first and pass its returned `sessionId` to later `dbx_execute_query` or `dbx_execute_batch` calls. Typical uses include: - `USE` or database-context changes - Temporary tables diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md index fb002a1213..9148471f6b 100644 --- a/packages/mcp-server/README.md +++ b/packages/mcp-server/README.md @@ -17,7 +17,7 @@ The MCP protocol, connection loading, SQL safety, schema access, Redis support, ## Features -- **17 MCP tools** for connection and database discovery, schemas, SQL, Redis, sessions, messages, and DBX UI integration +- **18 MCP tools** for connection and database discovery, schemas, SQL, Redis, sessions, messages, and DBX UI integration - **Precompiled native binaries** with no local Rust, Cargo, Python, or C/C++ build requirement - **No `better-sqlite3` runtime dependency** and no Node native-addon ABI coupling - **Local, Web, and Docker modes** using the same tool interface @@ -153,6 +153,7 @@ Ask the MCP client to: | `dbx_get_routine_source` | Return the source of a stored procedure or function by name, with an optional `signature` for overloaded names | | `dbx_get_schema_context` | Return compact schema context suitable for an AI model | | `dbx_execute_query` | Execute SQL or a supported MongoDB shell command, returning at most 100 rows | +| `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_close_session` | Close a session and release its pinned connection resources | | `dbx_execute_redis_command` | Execute a Redis command | @@ -435,7 +436,7 @@ MCP 协议、连接读取、SQL 安全检查、Schema、Redis、MongoDB、Web ### 主要能力 -- 17 个 MCP 工具，涵盖连接和数据库发现、Schema、SQL、Redis、会话、消息队列和 DBX 桌面集成 +- 18 个 MCP 工具，涵盖连接和数据库发现、Schema、SQL、Redis、会话、消息队列和 DBX 桌面集成 - 不依赖 `better-sqlite3`，没有 Node 原生模块 ABI 问题 - 支持本地 DBX、DBX Web 和 Docker - 可选 Streamable HTTP 传输，使用 Bearer Token 保护；stdio 仍为默认方式 @@ -526,6 +527,7 @@ MCP 配置： | `dbx_describe_table` | 获取字段和表结构 | | `dbx_get_schema_context` | 获取适合 AI 使用的紧凑 Schema 上下文 | | `dbx_execute_query` | 执行 SQL 或支持的 MongoDB Shell 命令，最多返回 100 行 | +| `dbx_execute_batch` | 一次执行包含多条语句的 SQL 脚本，按语句返回结果（多语句脚本搭配 `use_transaction` 时返回单个合并结果） | | `dbx_open_session` | 为 SQL 连接打开固定后端连接的有状态查询会话 | | `dbx_close_session` | 关闭会话并释放固定连接资源 | | `dbx_execute_redis_command` | 执行 Redis 命令 |