fix: preserve explicit TOP and OFFSET/FETCH limits

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

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<String> {
     let mut words = Vec::new();
     let mut current = String::new();
@@ -450,7 +481,9 @@ fn top_level_words(statement: &str) -> Vec<String> {
                 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<char> = 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(' '),

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!(

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();