diff --git a/src/rewrite/limit_offset.rs b/src/rewrite/limit_offset.rs new file mode 100644 index 0000000..2a5c15e --- /dev/null +++ b/src/rewrite/limit_offset.rs @@ -0,0 +1,59 @@ +// SPDX-FileCopyrightText: 2026 Stackable GmbH +// SPDX-License-Identifier: OSL-3.0 +use sqlparser::ast::{Fetch, LimitClause, Query, VisitorMut}; +use std::ops::ControlFlow; + +/// Reorders `LIMIT n OFFSET m` into a Trino-compatible form. +/// +/// PostgreSQL accepts `LIMIT n OFFSET m`, but Trino's grammar requires the +/// offset to come *before* the row-limiting clause (`OFFSET m LIMIT n` or +/// `OFFSET m FETCH FIRST n ROWS ONLY`). sqlparser's `Display` for +/// [`LimitClause::LimitOffset`] always writes `LIMIT` before `OFFSET` +/// regardless of input order, so a plain round-trip keeps the order Trino +/// rejects — a `VisitorMut` on expressions cannot fix it. +/// +/// Instead we exploit `Query`'s field render order: `limit_clause` is emitted +/// before `fetch`. So we leave the `OFFSET` in the limit clause and move the +/// `LIMIT` value into a `FETCH FIRST n ROWS ONLY` clause. The result renders as +/// `... OFFSET m FETCH FIRST n ROWS ONLY`, which is valid Trino and +/// semantically identical to `LIMIT n OFFSET m`. Everything is built from AST +/// nodes — no raw-string manipulation (see the "AST, never raw strings" rule in +/// `AGENTS.md`). +/// +/// Using [`VisitorMut::post_visit_query`] means every `Query` node is handled, +/// including subqueries and CTEs, not just the top level. +pub struct LimitOffsetRewriter; + +impl VisitorMut for LimitOffsetRewriter { + type Break = (); + + fn post_visit_query(&mut self, query: &mut Query) -> ControlFlow<()> { + // Don't clobber a pre-existing FETCH (would be a malformed query anyway). + if query.fetch.is_some() { + return ControlFlow::Continue(()); + } + + // Only the plain `LIMIT OFFSET ` case: both present, no + // ClickHouse `LIMIT BY`. `LIMIT ALL OFFSET m` parses to `limit: None` + // (sqlparser drops `ALL`), so `.take()` yields `None` and we leave the + // bare `OFFSET m` untouched — Trino accepts that as-is. + let limit = match &mut query.limit_clause { + Some(LimitClause::LimitOffset { + limit, + offset: Some(_), + limit_by, + }) if limit_by.is_empty() => limit.take(), + _ => None, + }; + + if let Some(limit) = limit { + query.fetch = Some(Fetch { + with_ties: false, + percent: false, + quantity: Some(limit), + }); + } + + ControlFlow::Continue(()) + } +} diff --git a/src/rewrite/mod.rs b/src/rewrite/mod.rs index 2b266ab..ddc98c9 100644 --- a/src/rewrite/mod.rs +++ b/src/rewrite/mod.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: OSL-3.0 mod casts; mod functions; +mod limit_offset; mod predicates; use sqlparser::ast::VisitMut; @@ -22,6 +23,8 @@ use sqlparser::parser::Parser; /// - PostgreSQL type names are normalized to Trino equivalents /// - `ILIKE` becomes `lower(x) LIKE lower(pattern)` /// - PostgreSQL function names are mapped to Trino equivalents +/// - `LIMIT n OFFSET m` is reordered into Trino order +/// (`OFFSET m FETCH FIRST n ROWS ONLY`) /// /// If parsing fails (e.g. for `SET`, `SHOW`, `DISCARD` commands), the original /// SQL is returned unchanged. @@ -57,9 +60,11 @@ pub fn rewrite_sql(sql: &str) -> String { let mut cast_rewriter = casts::CastRewriter; let mut ilike_rewriter = predicates::ILikeRewriter; let mut fn_renamer = functions::FunctionRenamer; + let mut limit_offset_rewriter = limit_offset::LimitOffsetRewriter; let _ = stmt.visit(&mut cast_rewriter); let _ = stmt.visit(&mut ilike_rewriter); let _ = stmt.visit(&mut fn_renamer); + let _ = stmt.visit(&mut limit_offset_rewriter); stmt.to_string() } @@ -133,6 +138,36 @@ mod tests { must_contain: &["SELECT", "FROM"], must_not_contain: &[], }, + Case { + name: "LIMIT n OFFSET m → OFFSET m FETCH FIRST n (no bare LIMIT)", + input: "SELECT name FROM t ORDER BY name LIMIT 2 OFFSET 1", + must_contain: &["OFFSET 1", "FETCH FIRST 2"], + must_not_contain: &["LIMIT"], + }, + Case { + name: "LIMIT only is left unchanged", + input: "SELECT name FROM t LIMIT 5", + must_contain: &["LIMIT 5"], + must_not_contain: &["FETCH", "OFFSET"], + }, + Case { + name: "OFFSET only is left unchanged", + input: "SELECT name FROM t OFFSET 3", + must_contain: &["OFFSET 3"], + must_not_contain: &["FETCH", "LIMIT"], + }, + Case { + name: "LIMIT ALL OFFSET m → bare OFFSET (ALL dropped, no FETCH)", + input: "SELECT name FROM t LIMIT ALL OFFSET 4", + must_contain: &["OFFSET 4"], + must_not_contain: &["FETCH", "LIMIT", "ALL"], + }, + Case { + name: "subquery LIMIT+OFFSET is reordered too", + input: "SELECT * FROM (SELECT name FROM t ORDER BY name LIMIT 2 OFFSET 1) x", + must_contain: &["OFFSET 1", "FETCH FIRST 2"], + must_not_contain: &["LIMIT"], + }, ]; #[test] @@ -162,6 +197,20 @@ mod tests { assert_eq!(rewrite_sql(input), input); } + /// The reordered clause must place `OFFSET` before the row-limiting + /// `FETCH` — the whole point of the rewrite, which the substring-based + /// `Case` table cannot assert on its own. + #[test] + fn limit_offset_emits_offset_before_fetch() { + let result = rewrite_sql("SELECT name FROM t ORDER BY name LIMIT 2 OFFSET 1"); + let offset_at = result.find("OFFSET").expect("OFFSET present"); + let fetch_at = result.find("FETCH").expect("FETCH present"); + assert!( + offset_at < fetch_at, + "expected OFFSET before FETCH in: {result}" + ); + } + #[test] fn show_passes_through_non_empty() { let result = rewrite_sql("SHOW server_version"); diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 7bb8004..8f4cdf3 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -437,6 +437,22 @@ trino_tests!( "SELECT name FROM nation ORDER BY nationkey OFFSET 5 ROWS FETCH FIRST 3 ROWS ONLY", Check::Rows { min_rows: 3 } ), + // PostgreSQL-order `LIMIT n OFFSET m` — Trino rejects it verbatim; the + // rewriter reorders it into `OFFSET m FETCH FIRST n ROWS ONLY`. nation + // ordered by nationkey is ALGERIA(0), ARGENTINA(1), ...; offset 1 + // limit 1 must yield ARGENTINA. + ( + "pg-order limit offset", + "SELECT name FROM nation ORDER BY nationkey LIMIT 1 OFFSET 1", + Check::Value { value: "ARGENTINA" } + ), + // Same rewrite must apply inside a subquery (inner → ARGENTINA, BRAZIL; + // outer takes the first alphabetically). + ( + "pg-order limit offset in subquery", + "SELECT name FROM (SELECT name FROM nation ORDER BY nationkey LIMIT 2 OFFSET 1) t ORDER BY name LIMIT 1", + Check::Value { value: "ARGENTINA" } + ), ] );