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/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/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/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 17665d0bd..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,6 +106,7 @@ pub struct PreparedStatements { config: PreparedStatementsConfig, memory_used: usize, oids: Arc, + server_state: State, } #[cfg(test)] @@ -126,6 +128,7 @@ impl PreparedStatements { config: PreparedStatementsConfig::default(), memory_used: 0, oids, + server_state: State::Idle, } } @@ -135,6 +138,10 @@ impl PreparedStatements { self.config = config; } + pub(super) fn set_server_state(&mut self, state: State) { + self.server_state = state; + } + /// Current prepared statement settings. pub fn config(&self) -> PreparedStatementsConfig { self.config @@ -272,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 { @@ -302,15 +311,43 @@ impl PreparedStatements { self.state.add('3'); } } - ProtocolMessage::Prepare { name, .. } => { - if self.contains(name) { + ProtocolMessage::PrepareFromClient(prepare) => { + use crate::net::{CommandComplete, ReadyForQuery}; + if self.contains(prepare.name()) { + 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( + if self.server_state == State::TransactionError { + ReadyForQuery::error() + } else { + ReadyForQuery::in_transaction( + self.server_state == State::IdleInTransaction, + ) + } + .message()?, + ); + return Ok(HandleResult::Drop); + } else { + self.parses.push_back(prepare.name().to_owned()); + self.state.add(ExecutionCode::ReadyForQuery); + } + } + ProtocolMessage::EnsurePrepared(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. 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 5ab6e9e0a..7fd103c9f 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) => { @@ -680,6 +680,9 @@ impl Server { _ => (), } + self.prepared_statements + .set_server_state(self.stats.get_state()); + trace!("{:#?} <<< [{}]", message, self.addr()); Ok(message) @@ -1326,7 +1329,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}, @@ -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, RewritePlan}, + net::{Prepare, *}, }; use super::{Error, *}; @@ -2314,20 +2319,20 @@ pub mod test { #[tokio::test] async fn test_manual_prepared() { + crate::logger(); 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"); + let name = "test"; + let query = Bytes::from("SELECT 1::bigint".to_owned()); + let prepare = prep.insert_prepare(name, query.clone(), &RewritePlan::default()); + assert_eq!(prepare.name(), "__pgdog_1"); server .send( - &vec![ProtocolMessage::from(Query::new(format!( - "PREPARE {} AS {}", - parse.name(), - parse.query() - )))] + &vec![ProtocolMessage::Query(Query::new( + "PREPARE __pgdog_1 AS SELECT 1::bigint", + ))] .into(), ) .await @@ -4362,6 +4367,126 @@ pub mod test { ); } + #[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); + } + } + } + + #[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_effective_max_age_default_is_base() { let server = Server::default(); 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/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/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 + ); } 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..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,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::PrepareFromClient(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 { name, statement } if name == "__pgdog_1" && statement == "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!( @@ -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::PrepareFromClient(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. 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..af49a84b9 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/sharded_prepared.rs @@ -0,0 +1,165 @@ +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; +} + +#[tokio::test] +async fn test_cross_shard_prepared() { + let mut client = TestClient::new_sharded(Parameters::default()) + .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)", + )) + .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/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/mod.rs b/pgdog/src/frontend/mod.rs index 284b777b0..4109448c9 100644 --- a/pgdog/src/frontend/mod.rs +++ b/pgdog/src/frontend/mod.rs @@ -22,10 +22,10 @@ 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; -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/cache_key.rs b/pgdog/src/frontend/prepared_statements/cache_key.rs new file mode 100644 index 000000000..fae1cfad0 --- /dev/null +++ b/pgdog/src/frontend/prepared_statements/cache_key.rs @@ -0,0 +1,53 @@ +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 }, +} + +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 { + /// 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, .. } => Ok(from_utf8(&query[0..query.len() - 1])?), + Self::Simple { query } => Ok(from_utf8(query)?), // Simple queries are regular Rust strings. + } + } +} + +#[cfg(test)] +mod test { + use super::*; + + 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/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..f448fa5c4 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -1,112 +1,17 @@ -use bytes::Bytes; - use crate::{ - net::messages::{Parse, RowDescription}, + frontend::RewritePlan, + 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 +50,78 @@ 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 cache_key = CacheKey::Extended { query: parse.query_ref(), data_types: parse.data_types_ref(), }; - if let Some(entry) = self.statements.get_mut(&parse_key) { - if entry.used == 0 { - self.unused.remove(&entry.counter); - } - entry.used += 1; + if let Some(name) = self.reuse(&cache_key) { + return (false, name); + } - (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, - }, - ); + 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, + rewrite: None, + }, + cache_key: cache_key.clone(), + row_description: None, + }; - self.names.insert( - name.clone(), - Statement { - parse, - cache_key, - rewrite: None, - row_description: None, - }, - ); + self.insert_internal(&name, cache_key, statement); - (true, name) - } + (true, 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; - - 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. + /// Insert a statement prepared using the simple protocol into the global cache. + pub(super) fn insert_prepare( + &mut self, + query: Bytes, + rewrite_plan: &RewritePlan, + ) -> (bool, Prepare) { let cache_key = CacheKey::Simple { - query: parse.query_ref(), - unique: self.counter, + query: query.clone(), }; - self.statements.insert( - cache_key.clone(), - CachedStmt { - counter: self.counter, - used: 1, - }, - ); + if let Some(name) = self.reuse(&cache_key) { + return ( + false, + self.prepare_and_rewrite(&name) + .expect("prepared to be in cache if reuse is true") + .0, + ); + } - self.names.insert( - name.clone(), - Statement { - parse, - cache_key, - rewrite: None, - row_description: None, + let name = self.next_name(); + let prepare = Prepare { + name: Bytes::from(name.clone()), + query, + }; + + let statement = Statement { + stmt: StatementType::Prepare { + prepare: prepare.clone(), + rewrite_plan: Arc::new(rewrite_plan.clone()), }, - ); + row_description: None, + cache_key: cache_key.clone(), + }; - name + self.insert_internal(&name, cache_key, statement); + (true, prepare) } /// 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 +148,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).and_then(|p| p.parse().clone()) + } + + /// Get the [`Prepare`] message for a globally unique prepare statement name. + 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. @@ -282,15 +163,15 @@ impl GlobalCache { pub fn rewritten_parse(&self, name: &str) -> Option { self.names .get(name) - .map(|p| p.rewrite.clone().unwrap_or(p.parse.clone())) + .and_then(|p| p.rewritten_parse().clone().or(p.parse())) } /// 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.rewritten_parse().is_some()) .unwrap_or_default() } @@ -298,22 +179,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 +211,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 +238,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 +255,33 @@ 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 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; + + Some(entry.name()) + } else { + 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); } } @@ -402,7 +314,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 +352,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 { @@ -453,28 +365,30 @@ mod test { #[test] fn test_simple_prepared_is_never_shared() { let mut cache = GlobalCache::default(); + + let query = Bytes::from("PREPARE __pgdog_template_name AS SELECT $1"); let parse = Parse::named("client_stmt", "SELECT $1"); - let first = cache.insert_prepare(&parse); - let second = cache.insert_prepare(&parse); + let (_, first) = cache.insert_prepare(query.clone(), &RewritePlan::default()); + let (_, second) = cache.insert_prepare(query, &RewritePlan::default()); - assert_ne!(first, second); - assert_eq!(cache.len(), 2); - assert_eq!(used(&cache, &first), 1); - assert_eq!(used(&cache, &second), 1); + assert_eq!(first, second); + assert_eq!(cache.len(), 1); + 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_eq!(cache.len(), 3); + assert_ne!(extended, first.name()); + assert_ne!(extended, second.name()); + 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/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..b52a4b05a 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 crate::util::safe_sleep; use crate::{ - config::{PreparedStatements as PreparedStatementsLevel, config}, - net::{Parse, ProtocolMessage}, + config::PreparedStatements as PreparedStatementsLevel, + frontend::RewritePlan, + 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>, @@ -49,94 +54,117 @@ impl Default for PreparedStatements { } 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(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()); + + prepare } - /// 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_and_rewrite(&self, name: &str) -> Option<(Prepare, Arc)> { + self.local + .get(name) + .and_then(|name| self.global.read().prepare_and_rewrite(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 +174,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. + pub(super) fn close_all(&mut self) { if !self.local.is_empty() { let mut global = self.global.write(); @@ -159,42 +190,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..f92d42ad8 --- /dev/null +++ b/pgdog/src/frontend/prepared_statements/statement.rs @@ -0,0 +1,99 @@ +use std::sync::Arc; + +use crate::{frontend::RewritePlan, 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, + rewrite_plan: Arc, + }, +} + +impl MemoryUsage for StatementType { + fn memory_usage(&self) -> usize { + match self { + Self::Prepare { prepare, .. } => prepare.len(), + Self::Parse { parse, rewrite } => { + parse.len() + + rewrite + .as_ref() + .map(|rewrite| rewrite.len()) + .unwrap_or_default() + } + } + } +} + +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_and_rewrite(&self) -> Option<(Prepare, Arc)> { + match self.stmt { + StatementType::Prepare { + ref prepare, + ref rewrite_plan, + } => Some((prepare.clone(), rewrite_plan.clone())), + _ => None, + } + } + + pub(crate) fn rewritten_parse(&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) { + if let StatementType::Parse { + ref mut rewrite, .. + } = self.stmt + { + *rewrite = Some(parse.clone()) + } + } +} diff --git a/pgdog/src/frontend/router/context.rs b/pgdog/src/frontend/router/context.rs index df99253ef..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::Ast, - router::sharding::ResolvedLookups, + router::{Ast, parser::StatementParameters, sharding::ResolvedLookups}, }, - net::{Bind, Parameters}, + net::Parameters, }; -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct RouterContext<'a> { /// Bound parameters to the query. - pub bind: Option<&'a Bind>, + 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> { @@ -53,7 +48,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 { @@ -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/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/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..abe5eda6f 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,30 @@ impl Cache { Ok(entry) } + pub(crate) fn record(&self, query: &str) -> Result { + { + let mut guard = self.inner.lock(); + 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 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 +224,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..8fd361a71 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -111,4 +111,10 @@ pub enum Error { #[error("unmapped sharding key was specified")] UnmappedShardKey(String), + + #[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/limit.rs b/pgdog/src/frontend/router/parser/limit.rs index e8c43ff62..568ae3f77 100644 --- a/pgdog/src/frontend/router/parser/limit.rs +++ b/pgdog/src/frontend/router/parser/limit.rs @@ -3,8 +3,7 @@ use pg_raw_parse::{ nodes::{self, SelectStmt}, }; -use super::Error; -use crate::net::Bind; +use super::{Error, StatementParameters}; #[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 Bind>, + bind: Option>, } 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>) -> Self { Self { stmt, bind } } diff --git a/pgdog/src/frontend/router/parser/mod.rs b/pgdog/src/frontend/router/parser/mod.rs index e890795dc..06e185c38 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,8 +46,11 @@ use function::Function; pub use key::Key; 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/params.rs b/pgdog/src/frontend/router/parser/params.rs new file mode 100644 index 000000000..bc497999a --- /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(crate) struct ExecuteParams { + params: Vec, + format: Vec, +} + +impl ExecuteParams { + pub(crate) 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<'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) -> &'a [BindParameter] { + match self { + Self::Bind(bind) => bind.params_raw(), + Self::Execute(params) => ¶ms.params, + } + } + + pub(super) fn format_codes_raw(self) -> &'a [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..f1b379f30 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}, *, @@ -25,6 +25,7 @@ mod ddl; mod delete; mod explain; mod plugins; +mod prepare; mod select; mod set; mod set_config; @@ -348,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( @@ -430,6 +431,10 @@ impl QueryParser { )); } + Node::PrepareStmt(stmt) => self.prepare(stmt, context), + + Node::ExecuteStmt(stmt) => self.execute(stmt, context), + Node::ExplainStmt(stmt) => self.explain(&statement, stmt, context), Node::DiscardStmt { .. } => { @@ -450,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 @@ -519,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 new file mode 100644 index 000000000..d4502d30b --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -0,0 +1,80 @@ +use pg_raw_parse::nodes::{ExecuteStmt, PrepareStmt}; + +use crate::{ + frontend::{BufferedQuery, PreparedStatements}, + net::{PREPARE_TEMPLATE_NAME, Query}, +}; + +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(context.shards)).into(), + )), + _ => Route::write(ShardWithPriority::new_rr_not_executable( + (round_robin::next(context.shards)).into(), + )), + }; + + Ok(Command::Query(route)) + } + + /// An `EXECUTE` statement. + pub(super) fn execute( + &mut self, + stmt: &ExecuteStmt, + context: &mut QueryParserContext<'_>, + ) -> Result { + // 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 1: The prepared statements rewriter places it there. + // INVARIANT 2: The rewriter renamed the EXECUTE statement to its global name. + let (prepare, _) = PreparedStatements::global() + .read() + .prepare_and_rewrite(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() { + Some(Node::PrepareStmt(stmt)) => Some(stmt.query()), + stmt => stmt, + }; + + 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), + Some(Node::DeleteStmt(stmt)) => self.delete(stmt.into(), &mut context), + _ => Ok(Command::Query(Route::write( + ShardWithPriority::new_default_unset(Shard::All), + ))), + } + } +} diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs index 4d54d744e..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). @@ -258,9 +261,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/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/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..cab429a20 --- /dev/null +++ b/pgdog/src/frontend/router/parser/query/test/test_prepared.rs @@ -0,0 +1,54 @@ +use std::collections::HashSet; + +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(_))); +} + +#[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); +} 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 8ef1105d2..fceb960da 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -20,8 +20,8 @@ pub mod update; pub use error::Error; pub use insert::InsertSplit; -pub(crate) use plan::RewritePlan; -pub use simple_prepared::SimplePreparedResult; +pub use plan::RewritePlan; +pub(crate) use simple_prepared::PrepareExecute; pub(crate) use update::*; /// Statement rewrite engine context. @@ -104,32 +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::ExecuteStmt(_) | Node::ExplainStmt(_) => {} + | 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), - } + }; - // 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.prepares = prepared_result.prepares; + 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 @@ -162,6 +171,13 @@ impl<'a> StatementRewrite<'a> { self.limit_offset(&select, &mut plan); } + // Handle top-level PREPARE/EXECUTE statements. + 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()); } 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/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 6b2bbb9e8..b58958163 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -5,7 +5,9 @@ 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<(String, String)>, + 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,16 +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(|(name, statement)| ProtocolMessage::Prepare { - name: name.clone(), - statement: statement.clone(), - }) - .collect(); - request.messages.splice(0..0, prepends); + .for_each(|prepare| match prepare { + PrepareExecute::Prepare(prepare) => { + request.messages.clear(); + request.push(ProtocolMessage::PrepareFromClient(prepare.clone())); + } + PrepareExecute::Execute(prepare) => { + request + .messages + .splice(0..0, vec![ProtocolMessage::EnsurePrepared(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 34a854e7a..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,17 +1,30 @@ -use pg_raw_parse::{NodeMut, make::MemoryToken}; +use bytes::Bytes; +use pg_raw_parse::{ConstValue, NodeMut, make::MemoryToken, nodes::ExecuteStmtMut}; -use crate::frontend::PreparedStatements; -use crate::net::Parse; +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 { + /// 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<(String, String)>, + pub(crate) rewrites: Vec, } /// Result of rewriting a single PREPARE or EXECUTE SQL command. @@ -20,10 +33,10 @@ 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 { name: String, statement: String }, + Executed { prepare: Prepare }, } impl StatementRewrite<'_> { @@ -42,6 +55,7 @@ impl StatementRewrite<'_> { &mut self, node: NodeMut<'a, '_>, mem: MemoryToken<'a>, + plan: &mut RewritePlan, ) -> Result { let mut result = SimplePreparedResult::default(); @@ -49,12 +63,13 @@ impl StatementRewrite<'_> { return Ok(result); } - match rewrite_single_prepared(node, mem, self.prepared_statements)? { - SimplePreparedRewrite::Prepared => { + match rewrite_single_prepared(node, mem, self.prepared_statements, plan)? { + SimplePreparedRewrite::Prepared { prepare } => { + result.rewrites.push(PrepareExecute::Prepare(prepare)); result.rewritten = true; } - SimplePreparedRewrite::Executed { name, statement } => { - result.prepares.push((name, statement)); + SimplePreparedRewrite::Executed { prepare } => { + result.rewrites.push(PrepareExecute::Execute(prepare)); result.rewritten = true; } SimplePreparedRewrite::None => {} @@ -69,33 +84,34 @@ fn rewrite_single_prepared<'a>( node: NodeMut<'a, '_>, mem: MemoryToken<'a>, prepared_statements: &mut PreparedStatements, + plan: &RewritePlan, ) -> Result { match node { NodeMut::PrepareStmt(mut stmt) => { - let query = pg_raw_parse::deparse(stmt.query())?; + let client_name = stmt.name().expect("prepare must have a name").to_owned(); + + // Create a globally unique key using the query text + // with a hardcoded name. + stmt.set_name(Some(mem.copy_string(PREPARE_TEMPLATE_NAME))); + let query = Bytes::from(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned()); + + let prepare = prepared_statements.insert_prepare(&client_name, query, plan); - let mut parse = Parse::named( - stmt.name().expect("PREPARE always has a name"), - query.as_str(), - ); - prepared_statements.insert_prepare(&mut parse); - stmt.set_name(Some(mem.copy_string(parse.name()))); + stmt.set_name(Some(mem.copy_string(prepare.name()))); - Ok(SimplePreparedRewrite::Prepared) + Ok(SimplePreparedRewrite::Prepared { prepare }) } 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_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 { Err(Error::ExecuteMissingPrepare(stmt_name.to_owned())) } @@ -105,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}; @@ -112,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, @@ -160,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(); @@ -173,8 +275,21 @@ mod tests { !sql.contains("test_stmt"), "original name should be replaced: {sql}" ); - assert!(plan.prepares.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] @@ -187,11 +302,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.prepare_rewrites[0]; + match prepare { + PrepareExecute::Execute(prepare) => { + assert!(prepare.name().starts_with("__pgdog_")); + assert_eq!(prepare.query(), "PREPARE __pgdog_template_name AS SELECT 1"); + } - let (name, statement) = &plan.prepares[0]; - assert!(name.starts_with("__pgdog_")); - assert_eq!(statement, "SELECT 1"); + _ => panic!("expected PrepareExecute::Execute"), + } } #[test] @@ -208,7 +329,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] @@ -224,7 +345,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/frontend/router/parser/statement.rs b/pgdog/src/frontend/router/parser/statement.rs index 957a660e8..39cb8305e 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>) -> Option { match node { Node::A_Const(a) => a.val()?.numeric_value(), Node::TypeCast(c) => integer_arg(c.arg(), bind), @@ -232,7 +233,7 @@ use crate::{ ShardedTable, Tables, lookup, }, }, - net::{Bind, messages::Format, parameter::ParameterValue}, + net::{messages::Format, parameter::ParameterValue}, }; use pgdog_config::LookupResult; @@ -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 { @@ -439,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 { @@ -1089,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) } @@ -1168,7 +1167,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() } @@ -1181,7 +1180,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()) } @@ -2148,7 +2147,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() } @@ -2260,7 +2259,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() } @@ -2416,7 +2415,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() } @@ -2619,7 +2618,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/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 } 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 } diff --git a/pgdog/src/net/messages/mod.rs b/pgdog/src/net/messages/mod.rs index 733fd3e75..1bc78ad47 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_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 new file mode 100644 index 000000000..2578496bb --- /dev/null +++ b/pgdog/src/net/messages/prepare.rs @@ -0,0 +1,74 @@ +//! 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_unchecked; + +pub static PREPARE_TEMPLATE_NAME: &str = "__pgdog_template_name"; + +use super::prelude::*; +use super::{Message, Query}; + +#[derive(Debug, Clone, PartialEq)] +pub struct Prepare { + pub(crate) name: Bytes, + pub(crate) query: Bytes, +} + +impl Prepare { + pub fn query(&self) -> &str { + // SAFETY: We only support UTF-8. + unsafe { from_utf8_unchecked(&self.query) } + } + + pub fn name(&self) -> &str { + // SAFETY: We only support UTF-8. + unsafe { from_utf8_unchecked(&self.name) } + } + + 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 { + 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 = 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(query).to_bytes() + } +} diff --git a/pgdog/src/net/protocol_message.rs b/pgdog/src/net/protocol_message.rs index 955d5d12f..544f1e8f0 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,8 @@ pub enum ProtocolMessage { Bind(Bind), Parse(Parse), Describe(Describe), - Prepare { name: String, statement: String }, + EnsurePrepared(Prepare), + PrepareFromClient(Prepare), Execute(Execute), Close(Close), Query(Query), @@ -59,7 +62,8 @@ 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::EnsurePrepared(prepare) => prepare.len(), + Self::PrepareFromClient(prepare) => prepare.len(), Self::Execute(execute) => execute.len(), Self::Close(close) => close.len(), Self::Query(query) => query.len(), @@ -79,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::EnsurePrepared { .. } | Self::PrepareFromClient { .. } => 'Q', Self::Execute(execute) => execute.code(), Self::Close(close) => close.code(), Self::Query(query) => query.code(), @@ -119,9 +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 { statement, name } => { - Query::new(format!("PREPARE {} AS {}", name, statement)).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(), 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());