Support for prepared statements (#474)

* Start prepared statements

* parse

* Ok

* optional

* dont rewrite anonymous prepared stmts

* Dont rewrite anonymous prep statements

* hm?

* prep statements

* I see!

* comment

* Print config value

* Rewrite bind and add sqlx test

* fmt

* ok

* Fix

* Fix stats

* its late

* clean up PREPARE

dev/Dockerfile

@@ -1,4 +1,4 @@
-FROM rust:bullseye
+FROM rust:1.70-bullseye
 
 # Dependencies
 RUN apt-get update -y \

pgcat.toml

@@ -60,6 +60,9 @@ tcp_keepalives_count = 5
 # Number of seconds between keepalive packets.
 tcp_keepalives_interval = 5
 
+# Handle prepared statements.
+prepared_statements = true
+
 # Path to TLS Certificate file to use for TLS connections
 # tls_certificate = ".circleci/server.cert"
 # Path to TLS private key file to use for TLS connections

src/admin.rs

@@ -699,6 +699,8 @@ where
         ("bytes_sent", DataType::Numeric),
         ("bytes_received", DataType::Numeric),
         ("age_seconds", DataType::Numeric),
+        ("prepare_cache_hit", DataType::Numeric),
+        ("prepare_cache_miss", DataType::Numeric),
     ];
 
     let new_map = get_server_stats();
@@ -722,6 +724,14 @@ where
                 .duration_since(server.connect_time())
                 .as_secs()
                 .to_string(),
+            server
+                .prepared_hit_count
+                .load(Ordering::Relaxed)
+                .to_string(),
+            server
+                .prepared_miss_count
+                .load(Ordering::Relaxed)
+                .to_string(),
         ];
 
         res.put(data_row(&row));

src/client.rs

@@ -3,8 +3,9 @@ use crate::pool::BanReason;
 /// Handle clients by pretending to be a PostgreSQL server.
 use bytes::{Buf, BufMut, BytesMut};
 use log::{debug, error, info, trace, warn};
+use once_cell::sync::Lazy;
 use std::collections::HashMap;
-use std::sync::Arc;
+use std::sync::{atomic::AtomicUsize, Arc};
 use std::time::Instant;
 use tokio::io::{split, AsyncReadExt, BufReader, ReadHalf, WriteHalf};
 use tokio::net::TcpStream;
@@ -13,7 +14,9 @@ use tokio::sync::mpsc::Sender;
 
 use crate::admin::{generate_server_info_for_admin, handle_admin};
 use crate::auth_passthrough::refetch_auth_hash;
-use crate::config::{get_config, get_idle_client_in_transaction_timeout, Address, PoolMode};
+use crate::config::{
+    get_config, get_idle_client_in_transaction_timeout, get_prepared_statements, Address, PoolMode,
+};
 use crate::constants::*;
 use crate::messages::*;
 use crate::plugins::PluginOutput;
@@ -25,6 +28,11 @@ use crate::tls::Tls;
 
 use tokio_rustls::server::TlsStream;
 
+/// Incrementally count prepared statements
+/// to avoid random conflicts in places where the random number generator is weak.
+pub static PREPARED_STATEMENT_COUNTER: Lazy<Arc<AtomicUsize>> =
+    Lazy::new(|| Arc::new(AtomicUsize::new(0)));
+
 /// Type of connection received from client.
 enum ClientConnectionType {
     Startup,
@@ -93,6 +101,9 @@ pub struct Client<S, T> {
 
     /// Used to notify clients about an impending shutdown
     shutdown: Receiver<()>,
+
+    /// Prepared statements
+    prepared_statements: HashMap<String, Parse>,
 }
 
 /// Client entrypoint.
@@ -682,6 +693,7 @@ where
             application_name: application_name.to_string(),
             shutdown,
             connected_to_server: false,
+            prepared_statements: HashMap::new(),
         })
     }
 
@@ -716,6 +728,7 @@ where
             application_name: String::from("undefined"),
             shutdown,
             connected_to_server: false,
+            prepared_statements: HashMap::new(),
         })
     }
 
@@ -757,6 +770,10 @@ where
         // Result returned by one of the plugins.
         let mut plugin_output = None;
 
+        // Prepared statement being executed
+        let mut prepared_statement = None;
+        let mut will_prepare = false;
+
         // Our custom protocol loop.
         // We expect the client to either start a transaction with regular queries
         // or issue commands for our sharding and server selection protocol.
@@ -766,13 +783,16 @@ where
                 self.transaction_mode
             );
 
+            // Should we rewrite prepared statements and bind messages?
+            let mut prepared_statements_enabled = get_prepared_statements();
+
             // Read a complete message from the client, which normally would be
             // either a `Q` (query) or `P` (prepare, extended protocol).
             // We can parse it here before grabbing a server from the pool,
             // in case the client is sending some custom protocol messages, e.g.
             // SET SHARDING KEY TO 'bigint';
 
-            let message = tokio::select! {
+            let mut message = tokio::select! {
                 _ = self.shutdown.recv() => {
                     if !self.admin {
                         error_response_terminal(
@@ -800,7 +820,21 @@ where
                 // allocate a connection, we wouldn't be able to send back an error message
                 // to the client so we buffer them and defer the decision to error out or not
                 // to when we get the S message
-                'D' | 'E' => {
+                'D' => {
+                    if prepared_statements_enabled {
+                        let name;
+                        (name, message) = self.rewrite_describe(message).await?;
+
+                        if let Some(name) = name {
+                            prepared_statement = Some(name);
+                        }
+                    }
+
+                    self.buffer.put(&message[..]);
+                    continue;
+                }
+
+                'E' => {
                     self.buffer.put(&message[..]);
                     continue;
                 }
@@ -830,6 +864,11 @@ where
                 }
 
                 'P' => {
+                    if prepared_statements_enabled {
+                        (prepared_statement, message) = self.rewrite_parse(message)?;
+                        will_prepare = true;
+                    }
+
                     self.buffer.put(&message[..]);
 
                     if query_router.query_parser_enabled() {
@@ -846,6 +885,10 @@ where
                 }
 
                 'B' => {
+                    if prepared_statements_enabled {
+                        (prepared_statement, message) = self.rewrite_bind(message).await?;
+                    }
+
                     self.buffer.put(&message[..]);
 
                     if query_router.query_parser_enabled() {
@@ -1054,7 +1097,48 @@ where
             // If the client is in session mode, no more custom protocol
             // commands will be accepted.
             loop {
-                let message = match initial_message {
+                // Only check if we should rewrite prepared statements
+                // in session mode. In transaction mode, we check at the beginning of
+                // each transaction.
+                if !self.transaction_mode {
+                    prepared_statements_enabled = get_prepared_statements();
+                }
+
+                debug!("Prepared statement active: {:?}", prepared_statement);
+
+                // We are processing a prepared statement.
+                if let Some(ref name) = prepared_statement {
+                    debug!("Checking prepared statement is on server");
+                    // Get the prepared statement the server expects to see.
+                    let statement = match self.prepared_statements.get(name) {
+                        Some(statement) => {
+                            debug!("Prepared statement `{}` found in cache", name);
+                            statement
+                        }
+                        None => {
+                            return Err(Error::ClientError(format!(
+                                "prepared statement `{}` not found",
+                                name
+                            )))
+                        }
+                    };
+
+                    // Since it's already in the buffer, we don't need to prepare it on this server.
+                    if will_prepare {
+                        server.will_prepare(&statement.name);
+                        will_prepare = false;
+                    } else {
+                        // The statement is not prepared on the server, so we need to prepare it.
+                        if server.should_prepare(&statement.name) {
+                            server.prepare(statement).await?;
+                        }
+                    }
+
+                    // Done processing the prepared statement.
+                    prepared_statement = None;
+                }
+
+                let mut message = match initial_message {
                     None => {
                         trace!("Waiting for message inside transaction or in session mode");
 
@@ -1173,6 +1257,11 @@ where
                     // Parse
                     // The query with placeholders is here, e.g. `SELECT * FROM users WHERE email = $1 AND active = $2`.
                     'P' => {
+                        if prepared_statements_enabled {
+                            (prepared_statement, message) = self.rewrite_parse(message)?;
+                            will_prepare = true;
+                        }
+
                         if query_router.query_parser_enabled() {
                             if let Ok(ast) = QueryRouter::parse(&message) {
                                 if let Ok(output) = query_router.execute_plugins(&ast).await {
@@ -1187,12 +1276,25 @@ where
                     // Bind
                     // The placeholder's replacements are here, e.g. 'user@email.com' and 'true'
                     'B' => {
+                        if prepared_statements_enabled {
+                            (prepared_statement, message) = self.rewrite_bind(message).await?;
+                        }
+
                         self.buffer.put(&message[..]);
                     }
 
                     // Describe
                     // Command a client can issue to describe a previously prepared named statement.
                     'D' => {
+                        if prepared_statements_enabled {
+                            let name;
+                            (name, message) = self.rewrite_describe(message).await?;
+
+                            if let Some(name) = name {
+                                prepared_statement = Some(name);
+                            }
+                        }
+
                         self.buffer.put(&message[..]);
                     }
 
@@ -1235,7 +1337,7 @@ where
                         let first_message_code = (*self.buffer.get(0).unwrap_or(&0)) as char;
 
                         // Almost certainly true
-                        if first_message_code == 'P' {
+                        if first_message_code == 'P' && !prepared_statements_enabled {
                             // Message layout
                             // P followed by 32 int followed by null-terminated statement name
                             // So message code should be in offset 0 of the buffer, first character
@@ -1363,6 +1465,107 @@ where
         }
     }
 
+    /// Rewrite Parse (F) message to set the prepared statement name to one we control.
+    /// Save it into the client cache.
+    fn rewrite_parse(&mut self, message: BytesMut) -> Result<(Option<String>, BytesMut), Error> {
+        let parse: Parse = (&message).try_into()?;
+
+        let name = parse.name.clone();
+
+        // Don't rewrite anonymous prepared statements
+        if parse.anonymous() {
+            debug!("Anonymous prepared statement");
+            return Ok((None, message));
+        }
+
+        let parse = parse.rename();
+
+        debug!(
+            "Renamed prepared statement `{}` to `{}` and saved to cache",
+            name, parse.name
+        );
+
+        self.prepared_statements.insert(name.clone(), parse.clone());
+
+        Ok((Some(name), parse.try_into()?))
+    }
+
+    /// Rewrite the Bind (F) message to use the prepared statement name
+    /// saved in the client cache.
+    async fn rewrite_bind(
+        &mut self,
+        message: BytesMut,
+    ) -> Result<(Option<String>, BytesMut), Error> {
+        let bind: Bind = (&message).try_into()?;
+        let name = bind.prepared_statement.clone();
+
+        if bind.anonymous() {
+            debug!("Anonymous bind message");
+            return Ok((None, message));
+        }
+
+        match self.prepared_statements.get(&name) {
+            Some(prepared_stmt) => {
+                let bind = bind.reassign(prepared_stmt);
+
+                debug!("Rewrote bind `{}` to `{}`", name, bind.prepared_statement);
+
+                Ok((Some(name), bind.try_into()?))
+            }
+            None => {
+                debug!("Got bind for unknown prepared statement {:?}", bind);
+
+                error_response(
+                    &mut self.write,
+                    &format!(
+                        "prepared statement \"{}\" does not exist",
+                        bind.prepared_statement
+                    ),
+                )
+                .await?;
+
+                Err(Error::ClientError(format!(
+                    "Prepared statement `{}` doesn't exist",
+                    name
+                )))
+            }
+        }
+    }
+
+    /// Rewrite the Describe (F) message to use the prepared statement name
+    /// saved in the client cache.
+    async fn rewrite_describe(
+        &mut self,
+        message: BytesMut,
+    ) -> Result<(Option<String>, BytesMut), Error> {
+        let describe: Describe = (&message).try_into()?;
+        let name = describe.statement_name.clone();
+
+        if describe.anonymous() {
+            debug!("Anonymous describe");
+            return Ok((None, message));
+        }
+
+        match self.prepared_statements.get(&name) {
+            Some(prepared_stmt) => {
+                let describe = describe.rename(&prepared_stmt.name);
+
+                debug!(
+                    "Rewrote describe `{}` to `{}`",
+                    name, describe.statement_name
+                );
+
+                Ok((Some(name), describe.try_into()?))
+            }
+
+            None => {
+                debug!("Got describe for unknown prepared statement {:?}", describe);
+
+                Ok((None, message))
+            }
+        }
+    }
+
     /// Release the server from the client: it can't cancel its queries anymore.
     pub fn release(&self) {
         let mut guard = self.client_server_map.lock();

src/config.rs

@@ -320,6 +320,9 @@ pub struct General {
     pub auth_query: Option<String>,
     pub auth_query_user: Option<String>,
     pub auth_query_password: Option<String>,
+
+    #[serde(default)]
+    pub prepared_statements: bool,
 }
 
 impl General {
@@ -434,6 +437,7 @@ impl Default for General {
             server_lifetime: Self::default_server_lifetime(),
             server_round_robin: false,
             validate_config: true,
+            prepared_statements: false,
         }
     }
 }
@@ -1015,6 +1019,7 @@ impl Config {
             "Server TLS certificate verification: {}",
             self.general.verify_server_certificate
         );
+        info!("Prepared statements: {}", self.general.prepared_statements);
         info!(
             "Plugins: {}",
             match self.plugins {
@@ -1239,6 +1244,10 @@ pub fn get_idle_client_in_transaction_timeout() -> u64 {
         .idle_client_in_transaction_timeout
 }
 
+pub fn get_prepared_statements() -> bool {
+    (*(*CONFIG.load())).general.prepared_statements
+}
+
 /// Parse the configuration file located at the path.
 pub async fn parse(path: &str) -> Result<(), Error> {
     let mut contents = String::new();

src/errors.rs

@@ -26,6 +26,7 @@ pub enum Error {
     AuthPassthroughError(String),
     UnsupportedStatement,
     QueryRouterParserError(String),
+    QueryRouterError(String),
 }
 
 #[derive(Clone, PartialEq, Debug)]
@@ -121,3 +122,9 @@ impl std::fmt::Display for Error {
         }
     }
 }
+
+impl From<std::ffi::NulError> for Error {
+    fn from(err: std::ffi::NulError) -> Self {
+        Error::QueryRouterError(err.to_string())
+    }
+}

src/messages.rs

@@ -7,11 +7,15 @@ use socket2::{SockRef, TcpKeepalive};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpStream;
 
+use crate::client::PREPARED_STATEMENT_COUNTER;
 use crate::config::get_config;
 use crate::errors::Error;
+
 use std::collections::HashMap;
+use std::ffi::CString;
 use std::io::{BufRead, Cursor};
 use std::mem;
+use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 /// Postgres data type mappings
@@ -526,6 +530,13 @@ pub fn command_complete(command: &str) -> BytesMut {
     res
 }
 
+pub fn flush() -> BytesMut {
+    let mut bytes = BytesMut::new();
+    bytes.put_u8(b'H');
+    bytes.put_i32(4);
+    bytes
+}
+
 /// Write all data in the buffer to the TcpStream.
 pub async fn write_all<S>(stream: &mut S, buf: BytesMut) -> Result<(), Error>
 where
@@ -689,3 +700,285 @@ impl BytesMutReader for Cursor<&BytesMut> {
         }
     }
 }
+
+/// Parse (F) message.
+/// See: <https://www.postgresql.org/docs/current/protocol-message-formats.html>
+#[derive(Clone, Debug)]
+pub struct Parse {
+    code: char,
+    #[allow(dead_code)]
+    len: i32,
+    pub name: String,
+    pub generated_name: String,
+    query: String,
+    num_params: i16,
+    param_types: Vec<i32>,
+}
+
+impl TryFrom<&BytesMut> for Parse {
+    type Error = Error;
+
+    fn try_from(buf: &BytesMut) -> Result<Parse, Error> {
+        let mut cursor = Cursor::new(buf);
+        let code = cursor.get_u8() as char;
+        let len = cursor.get_i32();
+        let name = cursor.read_string()?;
+        let query = cursor.read_string()?;
+        let num_params = cursor.get_i16();
+        let mut param_types = Vec::new();
+
+        for _ in 0..num_params {
+            param_types.push(cursor.get_i32());
+        }
+
+        Ok(Parse {
+            code,
+            len,
+            name,
+            generated_name: prepared_statement_name(),
+            query,
+            num_params,
+            param_types,
+        })
+    }
+}
+
+impl TryFrom<Parse> for BytesMut {
+    type Error = Error;
+
+    fn try_from(parse: Parse) -> Result<BytesMut, Error> {
+        let mut bytes = BytesMut::new();
+
+        let name_binding = CString::new(parse.name)?;
+        let name = name_binding.as_bytes_with_nul();
+
+        let query_binding = CString::new(parse.query)?;
+        let query = query_binding.as_bytes_with_nul();
+
+        // Recompute length of the message.
+        let len = 4 // self
+            + name.len()
+            + query.len()
+            + 2
+            + 4 * parse.num_params as usize;
+
+        bytes.put_u8(parse.code as u8);
+        bytes.put_i32(len as i32);
+        bytes.put_slice(name);
+        bytes.put_slice(query);
+        bytes.put_i16(parse.num_params);
+        for param in parse.param_types {
+            bytes.put_i32(param);
+        }
+
+        Ok(bytes)
+    }
+}
+
+impl TryFrom<&Parse> for BytesMut {
+    type Error = Error;
+
+    fn try_from(parse: &Parse) -> Result<BytesMut, Error> {
+        parse.clone().try_into()
+    }
+}
+
+impl Parse {
+    pub fn rename(mut self) -> Self {
+        self.name = self.generated_name.to_string();
+        self
+    }
+
+    pub fn anonymous(&self) -> bool {
+        self.name.is_empty()
+    }
+}
+
+/// Bind (B) message.
+/// See: <https://www.postgresql.org/docs/current/protocol-message-formats.html>
+#[derive(Clone, Debug)]
+pub struct Bind {
+    code: char,
+    #[allow(dead_code)]
+    len: i64,
+    portal: String,
+    pub prepared_statement: String,
+    num_param_format_codes: i16,
+    param_format_codes: Vec<i16>,
+    num_param_values: i16,
+    param_values: Vec<(i32, BytesMut)>,
+    num_result_column_format_codes: i16,
+    result_columns_format_codes: Vec<i16>,
+}
+
+impl TryFrom<&BytesMut> for Bind {
+    type Error = Error;
+
+    fn try_from(buf: &BytesMut) -> Result<Bind, Error> {
+        let mut cursor = Cursor::new(buf);
+        let code = cursor.get_u8() as char;
+        let len = cursor.get_i32();
+        let portal = cursor.read_string()?;
+        let prepared_statement = cursor.read_string()?;
+        let num_param_format_codes = cursor.get_i16();
+        let mut param_format_codes = Vec::new();
+
+        for _ in 0..num_param_format_codes {
+            param_format_codes.push(cursor.get_i16());
+        }
+
+        let num_param_values = cursor.get_i16();
+        let mut param_values = Vec::new();
+
+        for _ in 0..num_param_values {
+            let param_len = cursor.get_i32();
+            let mut param = BytesMut::with_capacity(param_len as usize);
+            param.resize(param_len as usize, b'0');
+            cursor.copy_to_slice(&mut param);
+            param_values.push((param_len, param));
+        }
+
+        let num_result_column_format_codes = cursor.get_i16();
+        let mut result_columns_format_codes = Vec::new();
+
+        for _ in 0..num_result_column_format_codes {
+            result_columns_format_codes.push(cursor.get_i16());
+        }
+
+        Ok(Bind {
+            code,
+            len: len as i64,
+            portal,
+            prepared_statement,
+            num_param_format_codes,
+            param_format_codes,
+            num_param_values,
+            param_values,
+            num_result_column_format_codes,
+            result_columns_format_codes,
+        })
+    }
+}
+
+impl TryFrom<Bind> for BytesMut {
+    type Error = Error;
+
+    fn try_from(bind: Bind) -> Result<BytesMut, Error> {
+        let mut bytes = BytesMut::new();
+
+        let portal_binding = CString::new(bind.portal)?;
+        let portal = portal_binding.as_bytes_with_nul();
+
+        let prepared_statement_binding = CString::new(bind.prepared_statement)?;
+        let prepared_statement = prepared_statement_binding.as_bytes_with_nul();
+
+        let mut len = 4 // self
+            + portal.len()
+            + prepared_statement.len()
+            + 2 // num_param_format_codes
+            + 2 * bind.num_param_format_codes as usize // num_param_format_codes
+            + 2; // num_param_values
+
+        for (param_len, _) in &bind.param_values {
+            len += 4 + *param_len as usize;
+        }
+        len += 2; // num_result_column_format_codes
+        len += 2 * bind.num_result_column_format_codes as usize;
+
+        bytes.put_u8(bind.code as u8);
+        bytes.put_i32(len as i32);
+        bytes.put_slice(portal);
+        bytes.put_slice(prepared_statement);
+        bytes.put_i16(bind.num_param_format_codes);
+        for param_format_code in bind.param_format_codes {
+            bytes.put_i16(param_format_code);
+        }
+        bytes.put_i16(bind.num_param_values);
+        for (param_len, param) in bind.param_values {
+            bytes.put_i32(param_len);
+            bytes.put_slice(&param);
+        }
+        bytes.put_i16(bind.num_result_column_format_codes);
+        for result_column_format_code in bind.result_columns_format_codes {
+            bytes.put_i16(result_column_format_code);
+        }
+
+        Ok(bytes)
+    }
+}
+
+impl Bind {
+    pub fn reassign(mut self, parse: &Parse) -> Self {
+        self.prepared_statement = parse.name.clone();
+        self
+    }
+
+    pub fn anonymous(&self) -> bool {
+        self.prepared_statement.is_empty()
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Describe {
+    code: char,
+
+    #[allow(dead_code)]
+    len: i32,
+    target: char,
+    pub statement_name: String,
+}
+
+impl TryFrom<&BytesMut> for Describe {
+    type Error = Error;
+
+    fn try_from(bytes: &BytesMut) -> Result<Describe, Error> {
+        let mut cursor = Cursor::new(bytes);
+        let code = cursor.get_u8() as char;
+        let len = cursor.get_i32();
+        let target = cursor.get_u8() as char;
+        let statement_name = cursor.read_string()?;
+
+        Ok(Describe {
+            code,
+            len,
+            target,
+            statement_name,
+        })
+    }
+}
+
+impl TryFrom<Describe> for BytesMut {
+    type Error = Error;
+
+    fn try_from(describe: Describe) -> Result<BytesMut, Error> {
+        let mut bytes = BytesMut::new();
+        let statement_name_binding = CString::new(describe.statement_name)?;
+        let statement_name = statement_name_binding.as_bytes_with_nul();
+        let len = 4 + 1 + statement_name.len();
+
+        bytes.put_u8(describe.code as u8);
+        bytes.put_i32(len as i32);
+        bytes.put_u8(describe.target as u8);
+        bytes.put_slice(statement_name);
+
+        Ok(bytes)
+    }
+}
+
+impl Describe {
+    pub fn rename(mut self, name: &str) -> Self {
+        self.statement_name = name.to_string();
+        self
+    }
+
+    pub fn anonymous(&self) -> bool {
+        self.statement_name.is_empty()
+    }
+}
+
+pub fn prepared_statement_name() -> String {
+    format!(
+        "P_{}",
+        PREPARED_STATEMENT_COUNTER.fetch_add(1, Ordering::SeqCst)
+    )
+}

src/query_router.rs

@@ -331,7 +331,7 @@ impl QueryRouter {
         Some((command, value))
     }
 
-    pub fn parse(message: &BytesMut) -> Result<Vec<sqlparser::ast::Statement>, Error> {
+    pub fn parse(message: &BytesMut) -> Result<Vec<Statement>, Error> {
         let mut message_cursor = Cursor::new(message);
 
         let code = message_cursor.get_u8() as char;
@@ -348,12 +348,13 @@ impl QueryRouter {
             // Parse (prepared statement)
             'P' => {
                 // Reads statement name
-                message_cursor.read_string().unwrap();
+                let _name = message_cursor.read_string().unwrap();
 
                 // Reads query string
                 let query = message_cursor.read_string().unwrap();
 
                 debug!("Prepared statement: '{}'", query);
+
                 query
             }
 

src/server.rs

@@ -5,7 +5,7 @@ use fallible_iterator::FallibleIterator;
 use log::{debug, error, info, trace, warn};
 use parking_lot::{Mutex, RwLock};
 use postgres_protocol::message;
-use std::collections::HashMap;
+use std::collections::{BTreeSet, HashMap};
 use std::io::Read;
 use std::net::IpAddr;
 use std::sync::Arc;
@@ -198,6 +198,9 @@ pub struct Server {
 
     /// Should clean up dirty connections?
     cleanup_connections: bool,
+
+    /// Prepared statements
+    prepared_statements: BTreeSet<String>,
 }
 
 impl Server {
@@ -692,6 +695,7 @@ impl Server {
                             )),
                         },
                         cleanup_connections,
+                        prepared_statements: BTreeSet::new(),
                     };
 
                     server.set_name("pgcat").await?;
@@ -910,6 +914,43 @@ impl Server {
         Ok(bytes)
     }
 
+    pub fn will_prepare(&mut self, name: &str) {
+        debug!("Will prepare `{}`", name);
+
+        self.prepared_statements.insert(name.to_string());
+    }
+
+    pub fn should_prepare(&self, name: &str) -> bool {
+        let should_prepare = !self.prepared_statements.contains(name);
+
+        debug!("Should prepare `{}`: {}", name, should_prepare);
+
+        if should_prepare {
+            self.stats.prepared_cache_miss();
+        } else {
+            self.stats.prepared_cache_hit();
+        }
+
+        should_prepare
+    }
+
+    pub async fn prepare(&mut self, parse: &Parse) -> Result<(), Error> {
+        debug!("Preparing `{}`", parse.name);
+
+        let bytes: BytesMut = parse.try_into()?;
+        self.send(&bytes).await?;
+        self.send(&flush()).await?;
+
+        // Read and discard ParseComplete (B)
+        let _ = read_message(&mut self.stream).await?;
+
+        self.prepared_statements.insert(parse.name.to_string());
+
+        debug!("Prepared `{}`", parse.name);
+
+        Ok(())
+    }
+
     /// If the server is still inside a transaction.
     /// If the client disconnects while the server is in a transaction, we will clean it up.
     pub fn in_transaction(&self) -> bool {

src/stats/server.rs

@@ -47,6 +47,8 @@ pub struct ServerStats {
     pub transaction_count: Arc<AtomicU64>,
     pub query_count: Arc<AtomicU64>,
     pub error_count: Arc<AtomicU64>,
+    pub prepared_hit_count: Arc<AtomicU64>,
+    pub prepared_miss_count: Arc<AtomicU64>,
 }
 
 impl Default for ServerStats {
@@ -63,6 +65,8 @@ impl Default for ServerStats {
             query_count: Arc::new(AtomicU64::new(0)),
             error_count: Arc::new(AtomicU64::new(0)),
             reporter: get_reporter(),
+            prepared_hit_count: Arc::new(AtomicU64::new(0)),
+            prepared_miss_count: Arc::new(AtomicU64::new(0)),
         }
     }
 }
@@ -172,6 +176,7 @@ impl ServerStats {
         self.set_application(application_name.to_string());
         self.address.stats.query_count_add();
         self.address.stats.query_time_add(milliseconds);
+        self.query_count.fetch_add(1, Ordering::Relaxed);
     }
 
     /// Report a transaction executed by a client a server
@@ -198,4 +203,14 @@ impl ServerStats {
             .fetch_add(amount_bytes as u64, Ordering::Relaxed);
         self.address.stats.bytes_received_add(amount_bytes as u64);
     }
+
+    /// Report a prepared statement that already exists on the server.
+    pub fn prepared_cache_hit(&self) {
+        self.prepared_hit_count.fetch_add(1, Ordering::Relaxed);
+    }
+
+    /// Report a prepared statement that does not exist on the server yet.
+    pub fn prepared_cache_miss(&self) {
+        self.prepared_miss_count.fetch_add(1, Ordering::Relaxed);
+    }
 }

tests/rust/.gitignore

@@ -0,0 +1 @@
+target/

tests/rust/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "rust"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+sqlx = { version = "0.6.2", features = [ "runtime-tokio-rustls", "postgres", "json", "tls", "migrate", "time", "uuid", "ipnetwork"] }
+tokio = { version = "1", features = ["full"] }

tests/rust/src/main.rs

@@ -0,0 +1,29 @@
+#[tokio::main]
+async fn main() {
+    test_prepared_statements().await;
+}
+
+async fn test_prepared_statements() {
+    let pool = sqlx::postgres::PgPoolOptions::new()
+        .max_connections(5)
+        .connect("postgres://sharding_user:sharding_user@127.0.0.1:6432/sharded_db")
+        .await
+        .unwrap();
+
+    let mut handles = Vec::new();
+
+    for _ in 0..5 {
+        let pool = pool.clone();
+        let handle = tokio::task::spawn(async move {
+            for _ in 0..1000 {
+                sqlx::query("SELECT 1").fetch_all(&pool).await.unwrap();
+            }
+        });
+
+        handles.push(handle);
+    }
+
+    for handle in handles {
+        handle.await.unwrap();
+    }
+}