diff --git a/integration/rust/tests/integration/limit.rs b/integration/rust/tests/integration/limit.rs index 3bbcba93a..35fc04ca1 100644 --- a/integration/rust/tests/integration/limit.rs +++ b/integration/rust/tests/integration/limit.rs @@ -2,6 +2,41 @@ use crate::setup::admin_sqlx; use crate::setup::connections_sqlx; use sqlx::{Executor, Row, postgres::PgPool}; +#[tokio::test] +async fn limit_without_order_by_across_shards() -> Result<(), Box> { + let pools = connections_sqlx().await; + let mut transaction = pools[1].begin().await?; + + transaction.execute("TRUNCATE sharded").await?; + // Guarantee five rows on each shard so a per-shard LIMIT exceeds the global limit. + for shard in [0, 1] { + transaction + .execute( + format!( + "/* pgdog_shard: {shard} */ INSERT INTO sharded (id) VALUES (1), (2), (3), (4), (5)" + ) + .as_str(), + ) + .await?; + } + + for (limit, expected) in [(5_i64, 5), (0, 0), (100, 10)] { + let rows = sqlx::raw_sql(&format!("SELECT id FROM sharded LIMIT {limit}")) + .fetch_all(&mut *transaction) + .await?; + assert_eq!(rows.len(), expected, "simple protocol LIMIT {limit}"); + + let rows = sqlx::query("SELECT id FROM sharded LIMIT $1") + .bind(limit) + .fetch_all(&mut *transaction) + .await?; + assert_eq!(rows.len(), expected, "extended protocol LIMIT {limit}"); + } + + transaction.rollback().await?; + Ok(()) +} + #[tokio::test] async fn limit_across_shards() -> Result<(), Box> { let sharded = connections_sqlx().await.get(1).cloned().unwrap(); diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index eac6d65cd..a0ae3728e 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -333,6 +333,7 @@ impl Route { !self.order_by().is_empty() || !self.aggregate().is_empty() || self.distinct().is_some() + || self.limit().limit.is_some() || self.limit().offset.is_some() } @@ -741,7 +742,7 @@ mod test { }, None, ); - assert!(!route.requires_post_processing()); + assert!(route.requires_post_processing()); } #[test]