diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index cf3132398..f574ec87b 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -2809,6 +2809,35 @@ pub mod test { } } + /// Names Postgres itself thinks are prepared on this connection. + pub(crate) async fn prepared_in_postgres(server: &mut Server) -> Vec { + server + .fetch_all::("SELECT name FROM pg_prepared_statements") + .await + .unwrap() + } + + /// Run a named statement the way a client would, leaving it to the + /// connection to prepare it first if it isn't already. + pub(crate) async fn execute_prepared( + server: &mut Server, + name: &str, + param: &[u8], + ) -> Vec { + use crate::net::bind::Parameter; + + let request = ServerRequest { + messages: vec![ + Bind::new_params(name, &[Parameter::new(param)]).into(), + Execute::new().into(), + Sync::new().into(), + ], + expected: 1, + }; + + server.fetch_all::(request).await.unwrap() + } + #[tokio::test] async fn test_deallocate_all_clears_cache() { let mut server = test_server().await; diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 88d87e27f..9d50e84fc 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -238,17 +238,6 @@ impl GlobalCache { } } - /// Clear the global cache. Test-only: rolling the name counter back - /// would reuse global statement names. - #[cfg(test)] - pub fn reset(&mut self) { - self.statements.clear(); - self.names.clear(); - self.unused.clear(); - self.counter = 0; - self.versions = 0; - } - /// Get the query string stored in the global cache /// for the given globally unique prepared statement name. #[inline] diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index 7419cceb8..f41dacc83 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -194,10 +194,62 @@ pub fn run_maintenance() { #[cfg(test)] mod test { + use crate::backend::server::test::{execute_prepared, prepared_in_postgres, test_server}; use crate::net::messages::Bind; use super::*; + #[tokio::test] + async fn test_close_unused_does_not_reuse_names() { + let mut client = PreparedStatements::new(); + + let mut first = Parse::named("client_a", "SELECT $1::bigint"); + client.insert(&mut first); + let first_name = first.name(); + + let mut server = test_server().await; + assert_eq!(execute_prepared(&mut server, first_name, b"1").await, [1]); + assert_eq!(prepared_in_postgres(&mut server).await, [first_name]); + + client.close("client_a"); + PreparedStatements::global().write().close_unused(0); + assert!(PreparedStatements::global().read().is_empty()); + + assert_eq!(prepared_in_postgres(&mut server).await, [first_name]); + + let mut second = Parse::named("client_b", "SELECT $1::bigint + 100"); + client.insert(&mut second); + let second_name = second.name(); + assert_ne!(second_name, first_name); + + assert_eq!( + execute_prepared(&mut server, second_name, b"1").await, + [101] + ); + } + + #[tokio::test] + async fn test_close_unused_keeps_statements_clients_still_hold() { + let mut client = PreparedStatements::new(); + + let mut parse = Parse::named("client_a", "SELECT $1::bigint"); + client.insert(&mut parse); + let name = parse.name(); + + let mut warm = test_server().await; + assert_eq!(execute_prepared(&mut warm, name, b"1").await, [1]); + assert_eq!(prepared_in_postgres(&mut warm).await, [name]); + + PreparedStatements::global().write().close_unused(0); + + assert_eq!(execute_prepared(&mut warm, name, b"1").await, [1]); + + let mut cold = test_server().await; + assert!(prepared_in_postgres(&mut cold).await.is_empty()); + assert_eq!(execute_prepared(&mut cold, name, b"1").await, [1]); + assert_eq!(prepared_in_postgres(&mut cold).await, [name]); + } + #[test] fn test_maybe_rewrite() { let mut statements = PreparedStatements::default();