fix: rewrite limit / offset

Author: sweb

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 <expr> OFFSET <expr>` 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(())
+    }
+}

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");

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" }
+        ),
     ]
 );