fix(query): sort row-limited CTE results by wrapping the whole statement

This commit is contained in:
caiyq925-cell authored and GitHub committed 2026-10-28 00:52:07 +08:00
1 parent 5380331724
commit 7d957d824b
1 file changed
+179 -17
+179 -17
View File
@@ -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::<String>()
),
)
));
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::<String>()
),
)
));
}
// 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<DatabaseType>) -> 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<DatabaseType>) -> 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<DatabaseType>) -> 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)",