From b814b3d9befc931c4f5cd44d3e273c17f9ac47e1 Mon Sep 17 00:00:00 2001 From: Abeautifulsnow Date: Wed, 12 Aug 2026 00:19:57 +0800 Subject: [PATCH] fix(kingbase): handle TOP-bounded export pagination --- crates/dbx-core/src/query_result_export.rs | 242 ++++++++++++++++++--- crates/dbx-core/src/query_result_sql.rs | 230 ++++++++++++++++++++ 2 files changed, 440 insertions(+), 32 deletions(-) diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 6c00534abb..1df20bd87c 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -19,7 +19,8 @@ use crate::query::{ operation_budget_for_pool_key, QueryExecutionOptions, StreamProgressClock, QUERY_CANCELED, }; use crate::query_result_sql::{ - build_query_pagination_execution_plan, QueryPagination, QueryPaginationExecutionPlanOptions, + build_query_pagination_execution_plan, has_top_level_top, top_level_top_row_count, QueryPagination, + QueryPaginationExecutionPlanOptions, }; use crate::table_export::TableExportProgress; use crate::transfer::keyset_pagination_sql; @@ -461,6 +462,38 @@ fn supports_streaming_offset_pagination(request: &QueryResultExportRequest, page && !first_sql.trim().eq_ignore_ascii_case(second_sql.trim()) } +/// Enforceable in-memory row bound for a single-execution export, or `None` +/// when the query cannot be safely streamed in one shot without an Agent +/// cursor. Kingbase SQL Server compatibility mode TOP queries cannot be +/// rewritten with LIMIT/OFFSET; a concrete `TOP n` bounds the result and the +/// user's export row limit caps it further. Percentage TOP / `WITH TIES` have +/// no concrete row bound, so they are only single-execution-capable when a row +/// limit is configured. +fn single_execution_row_bound(request: &QueryResultExportRequest) -> Option { + if !has_top_level_top(&request.sql) { + return None; + } + match (top_level_top_row_count(&request.sql), request.row_limit) { + (Some(top), Some(row_limit)) => Some(top.min(row_limit)), + (Some(top), None) => Some(top), + (None, Some(row_limit)) => Some(row_limit), + (None, None) => None, + } +} + +fn single_execution_page_limit(request: &QueryResultExportRequest, page_size: usize) -> Option { + single_execution_row_bound(request).filter(|bound| *bound > 0 && *bound <= page_size.max(1)) +} + +/// True when a non-agent export can still stream this query by executing it +/// exactly once without exceeding one normal export page. Larger TOP/row-limit +/// bounds require an Agent result session; executing them in one response would +/// defeat streaming and recreate the large-result memory spike. +#[cfg(test)] +fn supports_single_execution_export(request: &QueryResultExportRequest, page_size: usize) -> bool { + single_execution_page_limit(request, page_size).is_some() +} + #[derive(Debug, Clone, PartialEq, Eq)] struct SafeKeysetCandidate { schema: Option, @@ -511,6 +544,7 @@ fn safe_keyset_candidate(sql: &str) -> Option { return None; }; if select.distinct.is_some() + || select.top.is_some() || !matches!(&select.group_by, GroupByExpr::Expressions(exprs, _) if exprs.is_empty()) || select.having.is_some() || select.selection.is_some() @@ -650,7 +684,16 @@ async fn export_query_result_core_inner( let mut offset: usize = 0; let mut wrote_text_header = false; let mut keyset_plan = build_keyset_plan(state, request).await; - if keyset_plan.is_none() && !request.use_agent_cursor && !supports_streaming_offset_pagination(request, page_size) { + // A Kingbase SQL Server compat TOP query that cannot be offset-paginated is + // exported with a single execution whose page size is the enforceable row + // bound (concrete TOP count and/or the configured export row limit), then it + // stops after the first response. + let single_execution_bound = single_execution_page_limit(request, page_size); + if keyset_plan.is_none() + && !request.use_agent_cursor + && !supports_streaming_offset_pagination(request, page_size) + && single_execution_bound.is_none() + { return Err(STREAMING_PAGINATION_UNSUPPORTED_ERROR.to_string()); } @@ -674,37 +717,45 @@ async fn export_query_result_core_inner( if matches!(remaining, Some(0)) { break; } - let this_page = remaining.map_or(page_size, |rem| rem.min(page_size)).max(1); - - let (sql_to_execute, plan_limit, use_agent_result_session) = if let Some(plan) = keyset_plan.as_ref() { - ( - keyset_pagination_sql( - &plan.columns, - &plan.table, - &plan.schema, - &request.database_type, - &plan.primary_keys, - &plan.last_pk_values, - this_page, - ), - this_page, - false, - ) + let this_page = if keyset_plan.is_none() && single_execution_bound.is_some() { + // Single execution covers the whole enforceable TOP/row-limit bound in + // one shot; never an unbounded i32::MAX page. + single_execution_bound.unwrap_or(page_size).max(1) } else { - let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { - sql: request.sql.clone(), - query_base_sql: request.query_base_sql.clone(), - database_type: Some(request.database_type), - pagination: QueryPagination { limit: this_page, offset, session_id: session_id.clone() }, - use_agent_cursor: request.use_agent_cursor, - first_page_uses_actual_sql: true, - }); - let Some(plan_limit) = plan.page_limit else { - return Err("Failed to build query pagination plan for export".to_string()); - }; - (plan.sql_to_execute, plan_limit, plan.use_agent_result_session) + remaining.map_or(page_size, |rem| rem.min(page_size)).max(1) }; + let (sql_to_execute, plan_limit, use_agent_result_session, single_execution) = + if let Some(plan) = keyset_plan.as_ref() { + ( + keyset_pagination_sql( + &plan.columns, + &plan.table, + &plan.schema, + &request.database_type, + &plan.primary_keys, + &plan.last_pk_values, + this_page, + ), + this_page, + false, + false, + ) + } else { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: request.sql.clone(), + query_base_sql: request.query_base_sql.clone(), + database_type: Some(request.database_type), + pagination: QueryPagination { limit: this_page, offset, session_id: session_id.clone() }, + use_agent_cursor: request.use_agent_cursor, + first_page_uses_actual_sql: true, + }); + let Some(plan_limit) = plan.page_limit else { + return Err("Failed to build query pagination plan for export".to_string()); + }; + (plan.sql_to_execute, plan_limit, plan.use_agent_result_session, plan.single_execution) + }; + let options = if use_agent_result_session { QueryExecutionOptions { page_size: Some(plan_limit), @@ -841,8 +892,13 @@ async fn export_query_result_core_inner( plan.pk_indices.iter().map(|&index| last_row.get(index).cloned().unwrap_or(Value::Null)).collect(); } } - let should_continue = - should_fetch_next_page(use_agent_result_session, result.has_more, fetched_row_count, row_count, plan_limit); + let should_continue = if single_execution { + // A single execution already streamed the full (TOP-bounded) result; + // there is no offset to advance to. + false + } else { + should_fetch_next_page(use_agent_result_session, result.has_more, fetched_row_count, row_count, plan_limit) + }; if cancel_token.as_ref().is_some_and(|token| token.is_cancelled()) || is_export_cancelled(&request.export_id).await { @@ -1996,6 +2052,128 @@ mod tests { assert!(!supports_streaming_offset_pagination(&oracle_req, 100)); } + #[test] + fn kingbase_non_keyset_top_export_falls_back_to_single_execution() { + // Regression for t8y2/dbx#5910: a non-keyset Kingbase SQL Server compat + // TOP query (e.g. a join) cannot be offset-paginated, so the export must + // stream it in a single execution rather than reject it. This mirrors the + // guard in export_query_result_core_inner: offset pagination says no, but + // single-execution support lets the export proceed. + let req = QueryResultExportRequest { + sql: "SELECT TOP 100 * FROM orders o JOIN customers c ON c.id = o.customer_id ORDER BY o.id".to_string(), + query_base_sql: "SELECT TOP 100 * FROM orders o JOIN customers c ON c.id = o.customer_id ORDER BY o.id" + .to_string(), + database_type: DatabaseType::Kingbase, + use_agent_cursor: false, + ..request("csv", Some(1000), None) + }; + + assert!(!supports_streaming_offset_pagination(&req, 100)); + assert!(supports_single_execution_export(&req, 100)); + // The enforceable bound is the concrete TOP count (100), not the row limit. + assert_eq!(single_execution_row_bound(&req), Some(100)); + } + + #[test] + fn kingbase_single_table_top_query_never_uses_keyset_and_is_bounded() { + // P0 regression: a simple `SELECT TOP 100 * FROM users` must not qualify + // for the keyset path (which reconstructs SQL and drops TOP), and its + // single-execution bound is exactly 100 — never an unbounded page. + let req = QueryResultExportRequest { + sql: "SELECT TOP 100 * FROM users".to_string(), + query_base_sql: "SELECT TOP 100 * FROM users".to_string(), + database_type: DatabaseType::Kingbase, + use_agent_cursor: false, + ..request("csv", None, None) + }; + + assert!(safe_keyset_candidate(&req.sql).is_none(), "TOP must not qualify for keyset"); + assert!(!supports_streaming_offset_pagination(&req, 100)); + assert_eq!(single_execution_row_bound(&req), Some(100)); + assert!(supports_single_execution_export(&req, 100)); + } + + #[test] + fn kingbase_percent_and_with_ties_need_a_row_limit_for_single_execution() { + // Percentage TOP and WITH TIES have no concrete row-count bound, so + // without a configured export row limit the single-execution fallback is + // unavailable (the export is rejected honestly instead of unbounded). + let percent_no_limit = QueryResultExportRequest { + sql: "SELECT TOP 10 PERCENT * FROM orders".to_string(), + query_base_sql: "SELECT TOP 10 PERCENT * FROM orders".to_string(), + database_type: DatabaseType::Kingbase, + use_agent_cursor: false, + ..request("csv", None, None) + }; + assert!(!supports_single_execution_export(&percent_no_limit, 100)); + assert_eq!(single_execution_row_bound(&percent_no_limit), None); + + let ties_no_limit = QueryResultExportRequest { + sql: "SELECT TOP (2) WITH TIES * FROM orders".to_string(), + query_base_sql: "SELECT TOP (2) WITH TIES * FROM orders".to_string(), + database_type: DatabaseType::Kingbase, + use_agent_cursor: false, + ..request("csv", None, None) + }; + assert!(!supports_single_execution_export(&ties_no_limit, 100)); + assert_eq!(single_execution_row_bound(&ties_no_limit), None); + + // With a configured row limit the same queries are capped by that limit. + let percent_with_limit = QueryResultExportRequest { row_limit: Some(5000), ..percent_no_limit }; + assert_eq!(single_execution_row_bound(&percent_with_limit), Some(5000)); + assert!(!supports_single_execution_export(&percent_with_limit, 100)); + } + + #[test] + fn kingbase_top_expression_export_requires_row_limit_or_cursor() { + // P1: TOP (100 + 1) returns 101 rows, so its bound must not be treated as + // 100 (which would silently truncate the export). Without a row limit the + // single-execution fallback is unavailable and the export is rejected. + let req = QueryResultExportRequest { + sql: "SELECT TOP (100 + 1) * FROM orders".to_string(), + query_base_sql: "SELECT TOP (100 + 1) * FROM orders".to_string(), + database_type: DatabaseType::Kingbase, + use_agent_cursor: false, + ..request("csv", None, None) + }; + assert!(!supports_single_execution_export(&req, 100)); + assert_eq!(single_execution_row_bound(&req), None); + + // A configured row limit gives the export an explicit cap. + let with_limit = QueryResultExportRequest { row_limit: Some(200), ..req }; + assert_eq!(single_execution_row_bound(&with_limit), Some(200)); + assert!(!supports_single_execution_export(&with_limit, 100)); + } + + #[test] + fn kingbase_without_top_uses_streaming_offset_pagination() { + let req = QueryResultExportRequest { + sql: "SELECT * FROM orders o JOIN customers c ON c.id = o.customer_id ORDER BY o.id".to_string(), + query_base_sql: "SELECT * FROM orders o JOIN customers c ON c.id = o.customer_id ORDER BY o.id".to_string(), + database_type: DatabaseType::Kingbase, + use_agent_cursor: false, + ..request("csv", Some(1000), None) + }; + + assert!(supports_streaming_offset_pagination(&req, 100)); + assert!(!supports_single_execution_export(&req, 100)); + } + + #[test] + fn kingbase_single_execution_never_exceeds_one_export_page() { + let req = QueryResultExportRequest { + sql: "SELECT TOP 1000 * FROM orders".to_string(), + query_base_sql: "SELECT TOP 1000 * FROM orders".to_string(), + database_type: DatabaseType::Kingbase, + use_agent_cursor: false, + ..request("csv", None, None) + }; + + assert_eq!(single_execution_row_bound(&req), Some(1000)); + assert_eq!(single_execution_page_limit(&req, 100), None); + assert!(supports_single_execution_export(&req, 1000)); + } + #[test] fn clickhouse_scalar_with_query_supports_streaming_pagination() { let sql = "WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept"; diff --git a/crates/dbx-core/src/query_result_sql.rs b/crates/dbx-core/src/query_result_sql.rs index b2e32f64e5..aea3a09c91 100644 --- a/crates/dbx-core/src/query_result_sql.rs +++ b/crates/dbx-core/src/query_result_sql.rs @@ -57,6 +57,12 @@ pub struct QueryPaginationExecutionPlan { #[serde(skip_serializing_if = "Option::is_none")] pub count_sql: Option, pub use_agent_result_session: bool, + /// True when the statement cannot be paginated server-side and must be + /// executed once with the whole result streamed back (single execution). + /// Only meaningful to in-process callers (query-result export); never + /// serialized to the frontend. + #[serde(skip)] + pub single_execution: bool, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -116,6 +122,7 @@ pub fn build_query_pagination_execution_plan( page_offset: None, count_sql: None, use_agent_result_session: false, + single_execution: false, }; let counted = build_count_query_sql(CountQuerySqlOptions { @@ -171,6 +178,17 @@ pub fn build_query_pagination_execution_plan( plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); plan.use_agent_result_session = true; + } else if options.database_type == Some(DatabaseType::Kingbase) + && single_selectable_statement(&options.sql).is_ok() + && has_top_level_top(&options.sql) + { + // Kingbase SQL Server compatibility mode rejects a statement that mixes a + // top-level TOP with a sibling LIMIT/OFFSET. Without an Agent cursor the + // query-result export executes the statement once and streams the whole + // result; the TOP clause already bounds the row count. + plan.page_limit = Some(options.pagination.limit); + plan.page_offset = Some(options.pagination.offset); + plan.single_execution = true; } plan } @@ -212,6 +230,13 @@ pub fn build_paginated_query_sql(options: PaginatedQuerySqlOptions) -> QuerySqlB TablePaginationStrategy::AgentMaxRows | TablePaginationStrategy::Unbounded => ok(format!("{statement};")), TablePaginationStrategy::IrisTop => ok(add_iris_top_limit(&statement, safe_limit)), TablePaginationStrategy::LimitOffset => { + // Kingbase SQL Server compatibility mode accepts TOP as a real clause. + // Appending LIMIT/OFFSET alongside a top-level TOP would be rejected by + // the server ("multiple TOP/LIMIT clauses not allowed"), so fall back to + // the Agent cursor / client-side row cap for such statements. + if options.database_type == Some(DatabaseType::Kingbase) && has_top_level_top(&statement) { + return err("unsupported"); + } let dedup_count = dedup_projection_count_without_order_by(&options.original_sql); ok(add_standard_limit(&statement, options.database_type, safe_limit, safe_offset, dedup_count)) } @@ -1081,6 +1106,60 @@ fn has_top_level_limit(sql: &str) -> bool { top_level_sql_tokens(sql).iter().any(|token| token.text == "LIMIT") } +/// True when the statement has a top-level `TOP` clause (SQL Server dialect). +/// Kingbase's SQL Server compatibility mode treats TOP as a real clause, so a +/// statement that already bounds rows with TOP must not receive a sibling LIMIT. +pub(crate) fn has_top_level_top(sql: &str) -> bool { + top_level_sql_tokens(sql).iter().any(|token| token.text == "TOP") +} + +/// Concrete row-count bound of a top-level `TOP` clause when written as a +/// literal (`TOP n`, `TOP(n)`, `TOP (n)`). Returns `None` for percentage TOP +/// (`TOP n PERCENT`), `WITH TIES` (the server may return more than `n` rows), +/// parenthesized expressions (`TOP (100 + 1)`, `TOP (100 * 2)`), or when the +/// TOP clause has no literal at all. A parenthesized form is only accepted when +/// it is exactly one integer literal followed by `)`, so the returned bound is +/// always exact — never a silent under-count. +pub(crate) fn top_level_top_row_count(sql: &str) -> Option { + let tokens = top_level_sql_tokens(sql); + let top_index = tokens.iter().position(|token| token.text == "TOP")?; + let top_token = &tokens[top_index]; + // A modifier keyword directly after TOP (ALL / DISTINCT) means the following + // literal is not a plain row-count bound. PERCENT / WITH TIES come after the + // literal and are handled by the check below. + if tokens.get(top_index + 1).is_some_and(|token| matches!(token.text.as_str(), "ALL" | "DISTINCT")) { + return None; + } + let mut cursor = skip_sql_whitespace(sql, top_token.start + top_token.text.len()); + let parenthesized = sql.get(cursor..)?.starts_with('('); + if parenthesized { + cursor = skip_sql_whitespace(sql, cursor + 1); + } + let count = parse_usize_literal(sql, &mut cursor)?; + let after = skip_sql_whitespace(sql, cursor); + if parenthesized { + // The parenthesized form must be exactly one integer literal followed by + // `)`. Anything else is an expression whose real bound we cannot know + // (e.g. TOP (100 + 1) returns 101 rows, not 100), so refuse to treat it + // as a bound. + if !sql.get(after..)?.starts_with(')') { + return None; + } + } + // `TOP n PERCENT` and `TOP n WITH TIES` (parenthesized or not) do not bound + // the row count to the literal, so reject those adjacent modifiers. A later + // table hint such as `FROM events WITH (NOLOCK)` is unrelated to TOP. + let after_paren = if parenthesized { skip_sql_whitespace(sql, after + 1) } else { after }; + let modifier_index = tokens.iter().position(|token| token.start >= after_paren); + if modifier_index.is_some_and(|index| { + tokens[index].text == "PERCENT" + || (tokens[index].text == "WITH" && tokens.get(index + 1).is_some_and(|token| token.text == "TIES")) + }) { + return None; + } + Some(count) +} + fn top_level_limit_row_count(sql: &str) -> Option { let token = top_level_sql_tokens(sql).into_iter().find(|token| token.text == "LIMIT")?; parse_standard_limit_row_count(sql, token.start + token.text.len()) @@ -3106,6 +3185,157 @@ WHERE u.id = picked.id; assert!(plan.use_agent_result_session); } + #[test] + fn kingbase_top_clause_falls_back_to_agent_cursor() { + for sql in [ + "SELECT TOP 100 * FROM events", + "SELECT TOP(100) * FROM events", + "SELECT TOP (100) * FROM events", + "SELECT TOP 10 PERCENT * FROM events", + "SELECT TOP (2) WITH TIES * FROM events", + "SELECT TOP 100 events.name FROM events ORDER BY events.name", + ] { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: sql.to_string(), + query_base_sql: sql.to_string(), + database_type: Some(DatabaseType::Kingbase), + pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert_eq!(plan.sql_to_execute, sql, "sql_to_execute should stay untouched for {sql}"); + assert!(plan.page_sql.is_none(), "no LIMIT rewrite for {sql}"); + assert_eq!(plan.page_limit, Some(500)); + assert_eq!(plan.page_offset, Some(0)); + assert!(plan.use_agent_result_session, "Agent cursor fallback for {sql}"); + } + } + + #[test] + fn kingbase_subquery_top_still_uses_server_pagination() { + // A TOP inside a derived table is not a top-level clause, so the outer + // statement can still be paginated with LIMIT/OFFSET. + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT * FROM (SELECT TOP 100 * FROM events) t".to_string(), + query_base_sql: "SELECT * FROM (SELECT TOP 100 * FROM events) t".to_string(), + database_type: Some(DatabaseType::Kingbase), + pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert_eq!(plan.sql_to_execute, "SELECT * FROM (SELECT TOP 100 * FROM events) t LIMIT 500;"); + assert_eq!(plan.page_sql, Some(plan.sql_to_execute.clone())); + assert_eq!(plan.page_limit, Some(500)); + assert_eq!(plan.page_offset, Some(0)); + assert!(!plan.use_agent_result_session); + } + + #[test] + fn kingbase_top_clause_without_agent_uses_single_execution() { + // Non-agent callers (query-result export with use_agent_cursor=false) + // cannot open an Agent result session, so the Kingbase-TOP plan must + // mark the query single-execution instead: original SQL unchanged, no + // page_sql, but bounded page limits the export loop can stream once. + for sql in [ + "SELECT TOP 100 * FROM events", + "SELECT TOP(100) * FROM events", + "SELECT TOP 10 PERCENT * FROM events", + "SELECT TOP (2) WITH TIES * FROM events", + "SELECT TOP 100 events.name FROM events JOIN users ON users.id = events.id", + ] { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: sql.to_string(), + query_base_sql: sql.to_string(), + database_type: Some(DatabaseType::Kingbase), + pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, + use_agent_cursor: false, + first_page_uses_actual_sql: true, + }); + + assert_eq!(plan.sql_to_execute, sql, "sql_to_execute should stay untouched for {sql}"); + assert!(plan.page_sql.is_none(), "no LIMIT rewrite for {sql}"); + assert_eq!(plan.page_limit, Some(500)); + assert_eq!(plan.page_offset, Some(0)); + assert!(!plan.use_agent_result_session); + assert!(plan.single_execution, "single-execution fallback for {sql}"); + } + } + + #[test] + fn kingbase_top_with_agent_still_uses_agent_session() { + // Grid/query path: with an Agent cursor available the Kingbase-TOP query + // keeps the existing bounded Agent-session fallback (not single-execution). + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT TOP 100 * FROM events".to_string(), + query_base_sql: "SELECT TOP 100 * FROM events".to_string(), + database_type: Some(DatabaseType::Kingbase), + pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert_eq!(plan.sql_to_execute, "SELECT TOP 100 * FROM events"); + assert!(plan.page_sql.is_none()); + assert_eq!(plan.page_limit, Some(500)); + assert_eq!(plan.page_offset, Some(0)); + assert!(plan.use_agent_result_session); + assert!(!plan.single_execution); + } + + #[test] + fn top_level_top_row_count_extracts_only_concrete_bounds() { + assert_eq!(top_level_top_row_count("SELECT TOP 100 * FROM events"), Some(100)); + assert_eq!(top_level_top_row_count("SELECT TOP(100) * FROM events"), Some(100)); + assert_eq!(top_level_top_row_count("SELECT TOP (100) * FROM events"), Some(100)); + // Whitespace-only inside the parens is still a single literal bound. + assert_eq!(top_level_top_row_count("SELECT TOP ( 100 ) * FROM events"), Some(100)); + assert_eq!(top_level_top_row_count("SELECT TOP 100 events.name FROM events ORDER BY events.name"), Some(100)); + assert_eq!(top_level_top_row_count("SELECT TOP 100 * FROM events LIMIT 5"), Some(100)); + + // Parenthesized expressions have a real bound different from the leading + // digits; refusing them (None) beats silently under-counting the export. + assert_eq!(top_level_top_row_count("SELECT TOP (100 + 1) * FROM events"), None); + assert_eq!(top_level_top_row_count("SELECT TOP (100 * 2) * FROM events"), None); + assert_eq!(top_level_top_row_count("SELECT TOP (100 - 1) * FROM events"), None); + assert_eq!(top_level_top_row_count("SELECT TOP (len(events.name)) * FROM events"), None); + + // Percentage TOP and WITH TIES are not concrete row-count bounds, with or + // without parentheses. + assert_eq!(top_level_top_row_count("SELECT TOP 10 PERCENT * FROM events"), None); + assert_eq!(top_level_top_row_count("SELECT TOP (10) PERCENT * FROM events"), None); + assert_eq!(top_level_top_row_count("SELECT TOP (2) WITH TIES * FROM events"), None); + assert_eq!(top_level_top_row_count("SELECT TOP 2 WITH TIES * FROM events"), None); + assert_eq!(top_level_top_row_count("SELECT TOP 2 * FROM events WITH (NOLOCK)"), Some(2)); + + // No TOP clause at all. + assert_eq!(top_level_top_row_count("SELECT * FROM events"), None); + // TOP only inside a subquery is not a top-level clause. + assert_eq!(top_level_top_row_count("SELECT * FROM (SELECT TOP 5 * FROM events) t"), None); + } + + #[test] + fn non_kingbase_limit_offset_is_unaffected_by_top_keyword() { + // A column literally named `top` must not be mistaken for a TOP clause + // on dialects where TOP is not a clause keyword. Force the server-side + // LimitOffset path (no Agent cursor) to exercise the rewrite. + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT top, name FROM users".to_string(), + query_base_sql: "SELECT top, name FROM users".to_string(), + database_type: Some(DatabaseType::Postgres), + pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, + use_agent_cursor: false, + first_page_uses_actual_sql: false, + }); + + assert_eq!(plan.sql_to_execute, "SELECT top, name FROM users LIMIT 500;"); + assert_eq!(plan.page_sql, Some(plan.sql_to_execute.clone())); + assert_eq!(plan.page_limit, Some(500)); + assert_eq!(plan.page_offset, Some(0)); + assert!(!plan.use_agent_result_session); + } + #[test] fn export_agent_cursor_first_page_keeps_actual_sorted_sql() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions {