Update client cache name if pool renames prepared statement

Author: zainkabani

src/client.rs

@@ -1316,7 +1316,7 @@ where
                         {
                             match protocol_data {
                                 ExtendedProtocolData::Parse { data, metadata } => {
-                                    let (parse, hash) = match metadata {
+                                    let (client_given_name, parse, hash) = match metadata {
                                         Some(metadata) => metadata,
                                         None => {
                                             let first_char_in_name = *data.get(5).unwrap_or(&0);
@@ -1350,7 +1350,13 @@ where
                                         // TODO: Consider adding the close logic that this function can send for eviction to the client buffer instead
                                         // In this case we don't want to send the parse message to the server since the client is sending it
                                         self.register_parse_to_server_cache(
-                                            false, hash, &parse, &pool, server, &address,
+                                            false,
+                                            client_given_name,
+                                            hash,
+                                            parse,
+                                            &pool,
+                                            server,
+                                            &address,
                                         )
                                         .await?;
 
@@ -1651,24 +1657,32 @@ where
     /// Makes sure the the checked out server has the prepared statement and sends it to the server if it doesn't
     async fn ensure_prepared_statement_is_on_server(
         &mut self,
-        client_name: String,
+        client_given_name: String,
         pool: &ConnectionPool,
         server: &mut Server,
         address: &Address,
     ) -> Result<(), Error> {
-        match self.prepared_statements.get(&client_name) {
+        match self.prepared_statements.get(&client_given_name) {
             Some((parse, hash)) => {
                 debug!("Prepared statement `{}` found in cache", parse.name);
                 // In this case we want to send the parse message to the server
                 // since pgcat is initiating the prepared statement on this specific server
-                self.register_parse_to_server_cache(true, *hash, parse, pool, server, address)
-                    .await?;
+                self.register_parse_to_server_cache(
+                    true,
+                    client_given_name,
+                    *hash,
+                    parse.clone(),
+                    pool,
+                    server,
+                    address,
+                )
+                .await?;
             }
 
             None => {
                 return Err(Error::ClientError(format!(
                     "prepared statement `{}` not found",
-                    client_name
+                    client_given_name
                 )))
             }
         };
@@ -1679,21 +1693,34 @@ where
     /// Register the parse to the server cache and send it to the server if requested (ie. requested by pgcat)
     ///
     /// Also updates the pool LRU that this parse was used recently
+    #[allow(clippy::too_many_arguments)]
     async fn register_parse_to_server_cache(
-        &self,
+        &mut self,
         should_send_parse_to_server: bool,
+        client_given_name: String,
         hash: u64,
-        parse: &Arc<Parse>,
+        mut parse: Arc<Parse>,
         pool: &ConnectionPool,
         server: &mut Server,
         address: &Address,
     ) -> Result<(), Error> {
         // We want to update this in the LRU to know this was recently used and add it if it isn't there already
         // This could be the case if it was evicted or if doesn't exist (ie. we reloaded and it got removed)
-        pool.register_parse_to_cache(hash, parse);
+        if let Some(new_parse) = pool.register_parse_to_cache(hash, &parse) {
+            // If the pool has renamed this parse, we need to update the client cache with the new name
+            if parse.name != new_parse.name {
+                warn!(
+                    "Pool renamed prepared statement prepared statement `{}` to `{}` saving new name to client cache",
+                    parse.name, new_parse.name
+                );
+            }
+            self.prepared_statements
+                .insert(client_given_name.clone(), (new_parse.clone(), hash));
+            parse = new_parse;
+        };
 
         if let Err(err) = server
-            .register_prepared_statement(parse, should_send_parse_to_server)
+            .register_prepared_statement(&parse, should_send_parse_to_server)
             .await
         {
             pool.ban(address, BanReason::MessageSendFailed, Some(&self.stats));
@@ -1741,12 +1768,12 @@ where
         );
 
         self.prepared_statements
-            .insert(client_given_name, (new_parse.clone(), hash));
+            .insert(client_given_name.clone(), (new_parse.clone(), hash));
 
         self.extended_protocol_data_buffer
             .push_back(ExtendedProtocolData::create_new_parse(
                 new_parse.as_ref().try_into()?,
-                Some((new_parse.clone(), hash)),
+                Some((client_given_name, new_parse.clone(), hash)),
             ));
 
         Ok(())

src/messages.rs

@@ -770,7 +770,7 @@ impl BytesMutReader for BytesMut {
 pub enum ExtendedProtocolData {
     Parse {
         data: BytesMut,
-        metadata: Option<(Arc<Parse>, u64)>,
+        metadata: Option<(String, Arc<Parse>, u64)>,
     },
     Bind {
         data: BytesMut,
@@ -786,7 +786,7 @@ pub enum ExtendedProtocolData {
 }
 
 impl ExtendedProtocolData {
-    pub fn create_new_parse(data: BytesMut, metadata: Option<(Arc<Parse>, u64)>) -> Self {
+    pub fn create_new_parse(data: BytesMut, metadata: Option<(String, Arc<Parse>, u64)>) -> Self {
         Self::Parse { data, metadata }
     }