From ad81f02db4ecc56e27f688062f49e6f10175a209 Mon Sep 17 00:00:00 2001 From: zipg Date: Wed, 5 Aug 2026 15:53:16 +0800 Subject: [PATCH] fix(sqlserver): show PRINT messages --- .../__tests__/useSqlExecution.spec.ts | 94 ++++ .../src/composables/useSqlExecution.ts | 8 +- apps/desktop/src/types/database.ts | 2 + crates/dbx-core/src/db/sqlserver.rs | 448 +++++++++++++++--- crates/dbx-core/src/query.rs | 48 +- .../tests/sqlserver_server_messages.rs | 124 ++++- crates/dbx-web/src/routes/query.rs | 1 + 7 files changed, 638 insertions(+), 87 deletions(-) diff --git a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts index 9a07d07362..6ba2c56e0b 100644 --- a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts +++ b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts @@ -170,6 +170,100 @@ describe("useSqlExecution", () => { expect(activeOutputView.value).toBe("summary"); }); + it("shows SQL Server PRINT messages instead of leaving them behind the execution summary", async () => { + const sql = `IF 1 = 1 +BEGIN + PRINT 'x'; +END +ELSE +BEGIN + PRINT 'y'; +END +GO`; + const activeTab = ref({ ...queryTab("dbx_sqlserver_demo"), sql }); + const activeConnection = ref(connection("sqlserver")); + const activeOutputView = ref<"result" | "summary" | "explain" | "chart">("result"); + const queryStore = useQueryStore(); + vi.spyOn(queryStore, "executeCurrentSql").mockImplementation(async () => { + if (activeTab.value) activeTab.value.result = { columns: ["Message"], column_types: ["nvarchar"], rows: [["x"]], affected_rows: 0, execution_time_ms: 1, server_message: true }; + }); + vi.spyOn(useHistoryStore(), "add").mockResolvedValue(undefined); + + const execution = useSqlExecution({ + activeTab: computed(() => activeTab.value), + activeConnection: computed(() => activeConnection.value), + executableSql: computed(() => sql), + activeOutputView, + }); + + await execution.tryExecute(); + + expect(activeOutputView.value).toBe("result"); + expect(activeTab.value?.result?.rows).toEqual([["x"]]); + }); + + it("selects a trailing SQL Server PRINT result after a data result", async () => { + const sql = "SELECT 1 AS value; PRINT N'x';"; + const activeTab = ref({ ...queryTab("dbx_sqlserver_demo"), sql }); + const activeConnection = ref(connection("sqlserver")); + const activeOutputView = ref<"result" | "summary" | "explain" | "chart">("result"); + const queryStore = useQueryStore(); + const setActiveResultIndex = vi.spyOn(queryStore, "setActiveResultIndex").mockImplementation((_id, index) => { + if (!activeTab.value?.results) return; + activeTab.value.activeResultIndex = index; + activeTab.value.result = activeTab.value.results[index]; + }); + vi.spyOn(queryStore, "executeCurrentSql").mockImplementation(async () => { + if (!activeTab.value) return; + const dataResult = { columns: ["value"], column_types: ["int"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }; + const messageResult = { columns: ["Message"], column_types: ["nvarchar"], rows: [["x"]], affected_rows: 0, execution_time_ms: 1, server_message: true }; + activeTab.value.results = [dataResult, messageResult]; + activeTab.value.activeResultIndex = 0; + activeTab.value.result = dataResult; + }); + vi.spyOn(useHistoryStore(), "add").mockResolvedValue(undefined); + + const execution = useSqlExecution({ + activeTab: computed(() => activeTab.value), + activeConnection: computed(() => activeConnection.value), + executableSql: computed(() => sql), + activeOutputView, + }); + + await execution.tryExecute(); + + expect(activeOutputView.value).toBe("result"); + expect(setActiveResultIndex).toHaveBeenCalledWith("tab-1", 1); + expect(activeTab.value?.activeResultIndex).toBe(1); + expect(activeTab.value?.result?.server_message).toBe(true); + expect(activeTab.value?.result?.rows).toEqual([["x"]]); + }); + + it("keeps the summary for ordinary SQL Server data aliased as Message", async () => { + const sql = `DECLARE @value nvarchar(1) = N'x'; +SELECT @value AS Message;`; + const activeTab = ref({ ...queryTab("dbx_sqlserver_demo"), sql }); + const activeConnection = ref(connection("sqlserver")); + const activeOutputView = ref<"result" | "summary" | "explain" | "chart">("result"); + const queryStore = useQueryStore(); + vi.spyOn(queryStore, "executeCurrentSql").mockImplementation(async () => { + if (activeTab.value) activeTab.value.result = { columns: ["Message"], column_types: ["nvarchar"], rows: [["x"]], affected_rows: 0, execution_time_ms: 1 }; + }); + vi.spyOn(useHistoryStore(), "add").mockResolvedValue(undefined); + + const execution = useSqlExecution({ + activeTab: computed(() => activeTab.value), + activeConnection: computed(() => activeConnection.value), + executableSql: computed(() => sql), + activeOutputView, + }); + + await execution.tryExecute(); + + expect(activeOutputView.value).toBe("summary"); + expect(activeTab.value?.result?.rows).toEqual([["x"]]); + }); + it("forwards execute-in-new-result-tab intent to the query store", async () => { const sql = "SELECT * FROM users"; const activeTab = ref({ ...queryTab("app"), sql }); diff --git a/apps/desktop/src/composables/useSqlExecution.ts b/apps/desktop/src/composables/useSqlExecution.ts index 998be9f9e9..e0b13328b2 100644 --- a/apps/desktop/src/composables/useSqlExecution.ts +++ b/apps/desktop/src/composables/useSqlExecution.ts @@ -228,7 +228,13 @@ export function useSqlExecution(deps: { ...(options.openInNewResultTab ? { openInNewResultTab: true } : {}), }); if (producedResult === false) return; - if (tab.result && !tab.result.columns.length && !tab.results?.some((result) => result.columns.length > 0)) { + const sqlServerMessageResultIndex = executionDatabaseType === "sqlserver" ? tab.results?.findIndex((result) => result.server_message === true) : undefined; + if (sqlServerMessageResultIndex !== undefined && sqlServerMessageResultIndex >= 0) { + queryStore.setActiveResultIndex(tab.id, sqlServerMessageResultIndex); + deps.activeOutputView.value = "result"; + } else if (executionDatabaseType === "sqlserver" && tab.result?.server_message === true) { + deps.activeOutputView.value = "result"; + } else if (tab.result && !tab.result.columns.length && !tab.results?.some((result) => result.columns.length > 0)) { deps.activeOutputView.value = "summary"; } const elapsed = Date.now() - start; diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index a280b029e0..9e67253118 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -580,6 +580,8 @@ export interface QueryResult { appended_from_row_count?: number; /** Set for synthesized query execution failures. */ execution_error?: true; + /** Set only for SQL Server informational messages emitted by the backend. */ + server_message?: true; /** Structured backend error; authoritative when execution_error is true. */ error?: BackendError; /** Zero-based index of the submitted statement that produced this result. */ diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index e8e599c809..512e7914d0 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -374,34 +374,59 @@ fn column_types_from_metadata(metadata: &tiberius::ResultMetadata) -> Vec>>, +#[derive(Debug, Clone)] +pub(crate) struct SqlServerBatchResult { + pub result: QueryResult, + pub server_message: bool, } -impl Layer for SqlServerMessageLayer +#[derive(Debug, Clone, PartialEq, Eq)] +enum SqlServerTdsEvent { + Message(String), + Done(Option), +} + +#[derive(Clone, Default)] +struct SqlServerTdsEventLayer { + events: StdArc>>, +} + +impl SqlServerTdsEventLayer { + fn take_events(&self) -> Vec { + self.events.lock().map(|mut events| std::mem::take(&mut *events)).unwrap_or_default() + } +} + +impl Layer 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 { + 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(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::().ok()).unwrap_or(0)) + } else { + Some(0) + }; + Some(SqlServerTdsEvent::Done(rows)) +} + async fn capture_sqlserver_messages(future: F) -> (T, Vec) where F: Future, { - 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) -> QueryResult { +fn query_result_with_server_messages(result: QueryResult, messages: Vec) -> QueryResult { + query_result_with_server_messages_metadata(result, messages).result +} + +fn query_result_with_server_messages_metadata(mut result: QueryResult, messages: Vec) -> 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, start: Instant) -> Option { +fn server_messages_query_result(messages: Vec, start: Instant) -> Option { 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, result: Option, + result: Option, + 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, + current: &mut Option, + events: Vec, + 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, result_offset: usize, -) -> Result, String> { + event_layer: SqlServerTdsEventLayer, +) -> Result, String> { let row_limit = query_result_row_limit(max_rows); let mut results = Vec::new(); let mut current: Option = 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, ) -> Result, 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, +) -> Result, 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, ) -> Result, 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, +) -> Result, 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(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(0b10000, Count)", + ), + Some(SqlServerTdsEvent::Done(Some(0))) + ); + assert_eq!( + sqlserver_done_trace_event( + super::TIBERIUS_DONE_PROC_TOKEN_EVENT_LINE, + "Done with status BitFlags(0b0)", + ), + None + ); + assert_eq!( + sqlserver_done_trace_event( + super::TIBERIUS_DONE_TOKEN_EVENT_LINE, + "Done with status BitFlags(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] diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index a43faf2aeb..debe5dc000 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -144,6 +144,7 @@ fn query_error_with_omitted_sql_context(error: &str, _sql: &str) -> String { /// `execution_error` is emitted for synthesized per-statement errors so clients /// can distinguish them from a successful result column named `Error`. /// `statement_index` is emitted only after a concrete statement starts running. +/// `server_message` is emitted only for SQL Server TDS informational messages. #[derive(Debug, Clone, Serialize)] pub struct ExecuteMultiResult { #[serde(flatten)] @@ -154,6 +155,8 @@ pub struct ExecuteMultiResult { pub statement_index: Option, #[serde(skip_serializing_if = "Option::is_none")] pub error: Option, + #[serde(skip_serializing_if = "is_false")] + pub server_message: bool, } #[derive(Debug, Clone, PartialEq)] @@ -193,12 +196,12 @@ fn report_execute_multi_progress( impl ExecuteMultiResult { fn execution_error(result: db::QueryResult) -> Self { let error = error_from_query_result(&result); - Self { result, execution_error: true, statement_index: None, error } + Self { result, execution_error: true, statement_index: None, error, server_message: false } } fn execution_error_with_index(result: db::QueryResult, statement_index: usize) -> Self { let error = error_from_query_result(&result); - Self { result, execution_error: true, statement_index: Some(statement_index), error } + Self { result, execution_error: true, statement_index: Some(statement_index), error, server_message: false } } fn execution_error_with_backend( @@ -206,11 +209,17 @@ impl ExecuteMultiResult { statement_index: Option, error: crate::backend_error::BackendError, ) -> Self { - Self { result, execution_error: true, statement_index, error: Some(error) } + Self { result, execution_error: true, statement_index, error: Some(error), server_message: false } } fn success_with_index(result: db::QueryResult, statement_index: usize) -> Self { - Self { result, execution_error: false, statement_index: Some(statement_index), error: None } + Self { + result, + execution_error: false, + statement_index: Some(statement_index), + error: None, + server_message: false, + } } pub fn without_error_detail(mut self) -> Self { @@ -225,7 +234,19 @@ impl ExecuteMultiResult { impl From for ExecuteMultiResult { fn from(result: db::QueryResult) -> Self { - Self { result, execution_error: false, statement_index: None, error: None } + Self { result, execution_error: false, statement_index: None, error: None, server_message: false } + } +} + +impl From for ExecuteMultiResult { + fn from(result: db::sqlserver::SqlServerBatchResult) -> Self { + Self { + result: result.result, + execution_error: false, + statement_index: None, + error: None, + server_message: result.server_message, + } } } @@ -2431,7 +2452,7 @@ fn empty_query_result(execution_time_ms: u128) -> db::QueryResult { } } -fn sqlserver_batch_results(results: Vec) -> Vec { +fn sqlserver_batch_results(results: Vec) -> Vec { results.into_iter().map(ExecuteMultiResult::from).collect() } @@ -2494,9 +2515,9 @@ async fn execute_multi_sqlserver( let execution = async { if execution_mode == QueryExecutionMode::Simple { - db::sqlserver::execute_simple_batch_with_max_rows(&mut client_guard, batch, max_rows).await + db::sqlserver::execute_simple_batch_with_max_rows_metadata(&mut client_guard, batch, max_rows).await } else { - db::sqlserver::execute_batch_with_max_rows(&mut client_guard, batch, max_rows).await + db::sqlserver::execute_batch_with_max_rows_metadata(&mut client_guard, batch, max_rows).await } }; let result = wait_for_result_opt(cancel_token.clone(), query_timeout, execution).await; @@ -4954,6 +4975,7 @@ for line in sys.stdin: let success = serde_json::to_value(ExecuteMultiResult::from(empty_query_result(0))).unwrap(); assert!(success.get("execution_error").is_none()); assert!(success.get("statement_index").is_none()); + assert!(success.get("server_message").is_none()); let mut error_column = empty_query_result(0); error_column.columns = vec!["Error".to_string()]; @@ -4987,9 +5009,17 @@ for line in sys.stdin: fn sqlserver_batch_results_do_not_claim_statement_indexes() { assert_eq!(split_sql_batches("SELECT 1; SELECT 2;").len(), 1); - let results = sqlserver_batch_results(vec![empty_query_result(1), empty_query_result(2)]); + let results = sqlserver_batch_results(vec![ + db::sqlserver::SqlServerBatchResult { result: empty_query_result(1), server_message: false }, + db::sqlserver::SqlServerBatchResult { result: empty_query_result(2), server_message: true }, + ]); assert_eq!(results.iter().map(|result| result.statement_index).collect::>(), vec![None, None]); + assert!(!results[0].server_message); + assert!(results[1].server_message); + + let serialized = serde_json::to_value(&results[1]).unwrap(); + assert_eq!(serialized.get("server_message"), Some(&serde_json::Value::Bool(true))); } #[test] diff --git a/crates/dbx-core/tests/sqlserver_server_messages.rs b/crates/dbx-core/tests/sqlserver_server_messages.rs index c851f1f5b3..4e2908bee0 100644 --- a/crates/dbx-core/tests/sqlserver_server_messages.rs +++ b/crates/dbx-core/tests/sqlserver_server_messages.rs @@ -1,24 +1,21 @@ use dbx_core::db::sqlserver; use std::time::Duration; -#[tokio::test] -#[ignore = "requires DBX_TEST_SQLSERVER_HOST and DBX_TEST_SQLSERVER_PASSWORD"] -async fn sqlserver_dbcc_messages_do_not_replace_results_or_errors() { +async fn connect_sqlserver() -> sqlserver::SqlServerClient { let host = std::env::var("DBX_TEST_SQLSERVER_HOST").expect("DBX_TEST_SQLSERVER_HOST"); let port = std::env::var("DBX_TEST_SQLSERVER_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(1433); let user = std::env::var("DBX_TEST_SQLSERVER_USER").unwrap_or_else(|_| "sa".to_string()); let password = std::env::var("DBX_TEST_SQLSERVER_PASSWORD").expect("DBX_TEST_SQLSERVER_PASSWORD"); - let mut client = sqlserver::connect_with_port_explicit( - &host, - port, - true, - &user, - &password, - Some("master"), - Duration::from_secs(15), - ) - .await - .expect("connect to SQL Server"); + let database = std::env::var("DBX_TEST_SQLSERVER_DATABASE").unwrap_or_else(|_| "master".to_string()); + sqlserver::connect_with_port_explicit(&host, port, true, &user, &password, Some(&database), Duration::from_secs(15)) + .await + .expect("connect to SQL Server") +} + +#[tokio::test] +#[ignore = "requires DBX_TEST_SQLSERVER_HOST and DBX_TEST_SQLSERVER_PASSWORD"] +async fn sqlserver_dbcc_messages_do_not_replace_results_or_errors() { + let mut client = connect_sqlserver().await; let table = "dbo.dbx_issue_3583_messages"; let setup = format!( @@ -54,6 +51,8 @@ async fn sqlserver_dbcc_messages_do_not_replace_results_or_errors() { assert_eq!(dml.affected_rows, 1); assert!(dml.columns.is_empty()); + sqlserver::execute_query(&mut client, &format!("DROP TABLE {table}")).await.expect("clean up DBCC fixture"); + let use_database = sqlserver::execute_query(&mut client, "USE master").await.expect("execute database context change"); assert_eq!(use_database.columns, vec!["Message"]); @@ -70,6 +69,101 @@ async fn sqlserver_dbcc_messages_do_not_replace_results_or_errors() { .await .expect_err("real SQL Server errors must remain failures"); assert!(!error.trim().is_empty()); +} - sqlserver::execute_query(&mut client, &format!("DROP TABLE {table}")).await.expect("clean up DBCC fixture"); +#[tokio::test] +#[ignore = "requires DBX_TEST_SQLSERVER_HOST and DBX_TEST_SQLSERVER_PASSWORD"] +async fn sqlserver_print_messages_keep_tds_order_around_results_and_counts() { + let mut client = connect_sqlserver().await; + + let result_order = sqlserver::execute_batch( + &mut client, + "SET NOCOUNT ON; \ + PRINT N'before'; \ + SELECT CAST(1 AS int) AS first; \ + PRINT N'between'; \ + SELECT CAST(2 AS int) AS second; \ + PRINT N'after one'; \ + PRINT N'after two';", + ) + .await + .expect("execute PRINT and result-set batch"); + + assert_eq!(result_order.len(), 5, "unexpected ordered results: {result_order:?}"); + assert_eq!(result_order[0].columns, vec!["Message"]); + assert_eq!(result_order[0].rows, vec![vec![serde_json::json!("before")]]); + assert_eq!(result_order[1].columns, vec!["first"]); + assert_eq!(result_order[2].rows, vec![vec![serde_json::json!("between")]]); + assert_eq!(result_order[3].columns, vec!["second"]); + assert_eq!(result_order[4].rows, vec![vec![serde_json::json!("after one")], vec![serde_json::json!("after two")]]); + + sqlserver::execute_simple_batch_with_max_rows( + &mut client, + "SET NOCOUNT ON; \ + CREATE TABLE #dbx_print_order (id int NOT NULL PRIMARY KEY, value int NOT NULL); \ + INSERT INTO #dbx_print_order VALUES (1, 0), (2, 0); \ + SET NOCOUNT OFF;", + None, + ) + .await + .expect("create PRINT ordering fixture"); + let count_order = sqlserver::execute_batch( + &mut client, + "PRINT N'before update'; \ + UPDATE #dbx_print_order SET value = value + 1; \ + PRINT N'after update';", + ) + .await + .expect("execute PRINT and update-count batch"); + + assert_eq!(count_order.len(), 3, "unexpected message/count order: {count_order:?}"); + assert_eq!(count_order[0].rows, vec![vec![serde_json::json!("before update")]]); + assert_eq!(count_order[1].affected_rows, 2); + assert_eq!(count_order[2].rows, vec![vec![serde_json::json!("after update")]]); +} + +#[tokio::test] +#[ignore = "requires DBX_TEST_SQLSERVER_HOST and DBX_TEST_SQLSERVER_PASSWORD"] +async fn sqlserver_default_no_count_keeps_select_and_dml_result_shape() { + let mut client = connect_sqlserver().await; + + let single = sqlserver::execute_batch(&mut client, "SELECT CAST(1 AS int) AS selected") + .await + .expect("execute one SELECT with NOCOUNT OFF"); + assert_eq!(single.len(), 1, "SELECT count must not become a separate result: {single:?}"); + assert_eq!(single[0].columns, vec!["selected"]); + + let multiple = + sqlserver::execute_batch(&mut client, "SELECT CAST(1 AS int) AS first; SELECT CAST(2 AS int) AS second;") + .await + .expect("execute multiple SELECT statements with NOCOUNT OFF"); + assert_eq!(multiple.len(), 2, "SELECT counts must not become separate results: {multiple:?}"); + assert_eq!(multiple[0].columns, vec!["first"]); + assert_eq!(multiple[1].columns, vec!["second"]); + + sqlserver::execute_simple_batch_with_max_rows( + &mut client, + "SET NOCOUNT ON; \ + CREATE TABLE #dbx_select_dml_order (id int NOT NULL PRIMARY KEY, value int NOT NULL); \ + INSERT INTO #dbx_select_dml_order VALUES (1, 0), (2, 0); \ + SET NOCOUNT OFF;", + None, + ) + .await + .expect("create SELECT and DML ordering fixture"); + let mixed = sqlserver::execute_batch( + &mut client, + "SELECT CAST(1 AS int) AS first; \ + UPDATE #dbx_select_dml_order SET value = value + 1; \ + SELECT CAST(2 AS int) AS second; \ + DELETE FROM #dbx_select_dml_order;", + ) + .await + .expect("execute SELECT and DML batch with NOCOUNT OFF"); + + assert_eq!(mixed.len(), 4, "unexpected SELECT and DML result order: {mixed:?}"); + assert_eq!(mixed[0].columns, vec!["first"]); + assert_eq!(mixed[1].affected_rows, 2); + assert_eq!(mixed[2].columns, vec!["second"]); + assert_eq!(mixed[3].affected_rows, 2); } diff --git a/crates/dbx-web/src/routes/query.rs b/crates/dbx-web/src/routes/query.rs index 82bc909896..81374e514c 100644 --- a/crates/dbx-web/src/routes/query.rs +++ b/crates/dbx-web/src/routes/query.rs @@ -1033,6 +1033,7 @@ mod tests { error: Some(dbx_core::backend_error::BackendError::from_sql_detail( "relation customer_orders does not exist", )), + server_message: false, }; let payload = serde_json::to_value(execute_multi_response(vec![result]).0).unwrap();