From 8825719fffc74f505bcd05224f4ad93a521a66d4 Mon Sep 17 00:00:00 2001 From: Andrea Debernardi Date: Thu, 24 Sep 2026 17:23:30 +0200 Subject: [PATCH] fix: preserve explicit SQL Server row limits (#26) --- CHANGELOG.md | 6 ++++ src/driver/helpers.rs | 49 +++++++++++++++++++++++++++++-- src/driver/helpers/tests.rs | 46 ++++++++++++++++++++++++++++++ tests/live_db.rs | 57 +++++++++++++++++++++++++++++++++++++ 4 files changed, 155 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ffa78e..23b80df 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- Preserve explicit outer `TOP` and `OFFSET/FETCH` limits instead of adding + incompatible automatic pagination (#26). Limits inside CTEs and subqueries + still allow pagination of the outer query. + ## [1.0.0-beta.2] - 2026-09-21 ### Added diff --git a/src/driver/helpers.rs b/src/driver/helpers.rs index acb72dd..419a1d2 100644 --- a/src/driver/helpers.rs +++ b/src/driver/helpers.rs @@ -348,6 +348,8 @@ pub fn query_returns_result_set(query: &str) -> bool { .any(|statement| statement_returns_result_set(statement)) } +/// Only add host pagination to a single result-bearing statement without an +/// explicit outer row limit. Inner TOP/OFFSET clauses do not limit its result. pub fn query_can_be_paginated(query: &str) -> bool { let statements = top_level_statements(&code_mask(query)); statements.len() == 1 && statement_can_be_paginated(&statements[0]) @@ -390,7 +392,10 @@ pub fn build_paginated_query(query: &str, page_size: u32, page: u32) -> String { fn statement_can_be_paginated(statement: &str) -> bool { let words = top_level_words(statement); statement_operation(&words).is_some_and(|(operation_index, operation)| match operation { - "SELECT" => !select_has_top_level_into(&words, operation_index), + "SELECT" => { + !select_has_top_level_into(&words, operation_index) + && !select_has_top_level_row_limit(&words, operation_index) + } "VALUES" => true, _ => false, }) @@ -434,6 +439,32 @@ fn select_has_top_level_into(words: &[String], operation_index: usize) -> bool { .any(|word| word == "INTO") } +fn select_has_top_level_row_limit(words: &[String], operation_index: usize) -> bool { + let words = &words[operation_index + 1..]; + // TOP is reserved; quoted identifiers, literals, comments and nested + // scopes have already been removed. Also preserve TOP in set operands. + if words.iter().any(|word| word == "TOP") { + return true; + } + + // OFFSET is not reserved, so merely selecting a column named `offset` + // must not disable pagination. A paging clause follows ORDER BY and + // terminates its count with ROW/ROWS (FETCH is optional). + let Some(order_index) = words + .windows(2) + .position(|pair| pair[0] == "ORDER" && pair[1] == "BY") + else { + return false; + }; + let order_words = &words[order_index + 2..]; + order_words.iter().enumerate().any(|(index, word)| { + word == "OFFSET" + && order_words[index + 1..] + .iter() + .any(|word| matches!(word.as_str(), "ROW" | "ROWS")) + }) +} + fn top_level_words(statement: &str) -> Vec { let mut words = Vec::new(); let mut current = String::new(); @@ -450,7 +481,9 @@ fn top_level_words(statement: &str) -> Vec { depth = depth.saturating_sub(1); current.clear(); } - _ if depth == 0 && (character.is_alphanumeric() || character == '_') => { + _ if depth == 0 + && (character.is_alphanumeric() || matches!(character, '_' | '@' | '#' | '$')) => + { current.push(character.to_ascii_uppercase()); } _ if depth == 0 && !current.is_empty() => { @@ -501,6 +534,7 @@ fn code_mask(query: &str) -> String { let characters: Vec = query.chars().collect(); let mut masked = String::with_capacity(query.len()); let mut state = State::Normal; + let mut block_depth = 0_u32; let mut position = 0; while position < characters.len() { let character = characters[position]; @@ -526,6 +560,7 @@ fn code_mask(query: &str) -> String { } ('/', Some('*')) => { state = State::BlockComment; + block_depth = 1; masked.push_str(" "); position += 1; } @@ -562,9 +597,17 @@ fn code_mask(query: &str) -> String { masked.push(character); state = State::Normal; } + State::BlockComment if character == '/' && next == Some('*') => { + masked.push_str(" "); + block_depth += 1; + position += 1; + } State::BlockComment if character == '*' && next == Some('/') => { masked.push_str(" "); - state = State::Normal; + block_depth -= 1; + if block_depth == 0 { + state = State::Normal; + } position += 1; } _ => masked.push(' '), diff --git a/src/driver/helpers/tests.rs b/src/driver/helpers/tests.rs index c57f09e..07193ba 100644 --- a/src/driver/helpers/tests.rs +++ b/src/driver/helpers/tests.rs @@ -435,6 +435,52 @@ fn affected_rows_are_only_reported_for_final_dml_statement() { assert!(!query_reports_affected_rows("DROP TABLE dbo.items")); } +#[test] +fn classifier_preserves_explicit_outer_row_limits() { + for query in [ + "SELECT TOP 100 * FROM users", + "select top(100) * from users;", + "SELECT DISTINCT TOP (10) id FROM users", + "SELECT ALL TOP (@count) id FROM users", + "SELECT TOP (10) PERCENT WITH TIES id FROM users ORDER BY id", + "SELECT /* limit */ TOP /* count */ (100) * FROM users", + ";WITH source AS (SELECT id FROM users) SELECT TOP (2) id FROM source", + "SELECT id FROM users ORDER BY id OFFSET 10 ROWS", + "SELECT id FROM users ORDER BY id OFFSET (10) ROW FETCH NEXT (5) ROWS ONLY", + "SELECT id FROM users ORDER BY id OFFSET @skip ROWS FETCH FIRST @take ROW ONLY", + "SELECT id FROM users ORDER BY id OFFSET (SELECT 10) ROWS", + "WITH source AS (SELECT id FROM users) SELECT id FROM source ORDER BY id OFFSET 0 ROWS FETCH NEXT 10 ROWS ONLY; -- done", + "SELECT id FROM users ORDER BY id OFFSET /* skip */ 0 ROWS /* take */ FETCH NEXT 10 ROWS ONLY", + "SELECT TOP (1) id FROM users UNION ALL SELECT id FROM admins", + ] { + assert!(query_returns_result_set(query), "lost result set: {query}"); + assert!(!query_can_be_paginated(query), "paginated: {query}"); + assert!(!query_reports_affected_rows(query), "lost SELECT: {query}"); + } +} + +#[test] +fn classifier_still_paginates_when_row_limits_are_only_nested_or_masked() { + for query in [ + "SELECT * FROM (SELECT TOP (5) * FROM users ORDER BY id) AS recent", + "WITH recent AS (SELECT TOP (5) id FROM users) SELECT id FROM recent", + "WITH recent AS (SELECT id FROM users ORDER BY id OFFSET 0 ROWS FETCH NEXT 5 ROWS ONLY) SELECT id FROM recent", + "SELECT (SELECT TOP (1) id FROM users) AS first_id", + "SELECT * FROM (SELECT id FROM users ORDER BY id OFFSET 0 ROWS) AS recent", + "SELECT 'TOP (100) OFFSET 0 ROWS FETCH NEXT 5 ROWS ONLY' AS label", + "SELECT [TOP], [OFFSET], [FETCH] FROM [users]", + "SELECT \"TOP\", \"OFFSET\", \"FETCH\" FROM users", + "SELECT id FROM users -- TOP (100) OFFSET 0 ROWS FETCH NEXT 5 ROWS ONLY", + "SELECT id FROM users /* TOP (100) OFFSET 0 ROWS FETCH NEXT 5 ROWS ONLY */", + "SELECT id /* outer /* nested */ TOP (100) */ FROM users", + "SELECT @top, @offset, @fetch FROM users", + "SELECT offset FROM users ORDER BY offset", + ] { + assert!(query_returns_result_set(query), "lost result set: {query}"); + assert!(query_can_be_paginated(query), "not paginated: {query}"); + } +} + #[test] fn paginated_query_adds_order_when_missing() { assert_eq!( diff --git a/tests/live_db.rs b/tests/live_db.rs index 5fa4c70..935c883 100644 --- a/tests/live_db.rs +++ b/tests/live_db.rs @@ -1314,6 +1314,63 @@ fn pagination_and_batch_semantics_cover_ordered_unordered_cte_and_dml() { assert_eq!(batch[3]["result"]["pagination"]["has_more"], true); } +#[test] +fn explicit_row_limits_bypass_host_pagination_in_single_and_batch_queries() { + let mut plugin = Plugin::with_scratch_database(); + let source = "(VALUES (1), (2), (3), (4), (5)) AS source(id)"; + let cases = [ + ( + format!("SELECT TOP 3 id FROM {source} ORDER BY id"), + json!([[1], [2], [3]]), + ), + ( + format!("SELECT DISTINCT TOP (3) id FROM {source} ORDER BY id; -- keep limit"), + json!([[1], [2], [3]]), + ), + ( + format!("SELECT TOP (60) PERCENT id FROM {source} ORDER BY id"), + json!([[1], [2], [3]]), + ), + ( + format!("SELECT TOP (3) WITH TIES id FROM {source} ORDER BY id"), + json!([[1], [2], [3]]), + ), + ( + format!("WITH cte AS (SELECT id FROM {source}) SELECT TOP (3) id FROM cte ORDER BY id"), + json!([[1], [2], [3]]), + ), + ( + format!("SELECT id FROM {source} ORDER BY id OFFSET 2 ROWS"), + json!([[3], [4], [5]]), + ), + ( + format!("SELECT id FROM {source} ORDER BY id OFFSET 1 ROWS FETCH NEXT 3 ROWS ONLY"), + json!([[2], [3], [4]]), + ), + ]; + + for (query, expected_rows) in &cases { + let result = plugin.call_ok( + "execute_query", + json!({ "params": connection_params(), "query": query, "limit": 1, "page": 2 }), + ); + assert_eq!(result["rows"], *expected_rows, "{query}"); + assert_eq!(result["pagination"], Value::Null, "{query}"); + assert_eq!(result["truncated"], false, "{query}"); + } + + let queries: Vec<_> = cases.iter().map(|(query, _)| query).collect(); + let batch = plugin.call_ok( + "execute_query_batch", + json!({ "params": connection_params(), "queries": queries, "limit": 1, "page": 2 }), + ); + for (index, (query, expected_rows)) in cases.iter().enumerate() { + assert_eq!(batch[index]["result"]["rows"], *expected_rows, "{query}"); + assert_eq!(batch[index]["result"]["pagination"], Value::Null, "{query}"); + assert_eq!(batch[index]["result"]["truncated"], false, "{query}"); + } +} + #[test] fn million_row_query_is_bounded_and_marks_truncation() { let mut plugin = Plugin::with_scratch_database();