Cosmetic fixes (#64)

* Cosmetic fixes

* fix test

Author: levkk

src/client.rs

@@ -1689,7 +1689,7 @@ where
         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 go removed)
+        // 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 Err(err) = server
@@ -1704,8 +1704,7 @@ where
     }
 
     /// Register and rewrite the parse statement to the clients statement cache
-    /// and also the pool's statement cache. Add it to extended protocol data
-    /// Returns True if the parse was registered, False otherwise
+    /// and also the pool's statement cache. Add it to extended protocol data.
     fn buffer_parse(&mut self, message: BytesMut, pool: &ConnectionPool) -> Result<(), Error> {
         // Avoid parsing if prepared statements not enabled
         let client_given_name = match self.prepared_statements_enabled {
@@ -1755,7 +1754,6 @@ where
 
     /// Rewrite the Bind (F) message to use the prepared statement name
     /// saved in the client cache.
-    /// Returns True if the bind is a named prepared statement, False otherwise
     async fn buffer_bind(&mut self, message: BytesMut) -> Result<(), Error> {
         // Avoid parsing if prepared statements not enabled
         let client_given_name = match self.prepared_statements_enabled {

src/messages.rs

@@ -1322,7 +1322,7 @@ impl Display for PgErrorMsg {
 }
 
 impl PgErrorMsg {
-    pub fn parse(error_msg: Vec<u8>) -> Result<PgErrorMsg, Error> {
+    pub fn parse(error_msg: &[u8]) -> Result<PgErrorMsg, Error> {
         let mut out = PgErrorMsg {
             severity_localized: "".to_string(),
             severity: "".to_string(),
@@ -1470,7 +1470,7 @@ mod tests {
 
         info!(
             "full message: {}",
-            PgErrorMsg::parse(complete_msg.clone()).unwrap()
+            PgErrorMsg::parse(&complete_msg).unwrap()
         );
         assert_eq!(
             PgErrorMsg {
@@ -1493,7 +1493,7 @@ mod tests {
                 line: Some(335),
                 routine: Some(routine_msg.to_string()),
             },
-            PgErrorMsg::parse(complete_msg).unwrap()
+            PgErrorMsg::parse(&complete_msg).unwrap()
         );
 
         let mut only_mandatory_msg = vec![];
@@ -1503,7 +1503,7 @@ mod tests {
         only_mandatory_msg.extend(field('M', message));
         only_mandatory_msg.extend(field('D', detail_msg));
 
-        let err_fields = PgErrorMsg::parse(only_mandatory_msg.clone()).unwrap();
+        let err_fields = PgErrorMsg::parse(&only_mandatory_msg).unwrap();
         info!("only mandatory fields: {}", &err_fields);
         error!(
             "server error: {}: {}",
@@ -1530,7 +1530,7 @@ mod tests {
                 line: None,
                 routine: None,
             },
-            PgErrorMsg::parse(only_mandatory_msg).unwrap()
+            PgErrorMsg::parse(&only_mandatory_msg).unwrap()
         );
     }
 }

src/server.rs

@@ -8,7 +8,6 @@ use once_cell::sync::Lazy;
 use parking_lot::{Mutex, RwLock};
 use postgres_protocol::message;
 use std::collections::{HashMap, HashSet};
-use std::io::Cursor;
 use std::mem;
 use std::net::IpAddr;
 use std::num::NonZeroUsize;
@@ -720,7 +719,7 @@ impl Server {
                                 }
                             };
 
-                            let fields = match PgErrorMsg::parse(error) {
+                            let fields = match PgErrorMsg::parse(&error) {
                                 Ok(f) => f,
                                 Err(err) => {
                                     return Err(err);
@@ -970,40 +969,18 @@ impl Server {
                     if self.in_copy_mode {
                         self.in_copy_mode = false;
                     }
-                    // TODO: consider logging a warning here
 
                     if self.prepared_statement_enabled {
-                        // It's probably okay to parse this since we don't expect to get ErrorResponses too frequently
-                        let mut cursor = Cursor::new(&message);
-
-                        loop {
-                            let code = cursor.get_u8();
-                            match cursor.read_string() {
-                                Ok(content) => {
-                                    println!(
-                                        "GOT CODE: {} with content: {}",
-                                        code as char, content
-                                    );
-                                    // This is allowed to be what looks like an infinite loop
-                                    // because the 'M' message is always present
-                                    if code as char == 'M' {
-                                        if content == "cached plan must not change result type" {
-                                            // This will still result in an error to the client, but this server connection
-                                            // will be dropped, and will not bleed into the pool.
-                                            // TODO: Other ideas to solve issues with DDL changes when using prepared statements
-                                            //  - Recreate connection pool to force recreation of server connections
-                                            //  - Just close the prepared statement instead of dropping the connection
-                                            //  - Implement a retry so the client doesn't see an error
-                                            self.mark_bad();
-                                        }
-                                        break;
-                                    }
-                                }
-                                Err(_) => {
-                                    warn!("Encountered an error while parsing ErrorResponse");
-                                    break;
-                                }
-                            }
+                        let error_message = crate::messages::PgErrorMsg::parse(&message)?;
+                        if error_message.message == "cached plan must not change result type" {
+                            warn!("Server {:?} changed schema, dropping connection to clean up prepared statements", self.address);
+                            // This will still result in an error to the client, but this server connection
+                            // will be dropped, and will not bleed into the pool.
+                            // TODO: Other ideas to solve issues with DDL changes when using prepared statements
+                            //  - Recreate connection pool to force recreation of server connections
+                            //  - Just close the prepared statement instead of dropping the connection
+                            //  - Implement a retry so the client doesn't see an error
+                            self.mark_bad();
                         }
                     }
                 }

tests/pgbench/simple.sql

@@ -10,10 +10,6 @@
 
 \set shard random(0, 2)
 
-SET SHARD TO :shard;
-
-SET SERVER ROLE TO 'auto';
-
 BEGIN;
 
 UPDATE pgbench_accounts SET abalance = abalance + :delta WHERE aid = :aid;
@@ -26,14 +22,6 @@ UPDATE pgbench_branches SET bbalance = bbalance + :delta WHERE bid = :bid;
 
 INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);
 
-END;
+SELECT * FROM pgbench_accounts;
 
-SET SHARDING KEY TO :aid;
-
--- Read load balancing
-SELECT abalance FROM pgbench_accounts WHERE aid = :aid;
-
-SET SERVER ROLE TO 'replica';
-
--- Read load balancing
-SELECT abalance FROM pgbench_accounts WHERE aid = :aid;
+END;