diff --git a/AGENTS.md b/AGENTS.md index e62fbbe..048cea0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -10,7 +10,8 @@ forwarding queries to Trino's REST API on the backend. Lets PG clients - The intercept layer answers `SET`, `SHOW`, `BEGIN`/`COMMIT`, and `pg_catalog` queries locally; everything else is forwarded. - The SQL rewriter transforms PG-dialect SQL to Trino-compatible SQL - (`::cast`, `ILIKE`, function-name remaps, type-name normalisation). + (`::cast`, `ILIKE`, function-name remaps, type-name normalisation, + `LIMIT`/`OFFSET` reordering). - The Trino backend forwards rewritten queries via the REST API and streams result pages back as PG wire-protocol DataRows. - Catalog emulation fakes `pg_type`, `pg_class`, `pg_attribute`, and a few @@ -33,7 +34,7 @@ forwarding queries to Trino's REST API on the backend. Lets PG clients - `cancel.rs` — PG `CancelRequest` to Trino `DELETE /v1/query/{id}` - `session.rs` — per-connection state, cancel registry, portal cache - `catalog/` — `pg_catalog` emulation (`pg_type`, `pg_class`, `pg_attribute`, stubs) - - `rewrite/` — SQL rewriting (casts, predicates, functions) + - `rewrite/` — SQL rewriting (casts, predicates, functions, limit/offset) - `types.rs` — Trino-to-PG type mapping and value encoding - `trino_stream.rs` — streaming bridge (poll Trino, yield PG DataRow) - `error_mapping.rs` — Trino errors to PG SQLSTATE codes diff --git a/README.md b/README.md index b5e1b36..7cbe034 100644 --- a/README.md +++ b/README.md @@ -155,6 +155,10 @@ break Trino's strict-PostgreSQL parsing: `log` → `log10`, `trunc` → `truncate`. - PostgreSQL type names map to Trino: `text` → `VARCHAR`, `int4` → `INTEGER`, and so on. +- `LIMIT n OFFSET m` is reordered into the order Trino's grammar + requires: `OFFSET m FETCH FIRST n ROWS ONLY`. `LIMIT 0 OFFSET m` + becomes a bare `LIMIT 0` — Trino rejects `FETCH FIRST 0 ROWS ONLY`, + and the offset makes no difference to an empty result. A handful of queries are intercepted and answered locally rather than forwarded: diff --git a/src/rewrite/limit_offset.rs b/src/rewrite/limit_offset.rs new file mode 100644 index 0000000..991b687 --- /dev/null +++ b/src/rewrite/limit_offset.rs @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: 2026 Stackable GmbH +// SPDX-License-Identifier: OSL-3.0 +use sqlparser::ast::{Expr, Fetch, LimitClause, Query, Value, 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`). +/// +/// `LIMIT 0` is the exception: Trino accepts `LIMIT 0` but rejects +/// `FETCH FIRST 0 ROWS ONLY`, so that case is rewritten to a bare `LIMIT 0` +/// instead (see [`is_zero_literal`]). +/// +/// 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, + }; + + let Some(limit) = limit else { + return ControlFlow::Continue(()); + }; + + // Trino rejects `FETCH FIRST 0 ROWS ONLY` ("FETCH FIRST row count must + // be positive"), while `LIMIT 0` is accepted and returns no rows — + // PostgreSQL accepts both. `LIMIT 0 OFFSET m` is empty for every `m`, + // so we drop the `OFFSET` and keep a bare `LIMIT 0`, which needs no + // reordering. Power BI issues `LIMIT 0` to probe result schemas, so + // this path is hit in practice. + if is_zero_literal(&limit) { + query.limit_clause = Some(LimitClause::LimitOffset { + limit: Some(limit), + offset: None, + limit_by: Vec::new(), + }); + return ControlFlow::Continue(()); + } + + query.fetch = Some(Fetch { + with_ties: false, + percent: false, + quantity: Some(limit), + }); + + ControlFlow::Continue(()) + } +} + +/// Whether `expr` is a numeric literal equal to zero. +/// +/// Only literals are recognised — a placeholder or expression that happens to +/// evaluate to zero still becomes a `FETCH` clause, which Trino rejects at +/// analysis time. Nothing we can do about that without evaluating the +/// expression ourselves. +fn is_zero_literal(expr: &Expr) -> bool { + let Expr::Value(value) = expr else { + return false; + }; + match &value.value { + Value::Number(n, _) => n.parse::().is_ok_and(|n| n == 0.0), + _ => false, + } +} diff --git a/src/rewrite/mod.rs b/src/rewrite/mod.rs index 2b266ab..f139787 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,50 @@ 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 { + // Trino rejects `FETCH FIRST 0 ROWS ONLY`; `LIMIT 0 OFFSET m` is + // empty for every `m`, so the OFFSET is dropped as well. + name: "LIMIT 0 OFFSET m → bare LIMIT 0 (no FETCH)", + input: "SELECT name FROM t ORDER BY name LIMIT 0 OFFSET 1", + must_contain: &["LIMIT 0"], + must_not_contain: &["FETCH", "OFFSET"], + }, + Case { + name: "LIMIT 0 without OFFSET is left unchanged", + input: "SELECT name FROM t LIMIT 0", + must_contain: &["LIMIT 0"], + must_not_contain: &["FETCH", "OFFSET"], + }, + 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 +211,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..1e60825 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -437,6 +437,35 @@ 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" } + ), + // `LIMIT 0` must not become `FETCH FIRST 0 ROWS ONLY` — Trino rejects + // that with "FETCH FIRST row count must be positive". Both cases must + // execute and return no rows (Power BI probes schemas this way). + ( + "pg-order limit zero offset", + "SELECT name FROM nation ORDER BY nationkey LIMIT 0 OFFSET 1", + Check::Rows { min_rows: 0 } + ), + ( + "limit zero", + "SELECT name FROM nation LIMIT 0", + Check::Rows { min_rows: 0 } + ), ] );