feat: multi-shard LIMIT+OFFSET in-flight re-write in PREPARE+EXECUTE

integration/rust/tests/integration/simple_prepared.rs

@@ -32,3 +32,159 @@ async fn test_simple_prepared_ttl() {
 
     assert_eq!(test_return.try_get::<i32, &str>("?column?").unwrap(), 1);
 }
+
+/// <https://github.com/pgdogdev/pgdog/issues/1383>
+/// TODO: will need to support extended-protocol `Bind`s later for the re-write
+#[tokio::test]
+async fn test_simple_prepared_limit() {
+    let mut conn =
+        sqlx::PgConnection::connect("postgres://pgdog:pgdog@127.0.0.1:6432/pgdog_sharded")
+            .await
+            .unwrap();
+
+    // Clear out the table first.
+    sqlx::raw_sql("TRUNCATE sharded")
+        .execute(&mut conn)
+        .await
+        .unwrap();
+
+    // Insert 55 rows.
+    for i in 1..=55 {
+        sqlx::query("INSERT INTO sharded(id) VALUES ($1)")
+            .bind(i)
+            .execute(&mut conn)
+            .await
+            .unwrap();
+    }
+
+    // Test limit 5 offset 10 (one ParamRef)
+    {
+        // Write the PREPARE / EXECUTE
+        sqlx::raw_sql("PREPARE stmt AS SELECT * FROM sharded ORDER BY id DESC LIMIT 5 OFFSET $1")
+            .execute(&mut conn)
+            .await
+            .unwrap();
+
+        let rows = sqlx::raw_sql("EXECUTE stmt(10)")
+            .fetch_all(&mut conn)
+            .await
+            .unwrap();
+
+        // There's 55 rows. We're offsetting by 10 with a limit of 5 in reverse (DESC).
+        // Therefore, we start at 45, and go down from there.
+        assert_eq!(
+            rows.iter()
+                .map(|row| row.get::<i64, &str>("id"))
+                .collect::<Vec<_>>(),
+            vec![45, 44, 43, 42, 41]
+        );
+    }
+
+    // Test limit 10 offset 5 (two ParamRefs);
+    // order Limit ref before Offset (out-of-order left-to-right refs)
+    {
+        // This also resolves to the same cached entry $2, $1 as the one ABOVE.
+        //
+        // However, since PgDog uses the pre-re-written `Query` (the one we send here) as the `CacheKey`,
+        // the two statements will resolve differently.
+        //
+        // If we didn't use the pre-re-written `Query`, it would re-use the LIMIT 5 from last time,
+        // and this would return an incorrect response.
+        sqlx::raw_sql("PREPARE stmt2 AS SELECT * FROM sharded ORDER BY id DESC LIMIT $2 OFFSET $1")
+            .execute(&mut conn)
+            .await
+            .unwrap();
+
+        // LIMIT 10 OFFSET 5
+        let rows = sqlx::raw_sql("EXECUTE stmt2(5, 10)")
+            .fetch_all(&mut conn)
+            .await
+            .unwrap();
+
+        // There's 55 rows. We're offsetting by 5 with a limit of 10 in reverse (DESC).
+        // Therefore, we start at 50, and go down from there.
+        assert_eq!(
+            rows.iter()
+                .map(|row| row.get::<i64, &str>("id"))
+                .collect::<Vec<_>>(),
+            vec![50, 49, 48, 47, 46, 45, 44, 43, 42, 41]
+        );
+    }
+
+    // Try normal order $1, $2
+    {
+        // Write the PREPARE / EXECUTE
+        sqlx::raw_sql("PREPARE stmt2 AS SELECT * FROM sharded ORDER BY id DESC LIMIT $1 OFFSET $2")
+            .execute(&mut conn)
+            .await
+            .unwrap();
+
+        let rows = sqlx::raw_sql("EXECUTE stmt2(5, 10)")
+            .fetch_all(&mut conn)
+            .await
+            .unwrap();
+
+        // There's 55 rows. We're offsetting by 10 with a limit of 5 in reverse (DESC).
+        // Therefore, we start at 45, and go down from there.
+        assert_eq!(
+            rows.iter()
+                .map(|row| row.get::<i64, &str>("id"))
+                .collect::<Vec<_>>(),
+            vec![45, 44, 43, 42, 41]
+        );
+    }
+
+    // Test limit 10 offset 5 (no ParamRefs; all A_Const nodes)
+    {
+        // Write the PREPARE / EXECUTE
+        sqlx::raw_sql("PREPARE stmt3 AS SELECT * FROM sharded ORDER BY id DESC LIMIT 10 OFFSET 5")
+            .execute(&mut conn)
+            .await
+            .unwrap();
+
+        let rows = sqlx::raw_sql("EXECUTE stmt3")
+            .fetch_all(&mut conn)
+            .await
+            .unwrap();
+
+        // There's 55 rows. We're offsetting by 5 with a limit of 10 in reverse (DESC).
+        // Therefore, we start at 50, and go down from there.
+        assert_eq!(
+            rows.iter()
+                .map(|row| row.get::<i64, &str>("id"))
+                .collect::<Vec<_>>(),
+            vec![50, 49, 48, 47, 46, 45, 44, 43, 42, 41]
+        );
+    }
+
+    // Lets also test with an un-related param (WHERE id < $2)
+    {
+        // Write the PREPARE / EXECUTE
+        sqlx::raw_sql("PREPARE stmt4 AS SELECT * FROM sharded WHERE id < $2 ORDER BY id DESC LIMIT $3 OFFSET $1")
+            .execute(&mut conn)
+            .await
+            .unwrap();
+
+        // [offset, WHERE id <, limit]
+        let rows = sqlx::raw_sql("EXECUTE stmt4(5, 25, 10)")
+            .fetch_all(&mut conn)
+            .await
+            .unwrap();
+
+        // There's 55 rows. We're offsetting by 5 with a limit of 10 in reverse (DESC).
+        // We also filter out any >= $2 (25)
+        // Therefore, we start at 19 (24 - 5), and go down from there.
+        assert_eq!(
+            rows.iter()
+                .map(|row| row.get::<i64, &str>("id"))
+                .collect::<Vec<_>>(),
+            vec![19, 18, 17, 16, 15, 14, 13, 12, 11, 10]
+        );
+    }
+
+    // Clean-up (clear again)
+    sqlx::raw_sql("TRUNCATE sharded")
+        .execute(&mut conn)
+        .await
+        .unwrap();
+}

pgdog/src/backend/server.rs

@@ -2260,7 +2260,7 @@ pub(crate) mod test {
         let mut prep = PreparedStatements::new();
         let name = "test";
         let query = Bytes::from("SELECT 1::bigint".to_owned());
-        let prepare = prep.insert_prepare(name, query.clone(), &RewritePlan::default());
+        let prepare = prep.insert_prepare(name, query.clone(), None, &RewritePlan::default(), None);
         assert_eq!(prepare.name(), "__pgdog_1");
 
         server

pgdog/src/frontend/prepared_statements/cache_key.rs

@@ -12,7 +12,6 @@ use super::prelude::*;
 /// A `Simple` key comes from SQL `PREPARE` and matches nothing but itself.
 /// Its declared argument types are not captured, so two of those
 /// statements are never known to be the same.
-///
 #[derive(Debug, Clone, PartialEq, Hash, Eq)]
 pub(crate) enum CacheKey {
     Extended { query: Bytes, data_types: Bytes },
@@ -33,7 +32,7 @@ impl CacheKey {
     pub(crate) fn query(&self) -> Result<&str, crate::net::Error> {
         match self {
             Self::Extended { query, .. } => Ok(from_utf8(&query[0..query.len() - 1])?),
-            Self::Simple { query } => Ok(from_utf8(query)?), // Simple queries are regular Rust strings.
+            Self::Simple { query, .. } => Ok(from_utf8(query)?), // Simple queries are regular Rust strings.
         }
     }
 }

pgdog/src/frontend/prepared_statements/global_cache.rs

@@ -81,13 +81,22 @@ impl GlobalCache {
     }
 
     /// Insert a statement prepared using the simple protocol into the global cache.
+    /// `original_query` is used as the `CacheKey`
+    /// If `rewritten_query` is...
+    ///     - Some(..): `rewritten_query` is sent to Postgres as the PREPARE inner-query.
+    ///     - None: `original_query` is sent to Postgres as the PREPARE inner-query.
     pub(super) fn insert_prepare(
         &mut self,
-        query: Bytes,
+        original_query: Bytes,
+        rewritten_query: Option<Bytes>,
+        // TODO: I think we should just pass `unique_ids` in here by itself.
+        //       Otherwise, it could be easily confused to want
+        //       to use `RewritePlan` for `offset_plan` too (which isn't possible; see comment below)
         rewrite_plan: &RewritePlan,
+        offset_plan: Option<OffsetPlan>,
     ) -> (bool, Prepare) {
         let cache_key = CacheKey::Simple {
-            query: query.clone(),
+            query: original_query.clone(),
         };
 
         if let Some(name) = self.reuse(&cache_key) {
@@ -101,13 +110,17 @@ impl GlobalCache {
         let name = self.next_name();
         let prepare = Prepare {
             name: Bytes::from(name.clone()),
-            query,
+            query: rewritten_query.unwrap_or(original_query),
         };
 
         let statement = Statement {
             stmt: StatementType::Prepare {
                 prepare: prepare.clone(),
                 unique_ids: rewrite_plan.unique_ids,
+                // The reason this isn't using [`rewrite_plan.offset`] is that in `rewrite_single_prepared`,
+                // for `PrepareStmt`, we don't set `offset` on`RewritePlan` yet. We only attach `offset`
+                // to the plan for `ExecuteStmt`, and we need access to `OffsetPlan` for both here.
+                offset_plan,
             },
             row_description: None,
             cache_key: cache_key.clone(),
@@ -146,10 +159,13 @@ impl GlobalCache {
     /// Get the [`Prepare`] message for a globally unique prepare statement name.
     pub(crate) fn prepare(&self, name: &str) -> Option<Prepare> {
         self.prepare_and_unique_ids(name)
-            .map(|(prepare, _)| prepare)
+            .map(|(prepare, _, _)| prepare)
     }
 
-    pub(crate) fn prepare_and_unique_ids(&self, name: &str) -> Option<(Prepare, u16)> {
+    pub(crate) fn prepare_and_unique_ids(
+        &self,
+        name: &str,
+    ) -> Option<(Prepare, u16, Option<OffsetPlan>)> {
         self.names
             .get(name)
             .and_then(|p| p.prepare_and_unique_ids())
@@ -376,8 +392,8 @@ mod test {
         let query = Bytes::from("PREPARE __pgdog_template_name AS SELECT $1");
         let parse = Parse::named("client_stmt", "SELECT $1");
 
-        let (_, first) = cache.insert_prepare(query.clone(), &RewritePlan::default());
-        let (_, second) = cache.insert_prepare(query, &RewritePlan::default());
+        let (_, first) = cache.insert_prepare(query.clone(), None, &RewritePlan::default(), None);
+        let (_, second) = cache.insert_prepare(query, None, &RewritePlan::default(), None);
 
         assert_eq!(first, second);
         assert_eq!(cache.len(), 1);

pgdog/src/frontend/prepared_statements/mod.rs

@@ -8,7 +8,7 @@ use parking_lot::RwLock;
 
 use crate::{
     config::PreparedStatementsLevel,
-    frontend::RewritePlan,
+    frontend::{RewritePlan, router::parser::rewrite::statement::offset::OffsetPlan},
     net::{Parse, Prepare, ProtocolMessage},
 };
 
@@ -111,22 +111,26 @@ impl PreparedStatements {
     }
 
     /// Insert PREPARE statement into the cache.
-    ///
-    /// # Arguments
-    ///
-    /// - `parse`: [`Parse`] message, with the prepared statement named by the client.
-    ///
-    /// # Return
-    ///
-    /// Nothing, but the message is renamed to a unique, global name.
-    ///
     pub(crate) fn insert_prepare(
         &mut self,
         name: &str,
-        query: Bytes,
+        original_query: Bytes,
+        rewritten_query: Option<Bytes>,
+        // TODO: I think we should just pass `unique_ids` in here by itself.
+        //       Otherwise, it could be easily confused to want
+        //       to use `RewritePlan` for `offset_plan` too (which isn't possible; see comment below)
         rewrite_plan: &RewritePlan,
+        // Needs to be separate from `RewritePlan`. See comment in `global_cache.rs`.
+        offset_plan: Option<OffsetPlan>,
     ) -> Prepare {
-        let (_new, prepare) = { self.global.write().insert_prepare(query, rewrite_plan) };
+        let (_new, prepare) = {
+            self.global.write().insert_prepare(
+                original_query,
+                rewritten_query,
+                rewrite_plan,
+                offset_plan,
+            )
+        };
 
         self.insert_internal(name, prepare.name());
 
@@ -140,7 +144,10 @@ impl PreparedStatements {
     }
 
     /// Get a globally unique [`Prepare`] message using the client name as key.
-    pub(crate) fn prepare_and_unique_ids(&self, name: &str) -> Option<(Prepare, u16)> {
+    pub(crate) fn prepare_and_unique_ids(
+        &self,
+        name: &str,
+    ) -> Option<(Prepare, u16, Option<OffsetPlan>)> {
         self.local
             .get(name)
             .and_then(|name| self.global.read().prepare_and_unique_ids(name))

pgdog/src/frontend/prepared_statements/statement.rs

@@ -1,4 +1,7 @@
-use crate::{net::Prepare, stats::memory::MemoryUsage};
+use crate::{
+    frontend::router::parser::rewrite::statement::offset::OffsetPlan, net::Prepare,
+    stats::memory::MemoryUsage,
+};
 
 use super::prelude::*;
 
@@ -24,6 +27,11 @@ pub(crate) enum StatementType {
         /// [`Self::prepare`] was previously rewritten to replace those calls
         /// with bind parameter placeholder numbered after all others
         unique_ids: u16,
+
+        /// Used to keep track of LIMIT + OFFSET queries (stemming from Prepare),
+        /// where we have to re-write `A_Const` nodes with `ParamRefs`, so that we can dynamically
+        /// modify limit/offset values before execution if it ends up being cross-shard.
+        offset_plan: Option<OffsetPlan>,
     },
 }
 
@@ -63,12 +71,13 @@ impl Statement {
         }
     }
 
-    pub(super) fn prepare_and_unique_ids(&self) -> Option<(Prepare, u16)> {
+    pub(super) fn prepare_and_unique_ids(&self) -> Option<(Prepare, u16, Option<OffsetPlan>)> {
         match &self.stmt {
             StatementType::Prepare {
                 prepare,
                 unique_ids,
-            } => Some((prepare.clone(), *unique_ids)),
+                offset_plan,
+            } => Some((prepare.clone(), *unique_ids, offset_plan.clone())),
             _ => None,
         }
     }

pgdog/src/frontend/router/parser/limit.rs

@@ -5,7 +5,7 @@ use pg_raw_parse::{
 
 use super::{Error, StatementParameters};
 
-#[derive(Debug, Clone, Copy, Default, PartialEq)]
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
 pub(crate) struct Limit {
     pub(crate) limit: Option<usize>,
     pub(crate) offset: Option<usize>,

pgdog/src/frontend/router/parser/rewrite/statement/error.rs

@@ -40,4 +40,7 @@ pub(crate) enum Error {
 
     #[error("prepared statement '{0}' does not exist")]
     ExecuteMissingPrepare(String),
+
+    #[error("missing or invalid parameters in Execute")]
+    IncorrectExecuteParameters,
 }

pgdog/src/frontend/router/parser/rewrite/statement/mod.rs

@@ -1,12 +1,11 @@
 //! Statement rewriter.
 
-use pg_raw_parse::{Node, NodeMut, make, nodes, transform, walk};
-
 use crate::backend::ShardingSchema;
 use crate::backend::schema::Schema;
 use crate::frontend::PreparedStatements;
 use crate::frontend::router::parser::AstContext;
 use crate::net::parameter::ParameterValue;
+use pg_raw_parse::{Node, NodeMut, make, nodes, transform, walk};
 
 pub(crate) mod aggregate;
 pub(crate) mod auto_id;