diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index cf5d1e9afc..f2229e0f3a 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -1222,6 +1222,21 @@ pub async fn execute_batch_with_max_rows( max_rows: Option, ) -> Result, String> { let start = Instant::now(); + if sqlserver_batch_can_use_execute(sql) { + let result = sqlserver_driver_result(client.execute(sql, &[])).await?; + return Ok(vec![QueryResult { + columns: vec![], + column_types: Vec::new(), + column_sortables: vec![], + rows: vec![], + affected_rows: result.rows_affected().iter().sum::(), + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + session_id: None, + has_more: false, + }]); + } + if is_single_sqlserver_select(sql) { if let Ok(Some(query_sql)) = spatial_safe_sqlserver_query(client, sql).await { let stream = sqlserver_driver_result(client.query(query_sql.as_str(), &[])).await?; @@ -1250,6 +1265,17 @@ pub async fn execute_batch_with_max_rows( Ok(results) } +fn sqlserver_batch_can_use_execute(sql: &str) -> bool { + !requires_simple_query_batch(sql) + && !starts_with_executable_sql_keyword(sql, &["SELECT", "EXEC", "WITH", "TABLE"]) + && !sqlserver_dml_output_returns_rows(sql) +} + +fn sqlserver_dml_output_returns_rows(sql: &str) -> bool { + starts_with_executable_sql_keyword(sql, &["INSERT", "UPDATE", "DELETE", "MERGE"]) + && 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() { @@ -1328,7 +1354,8 @@ fn first_sql_tokens(sql: &str, limit: usize) -> Vec { mod tests { use super::{ build_spatial_safe_sqlserver_query, is_sqlserver_spatial_column, requires_simple_query_batch, - sqlserver_cell_to_json, sqlserver_columns_sql, sqlserver_indexes_sql, sqlserver_list_objects_sql, + sqlserver_batch_can_use_execute, sqlserver_cell_to_json, sqlserver_columns_sql, + sqlserver_dml_output_returns_rows, sqlserver_indexes_sql, sqlserver_list_objects_sql, sqlserver_table_comment_sql, SqlServerDescribedColumn, SqlServerResultSet, }; use chrono::NaiveDate; @@ -1384,6 +1411,25 @@ mod tests { assert!(!requires_simple_query_batch("UPDATE dbo.t SET id = 1;")); } + #[test] + fn sqlserver_cud_batches_use_execute_for_affected_rows() { + assert!(sqlserver_batch_can_use_execute("UPDATE dbo.users SET active = 0 WHERE id = 1;")); + assert!(sqlserver_batch_can_use_execute("INSERT INTO dbo.users(id) VALUES (1);")); + assert!(sqlserver_batch_can_use_execute("DELETE FROM dbo.users WHERE id = 1;")); + assert!(sqlserver_batch_can_use_execute( + "MERGE dbo.t AS t USING dbo.s AS s ON t.id = s.id WHEN MATCHED THEN UPDATE SET name = s.name;" + )); + } + + #[test] + fn sqlserver_result_returning_batches_keep_simple_query_path() { + assert!(!sqlserver_batch_can_use_execute("SELECT * FROM dbo.users;")); + assert!(!sqlserver_batch_can_use_execute("EXEC dbo.list_users;")); + assert!(!sqlserver_batch_can_use_execute("WITH cte AS (SELECT 1 AS id) SELECT * FROM cte;")); + assert!(!sqlserver_batch_can_use_execute("UPDATE dbo.users SET active = 0 OUTPUT inserted.id WHERE id = 1;")); + assert!(sqlserver_dml_output_returns_rows("DELETE FROM dbo.users OUTPUT deleted.id WHERE id = 1;")); + } + #[test] fn sqlserver_user_query_paths_do_not_collect_full_results_before_limiting() { let source = include_str!("sqlserver.rs");