fix(sqlserver): show PRINT messages

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

@@ -374,34 +374,59 @@ fn column_types_from_metadata(metadata: &tiberius::ResultMetadata) -> Vec<String
 
 const SQLSERVER_MESSAGE_COLUMN: &str = "Message";
 // Tiberius 0.12.3 emits both user-visible INFO tokens and internal ENVCHANGE
-// tokens at the same tracing target/level. Keep only the TokenInfo callsite;
-// update this guard together with the pinned driver when Tiberius is upgraded.
+// tokens at the same tracing target/level. DONE tokens are also consumed before
+// QueryStream yields its next public item. Keep these callsite guards in sync
+// with the pinned driver when Tiberius is upgraded.
 const TIBERIUS_INFO_TOKEN_EVENT_LINE: u32 = 194;
+const TIBERIUS_DONE_TOKEN_EVENT_LINE: u32 = 153;
+const TIBERIUS_DONE_PROC_TOKEN_EVENT_LINE: u32 = 159;
+const TIBERIUS_DONE_IN_PROC_TOKEN_EVENT_LINE: u32 = 165;
 
-#[derive(Clone, Default)]
-struct SqlServerMessageLayer {
-    messages: StdArc<StdMutex<Vec<String>>>,
+#[derive(Debug, Clone)]
+pub(crate) struct SqlServerBatchResult {
+    pub result: QueryResult,
+    pub server_message: bool,
 }
 
-impl<S> Layer<S> for SqlServerMessageLayer
+#[derive(Debug, Clone, PartialEq, Eq)]
+enum SqlServerTdsEvent {
+    Message(String),
+    Done(Option<u64>),
+}
+
+#[derive(Clone, Default)]
+struct SqlServerTdsEventLayer {
+    events: StdArc<StdMutex<Vec<SqlServerTdsEvent>>>,
+}
+
+impl SqlServerTdsEventLayer {
+    fn take_events(&self) -> Vec<SqlServerTdsEvent> {
+        self.events.lock().map(|mut events| std::mem::take(&mut *events)).unwrap_or_default()
+    }
+}
+
+impl<S> Layer<S> for SqlServerTdsEventLayer
 where
     S: Subscriber,
 {
     fn on_event(&self, event: &Event<'_>, _ctx: LayerContext<'_, S>) {
         let metadata = event.metadata();
-        if metadata.level() != &Level::INFO
-            || metadata.target() != "tiberius::tds::stream::token"
-            || metadata.line() != Some(TIBERIUS_INFO_TOKEN_EVENT_LINE)
-        {
+        if metadata.target() != "tiberius::tds::stream::token" {
             return;
         }
 
         let mut visitor = SqlServerMessageVisitor::default();
         event.record(&mut visitor);
-        if let Some(message) = visitor.message.filter(|message| !message.trim().is_empty()) {
-            if let Ok(mut messages) = self.messages.lock() {
-                messages.push(message);
-            }
+        let Some(message) = visitor.message.filter(|message| !message.trim().is_empty()) else {
+            return;
+        };
+        let captured = match (metadata.level(), metadata.line()) {
+            (&Level::INFO, Some(TIBERIUS_INFO_TOKEN_EVENT_LINE)) => Some(SqlServerTdsEvent::Message(message)),
+            (&Level::TRACE, Some(line)) => sqlserver_done_trace_event(line, &message),
+            _ => None,
+        };
+        if let (Some(captured), Ok(mut events)) = (captured, self.events.lock()) {
+            events.push(captured);
         }
     }
 }
@@ -419,36 +444,69 @@ impl tracing::field::Visit for SqlServerMessageVisitor {
     }
 }
 
+fn sqlserver_done_trace_event(line: u32, message: &str) -> Option<SqlServerTdsEvent> {
+    if !matches!(
+        line,
+        TIBERIUS_DONE_TOKEN_EVENT_LINE | TIBERIUS_DONE_PROC_TOKEN_EVENT_LINE | TIBERIUS_DONE_IN_PROC_TOKEN_EVENT_LINE
+    ) {
+        return None;
+    }
+    if line == TIBERIUS_DONE_PROC_TOKEN_EVENT_LINE && message.contains("BitFlags<DoneStatus>(0b0)") {
+        return None;
+    }
+    let rows = if !message.contains("Count") {
+        None
+    } else if message.ends_with("(1 row left)") {
+        Some(1)
+    } else if let Some(prefix) = message.strip_suffix(" rows left)") {
+        Some(prefix.rsplit_once('(').and_then(|(_, rows)| rows.parse::<u64>().ok()).unwrap_or(0))
+    } else {
+        Some(0)
+    };
+    Some(SqlServerTdsEvent::Done(rows))
+}
+
 async fn capture_sqlserver_messages<F, T>(future: F) -> (T, Vec<String>)
 where
     F: Future<Output = T>,
 {
-    let layer = SqlServerMessageLayer::default();
-    let messages = layer.messages.clone();
+    let layer = SqlServerTdsEventLayer::default();
+    let captured = layer.clone();
     // Tiberius consumes TDS INFO tokens internally and exposes them only as
     // tracing events. Scope collection to this future to isolate connections.
     let output = future.with_subscriber(tracing_subscriber::registry().with(layer)).await;
-    let messages = messages.lock().map(|messages| messages.clone()).unwrap_or_default();
+    let messages = captured
+        .take_events()
+        .into_iter()
+        .filter_map(|event| match event {
+            SqlServerTdsEvent::Message(message) => Some(message),
+            SqlServerTdsEvent::Done(_) => None,
+        })
+        .collect();
     (output, messages)
 }
 
-fn query_result_with_server_messages(mut result: QueryResult, messages: Vec<String>) -> QueryResult {
+fn query_result_with_server_messages(result: QueryResult, messages: Vec<String>) -> QueryResult {
+    query_result_with_server_messages_metadata(result, messages).result
+}
+
+fn query_result_with_server_messages_metadata(mut result: QueryResult, messages: Vec<String>) -> SqlServerBatchResult {
     if messages.is_empty() || !result.columns.is_empty() || !result.rows.is_empty() {
-        return result;
+        return SqlServerBatchResult { result, server_message: false };
     }
 
     result.columns = vec![SQLSERVER_MESSAGE_COLUMN.to_string()];
     result.column_types = vec!["nvarchar".to_string()];
     result.rows = messages.into_iter().map(|message| vec![serde_json::Value::String(message)]).collect();
-    result
+    SqlServerBatchResult { result, server_message: true }
 }
 
-fn server_messages_query_result(messages: Vec<String>, start: Instant) -> Option<QueryResult> {
+fn server_messages_query_result(messages: Vec<String>, start: Instant) -> Option<SqlServerBatchResult> {
     if messages.is_empty() {
         return None;
     }
 
-    Some(query_result_with_server_messages(
+    Some(query_result_with_server_messages_metadata(
         QueryResult {
             columns: vec![],
             column_types: vec![],
@@ -1229,21 +1287,91 @@ fn push_sqlserver_result_set(results: &mut Vec<QueryResult>, result: Option<SqlS
     }
 }
 
-async fn collect_result_sets_limited(
+fn push_sqlserver_ordered_result_set(
+    results: &mut Vec<SqlServerBatchResult>,
+    result: Option<SqlServerResultSet>,
+    start: Instant,
+) {
+    let mut query_results = Vec::with_capacity(1);
+    push_sqlserver_result_set(&mut query_results, result, start);
+    results.extend(query_results.into_iter().map(|result| SqlServerBatchResult { result, server_message: false }));
+}
+
+fn push_sqlserver_ordered_events(
+    results: &mut Vec<SqlServerBatchResult>,
+    current: &mut Option<SqlServerResultSet>,
+    events: Vec<SqlServerTdsEvent>,
+    start: Instant,
+) {
+    if events.is_empty() {
+        return;
+    }
+    // The first DONE after an active result set closes that SELECT. Its row
+    // count is not a standalone DML result.
+    let mut result_set_done_pending = current.is_some();
+    push_sqlserver_ordered_result_set(results, current.take(), start);
+
+    let mut messages = Vec::new();
+    for event in events {
+        match event {
+            SqlServerTdsEvent::Message(message) => messages.push(message),
+            SqlServerTdsEvent::Done(affected_rows) => {
+                if result_set_done_pending {
+                    result_set_done_pending = false;
+                    continue;
+                }
+                let Some(affected_rows) = affected_rows else {
+                    continue;
+                };
+                if let Some(message_result) = server_messages_query_result(std::mem::take(&mut messages), start) {
+                    results.push(message_result);
+                }
+                results.push(SqlServerBatchResult {
+                    result: QueryResult {
+                        columns: vec![],
+                        column_types: vec![],
+                        column_sortables: vec![],
+                        spatial_columns: vec![],
+                        spatial_values: vec![],
+                        rows: vec![],
+                        affected_rows,
+                        execution_time_ms: start.elapsed().as_millis(),
+                        truncated: false,
+                        session_id: None,
+                        has_more: false,
+                        elasticsearch_raw_body: None,
+                    },
+                    server_message: false,
+                });
+            }
+        }
+    }
+    if let Some(message_result) = server_messages_query_result(messages, start) {
+        results.push(message_result);
+    }
+}
+
+async fn collect_ordered_result_sets_limited(
     mut stream: QueryStream<'_>,
     start: Instant,
     max_rows: Option<usize>,
     result_offset: usize,
-) -> Result<Vec<QueryResult>, String> {
+    event_layer: SqlServerTdsEventLayer,
+) -> Result<Vec<SqlServerBatchResult>, String> {
     let row_limit = query_result_row_limit(max_rows);
     let mut results = Vec::new();
     let mut current: Option<SqlServerResultSet> = None;
     let mut saw_result_set = false;
 
-    while let Some(item) = stream.try_next().await.map_err(|e| e.to_string())? {
+    loop {
+        let item = stream.try_next().await.map_err(|e| e.to_string())?;
+        push_sqlserver_ordered_events(&mut results, &mut current, event_layer.take_events(), start);
+        let Some(item) = item else {
+            break;
+        };
         match item {
             QueryItem::Metadata(metadata) => {
-                push_sqlserver_result_set(&mut results, current.take(), start);
+                push_sqlserver_ordered_result_set(&mut results, current.take(), start);
                 let remaining_offset = if saw_result_set { 0 } else { result_offset };
                 saw_result_set = true;
                 current = Some(SqlServerResultSet {
@@ -1276,7 +1404,7 @@ async fn collect_result_sets_limited(
         }
     }
 
-    push_sqlserver_result_set(&mut results, current, start);
+    push_sqlserver_ordered_result_set(&mut results, current, start);
     Ok(results)
 }
 
@@ -2556,12 +2684,22 @@ pub async fn execute_batch_with_max_rows(
     sql: &str,
     max_rows: Option<usize>,
 ) -> Result<Vec<QueryResult>, String> {
+    execute_batch_with_max_rows_metadata(client, sql, max_rows)
+        .await
+        .map(|results| results.into_iter().map(|result| result.result).collect())
+}
+
+pub(crate) async fn execute_batch_with_max_rows_metadata(
+    client: &mut SqlServerClient,
+    sql: &str,
+    max_rows: Option<usize>,
+) -> Result<Vec<SqlServerBatchResult>, String> {
     let start = Instant::now();
     let result_offset = crate::query_result_sql::sqlserver_result_offset(sql);
     if sqlserver_batch_can_use_execute(sql) {
         let (result, messages) = capture_sqlserver_messages(sqlserver_driver_result(client.execute(sql, &[]))).await;
         let result = result?;
-        return Ok(vec![query_result_with_server_messages(
+        return Ok(vec![query_result_with_server_messages_metadata(
             QueryResult {
                 columns: vec![],
                 column_types: Vec::new(),
@@ -2596,8 +2734,8 @@ pub async fn execute_batch_with_max_rows(
                 })
                 .await;
                 return result.map(|result| {
-                    let mut result = query_result_with_server_messages(result, messages);
-                    strip_dbx_sqlserver_row_number_column(&mut result, sql);
+                    let mut result = query_result_with_server_messages_metadata(result, messages);
+                    strip_dbx_sqlserver_row_number_column(&mut result.result, sql);
                     vec![result]
                 });
             }
@@ -2605,7 +2743,7 @@ pub async fn execute_batch_with_max_rows(
             Ok(None) | Err(_) => {}
         }
     }
-    execute_simple_batch_with_max_rows(client, sql, max_rows).await
+    execute_simple_batch_with_max_rows_metadata(client, sql, max_rows).await
 }
 
 /// Execute a SQL Server batch directly through TDS simple-query mode.
@@ -2618,34 +2756,49 @@ pub async fn execute_simple_batch_with_max_rows(
     sql: &str,
     max_rows: Option<usize>,
 ) -> Result<Vec<QueryResult>, String> {
+    execute_simple_batch_with_max_rows_metadata(client, sql, max_rows)
+        .await
+        .map(|results| results.into_iter().map(|result| result.result).collect())
+}
+
+pub(crate) async fn execute_simple_batch_with_max_rows_metadata(
+    client: &mut SqlServerClient,
+    sql: &str,
+    max_rows: Option<usize>,
+) -> Result<Vec<SqlServerBatchResult>, String> {
     let start = Instant::now();
     let result_offset = crate::query_result_sql::sqlserver_result_offset(sql);
-    let (results, messages) = capture_sqlserver_messages(async {
+    let layer = SqlServerTdsEventLayer::default();
+    let captured = layer.clone();
+    let results = async {
         let stream = sqlserver_driver_result(client.simple_query(sql)).await?;
-        sqlserver_driver_result(collect_result_sets_limited(stream, start, max_rows, result_offset)).await
-    })
+        sqlserver_driver_result(collect_ordered_result_sets_limited(stream, start, max_rows, result_offset, captured))
+            .await
+    }
+    .with_subscriber(tracing_subscriber::registry().with(layer))
     .await;
     let mut results = results?;
     for result in &mut results {
-        strip_dbx_sqlserver_row_number_column(result, sql);
+        strip_dbx_sqlserver_row_number_column(&mut result.result, sql);
     }
 
-    if let Some(message_result) = server_messages_query_result(messages, start) {
-        results.push(message_result);
-    } else if results.is_empty() {
-        results.push(QueryResult {
-            columns: vec![],
-            column_types: Vec::new(),
-            column_sortables: vec![],
-            spatial_columns: vec![],
-            spatial_values: vec![],
-            rows: vec![],
-            affected_rows: 0,
-            execution_time_ms: start.elapsed().as_millis(),
-            truncated: false,
-            session_id: None,
-            has_more: false,
-            elasticsearch_raw_body: None,
+    if results.is_empty() {
+        results.push(SqlServerBatchResult {
+            result: QueryResult {
+                columns: vec![],
+                column_types: Vec::new(),
+                column_sortables: vec![],
+                spatial_columns: vec![],
+                spatial_values: vec![],
+                rows: vec![],
+                affected_rows: 0,
+                execution_time_ms: start.elapsed().as_millis(),
+                truncated: false,
+                session_id: None,
+                has_more: false,
+                elasticsearch_raw_body: None,
+            },
+            server_message: false,
         });
     }
 
@@ -2701,6 +2854,7 @@ fn is_dbx_sqlserver_row_number_page_sql(sql: &str) -> bool {
 fn sqlserver_batch_can_use_execute(sql: &str) -> bool {
     let contains_transaction_control = contains_transaction_control(sql);
     !requires_simple_query_batch(sql)
+        && !sqlserver_batch_may_emit_info_messages(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.
@@ -2709,6 +2863,10 @@ fn sqlserver_batch_can_use_execute(sql: &str) -> bool {
         && !sqlserver_dml_output_returns_rows(sql)
 }
 
+fn sqlserver_batch_may_emit_info_messages(sql: &str) -> bool {
+    top_level_sqlserver_tokens(sql).iter().any(|token| matches!(token.text.as_str(), "PRINT" | "RAISERROR" | "DBCC"))
+}
+
 fn sqlserver_balanced_transaction_batch_can_use_execute(sql: &str) -> bool {
     let tokens = top_level_sqlserver_tokens(sql);
     if tokens.iter().any(|token| {
@@ -2905,18 +3063,19 @@ mod tests {
     use super::{
         build_sqlserver_unsafe_type_query, capture_sqlserver_messages, completion_context_from_query_result,
         decode_sqlserver_spatial_values, format_sqlserver_numeric, is_blocking_sqlserver_unsafe_probe_error,
-        is_sqlserver_spatial_column, is_sqlserver_variant_column, query_result_with_server_messages,
-        requires_simple_query_batch, restore_sqlserver_legacy_probe_output_names,
-        restore_sqlserver_spatial_column_types, sqlserver_batch_can_use_execute, sqlserver_bulk_token_row,
-        sqlserver_cell_to_json, sqlserver_columns_sql, sqlserver_completion_assistant_sql,
-        sqlserver_dml_output_returns_rows, sqlserver_filter_definition_error, sqlserver_hidden_schema_names,
-        sqlserver_indexes_sql, sqlserver_legacy_indexes_sql, sqlserver_legacy_probe, sqlserver_legacy_probe_with_nonce,
+        is_sqlserver_spatial_column, is_sqlserver_variant_column, push_sqlserver_ordered_events,
+        query_result_with_server_messages, query_result_with_server_messages_metadata, requires_simple_query_batch,
+        restore_sqlserver_legacy_probe_output_names, restore_sqlserver_spatial_column_types,
+        sqlserver_batch_can_use_execute, sqlserver_bulk_token_row, sqlserver_cell_to_json, sqlserver_columns_sql,
+        sqlserver_completion_assistant_sql, sqlserver_dml_output_returns_rows, sqlserver_done_trace_event,
+        sqlserver_filter_definition_error, sqlserver_hidden_schema_names, sqlserver_indexes_sql,
+        sqlserver_legacy_indexes_sql, sqlserver_legacy_probe, sqlserver_legacy_probe_with_nonce,
         sqlserver_list_objects_sql, sqlserver_list_schemas_sql, sqlserver_list_tables_sql,
         sqlserver_probe_explicit_alias, sqlserver_schema_name_predicate, sqlserver_spatial_marker,
         sqlserver_supports_session_database_switch, sqlserver_table_comment_sql, sqlserver_triggers_sql,
         sqlserver_visible_object_predicate, strip_dbx_sqlserver_row_number_column, SqlServerDescribedColumn,
-        SqlServerProbeOutputNameOverride, SqlServerResultSet, SqlServerSpatialColumn, SQLSERVER_COMPLETION_CONTEXT_SQL,
-        SQLSERVER_RESULT_TYPE_PROBE_SQL,
+        SqlServerProbeOutputNameOverride, SqlServerResultSet, SqlServerSpatialColumn, SqlServerTdsEvent,
+        SQLSERVER_COMPLETION_CONTEXT_SQL, SQLSERVER_RESULT_TYPE_PROBE_SQL,
     };
     use crate::types::{
         CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest, QueryResult,
@@ -2966,6 +3125,25 @@ mod tests {
         assert_eq!(result.columns, vec!["Message"]);
         assert_eq!(result.rows, vec![vec![serde_json::json!("DBCC execution completed")]]);
 
+        let message = query_result_with_server_messages_metadata(
+            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,
+            },
+            vec!["PRINT output".to_string()],
+        );
+        assert!(message.server_message);
+
         let select = QueryResult {
             columns: vec!["id".to_string()],
             column_types: vec!["int".to_string()],
@@ -2980,9 +3158,150 @@ mod tests {
             has_more: false,
             elasticsearch_raw_body: None,
         };
-        let result = query_result_with_server_messages(select, vec!["informational".to_string()]);
-        assert_eq!(result.columns, vec!["id"]);
-        assert_eq!(result.rows, vec![vec![serde_json::json!(1)]]);
+        let result = query_result_with_server_messages_metadata(select, vec!["informational".to_string()]);
+        assert!(!result.server_message);
+        assert_eq!(result.result.columns, vec!["id"]);
+        assert_eq!(result.result.rows, vec![vec![serde_json::json!(1)]]);
+    }
+
+    #[test]
+    fn sqlserver_done_trace_events_keep_statement_boundaries_and_counts() {
+        assert_eq!(
+            sqlserver_done_trace_event(
+                super::TIBERIUS_DONE_TOKEN_EVENT_LINE,
+                "Done with status BitFlags<DoneStatus>(0b10001, More | Count) (2 rows left)",
+            ),
+            Some(SqlServerTdsEvent::Done(Some(2)))
+        );
+        assert_eq!(
+            sqlserver_done_trace_event(
+                super::TIBERIUS_DONE_IN_PROC_TOKEN_EVENT_LINE,
+                "Done with status BitFlags<DoneStatus>(0b10000, Count)",
+            ),
+            Some(SqlServerTdsEvent::Done(Some(0)))
+        );
+        assert_eq!(
+            sqlserver_done_trace_event(
+                super::TIBERIUS_DONE_PROC_TOKEN_EVENT_LINE,
+                "Done with status BitFlags<DoneStatus>(0b0)",
+            ),
+            None
+        );
+        assert_eq!(
+            sqlserver_done_trace_event(
+                super::TIBERIUS_DONE_TOKEN_EVENT_LINE,
+                "Done with status BitFlags<DoneStatus>(0b1, More)",
+            ),
+            Some(SqlServerTdsEvent::Done(None))
+        );
+    }
+
+    #[test]
+    fn sqlserver_ordered_events_preserve_message_result_and_count_order() {
+        let start = Instant::now();
+        let mut results = Vec::new();
+        let mut current = None;
+
+        push_sqlserver_ordered_events(
+            &mut results,
+            &mut current,
+            vec![SqlServerTdsEvent::Message("before".to_string())],
+            start,
+        );
+        current = Some(SqlServerResultSet {
+            columns: vec!["first".to_string()],
+            column_types: vec!["int".to_string()],
+            rows: vec![vec![serde_json::json!(1)]],
+            truncated: false,
+            remaining_offset: 0,
+        });
+        push_sqlserver_ordered_events(
+            &mut results,
+            &mut current,
+            vec![SqlServerTdsEvent::Done(Some(1)), SqlServerTdsEvent::Message("between".to_string())],
+            start,
+        );
+        current = Some(SqlServerResultSet {
+            columns: vec!["second".to_string()],
+            column_types: vec!["int".to_string()],
+            rows: vec![vec![serde_json::json!(2)]],
+            truncated: false,
+            remaining_offset: 0,
+        });
+        push_sqlserver_ordered_events(
+            &mut results,
+            &mut current,
+            vec![
+                SqlServerTdsEvent::Done(Some(1)),
+                SqlServerTdsEvent::Message("after one".to_string()),
+                SqlServerTdsEvent::Message("after two".to_string()),
+                SqlServerTdsEvent::Done(Some(3)),
+            ],
+            start,
+        );
+
+        assert_eq!(results.len(), 6);
+        assert!(results[0].server_message);
+        assert_eq!(results[0].result.rows, vec![vec![serde_json::json!("before")]]);
+        assert_eq!(results[1].result.columns, vec!["first"]);
+        assert!(results[2].server_message);
+        assert_eq!(results[2].result.rows, vec![vec![serde_json::json!("between")]]);
+        assert_eq!(results[3].result.columns, vec!["second"]);
+        assert!(results[4].server_message);
+        assert_eq!(
+            results[4].result.rows,
+            vec![vec![serde_json::json!("after one")], vec![serde_json::json!("after two")]]
+        );
+        assert_eq!(results[5].result.affected_rows, 3);
+        assert!(!results[5].server_message);
+    }
+
+    #[test]
+    fn sqlserver_ordered_events_drop_select_counts_and_keep_following_dml_counts() {
+        let start = Instant::now();
+        let mut results = Vec::new();
+        let mut current = Some(SqlServerResultSet {
+            columns: vec!["selected".to_string()],
+            column_types: vec!["int".to_string()],
+            rows: vec![vec![serde_json::json!(1)]],
+            truncated: false,
+            remaining_offset: 0,
+        });
+
+        push_sqlserver_ordered_events(
+            &mut results,
+            &mut current,
+            vec![SqlServerTdsEvent::Done(Some(1)), SqlServerTdsEvent::Done(Some(2))],
+            start,
+        );
+
+        assert_eq!(results.len(), 2);
+        assert_eq!(results[0].result.columns, vec!["selected"]);
+        assert_eq!(results[1].result.affected_rows, 2);
+    }
+
+    #[test]
+    fn sqlserver_ordered_events_use_no_count_done_to_close_select_result() {
+        let start = Instant::now();
+        let mut results = Vec::new();
+        let mut current = Some(SqlServerResultSet {
+            columns: vec!["selected".to_string()],
+            column_types: vec!["int".to_string()],
+            rows: vec![vec![serde_json::json!(1)]],
+            truncated: false,
+            remaining_offset: 0,
+        });
+
+        push_sqlserver_ordered_events(
+            &mut results,
+            &mut current,
+            vec![SqlServerTdsEvent::Done(None), SqlServerTdsEvent::Done(Some(2))],
+            start,
+        );
+
+        assert_eq!(results.len(), 2);
+        assert_eq!(results[0].result.columns, vec!["selected"]);
+        assert_eq!(results[1].result.affected_rows, 2);
     }
 
     #[test]
@@ -3137,6 +3456,11 @@ mod tests {
         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;"
         ));
+        assert!(!sqlserver_batch_can_use_execute(
+            "PRINT N'before'; UPDATE dbo.users SET active = 0 WHERE id = 1; PRINT N'after';"
+        ));
+        assert!(!sqlserver_batch_can_use_execute("DBCC CHECKIDENT ('dbo.users', NORESEED);"));
+        assert!(!sqlserver_batch_can_use_execute("RAISERROR(N'progress', 0, 1) WITH NOWAIT;"));
     }
 
     #[test]