diff --git a/crates/dbx-sql-data/src/query_result_sql.rs b/crates/dbx-sql-data/src/query_result_sql.rs index c5f208c2a5..ae85cff63d 100644 --- a/crates/dbx-sql-data/src/query_result_sql.rs +++ b/crates/dbx-sql-data/src/query_result_sql.rs @@ -592,12 +592,24 @@ pub fn build_sorted_query_sql(options: SortedQuerySqlOptions) -> QuerySqlBuildRe let [Statement::Query(query)] = statements.as_slice() else { return err("unsupported"); }; - if query.limit_clause.is_some() + if options.column_index >= options.result_columns.len() { + return err("unsupported"); + } + // Appending ORDER BY after the statement cannot follow the user's + // row-limit clause (`... LIMIT 10 ORDER BY` is a syntax error), so a + // statement that entered this block on a row limit must not take the + // append path. A CTE on an engine that accepts WITH inside a derived + // table falls through to the wrapper sort instead; the Oracle family + // and SQL Server keep the historical rejection (#13639). Entry reasons + // other than the CTE itself (duplicate labels, OceanBase projection + // safety) cannot use the wrapper, so they stay rejected. + let has_row_limit = query.limit_clause.is_some() || query.fetch.is_some() || !query.locks.is_empty() || has_top_level_rownum(statement) - || has_top_level_top(statement) - || options.column_index >= options.result_columns.len() + || has_top_level_top(statement); + if has_row_limit + && (!is_cte || !derived_table_allows_cte(options.database_type) || derived_table_rejects_duplicate_labels) { return err("unsupported"); } @@ -608,18 +620,22 @@ pub fn build_sorted_query_sql(options: SortedQuerySqlOptions) -> QuerySqlBuildRe if order_by_index.is_none() && query.order_by.is_some() { return err("unsupported"); } - return ok(format!( - "{execution_hint}{}", - append_sql_suffix( - statement[..order_by_index.unwrap_or(statement.len())].trim_end(), - &format!( - "ORDER BY {} {}{};", - options.column_index + 1, - options.direction.as_sql(), - tie_breaker_indexes.iter().map(|index| format!(", {} ASC", index + 1)).collect::() - ), - ) - )); + if !has_row_limit { + return ok(format!( + "{execution_hint}{}", + append_sql_suffix( + statement[..order_by_index.unwrap_or(statement.len())].trim_end(), + &format!( + "ORDER BY {} {}{};", + options.column_index + 1, + options.direction.as_sql(), + tie_breaker_indexes.iter().map(|index| format!(", {} ASC", index + 1)).collect::() + ), + ) + )); + } + // A row-limited CTE on a wrapper-capable engine: fall through to the + // derived-table path below, which sorts the limited result set. } let uses_hive_subquery_syntax = @@ -759,6 +775,42 @@ fn unsupported_pagination_type(database_type: Option) -> bool { ) } +/// Whether the engine accepts a `WITH ... SELECT` statement inside a derived +/// table (`SELECT * FROM (WITH ... SELECT ...) t`), so a CTE can be sorted by +/// wrapping instead of appending ORDER BY to its main query. SQL Server, Db2 +/// and the Oracle family reject WITH in a derived table; PostgreSQL, MySQL, +/// SQLite and ClickHouse accept it. The Hive family also rejects it — "The +/// WITH clause is not supported within SubQuery Blocks" — which is why the +/// count path hoists the CTE out instead (`#8535`); engines not yet verified +/// against a live instance stay excluded so sorting degrades to the +/// "not sortable" message instead of a server error. +fn derived_table_allows_cte(database_type: Option) -> bool { + !matches!( + database_type, + Some( + DatabaseType::Oracle + | DatabaseType::OceanbaseOracle + | DatabaseType::Dameng + | DatabaseType::Yashandb + | DatabaseType::Gaussdb + | DatabaseType::Kingbase + | DatabaseType::Oscar + | DatabaseType::Xugu + | DatabaseType::SqlServer + | DatabaseType::Iris + | DatabaseType::SapHana + | DatabaseType::Hive + | DatabaseType::Argo + | DatabaseType::Impala + | DatabaseType::Db2 + | DatabaseType::Teradata + | DatabaseType::Firebird + | DatabaseType::Informix + | DatabaseType::Gbase8s + ) + ) +} + fn find_query_result_statement_at_cursor(sql: &str, cursor_pos: usize, database_type: Option) -> String { if matches!(database_type, Some(DatabaseType::Mysql | DatabaseType::Goldendb)) { find_statement_at_cursor_for_database(sql, cursor_pos, database_type.unwrap()) @@ -7975,17 +8027,127 @@ WHERE u.id = picked.id; ); } + #[test] + fn sorts_with_query_with_limit_by_wrapping_the_whole_statement() { + // A CTE whose main query ends in a row limit cannot take an appended + // ORDER BY (`... LIMIT 10 ORDER BY` is a syntax error), so engines that + // accept WITH inside a derived table wrap the whole statement and sort + // the limited result set (#13639). + let result = sorted_cte_sql( + "with tmp_fr as (select 1 as a) select a from tmp_fr limit 10", + Some(DatabaseType::Postgres), + &["a"], + 0, + QuerySortDirection::Asc, + Vec::new(), + ); + + assert_eq!( + result, + ok("SELECT * FROM (with tmp_fr as (select 1 as a) select a from tmp_fr limit 10) t(\"a\") ORDER BY \"a\" ASC;" + .to_string()) + ); + } + + #[test] + fn sorts_with_query_with_fetch_or_lock_by_wrapping_on_capable_engines() { + let result = sorted_cte_sql( + "WITH cte AS (SELECT 1 AS id) SELECT id FROM cte FETCH FIRST 10 ROWS ONLY", + Some(DatabaseType::Postgres), + &["id"], + 0, + QuerySortDirection::Desc, + Vec::new(), + ); + assert_eq!( + result, + ok("SELECT * FROM (WITH cte AS (SELECT 1 AS id) SELECT id FROM cte FETCH FIRST 10 ROWS ONLY) t(\"id\") ORDER BY \"id\" DESC;" + .to_string()) + ); + + // A locking clause rides along into the wrapper: the lock applies to + // the inner query and the outer sort keeps the limited rows (#13639). + let result = sorted_cte_sql( + "WITH cte AS (SELECT 1 AS id) SELECT id FROM cte LIMIT 5 FOR UPDATE", + Some(DatabaseType::Postgres), + &["id"], + 0, + QuerySortDirection::Asc, + Vec::new(), + ); + assert_eq!( + result, + ok("SELECT * FROM (WITH cte AS (SELECT 1 AS id) SELECT id FROM cte LIMIT 5 FOR UPDATE) t(\"id\") ORDER BY \"id\" ASC;" + .to_string()) + ); + } + + #[test] + fn derived_table_cte_exclusions_track_the_count_path_special_cases() { + // The count path already hoists the CTE out for Hive (#8535) and wraps + // Db2 specially because both engines reject `WITH` inside a derived + // table; the sort wrapper must not feed them that shape. + for database_type in [DatabaseType::Hive, DatabaseType::Argo, DatabaseType::Impala, DatabaseType::Db2] { + assert!( + !derived_table_allows_cte(Some(database_type)), + "{database_type:?} must stay excluded from CTE-in-derived-table sorting" + ); + } + for database_type in [DatabaseType::Postgres, DatabaseType::Mysql, DatabaseType::Sqlite] { + assert!(derived_table_allows_cte(Some(database_type))); + } + } + + #[test] + fn rejects_with_query_sorting_for_cte_rejecting_engines() { + // Hive rejects WITH inside subquery blocks outright and Db2's + // fullselect grammar has no nested WITH (SQL0199N), so a row-limited + // CTE stays unsupported there instead of producing a server error. + for (original_sql, database_type) in [ + ("WITH cte AS (SELECT 1 AS id) SELECT id FROM cte LIMIT 10", DatabaseType::Hive), + ("WITH cte AS (SELECT 1 AS id) SELECT id FROM cte FETCH FIRST 10 ROWS ONLY", DatabaseType::Db2), + ] { + let result = + sorted_cte_sql(original_sql, Some(database_type), &["id"], 0, QuerySortDirection::Asc, Vec::new()); + assert_eq!(result, err("unsupported"), "unexpected sort result for {database_type:?}"); + } + } + + #[test] + fn sorts_with_query_with_limit_by_wrapping_for_mysql() { + let result = sorted_cte_sql( + "with tmp as (select 1 as a) select a from tmp limit 10", + Some(DatabaseType::Mysql), + &["a"], + 0, + QuerySortDirection::Asc, + Vec::new(), + ); + + assert_eq!( + result, + ok("SELECT * FROM (with tmp as (select 1 as a) select a from tmp limit 10) t ORDER BY `a` ASC;" + .to_string()) + ); + } + #[test] fn rejects_with_query_sorting_when_order_by_cannot_be_appended() { for (original_sql, database_type, reason) in [ - ("WITH cte AS (SELECT 1 AS id) SELECT id FROM cte LIMIT 10", Some(DatabaseType::Postgres), "unsupported"), + // Engines that reject WITH inside a derived table have no wrapper + // fallback, so a row-limited CTE main query stays unsupported. + ("WITH cte AS (SELECT 1 AS id) SELECT id FROM cte LIMIT 10", Some(DatabaseType::Oracle), "unsupported"), ( "WITH cte AS (SELECT 1 AS id) SELECT id FROM cte FETCH FIRST 10 ROWS ONLY", Some(DatabaseType::Oracle), "unsupported", ), ("WITH cte AS (SELECT 1 AS id) SELECT TOP 10 id FROM cte", Some(DatabaseType::SqlServer), "unsupported"), - ("WITH cte AS (SELECT 1 AS id) SELECT id FROM cte FOR UPDATE", Some(DatabaseType::Mysql), "unsupported"), + ( + "WITH cte AS (SELECT 1 AS id) SELECT id FROM cte FOR UPDATE", + Some(DatabaseType::SqlServer), + "unsupported", + ), ("WITH cte AS (SELECT 1 AS id) SELECT id INTO #ids FROM cte", Some(DatabaseType::SqlServer), "with"), ( "WITH cte AS (SELECT 1 AS id) DELETE FROM users WHERE id IN (SELECT id FROM cte)",