fix(sqlserver): report affected rows for CUD batches

crates/dbx-core/src/db/sqlserver.rs

@@ -1222,6 +1222,21 @@ pub async fn execute_batch_with_max_rows(
     max_rows: Option<usize>,
 ) -> Result<Vec<QueryResult>, 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::<u64>(),
+            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<String> {
 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");