github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 05:20:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(sqlserver): preserve batch results and DML counts

Browse Source
This commit is contained in:
zipg
2026-08-03 15:23:12 +08:00
committed by GitHub
parent 147e92ff39
commit 02bad5d754
2 changed files with 128 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/db/sqlserver.rs
+55 -31
View File
@@ -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<usize>,
) -> Result<QueryResult, String> {
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");
crates/dbx-core/tests/sqlserver_batch_results.rs
+73
View File
@@ -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));
}