diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index dad6f01fad..9c24894efa 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -2496,30 +2496,8 @@ pub async fn execute_query_with_max_rows( let mut result = query_result_with_server_messages(result?, messages); strip_dbx_sqlserver_row_number_column(&mut result, sql); Ok(result) - } else if requires_simple_query_batch(sql) || contains_transaction_control(sql) { - let (result, messages) = capture_sqlserver_messages(async { - let stream = sqlserver_driver_result(client.simple_query(sql)).await?; - sqlserver_driver_result(collect_result_sets_limited(stream, start, max_rows)).await - }) - .await; - let _ = result?; - Ok(query_result_with_server_messages( - 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, - }, - messages, - )) + } else if !sqlserver_batch_can_use_execute(sql) { + execute_simple_batch_first_result_with_max_rows(client, sql, max_rows).await } else { let (result, messages) = capture_sqlserver_messages(sqlserver_driver_result(client.execute(sql, &[]))).await; let result = result?; @@ -2645,6 +2623,22 @@ pub async fn execute_simple_batch_with_max_rows( Ok(results) } +async fn execute_simple_batch_first_result_with_max_rows( + client: &mut SqlServerClient, + sql: &str, + max_rows: Option, +) -> Result { + let start = Instant::now(); + let (result, messages) = capture_sqlserver_messages(async { + let stream = sqlserver_driver_result(client.simple_query(sql)).await?; + sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, &[])).await + }) + .await; + let mut result = query_result_with_server_messages(result?, messages); + strip_dbx_sqlserver_row_number_column(&mut result, sql); + Ok(result) +} + fn strip_dbx_sqlserver_row_number_column(result: &mut QueryResult, sql: &str) { if !is_dbx_sqlserver_row_number_page_sql(sql) { return; @@ -2685,13 +2679,18 @@ fn sqlserver_batch_can_use_execute(sql: &str) -> bool { } fn sqlserver_batch_may_return_result_set(sql: &str) -> bool { - let tokens = top_level_sqlserver_tokens(sql); - let starts_with_cte_dml = tokens.first().is_some_and(|token| token.text == "WITH") - && tokens.iter().any(|token| matches!(token.text.as_str(), "INSERT" | "UPDATE" | "DELETE" | "MERGE")) - && crate::sql::split_sql_statements(sql).len() == 1; - tokens.iter().any(|token| { - matches!(token.text.as_str(), "SELECT" | "EXEC" | "EXECUTE" | "TABLE") - || (token.text == "WITH" && !starts_with_cte_dml) + crate::sql::split_sql_statements(sql).iter().any(|statement| { + let tokens = top_level_sqlserver_tokens(statement); + let starts_with_dml = starts_with_executable_sql_keyword(statement, &["INSERT", "UPDATE", "DELETE", "MERGE"]); + if starts_with_dml { + return false; + } + let starts_with_cte_dml = tokens.first().is_some_and(|token| token.text == "WITH") + && tokens.iter().any(|token| matches!(token.text.as_str(), "INSERT" | "UPDATE" | "DELETE" | "MERGE")); + tokens.iter().any(|token| { + matches!(token.text.as_str(), "SELECT" | "EXEC" | "EXECUTE" | "TABLE") + || (token.text == "WITH" && !starts_with_cte_dml) + }) }) } @@ -3053,6 +3052,12 @@ mod tests { 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( + "INSERT INTO dbo.user_archive(id, name) SELECT id, name FROM dbo.users WHERE active = 0;" + )); + assert!(!sqlserver_batch_can_use_execute( + "INSERT INTO dbo.user_archive(id) SELECT id FROM dbo.users; SELECT COUNT(*) FROM dbo.user_archive;" + )); 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;" @@ -3145,6 +3150,9 @@ mod tests { 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("DECLARE @id INT = 1; EXEC dbo.list_users @id;")); + assert!(!sqlserver_batch_can_use_execute( + "SET NOCOUNT OFF; DECLARE @rows TABLE (id INT); INSERT INTO @rows VALUES (1), (2); SELECT id FROM @rows;" + )); assert!(!sqlserver_batch_can_use_execute( "DECLARE @id INT = 1; CREATE TABLE #t(id INT); INSERT INTO #t VALUES (@id); SELECT id FROM #t;" )); @@ -3194,6 +3202,22 @@ mod tests { assert!(execute_query.contains("client.query(query.sql.as_str(), &[])")); } + #[test] + fn sqlserver_single_result_entrypoint_keeps_result_returning_batches() { + let source = include_str!("sqlserver.rs"); + let execute_query = source.split("pub async fn execute_query_with_max_rows").nth(1).unwrap(); + let execute_query = execute_query.split("pub async fn execute_batch").next().unwrap(); + + assert!(execute_query.contains("!sqlserver_batch_can_use_execute(sql)")); + assert!(execute_query.contains("execute_simple_batch_first_result_with_max_rows(client, sql, max_rows)")); + assert!(!execute_query.contains("execute_simple_batch_with_max_rows(client, sql, max_rows)")); + + let first_result = source.split("async fn execute_simple_batch_first_result_with_max_rows").nth(1).unwrap(); + let first_result = first_result.split("fn strip_dbx_sqlserver_row_number_column").next().unwrap(); + assert!(first_result.contains("collect_first_result_limited(stream, start, max_rows, &[])")); + assert!(!first_result.contains("collect_result_sets_limited")); + } + #[test] fn sqlserver_user_query_paths_do_not_collect_full_results_before_limiting() { let source = include_str!("sqlserver.rs"); diff --git a/crates/dbx-core/tests/sqlserver_batch_results.rs b/crates/dbx-core/tests/sqlserver_batch_results.rs new file mode 100644 index 0000000000..96a9830cc0 --- /dev/null +++ b/crates/dbx-core/tests/sqlserver_batch_results.rs @@ -0,0 +1,73 @@ +use dbx_core::db::sqlserver; +use std::time::Duration; + +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"); + sqlserver::connect_with_port_explicit(&host, port, true, &user, &password, Some("master"), 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_query_keeps_result_set_after_batch_setup_statements() { + let mut client = connect_sqlserver().await; + + let sql = "SET NOCOUNT OFF; \ + DECLARE @rows TABLE (id INT, label NVARCHAR(20)); \ + INSERT INTO @rows VALUES (1, N'first'), (2, N'second'); \ + SELECT id, label FROM @rows ORDER BY id;"; + let result = sqlserver::execute_query_with_max_rows(&mut client, sql, Some(100)) + .await + .expect("execute result-returning batch"); + + assert_eq!(result.columns, vec!["id", "label"]); + assert_eq!(result.rows.len(), 2); + assert_eq!(result.rows[0][0].as_i64(), Some(1)); + assert_eq!(result.rows[1][1].as_str(), Some("second")); + assert!(!result.truncated); + + let limited = sqlserver::execute_query_with_max_rows(&mut client, sql, Some(1)) + .await + .expect("execute limited result-returning batch"); + assert_eq!(limited.rows.len(), 1); + assert!(limited.truncated); +} + +#[tokio::test] +#[ignore = "requires DBX_TEST_SQLSERVER_HOST and DBX_TEST_SQLSERVER_PASSWORD"] +async fn sqlserver_insert_select_reports_affected_rows() { + let mut client = connect_sqlserver().await; + let sql = "DECLARE @target TABLE (id INT); \ + INSERT INTO @target(id) SELECT id FROM (VALUES (1), (2), (3)) AS source(id);"; + + let result = sqlserver::execute_query(&mut client, sql).await.expect("execute INSERT SELECT batch"); + + assert_eq!(result.affected_rows, 3); + assert!(result.columns.is_empty()); + assert!(result.rows.is_empty()); +} + +#[tokio::test] +#[ignore = "requires DBX_TEST_SQLSERVER_HOST and DBX_TEST_SQLSERVER_PASSWORD"] +async fn sqlserver_single_result_query_drains_later_results() { + let mut client = connect_sqlserver().await; + let sql = "SELECT CAST(1 AS INT) AS retained; \ + SELECT TOP (10000) object_id AS discarded FROM sys.all_objects ORDER BY object_id;"; + + let result = + sqlserver::execute_query_with_max_rows(&mut client, sql, Some(1)).await.expect("execute multi-result query"); + + assert_eq!(result.columns, vec!["retained"]); + assert_eq!(result.rows.len(), 1); + assert_eq!(result.rows[0][0].as_i64(), Some(1)); + assert!(!result.truncated); + + let follow_up = sqlserver::execute_query(&mut client, "SELECT CAST(2 AS INT) AS value") + .await + .expect("execute query after draining later results"); + assert_eq!(follow_up.rows[0][0].as_i64(), Some(2)); +}