diff --git a/integration/rust/tests/integration/simple_prepared.rs b/integration/rust/tests/integration/simple_prepared.rs index 82f552224..e76d9cd62 100644 --- a/integration/rust/tests/integration/simple_prepared.rs +++ b/integration/rust/tests/integration/simple_prepared.rs @@ -32,3 +32,159 @@ async fn test_simple_prepared_ttl() { assert_eq!(test_return.try_get::("?column?").unwrap(), 1); } + +/// +/// 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::("id")) + .collect::>(), + 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::("id")) + .collect::>(), + 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::("id")) + .collect::>(), + 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::("id")) + .collect::>(), + 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::("id")) + .collect::>(), + 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(); +} diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 981b14979..91261352f 100644 --- a/pgdog/src/backend/server.rs +++ b/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 diff --git a/pgdog/src/frontend/prepared_statements/cache_key.rs b/pgdog/src/frontend/prepared_statements/cache_key.rs index f64822954..e69118ca2 100644 --- a/pgdog/src/frontend/prepared_statements/cache_key.rs +++ b/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. } } } diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 39269cfae..4908b10db 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/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, + // 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, ) -> (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 { 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)> { 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); diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index 3e5e4bdf3..3e8b40654 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/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, + // 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, ) -> 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)> { self.local .get(name) .and_then(|name| self.global.read().prepare_and_unique_ids(name)) diff --git a/pgdog/src/frontend/prepared_statements/statement.rs b/pgdog/src/frontend/prepared_statements/statement.rs index 84acbdc28..bfc8fc402 100644 --- a/pgdog/src/frontend/prepared_statements/statement.rs +++ b/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, }, } @@ -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)> { match &self.stmt { StatementType::Prepare { prepare, unique_ids, - } => Some((prepare.clone(), *unique_ids)), + offset_plan, + } => Some((prepare.clone(), *unique_ids, offset_plan.clone())), _ => None, } } diff --git a/pgdog/src/frontend/router/parser/limit.rs b/pgdog/src/frontend/router/parser/limit.rs index 568ae3f77..d8602d1d8 100644 --- a/pgdog/src/frontend/router/parser/limit.rs +++ b/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, pub(crate) offset: Option, diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs index e29203f64..c131e8c92 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs +++ b/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, } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index 22a8a47f5..4f2555f46 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/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; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs index 3bd665e43..8ce84b0d2 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs @@ -1,4 +1,6 @@ -use pg_raw_parse::{ConstValue, Node, Owned, StmtList, nodes}; +use std::ops::Deref; + +use pg_raw_parse::{ConstValue, Node, Owned, StmtList, deparse, nodes}; use crate::frontend::ClientRequest; use crate::frontend::router::parser::Limit; @@ -7,11 +9,12 @@ use crate::net::messages::bind::{Format, Parameter}; use super::*; -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub(crate) struct OffsetPlan { pub(crate) limit: Limit, pub(crate) limit_param: usize, pub(crate) offset_param: usize, + pub(crate) prepare_execute: bool, } impl OffsetPlan { @@ -25,6 +28,10 @@ impl OffsetPlan { return Ok(()); } + if self.prepare_execute { + return self.handle_prepare_execute(request); + } + // Resolve actual values: use literal if known, otherwise read from Bind. let mut limit_val = self.limit.limit; let mut offset_val = self.limit.offset; @@ -106,8 +113,80 @@ impl OffsetPlan { Ok(()) } + + /// `apply_after_parser` helper method for handling Prepare + Execute cases, where + /// we need to re-write limit / offset for multi-shard queries upon execution. + fn handle_prepare_execute(&self, request: &mut ClientRequest) -> Result<(), Error> { + // Assert expectations of what should've happened before this method was called + // in case something beforehand is changed in the future. + assert!( + self.prepare_execute, + "self.prepare_execute was checked before method call" + ); + + let route = request + .route + .as_mut() + .expect("route.is_some() was checked before method call"); + + assert!( + route.is_cross_shard(), + "route.is_cross_shard() was checked before method call" + ); + + let node = &mut request.ast; + let node = node.as_mut().ok_or(Error::MissingAst)?; + let node = node.ast.first().ok_or(Error::MissingAst)?; + + let pg_raw_parse::Node::ExecuteStmt(execute) = node.stmt() else { + unreachable!("The query must be ExecuteStmt to have reached here."); + }; + + let new_execute = pg_raw_parse::make::owned(|mem| { + let mut execute_unique = mem.make_unique(execute); + + let mut mutable_execute_unique = execute_unique.as_mut(); + let mut params = mutable_execute_unique.params_mut(); + + let new_limit = self.limit.limit.unwrap_or(0) + self.limit.offset.unwrap_or(0); + + // These are guarenteed to be `ParamRefs` because of our + // re-write for all `A_Const` nodes for the original `PreparedStmt` that we cached. + params.set( + self.limit_param - 1, + mem.make_a_const(ConstValue::Integer(new_limit as i32)) + .uncast(), + ); + params.set( + self.offset_param - 1, + mem.make_a_const(ConstValue::Integer(0i32)).uncast(), + ); + + execute_unique + }); + + let new_execute = new_execute.deref(); + let new_execute_sql = deparse(new_execute)?; + + // `ExecuteStmt` will be a Query, because this is simple-protocol. + // Replace with our re-written `ExecuteStmt` + // (replacing limit/offset with proper multi-shard vlaues) + for message in request.messages.iter_mut() { + if let ProtocolMessage::Query(query) = message { + query.set_query(new_execute_sql.as_str()); + } + } + + route.set_limit(Limit { + limit: self.limit.limit, + offset: self.limit.offset, + }); + + Ok(()) + } } +#[derive(Debug)] enum LimitValueInfo { Literal(usize), Param(usize), @@ -176,6 +255,7 @@ impl StatementRewrite<'_> { }, limit_param: limit_info.param_index(), offset_param: offset_info.param_index(), + prepare_execute: false, }); } } @@ -324,6 +404,7 @@ mod tests { }, limit_param: 0, offset_param: 0, + prepare_execute: false, }; let mut request = ClientRequest::from(vec![ProtocolMessage::Query(Query::new( "SELECT * FROM t LIMIT 10 OFFSET 5", @@ -353,6 +434,7 @@ mod tests { }, limit_param: 1, offset_param: 2, + prepare_execute: false, }; let mut request = ClientRequest::from(vec![ProtocolMessage::Bind(Bind::new_params( "", @@ -383,6 +465,7 @@ mod tests { }, limit_param: 0, offset_param: 0, + prepare_execute: false, }; let mut request = ClientRequest::from(vec![ProtocolMessage::Query(Query::new( "SELECT * FROM t LIMIT 10 OFFSET 5", @@ -407,6 +490,7 @@ mod tests { }, limit_param: 0, offset_param: 1, + prepare_execute: false, }; let mut request = ClientRequest::from(vec![ ProtocolMessage::Parse(Parse::named("s", "SELECT * FROM t LIMIT 10 OFFSET $1")), diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs index efa7713bc..495f74a71 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -1,8 +1,15 @@ use bytes::Bytes; -use pg_raw_parse::{ConstValue, NodeMut, make::MemoryToken, nodes::ExecuteStmtMut}; +use pg_raw_parse::{ + ConstValue, NodeMut, + make::MemoryToken, + nodes::{ExecuteStmtMut, ParamRef, PrepareStmtMut}, +}; use crate::{ - frontend::PreparedStatements, + frontend::{ + PreparedStatements, + router::parser::{Limit, rewrite::statement::offset::OffsetPlan}, + }, net::{PREPARE_TEMPLATE_NAME, Prepare}, unique_id::UniqueId, }; @@ -84,7 +91,7 @@ fn rewrite_single_prepared<'a>( node: NodeMut<'a, '_>, mem: MemoryToken<'a>, prepared_statements: &mut PreparedStatements, - plan: &RewritePlan, + plan: &mut RewritePlan, ) -> Result { match node { NodeMut::PrepareStmt(mut stmt) => { @@ -93,9 +100,27 @@ fn rewrite_single_prepared<'a>( // Create a globally unique key using the query text // with a hardcoded name. stmt.set_name(Some(mem.copy_string(PREPARE_TEMPLATE_NAME))); - let query = Bytes::from(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned()); - let prepare = prepared_statements.insert_prepare(&client_name, query, plan); + let original_query = Bytes::from(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned()); + + // Is the query a SELECT? Do we have both LIMIT and OFFSET in the SELECT? + let offset_plan: Option = create_offset_plan(mem, &mut stmt); + + let new_query = offset_plan + .as_ref() + .map(|_| { + pg_raw_parse::deparse(&*stmt) + .map(|deparse_result| Bytes::from(deparse_result.as_str().to_owned())) + }) + .transpose()?; + + let prepare = prepared_statements.insert_prepare( + &client_name, + original_query, + new_query, + plan, + offset_plan, + ); stmt.set_name(Some(mem.copy_string(prepare.name()))); @@ -104,9 +129,18 @@ fn rewrite_single_prepared<'a>( NodeMut::ExecuteStmt(mut stmt) => { let stmt_name = stmt.name().expect("EXECUTE always has name"); - if let Some((prepare, unique_ids)) = + + if let Some((prepare, unique_ids, offset_plan)) = prepared_statements.prepare_and_unique_ids(stmt_name) { + if let Some(mut offset_plan) = offset_plan { + // Note: This needs to be ordered before the offset_val/limit_val adjustment. + insert_offset_params(&mut stmt, mem, &offset_plan); + update_offset_plan_fields(&mut offset_plan, &mut stmt)?; + + plan.offset = Some(offset_plan); + } + // Rewrite EXECUTE statement to match the rewrite // we did on the PREPARE statement. insert_unique_ids(&mut stmt, mem, unique_ids)?; @@ -122,6 +156,205 @@ fn rewrite_single_prepared<'a>( } } +/// Helper method for `rewrite_single_prepared` for `ExecuteStatement` +/// Replace the cached OffsetPlan's un-resolved values +/// (originally `ParamRef` nodes; weren't `A_Const` nodes in `PrepareStmt`) +/// with the ones now provided within the `ExecuteStmt` +fn update_offset_plan_fields<'a>( + offset_plan: &mut OffsetPlan, + stmt: &mut ExecuteStmtMut<'a, '_>, +) -> Result<(), Error> { + if offset_plan.limit.offset.is_none() { + let pg_raw_parse::Node::A_Const(constant) = stmt + .params() + .get(offset_plan.offset_param - 1) + .ok_or(Error::IncorrectExecuteParameters)? + else { + return Err(Error::IncorrectExecuteParameters); + }; + + offset_plan.limit.offset = Some( + constant + .val() + .ok_or(Error::IncorrectExecuteParameters)? + .numeric_value::() + .ok_or(Error::IncorrectExecuteParameters)? as usize, + ); + } + + if offset_plan.limit.limit.is_none() { + let pg_raw_parse::Node::A_Const(constant) = stmt + .params() + .get(offset_plan.limit_param - 1) + .ok_or(Error::IncorrectExecuteParameters)? + else { + return Err(Error::IncorrectExecuteParameters); + }; + + offset_plan.limit.limit = Some( + constant + .val() + .ok_or(Error::IncorrectExecuteParameters)? + .numeric_value::() + .ok_or(Error::IncorrectExecuteParameters)? as usize, + ); + } + + Ok(()) +} + +/// Helper method for `rewrite_single_prepared` for `PreparedStatement` +/// to create an `OffsetPlan` based on the SELECT query inside +/// of the `PreparedStatement`, which allows us to store +/// this `OffsetPlan` in the Prepared Statement cache +/// and reference later for re-writes. +fn create_offset_plan<'a>( + mem: MemoryToken<'a>, + stmt: &mut PrepareStmtMut<'a, '_>, +) -> Option { + // Inner query must be SELECT + let pg_raw_parse::Node::SelectStmt(stmt_query) = stmt.query() else { + return None; + }; + + // Must have both LIMIT and OFFSET + if matches!(stmt_query.limit_count(), pg_raw_parse::Node::None) + || matches!(stmt_query.limit_offset(), pg_raw_parse::Node::None) + { + return None; + } + + // Count the `ParamRef` nodes in the query, so that we know what number to start at + // if we need to add some more. + let mut param_refs_count: usize = 0; + pg_raw_parse::walk::walk(stmt_query.into(), |node| { + if let pg_raw_parse::Node::ParamRef(_) = node { + param_refs_count += 1; + } + }); + + // Make a unique copy of the Client's statement to mutate + let mut unique_stmt = mem.make_unique(stmt_query); + let mut unique_stmt_mut = unique_stmt.as_mut(); + + // Replace `A_Const` nodes with `ParamRef` nodes. + // This allows us to dynamically change the LIMIT/OFFSET at execution time, + // if we have a multi-shard query. + let (limit_param, limit_val) = + if let pg_raw_parse::Node::A_Const(limit_count) = stmt_query.limit_count() { + param_refs_count += 1; + + let mut param_ref_count = mem.make_node::(); + param_ref_count.as_mut().set_number(param_refs_count as i32); + unique_stmt_mut.set_limit_count(param_ref_count.uncast()); + + let limit_val = limit_count + .val() + .and_then(|limit_val| limit_val.numeric_value::())?; + + (param_refs_count, Some(limit_val as usize)) + } else if let pg_raw_parse::Node::ParamRef(param_ref) = stmt_query.limit_count() { + (param_ref.number as usize, None) + } else { + return None; + }; + + let (offset_param, offset_val) = + if let pg_raw_parse::Node::A_Const(limit_offset) = stmt_query.limit_offset() { + param_refs_count += 1; + + let mut param_ref_offset = mem.make_node::(); + param_ref_offset + .as_mut() + .set_number(param_refs_count as i32); + unique_stmt_mut.set_limit_offset(param_ref_offset.uncast()); + + let limit_offset_val = limit_offset + .val() + .and_then(|limit_offset_val| limit_offset_val.numeric_value::())?; + + (param_refs_count, Some(limit_offset_val as usize)) + } else if let pg_raw_parse::Node::ParamRef(param_ref) = stmt_query.limit_offset() { + (param_ref.number as usize, None) + } else { + return None; + }; + + // Re-writes the Client's original statement with our version. + stmt.set_query(unique_stmt.uncast()); + + Some(OffsetPlan { + limit: Limit { + limit: limit_val, + offset: offset_val, + }, + limit_param, + offset_param, + prepare_execute: true, + }) +} + +/// Helper method for `rewrite_single_prepared` to handle injecting +/// cached constants to an `ExecuteStmt` which we previously stripped +/// from the `PrepareStmt`, so that we could dynamically re-write later +/// if `Route` resolves to multi-shard +fn insert_offset_params<'a>( + stmt: &mut ExecuteStmtMut<'a, '_>, + mem: MemoryToken<'a>, + offset_plan: &OffsetPlan, +) { + let offset_val = offset_plan.limit.offset; + let limit_val = offset_plan.limit.limit; + let offset_pos = offset_plan.offset_param; + let limit_pos = offset_plan.limit_param; + + let mut params = stmt.params_mut(); + + if let Some(offset_val) = offset_val + && let Some(limit_val) = limit_val + { + // In this case, both nodes from the `PrepareStmt` were `A_Const` + // Therefore, we need to use the cached values, and inject them + // into the `ExecuteStmt`'s `params`. + // + // The if statements are to determine which one goes first in params, + // which is based on the refs ($1, $2) we chose + // + // These are deterministically ordered (since we do them), but it's just one + // extra check to do this, and doesn't try to enforce an invariant. + + let limit_val_node = mem + .make_a_const(ConstValue::Integer(limit_val as i32)) + .uncast(); + let offset_val_node = mem + .make_a_const(ConstValue::Integer(offset_val as i32)) + .uncast(); + + let (first_node, second_node) = if offset_pos > limit_pos { + (limit_val_node, offset_val_node) + } else { + (offset_val_node, limit_val_node) + }; + + params.push(mem, first_node); + params.push(mem, second_node); + } else if let Some(offset_val) = offset_val { + // Only OFFSET was `A_Const` + stmt.params_mut().push( + mem, + mem.make_a_const(ConstValue::Integer(offset_val as i32)) + .uncast(), + ); + } else if let Some(limit_val) = limit_val { + // Only LIMIT was `A_Const` + stmt.params_mut().push( + mem, + mem.make_a_const(ConstValue::Integer(limit_val as i32)) + .uncast(), + ); + } +} + fn insert_unique_ids<'a>( stmt: &mut ExecuteStmtMut<'a, '_>, mem: MemoryToken<'a>, @@ -263,6 +496,127 @@ mod tests { assert_eq!(ids.len(), 3, "all appended IDs should be unique"); } + /// Prepares two statements, both using LIMIT/OFFSET. One is re-written, one isn't. + /// They should resolve to two different entries in the `GlobalCache` (two different `CacheKeys`) + /// and their respective `CachedStmt` entries should contain the correct `OffsetPlan` information. + /// Integration test `test_simple_prepared_limit` tests end-to-end functionality. + #[test] + fn test_rewrite_prepare_offset_limit_cache_differs() { + let saved_first_time_sql; + + let mut ctx = TestContext::new(); + { + let (sql, plan) = ctx + .rewrite("PREPARE test_stmt AS SELECT * FROM sharded LIMIT 5 OFFSET 10") + .unwrap(); + + saved_first_time_sql = sql.clone(); + + assert!( + sql.contains("__pgdog_"), + "PREPARE should be renamed to __pgdog_N, got: {sql}" + ); + assert!( + !sql.contains("test_stmt"), + "original name should be replaced: {sql}" + ); + assert_eq!(plan.prepare_rewrites.len(), 1); + assert!(plan.stmt.is_some()); + + let prepare = &plan.prepare_rewrites[0]; + match prepare { + PrepareExecute::Prepare(prepare) => { + assert!(prepare.name().starts_with("__pgdog_")); + assert_eq!( + prepare.query(), + "PREPARE __pgdog_template_name AS SELECT * FROM sharded LIMIT $1 OFFSET $2" + ); + } + + _ => panic!("expected PrepareExecute::Prepare"), + } + + // Verify 1 local, 1 global before we re-try with $1, $2 in the next block. + assert_eq!(ctx.ps.local.len(), 1); + assert_eq!(ctx.ps.global.read().len(), 1); + + // Verify the OffsetPlan is correct from the PreparedStatement name used. + let (fetched_prepare, _, offset_plan) = + ctx.ps.prepare_and_unique_ids("test_stmt").unwrap(); + let offset_plan = offset_plan.unwrap(); + assert_eq!( + fetched_prepare.query, + "PREPARE __pgdog_template_name AS SELECT * FROM sharded LIMIT $1 OFFSET $2" + ); + assert!(offset_plan.prepare_execute); + + // Verifies these numbers were actually saved in the cache. + // In the next block, verify these are NOT present (correctly creating a new global entry) + assert_eq!(offset_plan.limit.limit, Some(5)); + assert_eq!(offset_plan.limit.offset, Some(10)); + } + + // LIMIT $1 OFFSET $2; assert OffsetPlan uses None instead of Some(5), Some(10) + // + { + // Notice that this is the resolved statement the last segment was re-written to (for the cache key) + // We're using a different name here to assert they resolve differently and to different global statements. + let (sql, plan) = ctx + .rewrite("PREPARE test_stmt2 AS SELECT * FROM sharded LIMIT $1 OFFSET $2") + .unwrap(); + + // Ensures that the global names differ. (diff global cache entries, diff OffsetPlans) + // "PREPARE __pgdog_1 AS SELECT * FROM sharded LIMIT $1 OFFSET $2" + // "PREPARE __pgdog_2 AS SELECT * FROM sharded LIMIT $1 OFFSET $2" + assert_ne!(sql, saved_first_time_sql); + + assert!( + sql.contains("__pgdog_"), + "PREPARE should be renamed to __pgdog_N, got: {sql}" + ); + assert!( + !sql.contains("test_stmt2"), + "original name should be replaced: {sql}" + ); + assert_eq!(plan.prepare_rewrites.len(), 1); + assert!(plan.stmt.is_some()); + + let prepare = &plan.prepare_rewrites[0]; + match prepare { + PrepareExecute::Prepare(prepare) => { + assert!(prepare.name().starts_with("__pgdog_")); + assert_eq!( + prepare.query(), + "PREPARE __pgdog_template_name AS SELECT * FROM sharded LIMIT $1 OFFSET $2" + ); + } + + _ => panic!("expected PrepareExecute::Prepare"), + } + + // Now there are **TWO** local entries. + assert_eq!(ctx.ps.local.len(), 2); + // Now there are **TWO** global entries. + assert_eq!(ctx.ps.global.read().len(), 2); + + // Verify the OffsetPlan is correct using the PreparedStatement name used. + let (fetched_prepare, _, offset_plan) = + ctx.ps.prepare_and_unique_ids("test_stmt2").unwrap(); + let offset_plan = offset_plan.unwrap(); + assert_eq!( + fetched_prepare.query, + "PREPARE __pgdog_template_name AS SELECT * FROM sharded LIMIT $1 OFFSET $2" + ); + assert!(offset_plan.prepare_execute); + + // If we saw Some(5) and Some(10) here, they would have resolved to the last block (in the same context). + // Since they don't, that means it's correctly using the pre-re-written `Query` as the `CacheKey`. + // They're correctly not resolving to the same `CachedStmt`. + assert_eq!(offset_plan.limit.limit, None); + assert_eq!(offset_plan.limit.offset, None); + } + } + #[test] fn test_rewrite_prepare() { let mut ctx = TestContext::new();