use std::collections::HashSet; use std::ops::ControlFlow; use serde::{Deserialize, Serialize}; use crate::models::connection::DatabaseType; use crate::sql::{find_statement_at_cursor, find_statement_at_cursor_for_database}; use crate::sql_dialect::{ firebird_rows_clause, pagination_strategy, quote_iris_identifier, quote_table_identifier, PaginationContext, TablePaginationStrategy, }; use sqlparser::ast::{ visit_expressions, Expr, GroupByExpr, LimitClause, ObjectNamePart, OrderByKind, Select, SelectItem, SelectModifiers, SetExpr, Statement, TableFactor, Value, ValueWithSpan, }; use sqlparser::dialect::{ClickHouseDialect, GenericDialect, MsSqlDialect, MySqlDialect}; use sqlparser::parser::Parser; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct QuerySqlBuildResult { pub ok: bool, #[serde(skip_serializing_if = "Option::is_none")] pub sql: Option, #[serde(skip_serializing_if = "Option::is_none")] pub reason: Option, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct QueryPagination { pub limit: usize, pub offset: usize, #[serde(default, skip_serializing_if = "Option::is_none")] pub session_id: Option, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct QueryPaginationExecutionPlanOptions { pub sql: String, pub query_base_sql: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub database_type: Option, pub pagination: QueryPagination, pub use_agent_cursor: bool, #[serde(default)] pub first_page_uses_actual_sql: bool, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct QueryPaginationExecutionPlan { pub sql_to_execute: String, #[serde(skip_serializing_if = "Option::is_none")] pub page_sql: Option, #[serde(skip_serializing_if = "Option::is_none")] pub page_limit: Option, #[serde(skip_serializing_if = "Option::is_none")] pub page_offset: Option, #[serde(skip_serializing_if = "Option::is_none")] pub count_sql: Option, #[serde(skip_serializing_if = "Option::is_none")] pub exact_query_row_bound: Option, pub use_agent_result_session: bool, /// Trailing helper column added by DBX's ROWNUM pagination wrapper. #[serde(default, skip_serializing_if = "Option::is_none")] pub pagination_row_number_column: Option, /// 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)] #[serde(rename_all = "camelCase")] pub struct PaginatedQuerySqlOptions { pub original_sql: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub database_type: Option, pub limit: usize, pub offset: usize, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct CountQuerySqlOptions { pub original_sql: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub database_type: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum QuerySortDirection { Asc, Desc, } impl QuerySortDirection { fn as_sql(self) -> &'static str { match self { Self::Asc => "ASC", Self::Desc => "DESC", } } } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct SortedQuerySqlOptions { pub original_sql: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub database_type: Option, #[serde(default)] pub result_columns: Vec, pub column_index: usize, pub column: String, pub direction: QuerySortDirection, } pub fn build_query_pagination_execution_plan( options: QueryPaginationExecutionPlanOptions, ) -> QueryPaginationExecutionPlan { // Every page DBX generates for this query is derived from the same // user-written statement, so a literal LIMIT/TOP the user already wrote // bounds the query's total row count regardless of how large the // underlying table is — a cheap, exact upper bound that needs no // COUNT(*) execution. SQL Server's TOP is covered separately from the // standard LIMIT/OFFSET dialects (MySQL, Postgres, etc.) since they use // different clause syntax. let exact_query_row_bound = match pagination_strategy(options.database_type, PaginationContext::UserQuery) { TablePaginationStrategy::SqlServerTop => { // TOP bounds input rows, while FOR JSON/XML emit chunked output // rows — a TOP bound would misreport the result total. if has_top_level_for_output_clause(&options.query_base_sql) { None } else { top_level_top_row_count(&options.query_base_sql) } } TablePaginationStrategy::LimitOffset => top_level_limit_row_count(&options.query_base_sql), _ => None, }; let mut plan = QueryPaginationExecutionPlan { sql_to_execute: options.sql.clone(), page_sql: None, page_limit: None, page_offset: None, count_sql: None, exact_query_row_bound, use_agent_result_session: false, pagination_row_number_column: None, single_execution: false, }; let sql_server_cte = options.database_type == Some(DatabaseType::SqlServer) && starts_with_cte(&options.query_base_sql); if !sql_server_cte { let counted = build_count_query_sql(CountQuerySqlOptions { original_sql: options.query_base_sql.clone(), database_type: options.database_type, }); if counted.ok { plan.count_sql = counted.sql; } } if options.pagination.session_id.is_some() { plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); plan.use_agent_result_session = true; return plan; } if sql_server_cte { return plan; } // Existing top-level ROWNUM predicates/projections cannot be wrapped without // changing their semantics. Keep one cursor so later pages advance it. if options.database_type == Some(DatabaseType::OceanbaseOracle) && options.use_agent_cursor && has_top_level_rownum(&options.sql) { plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); plan.use_agent_result_session = true; return plan; } let can_use_first_page_cursor = options.use_agent_cursor && options.pagination.offset == 0; // HighGo, OceanBase Oracle, and Xugu can spend substantially more time // executing an unbounded query before the Agent exposes its first cursor // page. Prefer a bounded SQL query whenever it can be rewritten safely. // Independent pages of an unordered query do not have a stable row order; // callers should add ORDER BY when that matters. // Kingbase keeps the cursor for unordered queries to preserve its behavior. let prefer_server_pagination = match options.database_type { Some(DatabaseType::Highgo | DatabaseType::OceanbaseOracle | DatabaseType::Xugu) => true, Some(DatabaseType::Kingbase) => kingbase_server_pagination_is_stable(&options.query_base_sql), _ => false, }; if can_use_first_page_cursor && !prefer_server_pagination { if !options.first_page_uses_actual_sql && options.sql == options.query_base_sql { plan.sql_to_execute = options.query_base_sql; } plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); plan.use_agent_result_session = true; return plan; } let paginated = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: options.sql.clone(), database_type: options.database_type, limit: options.pagination.limit, offset: options.pagination.offset, }); if paginated.ok { plan.sql_to_execute = paginated.sql.clone().unwrap_or_default(); plan.page_sql = paginated.sql; plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); if options.pagination.offset > 0 && pagination_strategy(options.database_type, PaginationContext::UserQuery) == TablePaginationStrategy::Rownum && !has_top_level_rownum(&options.sql) { plan.pagination_row_number_column = single_selectable_statement(&options.sql, options.database_type) .ok() .map(|statement| rownum_pagination_column_name(&statement)); } if options.use_agent_cursor && matches!( pagination_strategy(options.database_type, PaginationContext::UserQuery), TablePaginationStrategy::AgentMaxRows | TablePaginationStrategy::Unbounded ) { // The dialect cannot rewrite the statement with server-side // pagination (Oracle user queries, generic JDBC), so the page // metadata alone cannot reach a non-zero offset. Keep the Agent // result session enabled and let the client consume through to // the requested offset (#8993). plan.use_agent_result_session = true; } } else if can_use_first_page_cursor && options.database_type != Some(DatabaseType::Highgo) { // Kingbase JDBC may buffer an entire result in auto-commit mode, so use // LIMIT/OFFSET whenever the statement can be rewritten safely. Keep the // Agent cursor as a bounded fallback for multi-statement or dialect- // specific SQL that the pagination parser cannot transform. HighGo does // not take this fallback: its JDBC driver may materialize the unbounded // result before cursor paging starts. Leaving the page metadata unset // routes it through regular execution and the configured JDBC maxRows // safeguard instead. if !options.first_page_uses_actual_sql && options.sql == options.query_base_sql { plan.sql_to_execute = options.query_base_sql; } plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); plan.use_agent_result_session = true; } else if matches!(options.database_type, Some(DatabaseType::Kingbase | DatabaseType::Xugu)) && single_selectable_statement(&options.sql, options.database_type).is_ok() && has_top_level_top(&options.sql) { // Kingbase SQL Server compatibility mode and Xugu reject statements mixing 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 } pub fn build_paginated_query_sql(options: PaginatedQuerySqlOptions) -> QuerySqlBuildResult { let Ok(statement) = single_selectable_statement(&options.original_sql, options.database_type) else { return err(single_statement_error_reason(&options.original_sql)); }; if unsupported_pagination_type(options.database_type) { return err("unsupported"); } let safe_limit = options.limit.max(1); let safe_offset = options.offset; if matches!(options.database_type, Some(DatabaseType::Elasticsearch | DatabaseType::Easysearch)) { // If the user wrote their own LIMIT, leave the SQL alone — they // explicitly bounded the result set and the front-end will paginate // client-side. Otherwise wrap with an explicit OFFSET (even when // 0) so the ES driver can tell a plan-wrapped query from one the // user wrote, which decides whether affected_rows should reflect // the index total or the row count we actually returned. if has_top_level_limit(&statement) { return err("unsupported"); } return ok(format!("{statement} LIMIT {safe_limit} OFFSET {safe_offset};")); } match pagination_strategy(options.database_type, PaginationContext::UserQuery) { TablePaginationStrategy::SqlServerTop => add_sql_server_offset_fetch(&statement, safe_limit, safe_offset) .map(ok) .unwrap_or_else(|| err("unsupported")), TablePaginationStrategy::QuestDbLimit => ok(add_questdb_limit(&statement, safe_limit, safe_offset)), TablePaginationStrategy::InformixFirst => ok(add_informix_first_limit(&statement, safe_limit, safe_offset)), TablePaginationStrategy::FirebirdRows => ok(add_firebird_rows_limit(&statement, safe_limit, safe_offset)), TablePaginationStrategy::Db2FetchFirst | TablePaginationStrategy::FetchFirst => { ok(add_fetch_first_limit(&statement, safe_limit, safe_offset)) } TablePaginationStrategy::Rownum => ok(add_rownum_limit(&statement, safe_limit, safe_offset)), TablePaginationStrategy::AgentMaxRows | TablePaginationStrategy::Unbounded => ok(format!("{statement};")), TablePaginationStrategy::IrisTop => ok(add_iris_top_limit(&statement, safe_limit)), TablePaginationStrategy::LimitOffset => { // Kingbase SQL Server compatibility mode and Xugu accept 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 matches!(options.database_type, Some(DatabaseType::Kingbase | DatabaseType::Xugu)) && has_top_level_top(&statement) { return err("unsupported"); } let dedup_order_by = dedup_projection_count_without_order_by(&options.original_sql); ok(add_standard_limit(&statement, options.database_type, safe_limit, safe_offset, dedup_order_by)) } } } pub fn build_count_query_sql(options: CountQuerySqlOptions) -> QuerySqlBuildResult { let Ok(statement) = single_selectable_statement(&options.original_sql, options.database_type) else { return err(single_statement_error_reason(&options.original_sql)); }; if unsupported_pagination_type(options.database_type) { return err("unsupported"); } let (execution_hint, statement) = split_leading_execution_hint(&statement, options.database_type); // A locking clause does not affect cardinality and cannot appear inside // every dialect's derived-table count query. PostgreSQL permits pagination // after the lock clause; decline counting that uncommon order rather than // accidentally dropping the user's explicit LIMIT/OFFSET. let tokens = top_level_sql_tokens(statement); let statement = if let Some(index) = locking_clause_index(&tokens) { if has_pagination_clause_after(&tokens, index) { return err("locking"); } statement[..index].trim_end().to_string() } else { statement.to_string() }; // ES SQL can't wrap a SELECT in `SELECT COUNT(*) FROM (...)` — the // driver already reports the true match count via affected_rows. if matches!(options.database_type, Some(DatabaseType::Elasticsearch | DatabaseType::Easysearch)) { return err("unsupported"); } if options.database_type == Some(DatabaseType::SqlServer) { return sql_server_count_sql(&statement) .map(|sql| ok(format!("{execution_hint}{sql}"))) .unwrap_or_else(|| err("unsupported")); } if options.database_type == Some(DatabaseType::Mysql) { return mysql_count_sql(&statement) .map(|sql| ok(format!("{execution_hint}{sql}"))) .unwrap_or_else(|| err("unsupported")); } if options.database_type == Some(DatabaseType::Iotdb) { // IoTDB Tree Model does not support derived tables. Count the selected // time series directly when doing so is guaranteed to preserve the // result cardinality; decline complex queries instead of issuing SQL // that IoTDB cannot parse or returning a misleading total. return iotdb_tree_count_sql(&statement) .map(|sql| ok(format!("{execution_hint}{sql}"))) .unwrap_or_else(|| err("unsupported")); } if options.database_type == Some(DatabaseType::Iris) { // IRIS JDBC can parameterize a derived-table alias as `:%qpar` during // prepare, even when the alias is unquoted. A simple single-table // query does not need a derived table to count rows, so inject COUNT // directly and leave the wrapper fallback for queries whose row // cardinality depends on the original projection. if let Some(sql) = iris_count_sql(&statement) { return ok(format!("{execution_hint}{sql}")); } } let alias = if options.database_type == Some(DatabaseType::Iris) { // With delimited identifiers disabled, Caché 2016 can parameterize a // double-quoted alias as a string literal, leaving a trailing :%qpar. "dbx_count".to_string() } else { quote_table_identifier(options.database_type, "dbx_count") }; if options.database_type == Some(DatabaseType::Hive) && starts_with_cte(&statement) { if let Some(main_query_start) = top_level_sql_tokens(&statement) .into_iter() .find(|token| matches!(token.text.as_str(), "SELECT" | "FROM")) .map(|token| token.start) { let (with_clause, main_query) = statement.split_at(main_query_start); return ok(format!( "{execution_hint}{with_clause}{}", derived_table_sql("SELECT COUNT(*) AS dbx_total_rows FROM", main_query, &format!("{alias};")) )); } } let wrapped_sql = match options.database_type { Some(DatabaseType::Iris) => iris_statement_for_derived_table(&statement), _ => statement, }; ok(format!( "{execution_hint}{}", derived_table_sql("SELECT COUNT(*) AS dbx_total_rows FROM", &wrapped_sql, &format!("{alias};")) )) } pub fn build_sorted_query_sql(options: SortedQuerySqlOptions) -> QuerySqlBuildResult { let base_sql = options.original_sql.trim(); if base_sql.is_empty() { return err("empty"); } let mut statement = find_query_result_statement_at_cursor(base_sql, 0, options.database_type) .trim() .trim_end_matches(';') .trim() .to_string(); if statement.is_empty() { return err("empty"); } let normalized_base_len = base_sql.trim_end_matches(';').trim().len(); if statement.len() != normalized_base_len { if options.database_type != Some(DatabaseType::Mysql) { return err("multi"); } statement = restore_leading_execution_hint(base_sql, &statement, options.database_type); if statement.len() != normalized_base_len { return err("multi"); } } let (execution_hint, statement) = if options.database_type == Some(DatabaseType::Mysql) { split_leading_execution_hint(&statement, options.database_type) } else { ("", statement.as_str()) }; if statement.trim_start().to_ascii_uppercase().starts_with("WITH") { return err("with"); } if !statement.trim_start().to_ascii_uppercase().starts_with("SELECT") { return err("not_select"); } let uses_hive_subquery_syntax = matches!(options.database_type, Some(DatabaseType::Argo | DatabaseType::Hive | DatabaseType::Impala)); if uses_hive_subquery_syntax { let mut column_names = HashSet::new(); if options.result_columns.iter().any(|column| !column_names.insert(column.to_lowercase())) { return err("unsupported"); } } let aliases = build_derived_column_aliases(&options.result_columns); // Caché/IRIS rejects derived-table column alias lists (`t(col, col)`) // outright (SQLCODE -25), regardless of delimited-identifier support. let use_derived_column_aliases = !uses_hive_subquery_syntax && options.database_type != Some(DatabaseType::Mysql) && options.database_type != Some(DatabaseType::ClickHouse) // Doris accepts the derived-table alias but not its column-name list. && options.database_type != Some(DatabaseType::Doris) && options.database_type != Some(DatabaseType::Sqlite) && options.database_type != Some(DatabaseType::DuckDb) && options.database_type != Some(DatabaseType::Dameng) && options.database_type != Some(DatabaseType::Oracle) && options.database_type != Some(DatabaseType::OceanbaseOracle) && options.database_type != Some(DatabaseType::SapHana) && options.database_type != Some(DatabaseType::Iris); let sort_alias = if use_derived_column_aliases { aliases .get(options.column_index) .or_else(|| { options .result_columns .iter() .position(|column| column == &options.column) .and_then(|index| aliases.get(index)) }) .cloned() .unwrap_or_else(|| fallback_alias(options.column_index)) } else { options.result_columns.get(options.column_index).cloned().unwrap_or_else(|| options.column.clone()) }; // Oracle-compatible derived tables do not accept a PostgreSQL-style // column alias list. Use the selected column position when duplicate // labels would otherwise make ORDER BY ambiguous. let use_sort_ordinal = !use_derived_column_aliases && matches!( options.database_type, Some( DatabaseType::Dameng | DatabaseType::Oracle | DatabaseType::OceanbaseOracle | DatabaseType::SapHana | DatabaseType::Iris ) ) && options.result_columns.get(options.column_index).is_some_and(|column| { options.result_columns.iter().filter(|candidate| candidate.eq_ignore_ascii_case(column)).count() > 1 }); let sort_reference = if use_sort_ordinal { (options.column_index + 1).to_string() } else if options.database_type == Some(DatabaseType::Iris) { // With delimited identifiers disabled, a quoted ORDER BY name becomes // a string literal on Caché and the sort silently degrades to a // constant. Ordinary names must be sent unquoted. quote_iris_identifier(&sort_alias, None) } else { quote_table_identifier(options.database_type, &sort_alias) }; let wrapped_statement = if options.database_type == Some(DatabaseType::SqlServer) { sql_server_statement_for_derived_table(statement) } else { statement.to_string() }; if use_derived_column_aliases { let alias_list = aliases .iter() .map(|alias| quote_table_identifier(options.database_type, alias)) .collect::>() .join(", "); ok(format!( "{execution_hint}SELECT * FROM ({wrapped_statement}) t({alias_list}) ORDER BY {sort_reference} {};", options.direction.as_sql() )) } else { ok(format!( "{execution_hint}SELECT * FROM ({wrapped_statement}) t ORDER BY {sort_reference} {};", options.direction.as_sql() )) } } fn ok(sql: String) -> QuerySqlBuildResult { QuerySqlBuildResult { ok: true, sql: Some(sql), reason: None } } fn err(reason: &str) -> QuerySqlBuildResult { QuerySqlBuildResult { ok: false, sql: None, reason: Some(reason.to_string()) } } fn unsupported_pagination_type(database_type: Option) -> bool { matches!( database_type, Some(DatabaseType::Neo4j | DatabaseType::MongoDb | DatabaseType::Redis | DatabaseType::Solr) ) } fn find_query_result_statement_at_cursor(sql: &str, cursor_pos: usize, database_type: Option) -> String { if database_type == Some(DatabaseType::Mysql) { find_statement_at_cursor_for_database(sql, cursor_pos, DatabaseType::Mysql) } else { find_statement_at_cursor(sql, cursor_pos) } } fn kingbase_server_pagination_is_stable(sql: &str) -> bool { let has_order_by = Parser::parse_sql(&GenericDialect {}, sql) .ok() .and_then(|statements| { let [Statement::Query(query)] = statements.as_slice() else { return None; }; Some(query.order_by.as_ref().is_some_and( |order_by| !matches!(&order_by.kind, OrderByKind::Expressions(expressions) if expressions.is_empty()), )) }) .unwrap_or(false); has_order_by || dedup_projection_count_without_order_by(sql).is_some() } fn single_selectable_statement(original_sql: &str, database_type: Option) -> Result { let base_sql = original_sql.trim(); if base_sql.is_empty() { return Err(()); } let extracted = find_query_result_statement_at_cursor(base_sql, 0, database_type) .trim() .trim_end_matches(';') .trim() .to_string(); if extracted.is_empty() { return Err(()); } if !single_statement_matches_base_sql(&extracted, base_sql) { return Err(()); } let statement = restore_leading_execution_hint(base_sql, &extracted, database_type); let statement_without_leading_comments = strip_leading_statement_comments(statement.trim_start_matches(';').trim_start()); let upper = statement_without_leading_comments.to_ascii_uppercase(); if upper.starts_with("WITH") { if !cte_main_statement_is_select(&statement) { return Err(()); } } else if !upper.starts_with("SELECT") { return Err(()); } if has_top_level_select_into(&statement) { return Err(()); } Ok(statement) } fn single_statement_matches_base_sql(statement: &str, base_sql: &str) -> bool { let normalized_statement = statement.trim().trim_end_matches(';').trim(); let normalized_base = base_sql.trim().trim_end_matches(';').trim(); if normalized_statement.len() == normalized_base.len() { return true; } let base_without_leading_comments = strip_leading_statement_comments(normalized_base).trim().trim_end_matches(';').trim(); normalized_statement == base_without_leading_comments } fn starts_with_cte(sql: &str) -> bool { let mut statement = sql; loop { let previous_len = statement.len(); statement = strip_leading_statement_comments(statement); statement = statement.trim_start_matches(';').trim_start(); if statement.len() == previous_len { break; } } sql_keyword_at(statement, 0, "WITH") } fn cte_main_statement_is_select(sql: &str) -> bool { let tokens = top_level_sql_tokens(sql); let mut index = match tokens.iter().position(|token| token.text == "WITH") { Some(index) => index + 1, None => return false, }; if tokens.get(index).is_some_and(|token| token.text == "RECURSIVE") { index += 1; } while let Some(token) = tokens.get(index) { if is_with_main_statement_keyword(&token.text) { return token.text == "SELECT"; } index += 1; } false } fn is_with_main_statement_keyword(token: &str) -> bool { matches!(token, "SELECT" | "INSERT" | "UPDATE" | "DELETE" | "MERGE") } fn single_statement_error_reason(original_sql: &str) -> &'static str { let base_sql = original_sql.trim(); if base_sql.is_empty() { return "empty"; } let statement = find_statement_at_cursor(base_sql, 0).trim().trim_end_matches(';').trim().to_string(); if statement.is_empty() { return "empty"; } if statement.len() != base_sql.trim_end_matches(';').trim().len() { return "multi"; } "not_select" } fn has_top_level_select_into(sql: &str) -> bool { let mut saw_select = false; for token in top_level_sql_tokens(sql) { if !saw_select { saw_select = token.text == "SELECT"; continue; } if token.text == "INTO" { return true; } } false } fn add_sql_server_offset_fetch(statement: &str, limit: usize, offset: usize) -> Option { // FOR JSON/XML produces an unnamed result value and cannot be projected // from a derived table used by DBX pagination. if has_top_level_for_output_clause(statement) { return None; } // 用户已写 OFFSET/FETCH 时必须原样保留,不能再注入 TOP(两者同块会被 SQL Server 拒绝)。 // 词法检测与 AST 检测任一命中即视为已有分页:词法扫描器在 # 临时表、 // 反斜杠字符串等场景会漏检,AST 检测负责把这些情况补上。 if has_top_level_offset_fetch_next(statement) || sql_server_ast_has_offset_or_fetch(statement) { return (offset == 0).then(|| statement.to_string()); } if has_top_level_select_top(statement) { return if sql_server_derived_table_projection_safe(statement) { Some(add_sql_server_existing_top_pagination(statement, limit, offset)) } else { (offset > 0).then(|| add_sql_server_rowcount_pagination(statement, limit, offset)) }; } // TOP injected into the first SELECT only bounds that branch of a UNION / // INTERSECT / EXCEPT, so the combined result comes out wrong. Bound the // whole statement instead. if has_top_level_set_operator(statement) { return Some(add_sql_server_rowcount_pagination(statement, limit, offset)); } 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(|| inject_sql_server_top(statement, limit)); } if offset == 0 { return Some(inject_sql_server_top(statement, limit)); } // The inner query may end with a line comment after removing ORDER BY. // Keep the derived-table closing parenthesis on a new line so it is not // swallowed by `--` / `#` comments. let statement_without_order = order_by_index .map(|index| statement_for_sql_suffix(statement[..index].trim_end())) .unwrap_or_else(|| statement_for_sql_suffix(statement)); if !sql_server_row_number_pagination_safe(statement) { return Some(add_sql_server_rowcount_pagination(statement, limit, offset)); } let row_number_order = order_by_index .map(|index| statement_for_sql_suffix(statement[index..].trim())) .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 Ok(statements) = Parser::parse_sql(&MsSqlDialect {}, 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; }; if !sql_server_derived_table_select_projection_safe(select) { return false; } let Some(order_by) = &query.order_by else { return true; }; let OrderByKind::Expressions(order_exprs) = &order_by.kind else { return false; }; let wildcard_projection = matches!(select.projection.as_slice(), [SelectItem::Wildcard(_)]); let output_names = select.projection.iter().filter_map(derived_projection_name).map(str::to_lowercase).collect::>(); order_exprs.iter().all(|order_expr| { if matches!(order_expr.expr, Expr::Value(_)) { return false; } !visit_expressions(&order_expr.expr, |expr| match expr { Expr::CompoundIdentifier(_) => ControlFlow::Break(()), Expr::Identifier(identifier) if !wildcard_projection && !output_names.contains(&identifier.value.to_lowercase()) => { ControlFlow::Break(()) } _ => ControlFlow::Continue(()), }) .is_break() }) } const SQLSERVER_RESULT_OFFSET_PREFIX: &str = "/*__dbx_result_offset="; const SQLSERVER_RESULT_OFFSET_SUFFIX: &str = "__*/"; fn add_sql_server_rowcount_pagination(statement: &str, limit: usize, offset: usize) -> String { let row_count = offset.saturating_add(limit); let escaped_statement = statement.replace('\'', "''"); // Keep duplicate result-column names intact while bounding the server response // on every SQL Server version supported by DBX. The dynamic batch scopes // SET ROWCOUNT to this execution instead of leaking it into the tab session. format!( "EXEC sys.sp_executesql N'SET ROWCOUNT {row_count}; {escaped_statement}'; {SQLSERVER_RESULT_OFFSET_PREFIX}{offset}{SQLSERVER_RESULT_OFFSET_SUFFIX}" ) } pub fn sqlserver_result_offset(sql: &str) -> usize { let sql = sql.trim_end(); if !sql.starts_with("EXEC sys.sp_executesql N'SET ROWCOUNT ") || !sql.ends_with(SQLSERVER_RESULT_OFFSET_SUFFIX) { return 0; } let Some(marker_index) = sql.rfind(SQLSERVER_RESULT_OFFSET_PREFIX) else { return 0; }; let value_start = marker_index + SQLSERVER_RESULT_OFFSET_PREFIX.len(); let value_end = sql.len() - SQLSERVER_RESULT_OFFSET_SUFFIX.len(); sql[value_start..value_end].parse().unwrap_or(0) } fn add_sql_server_existing_top_pagination(statement: &str, limit: usize, offset: usize) -> String { // Prefer the user's own trailing ORDER BY so wrapping the query in a // derived table does not silently override their requested ordering. // Fall back to the first projection column when it cannot be mapped to // the derived table's output columns (existing behavior). let row_number_order = sql_server_derived_pagination_order(statement) .unwrap_or_else(|| format!("ORDER BY {}", sql_server_default_pagination_order(statement))); if offset == 0 { let derived_statement = statement_for_sql_suffix(statement); return format!("SELECT TOP ({limit}) * FROM ({derived_statement}) [dbx_page] {row_number_order};"); } let end = offset + limit; let derived_statement = statement_for_sql_suffix(statement); format!( "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER ({row_number_order}) AS [__dbx_row_num] FROM ({derived_statement}) dbx_page_source) dbx_page WHERE [__dbx_row_num] > {offset} AND [__dbx_row_num] <= {end} ORDER BY [__dbx_row_num];" ) } /// Reuses the user's trailing ORDER BY as the pagination sort key when every /// expression can be expressed against the derived table's output columns. /// Returns None when the statement has no ORDER BY, cannot be parsed, or any /// sort expression is not an output column or ordinal, so callers keep their /// existing deterministic fallback ordering. fn sql_server_derived_pagination_order(statement: &str) -> Option { let dialect = MsSqlDialect {}; let Ok(statements) = Parser::parse_sql(&dialect, statement) else { return None; }; let [Statement::Query(query)] = statements.as_slice() else { return None; }; let SetExpr::Select(select) = query.body.as_ref() else { return None; }; let OrderByKind::Expressions(order_by_exprs) = &query.order_by.as_ref()?.kind else { return None; }; if order_by_exprs.is_empty() || !sql_server_derived_table_projection_safe(statement) { return None; } let output_columns = select.projection.iter().map(derived_projection_name).collect::>>()?; let column_names = output_columns.iter().map(|name| name.to_lowercase()).collect::>(); let mut parts = Vec::with_capacity(order_by_exprs.len()); for order_by in order_by_exprs { let column = sql_server_order_expr_output_column(&order_by.expr, &output_columns, &column_names)?; let direction = match order_by.options.asc { Some(true) => " ASC", Some(false) => " DESC", None => "", }; let mut part = column; part.push_str(direction); parts.push(part); } Some(format!("ORDER BY {}", parts.join(", "))) } /// Maps one ORDER BY expression to a reference that stays valid outside the /// derived table: output columns by name (case-insensitive, bracket/quoted /// identifiers already unquoted by the parser) or positional ordinals mapped /// to their corresponding output column. /// Anything else (functions, unprojected columns, ...) cannot be referenced /// from the wrapper and makes the caller fall back to its default ordering. fn sql_server_order_expr_output_column( expr: &Expr, output_columns: &[&str], column_names: &HashSet, ) -> Option { match expr { Expr::Identifier(identifier) => { let name = identifier.value.to_lowercase(); column_names .contains(&name) .then(|| quote_table_identifier(Some(DatabaseType::SqlServer), &identifier.value)) } Expr::CompoundIdentifier(identifiers) => { let last = identifiers.last()?; let name = last.value.to_lowercase(); column_names.contains(&name).then(|| quote_table_identifier(Some(DatabaseType::SqlServer), &last.value)) } Expr::Value(ValueWithSpan { value: Value::Number(number, _), .. }) => { let ordinal = number.parse::().ok()?.checked_sub(1)?; output_columns.get(ordinal).map(|name| quote_table_identifier(Some(DatabaseType::SqlServer), name)) } _ => None, } } fn sql_server_default_pagination_order(statement: &str) -> String { first_simple_sqlserver_projection_order_column(statement).unwrap_or_else(|| "(SELECT NULL)".to_string()) } fn first_simple_sqlserver_projection_order_column(statement: &str) -> Option { let sql = statement.trim(); let sql = &sql[skip_leading_sql_comments(sql, 0)..]; if sql.len() < 6 || !sql[..6].eq_ignore_ascii_case("SELECT") { return None; } let mut index = 6; index = skip_sql_whitespace(sql, index); if let Some(next) = skip_sql_keyword(sql, index, "DISTINCT").or_else(|| skip_sql_keyword(sql, index, "ALL")) { index = skip_sql_whitespace(sql, next); } if let Some(next) = skip_sql_keyword(sql, index, "TOP") { index = skip_sqlserver_top_clause(sql, next); } let projection_start = skip_sql_whitespace(sql, index); let projection_end = find_first_projection_end(sql, projection_start)?; let projection = sql[projection_start..projection_end].trim(); sql_server_derived_order_column_from_projection(projection) } fn skip_leading_sql_comments(sql: &str, mut index: usize) -> usize { loop { index = skip_sql_whitespace(sql, index); if sql[index..].starts_with("--") { index += 2; while index < sql.len() && next_char(sql, index) != '\n' { index += next_char(sql, index).len_utf8(); } continue; } if sql[index..].starts_with("/*") { index += 2; while index < sql.len() { let ch = next_char(sql, index); let next = next_char_at(sql, index + ch.len_utf8()); index += ch.len_utf8(); if ch == '*' && next == Some('/') { index += 1; break; } } continue; } return index; } } fn restore_leading_execution_hint(original_sql: &str, statement: &str, database_type: Option) -> String { if leading_execution_hint_prefix_for_database(statement, database_type).is_some() { return statement.to_string(); } let normalized = original_sql.trim().trim_end_matches(';').trim(); match leading_execution_hint_prefix_for_database(normalized, database_type) { Some(prefix) => format!("{prefix}{statement}"), None => statement.to_string(), } } fn split_leading_execution_hint(sql: &str, database_type: Option) -> (&str, &str) { match leading_execution_hint_prefix_for_database(sql, database_type) { Some(prefix) => (prefix, sql[prefix.len()..].trim_start()), None => ("", sql), } } fn leading_execution_hint_prefix_for_database(sql: &str, database_type: Option) -> Option<&str> { if let Some(prefix) = leading_execution_hint_prefix(sql) { return Some(prefix); } if database_type != Some(DatabaseType::Mysql) { return None; } let executable_start = skip_leading_sql_comments(sql, 0); let directive_start = crate::tdsql_mysql::leading_directive_start(sql, executable_start)?; (directive_start < executable_start).then(|| &sql[directive_start..executable_start]) } fn leading_execution_hint_prefix(sql: &str) -> Option<&str> { let mut index = 0; let mut hint_start = None; loop { index = skip_sql_whitespace(sql, index); let rest = &sql[index..]; if rest.starts_with("--") { index += 2; while index < sql.len() && next_char(sql, index) != '\n' { index += next_char(sql, index).len_utf8(); } continue; } if !rest.starts_with("/*") { return hint_start.map(|start| &sql[start..index]); } if hint_start.is_none() && matches!(next_char_at(sql, index + 2), Some('+' | '@' | '&')) { hint_start = Some(index); } index += 2; while index < sql.len() { let ch = next_char(sql, index); let next = next_char_at(sql, index + ch.len_utf8()); index += ch.len_utf8(); if ch == '*' && next == Some('/') { index += 1; break; } } } } fn strip_leading_statement_comments(sql: &str) -> &str { &sql[skip_leading_sql_comments(sql, 0)..] } fn skip_sqlserver_top_clause(sql: &str, index: usize) -> usize { let mut cursor = skip_sql_whitespace(sql, index); if next_char_at(sql, cursor) == Some('(') { cursor = skip_sql_parenthesized(sql, cursor); } else { while cursor < sql.len() && !next_char(sql, cursor).is_whitespace() && next_char(sql, cursor) != ',' { cursor += next_char(sql, cursor).len_utf8(); } } cursor = skip_sql_whitespace(sql, cursor); if let Some(next) = skip_sql_keyword(sql, cursor, "PERCENT") { cursor = skip_sql_whitespace(sql, next); } if let Some(next) = skip_sql_keyword(sql, cursor, "WITH") { let after_with = skip_sql_whitespace(sql, next); if let Some(after_ties) = skip_sql_keyword(sql, after_with, "TIES") { cursor = skip_sql_whitespace(sql, after_ties); } } cursor } fn find_first_projection_end(sql: &str, start: usize) -> Option { let mut index = start; let mut depth = 0usize; while index < sql.len() { let ch = next_char(sql, index); if matches!(ch, '\'' | '"' | '`') { index = skip_sql_quoted(sql, index, ch); continue; } if ch == '[' { index = skip_sql_bracket_identifier(sql, index); continue; } if ch == '(' { depth += 1; index += ch.len_utf8(); continue; } if ch == ')' { depth = depth.saturating_sub(1); index += ch.len_utf8(); continue; } if depth == 0 && ch == ',' { return Some(index); } if depth == 0 && sql_keyword_at(sql, index, "FROM") { return Some(index); } index += ch.len_utf8(); } None } fn is_simple_sqlserver_order_projection(projection: &str) -> bool { if projection.is_empty() || projection == "*" { return false; } let mut expect_part = true; let mut saw_part = false; let mut index = 0; while index < projection.len() { let ch = next_char(projection, index); if ch.is_whitespace() { return false; } if ch == '.' { if expect_part { return false; } expect_part = true; index += 1; continue; } if ch == '[' { if !expect_part { return false; } let next = skip_sql_bracket_identifier(projection, index); if next <= index + 1 || next > projection.len() { return false; } saw_part = true; expect_part = false; index = next; continue; } if is_sql_token_start(ch) { if !expect_part { return false; } index += ch.len_utf8(); while index < projection.len() && is_sql_token_part(next_char(projection, index)) { index += next_char(projection, index).len_utf8(); } saw_part = true; expect_part = false; continue; } return false; } saw_part && !expect_part } fn sql_server_derived_order_column_from_projection(projection: &str) -> Option { if !is_simple_sqlserver_order_projection(projection) { return None; } let last_part = last_sqlserver_identifier_part(projection); if last_part.starts_with('[') { return Some(last_part.to_string()); } Some(quote_table_identifier(Some(DatabaseType::SqlServer), last_part)) } fn last_sqlserver_identifier_part(projection: &str) -> &str { let mut last_start = 0; let mut index = 0; while index < projection.len() { let ch = next_char(projection, index); if ch == '[' { index = skip_sql_bracket_identifier(projection, index); continue; } if ch == '.' { last_start = index + 1; index += 1; continue; } index += ch.len_utf8(); } projection[last_start..].trim() } fn skip_sql_whitespace(sql: &str, mut index: usize) -> usize { while index < sql.len() && next_char(sql, index).is_whitespace() { index += next_char(sql, index).len_utf8(); } index } fn skip_sql_keyword(sql: &str, index: usize, keyword: &str) -> Option { sql_keyword_at(sql, index, keyword).then_some(index + keyword.len()) } fn sql_keyword_at(sql: &str, index: usize, keyword: &str) -> bool { let Some(candidate) = sql.get(index..index + keyword.len()) else { return false; }; if !candidate.eq_ignore_ascii_case(keyword) { return false; } let before_ok = index == 0 || !is_sql_token_part(next_char_before(sql, index)); let after = index + keyword.len(); let after_ok = after >= sql.len() || !is_sql_token_part(next_char(sql, after)); before_ok && after_ok } fn skip_sql_parenthesized(sql: &str, index: usize) -> usize { let mut cursor = index; let mut depth = 0usize; while cursor < sql.len() { let ch = next_char(sql, cursor); if matches!(ch, '\'' | '"' | '`') { cursor = skip_sql_quoted(sql, cursor, ch); continue; } if ch == '[' { cursor = skip_sql_bracket_identifier(sql, cursor); continue; } if ch == '(' { depth += 1; } else if ch == ')' { depth = depth.saturating_sub(1); cursor += ch.len_utf8(); if depth == 0 { return cursor; } continue; } cursor += ch.len_utf8(); } sql.len() } fn sql_server_derived_table_projection_safe(statement: &str) -> bool { let dialect = MsSqlDialect {}; 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; }; sql_server_derived_table_select_projection_safe(select) } fn sql_server_derived_table_select_projection_safe(select: &Select) -> bool { if matches!(select.projection.as_slice(), [SelectItem::Wildcard(_)]) { return select.from.len() == 1 && select.from[0].joins.is_empty(); } let mut column_names = HashSet::with_capacity(select.projection.len()); select.projection.iter().all(|item| { let Some(name) = derived_projection_name(item) else { return false; }; column_names.insert(name.to_lowercase()) }) } #[derive(Clone, Copy, PartialEq, Eq)] enum MysqlDerivedProjectionSafety { Safe, Ambiguous, Unknown, } fn mysql_count_sql(statement: &str) -> Option { let dialect = MySqlDialect {}; let Ok(mut statements) = Parser::parse_sql(&dialect, statement) else { return Some(mysql_wrapped_count_sql(statement)); }; let projection_safety = { let [Statement::Query(query)] = statements.as_slice() else { return None; }; mysql_derived_table_set_projection_safety(query.body.as_ref()) }; if projection_safety != MysqlDerivedProjectionSafety::Ambiguous { return Some(mysql_wrapped_count_sql(statement)); } let replacement_projection = match Parser::parse_sql(&dialect, "SELECT 1 AS dbx_count_value").ok()?.pop()? { Statement::Query(query) => match query.body.as_ref() { SetExpr::Select(select) => select.projection.clone(), _ => return None, }, _ => return None, }; { let [Statement::Query(query)] = statements.as_mut_slice() else { return None; }; let SetExpr::Select(select) = query.body.as_mut() else { return None; }; let group_by_is_empty = matches!(&select.group_by, GroupByExpr::Expressions(expressions, modifiers) if expressions.is_empty() && modifiers.is_empty()); if select.distinct.is_some() || select.into.is_some() || !group_by_is_empty || select.having.is_some() || !select.projection.iter().all(mysql_projection_item_is_row_preserving) { return None; } select.projection = replacement_projection; // An ORDER BY can refer to a removed output alias; ordering never // changes how many rows survive MySQL LIMIT/OFFSET. query.order_by = None; } Some(mysql_wrapped_count_sql(&statements.pop()?.to_string())) } fn mysql_wrapped_count_sql(statement: &str) -> String { let alias = quote_table_identifier(Some(DatabaseType::Mysql), "dbx_count"); derived_table_sql("SELECT COUNT(*) AS dbx_total_rows FROM", statement, &format!("{alias};")) } fn iotdb_tree_count_sql(statement: &str) -> Option { let dialect = GenericDialect {}; let mut statements = Parser::parse_sql(&dialect, statement).ok()?; let [Statement::Query(query)] = statements.as_mut_slice() else { return None; }; if query.with.is_some() || query.limit_clause.is_some() || query.fetch.is_some() || !query.locks.is_empty() || query.for_clause.is_some() { return None; } let SetExpr::Select(select) = query.body.as_mut() else { return None; }; let group_by_is_empty = matches!(&select.group_by, GroupByExpr::Expressions(expressions, modifiers) if expressions.is_empty() && modifiers.is_empty()); let tree_path = match select.from.as_slice() { [source] if source.joins.is_empty() => match &source.relation { TableFactor::Table { name, .. } => name, _ => return None, }, _ => return None, }; let tree_path_parts = tree_path.0.iter().map(ObjectNamePart::as_ident).collect::>>()?; if tree_path_parts.len() < 3 || !tree_path_parts[0].value.eq_ignore_ascii_case("root") { return None; } if select.distinct.is_some() || select.top.is_some() || select.exclude.is_some() || select.prewhere.is_some() || select.into.is_some() || !group_by_is_empty || select.having.is_some() || select.qualify.is_some() || !select.lateral_views.is_empty() || !select.optimizer_hints.is_empty() || select.select_modifiers.as_ref().is_some_and(SelectModifiers::is_any_set) || !select.connect_by.is_empty() || !select.cluster_by.is_empty() || !select.distribute_by.is_empty() || !select.sort_by.is_empty() || !select.named_window.is_empty() || select.value_table_mode.is_some() { return None; } let measurement = match select.projection.as_slice() { [SelectItem::UnnamedExpr(expr @ (Expr::Identifier(_) | Expr::CompoundIdentifier(_)))] => expr.to_string(), [SelectItem::ExprWithAlias { expr: expr @ (Expr::Identifier(_) | Expr::CompoundIdentifier(_)), .. }] => { expr.to_string() } _ => return None, }; // IoTDB's value-filter mode allows WHERE to reference a series that is // not in the SELECT list; `COUNT()` would then count only // the selected series' non-null points and understate the row total, so // only count when every WHERE reference is `time` or the measurement // itself (bare or full-path form). if let Some(selection) = select.selection.as_ref() { let allowed = |expr: &Expr| match expr { Expr::Identifier(identifier) => { identifier.value.eq_ignore_ascii_case("time") || identifier.value == measurement } Expr::CompoundIdentifier(parts) => { parts.iter().map(|part| part.value.as_str()).collect::>().join(".") == measurement } _ => true, }; if visit_expressions( selection, |expr| { if allowed(expr) { ControlFlow::Continue(()) } else { ControlFlow::Break(()) } }, ) .is_break() { return None; } } let count_projection = match Parser::parse_sql(&dialect, &format!("SELECT COUNT({measurement}) AS dbx_total_rows")).ok()?.pop()? { Statement::Query(query) => match query.body.as_ref() { SetExpr::Select(select) => select.projection.clone(), _ => return None, }, _ => return None, }; select.projection = count_projection; query.order_by = None; Some(format!("{query};")) } fn iris_count_sql(statement: &str) -> Option { let dialect = GenericDialect {}; let mut statements = Parser::parse_sql(&dialect, statement).ok()?; let [Statement::Query(query)] = statements.as_mut_slice() else { return None; }; if query.with.is_some() || query.limit_clause.is_some() || query.fetch.is_some() || !query.locks.is_empty() || query.for_clause.is_some() { return None; } let SetExpr::Select(select) = query.body.as_mut() else { return None; }; let group_by_is_empty = matches!(&select.group_by, GroupByExpr::Expressions(expressions, modifiers) if expressions.is_empty() && modifiers.is_empty()); if select.from.len() != 1 || !select.from[0].joins.is_empty() || !matches!(&select.from[0].relation, TableFactor::Table { .. }) || !matches!(select.projection.as_slice(), [SelectItem::Wildcard(_)] | [SelectItem::QualifiedWildcard(_, _)]) || select.distinct.is_some() || select.top.is_some() || select.exclude.is_some() || select.prewhere.is_some() || select.into.is_some() || !group_by_is_empty || select.having.is_some() || select.qualify.is_some() || !select.lateral_views.is_empty() || !select.optimizer_hints.is_empty() || select.select_modifiers.as_ref().is_some_and(SelectModifiers::is_any_set) || !select.connect_by.is_empty() || !select.cluster_by.is_empty() || !select.distribute_by.is_empty() || !select.sort_by.is_empty() || !select.named_window.is_empty() || select.value_table_mode.is_some() { return None; } let count_projection = match Parser::parse_sql(&dialect, "SELECT COUNT(*) AS dbx_total_rows").ok()?.pop()? { Statement::Query(query) => match query.body.as_ref() { SetExpr::Select(select) => select.projection.clone(), _ => return None, }, _ => return None, }; select.projection = count_projection; query.order_by = None; Some(format!("{query};")) } fn mysql_derived_table_set_projection_safety(set_expr: &SetExpr) -> MysqlDerivedProjectionSafety { match set_expr { SetExpr::Select(select) => mysql_derived_table_select_projection_safety(select), SetExpr::Query(query) => mysql_derived_table_set_projection_safety(query.body.as_ref()), SetExpr::SetOperation { left, .. } => mysql_derived_table_set_projection_safety(left), SetExpr::Values(_) | SetExpr::Table(_) => MysqlDerivedProjectionSafety::Safe, SetExpr::Insert(_) | SetExpr::Update(_) | SetExpr::Delete(_) | SetExpr::Merge(_) => { MysqlDerivedProjectionSafety::Unknown } } } fn mysql_derived_table_select_projection_safety(select: &Select) -> MysqlDerivedProjectionSafety { if select.projection.len() == 1 { return if matches!(select.projection.as_slice(), [SelectItem::Wildcard(_)]) && (select.from.len() != 1 || !select.from[0].joins.is_empty()) { MysqlDerivedProjectionSafety::Ambiguous } else { MysqlDerivedProjectionSafety::Safe }; } if select .projection .iter() .any(|item| matches!(item, SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _))) { return MysqlDerivedProjectionSafety::Ambiguous; } let mut column_names = HashSet::with_capacity(select.projection.len()); let mut unknown_name = false; for item in &select.projection { let Some(name) = derived_projection_name(item) else { unknown_name = true; continue; }; if !column_names.insert(name.to_lowercase()) { return MysqlDerivedProjectionSafety::Ambiguous; } } if unknown_name { MysqlDerivedProjectionSafety::Unknown } else { MysqlDerivedProjectionSafety::Safe } } fn mysql_projection_item_is_row_preserving(item: &SelectItem) -> bool { match item { SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _) => true, SelectItem::UnnamedExpr(Expr::Identifier(_) | Expr::CompoundIdentifier(_)) => true, SelectItem::ExprWithAlias { expr: Expr::Identifier(_) | Expr::CompoundIdentifier(_), .. } => true, SelectItem::UnnamedExpr(_) | SelectItem::ExprWithAlias { .. } | SelectItem::ExprWithAliases { .. } => false, } } fn sql_server_count_sql(statement: &str) -> Option { // FOR JSON/XML produces an unnamed result value that cannot be projected // from the derived table used by the count wrapper. if has_top_level_for_output_clause(statement) { return None; } let dialect = MsSqlDialect {}; let mut statements = Parser::parse_sql(&dialect, statement).ok()?; let derived_table_projection_safe = { let [Statement::Query(query)] = statements.as_slice() else { return None; }; let SetExpr::Select(select) = query.body.as_ref() else { return None; }; sql_server_derived_table_select_projection_safe(select) }; if derived_table_projection_safe { let alias = quote_table_identifier(Some(DatabaseType::SqlServer), "dbx_count"); let wrapped_sql = sql_server_statement_for_derived_table(statement); return Some(derived_table_sql("SELECT COUNT(*) AS dbx_total_rows FROM", &wrapped_sql, &format!("{alias};"))); } let count_projection = match Parser::parse_sql(&dialect, "SELECT COUNT(*) AS dbx_total_rows").ok()?.pop()? { Statement::Query(query) => match query.body.as_ref() { SetExpr::Select(select) => select.projection.clone(), _ => return None, }, _ => return None, }; { let [Statement::Query(query)] = statements.as_mut_slice() else { return None; }; if query.limit_clause.is_some() || query.fetch.is_some() || query.for_clause.is_some() || !query.locks.is_empty() { return None; } let SetExpr::Select(select) = query.body.as_mut() else { return None; }; let group_by_is_empty = matches!(&select.group_by, GroupByExpr::Expressions(expressions, modifiers) if expressions.is_empty() && modifiers.is_empty()); if select.distinct.is_some() || select.top.is_some() || select.into.is_some() || select.from.is_empty() || !group_by_is_empty || select.having.is_some() || !select .projection .iter() .all(|item| matches!(item, SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _))) { return None; } select.projection = count_projection; query.order_by = None; } Some(format!("{};", statements.pop()?)) } fn derived_projection_name(item: &SelectItem) -> Option<&str> { match item { SelectItem::ExprWithAlias { alias, .. } => Some(&alias.value), SelectItem::UnnamedExpr(Expr::Identifier(identifier)) if !identifier.value.starts_with('@') => { Some(&identifier.value) } SelectItem::UnnamedExpr(Expr::CompoundIdentifier(identifiers)) => { identifiers.last().map(|identifier| identifier.value.as_str()) } SelectItem::UnnamedExpr(_) | SelectItem::ExprWithAliases { .. } | SelectItem::QualifiedWildcard(_, _) | SelectItem::Wildcard(_) => None, } } /// Inserts TOP after add_sql_server_offset_fetch has ruled out an existing /// TOP or OFFSET/FETCH clause, avoiding a second AST parse on the first page. fn inject_sql_server_top(sql: &str, limit: usize) -> String { if sql.len() >= 6 && sql[..6].eq_ignore_ascii_case("SELECT") { let rest = &sql[6..]; if let Some((leading, after_modifier)) = strip_sql_server_select_modifier(rest, "DISTINCT") { return format!("SELECT{leading}DISTINCT TOP ({limit}){after_modifier}"); } if let Some((leading, after_modifier)) = strip_sql_server_select_modifier(rest, "ALL") { return format!("SELECT{leading}ALL TOP ({limit}){after_modifier}"); } format!("SELECT TOP ({limit}){rest}") } else { format!("SELECT TOP ({limit}) * FROM ({sql}) [dbx_page]") } } fn strip_sql_server_select_modifier<'a>(rest: &'a str, modifier: &str) -> Option<(&'a str, &'a str)> { let trimmed = rest.trim_start(); let leading_ws_len = rest.len() - trimmed.len(); let candidate = trimmed.get(..modifier.len())?; let after_modifier = trimmed.get(modifier.len()..)?; if !candidate.eq_ignore_ascii_case(modifier) { return None; } if after_modifier.chars().next().is_some_and(is_sql_token_part) { return None; } Some((&rest[..leading_ws_len], after_modifier)) } fn sql_server_statement_for_derived_table(statement: &str) -> String { if has_top_level_select_top(statement) || has_top_level_for_output_clause(statement) { return statement.to_string(); } statement_for_order_insensitive_derived_table(statement) } fn iris_statement_for_derived_table(statement: &str) -> String { statement_for_order_insensitive_derived_table(statement) } fn statement_for_order_insensitive_derived_table(statement: &str) -> String { let Some(order_by) = find_top_level_trailing_order_by(statement) else { return statement.to_string(); }; // Result ordering does not change COUNT cardinality, and stripping only // depth-zero ORDER BY keeps nested query semantics intact. statement[..order_by].trim_end().to_string() } fn add_informix_first_limit(statement: &str, limit: usize, offset: usize) -> String { if has_top_level_informix_row_limit(statement) { return format!("{statement};"); } let row_limit = if offset > 0 { format!("SKIP {offset} FIRST {limit}") } else { format!("FIRST {limit}") }; if statement.len() >= 6 && statement[..6].eq_ignore_ascii_case("SELECT") { let rest = &statement[6..]; if let Some((leading, after_modifier)) = strip_sql_server_select_modifier(rest, "DISTINCT") { return format!("SELECT{leading}DISTINCT {row_limit}{after_modifier};"); } if let Some((leading, after_modifier)) = strip_sql_server_select_modifier(rest, "UNIQUE") { return format!("SELECT{leading}UNIQUE {row_limit}{after_modifier};"); } if let Some((leading, after_modifier)) = strip_sql_server_select_modifier(rest, "ALL") { return format!("SELECT{leading}ALL {row_limit}{after_modifier};"); } return format!("SELECT {row_limit}{rest};"); } format!("SELECT {row_limit} * FROM ({statement}) dbx_page;") } fn add_iris_top_limit(statement: &str, limit: usize) -> String { if has_top_level_select_top(statement) { return format!("{statement};"); } if statement.len() >= 6 && statement[..6].eq_ignore_ascii_case("SELECT") { let rest = &statement[6..]; format!("SELECT TOP {limit}{rest};") } else { format!("SELECT TOP {limit} * FROM ({statement}) dbx_page;") } } fn add_questdb_limit(statement: &str, limit: usize, offset: usize) -> String { if has_top_level_limit(statement) { if offset > 0 { return add_outer_standard_limit(statement, Some(DatabaseType::Questdb), limit, offset, ""); } return format!("{statement};"); } let limit_sql = if offset > 0 { let upper_bound = offset + limit; format!("LIMIT {offset}, {upper_bound}") } else { format!("LIMIT {limit}") }; append_or_insert_before_locking(statement, &limit_sql) } 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 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 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 tokens = top_level_sql_tokens(sql); let limit_index = tokens.iter().position(|token| token.text == "LIMIT")?; let token = &tokens[limit_index]; let (count, suffix_start) = parse_standard_limit_row_count(sql, token.start + token.text.len())?; let suffix_start = skip_sql_whitespace(sql, suffix_start); if sql.get(suffix_start..)?.starts_with('%') { return None; } if tokens[limit_index + 1..].iter().enumerate().any(|(offset, token)| { token.text == "BY" || token.text == "PERCENT" || (token.text == "WITH" && tokens.get(limit_index + offset + 2).is_some_and(|next| next.text == "TIES")) }) { return None; } Some(count) } fn parse_standard_limit_row_count(sql: &str, start: usize) -> Option<(usize, usize)> { let mut cursor = skip_sql_whitespace(sql, start); let first = parse_usize_literal(sql, &mut cursor)?; cursor = skip_sql_whitespace(sql, cursor); if sql.get(cursor..)?.starts_with(',') { cursor = skip_sql_whitespace(sql, cursor + 1); let count = parse_usize_literal(sql, &mut cursor)?; return Some((count, cursor)); } Some((first, cursor)) } fn parse_usize_literal(sql: &str, cursor: &mut usize) -> Option { let start = *cursor; while *cursor < sql.len() && sql.as_bytes()[*cursor].is_ascii_digit() { *cursor += 1; } if *cursor == start { return None; } sql[start..*cursor].parse().ok() } fn has_top_level_informix_row_limit(sql: &str) -> bool { if has_top_level_limit(sql) { return true; } let tokens = top_level_sql_tokens(sql); let Some(select_index) = tokens.iter().position(|token| token.text == "SELECT") else { return false; }; let from_index = tokens .iter() .enumerate() .find(|(index, token)| *index > select_index && token.text == "FROM") .map(|(index, _)| index) .unwrap_or(tokens.len()); tokens[select_index + 1..from_index].iter().any(|token| token.text == "FIRST" || token.text == "SKIP") } fn has_top_level_firebird_row_limit(sql: &str) -> bool { if has_top_level_fetch_first(sql) { return true; } let tokens = top_level_sql_tokens(sql); if tokens.iter().any(|token| token.text == "ROWS") { return true; } let Some(select_index) = tokens.iter().position(|token| token.text == "SELECT") else { return false; }; let from_index = tokens .iter() .enumerate() .find(|(index, token)| *index > select_index && token.text == "FROM") .map(|(index, _)| index) .unwrap_or(tokens.len()); tokens[select_index + 1..from_index].iter().any(|token| token.text == "FIRST" || token.text == "SKIP") } fn has_top_level_fetch_first(sql: &str) -> bool { let tokens = top_level_sql_tokens(sql); tokens.windows(2).any(|w| w[0].text == "FETCH" && w[1].text == "FIRST") } fn has_top_level_rownum(sql: &str) -> bool { top_level_sql_tokens(sql).iter().any(|token| token.text == "ROWNUM") } fn has_top_level_offset_fetch_next(sql: &str) -> bool { let tokens = top_level_sql_tokens(sql); let has_offset = tokens.iter().any(|token| token.text == "OFFSET"); let has_fetch_next = tokens.windows(2).any(|w| w[0].text == "FETCH" && w[1].text == "NEXT"); has_offset && has_fetch_next } /// 用 sqlparser AST 判断 SQL Server 语句顶层是否已带 OFFSET 或 FETCH。 /// 词法扫描器(top_level_sql_tokens)存在三类漏检: /// 1. 把 `#`/`##` 临时表前缀当成注释跳过到行尾,丢掉同一行后面的 OFFSET/FETCH; /// 2. 要求 OFFSET 与 FETCH NEXT 同时出现,漏掉只写 `OFFSET n ROWS` 的合法语句; /// 3. 把字符串里的反斜杠当转义符,引号配对错乱后丢失后续词法。 /// /// AST 解析不受这些影响,因此作为补充检测,避免向已有分页的语句注入 TOP。 fn sql_server_ast_has_offset_or_fetch(statement: &str) -> bool { // 解析失败时返回 false,交由原有词法检测结果决定,不改变既有行为。 let Ok(statements) = Parser::parse_sql(&MsSqlDialect {}, statement) else { return false; }; let [Statement::Query(query)] = statements.as_slice() else { return false; }; // OFFSET 挂在 limit_clause 里(可能是 `OFFSET n ROWS` 单独出现,也可能与 FETCH 同时出现)。 let has_offset = matches!(&query.limit_clause, Some(LimitClause::LimitOffset { offset: Some(_), .. })); has_offset || query.fetch.is_some() } fn add_fetch_first_limit(statement: &str, limit: usize, offset: usize) -> String { if has_top_level_fetch_first(statement) { if offset > 0 { let alias = quote_table_identifier(None, "dbx_page"); return derived_table_sql( "SELECT * FROM", statement, &format!("{alias} OFFSET {offset} ROWS FETCH FIRST {limit} ROWS ONLY;"), ); } return format!("{statement};"); } let offset_sql = if offset > 0 { format!(" OFFSET {offset} ROWS") } else { String::new() }; append_or_insert_before_locking(statement, &format!("{offset_sql} FETCH FIRST {limit} ROWS ONLY")) } fn add_firebird_rows_limit(statement: &str, limit: usize, offset: usize) -> String { if has_top_level_firebird_row_limit(statement) { return format!("{statement};"); } let rows = firebird_rows_clause(limit, offset); append_sql_suffix(statement, &format!("{rows};")) } fn add_rownum_limit(statement: &str, limit: usize, offset: usize) -> String { if has_top_level_rownum(statement) { return format!("{statement};"); } if offset == 0 { return derived_table_sql("SELECT * FROM", statement, &format!("WHERE ROWNUM <= {limit};")); } let end = offset + limit; let column = rownum_pagination_column_name(statement); let inner = derived_table_sql( &format!("SELECT dbx_inner.*, ROWNUM AS \"{column}\" FROM"), statement, &format!("dbx_inner WHERE ROWNUM <= {end}"), ); format!("SELECT * FROM ({inner}) WHERE \"{column}\" > {offset};") } fn rownum_pagination_column_name(statement: &str) -> String { let sql = statement.to_ascii_lowercase(); let base = "__dbx_row_num"; let mut column = base.to_string(); let mut suffix = 0; while sql.contains(&column) { suffix += 1; column = format!("{base}_{suffix}"); } column } fn add_standard_limit( statement: &str, database_type: Option, limit: usize, offset: usize, dedup_order_by: Option>, ) -> String { let order_sql = dedup_order_by.as_deref().map_or(String::new(), format_positional_order_by); if has_top_level_limit(statement) { if !order_sql.is_empty() { // For dedup queries (DISTINCT / GROUP BY) without user ORDER BY, // wrap the query to guarantee deterministic LIMIT/OFFSET pagination. // The inner query preserves DISTINCT semantics; the outer query // adds ORDER BY on positional columns to ensure consistent row // ordering across pages in distributed databases like Doris. return add_outer_standard_limit(statement, database_type, limit, offset, &order_sql); } // A user/top-level LIMIT can still be wider than the selected grid page size. // Wrap it so the first page respects the UI page limit while preserving the user's cap. if offset > 0 || top_level_limit_row_count(statement).is_some_and(|row_count| row_count > limit) { return add_outer_standard_limit(statement, database_type, limit, offset, ""); } return format!("{statement};"); } let offset_sql = if offset > 0 { format!(" OFFSET {offset}") } else { String::new() }; let limit_sql = format!("{order_sql} LIMIT {limit}{offset_sql}"); if database_type == Some(DatabaseType::ClickHouse) { return add_clickhouse_limit(statement, &limit_sql); } append_or_insert_before_locking(statement, &limit_sql) } fn add_outer_standard_limit( statement: &str, database_type: Option, limit: usize, offset: usize, order_sql: &str, ) -> String { let alias = quote_table_identifier(database_type, "dbx_page"); derived_table_sql("SELECT * FROM", statement, &format!("{alias}{order_sql} LIMIT {limit} OFFSET {offset};")) } fn add_clickhouse_limit(statement: &str, limit_sql: &str) -> String { let limit_sql = limit_sql.trim(); let settings_index = clickhouse_settings_clause_index(statement); if let Some(index) = settings_index { let statement_before_settings = statement[..index].trim_end(); let settings_clause = statement[index..].trim_start(); return format!("{statement_before_settings} {limit_sql} {settings_clause};"); } append_or_insert_before_locking(statement, limit_sql) } /// Insert pagination before a top-level locking clause; SQL dialects require /// LIMIT/FETCH to precede FOR UPDATE, FOR SHARE, or LOCK IN SHARE MODE. fn append_or_insert_before_locking(statement: &str, clause: &str) -> String { let clause = clause.trim(); if let Some(index) = locking_clause_index(&top_level_sql_tokens(statement)) { let before = statement[..index].trim_end(); let after = statement[index..].trim_start(); let separator = if sql_suffix_needs_newline(before) { "\n" } else { " " }; return format!("{before}{separator}{clause} {after};"); } append_sql_suffix(statement, &format!("{clause};")) } const LOCKING_CLAUSE_PATTERNS: &[&[&str]] = &[ &["FOR", "UPDATE"], &["FOR", "SHARE"], &["FOR", "KEY", "SHARE"], &["FOR", "NO", "KEY", "UPDATE"], &["LOCK", "IN", "SHARE", "MODE"], ]; fn locking_clause_index(tokens: &[SqlToken]) -> Option { tokens.iter().enumerate().find_map(|(index, token)| { LOCKING_CLAUSE_PATTERNS .iter() .any(|pattern| token_sequence_matches(&tokens[index..], pattern)) .then_some(token.start) }) } fn token_sequence_matches(tokens: &[SqlToken], expected: &[&str]) -> bool { tokens.len() >= expected.len() && tokens.iter().zip(expected).all(|(token, expected)| token.text == *expected) } fn has_pagination_clause_after(tokens: &[SqlToken], index: usize) -> bool { tokens.iter().any(|token| token.start > index && matches!(token.text.as_str(), "LIMIT" | "OFFSET" | "FETCH")) } fn clickhouse_settings_clause_index(statement: &str) -> Option { let parsed_settings = Parser::parse_sql(&ClickHouseDialect {}, statement).ok().and_then(|statements| { let [Statement::Query(query)] = statements.as_slice() else { return None; }; Some(query.settings.is_some()) }); if parsed_settings == Some(false) { return None; } // Keep the lexical fallback for valid ClickHouse syntax that sqlparser does not yet support. top_level_sql_tokens(statement) .iter() .rev() .find(|token| token.text == "SETTINGS" && !is_qualified_identifier_part(statement, token)) .map(|token| token.start) } fn is_qualified_identifier_part(sql: &str, token: &SqlToken) -> bool { let token_end = token.start + token.text.len(); sql[..token.start].chars().rev().find(|ch| !ch.is_whitespace()) == Some('.') || sql[token_end..].chars().find(|ch| !ch.is_whitespace()) == Some('.') } fn derived_table_sql(prefix: &str, statement: &str, suffix: &str) -> String { format!("{prefix} ({}) {suffix}", statement_for_sql_suffix(statement)) } fn append_sql_suffix(statement: &str, suffix: &str) -> String { let separator = if sql_suffix_needs_newline(statement) { "\n" } else { " " }; format!("{}{separator}{}", statement.trim_end(), suffix.trim_start()) } fn statement_for_sql_suffix(statement: &str) -> String { let trimmed = statement.trim_end(); if sql_suffix_needs_newline(trimmed) { format!("{trimmed}\n") } else { trimmed.to_string() } } fn sql_suffix_needs_newline(sql: &str) -> bool { let Some(last_line_start) = sql.rfind(['\n', '\r']).map(|index| index + 1) else { return line_has_open_line_comment(sql); }; line_has_open_line_comment(&sql[last_line_start..]) } fn line_has_open_line_comment(line: &str) -> bool { let mut index = 0; while index < line.len() { let ch = next_char(line, index); let next = next_char_at(line, index + ch.len_utf8()); if matches!(ch, '\'' | '"' | '`') { index = skip_sql_quoted(line, index, ch); continue; } if ch == '[' { index = skip_sql_bracket_identifier(line, index); continue; } if ch == '/' && next == Some('*') { index += 2; while index < line.len() { let current = next_char(line, index); let following = next_char_at(line, index + current.len_utf8()); index += current.len_utf8(); if current == '*' && following == Some('/') { index += 1; break; } } continue; } if ch == '-' && next == Some('-') { return true; } if ch == '#' { return true; } index += ch.len_utf8(); } false } /// For dedup queries (DISTINCT / GROUP BY) without an ORDER BY clause, generate /// a positional `ORDER BY 1, 2, ..., N` clause so that LIMIT/OFFSET pagination /// returns deterministic results across pages. This is especially important for /// distributed databases (e.g. Doris, StarRocks) where tablet scan order varies /// between independent query executions. fn format_positional_order_by(positions: &[usize]) -> String { if positions.is_empty() { return String::new(); } let cols: Vec = positions.iter().map(|position| position.to_string()).collect(); format!(" ORDER BY {}", cols.join(", ")) } /// Detect dedup queries (SELECT DISTINCT, GROUP BY, HAVING) that lack a /// top-level ORDER BY clause. Returns the 1-based projection positions to sort /// on so that a positional ORDER BY can be injected for deterministic pagination. /// /// A GROUP BY query sorts on its grouping keys alone: the keys already identify /// an output row uniquely, and some engines reject sorting on aggregate outputs. /// /// Returns `None` for: /// - Non-SELECT queries /// - Queries without dedup semantics /// - Queries that already specify ORDER BY /// - Wildcard projections (`SELECT *`) /// - Parse failures fn dedup_projection_count_without_order_by(sql: &str) -> Option> { let dialect = GenericDialect {}; let statements = Parser::parse_sql(&dialect, sql).ok()?; let [Statement::Query(query)] = statements.as_slice() else { return None; }; // Reject if the query already has an ORDER BY clause. if query.order_by.is_some() { return None; } let SetExpr::Select(select) = query.body.as_ref() else { return None; }; let has_distinct = select.distinct.is_some(); let has_group_by = !matches!(&select.group_by, GroupByExpr::Expressions(exprs, _) if exprs.is_empty()); let has_having = select.having.is_some(); if !has_distinct && !has_group_by && !has_having { return None; } // Wildcard projections cannot be used with positional ORDER BY. if select.projection.len() == 1 && matches!(select.projection.first(), Some(SelectItem::Wildcard(_))) { return None; } let all_positions: Vec = (1..=select.projection.len()).collect(); if has_distinct { return Some(all_positions); } if let GroupByExpr::Expressions(exprs, _) = &select.group_by { if let Some(positions) = group_by_key_positions(exprs, &select.projection) { return Some(positions); } } Some(all_positions) } /// Map every GROUP BY key onto the 1-based position of the output column that /// exposes it. Returns `None` when a key is not projected (sorting on the keys /// alone is then impossible) or when a wildcard hides which position is which, /// so callers fall back to sorting on the full projection. fn group_by_key_positions(group_by: &[Expr], projection: &[SelectItem]) -> Option> { if group_by.is_empty() { return None; } if !projection.iter().all(|item| matches!(item, SelectItem::UnnamedExpr(_) | SelectItem::ExprWithAlias { .. })) { return None; } let mut positions: Vec = Vec::with_capacity(group_by.len()); for key in group_by { let position = group_by_key_position(key, projection)?; if !positions.contains(&position) { positions.push(position); } } positions.sort_unstable(); Some(positions) } fn group_by_key_position(key: &Expr, projection: &[SelectItem]) -> Option { // `GROUP BY 1` is already a projection position. if let Expr::Value(ValueWithSpan { value: Value::Number(number, _), .. }) = key { let position = number.parse::().ok()?; return (1..=projection.len()).contains(&position).then_some(position); } let key_sql = key.to_string(); projection .iter() .position(|item| match item { SelectItem::UnnamedExpr(expr) => expr.to_string() == key_sql, SelectItem::ExprWithAlias { expr, alias } => alias.value == key_sql || expr.to_string() == key_sql, _ => false, }) .map(|index| index + 1) } fn find_top_level_trailing_order_by(sql: &str) -> Option { let tokens = top_level_sql_tokens(sql); for index in (0..tokens.len().saturating_sub(1)).rev() { if tokens[index].text == "ORDER" && tokens.get(index + 1).is_some_and(|token| token.text == "BY") { return Some(tokens[index].start); } } None } fn has_top_level_select_top(sql: &str) -> bool { top_level_select_tokens_before_from(sql).iter().any(|token| token.text == "TOP") } fn has_top_level_set_operator(sql: &str) -> bool { top_level_sql_tokens(sql).iter().any(|token| matches!(token.text.as_str(), "UNION" | "INTERSECT" | "EXCEPT")) } fn has_top_level_select_distinct(sql: &str) -> bool { top_level_select_tokens_before_from(sql).iter().any(|token| token.text == "DISTINCT") } fn top_level_select_tokens_before_from(sql: &str) -> Vec { let tokens = top_level_sql_tokens(sql); let Some(select_index) = tokens.iter().position(|token| token.text == "SELECT") else { return Vec::new(); }; let from_index = tokens .iter() .enumerate() .find(|(index, token)| *index > select_index && token.text == "FROM") .map(|(index, _)| index) .unwrap_or(tokens.len()); tokens[select_index + 1..from_index].to_vec() } fn has_top_level_for_output_clause(sql: &str) -> bool { let tokens = top_level_sql_tokens(sql); tokens.iter().enumerate().any(|(index, token)| { token.text == "FOR" && tokens.get(index + 1).is_some_and(|next| next.text == "XML" || next.text == "JSON") }) } #[derive(Debug, Clone, PartialEq, Eq)] struct SqlToken { text: String, start: usize, } fn top_level_sql_tokens(sql: &str) -> Vec { let mut tokens = Vec::new(); let mut i = 0; let mut depth = 0usize; while i < sql.len() { let ch = next_char(sql, i); let next = next_char_at(sql, i + ch.len_utf8()); if ch == '-' && next == Some('-') { i += 2; while i < sql.len() && next_char(sql, i) != '\n' { i += next_char(sql, i).len_utf8(); } continue; } if ch == '#' { i += 1; while i < sql.len() && next_char(sql, i) != '\n' { i += next_char(sql, i).len_utf8(); } continue; } if ch == '/' && next == Some('*') { i += 2; while i < sql.len() { let current = next_char(sql, i); let following = next_char_at(sql, i + current.len_utf8()); if current == '*' && following == Some('/') { i += 2; break; } i += current.len_utf8(); } continue; } // PostgreSQL dollar-quoted bodies may contain arbitrary SQL keywords. // Skip them before scanning for top-level clauses such as FOR UPDATE. if ch == '$' { if let Some(end) = skip_sql_dollar_quoted(sql, i) { i = end; continue; } } if matches!(ch, '\'' | '"' | '`') { i = skip_sql_quoted(sql, i, ch); continue; } if ch == '[' { i = skip_sql_bracket_identifier(sql, i); continue; } if ch == '(' { depth += 1; i += ch.len_utf8(); continue; } if ch == ')' { depth = depth.saturating_sub(1); i += ch.len_utf8(); continue; } if depth == 0 && is_sql_token_start(ch) { let start = i; i += ch.len_utf8(); while i < sql.len() && is_sql_token_part(next_char(sql, i)) { i += next_char(sql, i).len_utf8(); } tokens.push(SqlToken { text: sql[start..i].to_ascii_uppercase(), start }); continue; } i += ch.len_utf8(); } tokens } fn skip_sql_dollar_quoted(sql: &str, pos: usize) -> Option { let tag_end_offset = sql.get(pos + 1..)?.find('$')?; let tag_end = pos + 1 + tag_end_offset; let tag = &sql[pos + 1..tag_end]; let valid_tag = tag.is_empty() || (tag.chars().next().is_some_and(|ch| ch.is_ascii_alphabetic() || ch == '_') && tag.chars().all(|ch| ch.is_ascii_alphanumeric() || ch == '_')); if !valid_tag { return None; } let delimiter = &sql[pos..=tag_end]; let content_start = tag_end + 1; sql.get(content_start..)?.find(delimiter).map(|closing_offset| content_start + closing_offset + delimiter.len()) } fn skip_sql_quoted(sql: &str, pos: usize, quote: char) -> usize { let mut i = pos + quote.len_utf8(); while i < sql.len() { let ch = next_char(sql, i); let next = next_char_at(sql, i + ch.len_utf8()); if ch == quote { if next == Some(quote) { i += ch.len_utf8() + quote.len_utf8(); continue; } return i + ch.len_utf8(); } if quote == '\'' && ch == '\\' { i += ch.len_utf8(); if i < sql.len() { i += next_char(sql, i).len_utf8(); } continue; } i += ch.len_utf8(); } sql.len() } fn skip_sql_bracket_identifier(sql: &str, pos: usize) -> usize { let mut i = pos + 1; while i < sql.len() { let ch = next_char(sql, i); let next = next_char_at(sql, i + ch.len_utf8()); if ch == ']' { if next == Some(']') { i += 2; continue; } return i + 1; } i += ch.len_utf8(); } sql.len() } fn is_sql_token_start(ch: char) -> bool { ch.is_ascii_alphabetic() || ch == '_' } fn is_sql_token_part(ch: char) -> bool { ch.is_ascii_alphanumeric() || matches!(ch, '_' | '$' | '#') } fn next_char(sql: &str, index: usize) -> char { sql[index..].chars().next().unwrap_or('\0') } fn next_char_before(sql: &str, index: usize) -> char { sql[..index].chars().next_back().unwrap_or('\0') } fn next_char_at(sql: &str, index: usize) -> Option { if index >= sql.len() { None } else { sql[index..].chars().next() } } fn build_derived_column_aliases(result_columns: &[String]) -> Vec { let mut seen: std::collections::HashMap = std::collections::HashMap::new(); result_columns .iter() .enumerate() .map(|(index, column)| { let base = normalize_alias_base(column, index); let count = seen.entry(base.clone()).and_modify(|value| *value += 1).or_insert(1); if *count == 1 { base } else { format!("{base}_{count}") } }) .collect() } fn normalize_alias_base(column: &str, index: usize) -> String { let compact = column.split_whitespace().collect::>().join("_"); let safe = compact .chars() .map(|ch| if ch.is_alphanumeric() || matches!(ch, '_' | '$') { ch } else { '_' }) .collect::() .trim_matches('_') .to_string(); if safe.is_empty() { fallback_alias(index) } else { safe } } fn fallback_alias(index: usize) -> String { format!("column_{}", index + 1) } #[cfg(test)] mod tests { use super::*; /// Query shape from issue #7832: a MySQL GROUP BY over a LEFT JOIN with an /// aggregated derived table, COUNT(DISTINCT IF(...)) in the projection, and /// inline `-- 中文` line comments. Locks in the invariants that keep DBX's /// derived page/count SQL row-count-identical to the user's statement: /// the statement splitter must keep it a single statement, the page SQL /// must preserve the GROUP BY while injecting deterministic pagination, /// and the count wrap must count the grouped result, not the raw join. #[test] fn mysql_group_by_with_distinct_if_and_comments_keeps_grouping_in_page_and_count_sql() { let sql = "SELECT\n base.brand_name\n ,base.stall_id\n ,base.floor\n ,COUNT(1) total_invite -- 邀约数量\n ,SUM(IFNULL(base.ver_status, 0)) sign_num -- 签到数量\n ,SUM(IFNULL(dr.draw_count, 0)) draw_num -- 抽奖次数\n ,SUM(IFNULL(dr.draw_user_num, 0)) draw_user_num -- 抽奖人数\n ,COUNT(distinct IF(base.ver_status = 1, base.mobile, null)) sign_and_draw_user_num\nFROM v_form_data_1786326962 base\nLEFT JOIN (\n SELECT id, SUM(IFNULL(hx_status, 0)) draw_count, COUNT(distinct IF(hx_status = 1, mobile, null)) draw_user_num\n FROM v_form_data_1786326962_coupon\n GROUP BY id\n) dr ON base.id = dr.id\nGROUP BY base.brand_name, base.stall_id, base.floor"; let statements = crate::sql::split_sql_statements_for_database(sql, DatabaseType::Mysql); assert_eq!(statements.len(), 1, "inline comments must not split the statement"); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::Mysql), pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); let page_sql = plan.page_sql.expect("pagination plan provides page SQL"); assert!(page_sql.to_uppercase().contains("GROUP BY")); assert!( page_sql.contains("ORDER BY 1, 2, 3 LIMIT 500;"), "dedup pagination must be appended after the GROUP BY, got: {page_sql}" ); let count_sql = plan.count_sql.expect("pagination plan provides count SQL"); assert!(count_sql.starts_with("SELECT COUNT(*) AS dbx_total_rows FROM (")); assert_eq!( count_sql.matches("GROUP BY").count(), 2, "the count wrap must keep both the outer and derived-table GROUP BY, got: {count_sql}" ); let mut variants = vec![ ("crlf", sql.replace('\n', "\r\n")), ("trailing semicolon", format!("{sql};")), ("trailing GROUP BY comment", format!("{sql} -- 分组")), ]; if let Some(stripped) = sql.strip_prefix("SELECT") { variants.push(("leading comment", format!("-- header\nSELECT{stripped}"))); } for (name, variant) in variants { let counted = build_count_query_sql(CountQuerySqlOptions { original_sql: variant.clone(), database_type: Some(DatabaseType::Mysql), }); let counted = counted.sql.unwrap_or_default(); assert_eq!(counted.matches("GROUP BY").count(), 2, "{name}: count wrap kept both GROUP BYs"); let statements = crate::sql::split_sql_statements_for_database(&variant, DatabaseType::Mysql); assert_eq!(statements.len(), 1, "{name}: still a single statement"); } } #[test] fn easysearch_uses_elasticsearch_sql_pagination_rules() { let paginated = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT name FROM products".to_string(), database_type: Some(DatabaseType::Easysearch), limit: 100, offset: 200, }); let counted = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT name FROM products".to_string(), database_type: Some(DatabaseType::Easysearch), }); assert_eq!(paginated.sql.as_deref(), Some("SELECT name FROM products LIMIT 100 OFFSET 200;")); assert_eq!(counted, err("unsupported")); } #[test] fn wraps_single_select_query_with_limit_and_offset() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id, name FROM users;".to_string(), database_type: Some(DatabaseType::Postgres), limit: 100, offset: 200, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT id, name FROM users LIMIT 100 OFFSET 200;"); } #[test] 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), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT TOP (100) id FROM users ORDER BY id DESC"); } #[test] fn sqlserver_qualified_order_by_uses_bounded_rowcount_for_later_pages() { let sql = "SELECT LEFT(d.lbbh, 2) AS dlbh, big.lbmc AS dlmc, LEFT(d.lbbh, 4) AS zlbh, middle.lbmc AS zlmc, LEFT(d.lbbh, 6) AS xlbh, small.lbmc AS xlmc, d.lbbh AS cxlbh, d.lbmc AS cxlmc, CASE WHEN d.tybz = '1' THEN '启用' ELSE '停用' END AS zt FROM T_BASE_WZLB AS d LEFT JOIN T_BASE_WZLB AS big ON big.lbbh = LEFT(d.lbbh, 2) AND big.TreeInfo_Layer = 1 LEFT JOIN T_BASE_WZLB AS middle ON middle.lbbh = LEFT(d.lbbh, 4) AND middle.TreeInfo_Layer = 2 LEFT JOIN T_BASE_WZLB AS small ON small.lbbh = LEFT(d.lbbh, 6) AND small.TreeInfo_Layer = 3 WHERE d.TreeInfo_IsDetail = 1 ORDER BY d.lbbh"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 500, offset: 500, }); let page_sql = result.sql.expect("build qualified order page"); assert!(page_sql.starts_with("EXEC sys.sp_executesql N'SET ROWCOUNT 1000; SELECT LEFT(d.lbbh")); assert!(page_sql.contains("ORDER BY d.lbbh'")); assert!(!page_sql.contains("ROW_NUMBER()")); assert_eq!(sqlserver_result_offset(&page_sql), 500); } #[test] fn sqlserver_output_alias_order_keeps_row_number_pagination() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT d.lbbh AS cxlbh FROM T_BASE_WZLB AS d ORDER BY cxlbh".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 500, offset: 500, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY cxlbh) AS [__dbx_row_num] FROM (SELECT d.lbbh AS cxlbh FROM T_BASE_WZLB AS d) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 500 AND [__dbx_row_num] <= 1000 ORDER BY [__dbx_row_num];" ); } #[test] fn sqlserver_unprojected_order_column_uses_bounded_rowcount() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT d.lbmc AS cxlmc FROM T_BASE_WZLB AS d ORDER BY lbbh".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 500, offset: 500, }); let page_sql = result.sql.expect("build hidden order column page"); assert!(page_sql.starts_with("EXEC sys.sp_executesql N'SET ROWCOUNT 1000;")); assert_eq!(sqlserver_result_offset(&page_sql), 500); } #[test] 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), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT TOP (100) COUNT(*) FROM TicketInfo"); } #[test] fn uses_sqlserver_top_for_unnamed_expression_on_first_page() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id + 1 FROM TicketInfo".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT TOP (100) id + 1 FROM TicketInfo"); } #[test] fn paginates_sqlserver_unnamed_expression_with_rowcount() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id + 1 FROM TicketInfo".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); let sql = result.sql.expect("build unnamed expression page"); assert!(sql.starts_with("EXEC sys.sp_executesql N'SET ROWCOUNT 200; SELECT id + 1 FROM TicketInfo'")); assert_eq!(sqlserver_result_offset(&sql), 100); } #[test] fn paginates_sqlserver_set_operations_with_rowcount_instead_of_limiting_first_branch() { // TOP injected into the first SELECT of INTERSECT/EXCEPT/UNION only bounds // that branch, so the combined result is wrong (e.g. 0 rows instead of the // real intersection). The whole statement must be bounded instead. for operator in ["INTERSECT", "EXCEPT", "UNION", "UNION ALL"] { let original_sql = format!("SELECT id FROM a {operator} SELECT id FROM b"); let first_page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: original_sql.clone(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); let sql = first_page.sql.expect("build set operation first page"); assert!(!sql.contains("TOP ("), "{operator}: {sql}"); assert_eq!( sql, format!("EXEC sys.sp_executesql N'SET ROWCOUNT 100; {original_sql}'; /*__dbx_result_offset=0__*/") ); let later_page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql, database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); let sql = later_page.sql.expect("build set operation later page"); assert!(sql.starts_with("EXEC sys.sp_executesql N'SET ROWCOUNT 200; "), "{operator}: {sql}"); assert_eq!(sqlserver_result_offset(&sql), 100); } } #[test] fn keeps_top_level_order_by_when_paginating_sqlserver_set_operations() { let original_sql = "SELECT id FROM a INTERSECT SELECT id FROM b ORDER BY id"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert_eq!( result.sql.unwrap(), format!("EXEC sys.sp_executesql N'SET ROWCOUNT 100; {original_sql}'; /*__dbx_result_offset=0__*/") ); } #[test] fn keeps_sqlserver_top_when_set_operator_is_only_inside_subquery() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM (SELECT id FROM a UNION SELECT id FROM b) t".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT TOP (100) id FROM (SELECT id FROM a UNION SELECT id FROM b) t"); } #[test] fn uses_sqlserver_top_for_distinct_queries_without_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT ProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT DISTINCT TOP (100) ProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data" ); } #[test] fn rejects_sqlserver_distinct_later_pages_without_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT ProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); assert_eq!(result, err("unsupported")); } #[test] 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), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT ALL TOP (100) ProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data"); } #[test] fn paginates_sqlserver_select_prefix_like_all_modifier() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT AllProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT TOP (100) AllProjectType FROM JDDR_sys_BasicConfig_ProjectInfo_Data"); } #[test] fn paginates_existing_sqlserver_top_clause_on_first_page() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT TOP 1000 * FROM TicketInfo".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT TOP (100) * FROM (SELECT TOP 1000 * FROM TicketInfo) [dbx_page] ORDER BY (SELECT NULL);" ); } #[test] fn paginates_existing_sqlserver_top_clause_for_later_pages() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT TOP 1000 * FROM TicketInfo".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS [__dbx_row_num] FROM (SELECT TOP 1000 * FROM TicketInfo) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];" ); } #[test] fn paginates_existing_sqlserver_top_clause_keeping_user_order_by_desc() { let sql = "SELECT top 10 [ID], [ProcInstID], [ActInstID], [ActInstDestID], [StartDate], [DestUser], [Status], [Data], [SerialNumber] FROM [dbo].[_WorkList] ORDER BY ID DESC"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT TOP (100) * FROM (SELECT top 10 [ID], [ProcInstID], [ActInstID], [ActInstDestID], [StartDate], [DestUser], [Status], [Data], [SerialNumber] FROM [dbo].[_WorkList] ORDER BY ID DESC) [dbx_page] ORDER BY [ID] DESC;" ); } #[test] fn paginates_existing_sqlserver_top_clause_later_page_keeping_user_order_by_desc() { let sql = "SELECT top 10 [ID], [ProcInstID] FROM [dbo].[_WorkList] ORDER BY ID DESC"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY [ID] DESC) AS [__dbx_row_num] FROM (SELECT top 10 [ID], [ProcInstID] FROM [dbo].[_WorkList] ORDER BY ID DESC) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];" ); } #[test] fn paginates_existing_sqlserver_top_clause_keeping_multi_column_user_order_by() { let sql = "SELECT top 100 [ID], [StartDate] FROM [dbo].[_WorkList] ORDER BY [StartDate] DESC, ID"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT TOP (100) * FROM (SELECT top 100 [ID], [StartDate] FROM [dbo].[_WorkList] ORDER BY [StartDate] DESC, ID) [dbx_page] ORDER BY [StartDate] DESC, [ID];" ); } #[test] fn paginates_existing_sqlserver_top_clause_keeping_ordinal_order_by() { let sql = "SELECT top 100 [ID], [Name] FROM [dbo].[_WorkList] ORDER BY 2 DESC"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT TOP (100) * FROM (SELECT top 100 [ID], [Name] FROM [dbo].[_WorkList] ORDER BY 2 DESC) [dbx_page] ORDER BY [Name] DESC;" ); } #[test] fn paginates_later_sqlserver_top_page_mapping_ordinal_order_to_output_column() { let sql = "SELECT top 100 [ID], [Name] FROM [dbo].[_WorkList] ORDER BY 2 DESC"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY [Name] DESC) AS [__dbx_row_num] FROM (SELECT top 100 [ID], [Name] FROM [dbo].[_WorkList] ORDER BY 2 DESC) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];" ); } #[test] fn paginates_sqlserver_top_clause_falling_back_for_out_of_range_ordinal() { let sql = "SELECT top 100 [ID], [Name] FROM [dbo].[_WorkList] ORDER BY 3 DESC"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY [ID]) AS [__dbx_row_num] FROM (SELECT top 100 [ID], [Name] FROM [dbo].[_WorkList] ORDER BY 3 DESC) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];" ); } #[test] fn paginates_existing_sqlserver_top_clause_falls_back_when_order_by_not_projected() { let sql = "SELECT top 100 [ID] FROM [dbo].[_WorkList] ORDER BY [SerialNumber]"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT TOP (100) * FROM (SELECT top 100 [ID] FROM [dbo].[_WorkList] ORDER BY [SerialNumber]) [dbx_page] ORDER BY [ID];" ); } #[test] fn sqlserver_top_query_later_page_keeps_page_metadata() { let sql = "SELECT TOP 1000 * FROM TicketInfo".to_string(); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.clone(), query_base_sql: sql, database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 100, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!( plan.sql_to_execute, "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS [__dbx_row_num] FROM (SELECT TOP 1000 * FROM TicketInfo) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];" ); assert_eq!(plan.page_sql, Some(plan.sql_to_execute.clone())); assert_eq!(plan.page_limit, Some(100)); assert_eq!(plan.page_offset, Some(100)); } #[test] fn sqlserver_unsafe_top_projections_use_rowcount_for_later_pages() { let queries = [ "SELECT TOP 100 AAA, * FROM BBB", "SELECT TOP 100 AAA, bbb AS aaa FROM BBB", "SELECT TOP 100 a.id, b.id FROM AAA a JOIN BBB b ON a.id = b.id", "SELECT TOP 100 a.*, b.* FROM AAA a JOIN BBB b ON a.id = b.id", "SELECT TOP 100 a.*, * FROM AAA a", "SELECT TOP 100 AAA + 1 FROM BBB", ]; for sql in queries { let first_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(first_page.sql_to_execute, sql); assert!(first_page.page_sql.is_none()); assert!(first_page.count_sql.is_none()); let later_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 100, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert!(later_page.sql_to_execute.starts_with("EXEC sys.sp_executesql N'SET ROWCOUNT 200; ")); assert_eq!(sqlserver_result_offset(&later_page.sql_to_execute), 100); assert_eq!(later_page.page_sql, Some(later_page.sql_to_execute.clone())); assert!(later_page.count_sql.is_none()); assert_eq!(later_page.page_limit, Some(100)); assert_eq!(later_page.page_offset, Some(100)); } } #[test] fn sqlserver_unsafe_projection_is_not_wrapped_for_count() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT TOP 100 AAA, * FROM BBB".to_string(), database_type: Some(DatabaseType::SqlServer), }); assert_eq!(result, err("unsupported")); } #[test] fn sqlserver_join_wildcard_injects_count_projection() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT * FROM WZ_CKGL_WZLLDSQ_DETAIL d LEFT JOIN WZ_CKGL_WZLLDSQ_MAIN AS m ON m.ID = d.ParentID" .to_string(), database_type: Some(DatabaseType::SqlServer), }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM WZ_CKGL_WZLLDSQ_DETAIL d LEFT JOIN WZ_CKGL_WZLLDSQ_MAIN AS m ON m.ID = d.ParentID;" ); } #[test] fn sqlserver_join_wildcard_count_removes_top_level_order_by() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT * FROM AAA a JOIN BBB b ON a.id = b.id ORDER BY a.id".to_string(), database_type: Some(DatabaseType::SqlServer), }); assert_eq!( result.sql.as_deref(), Some("SELECT COUNT(*) AS dbx_total_rows FROM AAA a JOIN BBB b ON a.id = b.id;") ); } #[test] fn sqlserver_join_wildcard_plan_exposes_count_without_changing_first_page() { let sql = "SELECT d.*, m.* FROM WZ_CKGL_WZLLDSQ_DETAIL d LEFT JOIN WZ_CKGL_WZLLDSQ_MAIN m ON m.ID = d.ParentID"; let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), 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 (500) d.*, m.* FROM WZ_CKGL_WZLLDSQ_DETAIL d LEFT JOIN WZ_CKGL_WZLLDSQ_MAIN m ON m.ID = d.ParentID" ); assert_eq!( plan.count_sql.as_deref(), Some( "SELECT COUNT(*) AS dbx_total_rows FROM WZ_CKGL_WZLLDSQ_DETAIL d LEFT JOIN WZ_CKGL_WZLLDSQ_MAIN m ON m.ID = d.ParentID;" ) ); } #[test] fn sqlserver_wildcard_count_rejects_semantic_modifiers() { for sql in [ "SELECT DISTINCT * FROM AAA a JOIN BBB b ON a.id = b.id", "SELECT TOP 10 * FROM AAA a JOIN BBB b ON a.id = b.id", "SELECT * INTO #joined FROM AAA a JOIN BBB b ON a.id = b.id", "SELECT * FROM AAA a JOIN BBB b ON a.id = b.id GROUP BY a.id", "SELECT * FROM AAA a JOIN BBB b ON a.id = b.id ORDER BY a.id OFFSET 10 ROWS FETCH NEXT 10 ROWS ONLY", ] { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), }); assert!(!result.ok, "must not rewrite {sql}"); } } #[test] fn sqlserver_join_wildcard_uses_bounded_rowcount_pagination() { let sql = "SELECT * FROM WZ_CKGL_WZLLDSQ_DETAIL d LEFT JOIN WZ_CKGL_WZLLDSQ_MAIN AS m ON m.ID = d.ParentID"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 500, offset: 500, }); assert!(result.ok); let sql = result.sql.unwrap(); assert_eq!( sql, "EXEC sys.sp_executesql N'SET ROWCOUNT 1000; SELECT * FROM WZ_CKGL_WZLLDSQ_DETAIL d LEFT JOIN WZ_CKGL_WZLLDSQ_MAIN AS m ON m.ID = d.ParentID'; /*__dbx_result_offset=500__*/" ); assert_eq!(sqlserver_result_offset(&sql), 500); } #[test] fn sqlserver_result_offset_ignores_user_sql_markers() { assert_eq!(sqlserver_result_offset("SELECT 1 /*__dbx_result_offset=500__*/"), 0); } #[test] fn sqlserver_rowcount_pagination_escapes_string_literals() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM detail d JOIN parent p ON p.id = d.parent_id WHERE p.label = N'O''Brien'" .to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); let sql = result.sql.expect("build rowcount pagination SQL"); assert!(sql.contains("WHERE p.label = N''O''''Brien''")); assert_eq!(sqlserver_result_offset(&sql), 100); } #[test] fn sqlserver_unique_top_projection_keeps_server_pagination() { let sql = "SELECT TOP 500 [id], [order_no] FROM [sales].[orders_10k]".to_string(); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.clone(), query_base_sql: sql, database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 100, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert!(plan.sql_to_execute.contains("ROW_NUMBER()")); assert_eq!(plan.page_sql, Some(plan.sql_to_execute.clone())); assert_eq!(plan.page_limit, Some(100)); assert_eq!(plan.page_offset, Some(100)); assert!(plan.count_sql.is_some()); } #[test] fn sqlserver_json_output_executes_original_sql_without_wrappers() { let sql = "SELECT TOP 10 * FROM [sales].[orders_10k] FOR JSON AUTO".to_string(); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.clone(), query_base_sql: sql.clone(), database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, sql); assert!(plan.page_sql.is_none()); assert!(plan.count_sql.is_none()); assert_eq!(plan.page_limit, None); assert_eq!(plan.page_offset, None); assert_eq!(plan.exact_query_row_bound, None); assert!(!plan.single_execution); assert_eq!( build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT TOP 10 * FROM [sales].[orders_10k] FOR JSON AUTO".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }), err("unsupported") ); assert_eq!( build_count_query_sql(CountQuerySqlOptions { original_sql: sql, database_type: Some(DatabaseType::SqlServer), }), err("unsupported") ); } #[test] fn sqlserver_json_output_without_top_also_skips_wrappers() { let sql = "SELECT id, name FROM users FOR JSON PATH".to_string(); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.clone(), query_base_sql: sql.clone(), database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, sql); assert!(plan.page_sql.is_none()); assert!(plan.count_sql.is_none()); assert_eq!(plan.exact_query_row_bound, None); } #[test] fn sqlserver_nested_json_output_in_subquery_still_paginates() { // FOR JSON nested inside a scalar subquery is not a top-level output // clause; the outer query is a normal projection and paginates. let sql = "SELECT id, (SELECT name FROM tags FOR JSON PATH) AS tags FROM users".to_string(); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.clone(), query_base_sql: sql, database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert!(plan.page_sql.is_some()); assert!(plan.count_sql.is_some()); } #[test] fn paginates_sqlserver_top_parenthesized_projection_query_by_first_column() { let sql = "SELECT TOP (500) [id], [order_no], [store_id], [product_id], [customer_name], [quantity], [amount], [order_status], [created_at] FROM [sales].[orders_10k]"; let first_page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); let second_page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); assert!(first_page.ok); assert!(second_page.ok); assert_eq!( first_page.sql.unwrap(), "SELECT TOP (100) * FROM (SELECT TOP (500) [id], [order_no], [store_id], [product_id], [customer_name], [quantity], [amount], [order_status], [created_at] FROM [sales].[orders_10k]) [dbx_page] ORDER BY [id];" ); assert_eq!( second_page.sql.unwrap(), "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY [id]) AS [__dbx_row_num] FROM (SELECT TOP (500) [id], [order_no], [store_id], [product_id], [customer_name], [quantity], [amount], [order_status], [created_at] FROM [sales].[orders_10k]) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];" ); } #[test] fn paginates_sqlserver_top_query_with_qualified_first_projection_by_exposed_column() { let sql = "select top 1 t1.FSUPPLIERID, kh.FNAME gysnm from GDWORKOUT t1 join SUPPLIER kh on t1.FSUPPLIERID=kh.FITEMID join GDWORKOUTS t2 on t1.fid=t2.fid"; let first_page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); let second_page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); assert!(first_page.ok); assert!(second_page.ok); assert_eq!( first_page.sql.unwrap(), "SELECT TOP (100) * FROM (select top 1 t1.FSUPPLIERID, kh.FNAME gysnm from GDWORKOUT t1 join SUPPLIER kh on t1.FSUPPLIERID=kh.FITEMID join GDWORKOUTS t2 on t1.fid=t2.fid) [dbx_page] ORDER BY [FSUPPLIERID];" ); assert_eq!( second_page.sql.unwrap(), "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY [FSUPPLIERID]) AS [__dbx_row_num] FROM (select top 1 t1.FSUPPLIERID, kh.FNAME gysnm from GDWORKOUT t1 join SUPPLIER kh on t1.FSUPPLIERID=kh.FITEMID join GDWORKOUTS t2 on t1.fid=t2.fid) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];" ); } #[test] fn paginates_sqlserver_top_query_with_dotted_bracket_identifier() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT TOP 1000 [order.id] FROM TicketInfo".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT TOP (100) * FROM (SELECT TOP 1000 [order.id] FROM TicketInfo) [dbx_page] ORDER BY [order.id];" ); } #[test] fn paginates_sqlserver_top_query_after_leading_comment_by_first_column() { let sql = "-- 测试\nSELECT TOP (500) [id], [order_no], [store_id], [product_id], [customer_name], [quantity], [amount], [order_status], [created_at] FROM [sales].[orders_10k]"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT dbx_page_source.*, ROW_NUMBER() OVER (ORDER BY [id]) AS [__dbx_row_num] FROM (SELECT TOP (500) [id], [order_no], [store_id], [product_id], [customer_name], [quantity], [amount], [order_status], [created_at] FROM [sales].[orders_10k]) dbx_page_source) dbx_page WHERE [__dbx_row_num] > 100 AND [__dbx_row_num] <= 200 ORDER BY [__dbx_row_num];" ); } #[test] fn sqlserver_later_page_keeps_derived_table_closing_after_comment_before_order_by() { let sql = "SELECT\n*\nFROM\ncode\n-- WHERE\n-- ccode = '1002'\nORDER BY\nccode;"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); let generated = result.sql.expect("build SQL Server page SQL"); assert!(generated.contains("FROM (SELECT\n*\nFROM\ncode\n-- WHERE\n-- ccode = '1002'\n) dbx_page_source")); assert!(!generated.contains("-- ccode = '1002') dbx_page_source")); } #[test] fn sqlserver_later_page_keeps_row_number_window_after_comment_after_order_by() { let sql = "SELECT\n*\nFROM\ncode\nORDER BY\nccode -- sort"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); let generated = result.sql.expect("build SQL Server page SQL"); assert!(generated.contains("ROW_NUMBER() OVER (ORDER BY\nccode -- sort\n) AS [__dbx_row_num]")); assert!(!generated.contains("-- sort) AS [__dbx_row_num]")); } #[test] fn sqlserver_cte_pagination_plan_executes_original_sql() { let sql = ";WITH ranked AS (SELECT id FROM dbo.users) SELECT * FROM ranked".to_string(); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.clone(), query_base_sql: sql.clone(), database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, sql); assert!(plan.page_sql.is_none()); assert!(plan.count_sql.is_none()); assert_eq!(plan.page_limit, None); assert_eq!(plan.page_offset, None); } #[test] fn sqlserver_cte_after_leading_comments_executes_original_sql() { for sql in [ "-- heading comment\nWITH staff AS (SELECT 1 AS id) SELECT id FROM staff;", "/* heading comment */\nWITH staff AS (SELECT 1 AS id) SELECT id FROM staff;", ";-- heading comment\nWITH staff AS (SELECT 1 AS id) SELECT id FROM staff;", "-- heading comment\n;WITH staff AS (SELECT 1 AS id) SELECT id FROM staff;", ] { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::SqlServer), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, sql); assert!(plan.page_sql.is_none()); assert!(plan.count_sql.is_none()); assert_eq!(plan.page_limit, None); assert_eq!(plan.page_offset, None); } } #[test] fn sqlserver_cte_count_query_is_not_wrapped_as_derived_table() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: ";WITH cte AS (SELECT 1 AS id) SELECT * FROM cte".to_string(), database_type: Some(DatabaseType::SqlServer), }); assert!(!result.ok); assert!(result.sql.is_none()); } #[test] fn postgres_cte_update_is_not_paginated() { let sql = r#" WITH available AS ( SELECT id FROM app_users WHERE deleted_at IS NULL ), picked AS ( SELECT id FROM available ORDER BY random() LIMIT 10 ) UPDATE app_users AS u SET subscription_type = 1 FROM picked WHERE u.id = picked.id; "#; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Postgres), limit: 100, offset: 0, }); assert_eq!(result, err("not_select")); } #[test] fn postgres_cte_update_pagination_plan_executes_original_sql() { let sql = "WITH picked AS (SELECT id FROM app_users LIMIT 10) UPDATE app_users SET subscription_type = 1 FROM picked WHERE app_users.id = picked.id".to_string(); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.clone(), query_base_sql: sql.clone(), database_type: Some(DatabaseType::Postgres), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, sql); assert!(plan.page_sql.is_none()); assert!(plan.count_sql.is_none()); assert_eq!(plan.page_limit, None); assert_eq!(plan.page_offset, None); } #[test] fn postgres_cte_select_still_paginates() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "WITH picked AS (SELECT id FROM app_users LIMIT 10) SELECT * FROM picked".to_string(), database_type: Some(DatabaseType::Postgres), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "WITH picked AS (SELECT id FROM app_users LIMIT 10) SELECT * FROM picked LIMIT 100;" ); } #[test] fn clickhouse_scalar_with_select_is_paginated() { let sql = "WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept"; let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::ClickHouse), limit: 50, offset: 100, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept LIMIT 50 OFFSET 100;" ); } #[test] fn clickhouse_settings_clause_is_paginated_before_settings() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM system.clusters SETTINGS max_execution_time = 0".to_string(), database_type: Some(DatabaseType::ClickHouse), limit: 50, offset: 100, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT * FROM system.clusters LIMIT 50 OFFSET 100 SETTINGS max_execution_time = 0;" ); } #[test] fn clickhouse_settings_table_is_not_treated_as_settings_clause() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM system.settings".to_string(), database_type: Some(DatabaseType::ClickHouse), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT * FROM system.settings LIMIT 100;"); } #[test] fn clickhouse_settings_identifier_is_not_treated_as_settings_clause() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT settings FROM (SELECT 1 AS settings)".to_string(), database_type: Some(DatabaseType::ClickHouse), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT settings FROM (SELECT 1 AS settings) LIMIT 100;"); } #[test] fn clickhouse_settings_identifier_keeps_trailing_settings_clause() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT settings FROM (SELECT 1 AS settings) SETTINGS max_threads = 1".to_string(), database_type: Some(DatabaseType::ClickHouse), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT settings FROM (SELECT 1 AS settings) LIMIT 100 SETTINGS max_threads = 1;" ); } #[test] fn clickhouse_query_plan_places_limit_before_settings() { let sql = "SELECT * FROM system.clusters SETTINGS max_execution_time = 0"; let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::ClickHouse), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, "SELECT * FROM system.clusters LIMIT 100 SETTINGS max_execution_time = 0;"); assert_eq!( plan.page_sql, Some("SELECT * FROM system.clusters LIMIT 100 SETTINGS max_execution_time = 0;".to_string()) ); assert_eq!(plan.page_limit, Some(100)); assert_eq!(plan.page_offset, Some(0)); } #[test] fn mysql_for_update_places_limit_before_locking_clause() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM `test`\nwhere id=1 for update".to_string(), database_type: Some(DatabaseType::Mysql), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT * FROM `test`\nwhere id=1 LIMIT 100 for update;"); } #[test] fn locking_query_plan_keeps_server_pagination_and_count() { let sql = "SELECT * FROM `test`\nwhere id=1 for update".to_string(); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.clone(), query_base_sql: sql.clone(), database_type: Some(DatabaseType::Mysql), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, "SELECT * FROM `test`\nwhere id=1 LIMIT 100 for update;"); assert_eq!(plan.page_sql, Some("SELECT * FROM `test`\nwhere id=1 LIMIT 100 for update;".to_string())); assert_eq!(plan.page_limit, Some(100)); assert_eq!(plan.page_offset, Some(0)); assert_eq!( plan.count_sql, Some("SELECT COUNT(*) AS dbx_total_rows FROM (SELECT * FROM `test`\nwhere id=1) `dbx_count`;".to_string()) ); assert!(!plan.use_agent_result_session); } #[test] fn locking_query_later_page_places_offset_before_locking_clause() { let sql = "SELECT * FROM t WHERE deleted = 0 FOR UPDATE SKIP LOCKED".to_string(); let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.clone(), query_base_sql: sql.clone(), database_type: Some(DatabaseType::Mysql), pagination: QueryPagination { limit: 100, offset: 100, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!( plan.sql_to_execute, "SELECT * FROM t WHERE deleted = 0 LIMIT 100 OFFSET 100 FOR UPDATE SKIP LOCKED;" ); assert_eq!(plan.page_limit, Some(100)); assert_eq!(plan.page_offset, Some(100)); } #[test] fn locking_query_count_removes_locking_clause() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT * FROM t WHERE deleted = 0 FOR UPDATE".to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql, Some("SELECT COUNT(*) AS dbx_total_rows FROM (SELECT * FROM t WHERE deleted = 0) `dbx_count`;".to_string()) ); } #[test] fn locking_query_count_preserves_postgres_limit_after_lock_by_declining_rewrite() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT * FROM t FOR UPDATE LIMIT 10".to_string(), database_type: Some(DatabaseType::Postgres), }); assert_eq!(result, err("locking")); } #[test] fn places_limit_before_supported_top_level_locking_clause_variants() { for (sql, expected) in [ ("SELECT * FROM t FOR UPDATE", "SELECT * FROM t LIMIT 100 FOR UPDATE;"), ("SELECT * FROM t FOR SHARE", "SELECT * FROM t LIMIT 100 FOR SHARE;"), ("SELECT * FROM t FOR KEY SHARE", "SELECT * FROM t LIMIT 100 FOR KEY SHARE;"), ("SELECT * FROM t FOR NO KEY UPDATE", "SELECT * FROM t LIMIT 100 FOR NO KEY UPDATE;"), ("SELECT * FROM t LOCK IN SHARE MODE", "SELECT * FROM t LIMIT 100 LOCK IN SHARE MODE;"), ] { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Postgres), limit: 100, offset: 0, }); assert_eq!(result.sql.as_deref(), Some(expected), "{sql}"); } } #[test] fn nested_for_update_does_not_block_outer_limit_append() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM (SELECT id FROM t FOR UPDATE) locked".to_string(), database_type: Some(DatabaseType::Mysql), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT * FROM (SELECT id FROM t FOR UPDATE) locked LIMIT 100;"); } #[test] fn for_xml_is_not_treated_as_locking_clause() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id, name FROM users FOR XML PATH('row')".to_string(), database_type: Some(DatabaseType::Mysql), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT id, name FROM users FOR XML PATH('row') LIMIT 100;"); } #[test] fn ordinary_select_still_appends_limit() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM t WHERE deleted = 0".to_string(), database_type: Some(DatabaseType::Mysql), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT * FROM t WHERE deleted = 0 LIMIT 100;"); } #[test] fn locking_keywords_inside_postgres_dollar_quote_are_not_rewritten() { let sql = "SELECT $$FOR UPDATE$$ AS message"; let paginated = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Postgres), limit: 100, offset: 0, }); let counted = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Postgres), }); assert_eq!(paginated.sql.as_deref(), Some("SELECT $$FOR UPDATE$$ AS message LIMIT 100;")); assert_eq!( counted.sql.as_deref(), Some("SELECT COUNT(*) AS dbx_total_rows FROM (SELECT $$FOR UPDATE$$ AS message) \"dbx_count\";") ); } #[test] fn locking_keywords_inside_mysql_hash_comment_are_not_rewritten() { let sql = "SELECT * FROM t\n# FOR UPDATE LIMIT 1"; let paginated = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Mysql), limit: 100, offset: 0, }); let counted = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!(paginated.sql.as_deref(), Some("SELECT * FROM t\n# FOR UPDATE LIMIT 1\nLIMIT 100;")); assert_eq!( counted.sql.as_deref(), Some("SELECT COUNT(*) AS dbx_total_rows FROM (SELECT * FROM t\n# FOR UPDATE LIMIT 1\n) `dbx_count`;") ); } #[test] fn clickhouse_scalar_with_select_can_be_counted() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept" .to_string(), database_type: Some(DatabaseType::ClickHouse), }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept) `dbx_count`;" ); } #[test] fn clickhouse_scalar_with_update_is_not_paginated() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "WITH 1 AS min_id UPDATE employees SET dept = 'sales' WHERE id = min_id".to_string(), database_type: Some(DatabaseType::ClickHouse), limit: 50, offset: 0, }); assert_eq!(result, err("not_select")); } #[test] fn wraps_sqlserver_select_with_unnamed_column() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT @@version".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT TOP (100) @@version"); } #[test] 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), limit: 100, offset: 300, }); assert_eq!( result.sql.unwrap(), "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 paginates_sqlserver_wildcard_later_pages_without_derived_tables() { 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" .to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 100, }); let sql = result.sql.expect("build wildcard page"); assert!(sql.starts_with("EXEC sys.sp_executesql N'SET ROWCOUNT 200; SELECT b.ProjectType,*")); assert_eq!(sqlserver_result_offset(&sql), 100); } #[test] fn sqlserver_first_page_preserves_existing_pagination_or_injects_top() { for (original_sql, expected_sql) in [ ( "SELECT * FROM TABLE_NAME ORDER BY id OFFSET 1 ROWS FETCH NEXT 10 ROWS ONLY", "SELECT * FROM TABLE_NAME ORDER BY id OFFSET 1 ROWS FETCH NEXT 10 ROWS ONLY", ), ("SELECT id FROM TABLE_NAME", "SELECT TOP (100) id FROM TABLE_NAME"), ] { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert!(result.ok, "{original_sql}"); assert_eq!(result.sql.as_deref(), Some(expected_sql), "{original_sql}"); } } // 临时表 #tmp 会让词法扫描器把同一行后面的 OFFSET/FETCH 当成注释丢掉, // 必须靠 AST 检测拦住 TOP 注入,否则 SQL Server 报“TOP 不能与 OFFSET 同用”。 #[test] fn keeps_sqlserver_offset_fetch_next_with_temp_table_on_same_line() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM #tmp ORDER BY id OFFSET 0 ROWS FETCH NEXT 62000 ROWS ONLY".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 500, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT * FROM #tmp ORDER BY id OFFSET 0 ROWS FETCH NEXT 62000 ROWS ONLY"); } // 全局临时表 ##tmp 同样不能被注入 TOP。 #[test] fn keeps_sqlserver_offset_fetch_next_with_global_temp_table() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM ##tmp ORDER BY id OFFSET 0 ROWS FETCH NEXT 62000 ROWS ONLY".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 500, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT * FROM ##tmp ORDER BY id OFFSET 0 ROWS FETCH NEXT 62000 ROWS ONLY"); } // 只写 OFFSET 不写 FETCH NEXT 也是 SQL Server 合法分页写法,同样不能注入 TOP。 #[test] fn keeps_sqlserver_offset_without_fetch_next() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM t ORDER BY id OFFSET 0 ROWS".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 500, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT * FROM t ORDER BY id OFFSET 0 ROWS"); } // 字符串里的反斜杠会让词法扫描器引号配对错乱,AST 检测需补位。 #[test] fn keeps_sqlserver_offset_fetch_next_with_backslash_string() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM t WHERE p = 'C:\\' ORDER BY id OFFSET 0 ROWS FETCH NEXT 10 ROWS ONLY" .to_string(), database_type: Some(DatabaseType::SqlServer), limit: 500, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT * FROM t WHERE p = 'C:\\' ORDER BY id OFFSET 0 ROWS FETCH NEXT 10 ROWS ONLY" ); } // UNION 顶层的 OFFSET/FETCH 也不能被注入 TOP。 #[test] fn keeps_sqlserver_offset_fetch_next_after_union() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM a UNION SELECT id FROM b ORDER BY id OFFSET 5 ROWS FETCH NEXT 10 ROWS ONLY" .to_string(), database_type: Some(DatabaseType::SqlServer), limit: 500, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT id FROM a UNION SELECT id FROM b ORDER BY id OFFSET 5 ROWS FETCH NEXT 10 ROWS ONLY" ); } #[test] fn oracle_pagination_skips_sql_clause() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users".to_string(), database_type: Some(DatabaseType::Oracle), limit: 100, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT id FROM users;"); } #[test] fn oceanbase_oracle_pagination_wraps_with_rownum() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users ORDER BY id".to_string(), database_type: Some(DatabaseType::OceanbaseOracle), limit: 100, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT * FROM (SELECT id FROM users ORDER BY id) WHERE ROWNUM <= 100;"); } #[test] fn oceanbase_oracle_pagination_wraps_offset_with_rownum_bounds() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users ORDER BY id".to_string(), database_type: Some(DatabaseType::OceanbaseOracle), limit: 100, offset: 200, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT dbx_inner.*, ROWNUM AS \"__dbx_row_num\" FROM (SELECT id FROM users ORDER BY id) dbx_inner WHERE ROWNUM <= 300) WHERE \"__dbx_row_num\" > 200;" ); } #[test] fn uses_fetch_first_pagination_for_db2() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users".to_string(), database_type: Some(DatabaseType::Db2), limit: 100, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT id FROM users FETCH FIRST 100 ROWS ONLY;"); } #[test] fn uses_skip_first_pagination_for_informix() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users WHERE active = 1".to_string(), database_type: Some(DatabaseType::Informix), limit: 50, offset: 100, }); assert_eq!(result.sql.unwrap(), "SELECT SKIP 100 FIRST 50 id FROM users WHERE active = 1;"); } #[test] fn informix_pagination_keeps_existing_first() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT FIRST 20 id FROM users".to_string(), database_type: Some(DatabaseType::Informix), limit: 50, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT FIRST 20 id FROM users;"); } #[test] fn uses_rows_pagination_for_firebird() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users WHERE active = 1 ORDER BY id".to_string(), database_type: Some(DatabaseType::Firebird), limit: 50, offset: 100, }); assert_eq!(result.sql.unwrap(), "SELECT id FROM users WHERE active = 1 ORDER BY id ROWS 101 TO 150;"); } #[test] fn firebird_pagination_keeps_existing_rows_clause() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users ROWS 20".to_string(), database_type: Some(DatabaseType::Firebird), limit: 50, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT id FROM users ROWS 20;"); } #[test] fn uses_mysql_style_alias_for_pagination() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users WHERE active = 1".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT id FROM users WHERE active = 1 LIMIT 50;"); } #[test] fn mysql_pagination_preserves_leading_ampersand_routing_hint() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "/*& tenant:'test' */\nSELECT id FROM users".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(result.sql.unwrap(), "/*& tenant:'test' */\nSELECT id FROM users LIMIT 50;"); } #[test] fn mysql_pagination_preserves_supported_leading_execution_hints_only() { for hint in ["/*+ MAX_EXECUTION_TIME(1000) */", "/*@global:true*/", "/*& tenant:'test' */"] { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: format!("{hint}\nSELECT id FROM users"), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(result.sql.unwrap(), format!("{hint}\nSELECT id FROM users LIMIT 50;")); } let ordinary_comment = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "/* report query */\nSELECT id FROM users".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(ordinary_comment.sql.unwrap(), "SELECT id FROM users LIMIT 50;"); let ordinary_comment_before_hint = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "-- report query\n/*& tenant:'test' */\nSELECT id FROM users".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(ordinary_comment_before_hint.sql.unwrap(), "/*& tenant:'test' */\nSELECT id FROM users LIMIT 50;"); } #[test] fn mysql_pagination_preserves_same_line_tdsql_directives() { for directive in ["/*sets:allsets*/", "/*master*/", "/*slave:set_1781591902_7*/", "/*future-route:anywhere*/"] { let first_page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: format!("{directive}select @@server_id;"), database_type: Some(DatabaseType::Mysql), limit: 100, offset: 0, }); let later_page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: format!("{directive}select @@server_id;"), database_type: Some(DatabaseType::Mysql), limit: 100, offset: 100, }); assert_eq!(first_page.sql.unwrap(), format!("{directive}select @@server_id LIMIT 100;")); assert_eq!(later_page.sql.unwrap(), format!("{directive}select @@server_id LIMIT 100 OFFSET 100;")); } } #[test] fn mysql_pagination_preserves_exact_proxy_directive() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "/*proxy*/\nSELECT id FROM users".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(result.sql.unwrap(), "/*proxy*/\nSELECT id FROM users LIMIT 50;"); } #[test] fn native_mysql_pagination_plan_preserves_exact_issue_directive() { let original_sql = "/*sets:allsets*/select @@server_id;"; let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: original_sql.to_string(), query_base_sql: original_sql.to_string(), database_type: Some(DatabaseType::Mysql), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, "/*sets:allsets*/select @@server_id LIMIT 100;"); assert_eq!(plan.page_sql.as_deref(), Some(plan.sql_to_execute.as_str())); assert_eq!(plan.page_limit, Some(100)); assert_eq!(plan.page_offset, Some(0)); assert_eq!( plan.count_sql.as_deref(), Some("/*sets:allsets*/SELECT COUNT(*) AS dbx_total_rows FROM (select @@server_id) `dbx_count`;") ); assert!(!plan.use_agent_result_session); } #[test] fn mysql_count_preserves_leading_ampersand_routing_hint() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "/*& tenant:'test' */\nSELECT id FROM users".to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.unwrap(), "/*& tenant:'test' */\nSELECT COUNT(*) AS dbx_total_rows FROM (SELECT id FROM users) `dbx_count`;" ); } #[test] fn mysql_count_preserves_directives_at_outermost_start() { for prefix in [ "/*sets:allsets*/", "/*master*/", "/*slave:set_1781591902_7*/", "/*future-route:anywhere*/", "/*proxy*/\n", "/*+ MAX_EXECUTION_TIME(1000) */\n", "/*@global:true*/\n", "/*& tenant:'test' */\n", ] { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: format!("{prefix}select @@server_id;"), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.unwrap(), format!("{prefix}SELECT COUNT(*) AS dbx_total_rows FROM (select @@server_id) `dbx_count`;") ); } } #[test] fn mysql_generated_queries_keep_standalone_comments_non_executable() { for original_sql in ["/* report query */\nSELECT id FROM users", "-- report query\nSELECT id FROM users"] { let page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); let count = build_count_query_sql(CountQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(DatabaseType::Mysql), }); let sort = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(DatabaseType::Mysql), result_columns: vec!["id".to_string()], column_index: 0, column: "id".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!(page.sql.unwrap(), "SELECT id FROM users LIMIT 50;"); assert_eq!( count.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT id FROM users) `dbx_count`;" ); assert_eq!(sort, err("multi")); } } #[test] fn tdsql_directives_remain_mysql_only_in_generated_queries() { let original_sql = "/*sets:allsets*/SELECT id FROM users"; let page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(DatabaseType::Postgres), limit: 50, offset: 0, }); let count = build_count_query_sql(CountQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(DatabaseType::Postgres), }); assert_eq!(page.sql.unwrap(), "SELECT id FROM users LIMIT 50;"); assert_eq!(count.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT id FROM users) \"dbx_count\";"); } #[test] fn mysql_generated_queries_reject_non_executable_and_multi_statement_inputs() { for original_sql in ["/*sets:allsets*/", "/*sets:allsets SELECT id FROM users"] { let page = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); let count = build_count_query_sql(CountQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(DatabaseType::Mysql), }); assert!(!page.ok, "{original_sql}"); assert!(page.sql.is_none(), "{original_sql}"); assert!(!count.ok, "{original_sql}"); assert!(count.sql.is_none(), "{original_sql}"); } let multi = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "/*sets:allsets*/SELECT 1; SELECT 2".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(multi, err("multi")); } #[test] fn mysql_pagination_does_not_wrap_duplicate_result_columns() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT p.id, t.id FROM table1 p LEFT JOIN table2 t ON p.f = t.f".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 100, }); assert_eq!( result.sql.unwrap(), "SELECT p.id, t.id FROM table1 p LEFT JOIN table2 t ON p.f = t.f LIMIT 50 OFFSET 100;" ); } #[test] fn mysql_pagination_keeps_limit_outside_trailing_line_comment() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT 1 AS id\n-- tail comment".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT 1 AS id\n-- tail comment\nLIMIT 50;"); } #[test] fn mysql_pagination_keeps_limit_outside_trailing_hash_comment() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT 1 AS id\n# tail comment".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT 1 AS id\n# tail comment\nLIMIT 50;"); } #[test] fn mysql_pagination_keeps_existing_top_level_limit() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users LIMIT 20;".to_string(), database_type: Some(DatabaseType::Mysql), limit: 50, offset: 0, }); assert_eq!(result.sql.unwrap(), "SELECT id FROM users LIMIT 20;"); } #[test] fn mysql_pagination_wraps_wide_existing_limit_on_first_page() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM dy_promotion_item WHERE create_time < '2026-06-01' LIMIT 10000;".to_string(), database_type: Some(DatabaseType::Mysql), limit: 500, offset: 0, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT * FROM dy_promotion_item WHERE create_time < '2026-06-01' LIMIT 10000) `dbx_page` LIMIT 500 OFFSET 0;" ); } #[test] fn mysql_pagination_wraps_comma_limit_row_count_on_first_page() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM users LIMIT 20, 10000;".to_string(), database_type: Some(DatabaseType::Mysql), limit: 500, offset: 0, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT * FROM users LIMIT 20, 10000) `dbx_page` LIMIT 500 OFFSET 0;" ); } #[test] fn mysql_pagination_wraps_existing_limit_for_later_pages() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * FROM dy_promotion_item WHERE create_time < '2026-06-01' LIMIT 10000;".to_string(), database_type: Some(DatabaseType::Mysql), limit: 1000, offset: 1000, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT * FROM dy_promotion_item WHERE create_time < '2026-06-01' LIMIT 10000) `dbx_page` LIMIT 1000 OFFSET 1000;" ); } #[test] fn standard_pagination_wraps_existing_limit_for_later_pages() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id FROM users LIMIT 20;".to_string(), database_type: Some(DatabaseType::Postgres), limit: 5, offset: 10, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT id FROM users LIMIT 20) \"dbx_page\" LIMIT 5 OFFSET 10;" ); } #[test] fn rejects_multiple_statements_for_pagination() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT 1; SELECT 2;".to_string(), database_type: Some(DatabaseType::Postgres), limit: 100, offset: 0, }); assert_eq!(result, err("multi")); } #[test] fn rejects_select_into_for_pagination() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT * INTO copy_users FROM users WHERE active = 1".to_string(), database_type: Some(DatabaseType::SqlServer), limit: 100, offset: 0, }); assert_eq!(result, err("not_select")); } #[test] fn builds_count_query() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "WITH cte AS (SELECT 1 AS id) SELECT * FROM cte".to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (WITH cte AS (SELECT 1 AS id) SELECT * FROM cte) `dbx_count`;" ); } #[test] fn hive_count_keeps_cte_outside_derived_table() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "WITH cte AS (SELECT 1 AS id) SELECT * FROM cte".to_string(), database_type: Some(DatabaseType::Hive), }); assert_eq!( result.sql.unwrap(), "WITH cte AS (SELECT 1 AS id) SELECT COUNT(*) AS dbx_total_rows FROM (SELECT * FROM cte) `dbx_count`;" ); } #[test] fn hive_count_keeps_from_style_query_inside_derived_table() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "WITH cte AS (SELECT 1 AS id) FROM cte SELECT *".to_string(), database_type: Some(DatabaseType::Hive), }); assert_eq!( result.sql.unwrap(), "WITH cte AS (SELECT 1 AS id) SELECT COUNT(*) AS dbx_total_rows FROM (FROM cte SELECT *) `dbx_count`;" ); } #[test] fn mysql_count_rewrites_ambiguous_join_projection() { for sql in [ "SELECT a.*, b.* FROM a JOIN b ON b.a_id = a.id ORDER BY b.id", "SELECT * FROM a JOIN b ON b.a_id = a.id ORDER BY b.id", ] { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.as_deref(), Some( "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT 1 AS dbx_count_value FROM a JOIN b ON b.a_id = a.id) `dbx_count`;" ), "{sql}" ); } } #[test] fn mysql_count_rewrites_duplicate_explicit_names() { for sql in [ "SELECT a.id AS id, b.id AS id FROM a JOIN b ON b.a_id = a.id", "SELECT a.`1111`, b.`1111` FROM a JOIN b ON b.a_id = a.id", ] { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.as_deref(), Some( "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT 1 AS dbx_count_value FROM a JOIN b ON b.a_id = a.id) `dbx_count`;" ), "{sql}" ); } } #[test] fn mysql_count_keeps_unique_projection_wrapper() { let sql = "SELECT a.id AS a_id, b.id AS b_id FROM a JOIN b ON b.a_id = a.id ORDER BY b.id"; let result = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.as_deref(), Some( "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT a.id AS a_id, b.id AS b_id FROM a JOIN b ON b.a_id = a.id ORDER BY b.id) `dbx_count`;" ) ); } #[test] fn mysql_count_rewrite_preserves_limit_and_offset() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT a.*, b.* FROM a JOIN b ON b.a_id = a.id ORDER BY FIELD(b.id, 11, 10) LIMIT 2 OFFSET 1" .to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.as_deref(), Some( "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT 1 AS dbx_count_value FROM a JOIN b ON b.a_id = a.id LIMIT 2 OFFSET 1) `dbx_count`;" ) ); } #[test] fn mysql_count_rejects_ambiguous_cardinality_dependent_queries() { for sql in [ "SELECT DISTINCT a.id AS id, b.id AS id FROM a JOIN b ON b.a_id = a.id", "SELECT a.id AS id, b.id AS id FROM a JOIN b ON b.a_id = a.id GROUP BY a.id, b.id", "SELECT a.id AS id, b.id AS id FROM a JOIN b ON b.a_id = a.id HAVING id > 0", "SELECT a.id AS id, b.id AS id FROM a JOIN b ON b.a_id = a.id UNION ALL SELECT c.id AS id, d.id AS id FROM c JOIN d ON d.c_id = c.id", "SELECT COUNT(*) AS id, SUM(a.id) AS id FROM a", ] { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!(result, err("unsupported"), "{sql}"); } } #[test] fn mysql_count_rejects_ambiguous_select_into() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT a.id AS id, b.id AS id INTO OUTFILE 'dump.tsv' FROM a JOIN b ON b.a_id = a.id" .to_string(), database_type: Some(DatabaseType::Mysql), }); assert!(!result.ok); assert!(result.sql.is_none()); } #[test] fn mysql_count_keeps_parse_failure_fallback() { let sql = "SELECT a.id AS id, FROM a"; let result = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.as_deref(), Some("SELECT COUNT(*) AS dbx_total_rows FROM (SELECT a.id AS id, FROM a) `dbx_count`;") ); } #[test] fn non_mysql_count_keeps_ambiguous_projection_wrapper() { let sql = "SELECT a.*, b.* FROM a JOIN b ON b.a_id = a.id ORDER BY b.id"; for database_type in [DatabaseType::Postgres, DatabaseType::Sqlite, DatabaseType::ClickHouse] { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(database_type), }); let alias = quote_table_identifier(Some(database_type), "dbx_count"); assert_eq!( result.sql, Some(format!("SELECT COUNT(*) AS dbx_total_rows FROM ({sql}) {alias};")), "{database_type:?}" ); } } #[test] fn count_query_keeps_wrapper_outside_trailing_line_comment() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT 1 AS id\n-- tail comment".to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT 1 AS id\n-- tail comment\n) `dbx_count`;" ); } #[test] fn count_query_keeps_wrapper_outside_trailing_hash_comment() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT 1 AS id\n# tail comment".to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT 1 AS id\n# tail comment\n) `dbx_count`;" ); } #[test] fn count_query_preserves_user_limit() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT * FROM dy_promotion_item WHERE create_time < '2026-06-01' LIMIT 10000".to_string(), database_type: Some(DatabaseType::Mysql), }); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT * FROM dy_promotion_item WHERE create_time < '2026-06-01' LIMIT 10000) `dbx_count`;" ); } #[test] fn count_query_preserves_user_limit_offset() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT * FROM users WHERE active = 1 LIMIT 100 OFFSET 50".to_string(), database_type: Some(DatabaseType::Postgres), }); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT * FROM users WHERE active = 1 LIMIT 100 OFFSET 50) \"dbx_count\";" ); } #[test] fn iotdb_tree_count_rewrites_single_series_without_derived_table() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT WGEN_GnTmpSta1 FROM root.dbx_time_preview.device WHERE time >= 1 ORDER BY time DESC" .to_string(), database_type: Some(DatabaseType::Iotdb), }); assert_eq!( result.sql.as_deref(), Some("SELECT COUNT(WGEN_GnTmpSta1) AS dbx_total_rows FROM root.dbx_time_preview.device WHERE time >= 1;") ); } #[test] fn iotdb_count_allows_value_filters_on_selected_series() { for sql in [ "SELECT s1 FROM root.db.device WHERE s1 > 10", "SELECT root.db.device.s1 FROM root.db.device WHERE root.db.device.s1 > 10", ] { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Iotdb), }); assert!(result.sql.is_some(), "{sql}"); } } #[test] fn iotdb_count_declines_queries_without_safe_tree_row_semantics() { for sql in [ "SELECT * FROM root.db.device", "SELECT s1, s2 FROM root.db.device", "SELECT COUNT(s1) FROM root.db.device", "SELECT s1 FROM root.db.device GROUP BY ([0, 100), 10ms)", "SELECT s1 FROM root.db.device WHERE s2 > 10", "SELECT s1 FROM root.db.device WHERE root.db.device.s2 > 10", "SELECT s1 FROM root.db.device LIMIT 10", "SELECT s1 FROM root.db.device, root.db.other", "SELECT value FROM metrics", ] { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: sql.to_string(), database_type: Some(DatabaseType::Iotdb), }); assert_eq!(result, err("unsupported"), "{sql}"); } } #[test] fn iotdb_pagination_plan_uses_tree_model_count_query() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT WGEN_GnTmpSta1 FROM root.dbx_time_preview.device".to_string(), query_base_sql: "SELECT WGEN_GnTmpSta1 FROM root.dbx_time_preview.device".to_string(), database_type: Some(DatabaseType::Iotdb), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!( plan.count_sql.as_deref(), Some("SELECT COUNT(WGEN_GnTmpSta1) AS dbx_total_rows FROM root.dbx_time_preview.device;") ); } #[test] fn iris_count_query_removes_top_level_order_by() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT id, appointment_time FROM patients WHERE status = ? ORDER BY appointment_time DESC" .to_string(), database_type: Some(DatabaseType::Iris), }); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT id, appointment_time FROM patients WHERE status = ?) dbx_count;" ); } #[test] fn iris_count_query_preserves_nested_order_by_and_parameters() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT * FROM (SELECT TOP ? id FROM visits WHERE status = ? ORDER BY created_at DESC) recent WHERE id > ? ORDER BY id" .to_string(), database_type: Some(DatabaseType::Iris), }); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT * FROM (SELECT TOP ? id FROM visits WHERE status = ? ORDER BY created_at DESC) recent WHERE id > ?) dbx_count;" ); } #[test] fn iris_count_query_without_order_by_is_unchanged() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT id FROM visits WHERE status = ?".to_string(), database_type: Some(DatabaseType::Iris), }); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT id FROM visits WHERE status = ?) dbx_count;" ); } #[test] fn iris_count_query_injects_count_for_legacy_cache_table() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT * FROM TATFY WHERE UpdateDate = '2026-07-30' AND UpdateTime > '08:00:00'".to_string(), database_type: Some(DatabaseType::Iris), }); assert_eq!( result.sql.unwrap(), "SELECT COUNT(*) AS dbx_total_rows FROM TATFY WHERE UpdateDate = '2026-07-30' AND UpdateTime > '08:00:00';" ); } #[test] fn iris_count_query_injects_count_for_simple_select() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "select * from PA_Adm".to_string(), database_type: Some(DatabaseType::Iris), }); assert_eq!(result.sql.as_deref(), Some("SELECT COUNT(*) AS dbx_total_rows FROM PA_Adm;")); } #[test] fn iris_count_query_preserves_simple_select_predicates_without_derived_alias() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT * FROM PA_Adm WHERE status = ? ORDER BY id DESC".to_string(), database_type: Some(DatabaseType::Iris), }); assert_eq!(result.sql.as_deref(), Some("SELECT COUNT(*) AS dbx_total_rows FROM PA_Adm WHERE status = ?;")); } #[test] fn iris_count_query_keeps_wrapper_for_top_limited_select() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT TOP 10 * FROM PA_Adm ORDER BY id DESC".to_string(), database_type: Some(DatabaseType::Iris), }); assert_eq!( result.sql.as_deref(), Some("SELECT COUNT(*) AS dbx_total_rows FROM (SELECT TOP 10 * FROM PA_Adm) dbx_count;") ); } #[test] fn iris_count_query_keeps_derived_wrapper_for_complex_selects() { let result = build_count_query_sql(CountQuerySqlOptions { original_sql: "SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id".to_string(), database_type: Some(DatabaseType::Iris), }); assert_eq!( result.sql.as_deref(), Some( "SELECT COUNT(*) AS dbx_total_rows FROM (SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id) dbx_count;" ) ); } #[test] fn builds_agent_cursor_pagination_plan() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events".to_string(), query_base_sql: "SELECT * FROM events".to_string(), database_type: Some(DatabaseType::Oracle), 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 events"); assert_eq!(plan.page_limit, Some(500)); assert_eq!(plan.page_offset, Some(0)); assert!(plan.page_sql.is_none()); assert!(plan.use_agent_result_session); } #[test] fn oracle_offset_jump_keeps_agent_result_session() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events".to_string(), query_base_sql: "SELECT * FROM events".to_string(), database_type: Some(DatabaseType::Oracle), pagination: QueryPagination { limit: 100, offset: 200, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); // Oracle user queries are Unbounded: the SQL stays unchanged, so the // client must consume the agent result session through to the offset. assert_eq!(plan.sql_to_execute, "SELECT * FROM events;"); assert_eq!(plan.page_limit, Some(100)); assert_eq!(plan.page_offset, Some(200)); assert!(plan.use_agent_result_session); } #[test] fn jdbc_offset_jump_keeps_agent_result_session() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events".to_string(), query_base_sql: "SELECT * FROM events".to_string(), database_type: Some(DatabaseType::Jdbc), pagination: QueryPagination { limit: 100, offset: 200, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert!(plan.use_agent_result_session); assert_eq!(plan.page_offset, Some(200)); } #[test] fn rewritten_offset_jump_does_not_use_agent_result_session() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events".to_string(), query_base_sql: "SELECT * FROM events".to_string(), database_type: Some(DatabaseType::Postgres), pagination: QueryPagination { limit: 100, offset: 200, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert!(plan.sql_to_execute.contains("LIMIT")); assert!(!plan.use_agent_result_session); } #[test] fn highgo_prefers_server_pagination_over_agent_cursor() { let first_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events".to_string(), query_base_sql: "SELECT * FROM events".to_string(), database_type: Some(DatabaseType::Highgo), pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!(first_page.sql_to_execute, "SELECT * FROM events LIMIT 500;"); assert_eq!(first_page.page_sql, Some(first_page.sql_to_execute.clone())); assert_eq!(first_page.page_limit, Some(500)); assert_eq!(first_page.page_offset, Some(0)); assert!(!first_page.use_agent_result_session); let second_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events".to_string(), query_base_sql: "SELECT * FROM events".to_string(), database_type: Some(DatabaseType::Highgo), pagination: QueryPagination { limit: 500, offset: 500, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!(second_page.sql_to_execute, "SELECT * FROM events LIMIT 500 OFFSET 500;"); assert_eq!(second_page.page_sql, Some(second_page.sql_to_execute.clone())); assert_eq!(second_page.page_limit, Some(500)); assert_eq!(second_page.page_offset, Some(500)); assert!(!second_page.use_agent_result_session); } #[test] fn xugu_prefers_server_pagination_over_agent_cursor() { let first_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events".to_string(), query_base_sql: "SELECT * FROM events".to_string(), database_type: Some(DatabaseType::Xugu), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!(first_page.sql_to_execute, "SELECT * FROM events LIMIT 100;"); assert_eq!(first_page.page_sql, Some(first_page.sql_to_execute.clone())); assert_eq!(first_page.page_limit, Some(100)); assert_eq!(first_page.page_offset, Some(0)); assert!(!first_page.use_agent_result_session); let second_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events".to_string(), query_base_sql: "SELECT * FROM events".to_string(), database_type: Some(DatabaseType::Xugu), pagination: QueryPagination { limit: 100, offset: 100, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!(second_page.sql_to_execute, "SELECT * FROM events LIMIT 100 OFFSET 100;"); assert_eq!(second_page.page_sql, Some(second_page.sql_to_execute.clone())); assert_eq!(second_page.page_limit, Some(100)); assert_eq!(second_page.page_offset, Some(100)); assert!(!second_page.use_agent_result_session); } #[test] fn xugu_keeps_agent_fallback_for_unrewritable_query() { let sql = "SELECT * FROM events; SELECT 1"; let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::Xugu), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, sql); assert!(plan.page_sql.is_none()); assert!(plan.use_agent_result_session); } #[test] fn xugu_top_clause_keeps_first_page_agent_cursor() { for sql in [ "SELECT TOP 10 * FROM events", "select top 10 * from events", "/* first page */ SELECT /* bounded */ TOP 10 * FROM events", ] { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::Xugu), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!(plan.sql_to_execute, sql); assert!(plan.page_sql.is_none()); assert_eq!(plan.page_limit, Some(100)); assert_eq!(plan.page_offset, Some(0)); assert!(plan.use_agent_result_session); assert!(!plan.single_execution); } } #[test] fn xugu_top_clause_without_agent_uses_single_execution() { let sql = "SELECT TOP 10 * FROM events"; let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::Xugu), pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: true, }); assert_eq!(plan.sql_to_execute, sql); assert!(plan.page_sql.is_none()); assert_eq!(plan.page_limit, Some(100)); assert_eq!(plan.page_offset, Some(0)); assert!(!plan.use_agent_result_session); assert!(plan.single_execution); } #[test] fn xugu_top_clause_rejects_sibling_limit_on_every_page() { for offset in [0, 100] { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT TOP 10 * FROM events".to_string(), database_type: Some(DatabaseType::Xugu), limit: 100, offset, }); assert!(!result.ok); assert!(result.sql.is_none()); assert_eq!(result.reason.as_deref(), Some("unsupported")); } } #[test] fn oceanbase_oracle_prefers_bounded_first_page_and_keeps_cursor_fallback() { let sql = "SELECT * FROM events"; let first_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::OceanbaseOracle), pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!(first_page.sql_to_execute, "SELECT * FROM (SELECT * FROM events) WHERE ROWNUM <= 500;"); assert_eq!(first_page.page_sql.as_deref(), Some(first_page.sql_to_execute.as_str())); assert_eq!(first_page.page_limit, Some(500)); assert_eq!(first_page.page_offset, Some(0)); assert!(!first_page.use_agent_result_session); assert!(first_page.count_sql.is_some()); assert!(first_page.pagination_row_number_column.is_none()); let next_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::OceanbaseOracle), pagination: QueryPagination { limit: 500, offset: 500, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert!(next_page.sql_to_execute.contains("ROWNUM <= 1000")); assert!(next_page.sql_to_execute.contains("\"__dbx_row_num\" > 500")); assert_eq!(next_page.page_sql.as_deref(), Some(next_page.sql_to_execute.as_str())); assert!(!next_page.use_agent_result_session); assert_eq!(next_page.pagination_row_number_column.as_deref(), Some("__dbx_row_num")); let unrewritable = "SELECT * FROM events; SELECT 1"; let fallback = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: unrewritable.to_string(), query_base_sql: unrewritable.to_string(), database_type: Some(DatabaseType::OceanbaseOracle), pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!(fallback.sql_to_execute, unrewritable); assert!(fallback.page_sql.is_none()); assert!(fallback.use_agent_result_session); } #[test] fn rownum_helper_avoids_names_in_the_original_query() { let sql = "SELECT LEVEL AS \"__dbx_row_num\" FROM DUAL CONNECT BY LEVEL <= 10"; let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::OceanbaseOracle), pagination: QueryPagination { limit: 5, offset: 5, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert_eq!(plan.pagination_row_number_column.as_deref(), Some("__dbx_row_num_1")); assert!(plan.sql_to_execute.contains("ROWNUM AS \"__dbx_row_num_1\"")); assert!(plan.sql_to_execute.contains("WHERE \"__dbx_row_num_1\" > 5")); assert!(plan.sql_to_execute.contains(sql)); } #[test] fn existing_rownum_keeps_a_cursor_and_does_not_mark_a_user_column_as_generated() { let sql = "SELECT ROWNUM AS \"__dbx_row_num\", name FROM events WHERE ROWNUM <= 1000"; let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::OceanbaseOracle), pagination: QueryPagination { limit: 500, offset: 500, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: false, }); assert!(plan.pagination_row_number_column.is_none()); assert!(plan.use_agent_result_session); assert_eq!(plan.page_offset, Some(500)); assert_eq!(plan.sql_to_execute, sql); } #[test] fn highgo_avoids_agent_cursor_for_unrewritable_queries() { let sql = "SELECT * FROM events; SELECT 1"; let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type: Some(DatabaseType::Highgo), 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); assert!(plan.page_sql.is_none()); assert!(plan.page_limit.is_none()); assert!(plan.page_offset.is_none()); assert!(!plan.use_agent_result_session); } #[test] fn kingbase_prefers_server_pagination_for_ordered_query() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events ORDER BY id".to_string(), query_base_sql: "SELECT * FROM events ORDER BY id".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 events ORDER BY id 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_uses_agent_cursor_for_unordered_query() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: "SELECT * FROM events".to_string(), query_base_sql: "SELECT * 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 * 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); } #[test] fn kingbase_falls_back_to_agent_cursor_for_unrewritable_queries() { let sql = "SELECT * FROM events; SELECT 1"; 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); 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); } #[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_unordered_subquery_top_uses_agent_cursor() { 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"); 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); } #[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 pagination_plan_exposes_only_safe_exact_row_bounds() { for (database_type, sql, expected) in [ (Some(DatabaseType::SqlServer), "SELECT TOP 100 * FROM events", Some(100)), (Some(DatabaseType::SqlServer), "SELECT TOP(50) * FROM events", Some(50)), (Some(DatabaseType::SqlServer), "SELECT TOP (250) * FROM events", Some(250)), (Some(DatabaseType::SqlServer), "SELECT TOP 10 PERCENT * FROM events", None), (Some(DatabaseType::SqlServer), "SELECT TOP (2) WITH TIES * FROM events", None), (Some(DatabaseType::SqlServer), "SELECT TOP (100 + 1) * FROM events", None), (Some(DatabaseType::SqlServer), "SELECT TOP (@row_count) * FROM events", None), (Some(DatabaseType::SqlServer), "SELECT * FROM (SELECT TOP 5 * FROM events) t", None), (Some(DatabaseType::SqlServer), "SELECT * FROM events", None), (Some(DatabaseType::Postgres), "SELECT TOP 100 * FROM events", None), (Some(DatabaseType::Kingbase), "SELECT TOP 100 * FROM events", None), // The standard LIMIT/OFFSET dialects (MySQL, Postgres, and every // other database that isn't given its own pagination_strategy // arm) bound the query the same way SQL Server's TOP does. (Some(DatabaseType::Postgres), "SELECT * FROM events LIMIT 100", Some(100)), (Some(DatabaseType::Mysql), "SELECT * FROM events LIMIT 50", Some(50)), (Some(DatabaseType::Mysql), "SELECT * FROM events LIMIT 100 OFFSET 100", Some(100)), (Some(DatabaseType::ClickHouse), "SELECT * FROM events LIMIT 10 BY user_id", None), (Some(DatabaseType::ClickHouse), "SELECT * FROM events LIMIT 10 OFFSET 5 BY user_id", None), (Some(DatabaseType::ClickHouse), "SELECT * FROM events LIMIT 10 WITH TIES", None), (Some(DatabaseType::DuckDb), "SELECT * FROM events LIMIT 10%", None), (Some(DatabaseType::DuckDb), "SELECT * FROM events LIMIT 10 PERCENT", None), (None, "SELECT * FROM events LIMIT 100", Some(100)), (Some(DatabaseType::Postgres), "SELECT * FROM events", None), ] { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { sql: sql.to_string(), query_base_sql: sql.to_string(), database_type, pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, use_agent_cursor: false, first_page_uses_actual_sql: false, }); assert_eq!(plan.exact_query_row_bound, expected, "unexpected bound for {database_type:?}: {sql}"); let serialized = serde_json::to_value(&plan).unwrap(); assert_eq!( serialized.get("exactQueryRowBound").and_then(serde_json::Value::as_u64), expected.map(|value| value as u64) ); } } #[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 { sql: "SELECT * FROM (SELECT * FROM events) t ORDER BY created_at DESC".to_string(), query_base_sql: "SELECT * FROM events".to_string(), database_type: Some(DatabaseType::Oracle), pagination: QueryPagination { limit: 500, offset: 0, session_id: None }, use_agent_cursor: true, first_page_uses_actual_sql: true, }); assert_eq!(plan.sql_to_execute, "SELECT * FROM (SELECT * FROM events) t ORDER BY created_at DESC"); assert_eq!(plan.page_limit, Some(500)); assert_eq!(plan.page_offset, Some(0)); assert!(plan.page_sql.is_none()); assert!(plan.use_agent_result_session); } #[test] fn builds_sorted_query_sql() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT c.id, m.id FROM t_campaign c LEFT JOIN t_campaign_mdf m ON m.campaign_id = c.id" .to_string(), database_type: Some(DatabaseType::Postgres), result_columns: vec!["id".to_string(), "id".to_string()], column_index: 1, column: "id".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT c.id, m.id FROM t_campaign c LEFT JOIN t_campaign_mdf m ON m.campaign_id = c.id) t(\"id\", \"id_2\") ORDER BY \"id_2\" ASC;" ); } #[test] fn builds_sorted_query_sql_for_first_result_column() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT iso3, year, gdp_pc FROM country_gdp".to_string(), database_type: Some(DatabaseType::Postgres), result_columns: vec!["iso3".to_string(), "year".to_string(), "gdp_pc".to_string()], column_index: 0, column: "iso3".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT iso3, year, gdp_pc FROM country_gdp) t(\"iso3\", \"year\", \"gdp_pc\") ORDER BY \"iso3\" ASC;" ); } #[test] fn builds_mysql_sorted_query_without_alias_list() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT * FROM admin LIMIT 100;".to_string(), database_type: Some(DatabaseType::Mysql), result_columns: vec![ "id".to_string(), "guid".to_string(), "role_guid".to_string(), "login_name".to_string(), "password".to_string(), ], column_index: 3, column: "login_name".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!(result.sql.unwrap(), "SELECT * FROM (SELECT * FROM admin LIMIT 100) t ORDER BY `login_name` ASC;"); } #[test] fn hive_family_sort_uses_derived_table_alias_without_column_list() { for database_type in [DatabaseType::Argo, DatabaseType::Hive, DatabaseType::Impala] { for direction in [QuerySortDirection::Asc, QuerySortDirection::Desc] { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT id, name FROM users;".to_string(), database_type: Some(database_type), result_columns: vec!["id".to_string(), "name".to_string()], column_index: 1, column: "name".to_string(), direction, }); assert_eq!( result, ok(format!("SELECT * FROM (SELECT id, name FROM users) t ORDER BY `name` {};", direction.as_sql())), "{database_type:?} {direction:?}" ); } } } #[test] fn hive_family_sort_preserves_star_queries_and_existing_ordering() { for database_type in [DatabaseType::Argo, DatabaseType::Hive, DatabaseType::Impala] { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: " SELECT * FROM users WHERE active = 1 ORDER BY name DESC LIMIT 100; ".to_string(), database_type: Some(database_type), result_columns: vec!["id".to_string(), "name".to_string()], column_index: 0, column: "id".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!( result, ok("SELECT * FROM (SELECT * FROM users WHERE active = 1 ORDER BY name DESC LIMIT 100) t ORDER BY `id` ASC;".to_string()), "{database_type:?}" ); } } #[test] fn hive_family_sort_quotes_result_aliases_without_renaming_them() { for database_type in [DatabaseType::Argo, DatabaseType::Hive, DatabaseType::Impala] { for (column, quoted_column) in [ ("display_name", "`display_name`"), ("order", "`order`"), ("display name", "`display name`"), ("a.b", "`a.b`"), ("姓名", "`姓名`"), ("user`name", "`user``name`"), ("\"name\"", "`\"name\"`"), ] { let original_sql = format!("SELECT id, name AS {quoted_column} FROM users"); let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: original_sql.clone(), database_type: Some(database_type), result_columns: vec!["id".to_string(), column.to_string()], column_index: 1, column: "name".to_string(), direction: QuerySortDirection::Desc, }); assert_eq!( result, ok(format!("SELECT * FROM ({original_sql}) t ORDER BY {quoted_column} DESC;")), "{database_type:?} {column:?}" ); } } } #[test] fn hive_family_sort_falls_back_to_requested_column_for_out_of_range_index() { for database_type in [DatabaseType::Argo, DatabaseType::Hive, DatabaseType::Impala] { for result_columns in [vec!["id".to_string(), "display name".to_string()], Vec::new()] { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT id, name AS `display name` FROM users".to_string(), database_type: Some(database_type), result_columns, column_index: usize::MAX, column: "display name".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!( result, ok("SELECT * FROM (SELECT id, name AS `display name` FROM users) t ORDER BY `display name` ASC;" .to_string()), "{database_type:?}" ); } } } #[test] fn hive_family_sort_rejects_duplicate_derived_column_names_without_using_ordinals() { for database_type in [DatabaseType::Argo, DatabaseType::Hive, DatabaseType::Impala] { for result_columns in [["id", "id", "name"], ["ID", "id", "name"], ["姓名", "姓名", "name"]] { for column_index in [0, 1, 2, usize::MAX] { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT a.id, b.id, a.name FROM a JOIN b ON a.id = b.id".to_string(), database_type: Some(database_type), result_columns: result_columns.iter().map(|column| column.to_string()).collect(), column_index, column: "name".to_string(), direction: QuerySortDirection::Desc, }); assert_eq!(result, err("unsupported"), "{database_type:?} {result_columns:?} {column_index}"); } } } } #[test] fn hive_family_sort_preserves_statement_rejection_boundaries() { for database_type in [DatabaseType::Argo, DatabaseType::Hive, DatabaseType::Impala] { for (original_sql, reason) in [ ("", "empty"), (" \n\t", "empty"), (";", "empty"), ("SELECT 1; SELECT 2;", "multi"), ("WITH cte AS (SELECT 1 AS id) SELECT id FROM cte", "with"), ("SHOW TABLES", "not_select"), ("UPDATE users SET id = 2", "not_select"), ] { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: original_sql.to_string(), database_type: Some(database_type), result_columns: vec!["id".to_string()], column_index: 0, column: "id".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!(result, err(reason), "{database_type:?} {original_sql:?}"); } } } #[test] fn hive_family_sort_fix_preserves_other_dialects_column_alias_lists() { for (database_type, expected_suffix) in [ (None, "t(\"id\", \"id_2\") ORDER BY \"id_2\" DESC;"), (Some(DatabaseType::Jdbc), "t(id, id_2) ORDER BY id_2 DESC;"), (Some(DatabaseType::Postgres), "t(\"id\", \"id_2\") ORDER BY \"id_2\" DESC;"), (Some(DatabaseType::Spark), "t(`id`, `id_2`) ORDER BY `id_2` DESC;"), (Some(DatabaseType::Kyuubi), "t(`id`, `id_2`) ORDER BY `id_2` DESC;"), (Some(DatabaseType::Databricks), "t(`id`, `id_2`) ORDER BY `id_2` DESC;"), ] { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT a.id, b.id FROM a JOIN b ON a.id = b.id".to_string(), database_type, result_columns: vec!["id".to_string(), "id".to_string()], column_index: 1, column: "id".to_string(), direction: QuerySortDirection::Desc, }); assert_eq!( result, ok(format!("SELECT * FROM (SELECT a.id, b.id FROM a JOIN b ON a.id = b.id) {expected_suffix}")), "{database_type:?}" ); } } #[test] fn builds_saphana_sorted_query_without_derived_column_alias_list() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT ID, NAME, AMOUNT FROM DBX_ISSUE_7274_SORT".to_string(), database_type: Some(DatabaseType::SapHana), result_columns: vec!["ID".to_string(), "NAME".to_string(), "AMOUNT".to_string()], column_index: 1, column: "NAME".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT ID, NAME, AMOUNT FROM DBX_ISSUE_7274_SORT) t ORDER BY \"NAME\" ASC;" ); } #[test] fn builds_iris_sorted_query_without_derived_column_alias_list() { // Caché/IRIS rejects `t(col, col)` derived alias lists (SQLCODE -25) and // a quoted ORDER BY name becomes a string literal when delimited // identifiers are disabled, so the wrap must stay alias-free and // unquoted (#8340). let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT ID, Name FROM SQLUser.CT_Country".to_string(), database_type: Some(DatabaseType::Iris), result_columns: vec!["ID".to_string(), "Name".to_string()], column_index: 1, column: "Name".to_string(), direction: QuerySortDirection::Asc, }); let sql = result.sql.unwrap(); assert!(!sql.contains(") t("), "derived column alias list must not be emitted: {sql}"); assert_eq!(sql, "SELECT * FROM (SELECT ID, Name FROM SQLUser.CT_Country) t ORDER BY Name ASC;"); } #[test] fn builds_iris_sorted_query_by_ordinal_for_duplicate_columns() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id".to_string(), database_type: Some(DatabaseType::Iris), result_columns: vec!["ID".to_string(), "id".to_string()], column_index: 1, column: "id".to_string(), direction: QuerySortDirection::Desc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id) t ORDER BY 2 DESC;" ); } #[test] fn mysql_sort_preserves_directives_at_outermost_start() { for prefix in [ "/*sets:allsets*/", "/*master*/", "/*slave:set_1781591902_7*/", "/*future-route:anywhere*/", "/*proxy*/\n", "/*+ MAX_EXECUTION_TIME(1000) */\n", "/*@global:true*/\n", "/*& tenant:'test' */\n", ] { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: format!("{prefix}select @@server_id;"), database_type: Some(DatabaseType::Mysql), result_columns: vec!["@@server_id".to_string()], column_index: 0, column: "@@server_id".to_string(), direction: QuerySortDirection::Asc, }); assert!(result.ok, "{prefix}: {result:?}"); assert_eq!( result.sql.unwrap(), format!("{prefix}SELECT * FROM (select @@server_id) t ORDER BY `@@server_id` ASC;") ); } } #[test] fn builds_clickhouse_sorted_query_without_alias_list() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT id, part_day, hid FROM events LIMIT 100".to_string(), database_type: Some(DatabaseType::ClickHouse), result_columns: vec!["id".to_string(), "part_day".to_string(), "hid".to_string()], column_index: 0, column: "id".to_string(), direction: QuerySortDirection::Desc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT id, part_day, hid FROM events LIMIT 100) t ORDER BY `id` DESC;" ); } #[test] fn builds_doris_sorted_query_without_alias_list() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT id, name FROM users LIMIT 100".to_string(), database_type: Some(DatabaseType::Doris), result_columns: vec!["id".to_string(), "name".to_string()], column_index: 0, column: "id".to_string(), direction: QuerySortDirection::Desc, }); assert_eq!(result.sql.unwrap(), "SELECT * FROM (SELECT id, name FROM users LIMIT 100) t ORDER BY `id` DESC;"); } #[test] fn builds_starrocks_sorted_query_with_alias_list() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT id, name FROM users LIMIT 100".to_string(), database_type: Some(DatabaseType::StarRocks), result_columns: vec!["id".to_string(), "name".to_string()], column_index: 0, column: "id".to_string(), direction: QuerySortDirection::Desc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT id, name FROM users LIMIT 100) t(`id`, `name`) ORDER BY `id` DESC;" ); } #[test] fn strips_sqlserver_order_by_for_sorted_query() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT id, name FROM users ORDER BY id DESC".to_string(), database_type: Some(DatabaseType::SqlServer), result_columns: vec!["id".to_string(), "name".to_string()], column_index: 1, column: "name".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT id, name FROM users) t([id], [name]) ORDER BY [name] ASC;" ); } #[test] fn builds_dameng_sorted_query_without_alias_list() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT id, name FROM users".to_string(), database_type: Some(DatabaseType::Dameng), result_columns: vec!["id".to_string(), "name".to_string()], column_index: 0, column: "id".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!(result.sql.unwrap(), "SELECT * FROM (SELECT id, name FROM users) t ORDER BY \"id\" ASC;"); } #[test] fn builds_dameng_sorted_query_by_ordinal_for_duplicate_columns() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id".to_string(), database_type: Some(DatabaseType::Dameng), result_columns: vec!["ID".to_string(), "id".to_string()], column_index: 1, column: "id".to_string(), direction: QuerySortDirection::Desc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id) t ORDER BY 2 DESC;" ); } #[test] fn builds_oracle_sorted_query_by_ordinal_for_duplicate_columns() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id".to_string(), database_type: Some(DatabaseType::Oracle), result_columns: vec!["ID".to_string(), "ID".to_string()], column_index: 0, column: "ID".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id) t ORDER BY 1 ASC;" ); } #[test] fn builds_saphana_sorted_query_by_ordinal_for_duplicate_columns() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id".to_string(), database_type: Some(DatabaseType::SapHana), result_columns: vec!["ID".to_string(), "ID".to_string()], column_index: 1, column: "ID".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT a.id, b.id FROM a JOIN b ON b.a_id = a.id) t ORDER BY 2 ASC;" ); } #[test] fn preserves_derived_column_aliases_for_generic_jdbc() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "SELECT id, name FROM users".to_string(), database_type: Some(DatabaseType::Jdbc), result_columns: vec!["id".to_string(), "name".to_string()], column_index: 0, column: "id".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!(result.sql.unwrap(), "SELECT * FROM (SELECT id, name FROM users) t(id, name) ORDER BY id ASC;"); } #[test] fn rejects_with_query_sorting() { let result = build_sorted_query_sql(SortedQuerySqlOptions { original_sql: "WITH cte AS (SELECT 1) SELECT * FROM cte".to_string(), database_type: Some(DatabaseType::Postgres), result_columns: vec!["id".to_string()], column_index: 0, column: "id".to_string(), direction: QuerySortDirection::Asc, }); assert_eq!(result, err("with")); } // ----------------------------------------------------------------------- // Dedup query ORDER BY injection (DISTINCT / GROUP BY) // ----------------------------------------------------------------------- #[test] fn dedup_count_detects_distinct_query_without_order_by() { assert_eq!(dedup_projection_count_without_order_by("SELECT DISTINCT a, b, c FROM t"), Some(vec![1, 2, 3])); } #[test] fn dedup_count_detects_group_by_query() { assert_eq!( dedup_projection_count_without_order_by("SELECT city, COUNT(*) FROM users GROUP BY city"), Some(vec![1]) ); } #[test] fn dedup_count_returns_none_for_plain_select() { assert_eq!(dedup_projection_count_without_order_by("SELECT a, b FROM t"), None); } #[test] fn dedup_count_returns_none_when_order_by_exists() { assert_eq!(dedup_projection_count_without_order_by("SELECT DISTINCT a, b FROM t ORDER BY a"), None); } #[test] fn dedup_count_returns_none_for_wildcard() { assert_eq!(dedup_projection_count_without_order_by("SELECT DISTINCT * FROM t"), None); } #[test] fn doris_distinct_query_first_page_gets_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT city, name FROM users".to_string(), database_type: Some(DatabaseType::Doris), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT DISTINCT city, name FROM users ORDER BY 1, 2 LIMIT 100;"); } #[test] fn doris_distinct_query_second_page_wraps_with_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT city, name FROM users".to_string(), database_type: Some(DatabaseType::Doris), limit: 100, offset: 100, }); assert!(result.ok); // No top-level LIMIT in the original query, so ORDER BY + LIMIT + OFFSET // are appended directly to the statement (no wrapping needed). assert_eq!(result.sql.unwrap(), "SELECT DISTINCT city, name FROM users ORDER BY 1, 2 LIMIT 100 OFFSET 100;"); } #[test] fn doris_group_by_query_first_page_gets_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT dept, SUM(salary) FROM employees GROUP BY dept".to_string(), database_type: Some(DatabaseType::Doris), limit: 50, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT dept, SUM(salary) FROM employees GROUP BY dept ORDER BY 1 LIMIT 50;"); } #[test] fn doris_plain_query_without_distinct_is_unaffected() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT id, name FROM users".to_string(), database_type: Some(DatabaseType::Doris), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT id, name FROM users LIMIT 100;"); } #[test] fn doris_distinct_with_existing_limit_first_page_keeps_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT a, b, c FROM t LIMIT 500".to_string(), database_type: Some(DatabaseType::Doris), limit: 100, offset: 0, }); assert!(result.ok); // The dedup query has a LIMIT, so it must be wrapped. // `add_outer_standard_limit` always includes OFFSET clause. assert_eq!( result.sql.unwrap(), "SELECT * FROM (SELECT DISTINCT a, b, c FROM t LIMIT 500) `dbx_page` ORDER BY 1, 2, 3 LIMIT 100 OFFSET 0;" ); } #[test] fn mysql_distinct_query_first_page_gets_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT city FROM users".to_string(), database_type: Some(DatabaseType::Mysql), limit: 200, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT DISTINCT city FROM users ORDER BY 1 LIMIT 200;"); } #[test] fn postgres_distinct_query_first_page_gets_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT city, name FROM users".to_string(), database_type: Some(DatabaseType::Postgres), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!(result.sql.unwrap(), "SELECT DISTINCT city, name FROM users ORDER BY 1, 2 LIMIT 100;"); } #[test] fn distinct_query_with_existing_order_by_is_not_modified() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT a, b FROM t ORDER BY a".to_string(), database_type: Some(DatabaseType::Doris), limit: 100, offset: 0, }); assert!(result.ok); // Already has ORDER BY, so only LIMIT is appended. assert_eq!(result.sql.unwrap(), "SELECT DISTINCT a, b FROM t ORDER BY a LIMIT 100;"); } // ----------------------------------------------------------------------- // Complex queries with aliases, expressions, subqueries // ----------------------------------------------------------------------- #[test] fn dedup_count_handles_aliases() { assert_eq!( dedup_projection_count_without_order_by("SELECT DISTINCT a AS x, b AS y, c AS z FROM t"), Some(vec![1, 2, 3]) ); } #[test] fn dedup_count_handles_expressions() { assert_eq!( dedup_projection_count_without_order_by( "SELECT DISTINCT a + b AS sum_col, CASE WHEN c > 0 THEN 'Y' ELSE 'N' END AS flag FROM t" ), Some(vec![1, 2]) ); } #[test] fn dedup_count_handles_aggregate_with_alias() { assert_eq!( dedup_projection_count_without_order_by( "SELECT city, COUNT(*) AS cnt, AVG(salary) AS avg_sal FROM users GROUP BY city" ), Some(vec![1]) ); } #[test] fn doris_distinct_with_alias_and_expression_gets_positional_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT a AS col1, b + c AS col2, CASE WHEN d > 0 THEN 'Y' ELSE 'N' END AS col3 FROM t" .to_string(), database_type: Some(DatabaseType::Doris), limit: 100, offset: 0, }); assert!(result.ok); // Positional ORDER BY 1, 2, 3 works regardless of aliases or expressions. assert_eq!( result.sql.unwrap(), "SELECT DISTINCT a AS col1, b + c AS col2, CASE WHEN d > 0 THEN 'Y' ELSE 'N' END AS col3 FROM t ORDER BY 1, 2, 3 LIMIT 100;" ); } #[test] fn doris_group_by_with_aggregate_alias_gets_positional_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT dept, SUM(salary) AS total, COUNT(*) AS head_count FROM emp GROUP BY dept" .to_string(), database_type: Some(DatabaseType::Doris), limit: 50, offset: 100, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT dept, SUM(salary) AS total, COUNT(*) AS head_count FROM emp GROUP BY dept ORDER BY 1 LIMIT 50 OFFSET 100;" ); } #[test] fn doris_distinct_with_subquery_column_gets_positional_order_by() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT DISTINCT name, (SELECT MAX(score) FROM scores s WHERE s.uid = u.id) AS max_score FROM users u" .to_string(), database_type: Some(DatabaseType::Doris), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT DISTINCT name, (SELECT MAX(score) FROM scores s WHERE s.uid = u.id) AS max_score FROM users u ORDER BY 1, 2 LIMIT 100;" ); } #[test] fn group_by_order_by_skips_clickhouse_aggregate_states() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions { original_sql: "SELECT city, dept, argMaxState(name, hired_at) AS latest_hire, sumState(salary) AS payroll FROM employees GROUP BY city, dept" .to_string(), database_type: Some(DatabaseType::ClickHouse), limit: 100, offset: 0, }); assert!(result.ok); assert_eq!( result.sql.unwrap(), "SELECT city, dept, argMaxState(name, hired_at) AS latest_hire, sumState(salary) AS payroll FROM employees GROUP BY city, dept ORDER BY 1, 2 LIMIT 100;" ); } #[test] fn group_by_order_by_matches_key_by_output_alias() { assert_eq!( dedup_projection_count_without_order_by( "SELECT SUM(salary) AS total, dept AS team FROM employees GROUP BY team" ), Some(vec![2]) ); } #[test] fn group_by_order_by_accepts_positional_keys() { assert_eq!( dedup_projection_count_without_order_by("SELECT city, dept, COUNT(*) FROM employees GROUP BY 1, 2"), Some(vec![1, 2]) ); } #[test] fn group_by_order_by_falls_back_when_key_is_not_projected() { assert_eq!( dedup_projection_count_without_order_by( "SELECT COUNT(*) AS cnt, SUM(salary) AS total FROM employees GROUP BY city" ), Some(vec![1, 2]) ); } #[test] fn group_by_order_by_falls_back_for_wildcard_projection() { assert_eq!( dedup_projection_count_without_order_by("SELECT e.*, COUNT(*) FROM employees e GROUP BY city"), Some(vec![1, 2]) ); } #[test] fn union_query_is_not_treated_as_dedup() { assert_eq!(dedup_projection_count_without_order_by("SELECT a FROM t1 UNION SELECT b FROM t2"), None); } }