From 4f1b98cd4b023d0ec5bce13f129ace99e9f6cf0e Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sat, 15 Aug 2026 14:14:28 -0700 Subject: [PATCH 01/22] feat: save simple prepared in global cache --- pgdog-postgres-types/src/lib.rs | 2 + pgdog-postgres-types/src/type_name.rs | 39 ++ pgdog/src/backend/prepared_statements.rs | 6 +- pgdog/src/backend/server.rs | 154 ++++---- pgdog/src/frontend/client/mod.rs | 1 - pgdog/src/frontend/client/query_engine/mod.rs | 2 +- .../test/close_parse_global_cache.rs | 2 +- .../test/rewrite_simple_prepared.rs | 2 +- pgdog/src/frontend/mod.rs | 2 +- .../frontend/prepared_statements/cache_key.rs | 54 +++ .../prepared_statements/cached_statement.rs | 26 ++ .../prepared_statements/global_cache.rs | 355 ++++++------------ .../prepared_statements/maintenance.rs | 28 ++ pgdog/src/frontend/prepared_statements/mod.rs | 199 +++++----- .../frontend/prepared_statements/prelude.rs | 7 + .../frontend/prepared_statements/rewrite.rs | 4 +- .../frontend/prepared_statements/statement.rs | 84 +++++ .../router/parser/rewrite/statement/plan.rs | 9 +- .../rewrite/statement/simple_prepared.rs | 58 +-- pgdog/src/net/messages/mod.rs | 2 + pgdog/src/net/messages/prepare.rs | 72 ++++ pgdog/src/net/protocol_message.rs | 10 +- 22 files changed, 678 insertions(+), 440 deletions(-) create mode 100644 pgdog-postgres-types/src/type_name.rs create mode 100644 pgdog/src/frontend/prepared_statements/cache_key.rs create mode 100644 pgdog/src/frontend/prepared_statements/cached_statement.rs create mode 100644 pgdog/src/frontend/prepared_statements/maintenance.rs create mode 100644 pgdog/src/frontend/prepared_statements/prelude.rs create mode 100644 pgdog/src/frontend/prepared_statements/statement.rs create mode 100644 pgdog/src/net/messages/prepare.rs diff --git a/pgdog-postgres-types/src/lib.rs b/pgdog-postgres-types/src/lib.rs index 59af72502..b6fa2e643 100644 --- a/pgdog-postgres-types/src/lib.rs +++ b/pgdog-postgres-types/src/lib.rs @@ -16,6 +16,7 @@ pub mod smallint; pub mod text; pub mod timestamp; pub mod timestamptz; +pub mod type_name; pub mod uuid; pub mod vector; @@ -32,3 +33,4 @@ pub use numeric::Numeric; pub use oid::Oid; pub use timestamp::Timestamp; pub use timestamptz::TimestampTz; +pub use type_name::TypeName; diff --git a/pgdog-postgres-types/src/type_name.rs b/pgdog-postgres-types/src/type_name.rs new file mode 100644 index 000000000..d417adead --- /dev/null +++ b/pgdog-postgres-types/src/type_name.rs @@ -0,0 +1,39 @@ +//! PostgreSQL type name. + +use std::fmt; +use std::sync::Arc; + +/// Name of a PostgreSQL data type. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct TypeName(pub Arc>); + +impl From> for TypeName { + fn from(names: Vec) -> Self { + Self(Arc::new(names)) + } +} + +impl FromIterator for TypeName { + fn from_iter>(iter: T) -> Self { + Self(Arc::new(iter.into_iter().collect())) + } +} + +impl fmt::Display for TypeName { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0.join(".")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn conversions() { + let data_type = TypeName::from(vec![String::from("pg_catalog"), String::from("int8")]); + + assert_eq!(*data_type.0, ["pg_catalog", "int8"]); + assert_eq!(data_type.to_string(), "pg_catalog.int8"); + } +} diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 17665d0bd..16369856f 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -302,11 +302,11 @@ impl PreparedStatements { self.state.add('3'); } } - ProtocolMessage::Prepare { name, .. } => { - if self.contains(name) { + ProtocolMessage::Prepare(prepare) => { + if self.contains(prepare.name()) { return Ok(HandleResult::Drop); } else { - self.parses.push_back(name.clone()); + self.parses.push_back(prepare.name().to_string()); self.state.add_ignore('C'); // Prepare turns into a Simple Query ('Q') so it expects a regular RFQ back. diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 5ab6e9e0a..fd8caef98 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -2312,83 +2312,83 @@ pub mod test { assert!(!server.has_more_messages()); } - #[tokio::test] - async fn test_manual_prepared() { - let mut server = test_server().await; - - let mut prep = PreparedStatements::new(); - let mut parse = Parse::named("test", "SELECT 1::bigint"); - prep.insert_prepare(&mut parse); - assert_eq!(parse.name(), "__pgdog_1"); - - server - .send( - &vec![ProtocolMessage::from(Query::new(format!( - "PREPARE {} AS {}", - parse.name(), - parse.query() - )))] - .into(), - ) - .await - .unwrap(); - for c in ['C', 'Z'] { - let msg = server.read().await.unwrap(); - assert_eq!(msg.code(), c); - } - assert!(server.sync_prepared()); - server.sync_prepared_statements().await.unwrap(); - assert!(server.prepared_statements.contains("__pgdog_1")); - - let describe = Describe::new_statement("__pgdog_1"); - let bind = Bind::new_statement("__pgdog_1"); - let execute = Execute::new(); - server - .send( - &vec![ - describe.clone().into(), - bind.into(), - execute.into(), - ProtocolMessage::from(Sync), - ] - .into(), - ) - .await - .unwrap(); - - for c in ['t', 'T', '2', 'D', 'C', 'Z'] { - let msg = server.read().await.unwrap(); - assert_eq!(c, msg.code()); - } - - let parse = Parse::named("__pgdog_1", "SELECT 2::bigint"); - let describe = describe.clone(); - - server - .send(&vec![parse.into(), describe.into(), ProtocolMessage::from(Flush)].into()) - .await - .unwrap(); - - for c in ['1', 't', 'T'] { - let msg = server.read().await.unwrap(); - assert_eq!(msg.code(), c); - } - - server - .send(&vec![ProtocolMessage::from(Query::new("EXECUTE __pgdog_1"))].into()) - .await - .unwrap(); - for c in ['T', 'D', 'C', 'Z'] { - let msg = server.read().await.unwrap(); - assert_eq!(c, msg.code()); - if c == 'D' { - let data_row = DataRow::from_bytes(msg.to_bytes()).unwrap(); - let result: i64 = data_row.get(0, Format::Text).unwrap(); - assert_eq!(result, 1); // We prepared SELECT 1, SELECT 2 is ignored. - } - } - assert!(server.done()); - } + // #[tokio::test] + // async fn test_manual_prepared() { + // let mut server = test_server().await; + + // let mut prep = PreparedStatements::new(); + // let mut parse = Parse::named("test", "SELECT 1::bigint"); + // prep.insert_prepare(&mut parse); + // assert_eq!(parse.name(), "__pgdog_1"); + + // server + // .send( + // &vec![ProtocolMessage::from(Query::new(format!( + // "PREPARE {} AS {}", + // parse.name(), + // parse.query() + // )))] + // .into(), + // ) + // .await + // .unwrap(); + // for c in ['C', 'Z'] { + // let msg = server.read().await.unwrap(); + // assert_eq!(msg.code(), c); + // } + // assert!(server.sync_prepared()); + // server.sync_prepared_statements().await.unwrap(); + // assert!(server.prepared_statements.contains("__pgdog_1")); + + // let describe = Describe::new_statement("__pgdog_1"); + // let bind = Bind::new_statement("__pgdog_1"); + // let execute = Execute::new(); + // server + // .send( + // &vec![ + // describe.clone().into(), + // bind.into(), + // execute.into(), + // ProtocolMessage::from(Sync), + // ] + // .into(), + // ) + // .await + // .unwrap(); + + // for c in ['t', 'T', '2', 'D', 'C', 'Z'] { + // let msg = server.read().await.unwrap(); + // assert_eq!(c, msg.code()); + // } + + // let parse = Parse::named("__pgdog_1", "SELECT 2::bigint"); + // let describe = describe.clone(); + + // server + // .send(&vec![parse.into(), describe.into(), ProtocolMessage::from(Flush)].into()) + // .await + // .unwrap(); + + // for c in ['1', 't', 'T'] { + // let msg = server.read().await.unwrap(); + // assert_eq!(msg.code(), c); + // } + + // server + // .send(&vec![ProtocolMessage::from(Query::new("EXECUTE __pgdog_1"))].into()) + // .await + // .unwrap(); + // for c in ['T', 'D', 'C', 'Z'] { + // let msg = server.read().await.unwrap(); + // assert_eq!(c, msg.code()); + // if c == 'D' { + // let data_row = DataRow::from_bytes(msg.to_bytes()).unwrap(); + // let result: i64 = data_row.get(0, Format::Text).unwrap(); + // assert_eq!(result, 1); // We prepared SELECT 1, SELECT 2 is ignored. + // } + // } + // assert!(server.done()); + // } #[tokio::test] async fn test_sync_params() { diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 1eb869705..8dc0e221c 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -715,7 +715,6 @@ impl Client { impl Drop for Client { fn drop(&mut self) { self.comms.disconnect(); - self.prepared_statements.close_all(); } } diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index d90994c3a..e607b3bfe 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -280,7 +280,7 @@ impl QueryEngine { self.stats.state = state; self.stats - .prepared_statements(context.prepared_statements.len_local()); + .prepared_statements(context.prepared_statements.num_statements()); self.stats.memory_used(context.memory_stats); self.comms.update_stats(self.stats); diff --git a/pgdog/src/frontend/client/query_engine/test/close_parse_global_cache.rs b/pgdog/src/frontend/client/query_engine/test/close_parse_global_cache.rs index 37785e383..82cc350d3 100644 --- a/pgdog/src/frontend/client/query_engine/test/close_parse_global_cache.rs +++ b/pgdog/src/frontend/client/query_engine/test/close_parse_global_cache.rs @@ -33,7 +33,7 @@ async fn test_close_parse_same_name_global_cache() { assert_eq!(cached_query, "SELECT $1"); // Verify the client's local cache - assert_eq!(client.client().prepared_statements.len_local(), 1); + assert_eq!(client.client().prepared_statements.num_statements(), 1); assert!( client .client() diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs index 08f87f829..9687ec329 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs @@ -43,7 +43,7 @@ async fn test_rewrite_prepare() { .await; assert!( - matches!(messages[0].clone(), ProtocolMessage::Prepare { name, statement } if name == "__pgdog_1" && statement == "SELECT $1, $2, $3") + matches!(messages[0].clone(), ProtocolMessage::Prepare(prepare) if prepare.name() == "__pgdog_1" && prepare.query() == "SELECT $1, $2, $3") ); assert!( diff --git a/pgdog/src/frontend/mod.rs b/pgdog/src/frontend/mod.rs index 284b777b0..915be3ab6 100644 --- a/pgdog/src/frontend/mod.rs +++ b/pgdog/src/frontend/mod.rs @@ -22,7 +22,7 @@ pub use client_request::ClientRequest; pub use comms::{ClientComms, Comms}; pub use connected_client::ConnectedClient; pub(crate) use error::Error; -pub use prepared_statements::{PreparedStatements, Rewrite}; +pub(crate) use prepared_statements::PreparedStatements; #[cfg(debug_assertions)] pub use query_logger::QueryLogger; pub(crate) use regex_parser::RegexParser; diff --git a/pgdog/src/frontend/prepared_statements/cache_key.rs b/pgdog/src/frontend/prepared_statements/cache_key.rs new file mode 100644 index 000000000..2d2346ac1 --- /dev/null +++ b/pgdog/src/frontend/prepared_statements/cache_key.rs @@ -0,0 +1,54 @@ +use pgdog_postgres_types::TypeName; + +use crate::stats::memory::MemoryUsage; + +use super::prelude::*; + +/// Prepared statements cache key. +/// +/// If two `Extended` keys match, it's effectively the same statement. +/// If they don't, e.g. client sent the same query but +/// with different data types, we can't re-use it and +/// need to plan a new one. +/// +/// 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 enum CacheKey { + Extended { + query: Bytes, + data_types: Bytes, + }, + Simple { + query: Bytes, + data_types: Vec, + }, +} + +impl MemoryUsage for CacheKey { + #[inline] + fn memory_usage(&self) -> usize { + // The Bytes alias the Parse in Statement, which counts them via Parse::len. + std::mem::size_of::() + } +} + +impl CacheKey { + pub(super) fn query_ref(&self) -> &Bytes { + match self { + Self::Extended { query, .. } => query, + Self::Simple { query, .. } => query, + } + } + + /// Get a UTF-8 encoded query string + /// stored in the cache. + pub(crate) fn query(&self) -> Result<&str, crate::net::Error> { + let query = self.query_ref(); + + // Postgres string. + Ok(from_utf8(&query[0..query.len() - 1])?) + } +} diff --git a/pgdog/src/frontend/prepared_statements/cached_statement.rs b/pgdog/src/frontend/prepared_statements/cached_statement.rs new file mode 100644 index 000000000..1b496a6df --- /dev/null +++ b/pgdog/src/frontend/prepared_statements/cached_statement.rs @@ -0,0 +1,26 @@ +use crate::stats::memory::MemoryUsage; + +pub(crate) type Counter = usize; + +pub(crate) fn global_name(counter: Counter) -> String { + format!("__pgdog_{}", counter) +} + +#[derive(Debug, Copy, Clone)] +pub struct CachedStmt { + pub counter: Counter, + pub used: usize, +} + +impl MemoryUsage for CachedStmt { + #[inline] + fn memory_usage(&self) -> usize { + self.counter.memory_usage() + self.used.memory_usage() + } +} + +impl CachedStmt { + pub fn name(&self) -> String { + global_name(self.counter) + } +} diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 4ebdb3908..72d8af3a7 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -1,112 +1,16 @@ -use bytes::Bytes; - use crate::{ - net::messages::{Parse, RowDescription}, + net::{ + Prepare, + messages::{Parse, RowDescription}, + }, stats::memory::MemoryUsage, }; -use std::{collections::hash_map::HashMap, str::from_utf8}; +use std::collections::hash_map::HashMap; +use bytes::Bytes; use fnv::FnvHashSet as HashSet; -/// Identity of a prepared statement inside the global cache. -pub type Counter = usize; - -// Format the globally unique prepared statement -// name based on the counter. -fn global_name(counter: Counter) -> String { - format!("__pgdog_{}", counter) -} - -#[derive(Debug, Clone)] -pub struct Statement { - parse: Parse, - rewrite: Option, - row_description: Option, - cache_key: CacheKey, -} - -impl MemoryUsage for Statement { - #[inline] - fn memory_usage(&self) -> usize { - self.parse.len() - + if let Some(row_description) = &self.row_description { - row_description.memory_usage() - } else { - 0 - } - + self.cache_key.memory_usage() - } -} - -impl Statement { - pub fn query(&self) -> &str { - self.parse.query() - } - - fn cache_key(&self) -> &CacheKey { - &self.cache_key - } -} - -/// Prepared statements cache key. -/// -/// If two `Extended` keys match, it's effectively the same statement. -/// If they don't, e.g. client sent the same query but -/// with different data types, we can't re-use it and -/// need to plan a new one. -/// -/// 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 enum CacheKey { - Extended { query: Bytes, data_types: Bytes }, - Simple { query: Bytes, unique: Counter }, -} - -impl MemoryUsage for CacheKey { - #[inline] - fn memory_usage(&self) -> usize { - // The Bytes alias the Parse in Statement, which counts them via Parse::len. - std::mem::size_of::() - } -} - -impl CacheKey { - fn query_ref(&self) -> &Bytes { - match self { - Self::Extended { query, .. } => query, - Self::Simple { query, .. } => query, - } - } - - pub fn query(&self) -> Result<&str, crate::net::Error> { - let query = self.query_ref(); - - // Postgres string. - Ok(from_utf8(&query[0..query.len() - 1])?) - } -} - -#[derive(Debug, Copy, Clone)] -pub struct CachedStmt { - pub counter: Counter, - pub used: usize, -} - -impl MemoryUsage for CachedStmt { - #[inline] - fn memory_usage(&self) -> usize { - self.counter.memory_usage() + self.used.memory_usage() - } -} - -impl CachedStmt { - pub fn name(&self) -> String { - global_name(self.counter) - } -} +use super::*; /// Global prepared statements cache. /// @@ -145,107 +49,60 @@ impl GlobalCache { /// /// If the statement exists, no entry is created /// and the global name is returned instead. - pub fn insert(&mut self, parse: &Parse) -> (bool, String) { - let parse_key = CacheKey::Extended { + pub(crate) fn insert(&mut self, parse: &Parse) -> (bool, String) { + let name = self.next_name(); + + let cache_key = CacheKey::Extended { query: parse.query_ref(), data_types: parse.data_types_ref(), }; + let statement = Statement { + stmt: StatementType::Parse { + parse: parse.clone(), + rewrite: None, + }, + cache_key: cache_key.clone(), + row_description: None, + }; - if let Some(entry) = self.statements.get_mut(&parse_key) { - if entry.used == 0 { - self.unused.remove(&entry.counter); - } - entry.used += 1; - - (false, global_name(entry.counter)) - } else { - self.counter += 1; - let name = global_name(self.counter); - // PERF: we explicitly create the new parse with renamed - // to reallocate the data and not to hold original buffer - // from pgdog/src/frontend/client/mod.rs - // smaller memory footprints and smaller allocations, since - // the client buffer is generally bigger than the query. - // Holding onto it would also fragment that buffer: the client - // keeps appending messages to it and can't free the middle, - // so the buffer grows monotonically. - let parse = parse.renamed(&name); - - let cache_key = CacheKey::Extended { - query: parse.query_ref(), - data_types: parse.data_types_ref(), - }; - - self.statements.insert( - cache_key.clone(), - CachedStmt { - counter: self.counter, - used: 1, - }, - ); - - self.names.insert( - name.clone(), - Statement { - parse, - cache_key, - rewrite: None, - row_description: None, - }, - ); + let new = self.insert_internal(&name, cache_key, statement); - (true, name) - } + (new, name) } - /// Insert a prepared statement into the global cache ignoring - /// duplicate check. - /// - /// SQL `PREPARE` gets a key of its own, so it is never handed to - /// a second client. It is tracked and evicted like any other statement. - pub fn insert_prepare(&mut self, parse: &Parse) -> String { - self.counter += 1; + /// Insert a statement prepared using the simple protocol into the global cache. + pub(super) fn insert_prepare( + &mut self, + query: Bytes, + data_types: Vec, + ) -> (bool, String) { + let name = self.next_name(); - let name = global_name(self.counter); - let parse = parse.renamed(&name); - // insert_anyway is used for the simple query PREPARE call - // and here the `unique` field based on counter defines - // that this statement won't be reused with other clients - // i.e. it'll always have `used <= 1` and will be closed - // only by the specific client that created it. - // The close happens when the client re-uses the same PREPARE - // name, or on client disconnect in the close_all call. - // TODO: a direct DEALLOCATE won't close it yet. let cache_key = CacheKey::Simple { - query: parse.query_ref(), - unique: self.counter, + query: query.clone(), + data_types: data_types.clone(), }; - self.statements.insert( - cache_key.clone(), - CachedStmt { - counter: self.counter, - used: 1, - }, - ); - - self.names.insert( - name.clone(), - Statement { - parse, - cache_key, - rewrite: None, - row_description: None, + let statement = Statement { + stmt: StatementType::Prepare { + prepare: Prepare { + name: Bytes::from(name.clone()), + query: query, + data_types, + }, }, - ); + row_description: None, + cache_key: cache_key.clone(), + }; - name + let new = self.insert_internal(&name, cache_key, statement); + (new, name) } /// Rewrite prepared statement in the global cache. - pub fn rewrite(&mut self, parse: &Parse) { + pub(crate) fn rewrite(&mut self, parse: &Parse) { if let Some(stmt) = self.names.get_mut(parse.name()) { - stmt.rewrite = Some(parse.clone()); + stmt.set_rewrite(parse); } } @@ -272,7 +129,12 @@ impl GlobalCache { /// It can be used to prepare this statement on a server connection /// or to inspect the original query. pub fn parse(&self, name: &str) -> Option { - self.names.get(name).map(|p| p.parse.clone()) + self.names.get(name).map(|p| p.parse().clone()).flatten() + } + + /// Get the [`Prepare`] message for a globally unique prepare statement name. + pub(super) fn prepare(&self, name: &str) -> Option { + self.names.get(name).map(|p| p.prepare()).flatten() } /// Get the rewritten Parse statement. @@ -282,15 +144,16 @@ impl GlobalCache { pub fn rewritten_parse(&self, name: &str) -> Option { self.names .get(name) - .map(|p| p.rewrite.clone().unwrap_or(p.parse.clone())) + .map(|p| p.rewrite().clone().or(p.parse())) + .flatten() } /// Returns true if this prepared statement has been /// rewritten by the rewrite engine. - pub fn is_rewritten(&self, name: &str) -> bool { + pub(crate) fn is_rewritten(&self, name: &str) -> bool { self.names .get(name) - .map(|p| p.rewrite.is_some()) + .map(|p| p.rewrite().is_some()) .unwrap_or_default() } @@ -298,22 +161,23 @@ impl GlobalCache { /// /// It can be used to decode results received from executing the prepared /// statement. - pub fn row_description(&self, name: &str) -> Option { + pub(crate) fn row_description(&self, name: &str) -> Option { self.names.get(name).and_then(|p| p.row_description.clone()) } /// Number of prepared statements in the local cache. - pub fn len(&self) -> usize { + pub(crate) fn len(&self) -> usize { self.statements.len() } /// True if the local cache is empty. - pub fn is_empty(&self) -> bool { + #[cfg(test)] + pub(crate) fn is_empty(&self) -> bool { self.len() == 0 } /// Close prepared statement. - pub fn close(&mut self, name: &str) { + pub(crate) fn close(&mut self, name: &str) { if let Some(statement) = self.names.get(name) { let key = statement.cache_key(); @@ -329,7 +193,7 @@ impl GlobalCache { /// Close unused statements until the cache is down to `capacity` entries; /// `0` removes everything not in use. Statements in use stay, and global /// names are never reused. - pub fn close_unused(&mut self, capacity: usize) -> usize { + pub(crate) fn close_unused(&mut self, capacity: usize) -> usize { let over = self.len().saturating_sub(capacity); // move out of unused to mutate it without borrowing the self to be able to call self.remove later @@ -356,6 +220,16 @@ impl GlobalCache { removed } + /// Get all prepared statements in the global cache, keyed by name. + pub(crate) fn names(&self) -> &HashMap { + &self.names + } + + /// Get all prepared statements in the global cache, keyed by global unique key. + pub(crate) fn statements(&self) -> &HashMap { + &self.statements + } + /// Remove statement from global cache. fn remove(&mut self, name: &str) { if let Some(stmt) = self.names.remove(name) { @@ -363,13 +237,30 @@ impl GlobalCache { } } - /// Get all prepared statements by name. - pub fn names(&self) -> &HashMap { - &self.names + fn next_name(&mut self) -> String { + self.counter += 1; + global_name(self.counter) } - pub fn statements(&self) -> &HashMap { - &self.statements + fn insert_internal(&mut self, name: &str, cache_key: CacheKey, statement: Statement) -> bool { + if let Some(entry) = self.statements.get_mut(&cache_key) { + if entry.used == 0 { + self.unused.remove(&entry.counter); + } + entry.used += 1; + + false + } else { + self.statements.insert( + cache_key.clone(), + CachedStmt { + counter: self.counter, + used: 1, + }, + ); + self.names.insert(name.to_owned(), statement); + true + } } } @@ -402,7 +293,7 @@ mod test { let stored = cache.names.get(&name).unwrap(); let map_key = cache.statements.keys().next().unwrap(); - let owned = stored.parse.query_ref(); + let owned = stored.parse().expect("parse").query_ref(); assert_eq!(owned.as_ptr(), stored.cache_key.query_ref().as_ptr()); assert_eq!(owned.as_ptr(), map_key.query_ref().as_ptr()); @@ -440,9 +331,9 @@ mod test { assert_eq!(entry.used, 0); assert!(cache.unused.contains(&1)); // __pgdog_1 - let name = cache.insert_prepare(&parse); - cache.close(&name); - assert!(cache.unused.contains(&2)); // __pgdog_2 + // let (_, name) = cache.insert_prepare(&parse); + // cache.close(&name); + // assert!(cache.unused.contains(&2)); // __pgdog_2 } fn used(cache: &GlobalCache, name: &str) -> usize { @@ -450,33 +341,33 @@ mod test { cache.statements.get(statement.cache_key()).unwrap().used } - #[test] - fn test_simple_prepared_is_never_shared() { - let mut cache = GlobalCache::default(); - let parse = Parse::named("client_stmt", "SELECT $1"); - - let first = cache.insert_prepare(&parse); - let second = cache.insert_prepare(&parse); - - assert_ne!(first, second); - assert_eq!(cache.len(), 2); - assert_eq!(used(&cache, &first), 1); - assert_eq!(used(&cache, &second), 1); - - // A Parse never re-uses a SQL PREPARE statement. - let (new, extended) = cache.insert(&parse); - assert!(new); - assert_ne!(extended, first); - assert_ne!(extended, second); - assert_eq!(cache.len(), 3); - - // A Parse re-uses another Parse. - let (new_again, shared) = cache.insert(&parse); - assert!(!new_again); - assert_eq!(shared, extended); - assert_eq!(cache.len(), 3); - assert_eq!(used(&cache, &extended), 2); - } + // #[test] + // fn test_simple_prepared_is_never_shared() { + // let mut cache = GlobalCache::default(); + // let parse = Parse::named("client_stmt", "SELECT $1"); + + // let (_, first) = cache.insert_prepare(&parse); + // let (_, second) = cache.insert_prepare(&parse); + + // assert_ne!(first, second); + // assert_eq!(cache.len(), 2); + // assert_eq!(used(&cache, &first), 1); + // assert_eq!(used(&cache, &second), 1); + + // // A Parse never re-uses a SQL PREPARE statement. + // let (new, extended) = cache.insert(&parse); + // assert!(new); + // assert_ne!(extended, first); + // assert_ne!(extended, second); + // assert_eq!(cache.len(), 3); + + // // A Parse re-uses another Parse. + // let (new_again, shared) = cache.insert(&parse); + // assert!(!new_again); + // assert_eq!(shared, extended); + // assert_eq!(cache.len(), 3); + // assert_eq!(used(&cache, &extended), 2); + // } #[test] fn test_remove_unused() { diff --git a/pgdog/src/frontend/prepared_statements/maintenance.rs b/pgdog/src/frontend/prepared_statements/maintenance.rs new file mode 100644 index 000000000..1978a0c0d --- /dev/null +++ b/pgdog/src/frontend/prepared_statements/maintenance.rs @@ -0,0 +1,28 @@ +use super::{PreparedStatements, prelude::*}; +use crate::tasks::{shutdown_signal, spawn}; + +use std::time::Duration; + +/// Run prepared statements maintenance task every second. +/// +/// Public because it's used in main.rs. +pub fn start_maintenance() { + spawn("prepared statements cache", async move { + debug!("prepared statements cache maintenance started"); + let shutdown = shutdown_signal(); + loop { + tokio::select! { + _ = safe_sleep(Duration::from_secs(1)) => {} + _ = shutdown.cancelled() => break, + } + run_maintenance(); + } + }); +} + +/// Check prepared statements cache for overflows +/// and remove any unused statements exceeding the limit. +fn run_maintenance() { + let capacity = config().config.general.prepared_statements_limit; + PreparedStatements::global().write().close_unused(capacity); +} diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index 0658c9025..6f4714996 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -1,33 +1,38 @@ //! Prepared statements cache. -use std::{collections::HashMap, sync::Arc, time::Duration}; +use std::{collections::HashMap, sync::Arc}; +use bytes::Bytes; use once_cell::sync::Lazy; use parking_lot::RwLock; -use tracing::debug; +use pgdog_postgres_types::TypeName; -use crate::util::safe_sleep; use crate::{ - config::{PreparedStatements as PreparedStatementsLevel, config}, - net::{Parse, ProtocolMessage}, + config::PreparedStatements as PreparedStatementsLevel, + net::{Parse, Prepare, ProtocolMessage}, }; +mod cache_key; +mod cached_statement; pub mod error; pub mod global_cache; +mod maintenance; +mod prelude; pub mod rewrite; - -pub use error::Error; -pub use global_cache::GlobalCache; -pub use rewrite::Rewrite; +pub mod statement; + +pub(crate) use cache_key::CacheKey; +pub(crate) use cached_statement::global_name; +pub(crate) use cached_statement::{CachedStmt, Counter}; +pub(crate) use error::Error; +pub(crate) use global_cache::GlobalCache; +// Maintenance tasks are spawned in main.rs. +pub use maintenance::*; +pub(crate) use rewrite::Rewrite; +pub(crate) use statement::{Statement, StatementType}; static CACHE: Lazy = Lazy::new(PreparedStatements::default); -/// Approximate memory used by a String. -#[inline] -fn str_mem(s: &str) -> usize { - s.len() + std::mem::size_of::() -} - #[derive(Clone, Debug)] pub struct PreparedStatements { pub(super) global: Arc>, @@ -48,95 +53,124 @@ impl Default for PreparedStatements { } } +impl Drop for PreparedStatements { + fn drop(&mut self) { + self.close_all(); + } +} + impl PreparedStatements { - /// New shared prepared statements cache. - pub fn new() -> Self { + /// New shared prepared statements cache instance. + /// + /// Has access to the global cache singleton. + /// + pub(crate) fn new() -> Self { CACHE.clone() } - /// Get global cache. - pub fn global() -> Arc> { + /// Get global prepared statements cache singleton. + pub(crate) fn global() -> Arc> { Self::new().global.clone() } - /// Maybe rewrite message. - pub fn maybe_rewrite(&mut self, message: &mut ProtocolMessage) -> Result<(), Error> { + /// Rewrite extended protocol messages to use global names. This allows multiple + /// clients to re-use the same statement prepared on a Postgres server. + /// + /// # Arguments + /// + /// * `message`: Any protocol message. Unsupported messages are not rewritten. + /// + pub(crate) fn maybe_rewrite(&mut self, message: &mut ProtocolMessage) -> Result<(), Error> { let mut rewrite = Rewrite::new(self); rewrite.rewrite(message)?; Ok(()) } - /// Register prepared statement with the global cache. - pub fn insert(&mut self, parse: &mut Parse) { + /// Register prepared statement with the global cache and rewrite it + /// to use the globally unique name. + /// + /// # Arguments + /// + /// * `parse`: [`Parse`] message. It will be renamed in-place. + /// + pub(super) fn insert(&mut self, parse: &mut Parse) { let (_new, name) = { self.global.write().insert(parse) }; let key = parse.name(); - let existed = self.local.insert(key.to_owned(), name.clone()); - // Client prepared it again because it got an error the first time. - // We can check if this is a new statement first, but this is an error - // condition which happens very infrequently, so we optimize for the happy path. - if let Some(old_value) = existed { - // Key already existed, only value changed. - self.memory_used = self.memory_used.saturating_sub(str_mem(&old_value)); - self.memory_used += str_mem(&name); - self.global.write().close(&old_value); - } else { - // New entry. - self.memory_used += str_mem(key) + str_mem(&name); - } + self.insert_internal(key, &name); parse.rename(&name) } - /// Insert statement into the cache bypassing duplicate checks. - pub fn insert_prepare(&mut self, parse: &mut Parse) { - let name = { self.global.write().insert_prepare(parse) }; - let key = parse.name(); - let existed = self.local.insert(key.to_owned(), name.clone()); + fn insert_internal(&mut self, local: &str, global: &str) { + let existed = self.local.insert(local.to_owned(), global.to_owned()); if let Some(old_value) = existed { // Key already existed, only value changed. self.memory_used = self.memory_used.saturating_sub(str_mem(&old_value)); - self.memory_used += str_mem(&name); + self.memory_used += str_mem(&local); self.global.write().close(&old_value); } else { // New entry. - self.memory_used += str_mem(key) + str_mem(&name); + self.memory_used += str_mem(local) + str_mem(&global); } + } - parse.rename(&name) + /// Insert statement into the cache bypassing duplicate checks. + /// + /// # 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 fn insert_prepare( + &mut self, + name: &str, + query: Bytes, + data_types: Vec, + ) -> String { + let (_new, global_name) = { self.global.write().insert_prepare(query, data_types) }; + + self.insert_internal(name, &global_name); + + global_name } - /// Get global statement counter. + /// Get the global unique name for a prepared statement + /// using the name the client gave us as key. pub fn name(&self, name: &str) -> Option<&String> { self.local.get(name) } - /// Get globally-prepared statement by local name. + /// Get a globally unique [`Parse`] message using the client name as key. pub fn parse(&self, name: &str) -> Option { self.local .get(name) .and_then(|name| self.global.read().parse(name)) } - /// Number of prepared statements in the local cache. - pub fn len_local(&self) -> usize { - self.local.len() - } - - /// Current prepared statements compatibility level. - #[cfg(test)] - pub fn level(&self) -> PreparedStatementsLevel { - self.level + /// Get a globally unique [`Prepare`] message using the client name as key. + pub(crate) fn prepare(&self, name: &str) -> Option { + self.local + .get(name) + .and_then(|name| self.global.read().prepare(name)) } - /// Is the local cache empty? - pub fn is_empty(&self) -> bool { - self.len_local() == 0 + /// Number of prepared statements in the client's cache. + pub(crate) fn num_statements(&self) -> usize { + self.local.len() } - /// Remove prepared statement from local cache. - pub fn close(&mut self, name: &str) { + /// Remove prepared statement from client's cache. + /// + /// # Arguments + /// + /// * `name`: Name of the prepared statement according to the client. + /// + pub(crate) fn close(&mut self, name: &str) { if let Some(global_name) = self.local.remove(name) { self.global.write().close(&global_name); self.memory_used = self @@ -146,7 +180,10 @@ impl PreparedStatements { } /// Close all prepared statements on this client. - pub fn close_all(&mut self) { + /// + /// This only happens when the client disconnects. This will update + /// the global usage counters of all of client's prepared statements. + fn close_all(&mut self) { if !self.local.is_empty() { let mut global = self.global.write(); @@ -159,42 +196,34 @@ impl PreparedStatements { self.memory_used = 0; } - /// How much memory is used, approx. - pub fn memory_used(&self) -> usize { + /// How much memory is used, approximately, by the prepared statements cache + /// for this client. + pub(crate) fn memory_used(&self) -> usize { self.memory_used } /// Set the prepared statements level. - pub fn set_level(&mut self, level: PreparedStatementsLevel) { + pub(crate) fn set_level(&mut self, level: PreparedStatementsLevel) { self.level = level; } } -/// Run prepared statements maintenance task -/// every second. -pub fn start_maintenance() { - crate::tasks::spawn("prepared statements cache", async move { - debug!("prepared statements cache maintenance started"); - let shutdown = crate::tasks::shutdown_signal(); - loop { - tokio::select! { - _ = safe_sleep(Duration::from_secs(1)) => {} - _ = shutdown.cancelled() => break, - } - run_maintenance(); - } - }); -} - -/// Check prepared statements cache for overflows -/// and remove any unused statements exceeding the limit. -pub fn run_maintenance() { - let capacity = config().config.general.prepared_statements_limit; - PreparedStatements::global().write().close_unused(capacity); +/// Approximate memory used by a String. +#[inline] +fn str_mem(s: &str) -> usize { + s.len() + std::mem::size_of::() } #[cfg(test)] mod test { + + impl PreparedStatements { + /// Current prepared statements compatibility level. + pub(crate) fn level(&self) -> PreparedStatementsLevel { + self.level + } + } + use crate::backend::Server; use crate::backend::server::test::{execute_prepared, prepared_in_postgres, test_server}; use crate::net::messages::Bind; diff --git a/pgdog/src/frontend/prepared_statements/prelude.rs b/pgdog/src/frontend/prepared_statements/prelude.rs new file mode 100644 index 000000000..2cd079e5e --- /dev/null +++ b/pgdog/src/frontend/prepared_statements/prelude.rs @@ -0,0 +1,7 @@ +pub(super) use super::CacheKey; +pub(super) use crate::config::config; +pub(super) use crate::net::{Parse, RowDescription}; +pub(super) use crate::util::*; +pub(super) use bytes::Bytes; +pub(super) use std::str::from_utf8; +pub(super) use tracing::*; diff --git a/pgdog/src/frontend/prepared_statements/rewrite.rs b/pgdog/src/frontend/prepared_statements/rewrite.rs index a20492334..eb310ac00 100644 --- a/pgdog/src/frontend/prepared_statements/rewrite.rs +++ b/pgdog/src/frontend/prepared_statements/rewrite.rs @@ -103,7 +103,7 @@ mod test { assert_eq!(describe.statement(), "__pgdog_1"); assert_eq!(describe.kind(), 'S'); - assert_eq!(statements.len_local(), 1); + assert_eq!(statements.num_statements(), 1); assert_eq!(statements.global.read().len(), 1); } @@ -120,7 +120,7 @@ mod test { assert!(!parse.anonymous()); assert_eq!(parse.query(), "SELECT * FROM users"); - assert_eq!(statements.len_local(), 1); + assert_eq!(statements.num_statements(), 1); assert_eq!(statements.global.read().len(), 1); } diff --git a/pgdog/src/frontend/prepared_statements/statement.rs b/pgdog/src/frontend/prepared_statements/statement.rs new file mode 100644 index 000000000..8cfaf1e9f --- /dev/null +++ b/pgdog/src/frontend/prepared_statements/statement.rs @@ -0,0 +1,84 @@ +use crate::{net::Prepare, stats::memory::MemoryUsage}; + +use super::prelude::*; + +#[derive(Debug, Clone)] +pub struct Statement { + pub(super) stmt: StatementType, + pub(super) row_description: Option, + pub(super) cache_key: CacheKey, +} + +#[derive(Debug, Clone)] +pub(crate) enum StatementType { + Parse { + parse: Parse, + rewrite: Option, + }, + + Prepare { + prepare: Prepare, + }, +} + +impl MemoryUsage for StatementType { + fn memory_usage(&self) -> usize { + 0 + } +} + +impl MemoryUsage for Statement { + #[inline] + fn memory_usage(&self) -> usize { + self.stmt.memory_usage() + + if let Some(row_description) = &self.row_description { + row_description.memory_usage() + } else { + 0 + } + + self.cache_key.memory_usage() + } +} + +impl Statement { + pub(crate) fn parse(&self) -> Option { + match self.stmt { + StatementType::Parse { ref parse, .. } => Some(parse.clone()), + _ => None, + } + } + + pub(super) fn prepare(&self) -> Option { + match self.stmt { + StatementType::Prepare { ref prepare } => Some(prepare.clone()), + _ => None, + } + } + + pub(crate) fn rewrite(&self) -> Option { + match self.stmt { + StatementType::Parse { ref rewrite, .. } => rewrite.clone(), + _ => None, + } + } + + pub(super) fn query(&self) -> &str { + match self.stmt { + StatementType::Parse { ref parse, .. } => parse.query(), + StatementType::Prepare { ref prepare } => prepare.query(), + } + } + + pub(super) fn cache_key(&self) -> &CacheKey { + &self.cache_key + } + + pub(super) fn set_rewrite(&mut self, parse: &Parse) { + match self.stmt { + StatementType::Parse { + ref mut rewrite, .. + } => *rewrite = Some(parse.clone()), + _ => (), + } + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 6b2bbb9e8..b2591686a 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -1,6 +1,6 @@ use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::messages::bind::{Format, Parameter}; -use crate::net::{Bind, Parse, ProtocolMessage, Query}; +use crate::net::{Bind, Parse, Prepare, ProtocolMessage, Query}; use crate::unique_id::UniqueId; use super::insert::build_split_requests; @@ -30,7 +30,7 @@ pub struct RewritePlan { /// Prepared statements to prepend to the client request. /// Each tuple contains (name, statement) for ProtocolMessage::Prepare. - pub(crate) prepares: Vec<(String, String)>, + pub(crate) prepares: Vec, /// Splitting of multi-tuple INSERT statements into /// multiple queries. @@ -122,10 +122,7 @@ impl RewritePlan { let prepends: Vec = self .prepares .iter() - .map(|(name, statement)| ProtocolMessage::Prepare { - name: name.clone(), - statement: statement.clone(), - }) + .map(|prepare| ProtocolMessage::Prepare(prepare.clone())) .collect(); request.messages.splice(0..0, prepends); } 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 34a854e7a..9c01c13c9 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -1,7 +1,8 @@ +use bytes::Bytes; use pg_raw_parse::{NodeMut, make::MemoryToken}; +use pgdog_postgres_types::TypeName; -use crate::frontend::PreparedStatements; -use crate::net::Parse; +use crate::{frontend::PreparedStatements, net::Prepare}; use super::{Error, StatementRewrite}; @@ -11,7 +12,7 @@ pub struct SimplePreparedResult { /// Whether any statement was rewritten. pub rewritten: bool, /// Prepared statements to prepend (name, statement) for EXECUTE rewrites. - pub prepares: Vec<(String, String)>, + pub prepares: Vec, } /// Result of rewriting a single PREPARE or EXECUTE SQL command. @@ -23,7 +24,7 @@ enum SimplePreparedRewrite { Prepared, /// EXECUTE statement was rewritten. Contains the global name and statement /// needed to prepend a ProtocolMessage::Prepare. - Executed { name: String, statement: String }, + Executed { prepare: Prepare }, } impl StatementRewrite<'_> { @@ -53,8 +54,8 @@ impl StatementRewrite<'_> { SimplePreparedRewrite::Prepared => { result.rewritten = true; } - SimplePreparedRewrite::Executed { name, statement } => { - result.prepares.push((name, statement)); + SimplePreparedRewrite::Executed { prepare } => { + result.prepares.push(prepare); result.rewritten = true; } SimplePreparedRewrite::None => {} @@ -72,30 +73,36 @@ fn rewrite_single_prepared<'a>( ) -> Result { match node { NodeMut::PrepareStmt(mut stmt) => { - let query = pg_raw_parse::deparse(stmt.query())?; + let query = Bytes::from(pg_raw_parse::deparse(stmt.query())?.as_str().to_owned()); + let data_types: Vec = stmt + .argtypes() + .iter() + .map(|data_type| { + data_type + .names() + .iter() + .filter_map(|name| name.sval().map(str::to_owned)) + .collect() + }) + .collect(); - let mut parse = Parse::named( - stmt.name().expect("PREPARE always has a name"), - query.as_str(), + let global_name = prepared_statements.insert_prepare( + stmt.name().expect("prepare must have a name"), + query, + data_types, ); - prepared_statements.insert_prepare(&mut parse); - stmt.set_name(Some(mem.copy_string(parse.name()))); + + stmt.set_name(Some(mem.copy_string(global_name.as_str()))); Ok(SimplePreparedRewrite::Prepared) } NodeMut::ExecuteStmt(mut stmt) => { let stmt_name = stmt.name().expect("EXECUTE always has name"); - let parse = prepared_statements.parse(stmt_name); - if let Some(parse) = parse { - let global_name = parse.name().to_string(); - let statement = parse.query().to_string(); - stmt.set_name(Some(mem.copy_string(&global_name))); - - Ok(SimplePreparedRewrite::Executed { - name: global_name, - statement, - }) + let prepare = prepared_statements.prepare(stmt_name); + if let Some(prepare) = prepare { + stmt.set_name(Some(mem.copy_string(prepare.name()))); + Ok(SimplePreparedRewrite::Executed { prepare }) } else { Err(Error::ExecuteMissingPrepare(stmt_name.to_owned())) } @@ -188,10 +195,11 @@ mod tests { "EXECUTE should use global name, got: {sql}" ); assert_eq!(plan.prepares.len(), 1); + panic!("fix test"); - let (name, statement) = &plan.prepares[0]; - assert!(name.starts_with("__pgdog_")); - assert_eq!(statement, "SELECT 1"); + // let (name, statement) = &plan.prepares[0]; + // assert!(name.starts_with("__pgdog_")); + // assert_eq!(statement, "SELECT 1"); } #[test] diff --git a/pgdog/src/net/messages/mod.rs b/pgdog/src/net/messages/mod.rs index 733fd3e75..58cafe611 100644 --- a/pgdog/src/net/messages/mod.rs +++ b/pgdog/src/net/messages/mod.rs @@ -31,6 +31,7 @@ pub mod parse; pub mod parse_complete; pub mod payload; pub mod prelude; +pub mod prepare; pub mod protocol_version; pub mod query; pub mod replication; @@ -70,6 +71,7 @@ pub use parameter_status::ParameterStatus; pub use parse::Parse; pub use parse_complete::ParseComplete; pub use payload::Payload; +pub use prepare::Prepare; pub use protocol_version::ProtocolVersion; pub use query::Query; pub use rfq::{ReadyForQuery, TransactionState}; diff --git a/pgdog/src/net/messages/prepare.rs b/pgdog/src/net/messages/prepare.rs new file mode 100644 index 000000000..28a3710e8 --- /dev/null +++ b/pgdog/src/net/messages/prepare.rs @@ -0,0 +1,72 @@ +//! Used for preparing statements sent over with the simple protocol. + +use std::str::from_utf8; + +use pgdog_postgres_types::TypeName; + +use super::prelude::*; +use super::{Message, Query}; + +#[derive(Debug, Clone, PartialEq)] +pub struct Prepare { + pub(crate) name: Bytes, + pub(crate) query: Bytes, + pub(crate) data_types: Vec, +} + +impl Prepare { + pub fn query(&self) -> &str { + from_utf8(&self.query).expect("prepare query to be utf-8") + } + + pub fn name(&self) -> &str { + from_utf8(&self.name).expect("prepare name to be utf-8") + } + + pub fn len(&self) -> usize { + self.name.len() + + self.query.len() + + self.data_types.iter().map(|t| t.0.len()).sum::() + } +} + +impl FromBytes for Prepare { + fn from_bytes(bytes: Bytes) -> Result { + unreachable!("Prepare must be constructed manually") + } +} + +impl Protocol for Prepare { + fn code(&self) -> char { + 'Q' + } + + fn message(&self) -> Result { + Ok(Message::new(self.to_bytes()).frontend()) + } + + fn streaming(&self) -> bool { + false + } +} + +impl ToBytes for Prepare { + fn to_bytes(&self) -> Bytes { + let query = from_utf8(&self.query).expect("prepare query to be utf-8"); + let name = from_utf8(&self.name).expect("prepare name to be utf-8"); + let data_types = if self.data_types.is_empty() { + "".into() + } else { + format!( + "({})", + self.data_types + .iter() + .map(|data_type| data_type.to_string()) + .collect::>() + .join(", ") + ) + }; + + Query::new(format!(r#"PREPARE "{}" {} AS {}"#, name, data_types, query)).to_bytes() + } +} diff --git a/pgdog/src/net/protocol_message.rs b/pgdog/src/net/protocol_message.rs index 955d5d12f..dafeeae5a 100644 --- a/pgdog/src/net/protocol_message.rs +++ b/pgdog/src/net/protocol_message.rs @@ -1,6 +1,8 @@ use bytes::Buf; use std::io::Cursor; +use crate::net::Prepare; + use super::{ Bind, Close, CopyData, CopyDone, CopyFail, Describe, Execute, Fastpath, Flush, FromBytes, Message, Parse, Protocol, Query, Sync, ToBytes, @@ -11,7 +13,7 @@ pub enum ProtocolMessage { Bind(Bind), Parse(Parse), Describe(Describe), - Prepare { name: String, statement: String }, + Prepare(Prepare), Execute(Execute), Close(Close), Query(Query), @@ -59,7 +61,7 @@ impl ProtocolMessage { Self::Bind(bind) => bind.len(), Self::Parse(parse) => parse.len(), Self::Describe(describe) => describe.len(), - Self::Prepare { statement, .. } => statement.len() + 1 + 1 + 4, // NULL + code + len + Self::Prepare(prepare) => prepare.len(), Self::Execute(execute) => execute.len(), Self::Close(close) => close.len(), Self::Query(query) => query.len(), @@ -119,9 +121,7 @@ impl ToBytes for ProtocolMessage { Self::Bind(bind) => bind.to_bytes(), Self::Parse(parse) => parse.to_bytes(), Self::Describe(describe) => describe.to_bytes(), - Self::Prepare { statement, name } => { - Query::new(format!("PREPARE {} AS {}", name, statement)).to_bytes() - } + Self::Prepare(prepare) => prepare.to_bytes(), Self::Execute(execute) => execute.to_bytes(), Self::Close(close) => close.to_bytes(), Self::Query(query) => query.to_bytes(), From 5fa99b654de9c8967c5569b575e85d3c000d0bcf Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sat, 15 Aug 2026 14:23:55 -0700 Subject: [PATCH 02/22] fix invariant --- .../prepared_statements/global_cache.rs | 53 ++++++++++++------- 1 file changed, 34 insertions(+), 19 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 72d8af3a7..662c16496 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -50,24 +50,33 @@ impl GlobalCache { /// If the statement exists, no entry is created /// and the global name is returned instead. pub(crate) fn insert(&mut self, parse: &Parse) -> (bool, String) { - let name = self.next_name(); + let cache_key = CacheKey::Extended { + query: parse.query_ref(), + data_types: parse.data_types_ref(), + }; + if let Some(name) = self.reuse(&cache_key) { + return (false, name); + } + + let name = self.next_name(); + let parse = parse.renamed(&name); let cache_key = CacheKey::Extended { query: parse.query_ref(), data_types: parse.data_types_ref(), }; let statement = Statement { stmt: StatementType::Parse { - parse: parse.clone(), + parse, rewrite: None, }, cache_key: cache_key.clone(), row_description: None, }; - let new = self.insert_internal(&name, cache_key, statement); + self.insert_internal(&name, cache_key, statement); - (new, name) + (true, name) } /// Insert a statement prepared using the simple protocol into the global cache. @@ -76,13 +85,16 @@ impl GlobalCache { query: Bytes, data_types: Vec, ) -> (bool, String) { - let name = self.next_name(); - let cache_key = CacheKey::Simple { query: query.clone(), data_types: data_types.clone(), }; + if let Some(name) = self.reuse(&cache_key) { + return (false, name); + } + + let name = self.next_name(); let statement = Statement { stmt: StatementType::Prepare { prepare: Prepare { @@ -95,8 +107,8 @@ impl GlobalCache { cache_key: cache_key.clone(), }; - let new = self.insert_internal(&name, cache_key, statement); - (new, name) + self.insert_internal(&name, cache_key, statement); + (true, name) } /// Rewrite prepared statement in the global cache. @@ -242,26 +254,29 @@ impl GlobalCache { global_name(self.counter) } - fn insert_internal(&mut self, name: &str, cache_key: CacheKey, statement: Statement) -> bool { + fn reuse(&mut self, cache_key: &CacheKey) -> Option { if let Some(entry) = self.statements.get_mut(&cache_key) { if entry.used == 0 { self.unused.remove(&entry.counter); } entry.used += 1; - false + Some(entry.name()) } else { - self.statements.insert( - cache_key.clone(), - CachedStmt { - counter: self.counter, - used: 1, - }, - ); - self.names.insert(name.to_owned(), statement); - true + None } } + + fn insert_internal(&mut self, name: &str, cache_key: CacheKey, statement: Statement) { + self.statements.insert( + cache_key, + CachedStmt { + counter: self.counter, + used: 1, + }, + ); + self.names.insert(name.to_owned(), statement); + } } #[cfg(test)] From f0df3ca55afeecc47d29a858f7927780708efdc4 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sat, 15 Aug 2026 14:31:49 -0700 Subject: [PATCH 03/22] save --- pgdog/src/backend/server.rs | 156 ++++++++++++++++++------------------ 1 file changed, 78 insertions(+), 78 deletions(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index fd8caef98..95205eba7 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -1326,7 +1326,7 @@ impl Drop for Server { pub mod test { use std::time::SystemTime; - use bytes::{BufMut, BytesMut}; + use bytes::{BufMut, Bytes, BytesMut}; use pgdog_stats::PreparedStatementsConfig; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, @@ -2312,83 +2312,83 @@ pub mod test { assert!(!server.has_more_messages()); } - // #[tokio::test] - // async fn test_manual_prepared() { - // let mut server = test_server().await; - - // let mut prep = PreparedStatements::new(); - // let mut parse = Parse::named("test", "SELECT 1::bigint"); - // prep.insert_prepare(&mut parse); - // assert_eq!(parse.name(), "__pgdog_1"); - - // server - // .send( - // &vec![ProtocolMessage::from(Query::new(format!( - // "PREPARE {} AS {}", - // parse.name(), - // parse.query() - // )))] - // .into(), - // ) - // .await - // .unwrap(); - // for c in ['C', 'Z'] { - // let msg = server.read().await.unwrap(); - // assert_eq!(msg.code(), c); - // } - // assert!(server.sync_prepared()); - // server.sync_prepared_statements().await.unwrap(); - // assert!(server.prepared_statements.contains("__pgdog_1")); - - // let describe = Describe::new_statement("__pgdog_1"); - // let bind = Bind::new_statement("__pgdog_1"); - // let execute = Execute::new(); - // server - // .send( - // &vec![ - // describe.clone().into(), - // bind.into(), - // execute.into(), - // ProtocolMessage::from(Sync), - // ] - // .into(), - // ) - // .await - // .unwrap(); - - // for c in ['t', 'T', '2', 'D', 'C', 'Z'] { - // let msg = server.read().await.unwrap(); - // assert_eq!(c, msg.code()); - // } - - // let parse = Parse::named("__pgdog_1", "SELECT 2::bigint"); - // let describe = describe.clone(); - - // server - // .send(&vec![parse.into(), describe.into(), ProtocolMessage::from(Flush)].into()) - // .await - // .unwrap(); - - // for c in ['1', 't', 'T'] { - // let msg = server.read().await.unwrap(); - // assert_eq!(msg.code(), c); - // } - - // server - // .send(&vec![ProtocolMessage::from(Query::new("EXECUTE __pgdog_1"))].into()) - // .await - // .unwrap(); - // for c in ['T', 'D', 'C', 'Z'] { - // let msg = server.read().await.unwrap(); - // assert_eq!(c, msg.code()); - // if c == 'D' { - // let data_row = DataRow::from_bytes(msg.to_bytes()).unwrap(); - // let result: i64 = data_row.get(0, Format::Text).unwrap(); - // assert_eq!(result, 1); // We prepared SELECT 1, SELECT 2 is ignored. - // } - // } - // assert!(server.done()); - // } + #[tokio::test] + async fn test_manual_prepared() { + crate::logger(); + let mut server = test_server().await; + + let mut prep = PreparedStatements::new(); + let name = "test"; + let query = Bytes::from("SELECT 1::bigint".to_owned()); + let name = prep.insert_prepare(name, query.clone(), vec![]); + assert_eq!(name, "__pgdog_1"); + + server + .send( + &vec![ProtocolMessage::Query(Query::new( + "PREPARE __pgdog_1 AS SELECT 1::bigint", + ))] + .into(), + ) + .await + .unwrap(); + for c in ['C', 'Z'] { + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), c); + } + assert!(server.sync_prepared()); + server.sync_prepared_statements().await.unwrap(); + assert!(server.prepared_statements.contains("__pgdog_1")); + + let describe = Describe::new_statement("__pgdog_1"); + let bind = Bind::new_statement("__pgdog_1"); + let execute = Execute::new(); + server + .send( + &vec![ + describe.clone().into(), + bind.into(), + execute.into(), + ProtocolMessage::from(Sync), + ] + .into(), + ) + .await + .unwrap(); + + for c in ['t', 'T', '2', 'D', 'C', 'Z'] { + let msg = server.read().await.unwrap(); + assert_eq!(c, msg.code()); + } + + let parse = Parse::named("__pgdog_1", "SELECT 2::bigint"); + let describe = describe.clone(); + + server + .send(&vec![parse.into(), describe.into(), ProtocolMessage::from(Flush)].into()) + .await + .unwrap(); + + for c in ['1', 't', 'T'] { + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), c); + } + + server + .send(&vec![ProtocolMessage::from(Query::new("EXECUTE __pgdog_1"))].into()) + .await + .unwrap(); + for c in ['T', 'D', 'C', 'Z'] { + let msg = server.read().await.unwrap(); + assert_eq!(c, msg.code()); + if c == 'D' { + let data_row = DataRow::from_bytes(msg.to_bytes()).unwrap(); + let result: i64 = data_row.get(0, Format::Text).unwrap(); + assert_eq!(result, 1); // We prepared SELECT 1, SELECT 2 is ignored. + } + } + assert!(server.done()); + } #[tokio::test] async fn test_sync_params() { From 53e0237d1053d7fefe79f3700dfb9307513592c2 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sat, 15 Aug 2026 20:05:07 -0700 Subject: [PATCH 04/22] stmt parameter --- pgdog/src/frontend/router/context.rs | 8 +- pgdog/src/frontend/router/parser/limit.rs | 6 +- pgdog/src/frontend/router/parser/mod.rs | 3 + pgdog/src/frontend/router/parser/params.rs | 132 ++++++++++++++++++ pgdog/src/frontend/router/parser/query/mod.rs | 5 + .../frontend/router/parser/query/prepare.rs | 37 +++++ pgdog/src/frontend/router/parser/statement.rs | 9 +- pgdog/src/net/messages/bind.rs | 9 +- 8 files changed, 196 insertions(+), 13 deletions(-) create mode 100644 pgdog/src/frontend/router/parser/params.rs create mode 100644 pgdog/src/frontend/router/parser/query/prepare.rs diff --git a/pgdog/src/frontend/router/context.rs b/pgdog/src/frontend/router/context.rs index df99253ef..ce16f342a 100644 --- a/pgdog/src/frontend/router/context.rs +++ b/pgdog/src/frontend/router/context.rs @@ -4,16 +4,16 @@ use crate::{ frontend::{ BufferedQuery, ClientRequest, client::{Sticky, TransactionType}, - router::Ast, router::sharding::ResolvedLookups, + router::{Ast, parser::StatementParameters}, }, - net::{Bind, Parameters}, + net::Parameters, }; #[derive(Debug)] pub struct RouterContext<'a> { /// Bound parameters to the query. - pub bind: Option<&'a Bind>, + pub bind: Option>, /// Query we're looking it. pub query: Option, /// Cluster configuration. @@ -53,7 +53,7 @@ impl<'a> RouterContext<'a> { sticky: Sticky, ) -> Result { let query = buffer.query()?; - let bind = buffer.parameters()?; + let bind = buffer.parameters()?.map(|bind| bind.into()); let copy_mode = buffer.is_copy(); Ok(Self { diff --git a/pgdog/src/frontend/router/parser/limit.rs b/pgdog/src/frontend/router/parser/limit.rs index e8c43ff62..508366f8e 100644 --- a/pgdog/src/frontend/router/parser/limit.rs +++ b/pgdog/src/frontend/router/parser/limit.rs @@ -3,7 +3,7 @@ use pg_raw_parse::{ nodes::{self, SelectStmt}, }; -use super::Error; +use super::{Error, StatementParameters}; use crate::net::Bind; #[derive(Debug, Clone, Copy, Default, PartialEq)] @@ -15,11 +15,11 @@ pub(crate) struct Limit { #[derive(Debug, Clone)] pub(crate) struct LimitClause<'a> { stmt: &'a SelectStmt, - bind: Option<&'a Bind>, + bind: Option<&'a StatementParameters<'a>>, } impl<'a> LimitClause<'a> { - pub(crate) fn new(stmt: &'a SelectStmt, bind: Option<&'a Bind>) -> Self { + pub(crate) fn new(stmt: &'a SelectStmt, bind: Option<&'a StatementParameters<'a>>) -> Self { Self { stmt, bind } } diff --git a/pgdog/src/frontend/router/parser/mod.rs b/pgdog/src/frontend/router/parser/mod.rs index e890795dc..2a5651401 100644 --- a/pgdog/src/frontend/router/parser/mod.rs +++ b/pgdog/src/frontend/router/parser/mod.rs @@ -19,6 +19,7 @@ pub mod key; mod limit; pub mod multi_tenant; pub mod order_by; +mod params; pub mod query; pub mod rewrite; pub mod route; @@ -45,6 +46,8 @@ use function::Function; pub use key::Key; pub(crate) use limit::{Limit, LimitClause}; pub use order_by::OrderBy; +pub(crate) use params::StatementParameters; +pub(super) use params::*; pub use query::QueryParser; pub use rewrite::{Assignment, AssignmentValue, StatementRewrite, StatementRewriteContext}; pub use route::{Route, Shard, ShardWithPriority, ShardsWithPriority}; diff --git a/pgdog/src/frontend/router/parser/params.rs b/pgdog/src/frontend/router/parser/params.rs new file mode 100644 index 000000000..785513e6d --- /dev/null +++ b/pgdog/src/frontend/router/parser/params.rs @@ -0,0 +1,132 @@ +//! Get parameter value from [`Bind`] or [`ExecuteStmt`]. + +use bytes::Bytes; +use pg_raw_parse::{ConstValue, Node, nodes::ExecuteStmt}; + +use crate::net::{ + Error, + bind::{Bind, Format, Parameter as BindParameter, ParameterWithFormat}, +}; + +#[derive(Debug)] +pub(super) struct ExecuteParams { + params: Vec, + format: Vec, +} + +impl ExecuteParams { + fn new(stmt: &ExecuteStmt) -> Self { + let mut params = vec![]; + + for param in stmt.params() { + let param = match param { + Node::A_Const(a_const) => match a_const.val() { + None => BindParameter::new_null(), + Some(ConstValue::String(text)) => { + let data = Bytes::from(text.to_string()); + BindParameter { + len: data.len() as i32, + data, + } + } + + Some(ConstValue::Integer(int)) => { + let data = Bytes::from(int.to_string()); + BindParameter { + len: data.len() as i32, + data, + } + } + + Some(ConstValue::Float(float)) => { + let data = Bytes::from(float.to_string()); + BindParameter { + len: data.len() as i32, + data, + } + } + + Some(ConstValue::Boolean(bool)) => { + let data = Bytes::from((if bool { "t" } else { "f" }).to_string()); + BindParameter { + len: data.len() as i32, + data, + } + } + + _ => BindParameter::new_null(), + }, + + _ => BindParameter::new_null(), + }; + + params.push(param); + } + + Self { + params, + format: vec![Format::Text], + } + } + + fn parameter(&self, index: usize) -> Option> { + self.params + .get(index) + .map(|parameter| ParameterWithFormat::new(parameter, Format::Text)) + } +} + +#[derive(Debug, Copy, Clone)] +pub(crate) enum StatementParameters<'a> { + Bind(&'a Bind), + Execute(&'a ExecuteParams), +} + +impl<'a> From<&'a Bind> for StatementParameters<'a> { + fn from(value: &'a Bind) -> Self { + Self::Bind(value) + } +} + +impl StatementParameters<'_> { + pub(super) fn parameter(&self, index: usize) -> Result>, Error> { + match self { + Self::Bind(bind) => bind.parameter(index), + Self::Execute(params) => Ok(params.parameter(index)), + } + } + + pub(super) fn params_raw(&self) -> &[BindParameter] { + match self { + Self::Bind(bind) => bind.params_raw(), + Self::Execute(params) => ¶ms.params, + } + } + + pub(super) fn format_codes_raw(&self) -> &[Format] { + match self { + Self::Bind(bind) => bind.format_codes_raw(), + Self::Execute(params) => ¶ms.format, + } + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_exec_parameter() { + let stmt = pg_raw_parse::parse("EXECUTE __pgdog_1 (1, 'test', now())").unwrap(); + let stmt_1 = stmt.stmts().next().unwrap(); + match stmt_1 { + Node::ExecuteStmt(execute) => { + let params = StatementParameters::Execute(&ExecuteParams::new(execute)); + let first = params.parameter(0).unwrap().unwrap().bigint().unwrap(); + assert_eq!(first, 1); + } + + _ => panic!("not an execute stmt"), + } + } +} diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 50538785f..6b890d192 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -25,6 +25,7 @@ mod ddl; mod delete; mod explain; mod plugins; +mod prepare; mod select; mod set; mod set_config; @@ -430,6 +431,10 @@ impl QueryParser { )); } + Node::PrepareStmt(stmt) => self.prepare(stmt), + + Node::ExecuteStmt(stmt) => self.execute(stmt, context), + Node::ExplainStmt(stmt) => self.explain(&statement, stmt, context), Node::DiscardStmt { .. } => { diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs new file mode 100644 index 000000000..894309c1b --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -0,0 +1,37 @@ +use pg_raw_parse::nodes::{ExecuteStmt, PrepareStmt}; + +use super::*; + +impl QueryParser { + // A `PREPARE` statement can be sent to any shard. + // + // The only distinction we make here is between reads and writes: `SELECT` queries are sent + // to a replica, while everything else is prepared on the primary. + // + pub(super) fn prepare( + &self, + stmt: &PrepareStmt, + // context: &mut QueryParserContext<'_>, + ) -> Result { + let query = stmt.query(); + + let route = match query { + Node::SelectStmt(_) => Route::read(ShardWithPriority::new_rr_not_executable( + round_robin::next().into(), + )), + _ => Route::write(ShardWithPriority::new_rr_not_executable( + round_robin::next().into(), + )), + }; + + Ok(Command::Query(route)) + } + + pub(super) fn execute( + &self, + stmt: &ExecuteStmt, + context: &mut QueryParserContext<'_>, + ) -> Result { + todo!() + } +} diff --git a/pgdog/src/frontend/router/parser/statement.rs b/pgdog/src/frontend/router/parser/statement.rs index 957a660e8..7335b4d25 100644 --- a/pgdog/src/frontend/router/parser/statement.rs +++ b/pgdog/src/frontend/router/parser/statement.rs @@ -2,6 +2,7 @@ use std::borrow::Cow; use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use super::StatementParameters; use crate::util::ResultControlFlowExt; use itertools::*; use pg_raw_parse::walk::Recurse; @@ -10,7 +11,7 @@ use std::ops::ControlFlow; fn advisory_locks_from_func_call( func: &nodes::FuncCall, - bind: Option<&Bind>, + bind: Option>, values_columns: Option<&ValuesColumns<'_>>, ) -> Vec { let mut name_parts = func.funcname().into_iter().filter_map(Node::as_str); @@ -136,7 +137,7 @@ fn collect_values_columns(stmt: &nodes::SelectStmt) -> Option> Some(values) } -fn integer_arg(node: Node<'_>, bind: Option<&Bind>) -> Option { +fn integer_arg(node: Node<'_>, bind: Option<&StatementParameters<'_>>) -> Option { match node { Node::A_Const(a) => a.val()?.numeric_value(), Node::TypeCast(c) => integer_arg(c.arg(), bind), @@ -355,7 +356,7 @@ pub struct SchemaLookupContext<'a> { pub struct StatementParser<'a, 'b, 'c> { stmt: pg_raw_parse::Node<'a>, - bind: Option<&'b Bind>, + bind: Option>, schema: &'b ShardingSchema, recorder: Option<&'c mut ExplainRecorder>, /// Optional schema lookup context for INSERT without column list. @@ -373,7 +374,7 @@ pub struct StatementParser<'a, 'b, 'c> { impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { pub(crate) fn new( stmt: Node<'a>, - bind: Option<&'b Bind>, + bind: Option>, schema: &'b ShardingSchema, recorder: Option<&'c mut ExplainRecorder>, ) -> Self { diff --git a/pgdog/src/net/messages/bind.rs b/pgdog/src/net/messages/bind.rs index acd6145c3..fe281b208 100644 --- a/pgdog/src/net/messages/bind.rs +++ b/pgdog/src/net/messages/bind.rs @@ -66,6 +66,11 @@ pub struct ParameterWithFormat<'a> { } impl<'a> ParameterWithFormat<'a> { + /// Create new parameter with format information. + pub(crate) fn new(parameter: &'a Parameter, format: Format) -> Self { + Self { parameter, format } + } + /// Get text representation if it's valid UTF-8. pub fn text(&self) -> Option<&str> { from_utf8(&self.parameter.data).ok() @@ -268,11 +273,11 @@ impl Bind { me } - pub fn params_raw(&self) -> &Vec { + pub fn params_raw(&self) -> &[Parameter] { &self.params } - pub fn format_codes_raw(&self) -> &Vec { + pub fn format_codes_raw(&self) -> &[Format] { &self.codes } From edf0627b0f0bd64b3430b70a87943dd2b4c6cf1e Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sat, 15 Aug 2026 22:23:31 -0700 Subject: [PATCH 05/22] sharding support for prepare --- .../backend/pool/connection/mirror/handler.rs | 3 +- .../prepared_statements/global_cache.rs | 58 ++++++++++--------- pgdog/src/frontend/router/context.rs | 35 +++++------ pgdog/src/frontend/router/parser/cache/ast.rs | 7 +-- .../router/parser/cache/cache_impl.rs | 29 +++++++--- pgdog/src/frontend/router/parser/context.rs | 1 + pgdog/src/frontend/router/parser/ee/mod.rs | 10 +--- pgdog/src/frontend/router/parser/error.rs | 3 + pgdog/src/frontend/router/parser/limit.rs | 5 +- pgdog/src/frontend/router/parser/mod.rs | 3 +- pgdog/src/frontend/router/parser/params.rs | 12 ++-- pgdog/src/frontend/router/parser/query/mod.rs | 14 ++--- .../frontend/router/parser/query/prepare.rs | 36 +++++++++++- .../frontend/router/parser/query/select.rs | 7 ++- .../router/parser/rewrite/statement/offset.rs | 4 +- .../rewrite/statement/simple_prepared.rs | 7 +-- pgdog/src/frontend/router/parser/statement.rs | 18 +++--- pgdog/src/net/messages/prepare.rs | 2 +- 18 files changed, 142 insertions(+), 112 deletions(-) diff --git a/pgdog/src/backend/pool/connection/mirror/handler.rs b/pgdog/src/backend/pool/connection/mirror/handler.rs index da49e5151..42ebbd108 100644 --- a/pgdog/src/backend/pool/connection/mirror/handler.rs +++ b/pgdog/src/backend/pool/connection/mirror/handler.rs @@ -201,7 +201,6 @@ mod tests { use super::*; use crate::backend::pool::ClusterMetrics; use parking_lot::Mutex; - use pgdog_config::QueryParserEngine; use std::sync::Arc; use tokio::sync::mpsc::{Receiver, channel}; @@ -500,7 +499,7 @@ mod tests { fn request_with_ast(query: &str) -> ClientRequest { use crate::frontend::router::Ast; - let ast = Ast::new_record(query, QueryParserEngine::PgQueryProtobuf).unwrap(); + let ast = Ast::new_record(query).unwrap(); ClientRequest { ast: Some(ast), ..Default::default() diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 662c16496..e2924a82a 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -145,7 +145,7 @@ impl GlobalCache { } /// Get the [`Prepare`] message for a globally unique prepare statement name. - pub(super) fn prepare(&self, name: &str) -> Option { + pub(crate) fn prepare(&self, name: &str) -> Option { self.names.get(name).map(|p| p.prepare()).flatten() } @@ -356,33 +356,35 @@ mod test { cache.statements.get(statement.cache_key()).unwrap().used } - // #[test] - // fn test_simple_prepared_is_never_shared() { - // let mut cache = GlobalCache::default(); - // let parse = Parse::named("client_stmt", "SELECT $1"); - - // let (_, first) = cache.insert_prepare(&parse); - // let (_, second) = cache.insert_prepare(&parse); - - // assert_ne!(first, second); - // assert_eq!(cache.len(), 2); - // assert_eq!(used(&cache, &first), 1); - // assert_eq!(used(&cache, &second), 1); - - // // A Parse never re-uses a SQL PREPARE statement. - // let (new, extended) = cache.insert(&parse); - // assert!(new); - // assert_ne!(extended, first); - // assert_ne!(extended, second); - // assert_eq!(cache.len(), 3); - - // // A Parse re-uses another Parse. - // let (new_again, shared) = cache.insert(&parse); - // assert!(!new_again); - // assert_eq!(shared, extended); - // assert_eq!(cache.len(), 3); - // assert_eq!(used(&cache, &extended), 2); - // } + #[test] + fn test_simple_prepared_is_never_shared() { + let mut cache = GlobalCache::default(); + + let query = Bytes::from("SELECT $1"); + let parse = Parse::named("client_stmt", "SELECT $1"); + + let (_, first) = cache.insert_prepare(query.clone(), vec![]); + let (_, second) = cache.insert_prepare(query, vec![]); + + assert_eq!(first, second); + assert_eq!(cache.len(), 1); + assert_eq!(used(&cache, &first), 2); + assert_eq!(used(&cache, &second), 2); + + // A Parse never re-uses a SQL PREPARE statement. + let (new, extended) = cache.insert(&parse); + assert!(new); + assert_ne!(extended, first); + assert_ne!(extended, second); + assert_eq!(cache.len(), 3); + + // A Parse re-uses another Parse. + let (new_again, shared) = cache.insert(&parse); + assert!(!new_again); + assert_eq!(shared, extended); + assert_eq!(cache.len(), 3); + assert_eq!(used(&cache, &extended), 2); + } #[test] fn test_remove_unused() { diff --git a/pgdog/src/frontend/router/context.rs b/pgdog/src/frontend/router/context.rs index ce16f342a..7d32a3520 100644 --- a/pgdog/src/frontend/router/context.rs +++ b/pgdog/src/frontend/router/context.rs @@ -4,44 +4,39 @@ use crate::{ frontend::{ BufferedQuery, ClientRequest, client::{Sticky, TransactionType}, - router::sharding::ResolvedLookups, - router::{Ast, parser::StatementParameters}, + router::{Ast, parser::StatementParameters, sharding::ResolvedLookups}, }, net::Parameters, }; -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct RouterContext<'a> { /// Bound parameters to the query. - pub bind: Option>, + pub(super) bind: Option>, /// Query we're looking it. - pub query: Option, + pub(super) query: Option, /// Cluster configuration. - pub cluster: &'a Cluster, + pub(super) cluster: &'a Cluster, /// Client parameters, e.g. search_path. - pub parameter_hints: ParameterHints<'a>, + pub(super) parameter_hints: ParameterHints<'a>, /// Client inside transaction, - pub transaction: Option, + pub(super) transaction: Option, /// Currently executing COPY statement. - pub copy_mode: bool, + pub(super) copy_mode: bool, /// Do we have an executable buffer? - pub executable: bool, + pub(super) executable: bool, /// Two-pc enabled - pub two_pc: bool, + pub(super) two_pc: bool, /// Sticky omnisharded index. - pub sticky: Sticky, - /// Extended protocol. - pub extended: bool, + pub(super) sticky: Sticky, /// AST. - pub ast: Option, + pub(super) ast: Option, /// Schema. - pub schema: Schema, - /// Original client request. - pub client_request: &'a ClientRequest, + pub(super) schema: Schema, /// Sharding key translations resolved for this statement. Routing /// reads these before the lookup cache, so a second routing pass /// after resolving lookups can't miss. - pub resolved_lookups: ResolvedLookups, + pub(super) resolved_lookups: ResolvedLookups, } impl<'a> RouterContext<'a> { @@ -65,11 +60,9 @@ impl<'a> RouterContext<'a> { executable: buffer.is_executable(), two_pc: cluster.two_pc_enabled(), sticky, - extended: matches!(query, Some(BufferedQuery::Prepared(_))) || bind.is_some(), query, ast: buffer.ast.clone(), schema: cluster.schema(), - client_request: buffer, resolved_lookups: ResolvedLookups::default(), }) } diff --git a/pgdog/src/frontend/router/parser/cache/ast.rs b/pgdog/src/frontend/router/parser/cache/ast.rs index c9675fffc..5c2465ddf 100644 --- a/pgdog/src/frontend/router/parser/cache/ast.rs +++ b/pgdog/src/frontend/router/parser/cache/ast.rs @@ -148,18 +148,15 @@ impl Ast { } /// Record new AST entry, without rewriting or comment-routing. - pub(crate) fn new_record( - query: &str, - query_parser_engine: QueryParserEngine, - ) -> Result { + pub(crate) fn new_record(query: &str) -> Result { let ast = pg_raw_parse::parse(query)?; Ok(Self { cached: true, comment_role: None, comment_shard: None, - query_parser_engine, comment_sharding_key: None, + query_parser_engine: QueryParserEngine::default(), inner: Arc::new(AstInner::new(ast.into_inner())), }) } diff --git a/pgdog/src/frontend/router/parser/cache/cache_impl.rs b/pgdog/src/frontend/router/parser/cache/cache_impl.rs index 8155c105d..3aa215bf8 100644 --- a/pgdog/src/frontend/router/parser/cache/cache_impl.rs +++ b/pgdog/src/frontend/router/parser/cache/cache_impl.rs @@ -1,7 +1,6 @@ use lru::LruCache; use once_cell::sync::Lazy; use pg_raw_parse::normalize::normalize; -use pgdog_config::QueryParserEngine; use std::collections::HashMap; use std::time::Duration; @@ -190,17 +189,31 @@ impl Cache { Ok(entry) } + pub(crate) fn record(&self, query: &str) -> Result { + { + let mut guard = self.inner.lock(); + if let Some(entry) = self.inner.lock().queries.get(query) { + guard.stats.hits += 1; + entry.stats.lock().hits += 1; + return Ok(entry.clone()); + } + } + + let entry = Ast::new_record(&query)?; + + let mut guard = self.inner.lock(); + guard.queries.put(query.into(), entry.clone()); + guard.stats.misses += 1; + + Ok(entry) + } + /// Record a query sent over the simple protocol, while removing parameters. /// /// Used by dry run mode to keep stats on what queries are routed correctly, /// and which are not. /// - pub fn record_normalized( - &self, - query: &str, - route: &Route, - query_parser_engine: QueryParserEngine, - ) -> Result<(), Error> { + pub fn record_normalized(&self, query: &str, route: &Route) -> Result<(), Error> { let normalized = normalize(query)?; { @@ -212,7 +225,7 @@ impl Cache { } } - let entry = Ast::new_record(&normalized, query_parser_engine)?; + let entry = Ast::new_record(&normalized)?; entry.update_stats(route); let mut guard = self.inner.lock(); diff --git a/pgdog/src/frontend/router/parser/context.rs b/pgdog/src/frontend/router/parser/context.rs index e1c0a7193..ef264a8ea 100644 --- a/pgdog/src/frontend/router/parser/context.rs +++ b/pgdog/src/frontend/router/parser/context.rs @@ -18,6 +18,7 @@ use super::Error; /// Contains a lot of info we collect from the router context /// and its inputs. /// +#[derive(Clone)] pub struct QueryParserContext<'a> { /// Cluster is read-only, i.e. has no primary. pub(super) read_only: bool, diff --git a/pgdog/src/frontend/router/parser/ee/mod.rs b/pgdog/src/frontend/router/parser/ee/mod.rs index 07d9092ca..1cf77e439 100644 --- a/pgdog/src/frontend/router/parser/ee/mod.rs +++ b/pgdog/src/frontend/router/parser/ee/mod.rs @@ -2,13 +2,7 @@ use pgdog_config::Role; -use crate::{ - frontend::router::{ - parser::Value, - parser::{Column, Shard}, - }, - net::Bind, -}; +use crate::frontend::router::parser::{Column, Shard, StatementParameters, Value}; #[derive(Debug, Default, Clone)] pub(crate) struct ParserHooks {} @@ -20,7 +14,7 @@ impl ParserHooks { _shard: &Shard, _column: &Column<'_>, _value: &Value, - _bind: &Option<&Bind>, + _params: Option>, ) { } diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs index 24ae38e37..edc297945 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -111,4 +111,7 @@ pub enum Error { #[error("unmapped sharding key was specified")] UnmappedShardKey(String), + + #[error("prepare statement can only be DML")] + PrepareNotDml, } diff --git a/pgdog/src/frontend/router/parser/limit.rs b/pgdog/src/frontend/router/parser/limit.rs index 508366f8e..568ae3f77 100644 --- a/pgdog/src/frontend/router/parser/limit.rs +++ b/pgdog/src/frontend/router/parser/limit.rs @@ -4,7 +4,6 @@ use pg_raw_parse::{ }; use super::{Error, StatementParameters}; -use crate::net::Bind; #[derive(Debug, Clone, Copy, Default, PartialEq)] pub(crate) struct Limit { @@ -15,11 +14,11 @@ pub(crate) struct Limit { #[derive(Debug, Clone)] pub(crate) struct LimitClause<'a> { stmt: &'a SelectStmt, - bind: Option<&'a StatementParameters<'a>>, + bind: Option>, } impl<'a> LimitClause<'a> { - pub(crate) fn new(stmt: &'a SelectStmt, bind: Option<&'a StatementParameters<'a>>) -> Self { + pub(crate) fn new(stmt: &'a SelectStmt, bind: Option>) -> Self { Self { stmt, bind } } diff --git a/pgdog/src/frontend/router/parser/mod.rs b/pgdog/src/frontend/router/parser/mod.rs index 2a5651401..5d3f79468 100644 --- a/pgdog/src/frontend/router/parser/mod.rs +++ b/pgdog/src/frontend/router/parser/mod.rs @@ -46,8 +46,7 @@ use function::Function; pub use key::Key; pub(crate) use limit::{Limit, LimitClause}; pub use order_by::OrderBy; -pub(crate) use params::StatementParameters; -pub(super) use params::*; +pub(crate) use params::*; pub use query::QueryParser; pub use rewrite::{Assignment, AssignmentValue, StatementRewrite, StatementRewriteContext}; pub use route::{Route, Shard, ShardWithPriority, ShardsWithPriority}; diff --git a/pgdog/src/frontend/router/parser/params.rs b/pgdog/src/frontend/router/parser/params.rs index 785513e6d..bc497999a 100644 --- a/pgdog/src/frontend/router/parser/params.rs +++ b/pgdog/src/frontend/router/parser/params.rs @@ -9,13 +9,13 @@ use crate::net::{ }; #[derive(Debug)] -pub(super) struct ExecuteParams { +pub(crate) struct ExecuteParams { params: Vec, format: Vec, } impl ExecuteParams { - fn new(stmt: &ExecuteStmt) -> Self { + pub(crate) fn new(stmt: &ExecuteStmt) -> Self { let mut params = vec![]; for param in stmt.params() { @@ -88,22 +88,22 @@ impl<'a> From<&'a Bind> for StatementParameters<'a> { } } -impl StatementParameters<'_> { - pub(super) fn parameter(&self, index: usize) -> Result>, Error> { +impl<'a> StatementParameters<'a> { + pub(super) fn parameter(self, index: usize) -> Result>, Error> { match self { Self::Bind(bind) => bind.parameter(index), Self::Execute(params) => Ok(params.parameter(index)), } } - pub(super) fn params_raw(&self) -> &[BindParameter] { + pub(super) fn params_raw(self) -> &'a [BindParameter] { match self { Self::Bind(bind) => bind.params_raw(), Self::Execute(params) => ¶ms.params, } } - pub(super) fn format_codes_raw(&self) -> &[Format] { + pub(super) fn format_codes_raw(self) -> &'a [Format] { match self { Self::Bind(bind) => bind.format_codes_raw(), Self::Execute(params) => ¶ms.format, diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 6b890d192..fca8ffb96 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -10,13 +10,13 @@ use crate::{ round_robin, sharding::{self, Centroids, ContextBuilder, ShardOrLookup}, }, - net::{ - messages::{Bind, Vector}, - parameter::ParameterValue, - }, + net::{messages::Vector, parameter::ParameterValue}, plugin::plugins, }; +#[cfg(test)] +use crate::net::messages::Bind; + use super::{ explain_trace::{ExplainRecorder, ExplainSummary}, *, @@ -524,11 +524,7 @@ impl QueryParser { // Record statement in cache with normalized parameters. if !statement.cached { let query = context.query()?.query(); - Cache::get().record_normalized( - query, - command.route(), - context.sharding_schema.query_parser_engine, - )?; + Cache::get().record_normalized(query, command.route())?; } Ok(command.dry_run()) } else { diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs index 894309c1b..f734e8401 100644 --- a/pgdog/src/frontend/router/parser/query/prepare.rs +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -1,5 +1,10 @@ use pg_raw_parse::nodes::{ExecuteStmt, PrepareStmt}; +use crate::{ + frontend::{BufferedQuery, PreparedStatements}, + net::Query, +}; + use super::*; impl QueryParser { @@ -27,11 +32,38 @@ impl QueryParser { Ok(Command::Query(route)) } + /// An `EXECUTE` statement. pub(super) fn execute( - &self, + &mut self, stmt: &ExecuteStmt, context: &mut QueryParserContext<'_>, ) -> Result { - todo!() + // Extract parameters from `EXECUTE` statement. + let params = ExecuteParams::new(stmt); + let stmt_params = StatementParameters::Execute(¶ms); + + // Create new parser context. + let mut context = context.clone(); + context.router_context.bind = Some(stmt_params); + + // Get the original query from the prepared + // statements cache. + // + // INVARIANT: The prepared statements rewriter places it there. + let stmt = PreparedStatements::global() + .read() + .prepare(stmt.name().expect("execute to have a name")) + .expect("rewriter did not record PREPARE statement"); + + let query = BufferedQuery::Query(Query::new(stmt.query())); + let ast = Cache::get().record(&query)?; + + match ast.ast.stmts().next() { + Some(Node::SelectStmt(stmt)) => self.select(&ast, stmt.into(), &mut context), + Some(Node::InsertStmt(stmt)) => self.insert(stmt.into(), &mut context), + Some(Node::UpdateStmt(stmt)) => self.update(stmt.into(), &mut context), + Some(Node::DeleteStmt(stmt)) => self.delete(stmt.into(), &mut context), + _ => Err(Error::PrepareNotDml), + } } } diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs index 4d54d744e..af301a941 100644 --- a/pgdog/src/frontend/router/parser/query/select.rs +++ b/pgdog/src/frontend/router/parser/query/select.rs @@ -258,9 +258,12 @@ impl QueryParser { /// # Arguments /// /// * `nodes`: List of parser-generated nodes from the ORDER BY clause. - /// * `params`: Bind parameters, if any. + /// * `params`: Statement parameters, if any. /// - fn select_sort(stmt: &nodes::SelectStmt, params: Option<&Bind>) -> Vec { + fn select_sort( + stmt: &nodes::SelectStmt, + params: Option>, + ) -> Vec { stmt.sort_clause() .into_iter() .filter_map(|sort_by| { diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs index 2f9440043..3bd665e43 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs @@ -192,7 +192,7 @@ mod tests { use crate::net::Parse; use crate::net::messages::Query; use crate::net::messages::bind::{Bind, Parameter}; - use pgdog_config::{QueryParserEngine, Rewrite}; + use pgdog_config::Rewrite; fn sharded_schema() -> ShardingSchema { ShardingSchema { @@ -233,7 +233,7 @@ mod tests { } fn make_ast(sql: &str) -> Ast { - Ast::new_record(sql, QueryParserEngine::PgQueryProtobuf).unwrap() + Ast::new_record(sql).unwrap() } fn run_limit_offset(sql: &str, schema: &ShardingSchema) -> RewritePlan { 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 9c01c13c9..efe8d7fc9 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -195,11 +195,10 @@ mod tests { "EXECUTE should use global name, got: {sql}" ); assert_eq!(plan.prepares.len(), 1); - panic!("fix test"); - // let (name, statement) = &plan.prepares[0]; - // assert!(name.starts_with("__pgdog_")); - // assert_eq!(statement, "SELECT 1"); + let prepare = &plan.prepares[0]; + assert!(prepare.name().starts_with("__pgdog_")); + assert_eq!(prepare.query(), "SELECT 1"); } #[test] diff --git a/pgdog/src/frontend/router/parser/statement.rs b/pgdog/src/frontend/router/parser/statement.rs index 7335b4d25..5c58c9842 100644 --- a/pgdog/src/frontend/router/parser/statement.rs +++ b/pgdog/src/frontend/router/parser/statement.rs @@ -137,7 +137,7 @@ fn collect_values_columns(stmt: &nodes::SelectStmt) -> Option> Some(values) } -fn integer_arg(node: Node<'_>, bind: Option<&StatementParameters<'_>>) -> Option { +fn integer_arg(node: Node<'_>, bind: Option>) -> Option { match node { Node::A_Const(a) => a.val()?.numeric_value(), Node::TypeCast(c) => integer_arg(c.arg(), bind), @@ -233,7 +233,7 @@ use crate::{ ShardedTable, Tables, lookup, }, }, - net::{Bind, messages::Format, parameter::ParameterValue}, + net::{messages::Format, parameter::ParameterValue}, }; use pgdog_config::LookupResult; @@ -440,7 +440,7 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { /// Record a sharding key match. fn record_sharding_key(&mut self, shard: &Shard, column: Column<'_>, value: &Value<'_>) { self.hooks - .record_sharding_key(shard, &column, value, &self.bind); + .record_sharding_key(shard, &column, value, self.bind); if let Some(recorder) = self.recorder.as_mut() { let col_str = if let Some(table) = column.table { @@ -1169,7 +1169,7 @@ mod test { let schema = test_schema(); let raw = pg_raw_parse::parse(stmt).unwrap(); let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::new(stmt, bind, &schema, None); + let mut parser = StatementParser::new(stmt, bind.map(Into::into), &schema, None); parser.shard() } @@ -1182,7 +1182,7 @@ mod test { ) -> (Option, Vec) { let raw = pg_raw_parse::parse(stmt).unwrap(); let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::new(stmt, bind, schema, None); + let mut parser = StatementParser::new(stmt, bind.map(Into::into), schema, None); let shard = parser.shard().unwrap(); (shard, parser.take_pending_lookups()) } @@ -2149,7 +2149,7 @@ mod test { }; let raw = pg_raw_parse::parse(stmt).unwrap(); let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::new(stmt, bind, &schema, None); + let mut parser = StatementParser::new(stmt, bind.map(Into::into), &schema, None); parser.shard() } @@ -2261,7 +2261,7 @@ mod test { }; let raw = pg_raw_parse::parse(stmt).unwrap(); let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::new(stmt, bind, &schema, None); + let mut parser = StatementParser::new(stmt, bind.map(Into::into), &schema, None); parser.shard() } @@ -2417,7 +2417,7 @@ mod test { }; let raw = pg_raw_parse::parse(stmt).unwrap(); let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::new(stmt, bind, &sharding_schema, None) + let mut parser = StatementParser::new(stmt, bind.map(Into::into), &sharding_schema, None) .with_schema_lookup(schema_lookup); parser.shard() } @@ -2620,7 +2620,7 @@ mod test { let schema = ShardingSchema::default(); let raw = pg_raw_parse::parse(query).unwrap(); let stmt = raw.stmts().next().unwrap(); - let mut parser = StatementParser::new(stmt, bind, &schema, None); + let mut parser = StatementParser::new(stmt, bind.map(Into::into), &schema, None); let mut v: Vec<_> = parser.extract_advisory_locks().iter().copied().collect(); v.sort_by_key(|l| (l.id, l.unlock)); v diff --git a/pgdog/src/net/messages/prepare.rs b/pgdog/src/net/messages/prepare.rs index 28a3710e8..8a17e283c 100644 --- a/pgdog/src/net/messages/prepare.rs +++ b/pgdog/src/net/messages/prepare.rs @@ -31,7 +31,7 @@ impl Prepare { } impl FromBytes for Prepare { - fn from_bytes(bytes: Bytes) -> Result { + fn from_bytes(_bytes: Bytes) -> Result { unreachable!("Prepare must be constructed manually") } } From bf6675fe5bb6ed29c174a45a6b7ae4bfa05774eb Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sat, 15 Aug 2026 23:24:39 -0700 Subject: [PATCH 06/22] todo --- pgdog/src/admin/show_prepared_statements.rs | 10 +++++++--- .../frontend/prepared_statements/global_cache.rs | 15 +++++++-------- pgdog/src/frontend/prepared_statements/mod.rs | 4 ++-- .../src/frontend/prepared_statements/statement.rs | 10 +++++----- .../frontend/router/parser/cache/cache_impl.rs | 5 ++--- pgdog/src/frontend/router/parser/error.rs | 3 +++ pgdog/src/frontend/router/parser/query/prepare.rs | 7 ++++--- .../parser/rewrite/statement/simple_prepared.rs | 2 ++ 8 files changed, 32 insertions(+), 24 deletions(-) diff --git a/pgdog/src/admin/show_prepared_statements.rs b/pgdog/src/admin/show_prepared_statements.rs index 2ef7c0404..1a035e4aa 100644 --- a/pgdog/src/admin/show_prepared_statements.rs +++ b/pgdog/src/admin/show_prepared_statements.rs @@ -33,7 +33,7 @@ impl Command for ShowPreparedStatements { ]; for (key, stmt) in statements.statements() { let name = stmt.name(); - let rewrite = statements.rewritten_parse(&name).ok_or(Error::Empty)?; + let rewrite = statements.rewritten_parse(&name); let rewritten = statements.is_rewritten(&name); let name_memory = statements .names() @@ -43,8 +43,12 @@ impl Command for ShowPreparedStatements { let mut dr = DataRow::new(); dr.add(stmt.name()) .add(key.query()?) - .add(if rewritten { - rewrite.query().to_data_row_column() + .add(if let Some(rewrite) = rewrite { + if rewritten { + rewrite.query().to_data_row_column() + } else { + Data::null() + } } else { Data::null() }) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index e2924a82a..36c730301 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -99,7 +99,7 @@ impl GlobalCache { stmt: StatementType::Prepare { prepare: Prepare { name: Bytes::from(name.clone()), - query: query, + query, data_types, }, }, @@ -141,12 +141,12 @@ impl GlobalCache { /// It can be used to prepare this statement on a server connection /// or to inspect the original query. pub fn parse(&self, name: &str) -> Option { - self.names.get(name).map(|p| p.parse().clone()).flatten() + self.names.get(name).and_then(|p| p.parse().clone()) } /// Get the [`Prepare`] message for a globally unique prepare statement name. pub(crate) fn prepare(&self, name: &str) -> Option { - self.names.get(name).map(|p| p.prepare()).flatten() + self.names.get(name).and_then(|p| p.prepare()) } /// Get the rewritten Parse statement. @@ -156,8 +156,7 @@ impl GlobalCache { pub fn rewritten_parse(&self, name: &str) -> Option { self.names .get(name) - .map(|p| p.rewrite().clone().or(p.parse())) - .flatten() + .and_then(|p| p.rewrite().clone().or(p.parse())) } /// Returns true if this prepared statement has been @@ -255,7 +254,7 @@ impl GlobalCache { } fn reuse(&mut self, cache_key: &CacheKey) -> Option { - if let Some(entry) = self.statements.get_mut(&cache_key) { + if let Some(entry) = self.statements.get_mut(cache_key) { if entry.used == 0 { self.unused.remove(&entry.counter); } @@ -376,13 +375,13 @@ mod test { assert!(new); assert_ne!(extended, first); assert_ne!(extended, second); - assert_eq!(cache.len(), 3); + assert_eq!(cache.len(), 2); // A Parse re-uses another Parse. let (new_again, shared) = cache.insert(&parse); assert!(!new_again); assert_eq!(shared, extended); - assert_eq!(cache.len(), 3); + assert_eq!(cache.len(), 2); assert_eq!(used(&cache, &extended), 2); } diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index 6f4714996..156a9bd81 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -108,11 +108,11 @@ impl PreparedStatements { if let Some(old_value) = existed { // Key already existed, only value changed. self.memory_used = self.memory_used.saturating_sub(str_mem(&old_value)); - self.memory_used += str_mem(&local); + self.memory_used += str_mem(local); self.global.write().close(&old_value); } else { // New entry. - self.memory_used += str_mem(local) + str_mem(&global); + self.memory_used += str_mem(local) + str_mem(global); } } diff --git a/pgdog/src/frontend/prepared_statements/statement.rs b/pgdog/src/frontend/prepared_statements/statement.rs index 8cfaf1e9f..3630a054a 100644 --- a/pgdog/src/frontend/prepared_statements/statement.rs +++ b/pgdog/src/frontend/prepared_statements/statement.rs @@ -74,11 +74,11 @@ impl Statement { } pub(super) fn set_rewrite(&mut self, parse: &Parse) { - match self.stmt { - StatementType::Parse { - ref mut rewrite, .. - } => *rewrite = Some(parse.clone()), - _ => (), + if let StatementType::Parse { + ref mut rewrite, .. + } = self.stmt + { + *rewrite = Some(parse.clone()) } } } diff --git a/pgdog/src/frontend/router/parser/cache/cache_impl.rs b/pgdog/src/frontend/router/parser/cache/cache_impl.rs index 3aa215bf8..abe5eda6f 100644 --- a/pgdog/src/frontend/router/parser/cache/cache_impl.rs +++ b/pgdog/src/frontend/router/parser/cache/cache_impl.rs @@ -192,14 +192,13 @@ impl Cache { pub(crate) fn record(&self, query: &str) -> Result { { let mut guard = self.inner.lock(); - if let Some(entry) = self.inner.lock().queries.get(query) { - guard.stats.hits += 1; + if let Some(entry) = guard.queries.get_mut(query) { entry.stats.lock().hits += 1; return Ok(entry.clone()); } } - let entry = Ast::new_record(&query)?; + let entry = Ast::new_record(query)?; let mut guard = self.inner.lock(); guard.queries.put(query.into(), entry.clone()); diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs index edc297945..8fd361a71 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -114,4 +114,7 @@ pub enum Error { #[error("prepare statement can only be DML")] PrepareNotDml, + + #[error("execute requires prepared statements to be set to full")] + ExecuteRequiresFull, } diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs index f734e8401..fe39081ac 100644 --- a/pgdog/src/frontend/router/parser/query/prepare.rs +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -49,17 +49,18 @@ impl QueryParser { // Get the original query from the prepared // statements cache. // - // INVARIANT: The prepared statements rewriter places it there. + // INVARIANT 1: The prepared statements rewriter places it there. + // INVARIANT 2: The rewriter renamed the EXECUTE statement to its global name. let stmt = PreparedStatements::global() .read() .prepare(stmt.name().expect("execute to have a name")) - .expect("rewriter did not record PREPARE statement"); + .ok_or(Error::ExecuteRequiresFull)?; let query = BufferedQuery::Query(Query::new(stmt.query())); let ast = Cache::get().record(&query)?; match ast.ast.stmts().next() { - Some(Node::SelectStmt(stmt)) => self.select(&ast, stmt.into(), &mut context), + Some(Node::SelectStmt(stmt)) => self.select(&ast, stmt, &mut context), Some(Node::InsertStmt(stmt)) => self.insert(stmt.into(), &mut context), Some(Node::UpdateStmt(stmt)) => self.update(stmt.into(), &mut context), Some(Node::DeleteStmt(stmt)) => self.delete(stmt.into(), &mut context), 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 efe8d7fc9..cca6b6437 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -74,10 +74,12 @@ fn rewrite_single_prepared<'a>( match node { NodeMut::PrepareStmt(mut stmt) => { let query = Bytes::from(pg_raw_parse::deparse(stmt.query())?.as_str().to_owned()); + let data_types: Vec = stmt .argtypes() .iter() .map(|data_type| { + // data_type.typmods().iter().map(|tmod| tmod.) data_type .names() .iter() From fc209d706bfa08ad5b6d34dc0c54e34b5b5c1d3e Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 08:59:34 -0700 Subject: [PATCH 07/22] support for prepared --- pgdog/src/backend/prepared_statements.rs | 12 ++++ .../router/parser/rewrite/statement/mod.rs | 4 +- .../router/parser/rewrite/statement/plan.rs | 29 ++++++---- .../rewrite/statement/simple_prepared.rs | 56 ++++++++++++------- pgdog/src/net/protocol_message.rs | 5 +- 5 files changed, 74 insertions(+), 32 deletions(-) diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 16369856f..16f1090f0 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -302,6 +302,18 @@ impl PreparedStatements { self.state.add('3'); } } + ProtocolMessage::PrepareClient(prepare) => { + use crate::net::{CommandComplete, ReadyForQuery}; + if self.contains(prepare.name()) { + self.state + .add_simulated(CommandComplete::from_str("PREPARE").message()?); + self.state + .add_simulated(ReadyForQuery::in_transaction(false).message()?); + return Ok(HandleResult::Drop); + } else { + self.state.add(ExecutionCode::ReadyForQuery); + } + } ProtocolMessage::Prepare(prepare) => { if self.contains(prepare.name()) { return Ok(HandleResult::Drop); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index 8ef1105d2..b91228b5f 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -21,7 +21,7 @@ pub mod update; pub use error::Error; pub use insert::InsertSplit; pub(crate) use plan::RewritePlan; -pub use simple_prepared::SimplePreparedResult; +pub(crate) use simple_prepared::PrepareExecute; pub(crate) use update::*; /// Statement rewrite engine context. @@ -122,7 +122,7 @@ impl<'a> StatementRewrite<'a> { let prepared_result = self.rewrite_simple_prepared(stmt.stmt_mut(), mem)?; if prepared_result.rewritten { self.rewritten = true; - plan.prepares = prepared_result.prepares; + plan.prepare_rewrites = prepared_result.rewrites; } // Inject pgdog.unique_id() for missing BIGINT primary keys. diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index b2591686a..2d9c05551 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -1,11 +1,13 @@ use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::messages::bind::{Format, Parameter}; -use crate::net::{Bind, Parse, Prepare, ProtocolMessage, Query}; +use crate::net::{Bind, Parse, ProtocolMessage, Query}; use crate::unique_id::UniqueId; use super::insert::build_split_requests; use super::offset::OffsetPlan; -use super::{Error, InsertSplit, ShardingKeyUpdate, aggregate::AggregateRewritePlan}; +use super::{ + Error, InsertSplit, PrepareExecute, ShardingKeyUpdate, aggregate::AggregateRewritePlan, +}; /// Statement rewrite plan. /// @@ -30,7 +32,7 @@ pub struct RewritePlan { /// Prepared statements to prepend to the client request. /// Each tuple contains (name, statement) for ProtocolMessage::Prepare. - pub(crate) prepares: Vec, + pub(crate) prepare_rewrites: Vec, /// Splitting of multi-tuple INSERT statements into /// multiple queries. @@ -74,7 +76,7 @@ impl RewritePlan { self.unique_ids == 0 && self.auto_id_injected == 0 && self.stmt.is_none() - && self.prepares.is_empty() + && self.prepare_rewrites.is_empty() && self.insert_split.is_empty() && self.aggregates.is_noop() && self.sharding_key_update.is_none() @@ -118,13 +120,20 @@ impl RewritePlan { /// Apply the rewrite plan to a ClientRequest. pub(crate) fn apply(&self, request: &mut ClientRequest) -> Result { // Prepend any required Prepare messages for EXECUTE statements. - if !self.prepares.is_empty() { - let prepends: Vec = self - .prepares + if !self.prepare_rewrites.is_empty() { + self.prepare_rewrites .iter() - .map(|prepare| ProtocolMessage::Prepare(prepare.clone())) - .collect(); - request.messages.splice(0..0, prepends); + .for_each(|prepare| match prepare { + PrepareExecute::Prepare(prepare) => { + request.messages.clear(); + request.push(ProtocolMessage::PrepareClient(prepare.clone())); + } + PrepareExecute::Execute(prepare) => { + request + .messages + .splice(0..0, vec![ProtocolMessage::Prepare(prepare.clone())]); + } + }); } for message in request.messages.iter_mut() { 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 cca6b6437..142010aab 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -6,13 +6,22 @@ use crate::{frontend::PreparedStatements, net::Prepare}; use super::{Error, StatementRewrite}; +#[derive(Debug, Clone)] +pub(crate) enum PrepareExecute { + /// PREPARE statement sent by client + Prepare(Prepare), + /// EXECUTE statement sent by client and may require + /// a PREPARE first. + Execute(Prepare), +} + /// Result of rewriting all PREPARE/EXECUTE statements in a query. #[derive(Debug, Clone, Default)] -pub struct SimplePreparedResult { +pub(crate) struct SimplePreparedResult { /// Whether any statement was rewritten. - pub rewritten: bool, + pub(crate) rewritten: bool, /// Prepared statements to prepend (name, statement) for EXECUTE rewrites. - pub prepares: Vec, + pub(crate) rewrites: Vec, } /// Result of rewriting a single PREPARE or EXECUTE SQL command. @@ -21,7 +30,7 @@ enum SimplePreparedRewrite { /// Node was not a PREPARE or EXECUTE statement. None, /// PREPARE statement was rewritten. - Prepared, + Prepared { prepare: Prepare }, /// EXECUTE statement was rewritten. Contains the global name and statement /// needed to prepend a ProtocolMessage::Prepare. Executed { prepare: Prepare }, @@ -51,11 +60,12 @@ impl StatementRewrite<'_> { } match rewrite_single_prepared(node, mem, self.prepared_statements)? { - SimplePreparedRewrite::Prepared => { + SimplePreparedRewrite::Prepared { prepare } => { + result.rewrites.push(PrepareExecute::Prepare(prepare)); result.rewritten = true; } SimplePreparedRewrite::Executed { prepare } => { - result.prepares.push(prepare); + result.rewrites.push(PrepareExecute::Execute(prepare)); result.rewritten = true; } SimplePreparedRewrite::None => {} @@ -88,15 +98,17 @@ fn rewrite_single_prepared<'a>( }) .collect(); - let global_name = prepared_statements.insert_prepare( - stmt.name().expect("prepare must have a name"), - query, - data_types, - ); + let name = stmt.name().expect("prepare must have a name").to_owned(); + + let global_name = prepared_statements.insert_prepare(&name, query, data_types); stmt.set_name(Some(mem.copy_string(global_name.as_str()))); - Ok(SimplePreparedRewrite::Prepared) + Ok(SimplePreparedRewrite::Prepared { + prepare: prepared_statements + .prepare(&name) + .expect("global cache missing prepare we just inserted"), + }) } NodeMut::ExecuteStmt(mut stmt) => { @@ -182,7 +194,7 @@ mod tests { !sql.contains("test_stmt"), "original name should be replaced: {sql}" ); - assert!(plan.prepares.is_empty()); + assert!(plan.prepare_rewrites.is_empty()); assert!(plan.stmt.is_some()); } @@ -196,11 +208,17 @@ mod tests { sql.contains("__pgdog_"), "EXECUTE should use global name, got: {sql}" ); - assert_eq!(plan.prepares.len(), 1); + assert_eq!(plan.prepare_rewrites.len(), 1); - let prepare = &plan.prepares[0]; - assert!(prepare.name().starts_with("__pgdog_")); - assert_eq!(prepare.query(), "SELECT 1"); + let prepare = &plan.prepare_rewrites[0]; + match prepare { + PrepareExecute::Execute(prepare) => { + assert!(prepare.name().starts_with("__pgdog_")); + assert_eq!(prepare.query(), "SELECT 1"); + } + + _ => panic!("expected PrepareExecute::Execute"), + } } #[test] @@ -217,7 +235,7 @@ mod tests { sql.contains("(1, 'hello')"), "EXECUTE params should be preserved, got: {sql}" ); - assert_eq!(plan.prepares.len(), 1); + assert_eq!(plan.prepare_rewrites.len(), 1); } #[test] @@ -233,7 +251,7 @@ mod tests { let (sql, plan) = ctx.rewrite("SELECT 1, 2, 3").unwrap(); assert_eq!(sql, "SELECT 1, 2, 3"); - assert!(plan.prepares.is_empty()); + assert!(plan.prepare_rewrites.is_empty()); assert!(plan.stmt.is_none()); } } diff --git a/pgdog/src/net/protocol_message.rs b/pgdog/src/net/protocol_message.rs index dafeeae5a..deeceda86 100644 --- a/pgdog/src/net/protocol_message.rs +++ b/pgdog/src/net/protocol_message.rs @@ -14,6 +14,7 @@ pub enum ProtocolMessage { Parse(Parse), Describe(Describe), Prepare(Prepare), + PrepareClient(Prepare), Execute(Execute), Close(Close), Query(Query), @@ -62,6 +63,7 @@ impl ProtocolMessage { Self::Parse(parse) => parse.len(), Self::Describe(describe) => describe.len(), Self::Prepare(prepare) => prepare.len(), + Self::PrepareClient(prepare) => prepare.len(), Self::Execute(execute) => execute.len(), Self::Close(close) => close.len(), Self::Query(query) => query.len(), @@ -81,7 +83,7 @@ impl Protocol for ProtocolMessage { Self::Bind(bind) => bind.code(), Self::Parse(parse) => parse.code(), Self::Describe(describe) => describe.code(), - Self::Prepare { .. } => 'Q', + Self::Prepare { .. } | Self::PrepareClient { .. } => 'Q', Self::Execute(execute) => execute.code(), Self::Close(close) => close.code(), Self::Query(query) => query.code(), @@ -122,6 +124,7 @@ impl ToBytes for ProtocolMessage { Self::Parse(parse) => parse.to_bytes(), Self::Describe(describe) => describe.to_bytes(), Self::Prepare(prepare) => prepare.to_bytes(), + Self::PrepareClient(prepare) => prepare.to_bytes(), Self::Execute(execute) => execute.to_bytes(), Self::Close(close) => close.to_bytes(), Self::Query(query) => query.to_bytes(), From 5cdec64932ddb9311e48add4d4aa2bbf35f8dc25 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 09:42:59 -0700 Subject: [PATCH 08/22] use correct key --- pgdog/src/backend/server.rs | 4 +- .../frontend/prepared_statements/cache_key.rs | 12 +----- .../prepared_statements/global_cache.rs | 40 +++++++++---------- pgdog/src/frontend/prepared_statements/mod.rs | 14 ++----- .../frontend/router/parser/query/prepare.rs | 4 +- .../rewrite/statement/simple_prepared.rs | 36 ++++++----------- pgdog/src/net/messages/mod.rs | 2 +- pgdog/src/net/messages/prepare.rs | 40 ++++++++----------- 8 files changed, 61 insertions(+), 91 deletions(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 95205eba7..17b57a6cf 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -2320,8 +2320,8 @@ pub mod test { let mut prep = PreparedStatements::new(); let name = "test"; let query = Bytes::from("SELECT 1::bigint".to_owned()); - let name = prep.insert_prepare(name, query.clone(), vec![]); - assert_eq!(name, "__pgdog_1"); + let prepare = prep.insert_prepare(name, query.clone()); + assert_eq!(prepare.name(), "__pgdog_1"); server .send( diff --git a/pgdog/src/frontend/prepared_statements/cache_key.rs b/pgdog/src/frontend/prepared_statements/cache_key.rs index 2d2346ac1..b498bafa8 100644 --- a/pgdog/src/frontend/prepared_statements/cache_key.rs +++ b/pgdog/src/frontend/prepared_statements/cache_key.rs @@ -1,5 +1,3 @@ -use pgdog_postgres_types::TypeName; - use crate::stats::memory::MemoryUsage; use super::prelude::*; @@ -17,14 +15,8 @@ use super::prelude::*; /// #[derive(Debug, Clone, PartialEq, Hash, Eq)] pub enum CacheKey { - Extended { - query: Bytes, - data_types: Bytes, - }, - Simple { - query: Bytes, - data_types: Vec, - }, + Extended { query: Bytes, data_types: Bytes }, + Simple { query: Bytes }, } impl MemoryUsage for CacheKey { diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 36c730301..04af90381 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -80,35 +80,35 @@ impl GlobalCache { } /// Insert a statement prepared using the simple protocol into the global cache. - pub(super) fn insert_prepare( - &mut self, - query: Bytes, - data_types: Vec, - ) -> (bool, String) { + pub(super) fn insert_prepare(&mut self, query: Bytes) -> (bool, Prepare) { let cache_key = CacheKey::Simple { query: query.clone(), - data_types: data_types.clone(), }; if let Some(name) = self.reuse(&cache_key) { - return (false, name); + return ( + false, + self.prepare(&name) + .expect("prepared to be in cache if reuse is true"), + ); } let name = self.next_name(); + let prepare = Prepare { + name: Bytes::from(name.clone()), + query, + }; + let statement = Statement { stmt: StatementType::Prepare { - prepare: Prepare { - name: Bytes::from(name.clone()), - query, - data_types, - }, + prepare: prepare.clone(), }, row_description: None, cache_key: cache_key.clone(), }; self.insert_internal(&name, cache_key, statement); - (true, name) + (true, prepare) } /// Rewrite prepared statement in the global cache. @@ -359,22 +359,22 @@ mod test { fn test_simple_prepared_is_never_shared() { let mut cache = GlobalCache::default(); - let query = Bytes::from("SELECT $1"); + 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(), vec![]); - let (_, second) = cache.insert_prepare(query, vec![]); + let (_, first) = cache.insert_prepare(query.clone()); + let (_, second) = cache.insert_prepare(query); assert_eq!(first, second); assert_eq!(cache.len(), 1); - assert_eq!(used(&cache, &first), 2); - assert_eq!(used(&cache, &second), 2); + assert_eq!(used(&cache, &first.name()), 2); + assert_eq!(used(&cache, &second.name()), 2); // A Parse never re-uses a SQL PREPARE statement. let (new, extended) = cache.insert(&parse); assert!(new); - assert_ne!(extended, first); - assert_ne!(extended, second); + assert_ne!(extended, first.name()); + assert_ne!(extended, second.name()); assert_eq!(cache.len(), 2); // A Parse re-uses another Parse. diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index 156a9bd81..1d350ac25 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -5,7 +5,6 @@ use std::{collections::HashMap, sync::Arc}; use bytes::Bytes; use once_cell::sync::Lazy; use parking_lot::RwLock; -use pgdog_postgres_types::TypeName; use crate::{ config::PreparedStatements as PreparedStatementsLevel, @@ -126,17 +125,12 @@ impl PreparedStatements { /// /// Nothing, but the message is renamed to a unique, global name. /// - pub fn insert_prepare( - &mut self, - name: &str, - query: Bytes, - data_types: Vec, - ) -> String { - let (_new, global_name) = { self.global.write().insert_prepare(query, data_types) }; + pub(crate) fn insert_prepare(&mut self, name: &str, query: Bytes) -> Prepare { + let (_new, prepare) = { self.global.write().insert_prepare(query) }; - self.insert_internal(name, &global_name); + self.insert_internal(name, prepare.name()); - global_name + prepare } /// Get the global unique name for a prepared statement diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs index fe39081ac..5f6d5c3b1 100644 --- a/pgdog/src/frontend/router/parser/query/prepare.rs +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -64,7 +64,9 @@ impl QueryParser { Some(Node::InsertStmt(stmt)) => self.insert(stmt.into(), &mut context), Some(Node::UpdateStmt(stmt)) => self.update(stmt.into(), &mut context), Some(Node::DeleteStmt(stmt)) => self.delete(stmt.into(), &mut context), - _ => Err(Error::PrepareNotDml), + _ => Ok(Command::Query(Route::write( + ShardWithPriority::new_default_unset(Shard::All), + ))), } } } 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 142010aab..201e616fc 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,10 @@ use bytes::Bytes; use pg_raw_parse::{NodeMut, make::MemoryToken}; -use pgdog_postgres_types::TypeName; -use crate::{frontend::PreparedStatements, net::Prepare}; +use crate::{ + frontend::PreparedStatements, + net::{PREPARE_TEMPLATE_NAME, Prepare}, +}; use super::{Error, StatementRewrite}; @@ -83,32 +85,18 @@ fn rewrite_single_prepared<'a>( ) -> Result { match node { NodeMut::PrepareStmt(mut stmt) => { - let query = Bytes::from(pg_raw_parse::deparse(stmt.query())?.as_str().to_owned()); + let client_name = stmt.name().expect("prepare must have a name").to_owned(); - let data_types: Vec = stmt - .argtypes() - .iter() - .map(|data_type| { - // data_type.typmods().iter().map(|tmod| tmod.) - data_type - .names() - .iter() - .filter_map(|name| name.sval().map(str::to_owned)) - .collect() - }) - .collect(); + // 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 name = stmt.name().expect("prepare must have a name").to_owned(); + let prepare = prepared_statements.insert_prepare(&client_name, query); - let global_name = prepared_statements.insert_prepare(&name, query, data_types); + stmt.set_name(Some(mem.copy_string(prepare.name()))); - stmt.set_name(Some(mem.copy_string(global_name.as_str()))); - - Ok(SimplePreparedRewrite::Prepared { - prepare: prepared_statements - .prepare(&name) - .expect("global cache missing prepare we just inserted"), - }) + Ok(SimplePreparedRewrite::Prepared { prepare }) } NodeMut::ExecuteStmt(mut stmt) => { diff --git a/pgdog/src/net/messages/mod.rs b/pgdog/src/net/messages/mod.rs index 58cafe611..1bc78ad47 100644 --- a/pgdog/src/net/messages/mod.rs +++ b/pgdog/src/net/messages/mod.rs @@ -71,7 +71,7 @@ pub use parameter_status::ParameterStatus; pub use parse::Parse; pub use parse_complete::ParseComplete; pub use payload::Payload; -pub use prepare::Prepare; +pub use prepare::{PREPARE_TEMPLATE_NAME, Prepare}; pub use protocol_version::ProtocolVersion; pub use query::Query; pub use rfq::{ReadyForQuery, TransactionState}; diff --git a/pgdog/src/net/messages/prepare.rs b/pgdog/src/net/messages/prepare.rs index 8a17e283c..23e48ec51 100644 --- a/pgdog/src/net/messages/prepare.rs +++ b/pgdog/src/net/messages/prepare.rs @@ -1,8 +1,11 @@ //! Used for preparing statements sent over with the simple protocol. +//! +//! This is a fake message. It contains a PREPARE query sent using the simple +//! protocol. -use std::str::from_utf8; +use std::str::from_utf8_unchecked; -use pgdog_postgres_types::TypeName; +pub static PREPARE_TEMPLATE_NAME: &str = "__pgdog_template_name"; use super::prelude::*; use super::{Message, Query}; @@ -11,22 +14,21 @@ use super::{Message, Query}; pub struct Prepare { pub(crate) name: Bytes, pub(crate) query: Bytes, - pub(crate) data_types: Vec, } impl Prepare { pub fn query(&self) -> &str { - from_utf8(&self.query).expect("prepare query to be utf-8") + // SAFETY: We only support UTF-8. + unsafe { from_utf8_unchecked(&self.query) } } pub fn name(&self) -> &str { - from_utf8(&self.name).expect("prepare name to be utf-8") + // SAFETY: We only support UTF-8. + unsafe { from_utf8_unchecked(&self.name) } } pub fn len(&self) -> usize { - self.name.len() - + self.query.len() - + self.data_types.iter().map(|t| t.0.len()).sum::() + self.name.len() + self.query.len() } } @@ -52,21 +54,13 @@ impl Protocol for Prepare { impl ToBytes for Prepare { fn to_bytes(&self) -> Bytes { - let query = from_utf8(&self.query).expect("prepare query to be utf-8"); - let name = from_utf8(&self.name).expect("prepare name to be utf-8"); - let data_types = if self.data_types.is_empty() { - "".into() - } else { - format!( - "({})", - self.data_types - .iter() - .map(|data_type| data_type.to_string()) - .collect::>() - .join(", ") - ) - }; + let query = self.query(); + let name = self.name(); + // This is safe because the statement looks like this: + // PREPARE __pgdog_template_name AS [...] + // so the template name will always match first. + let query = query.replacen(PREPARE_TEMPLATE_NAME, name, 1); - Query::new(format!(r#"PREPARE "{}" {} AS {}"#, name, data_types, query)).to_bytes() + Query::new(query).to_bytes() } } From c16ec72d60216bfbaf341384408d749d5673e074 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 09:52:36 -0700 Subject: [PATCH 09/22] fix tests --- .../query_engine/test/rewrite_simple_prepared.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs index 9687ec329..911566372 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs @@ -32,8 +32,9 @@ async fn test_rewrite_prepare() { .await; assert!( - matches!(messages[0].clone(), ProtocolMessage::Query(query) if query.query() == "PREPARE __pgdog_1 AS SELECT $1, $2, $3"), - "expected rewritten prepared statement" + matches!(messages[0].clone(), ProtocolMessage::PrepareClient(prepare) if prepare.query() == "PREPARE __pgdog_template_name AS SELECT $1, $2, $3"), + "expected rewritten prepared statement: {:#?}", + messages, ); let messages = run_test( @@ -43,7 +44,7 @@ async fn test_rewrite_prepare() { .await; assert!( - matches!(messages[0].clone(), ProtocolMessage::Prepare(prepare) if prepare.name() == "__pgdog_1" && prepare.query() == "SELECT $1, $2, $3") + matches!(messages[0].clone(), ProtocolMessage::Prepare(prepare) if prepare.name() == "__pgdog_1" && prepare.query() == "PREPARE __pgdog_template_name AS SELECT $1, $2, $3") ); assert!( @@ -54,7 +55,7 @@ async fn test_rewrite_prepare() { fn rewritten_query(messages: &[ProtocolMessage]) -> String { match &messages[0] { - ProtocolMessage::Query(query) => query.query().to_string(), + ProtocolMessage::PrepareClient(prepare) => prepare.query().to_string(), other => panic!("expected Query, got {other:#?}"), } } @@ -97,7 +98,7 @@ async fn test_reprepare_releases_previous_statement() { assert!(second.starts_with("PREPARE __pgdog_")); assert!(second.ends_with(" AS SELECT $1::bigint + 1")); - // Re-PREPARE mints a new global name. + // PREPARE-ing a different statement with the same name creates a new global name. assert_eq!(global.read().len(), 2); // Only the statement the client replaced is evictable. From 44d902071f900e7890ab9472126b3e49ca1aa5db Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 10:36:39 -0700 Subject: [PATCH 10/22] tests --- .../frontend/client/query_engine/test/mod.rs | 1 + .../query_engine/test/sharded_prepared.rs | 101 ++++++++++++++++++ pgdog/src/frontend/client/test/test_client.rs | 8 ++ .../frontend/router/parser/query/prepare.rs | 10 +- .../frontend/router/parser/query/test/mod.rs | 1 + .../router/parser/query/test/setup.rs | 7 ++ .../router/parser/query/test/test_prepared.rs | 32 ++++++ pgdog/src/net/messages/prepare.rs | 8 ++ 8 files changed, 165 insertions(+), 3 deletions(-) create mode 100644 pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs create mode 100644 pgdog/src/frontend/router/parser/query/test/test_prepared.rs diff --git a/pgdog/src/frontend/client/query_engine/test/mod.rs b/pgdog/src/frontend/client/query_engine/test/mod.rs index efea13559..560851b50 100644 --- a/pgdog/src/frontend/client/query_engine/test/mod.rs +++ b/pgdog/src/frontend/client/query_engine/test/mod.rs @@ -34,6 +34,7 @@ mod schema_changed; mod set; mod set_schema_sharding; mod sharded; +mod sharded_prepared; mod spliced; mod test_omnisharded; mod transaction_state; diff --git a/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs b/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs new file mode 100644 index 000000000..4fe27bde6 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs @@ -0,0 +1,101 @@ +use crate::net::{CommandComplete, DataRow, Message}; + +use super::prelude::*; +use super::*; + +async fn query(client: &mut TestClient, sql: impl ToString) -> Vec { + client.send_simple(Query::new(sql)).await; + client.read_until('Z').await.unwrap() +} + +fn assert_command(messages: &[Message], expected: &str) { + let message = messages + .iter() + .find(|message| message.code() == 'C') + .unwrap(); + let command = CommandComplete::try_from(message.clone()).unwrap(); + assert_eq!(command.command(), expected); +} + +async fn prepared_crud(client: &mut TestClient, table: &str, id: i64, reads: usize) { + query(client, format!("DELETE FROM {table} WHERE id = {id}")).await; + + query( + client, + format!("PREPARE insert_stmt AS INSERT INTO {table} (id, value) VALUES ($1, $2)"), + ) + .await; + let messages = query(client, format!("EXECUTE insert_stmt({id}, 'inserted')")).await; + assert_command(&messages, "INSERT 0 1"); + + query( + client, + format!("PREPARE select_stmt AS SELECT value FROM {table} WHERE id = $1"), + ) + .await; + for _ in 0..reads { + let messages = query(client, format!("EXECUTE select_stmt({id})")).await; + let row = messages + .iter() + .find(|message| message.code() == 'D') + .unwrap(); + let row = DataRow::try_from(row.clone()).unwrap(); + assert_eq!(row.get_text(0).as_deref(), Some("inserted")); + } + + query( + client, + format!("PREPARE update_stmt AS UPDATE {table} SET value = $2 WHERE id = $1"), + ) + .await; + let messages = query(client, format!("EXECUTE update_stmt({id}, 'updated')")).await; + assert_command(&messages, "UPDATE 1"); + + for _ in 0..reads { + let messages = query(client, format!("EXECUTE select_stmt({id})")).await; + let row = messages + .iter() + .find(|message| message.code() == 'D') + .unwrap(); + let row = DataRow::try_from(row.clone()).unwrap(); + assert_eq!(row.get_text(0).as_deref(), Some("updated")); + } + + query( + client, + format!("PREPARE delete_stmt AS DELETE FROM {table} WHERE id = $1"), + ) + .await; + let messages = query(client, format!("EXECUTE delete_stmt({id})")).await; + assert_command(&messages, "DELETE 1"); + + for _ in 0..reads { + let messages = query(client, format!("EXECUTE select_stmt({id})")).await; + assert!(!messages.iter().any(|message| message.code() == 'D')); + assert_command(&messages, "SELECT 0"); + } +} + +#[tokio::test] +async fn test_sharded_prepared_crud() { + crate::logger(); + + let mut client = TestClient::new_sharded(Parameters::default()) + .await + .with_full_prepared_statements(); + let id = client.random_id_for_shard(1); + + prepared_crud(&mut client, "sharded", id, 1).await; +} + +#[tokio::test] +async fn test_omnisharded_prepared_crud() { + crate::logger(); + + let mut client = TestClient::new_sharded(Parameters::default()) + .await + .with_full_prepared_statements(); + let id = client.random_id_for_shard(1); + + prepared_crud(&mut client, "sharded_omni", id, 2).await; +} diff --git a/pgdog/src/frontend/client/test/test_client.rs b/pgdog/src/frontend/client/test/test_client.rs index f026ddb9c..5d569d439 100644 --- a/pgdog/src/frontend/client/test/test_client.rs +++ b/pgdog/src/frontend/client/test/test_client.rs @@ -195,6 +195,14 @@ impl TestClient { Self::new(params).await } + pub(crate) fn with_full_prepared_statements(self) -> Self { + let mut config = config().deref().clone(); + config.config.general.prepared_statements = pgdog_config::PreparedStatements::Full; + set(config).unwrap(); + reload_from_existing().unwrap(); + self + } + /// Send message to client. pub(crate) async fn send(&mut self, message: impl Protocol) { send_message(&mut self.conn, message).await; diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs index 5f6d5c3b1..6da36379e 100644 --- a/pgdog/src/frontend/router/parser/query/prepare.rs +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -51,15 +51,19 @@ impl QueryParser { // // INVARIANT 1: The prepared statements rewriter places it there. // INVARIANT 2: The rewriter renamed the EXECUTE statement to its global name. - let stmt = PreparedStatements::global() + let prepare = PreparedStatements::global() .read() .prepare(stmt.name().expect("execute to have a name")) .ok_or(Error::ExecuteRequiresFull)?; - let query = BufferedQuery::Query(Query::new(stmt.query())); + let query = BufferedQuery::Query(Query::new(prepare.query())); let ast = Cache::get().record(&query)?; + let stmt = match ast.ast.stmts().next() { + Some(Node::PrepareStmt(stmt)) => Some(stmt.query()), + stmt => stmt, + }; - match ast.ast.stmts().next() { + match stmt { Some(Node::SelectStmt(stmt)) => self.select(&ast, stmt, &mut context), Some(Node::InsertStmt(stmt)) => self.insert(stmt.into(), &mut context), Some(Node::UpdateStmt(stmt)) => self.update(stmt.into(), &mut context), diff --git a/pgdog/src/frontend/router/parser/query/test/mod.rs b/pgdog/src/frontend/router/parser/query/test/mod.rs index e9e744c85..95665bf4c 100644 --- a/pgdog/src/frontend/router/parser/query/test/mod.rs +++ b/pgdog/src/frontend/router/parser/query/test/mod.rs @@ -30,6 +30,7 @@ pub mod test_explain; pub mod test_functions; pub mod test_insert; pub mod test_prefer_primary; +pub mod test_prepared; pub mod test_rr; pub mod test_schema_sharding; pub mod test_search_path; diff --git a/pgdog/src/frontend/router/parser/query/test/setup.rs b/pgdog/src/frontend/router/parser/query/test/setup.rs index eac948163..f9fd83c2d 100644 --- a/pgdog/src/frontend/router/parser/query/test/setup.rs +++ b/pgdog/src/frontend/router/parser/query/test/setup.rs @@ -105,6 +105,13 @@ impl QueryParserTest { self } + /// Enable rewriting of simple-protocol PREPARE/EXECUTE statements. + pub(crate) fn with_full_prepared_statements(mut self) -> Self { + self.prepared + .set_level(pgdog_config::PreparedStatements::Full); + self + } + /// Replace the sharded tables configuration on the cluster. pub(crate) fn with_sharded_tables( mut self, diff --git a/pgdog/src/frontend/router/parser/query/test/test_prepared.rs b/pgdog/src/frontend/router/parser/query/test/test_prepared.rs new file mode 100644 index 000000000..957e4a5b5 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_prepared.rs @@ -0,0 +1,32 @@ +use crate::frontend::router::parser::Shard; + +use super::setup::*; + +#[test] +fn test_prepared() { + let mut test = QueryParserTest::new().with_full_prepared_statements(); + + let command = test.execute(vec![ + Query::new("PREPARE __stmt_1 AS SELECT * FROM sharded WHERE id = $1").into(), + ]); + assert!(command.route().is_read()); + assert!(matches!(command.route().shard(), Shard::Direct(_))); + + let command = test.execute(vec![Query::new("EXECUTE __stmt_1(11)").into()]); + assert!(command.route().is_read()); + assert_eq!(command.route().shard(), &Shard::Direct(1)); +} + +#[test] +fn test_prepared_omnisharded_table() { + let mut test = QueryParserTest::new().with_full_prepared_statements(); + + test.execute(vec![ + Query::new("PREPARE __stmt_1 AS SELECT * FROM sharded_omni WHERE id = $1").into(), + ]); + + let command = test.execute(vec![Query::new("EXECUTE __stmt_1(11)").into()]); + assert!(command.route().is_read()); + assert!(command.route().is_omnisharded()); + assert!(matches!(command.route().shard(), Shard::Direct(_))); +} diff --git a/pgdog/src/net/messages/prepare.rs b/pgdog/src/net/messages/prepare.rs index 23e48ec51..2578496bb 100644 --- a/pgdog/src/net/messages/prepare.rs +++ b/pgdog/src/net/messages/prepare.rs @@ -30,6 +30,14 @@ impl Prepare { pub fn len(&self) -> usize { self.name.len() + self.query.len() } + + #[cfg(test)] + pub fn new(name: &str, query: &str) -> Self { + Self { + name: Bytes::from(name.to_owned()), + query: Bytes::from(query.to_owned()), + } + } } impl FromBytes for Prepare { From 01a9579dddde558b9713af3b1b4c9a54cfd5912c Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 10:40:52 -0700 Subject: [PATCH 11/22] rename --- pgdog/src/backend/prepared_statements.rs | 4 ++-- .../test/rewrite_simple_prepared.rs | 6 +++--- .../router/parser/rewrite/statement/plan.rs | 4 ++-- .../parser/rewrite/statement/simple_prepared.rs | 17 +++++++++++++++-- pgdog/src/net/protocol_message.rs | 14 +++++++------- 5 files changed, 29 insertions(+), 16 deletions(-) diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 16f1090f0..7d2c1ba31 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -302,7 +302,7 @@ impl PreparedStatements { self.state.add('3'); } } - ProtocolMessage::PrepareClient(prepare) => { + ProtocolMessage::PrepareFromClient(prepare) => { use crate::net::{CommandComplete, ReadyForQuery}; if self.contains(prepare.name()) { self.state @@ -314,7 +314,7 @@ impl PreparedStatements { self.state.add(ExecutionCode::ReadyForQuery); } } - ProtocolMessage::Prepare(prepare) => { + ProtocolMessage::EnsurePrepared(prepare) => { if self.contains(prepare.name()) { return Ok(HandleResult::Drop); } else { diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs index 911566372..f2c670b49 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs @@ -32,7 +32,7 @@ async fn test_rewrite_prepare() { .await; assert!( - matches!(messages[0].clone(), ProtocolMessage::PrepareClient(prepare) if prepare.query() == "PREPARE __pgdog_template_name AS SELECT $1, $2, $3"), + matches!(messages[0].clone(), ProtocolMessage::PrepareFromClient(prepare) if prepare.query() == "PREPARE __pgdog_template_name AS SELECT $1, $2, $3"), "expected rewritten prepared statement: {:#?}", messages, ); @@ -44,7 +44,7 @@ async fn test_rewrite_prepare() { .await; assert!( - matches!(messages[0].clone(), ProtocolMessage::Prepare(prepare) if prepare.name() == "__pgdog_1" && prepare.query() == "PREPARE __pgdog_template_name AS SELECT $1, $2, $3") + matches!(messages[0].clone(), ProtocolMessage::EnsurePrepared(prepare) if prepare.name() == "__pgdog_1" && prepare.query() == "PREPARE __pgdog_template_name AS SELECT $1, $2, $3") ); assert!( @@ -55,7 +55,7 @@ async fn test_rewrite_prepare() { fn rewritten_query(messages: &[ProtocolMessage]) -> String { match &messages[0] { - ProtocolMessage::PrepareClient(prepare) => prepare.query().to_string(), + ProtocolMessage::PrepareFromClient(prepare) => prepare.query().to_string(), other => panic!("expected Query, got {other:#?}"), } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 2d9c05551..b58958163 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -126,12 +126,12 @@ impl RewritePlan { .for_each(|prepare| match prepare { PrepareExecute::Prepare(prepare) => { request.messages.clear(); - request.push(ProtocolMessage::PrepareClient(prepare.clone())); + request.push(ProtocolMessage::PrepareFromClient(prepare.clone())); } PrepareExecute::Execute(prepare) => { request .messages - .splice(0..0, vec![ProtocolMessage::Prepare(prepare.clone())]); + .splice(0..0, vec![ProtocolMessage::EnsurePrepared(prepare.clone())]); } }); } 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 201e616fc..b1deee6a1 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -182,8 +182,21 @@ mod tests { !sql.contains("test_stmt"), "original name should be replaced: {sql}" ); - assert!(plan.prepare_rewrites.is_empty()); + 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 $1, $2" + ); + } + + _ => panic!("expected PrepareExecute::Prepare"), + } } #[test] @@ -202,7 +215,7 @@ mod tests { match prepare { PrepareExecute::Execute(prepare) => { assert!(prepare.name().starts_with("__pgdog_")); - assert_eq!(prepare.query(), "SELECT 1"); + assert_eq!(prepare.query(), "PREPARE __pgdog_template_name AS SELECT 1"); } _ => panic!("expected PrepareExecute::Execute"), diff --git a/pgdog/src/net/protocol_message.rs b/pgdog/src/net/protocol_message.rs index deeceda86..544f1e8f0 100644 --- a/pgdog/src/net/protocol_message.rs +++ b/pgdog/src/net/protocol_message.rs @@ -13,8 +13,8 @@ pub enum ProtocolMessage { Bind(Bind), Parse(Parse), Describe(Describe), - Prepare(Prepare), - PrepareClient(Prepare), + EnsurePrepared(Prepare), + PrepareFromClient(Prepare), Execute(Execute), Close(Close), Query(Query), @@ -62,8 +62,8 @@ impl ProtocolMessage { Self::Bind(bind) => bind.len(), Self::Parse(parse) => parse.len(), Self::Describe(describe) => describe.len(), - Self::Prepare(prepare) => prepare.len(), - Self::PrepareClient(prepare) => prepare.len(), + Self::EnsurePrepared(prepare) => prepare.len(), + Self::PrepareFromClient(prepare) => prepare.len(), Self::Execute(execute) => execute.len(), Self::Close(close) => close.len(), Self::Query(query) => query.len(), @@ -83,7 +83,7 @@ impl Protocol for ProtocolMessage { Self::Bind(bind) => bind.code(), Self::Parse(parse) => parse.code(), Self::Describe(describe) => describe.code(), - Self::Prepare { .. } | Self::PrepareClient { .. } => 'Q', + Self::EnsurePrepared { .. } | Self::PrepareFromClient { .. } => 'Q', Self::Execute(execute) => execute.code(), Self::Close(close) => close.code(), Self::Query(query) => query.code(), @@ -123,8 +123,8 @@ impl ToBytes for ProtocolMessage { Self::Bind(bind) => bind.to_bytes(), Self::Parse(parse) => parse.to_bytes(), Self::Describe(describe) => describe.to_bytes(), - Self::Prepare(prepare) => prepare.to_bytes(), - Self::PrepareClient(prepare) => prepare.to_bytes(), + Self::EnsurePrepared(prepare) => prepare.to_bytes(), + Self::PrepareFromClient(prepare) => prepare.to_bytes(), Self::Execute(execute) => execute.to_bytes(), Self::Close(close) => close.to_bytes(), Self::Query(query) => query.to_bytes(), From 312a6cbc116889f2daff697f745adefe35c82e16 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 12:03:00 -0700 Subject: [PATCH 12/22] fmt --- pgdog/src/frontend/prepared_statements/global_cache.rs | 4 ++-- pgdog/src/util.rs | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 04af90381..f198af846 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -367,8 +367,8 @@ mod test { assert_eq!(first, second); assert_eq!(cache.len(), 1); - assert_eq!(used(&cache, &first.name()), 2); - assert_eq!(used(&cache, &second.name()), 2); + assert_eq!(used(&cache, first.name()), 2); + assert_eq!(used(&cache, second.name()), 2); // A Parse never re-uses a SQL PREPARE statement. let (new, extended) = cache.insert(&parse); diff --git a/pgdog/src/util.rs b/pgdog/src/util.rs index e030665c6..023ed2b11 100644 --- a/pgdog/src/util.rs +++ b/pgdog/src/util.rs @@ -533,6 +533,7 @@ mod test { } #[test] + #[pgdog_macros::flaky] fn test_node_id_error() { let _guard = remove_env_var("NODE_ID"); assert!(node_id().is_err()); From 99526bea88a584c6c26e064987f6c25c6978d16c Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 12:36:58 -0700 Subject: [PATCH 13/22] aha! --- pgdog-postgres-types/src/lib.rs | 2 - pgdog-postgres-types/src/type_name.rs | 39 ------------ pgdog/src/backend/prepared_statements.rs | 13 +++- pgdog/src/backend/server.rs | 58 +++++++++++++++++- .../query_engine/test/sharded_prepared.rs | 60 +++++++++++++++++++ .../frontend/router/parser/query/prepare.rs | 6 +- 6 files changed, 132 insertions(+), 46 deletions(-) delete mode 100644 pgdog-postgres-types/src/type_name.rs diff --git a/pgdog-postgres-types/src/lib.rs b/pgdog-postgres-types/src/lib.rs index b6fa2e643..59af72502 100644 --- a/pgdog-postgres-types/src/lib.rs +++ b/pgdog-postgres-types/src/lib.rs @@ -16,7 +16,6 @@ pub mod smallint; pub mod text; pub mod timestamp; pub mod timestamptz; -pub mod type_name; pub mod uuid; pub mod vector; @@ -33,4 +32,3 @@ pub use numeric::Numeric; pub use oid::Oid; pub use timestamp::Timestamp; pub use timestamptz::TimestampTz; -pub use type_name::TypeName; diff --git a/pgdog-postgres-types/src/type_name.rs b/pgdog-postgres-types/src/type_name.rs deleted file mode 100644 index d417adead..000000000 --- a/pgdog-postgres-types/src/type_name.rs +++ /dev/null @@ -1,39 +0,0 @@ -//! PostgreSQL type name. - -use std::fmt; -use std::sync::Arc; - -/// Name of a PostgreSQL data type. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct TypeName(pub Arc>); - -impl From> for TypeName { - fn from(names: Vec) -> Self { - Self(Arc::new(names)) - } -} - -impl FromIterator for TypeName { - fn from_iter>(iter: T) -> Self { - Self(Arc::new(iter.into_iter().collect())) - } -} - -impl fmt::Display for TypeName { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(&self.0.join(".")) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn conversions() { - let data_type = TypeName::from(vec![String::from("pg_catalog"), String::from("int8")]); - - assert_eq!(*data_type.0, ["pg_catalog", "int8"]); - assert_eq!(data_type.to_string(), "pg_catalog.int8"); - } -} diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 7d2c1ba31..0cce54ac9 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -105,6 +105,7 @@ pub struct PreparedStatements { config: PreparedStatementsConfig, memory_used: usize, oids: Arc, + in_transaction: bool, } #[cfg(test)] @@ -126,6 +127,7 @@ impl PreparedStatements { config: PreparedStatementsConfig::default(), memory_used: 0, oids, + in_transaction: false, } } @@ -135,6 +137,10 @@ impl PreparedStatements { self.config = config; } + pub(super) fn set_in_transaction(&mut self, in_transaction: bool) { + self.in_transaction = in_transaction; + } + /// Current prepared statement settings. pub fn config(&self) -> PreparedStatementsConfig { self.config @@ -307,10 +313,12 @@ impl PreparedStatements { if self.contains(prepare.name()) { self.state .add_simulated(CommandComplete::from_str("PREPARE").message()?); - self.state - .add_simulated(ReadyForQuery::in_transaction(false).message()?); + self.state.add_simulated( + ReadyForQuery::in_transaction(self.in_transaction).message()?, + ); return Ok(HandleResult::Drop); } else { + self.parses.push_back(prepare.name().to_owned()); self.state.add(ExecutionCode::ReadyForQuery); } } @@ -323,6 +331,7 @@ impl PreparedStatements { // Prepare turns into a Simple Query ('Q') so it expects a regular RFQ back. self.state.add_ignore(ExecutionCode::ReadyForQuery); + self.parses.push_back(prepare.name().to_owned()); return Ok(HandleResult::Forward); } } diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 17b57a6cf..258ae392d 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -472,6 +472,9 @@ impl Server { self.in_transaction = true; } + self.prepared_statements + .set_in_transaction(self.in_transaction); + for message in client_request.messages.iter() { self.send_one(message).await?; } @@ -1334,8 +1337,10 @@ pub mod test { }; use crate::{ - backend::pool::token_cache::TokenCache, config::Memory, frontend::PreparedStatements, - net::*, + backend::pool::token_cache::TokenCache, + config::Memory, + frontend::PreparedStatements, + net::{Prepare, *}, }; use super::{Error, *}; @@ -4378,6 +4383,55 @@ pub mod test { assert_eq!(Duration::from_millis(1500), server.effective_max_age(base)); } + #[tokio::test] + async fn test_prepare_from_client() { + let mut server = test_server().await; + + // The last 2 will be simulated + // and we won't receive a "prepared statement already exists" error. + for _ in 0..3 { + server + .send( + &vec![ProtocolMessage::PrepareFromClient(Prepare::new( + "__stmt_1", + "PREPARE __pgdog_template_name AS SELECT $1", + ))] + .into(), + ) + .await + .unwrap(); + + for c in ['C', 'Z'] { + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), c); + } + } + + assert!(server.prepared_statements_mut().contains("__stmt_1")); + } + + #[tokio::test] + async fn test_prepared_execute() { + let mut server = test_server().await; + + for _ in 0..3 { + let req = vec![ + ProtocolMessage::EnsurePrepared(Prepare::new( + "__stmt_1", + "PREPARE __pgdog_template_name (int) AS SELECT $1", + )), + ProtocolMessage::Query(Query::new("EXECUTE __stmt_1 (1)")), + ]; + + server.send(&req.into()).await.unwrap(); + + for c in ['T', 'D', 'C', 'Z'] { + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), c); + } + } + } + #[test] fn test_apply_lifetime_jitter_zero_is_noop() { let mut server = Server::default(); diff --git a/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs b/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs index 4fe27bde6..002d589f1 100644 --- a/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs +++ b/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs @@ -99,3 +99,63 @@ async fn test_omnisharded_prepared_crud() { prepared_crud(&mut client, "sharded_omni", id, 2).await; } + +#[tokio::test] +async fn test_cross_shard_prepared() { + let mut client = TestClient::new_sharded(Parameters::default()) + .await + .with_full_prepared_statements(); + + client + .send_simple(Query::new( + "PREPARE __stmt_1 AS INSERT INTO sharded (id, value) VALUES ($1, $2)", + )) + .await; + client.read_until('Z').await.unwrap(); + client + .send_simple(Query::new( + "PREPARE __stmt_2 AS SELECT id, value FROM sharded ORDER BY id DESC", + )) + .await; + client.read_until('Z').await.unwrap(); + + for shard in 0..20 { + let id = client.random_id_for_shard(shard % 2); + client + .send_simple(Query::new(format!("EXECUTE __stmt_1({}, 'value')", id))) + .await; + client.read_until('Z').await.unwrap(); + } + + client.send_simple(Query::new("EXECUTE __stmt_2")).await; + let msgs = client.read_until('Z').await.unwrap(); + + client + .send_simple(Query::new("PREPARE __stmt_3 AS DELETE FROM sharded")) + .await; + client.read_until('Z').await.unwrap(); + + client.send_simple(Query::new("EXECUTE __stmt_3")).await; + client.read_until('Z').await.unwrap(); + + assert_eq!(msgs.iter().filter(|m| m.code() == 'D').count(), 20); + let ids = msgs + .iter() + .filter(|m| m.code() == 'D') + .map(|m| { + DataRow::try_from(m.clone()) + .unwrap() + .get_int(0, true) + .unwrap() + }) + .collect::>(); + let mut sorted = ids.clone(); + sorted.sort(); + + assert_eq!(sorted.into_iter().rev().collect::>(), ids); + + client.send_simple(Query::new("EXECUTE __stmt_2")).await; + let msgs = client.read_until('Z').await.unwrap(); + + assert_eq!(msgs.iter().filter(|m| m.code() == 'D').count(), 0); +} diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs index 6da36379e..92cea2a7c 100644 --- a/pgdog/src/frontend/router/parser/query/prepare.rs +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -2,7 +2,7 @@ use pg_raw_parse::nodes::{ExecuteStmt, PrepareStmt}; use crate::{ frontend::{BufferedQuery, PreparedStatements}, - net::Query, + net::{PREPARE_TEMPLATE_NAME, Query}, }; use super::*; @@ -56,6 +56,10 @@ impl QueryParser { .prepare(stmt.name().expect("execute to have a name")) .ok_or(Error::ExecuteRequiresFull)?; + // Make sure we never store unique statement names + // in the cache! + debug_assert!(prepare.query().contains(PREPARE_TEMPLATE_NAME)); + let query = BufferedQuery::Query(Query::new(prepare.query())); let ast = Cache::get().record(&query)?; let stmt = match ast.ast.stmts().next() { From e0d26c66d99ecb93538f6bf7e46c2f9c4c197419 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 12:50:38 -0700 Subject: [PATCH 14/22] fix test --- .../src/frontend/client/query_engine/test/sharded_prepared.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs b/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs index 002d589f1..af49a84b9 100644 --- a/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs +++ b/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs @@ -106,6 +106,10 @@ async fn test_cross_shard_prepared() { .await .with_full_prepared_statements(); + // Clean up after bad tests. + client.send_simple(Query::new("DELETE FROM sharded")).await; + client.read_until('Z').await.unwrap(); + client .send_simple(Query::new( "PREPARE __stmt_1 AS INSERT INTO sharded (id, value) VALUES ($1, $2)", From 6a1db8a1d230f418aaee67ea36a82e6d6d55c1ed Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 21:18:48 -0700 Subject: [PATCH 15/22] save --- pgdog/src/backend/prepared_statements.rs | 32 ++++++-- pgdog/src/backend/server.rs | 77 ++++++++++++++++++- .../frontend/prepared_statements/cache_key.rs | 25 +++--- .../frontend/prepared_statements/statement.rs | 11 ++- 4 files changed, 124 insertions(+), 21 deletions(-) diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 0cce54ac9..1082d39dc 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -5,7 +5,6 @@ use std::{ time::{Duration, Instant}, }; -use crate::util::time::deadline; use crate::{ frontend::{self, prepared_statements::GlobalCache}, net::{ @@ -13,7 +12,9 @@ use crate::{ ToBytes, messages::{ParameterDescription, RowDescription, parse::Parse}, }, + state::State, }; +use crate::{net::ErrorResponse, util::time::deadline}; use parking_lot::RwLock; use pgdog_stats::PreparedStatementsConfig; @@ -105,7 +106,7 @@ pub struct PreparedStatements { config: PreparedStatementsConfig, memory_used: usize, oids: Arc, - in_transaction: bool, + server_state: State, } #[cfg(test)] @@ -127,7 +128,7 @@ impl PreparedStatements { config: PreparedStatementsConfig::default(), memory_used: 0, oids, - in_transaction: false, + server_state: State::Idle, } } @@ -137,8 +138,8 @@ impl PreparedStatements { self.config = config; } - pub(super) fn set_in_transaction(&mut self, in_transaction: bool) { - self.in_transaction = in_transaction; + pub(super) fn set_server_state(&mut self, state: State) { + self.server_state = state; } /// Current prepared statement settings. @@ -278,6 +279,8 @@ impl PreparedStatements { if !parse.anonymous() { if self.contains(parse.name()) { + // TODO(lev): perform the same in errored transaction check + // as we do for PREPARE below. self.state.add_simulated(ParseComplete.message()?); return Ok(HandleResult::Drop); } else { @@ -311,10 +314,23 @@ impl PreparedStatements { ProtocolMessage::PrepareFromClient(prepare) => { use crate::net::{CommandComplete, ReadyForQuery}; if self.contains(prepare.name()) { - self.state - .add_simulated(CommandComplete::from_str("PREPARE").message()?); + if self.server_state == State::TransactionError { + self.state + .add_simulated(ErrorResponse::in_failed_transaction().message()?); + } else { + self.state + .add_simulated(CommandComplete::from_str("PREPARE").message()?); + } + self.state.add_simulated( - ReadyForQuery::in_transaction(self.in_transaction).message()?, + if self.server_state == State::TransactionError { + ReadyForQuery::error() + } else { + ReadyForQuery::in_transaction( + self.server_state == State::IdleInTransaction, + ) + } + .message()?, ); return Ok(HandleResult::Drop); } else { diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 258ae392d..5221b0e90 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -472,9 +472,6 @@ impl Server { self.in_transaction = true; } - self.prepared_statements - .set_in_transaction(self.in_transaction); - for message in client_request.messages.iter() { self.send_one(message).await?; } @@ -683,6 +680,9 @@ impl Server { _ => (), } + self.prepared_statements + .set_server_state(self.stats.get_state()); + trace!("{:#?} <<< [{}]", message, self.addr()); Ok(message) @@ -4432,6 +4432,77 @@ pub mod test { } } + #[tokio::test] + async fn test_prepare_in_transaction() { + let mut server = test_server().await; + + server.execute("BEGIN").await.unwrap(); + + for _ in 0..3 { + server + .send( + &vec![ProtocolMessage::PrepareFromClient(Prepare::new( + "__stmt_1", + "PREPARE __pgdog_template_name AS SELECT $1", + ))] + .into(), + ) + .await + .unwrap(); + + let cmd = server.read().await.unwrap(); + assert_eq!(cmd.code(), 'C'); + let rfq = server.read().await.unwrap(); + assert!(rfq.in_transaction()); + } + + server.execute("ROLLBACK").await.unwrap(); + } + + #[tokio::test] + async fn test_prepare_in_transaction_error() { + let mut server = test_server().await; + + server + .send( + &vec![ProtocolMessage::PrepareFromClient(Prepare::new( + "__stmt_1", + "PREPARE __pgdog_template_name AS SELECT $1", + ))] + .into(), + ) + .await + .unwrap(); + + let cmd = server.read().await.unwrap(); + assert_eq!(cmd.code(), 'C'); + let rfq = server.read().await.unwrap(); + assert!(!rfq.in_transaction()); + + server.execute("BEGIN").await.unwrap(); + + let _ = server.execute("SELECT asd").await; + + for _ in 0..3 { + server + .send( + &vec![ProtocolMessage::PrepareFromClient(Prepare::new( + "__stmt_1", + "PREPARE __pgdog_template_name AS SELECT $1", + ))] + .into(), + ) + .await + .unwrap(); + + let err = ErrorResponse::try_from(server.read().await.unwrap()).unwrap(); + assert_eq!(err.code, "25P02"); + + let rfq = ReadyForQuery::try_from(server.read().await.unwrap()).unwrap(); + assert!(rfq.is_transaction_aborted()); + } + } + #[test] fn test_apply_lifetime_jitter_zero_is_noop() { let mut server = Server::default(); diff --git a/pgdog/src/frontend/prepared_statements/cache_key.rs b/pgdog/src/frontend/prepared_statements/cache_key.rs index b498bafa8..145cd6037 100644 --- a/pgdog/src/frontend/prepared_statements/cache_key.rs +++ b/pgdog/src/frontend/prepared_statements/cache_key.rs @@ -28,19 +28,26 @@ impl MemoryUsage for CacheKey { } impl CacheKey { - pub(super) fn query_ref(&self) -> &Bytes { + /// Get a UTF-8 encoded query string + /// stored in the cache. + pub(crate) fn query(&self) -> Result<&str, crate::net::Error> { match self { - Self::Extended { query, .. } => query, - Self::Simple { query, .. } => query, + Self::Extended { query, .. } => Ok(from_utf8(&query[0..query.len() - 1])?), + Self::Simple { query } => Ok(from_utf8(&query)?), // Simple queries are regular Rust strings. } } +} - /// Get a UTF-8 encoded query string - /// stored in the cache. - pub(crate) fn query(&self) -> Result<&str, crate::net::Error> { - let query = self.query_ref(); +#[cfg(test)] +mod test { + use super::*; - // Postgres string. - Ok(from_utf8(&query[0..query.len() - 1])?) + impl CacheKey { + pub(crate) fn query_ref(&self) -> &Bytes { + match self { + Self::Extended { query, .. } => query, + Self::Simple { query, .. } => query, + } + } } } diff --git a/pgdog/src/frontend/prepared_statements/statement.rs b/pgdog/src/frontend/prepared_statements/statement.rs index 3630a054a..21b7b69bd 100644 --- a/pgdog/src/frontend/prepared_statements/statement.rs +++ b/pgdog/src/frontend/prepared_statements/statement.rs @@ -23,7 +23,16 @@ pub(crate) enum StatementType { impl MemoryUsage for StatementType { fn memory_usage(&self) -> usize { - 0 + match self { + Self::Prepare { prepare } => prepare.len(), + Self::Parse { parse, rewrite } => { + parse.len() + + rewrite + .as_ref() + .map(|rewrite| rewrite.len()) + .unwrap_or_default() + } + } } } From f01d710c730a82c207ebf25ccb69a96c4456e11c Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 21:23:04 -0700 Subject: [PATCH 16/22] move test --- pgdog/src/backend/server.rs | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 5221b0e90..0a2751288 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -4367,22 +4367,6 @@ pub mod test { ); } - #[test] - fn test_effective_max_age_default_is_base() { - let server = Server::default(); - let base = Duration::from_secs(60); - assert_eq!(base, server.effective_max_age(base)); - } - - #[test] - fn test_effective_max_age_uses_stored_value() { - let mut server = Server::default(); - server.set_max_age(Duration::from_millis(1500)); - let base = Duration::from_secs(60); - // Stored value wins regardless of the supplied base. - assert_eq!(Duration::from_millis(1500), server.effective_max_age(base)); - } - #[tokio::test] async fn test_prepare_from_client() { let mut server = test_server().await; @@ -4503,6 +4487,22 @@ pub mod test { } } + #[test] + fn test_effective_max_age_default_is_base() { + let server = Server::default(); + let base = Duration::from_secs(60); + assert_eq!(base, server.effective_max_age(base)); + } + + #[test] + fn test_effective_max_age_uses_stored_value() { + let mut server = Server::default(); + server.set_max_age(Duration::from_millis(1500)); + let base = Duration::from_secs(60); + // Stored value wins regardless of the supplied base. + assert_eq!(Duration::from_millis(1500), server.effective_max_age(base)); + } + #[test] fn test_apply_lifetime_jitter_zero_is_noop() { let mut server = Server::default(); From de027de8232e412d17be5edec919f234cf1d6c95 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 21:34:36 -0700 Subject: [PATCH 17/22] process simulated messages through state manager --- pgdog/src/backend/server.rs | 2 +- pgdog/src/frontend/prepared_statements/cache_key.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 0a2751288..9e25ef78b 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -569,7 +569,7 @@ impl Server { if let Some(message) = self.prepared_statements.state_mut().get_simulated() { // INVARIANT: omni dedup in multi_shard relies on this being process-unique; // never substitute a non-unique value here. - return Ok(message.backend(self.id)); + break message.backend(self.id); } match self.stream_buffer.read(self.stream.as_mut().unwrap()).await { Ok(message) => { diff --git a/pgdog/src/frontend/prepared_statements/cache_key.rs b/pgdog/src/frontend/prepared_statements/cache_key.rs index 145cd6037..fae1cfad0 100644 --- a/pgdog/src/frontend/prepared_statements/cache_key.rs +++ b/pgdog/src/frontend/prepared_statements/cache_key.rs @@ -33,7 +33,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. } } } From 6b8cbe2acc0938993b8bc74d93a1bd7d5e97f18d Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 21:59:30 -0700 Subject: [PATCH 18/22] tests --- .../client/query_engine/test/replicas.rs | 31 ++++++++++++++++--- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/test/replicas.rs b/pgdog/src/frontend/client/query_engine/test/replicas.rs index 29dff7e03..d2f790691 100644 --- a/pgdog/src/frontend/client/query_engine/test/replicas.rs +++ b/pgdog/src/frontend/client/query_engine/test/replicas.rs @@ -1,6 +1,7 @@ use crate::{ backend::databases::databases, config::Role, + net::ParseComplete, net::{Parameters, ToBytes}, }; @@ -69,7 +70,13 @@ async fn test_round_robin_with_replicas() { assert_eq!(state.stats.counts.healthchecks, idle); // Parse count depends on number of idle connections (prepared statement sync). assert!(state.stats.counts.parse_count >= idle); - assert!(state.stats.counts.parse_count <= idle + 1); + assert_eq!( + state.stats.counts.parse_count, + 13, // We are counting simulated messages too + "{} <= {}", + state.stats.counts.parse_count, + idle + 1 + ); pool_recv -= (healthcheck_len_recv * state.stats.counts.healthchecks) as isize; } Role::Replica => { @@ -78,7 +85,13 @@ async fn test_round_robin_with_replicas() { assert_eq!(state.stats.counts.bind_count, 13); assert_eq!(state.stats.counts.rollbacks, 0); assert!(state.stats.counts.healthchecks <= idle + 1); - assert!(state.stats.counts.parse_count <= idle + 1); + assert_eq!( + state.stats.counts.parse_count, + 13, // We are counting simulated messages. + "{} <= {}", + state.stats.counts.parse_count, + idle + 1 + ); pool_sent -= (healthcheck_len_sent * state.stats.counts.healthchecks) as isize; } Role::Auto => unreachable!("role auto"), @@ -94,6 +107,16 @@ async fn test_round_robin_with_replicas() { ); } - assert!(pool_sent <= len_sent as isize); - assert!(pool_recv <= len_recv as isize); + assert!( + pool_sent <= len_sent as isize, + "{} <= {}", + pool_sent, + len_sent + ); + assert!( + pool_recv <= len_recv as isize + (ParseComplete.to_bytes().len() as isize * 12), // We count simulated length now + "{} <= {}", + pool_recv, + len_recv + ); } From ebcc0c76f2651e8877e08ecd15bd029f29c1fefc Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Sun, 16 Aug 2026 22:15:33 -0700 Subject: [PATCH 19/22] Drop: Clone shouldnt mutate global state --- pgdog/src/frontend/client/mod.rs | 1 + pgdog/src/frontend/prepared_statements/mod.rs | 8 +------- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 8dc0e221c..1eb869705 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -715,6 +715,7 @@ impl Client { impl Drop for Client { fn drop(&mut self) { self.comms.disconnect(); + self.prepared_statements.close_all(); } } diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index 1d350ac25..a4b05e948 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -52,12 +52,6 @@ impl Default for PreparedStatements { } } -impl Drop for PreparedStatements { - fn drop(&mut self) { - self.close_all(); - } -} - impl PreparedStatements { /// New shared prepared statements cache instance. /// @@ -177,7 +171,7 @@ impl PreparedStatements { /// /// This only happens when the client disconnects. This will update /// the global usage counters of all of client's prepared statements. - fn close_all(&mut self) { + pub(super) fn close_all(&mut self) { if !self.local.is_empty() { let mut global = self.global.write(); From 158f76112e0b5431539f274da7953f2752cabd6d Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Mon, 17 Aug 2026 14:57:47 -0700 Subject: [PATCH 20/22] woops --- pgdog/src/backend/pool/cluster.rs | 2 +- pgdog/src/frontend/router/parser/query/mod.rs | 6 +++--- .../frontend/router/parser/query/prepare.rs | 6 +++--- .../frontend/router/parser/query/select.rs | 11 +++++++---- .../src/frontend/router/parser/query/show.rs | 2 +- .../router/parser/rewrite/statement/mod.rs | 19 +++++++++++-------- pgdog/src/frontend/router/parser/statement.rs | 4 +--- pgdog/src/frontend/router/round_robin.rs | 4 ++-- 8 files changed, 29 insertions(+), 25 deletions(-) diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 238eb3e9f..0fd2e008f 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -767,7 +767,7 @@ impl Cluster { ) -> Result, crate::backend::Error> { let shard = self .shards - .get(round_robin::next() % self.shards.len().max(1)) + .get(round_robin::next(self.shards.len().max(1))) .ok_or(crate::backend::pool::Error::NoDatabases)?; let mut server = shard.primary_or_replica(&Request::default()).await?; diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index fca8ffb96..f1b379f30 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -349,7 +349,7 @@ impl QueryParser { context .shards_calculator .push(ShardWithPriority::new_rr_empty_query(Shard::Direct( - round_robin::next() % context.shards, + round_robin::next(context.shards), ))); // Send empty query to any shard. return Ok(Command::Query(Route::read( @@ -431,7 +431,7 @@ impl QueryParser { )); } - Node::PrepareStmt(stmt) => self.prepare(stmt), + Node::PrepareStmt(stmt) => self.prepare(stmt, context), Node::ExecuteStmt(stmt) => self.execute(stmt, context), @@ -455,7 +455,7 @@ impl QueryParser { context .shards_calculator .push(ShardWithPriority::new_rr_not_executable(Shard::Direct( - round_robin::next() % context.shards, + round_robin::next(context.shards), ))); // Since this query isn't executable and we decided diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs index 92cea2a7c..f570282e3 100644 --- a/pgdog/src/frontend/router/parser/query/prepare.rs +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -16,16 +16,16 @@ impl QueryParser { pub(super) fn prepare( &self, stmt: &PrepareStmt, - // context: &mut QueryParserContext<'_>, + context: &mut QueryParserContext<'_>, ) -> Result { let query = stmt.query(); let route = match query { Node::SelectStmt(_) => Route::read(ShardWithPriority::new_rr_not_executable( - round_robin::next().into(), + (round_robin::next(context.shards)).into(), )), _ => Route::write(ShardWithPriority::new_rr_not_executable( - round_robin::next().into(), + (round_robin::next(context.shards)).into(), )), }; diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs index af301a941..493742466 100644 --- a/pgdog/src/frontend/router/parser/query/select.rs +++ b/pgdog/src/frontend/router/parser/query/select.rs @@ -107,7 +107,7 @@ impl QueryParser { // SELECT NOW(), SELECT 1 if shards.is_empty() && stmt.from_clause().is_empty() { - let shard = Shard::Direct(round_robin::next() % context.shards); + let shard = Shard::Direct(round_robin::next(context.shards)); if let Some(recorder) = self.recorder_mut() { recorder.record_entry(Some(shard.clone()), "SELECT omnishard no table".to_string()); @@ -210,12 +210,15 @@ impl QueryParser { .tables() .is_omnisharded_sticky_default() { - (context.router_context.sticky.omni_index, "sticky") + ( + context.router_context.sticky.omni_index % context.shards, + "sticky", + ) } else { - (round_robin::next(), "round robin") + (round_robin::next(context.shards), "round robin") }; - let shard = Shard::Direct(rr_index % context.shards); + let shard = Shard::Direct(rr_index); // Routed to a single shard via the omnisharded-by-default path // (non-sharded tables, including system catalogs). diff --git a/pgdog/src/frontend/router/parser/query/show.rs b/pgdog/src/frontend/router/parser/query/show.rs index 731cfb9a6..c4050fbcd 100644 --- a/pgdog/src/frontend/router/parser/query/show.rs +++ b/pgdog/src/frontend/router/parser/query/show.rs @@ -18,7 +18,7 @@ impl QueryParser { context .shards_calculator .push(ShardWithPriority::new_rr_no_table(Shard::Direct( - round_robin::next() % context.shards, + round_robin::next(context.shards), ))); let route = Route::write(context.shards_calculator.shard().clone()) .with_read(context.read_only); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index b91228b5f..45d94f5f7 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -113,18 +113,14 @@ impl<'a> StatementRewrite<'a> { plan.params = plan.params.max(param.number as u16) } }), - Node::PrepareStmt(_) | Node::ExecuteStmt(_) | Node::ExplainStmt(_) => {} + Node::PrepareStmt(_) => { + // Will use parameters for replacing args, not materialize values. + self.extended = true; + } // We can't do anything with DDL statements _ => return Ok(plan), } - // Handle top-level PREPARE/EXECUTE statements. - let prepared_result = self.rewrite_simple_prepared(stmt.stmt_mut(), mem)?; - if prepared_result.rewritten { - self.rewritten = true; - plan.prepare_rewrites = prepared_result.rewrites; - } - // Inject pgdog.unique_id() for missing BIGINT primary keys. // This must run BEFORE the unique_id rewriter so the injected // function calls get processed. @@ -166,6 +162,13 @@ impl<'a> StatementRewrite<'a> { plan.stmt = Some(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned()); } + // Handle top-level PREPARE/EXECUTE statements. + let prepared_result = self.rewrite_simple_prepared(stmt.stmt_mut(), mem)?; + if prepared_result.rewritten { + self.rewritten = true; + plan.prepare_rewrites = prepared_result.rewrites; + } + if let Node::InsertStmt(insert) = stmt.stmt() { self.split_insert(insert, &mut plan)?; } diff --git a/pgdog/src/frontend/router/parser/statement.rs b/pgdog/src/frontend/router/parser/statement.rs index 5c58c9842..39cb8305e 100644 --- a/pgdog/src/frontend/router/parser/statement.rs +++ b/pgdog/src/frontend/router/parser/statement.rs @@ -1090,9 +1090,7 @@ impl<'a, 'b: 'a, 'c> StatementParser<'a, 'b, 'c> { if let Some(table) = ctx.table && Tables::new(self.schema).sharded(table).is_some() { - Ok(Some(Shard::Direct( - round_robin::next() % self.schema.shards, - ))) + Ok(Some(Shard::Direct(round_robin::next(self.schema.shards)))) } else { Ok(None) } diff --git a/pgdog/src/frontend/router/round_robin.rs b/pgdog/src/frontend/router/round_robin.rs index 6cc44ed7d..008453211 100644 --- a/pgdog/src/frontend/router/round_robin.rs +++ b/pgdog/src/frontend/router/round_robin.rs @@ -4,6 +4,6 @@ use std::sync::atomic::{AtomicUsize, Ordering}; static ROUND_ROBIN: Lazy = Lazy::new(|| AtomicUsize::new(0)); /// Get next round robin number. -pub fn next() -> usize { - ROUND_ROBIN.fetch_add(1, Ordering::Relaxed) +pub fn next(shards: usize) -> usize { + ROUND_ROBIN.fetch_add(1, Ordering::Relaxed) % shards } From 0038a5e1a1c2b737cd4ecf8918b60014b5983b2a Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Mon, 17 Aug 2026 15:48:25 -0700 Subject: [PATCH 21/22] support unique id and id injection in prepare --- pgdog/src/backend/server.rs | 4 +- pgdog/src/frontend/mod.rs | 2 +- .../prepared_statements/global_cache.rs | 25 +++-- pgdog/src/frontend/prepared_statements/mod.rs | 14 ++- .../frontend/prepared_statements/statement.rs | 18 ++- pgdog/src/frontend/router/mod.rs | 2 +- pgdog/src/frontend/router/parser/mod.rs | 4 +- .../frontend/router/parser/query/prepare.rs | 4 +- .../parser/rewrite/statement/auto_id.rs | 61 +++++++++- .../router/parser/rewrite/statement/mod.rs | 43 ++++--- .../rewrite/statement/simple_prepared.rs | 105 +++++++++++++++++- 11 files changed, 233 insertions(+), 49 deletions(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 9e25ef78b..7fd103c9f 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -1339,7 +1339,7 @@ pub mod test { use crate::{ backend::pool::token_cache::TokenCache, config::Memory, - frontend::PreparedStatements, + frontend::{PreparedStatements, RewritePlan}, net::{Prepare, *}, }; @@ -2325,7 +2325,7 @@ pub 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()); + let prepare = prep.insert_prepare(name, query.clone(), &RewritePlan::default()); assert_eq!(prepare.name(), "__pgdog_1"); server diff --git a/pgdog/src/frontend/mod.rs b/pgdog/src/frontend/mod.rs index 915be3ab6..4109448c9 100644 --- a/pgdog/src/frontend/mod.rs +++ b/pgdog/src/frontend/mod.rs @@ -26,6 +26,6 @@ pub(crate) use prepared_statements::PreparedStatements; #[cfg(debug_assertions)] pub use query_logger::QueryLogger; pub(crate) use regex_parser::RegexParser; -pub use router::{Command, Router, SetParam}; +pub use router::{Command, RewritePlan, Router, SetParam}; pub use router::{RouterContext, SearchPath}; pub use stats::Stats; diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index f198af846..f448fa5c4 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -1,4 +1,5 @@ use crate::{ + frontend::RewritePlan, net::{ Prepare, messages::{Parse, RowDescription}, @@ -80,7 +81,11 @@ impl GlobalCache { } /// Insert a statement prepared using the simple protocol into the global cache. - pub(super) fn insert_prepare(&mut self, query: Bytes) -> (bool, Prepare) { + pub(super) fn insert_prepare( + &mut self, + query: Bytes, + rewrite_plan: &RewritePlan, + ) -> (bool, Prepare) { let cache_key = CacheKey::Simple { query: query.clone(), }; @@ -88,8 +93,9 @@ impl GlobalCache { if let Some(name) = self.reuse(&cache_key) { return ( false, - self.prepare(&name) - .expect("prepared to be in cache if reuse is true"), + self.prepare_and_rewrite(&name) + .expect("prepared to be in cache if reuse is true") + .0, ); } @@ -102,6 +108,7 @@ impl GlobalCache { let statement = Statement { stmt: StatementType::Prepare { prepare: prepare.clone(), + rewrite_plan: Arc::new(rewrite_plan.clone()), }, row_description: None, cache_key: cache_key.clone(), @@ -145,8 +152,8 @@ impl GlobalCache { } /// Get the [`Prepare`] message for a globally unique prepare statement name. - pub(crate) fn prepare(&self, name: &str) -> Option { - self.names.get(name).and_then(|p| p.prepare()) + pub(crate) fn prepare_and_rewrite(&self, name: &str) -> Option<(Prepare, Arc)> { + self.names.get(name).and_then(|p| p.prepare_and_rewrite()) } /// Get the rewritten Parse statement. @@ -156,7 +163,7 @@ impl GlobalCache { pub fn rewritten_parse(&self, name: &str) -> Option { self.names .get(name) - .and_then(|p| p.rewrite().clone().or(p.parse())) + .and_then(|p| p.rewritten_parse().clone().or(p.parse())) } /// Returns true if this prepared statement has been @@ -164,7 +171,7 @@ impl GlobalCache { pub(crate) fn is_rewritten(&self, name: &str) -> bool { self.names .get(name) - .map(|p| p.rewrite().is_some()) + .map(|p| p.rewritten_parse().is_some()) .unwrap_or_default() } @@ -362,8 +369,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()); - let (_, second) = cache.insert_prepare(query); + let (_, first) = cache.insert_prepare(query.clone(), &RewritePlan::default()); + let (_, second) = cache.insert_prepare(query, &RewritePlan::default()); 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 a4b05e948..b52a4b05a 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -8,6 +8,7 @@ use parking_lot::RwLock; use crate::{ config::PreparedStatements as PreparedStatementsLevel, + frontend::RewritePlan, net::{Parse, Prepare, ProtocolMessage}, }; @@ -119,8 +120,13 @@ impl PreparedStatements { /// /// Nothing, but the message is renamed to a unique, global name. /// - pub(crate) fn insert_prepare(&mut self, name: &str, query: Bytes) -> Prepare { - let (_new, prepare) = { self.global.write().insert_prepare(query) }; + pub(crate) fn insert_prepare( + &mut self, + name: &str, + query: Bytes, + rewrite_plan: &RewritePlan, + ) -> Prepare { + let (_new, prepare) = { self.global.write().insert_prepare(query, rewrite_plan) }; self.insert_internal(name, prepare.name()); @@ -141,10 +147,10 @@ impl PreparedStatements { } /// Get a globally unique [`Prepare`] message using the client name as key. - pub(crate) fn prepare(&self, name: &str) -> Option { + pub(crate) fn prepare_and_rewrite(&self, name: &str) -> Option<(Prepare, Arc)> { self.local .get(name) - .and_then(|name| self.global.read().prepare(name)) + .and_then(|name| self.global.read().prepare_and_rewrite(name)) } /// Number of prepared statements in the client's cache. diff --git a/pgdog/src/frontend/prepared_statements/statement.rs b/pgdog/src/frontend/prepared_statements/statement.rs index 21b7b69bd..f92d42ad8 100644 --- a/pgdog/src/frontend/prepared_statements/statement.rs +++ b/pgdog/src/frontend/prepared_statements/statement.rs @@ -1,4 +1,6 @@ -use crate::{net::Prepare, stats::memory::MemoryUsage}; +use std::sync::Arc; + +use crate::{frontend::RewritePlan, net::Prepare, stats::memory::MemoryUsage}; use super::prelude::*; @@ -18,13 +20,14 @@ pub(crate) enum StatementType { Prepare { prepare: Prepare, + rewrite_plan: Arc, }, } impl MemoryUsage for StatementType { fn memory_usage(&self) -> usize { match self { - Self::Prepare { prepare } => prepare.len(), + Self::Prepare { prepare, .. } => prepare.len(), Self::Parse { parse, rewrite } => { parse.len() + rewrite @@ -57,14 +60,17 @@ impl Statement { } } - pub(super) fn prepare(&self) -> Option { + pub(super) fn prepare_and_rewrite(&self) -> Option<(Prepare, Arc)> { match self.stmt { - StatementType::Prepare { ref prepare } => Some(prepare.clone()), + StatementType::Prepare { + ref prepare, + ref rewrite_plan, + } => Some((prepare.clone(), rewrite_plan.clone())), _ => None, } } - pub(crate) fn rewrite(&self) -> Option { + pub(crate) fn rewritten_parse(&self) -> Option { match self.stmt { StatementType::Parse { ref rewrite, .. } => rewrite.clone(), _ => None, @@ -74,7 +80,7 @@ impl Statement { pub(super) fn query(&self) -> &str { match self.stmt { StatementType::Parse { ref parse, .. } => parse.query(), - StatementType::Prepare { ref prepare } => prepare.query(), + StatementType::Prepare { ref prepare, .. } => prepare.query(), } } diff --git a/pgdog/src/frontend/router/mod.rs b/pgdog/src/frontend/router/mod.rs index 3b3d38299..0e2b2a6e4 100644 --- a/pgdog/src/frontend/router/mod.rs +++ b/pgdog/src/frontend/router/mod.rs @@ -14,7 +14,7 @@ pub use copy::CopyRow; pub use error::Error; use lazy_static::lazy_static; use parser::Shard; -pub use parser::{Ast, AstQuery, Command, QueryParser, Route, SetParam}; +pub use parser::{Ast, AstQuery, Command, QueryParser, RewritePlan, Route, SetParam}; use crate::frontend::router::parser::ShardWithPriority; diff --git a/pgdog/src/frontend/router/parser/mod.rs b/pgdog/src/frontend/router/parser/mod.rs index 5d3f79468..06e185c38 100644 --- a/pgdog/src/frontend/router/parser/mod.rs +++ b/pgdog/src/frontend/router/parser/mod.rs @@ -48,7 +48,9 @@ pub(crate) use limit::{Limit, LimitClause}; pub use order_by::OrderBy; pub(crate) use params::*; pub use query::QueryParser; -pub use rewrite::{Assignment, AssignmentValue, StatementRewrite, StatementRewriteContext}; +pub use rewrite::{ + Assignment, AssignmentValue, StatementRewrite, StatementRewriteContext, statement::RewritePlan, +}; pub use route::{Route, Shard, ShardWithPriority, ShardsWithPriority}; pub use schema::Schema; pub(crate) use sequence::Sequence; diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs index f570282e3..d4502d30b 100644 --- a/pgdog/src/frontend/router/parser/query/prepare.rs +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -51,9 +51,9 @@ impl QueryParser { // // INVARIANT 1: The prepared statements rewriter places it there. // INVARIANT 2: The rewriter renamed the EXECUTE statement to its global name. - let prepare = PreparedStatements::global() + let (prepare, _) = PreparedStatements::global() .read() - .prepare(stmt.name().expect("execute to have a name")) + .prepare_and_rewrite(stmt.name().expect("execute to have a name")) .ok_or(Error::ExecuteRequiresFull)?; // Make sure we never store unique statement names diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs b/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs index c56a074bf..f045f5b7d 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs @@ -200,6 +200,7 @@ mod tests { use crate::backend::schema::columns::StatsColumn as SchemaColumn; use crate::backend::schema::{Relation, Schema}; use crate::backend::{ShardedTables, ShardingSchema}; + use crate::config::PreparedStatements as PreparedStatementsLevel; use crate::frontend::PreparedStatements; use crate::frontend::router::parser::StatementRewriteContext; use crate::test_utils::set_env_var; @@ -488,14 +489,23 @@ mod tests { sql: &str, db_schema: &Schema, schema: &ShardingSchema, + ) -> Result<(String, RewritePlan), Error> { + let mut prepared = PreparedStatements::default(); + rewrite_sql_with_prepared_statements(sql, db_schema, schema, &mut prepared) + } + + fn rewrite_sql_with_prepared_statements( + sql: &str, + db_schema: &Schema, + schema: &ShardingSchema, + prepared: &mut PreparedStatements, ) -> Result<(String, RewritePlan), Error> { let _guard = set_env_var("NODE_ID", "pgdog-1"); let ast = pg_raw_parse::parse(sql).unwrap(); - let mut prepared = PreparedStatements::default(); let mut rewriter = StatementRewrite::new(StatementRewriteContext { extended: false, prepared: false, - prepared_statements: &mut prepared, + prepared_statements: prepared, schema, db_schema, user: "", @@ -511,6 +521,53 @@ mod tests { Ok((sql, plan)) } + #[test] + fn test_prepare_execute_rewrite_injects_auto_id() { + let db_schema = make_schema_with_bigint_pk(); + let schema = sharding_schema_with_mode(RewriteMode::Rewrite); + let mut prepared = PreparedStatements::default(); + prepared.set_level(PreparedStatementsLevel::Full); + + let (prepare_sql, prepare_plan) = rewrite_sql_with_prepared_statements( + "PREPARE stmt(text) AS INSERT INTO users (name) VALUES ($1)", + &db_schema, + &schema, + &mut prepared, + ) + .unwrap(); + + assert_eq!(prepare_plan.params, 1); + assert_eq!(prepare_plan.auto_id_injected, 1); + assert_eq!(prepare_plan.unique_ids, 1); + assert!(prepare_sql.contains("(name, id)")); + assert!(prepare_sql.contains("$2::bigint")); + + let (execute_sql, _) = rewrite_sql_with_prepared_statements( + "EXECUTE stmt('alice')", + &db_schema, + &schema, + &mut prepared, + ) + .unwrap(); + let ast = pg_raw_parse::parse(&execute_sql).unwrap(); + let Node::ExecuteStmt(execute) = ast.stmts().next().unwrap() else { + panic!("expected EXECUTE statement"); + }; + + assert_eq!(execute.params().len(), 2); + assert!(matches!( + execute.params().first(), + Some(Node::A_Const(value)) + if matches!(value.val(), Some(pg_raw_parse::ConstValue::String("alice"))) + )); + assert!(matches!( + execute.params().get(1), + Some(Node::A_Const(value)) + if matches!(value.val(), Some(pg_raw_parse::ConstValue::Float(id)) + if id.parse::().is_ok()) + )); + } + #[test] fn test_rewrite_omni_skips_sharded_table() { let db_schema = make_schema_with_bigint_pk(); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index 45d94f5f7..fceb960da 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -20,7 +20,7 @@ pub mod update; pub use error::Error; pub use insert::InsertSplit; -pub(crate) use plan::RewritePlan; +pub use plan::RewritePlan; pub(crate) use simple_prepared::PrepareExecute; pub(crate) use update::*; @@ -104,28 +104,41 @@ impl<'a> StatementRewrite<'a> { ) -> Result { let mut plan = RewritePlan::default(); - match stmt.stmt() { + let node = stmt.stmt(); + let parameterized_stmt = match node { Node::InsertStmt(_) | Node::SelectStmt(_) | Node::UpdateStmt(_) - | Node::DeleteStmt(_) => walk::walk(stmt.stmt(), |node| { - if let Node::ParamRef(param) = node { - plan.params = plan.params.max(param.number as u16) - } - }), - Node::PrepareStmt(_) => { + | Node::DeleteStmt(_) => Some(node), + Node::PrepareStmt(prepare) => { // Will use parameters for replacing args, not materialize values. self.extended = true; + Some(prepare.query()) } + Node::ExecuteStmt(_) | Node::ExplainStmt(_) => None, // We can't do anything with DDL statements _ => return Ok(plan), + }; + + if let Some(parameterized_stmt) = parameterized_stmt { + walk::walk(parameterized_stmt, |node| { + if let Node::ParamRef(param) = node { + plan.params = plan.params.max(param.number as u16) + } + }); } // Inject pgdog.unique_id() for missing BIGINT primary keys. // This must run BEFORE the unique_id rewriter so the injected // function calls get processed. - if let NodeMut::InsertStmt(insert) = stmt.stmt_mut() { - self.inject_auto_id(insert, mem, &mut plan)?; + match stmt.stmt_mut() { + NodeMut::InsertStmt(insert) => self.inject_auto_id(insert, mem, &mut plan)?, + NodeMut::PrepareStmt(mut prepare) => { + if let NodeMut::InsertStmt(insert) = prepare.query_mut() { + self.inject_auto_id(insert, mem, &mut plan)?; + } + } + _ => {} } // Track the next parameter number to use @@ -158,17 +171,17 @@ impl<'a> StatementRewrite<'a> { self.limit_offset(&select, &mut plan); } - if self.rewritten { - plan.stmt = Some(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned()); - } - // Handle top-level PREPARE/EXECUTE statements. - let prepared_result = self.rewrite_simple_prepared(stmt.stmt_mut(), mem)?; + let prepared_result = self.rewrite_simple_prepared(stmt.stmt_mut(), mem, &mut plan)?; if prepared_result.rewritten { self.rewritten = true; plan.prepare_rewrites = prepared_result.rewrites; } + if self.rewritten { + plan.stmt = Some(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned()); + } + if let Node::InsertStmt(insert) = stmt.stmt() { self.split_insert(insert, &mut plan)?; } 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 b1deee6a1..87454b188 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -1,12 +1,13 @@ use bytes::Bytes; -use pg_raw_parse::{NodeMut, make::MemoryToken}; +use pg_raw_parse::{ConstValue, NodeMut, make::MemoryToken, nodes::ExecuteStmtMut}; use crate::{ frontend::PreparedStatements, net::{PREPARE_TEMPLATE_NAME, Prepare}, + unique_id::UniqueId, }; -use super::{Error, StatementRewrite}; +use super::{Error, RewritePlan, StatementRewrite}; #[derive(Debug, Clone)] pub(crate) enum PrepareExecute { @@ -54,6 +55,7 @@ impl StatementRewrite<'_> { &mut self, node: NodeMut<'a, '_>, mem: MemoryToken<'a>, + plan: &mut RewritePlan, ) -> Result { let mut result = SimplePreparedResult::default(); @@ -61,7 +63,7 @@ impl StatementRewrite<'_> { return Ok(result); } - match rewrite_single_prepared(node, mem, self.prepared_statements)? { + match rewrite_single_prepared(node, mem, self.prepared_statements, plan)? { SimplePreparedRewrite::Prepared { prepare } => { result.rewrites.push(PrepareExecute::Prepare(prepare)); result.rewritten = true; @@ -82,6 +84,7 @@ fn rewrite_single_prepared<'a>( node: NodeMut<'a, '_>, mem: MemoryToken<'a>, prepared_statements: &mut PreparedStatements, + plan: &RewritePlan, ) -> Result { match node { NodeMut::PrepareStmt(mut stmt) => { @@ -92,7 +95,7 @@ fn rewrite_single_prepared<'a>( 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); + let prepare = prepared_statements.insert_prepare(&client_name, query, plan); stmt.set_name(Some(mem.copy_string(prepare.name()))); @@ -101,8 +104,12 @@ fn rewrite_single_prepared<'a>( NodeMut::ExecuteStmt(mut stmt) => { let stmt_name = stmt.name().expect("EXECUTE always has name"); - let prepare = prepared_statements.prepare(stmt_name); - if let Some(prepare) = prepare { + let prepare_and_rewrite = prepared_statements.prepare_and_rewrite(stmt_name); + if let Some((prepare, rewrite_plan)) = prepare_and_rewrite { + // Rewrite EXECUTE statement to match the rewrite + // we did on the PREPARE statement. + apply_prepare_rewrite_plan(&mut stmt, mem, &rewrite_plan)?; + stmt.set_name(Some(mem.copy_string(prepare.name()))); Ok(SimplePreparedRewrite::Executed { prepare }) } else { @@ -114,6 +121,23 @@ fn rewrite_single_prepared<'a>( } } +fn apply_prepare_rewrite_plan<'a>( + stmt: &mut ExecuteStmtMut<'a, '_>, + mem: MemoryToken<'a>, + plan: &RewritePlan, +) -> Result<(), Error> { + for _ in 0..plan.unique_ids { + let unique_id = UniqueId::generator()?.next_id(); + stmt.params_mut().push( + mem, + mem.make_a_const(ConstValue::Float(&unique_id.to_string())) + .uncast(), + ); + } + + Ok(()) +} + #[cfg(test)] mod tests { use super::super::{RewritePlan, StatementRewrite, StatementRewriteContext}; @@ -121,7 +145,10 @@ mod tests { use crate::backend::ShardingSchema; use crate::backend::schema::Schema; use crate::config::PreparedStatements as PreparedStatementsLevel; + use crate::test_utils::set_env_var; + use pg_raw_parse::Node; use pgdog_config::Rewrite; + use std::collections::HashSet; struct TestContext { ps: PreparedStatements, @@ -169,6 +196,72 @@ mod tests { } } + fn apply_plan(sql: &str, plan: &RewritePlan) -> Result { + let stmt = pg_raw_parse::parse(sql)?; + let ast = pg_raw_parse::make::try_owned(|mem| { + let mut copy = mem.make_unique(&*stmt.into_inner()); + let mut raw_stmt = copy + .as_mut() + .into_iter() + .next() + .expect("query must contain a statement"); + let NodeMut::ExecuteStmt(mut execute) = raw_stmt.stmt_mut() else { + panic!("expected EXECUTE statement"); + }; + + apply_prepare_rewrite_plan(&mut execute, mem, plan)?; + Ok::<_, Error>(copy) + })?; + + Ok(pg_raw_parse::deparse_stmts(&*ast)?) + } + + #[test] + fn test_apply_prepare_rewrite_plan_no_unique_ids() { + let sql = apply_plan("EXECUTE stmt(1, 'hello')", &RewritePlan::default()).unwrap(); + + assert_eq!(sql, "EXECUTE stmt(1, 'hello')"); + } + + #[test] + fn test_apply_prepare_rewrite_plan_appends_unique_ids() { + let _guard = set_env_var("NODE_ID", "pgdog-1"); + let plan = RewritePlan { + unique_ids: 3, + ..Default::default() + }; + let sql = apply_plan("EXECUTE stmt(42)", &plan).unwrap(); + let ast = pg_raw_parse::parse(&sql).unwrap(); + let Node::ExecuteStmt(execute) = ast.stmts().next().unwrap() else { + panic!("expected EXECUTE statement"); + }; + + assert_eq!(execute.params().len(), 4); + assert!(matches!( + execute.params().first(), + Some(Node::A_Const(value)) + if matches!(value.val(), Some(ConstValue::Integer(42))) + )); + + let ids: HashSet<_> = execute + .params() + .iter() + .skip(1) + .map(|param| { + let Node::A_Const(value) = param else { + panic!("expected unique ID to be a constant"); + }; + let Some(ConstValue::Float(value)) = value.val() else { + panic!("expected unique ID to be a numeric literal"); + }; + + value.parse::().expect("unique ID must be an i64") + }) + .collect(); + + assert_eq!(ids.len(), 3, "all appended IDs should be unique"); + } + #[test] fn test_rewrite_prepare() { let mut ctx = TestContext::new(); From 42634de276c4159a7534e0bcecf069021feee639 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Mon, 17 Aug 2026 16:07:28 -0700 Subject: [PATCH 22/22] test --- .../router/parser/query/test/test_prepared.rs | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/pgdog/src/frontend/router/parser/query/test/test_prepared.rs b/pgdog/src/frontend/router/parser/query/test/test_prepared.rs index 957e4a5b5..cab429a20 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_prepared.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_prepared.rs @@ -1,3 +1,5 @@ +use std::collections::HashSet; + use crate::frontend::router::parser::Shard; use super::setup::*; @@ -30,3 +32,23 @@ fn test_prepared_omnisharded_table() { assert!(command.route().is_omnisharded()); assert!(matches!(command.route().shard(), Shard::Direct(_))); } + +#[test] +fn test_prepare_uses_round_robin_across_calls() { + let mut test = QueryParserTest::new(); + let mut shards = HashSet::new(); + + for name in ["stmt_1", "stmt_2", "stmt_3", "stmt_4"] { + let command = test.execute(vec![ + Query::new(format!("PREPARE {name} AS SELECT 1")).into(), + ]); + + assert!(command.route().is_read()); + let Shard::Direct(shard) = command.route().shard() else { + panic!("PREPARE should route to a direct shard"); + }; + shards.insert(*shard); + } + + assert_eq!(shards.len(), 2); +}