xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(sqlserver): support 2008 query pagination

This commit is contained in:
t8y2 2026-06-26 12:03:06 +08:00
parent 476c5cbc09
commit c3d47ece2c
2 changed files with 122 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

74
crates/dbx-core/src/db/sqlserver.rs
View File

@ -1520,7 +1520,9 @@ pub async fn execute_query_with_max_rows(
Ok(None) | Err(_) => sql.to_string(),
};
let stream = sqlserver_driver_result(client.query(query_sql.as_str(), &[])).await?;
sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows)).await
let mut result = sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows)).await?;
strip_dbx_sqlserver_row_number_column(&mut result, sql);
Ok(result)
} else if requires_simple_query_batch(sql) || is_transaction_control(sql) {
let stream = sqlserver_driver_result(client.simple_query(sql)).await?;
let _ = sqlserver_driver_result(collect_result_sets_limited(stream, start, max_rows)).await?;
@ -1579,13 +1581,19 @@ pub async fn execute_batch_with_max_rows(
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?;
return sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows))
.await
.map(|result| vec![result]);
return sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows)).await.map(
|mut result| {
strip_dbx_sqlserver_row_number_column(&mut result, sql);
vec![result]
},
);
}
}
let stream = sqlserver_driver_result(client.simple_query(sql)).await?;
let mut results = sqlserver_driver_result(collect_result_sets_limited(stream, start, max_rows)).await?;
for result in &mut results {
strip_dbx_sqlserver_row_number_column(result, sql);
}
if results.is_empty() {
results.push(QueryResult {
@ -1604,6 +1612,35 @@ pub async fn execute_batch_with_max_rows(
Ok(results)
}
fn strip_dbx_sqlserver_row_number_column(result: &mut QueryResult, sql: &str) {
if !is_dbx_sqlserver_row_number_page_sql(sql) {
return;
}
if !result.columns.last().is_some_and(|column| column.eq_ignore_ascii_case("__dbx_row_num")) {
return;
}
result.columns.pop();
if result.column_types.len() > result.columns.len() {
result.column_types.pop();
}
if result.column_sortables.len() > result.columns.len() {
result.column_sortables.pop();
}
for row in &mut result.rows {
if row.len() > result.columns.len() {
row.pop();
}
}
}
fn is_dbx_sqlserver_row_number_page_sql(sql: &str) -> bool {
let normalized = sql.to_ascii_uppercase();
normalized.contains("ROW_NUMBER() OVER")
&& normalized.contains("[__DBX_ROW_NUM]")
&& normalized.contains("DBX_PAGE_SOURCE.*")
}
fn sqlserver_batch_can_use_execute(sql: &str) -> bool {
!requires_simple_query_batch(sql)
&& !sqlserver_batch_may_return_result_set(sql)
@ -1705,9 +1742,12 @@ mod tests {
sqlserver_batch_can_use_execute, sqlserver_cell_to_json, sqlserver_columns_sql,
sqlserver_completion_assistant_sql, sqlserver_dml_output_returns_rows, sqlserver_hidden_schema_names,
sqlserver_indexes_sql, sqlserver_list_objects_sql, sqlserver_list_schemas_sql, sqlserver_list_tables_sql,
sqlserver_table_comment_sql, sqlserver_visible_object_predicate, SqlServerDescribedColumn, SqlServerResultSet,
sqlserver_table_comment_sql, sqlserver_visible_object_predicate, strip_dbx_sqlserver_row_number_column,
SqlServerDescribedColumn, SqlServerResultSet,
};
use crate::types::{
CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest, QueryResult,
};
use crate::types::{CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest};
use chrono::NaiveDate;
use std::time::Instant;
use tiberius::{ColumnData, IntoSql};
@ -2062,6 +2102,28 @@ mod tests {
assert_eq!(sqlserver_cell_to_json(&cell), serde_json::json!("0x000000000001cf53"));
}
#[test]
fn sqlserver_strips_generated_row_number_pagination_column() {
let sql = "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS [__dbx_row_num] FROM (SELECT id FROM users) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];";
let mut result = QueryResult {
columns: vec!["id".to_string(), "__dbx_row_num".to_string()],
column_types: vec!["int".to_string(), "bigint".to_string()],
column_sortables: vec![],
rows: vec![vec![serde_json::json!(42), serde_json::json!(101)]],
affected_rows: 0,
execution_time_ms: 1,
truncated: false,
session_id: None,
has_more: false,
};
strip_dbx_sqlserver_row_number_column(&mut result, sql);
assert_eq!(result.columns, vec!["id"]);
assert_eq!(result.column_types, vec!["int"]);
assert_eq!(result.rows, vec![vec![serde_json::json!(42)]]);
}
#[test]
fn sqlserver_detects_geometry_result_columns() {
assert!(is_sqlserver_spatial_column(&SqlServerDescribedColumn {

89
crates/dbx-core/src/query_result_sql.rs
View File

@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize};
use crate::models::connection::DatabaseType;
use crate::sql::find_statement_at_cursor;
use crate::sql_dialect::{pagination_strategy, quote_table_identifier, PaginationContext, TablePaginationStrategy};
use sqlparser::ast::{GroupByExpr, SelectItem, SetExpr, Statement};
use sqlparser::ast::{Expr, GroupByExpr, SelectItem, SetExpr, Statement};
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser;
@ -405,13 +405,50 @@ fn add_sql_server_offset_fetch(statement: &str, limit: usize, offset: usize) ->
return (offset == 0).then(|| statement.to_string());
}
let has_order_by = find_top_level_trailing_order_by(statement).is_some();
if !has_order_by && has_top_level_select_distinct(statement) {
let order_by_index = find_top_level_trailing_order_by(statement);
if order_by_index.is_none() && has_top_level_select_distinct(statement) {
return (offset == 0).then(|| add_sql_server_top(statement, limit));
}
let order_by = if has_order_by { String::new() } else { " ORDER BY (SELECT NULL)".to_string() };
Some(format!("{statement}{order_by} OFFSET {offset} ROWS FETCH NEXT {limit} ROWS ONLY;"))
if offset == 0 {
return Some(add_sql_server_top(statement, limit));
}
let statement_without_order = order_by_index.map(|index| statement[..index].trim_end()).unwrap_or(statement);
if !sql_server_row_number_pagination_safe(statement_without_order) {
return None;
}
let row_number_order = order_by_index
.map(|index| statement[index..].trim().to_string())
.unwrap_or_else(|| "ORDER BY (SELECT NULL)".to_string());
let end = offset + limit;
Some(format!(
"SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER ({row_number_order}) AS [__dbx_row_num] FROM ({statement_without_order}) dbx_page_source) dbx_page WHERE [__dbx_row_num] > {offset} AND [__dbx_row_num] <= {end} ORDER BY [__dbx_row_num];"
))
}
fn sql_server_row_number_pagination_safe(statement: &str) -> bool {
let dialect = GenericDialect {};
let Ok(statements) = Parser::parse_sql(&dialect, statement) else {
return false;
};
let [Statement::Query(query)] = statements.as_slice() else {
return false;
};
let SetExpr::Select(select) = query.body.as_ref() else {
return false;
};
select.projection.iter().all(|item| match item {
SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _) => false,
SelectItem::UnnamedExpr(expr) => sql_server_derived_projection_has_name(expr),
SelectItem::ExprWithAlias { .. } | SelectItem::ExprWithAliases { .. } => true,
})
}
fn sql_server_derived_projection_has_name(expr: &Expr) -> bool {
matches!(expr, Expr::Identifier(_) | Expr::CompoundIdentifier(_))
}
fn add_sql_server_top(sql: &str, limit: usize) -> String {
@ -883,7 +920,7 @@ mod tests {
}
#[test]
fn uses_sqlserver_offset_fetch_pagination_for_first_page() {
fn uses_sqlserver_top_pagination_for_first_page() {
let result = build_paginated_query_sql(PaginatedQuerySqlOptions {
original_sql: "SELECT id FROM users ORDER BY id DESC".to_string(),
database_type: Some(DatabaseType::SqlServer),
@ -892,14 +929,11 @@ mod tests {
});
assert!(result.ok);
assert_eq!(
result.sql.unwrap(),
"SELECT id FROM users ORDER BY id DESC OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY;"
);
assert_eq!(result.sql.unwrap(), "SELECT TOP (100) id FROM users ORDER BY id DESC");
}
#[test]
fn uses_sqlserver_offset_fetch_for_count_queries_without_derived_table() {
fn uses_sqlserver_top_for_count_queries_without_derived_table() {
let result = build_paginated_query_sql(PaginatedQuerySqlOptions {
original_sql: "SELECT COUNT(*) FROM TicketInfo".to_string(),
database_type: Some(DatabaseType::SqlServer),
@ -908,10 +942,7 @@ mod tests {
});
assert!(result.ok);
assert_eq!(
result.sql.unwrap(),
"SELECT COUNT(*) FROM TicketInfo ORDER BY (SELECT NULL) OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY;"
);
assert_eq!(result.sql.unwrap(), "SELECT TOP (100) COUNT(*) FROM TicketInfo");
}
#[test]
@ -943,7 +974,7 @@ mod tests {
}
#[test]
fn paginates_sqlserver_all_queries_with_offset_fetch() {
fn paginates_sqlserver_all_queries_with_top() {
let result = build_paginated_query_sql(PaginatedQuerySqlOptions {
original_sql: "SELECT ALL ProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data".to_string(),
database_type: Some(DatabaseType::SqlServer),
@ -952,10 +983,7 @@ mod tests {
});
assert!(result.ok);
assert_eq!(
result.sql.unwrap(),
"SELECT ALL ProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data ORDER BY (SELECT NULL) OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY;"
);
assert_eq!(result.sql.unwrap(), "SELECT ALL TOP (100) ProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data");
}
#[test]
@ -968,10 +996,7 @@ mod tests {
});
assert!(result.ok);
assert_eq!(
result.sql.unwrap(),
"SELECT AllProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data ORDER BY (SELECT NULL) OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY;"
);
assert_eq!(result.sql.unwrap(), "SELECT TOP (100) AllProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data");
}
#[test]
@ -1092,14 +1117,11 @@ WHERE u.id = picked.id;
});
assert!(result.ok);
assert_eq!(
result.sql.unwrap(),
"SELECT @@version ORDER BY (SELECT NULL) OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY;"
);
assert_eq!(result.sql.unwrap(), "SELECT TOP (100) @@version");
}
#[test]
fn uses_sqlserver_offset_fetch_pagination_for_later_pages() {
fn uses_sqlserver_row_number_pagination_for_later_pages() {
let result = build_paginated_query_sql(PaginatedQuerySqlOptions {
original_sql: "SELECT id FROM users".to_string(),
database_type: Some(DatabaseType::SqlServer),
@ -1109,12 +1131,12 @@ WHERE u.id = picked.id;
assert_eq!(
result.sql.unwrap(),
"SELECT id FROM users ORDER BY (SELECT NULL) OFFSET 300 ROWS FETCH NEXT 100 ROWS ONLY;"
"SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS [__dbx_row_num] FROM (SELECT id FROM users) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 300 AND [__dbx_row_num] <= 400 ORDER BY [__dbx_row_num];"
);
}
#[test]
fn sqlserver_join_pagination_does_not_wrap_duplicate_columns() {
fn rejects_sqlserver_wildcard_later_pages_to_avoid_duplicate_columns() {
let result = build_paginated_query_sql(PaginatedQuerySqlOptions {
original_sql:
"SELECT b.ProjectType,* FROM VesselBusinessOpportunity a LEFT JOIN JDDR_sys_BasicConfig_ProjectInfo_Data b ON a.ProjectID = b.ID"
@ -1124,10 +1146,7 @@ WHERE u.id = picked.id;
offset: 100,
});
assert_eq!(
result.sql.unwrap(),
"SELECT b.ProjectType,* FROM VesselBusinessOpportunity a LEFT JOIN JDDR_sys_BasicConfig_ProjectInfo_Data b ON a.ProjectID = b.ID ORDER BY (SELECT NULL) OFFSET 100 ROWS FETCH NEXT 100 ROWS ONLY;"
);
assert_eq!(result, err("unsupported"));
}
#[test]