From 8c6a4d60b295533c6abcd8dfd5b26f8e4ddeb958 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 14 Jul 2026 20:38:08 +0800 Subject: [PATCH] fix(sqlserver): preserve explicit transaction batches --- crates/dbx-core/src/db/sqlserver.rs | 57 ++++++++++++++----- .../tests/live_sqlserver_completion.rs | 43 ++++++++++++++ 2 files changed, 86 insertions(+), 14 deletions(-) diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index d5157e5d30..823f69e36e 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -1664,7 +1664,7 @@ pub async fn execute_query_with_max_rows( let mut result = sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows)).await?; strip_dbx_sqlserver_row_number_column(&mut result, sql); Ok(result) - } else if requires_simple_query_batch(sql) || is_transaction_control(sql) { + } else if requires_simple_query_batch(sql) || contains_transaction_control(sql) { let stream = sqlserver_driver_result(client.simple_query(sql)).await?; let _ = sqlserver_driver_result(collect_result_sets_limited(stream, start, max_rows)).await?; Ok(QueryResult { @@ -1784,6 +1784,10 @@ fn is_dbx_sqlserver_row_number_page_sql(sql: &str) -> bool { fn sqlserver_batch_can_use_execute(sql: &str) -> bool { !requires_simple_query_batch(sql) + // Tiberius executes this path through an RPC. SQL Server rejects an RPC + // that changes @@TRANCOUNT with error 266, while a regular batch is allowed + // to leave an explicit transaction open for a later COMMIT or ROLLBACK. + && !contains_transaction_control(sql) && !sqlserver_batch_may_return_result_set(sql) && !sqlserver_dml_output_returns_rows(sql) } @@ -1798,19 +1802,26 @@ fn sqlserver_dml_output_returns_rows(sql: &str) -> bool { && first_sql_tokens(sql, 64).iter().any(|token| token.eq_ignore_ascii_case("OUTPUT")) } -fn is_transaction_control(sql: &str) -> bool { - let tokens = first_sql_tokens(sql, 2); - if tokens.is_empty() { - return false; - } - let first = &tokens[0]; - if first.eq_ignore_ascii_case("COMMIT") || first.eq_ignore_ascii_case("ROLLBACK") { - return true; - } - if first.eq_ignore_ascii_case("BEGIN") { - return tokens.get(1).is_some_and(|t| t.eq_ignore_ascii_case("TRANSACTION") || t.eq_ignore_ascii_case("TRAN")); - } - false +fn contains_transaction_control(sql: &str) -> bool { + let tokens = top_level_sqlserver_tokens(sql); + tokens.iter().enumerate().any(|(index, token)| { + if matches!(token.text.as_str(), "COMMIT" | "ROLLBACK") { + return true; + } + if token.text == "BEGIN" { + if tokens.get(index + 1).is_some_and(|next| matches!(next.text.as_str(), "TRANSACTION" | "TRAN")) { + return true; + } + if tokens.get(index + 1).is_some_and(|next| next.text == "DISTRIBUTED") + && tokens.get(index + 2).is_some_and(|next| matches!(next.text.as_str(), "TRANSACTION" | "TRAN")) + { + return true; + } + } + token.text == "SET" + && tokens.get(index + 1).is_some_and(|next| next.text == "IMPLICIT_TRANSACTIONS") + && tokens.get(index + 2).is_some_and(|next| next.text == "ON") + }) } fn requires_simple_query_batch(sql: &str) -> bool { @@ -1995,6 +2006,24 @@ mod tests { )); } + #[test] + fn sqlserver_transaction_batches_keep_simple_query_path() { + assert!(!sqlserver_batch_can_use_execute("BEGIN TRANSACTION\nUPDATE dbo.users SET active = 0 WHERE id = 1;")); + assert!(!sqlserver_batch_can_use_execute("UPDATE dbo.users SET active = 0 WHERE id = 1;\nCOMMIT;")); + assert!(!sqlserver_batch_can_use_execute("ROLLBACK TRANSACTION;")); + assert!(!sqlserver_batch_can_use_execute( + "BEGIN DISTRIBUTED TRANSACTION\nUPDATE dbo.users SET active = 0 WHERE id = 1;" + )); + assert!(!sqlserver_batch_can_use_execute( + "SET IMPLICIT_TRANSACTIONS ON\nUPDATE dbo.users SET active = 0 WHERE id = 1;" + )); + assert!(sqlserver_batch_can_use_execute("BEGIN TRY\nUPDATE dbo.users SET active = 0 WHERE id = 1;\nEND TRY")); + assert!(sqlserver_batch_can_use_execute( + "UPDATE dbo.users SET note = 'BEGIN TRANSACTION; ROLLBACK' WHERE id = 1;" + )); + assert!(sqlserver_batch_can_use_execute("UPDATE dbo.users SET [rollback] = 1 WHERE id = 1;")); + } + #[test] fn sqlserver_result_returning_batches_keep_simple_query_path() { assert!(!sqlserver_batch_can_use_execute("SELECT * FROM dbo.users;")); diff --git a/crates/dbx-core/tests/live_sqlserver_completion.rs b/crates/dbx-core/tests/live_sqlserver_completion.rs index 7f8333f282..5d910e7518 100644 --- a/crates/dbx-core/tests/live_sqlserver_completion.rs +++ b/crates/dbx-core/tests/live_sqlserver_completion.rs @@ -95,6 +95,49 @@ async fn live_sqlserver_execute_query_creates_schema() { assert!(schemas.expect("list schemas").contains(&schema)); } +#[tokio::test] +#[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD pointing at a writable SQL Server database"] +async fn live_sqlserver_explicit_transaction_batch_can_rollback() { + let database = std::env::var("DBX_LIVE_SQLSERVER_DATABASE").unwrap_or_else(|_| "tempdb".to_string()); + let host = std::env::var("DBX_LIVE_SQLSERVER_HOST").unwrap_or_else(|_| "127.0.0.1".to_string()); + let port = std::env::var("DBX_LIVE_SQLSERVER_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(1433); + let user = std::env::var("DBX_LIVE_SQLSERVER_USER").unwrap_or_else(|_| "sa".to_string()); + let password = std::env::var("DBX_LIVE_SQLSERVER_PASSWORD").expect("DBX_LIVE_SQLSERVER_PASSWORD"); + let mut client = + dbx_core::db::sqlserver::connect(&host, port, &user, &password, Some(&database), None, Duration::from_secs(10)) + .await + .expect("connect SQL Server"); + + let table = format!("dbx_txn_{}", uuid::Uuid::new_v4().simple()); + dbx_core::db::sqlserver::execute_query( + &mut client, + &format!("CREATE TABLE dbo.[{table}] (id INT PRIMARY KEY, name NVARCHAR(50)); INSERT INTO dbo.[{table}] VALUES (50, N'old');"), + ) + .await + .expect("create transaction test table"); + + let execution = dbx_core::db::sqlserver::execute_batch( + &mut client, + &format!("BEGIN TRANSACTION\nUPDATE dbo.[{table}] SET name = N'changed' WHERE id = 50;"), + ) + .await; + let transaction_count = + dbx_core::db::sqlserver::execute_query(&mut client, "SELECT @@TRANCOUNT AS transaction_count").await; + let rollback = dbx_core::db::sqlserver::execute_batch(&mut client, "ROLLBACK TRANSACTION").await; + let verify = + dbx_core::db::sqlserver::execute_query(&mut client, &format!("SELECT name FROM dbo.[{table}] WHERE id = 50")) + .await; + let cleanup = dbx_core::db::sqlserver::execute_query(&mut client, &format!("DROP TABLE dbo.[{table}]")).await; + + execution.expect("execute explicit transaction batch without error 266"); + let transaction_count = transaction_count.expect("read transaction count"); + assert_eq!(transaction_count.rows, vec![vec![serde_json::json!(1)]]); + rollback.expect("rollback explicit transaction"); + let verify = verify.expect("verify rolled-back value"); + assert_eq!(verify.rows, vec![vec![serde_json::json!("old")]]); + cleanup.expect("drop transaction test table"); +} + #[tokio::test] #[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD pointing at a writable SQL Server database"] async fn live_sqlserver_table_structure_default_changes_drop_existing_constraints() {