Enable transactional pool mode configuration

Author: pimeys

postgres/src/config.rs

@@ -433,6 +433,20 @@ impl Config {
         self
     }
 
+    /// When enabled, the client skips all internal caching for statements,
+    /// allowing usage with a connection pool with transaction mode.
+    ///
+    /// Defaults to `false`.
+    pub fn transaction_pool_mode(&mut self, enable: bool) -> &mut Config {
+        self.config.transaction_pool_mode(enable);
+        self
+    }
+
+    /// Gets the transaction pool mode status.
+    pub fn get_transaction_pool_mode(&self) -> bool {
+        self.config.get_transaction_pool_mode()
+    }
+
     /// Opens a connection to a PostgreSQL database.
     pub fn connect<T>(&self, tls: T) -> Result<Client, Error>
     where

tokio-postgres/src/client.rs

@@ -83,6 +83,7 @@ struct CachedTypeInfo {
 
 pub struct InnerClient {
     sender: mpsc::UnboundedSender<Request>,
+    transaction_pool_mode: bool,
     cached_typeinfo: Mutex<CachedTypeInfo>,
 
     /// A buffer to use when writing out postgres commands.
@@ -108,31 +109,39 @@ impl InnerClient {
     }
 
     pub fn set_typeinfo(&self, statement: &Statement) {
-        self.cached_typeinfo.lock().typeinfo = Some(statement.clone());
+        if !self.transaction_pool_mode {
+            self.cached_typeinfo.lock().typeinfo = Some(statement.clone());
+        }
     }
 
     pub fn typeinfo_composite(&self) -> Option<Statement> {
         self.cached_typeinfo.lock().typeinfo_composite.clone()
     }
 
     pub fn set_typeinfo_composite(&self, statement: &Statement) {
-        self.cached_typeinfo.lock().typeinfo_composite = Some(statement.clone());
+        if !self.transaction_pool_mode {
+            self.cached_typeinfo.lock().typeinfo_composite = Some(statement.clone());
+        }
     }
 
     pub fn typeinfo_enum(&self) -> Option<Statement> {
         self.cached_typeinfo.lock().typeinfo_enum.clone()
     }
 
     pub fn set_typeinfo_enum(&self, statement: &Statement) {
-        self.cached_typeinfo.lock().typeinfo_enum = Some(statement.clone());
+        if !self.transaction_pool_mode {
+            self.cached_typeinfo.lock().typeinfo_enum = Some(statement.clone());
+        }
     }
 
     pub fn type_(&self, oid: Oid) -> Option<Type> {
         self.cached_typeinfo.lock().types.get(&oid).cloned()
     }
 
     pub fn set_type(&self, oid: Oid, type_: &Type) {
-        self.cached_typeinfo.lock().types.insert(oid, type_.clone());
+        if !self.transaction_pool_mode {
+            self.cached_typeinfo.lock().types.insert(oid, type_.clone());
+        }
     }
 
     pub fn clear_type_cache(&self) {
@@ -190,12 +199,14 @@ impl Client {
         ssl_mode: SslMode,
         process_id: i32,
         secret_key: i32,
+        transaction_pool_mode: bool,
     ) -> Client {
         Client {
             inner: Arc::new(InnerClient {
                 sender,
                 cached_typeinfo: Default::default(),
                 buffer: Default::default(),
+                transaction_pool_mode,
             }),
             #[cfg(feature = "runtime")]
             socket_config: None,

tokio-postgres/src/config.rs

@@ -210,6 +210,7 @@ pub struct Config {
     pub(crate) channel_binding: ChannelBinding,
     pub(crate) load_balance_hosts: LoadBalanceHosts,
     pub(crate) max_backend_message_size: Option<usize>,
+    pub(crate) transaction_pool_mode: bool,
 }
 
 impl Default for Config {
@@ -244,6 +245,7 @@ impl Config {
             channel_binding: ChannelBinding::Prefer,
             load_balance_hosts: LoadBalanceHosts::Disable,
             max_backend_message_size: None,
+            transaction_pool_mode: false,
         }
     }
 
@@ -511,6 +513,20 @@ impl Config {
         self.channel_binding
     }
 
+    /// When enabled, the client skips all internal caching for statements,
+    /// allowing usage with a connection pool with transaction mode.
+    ///
+    /// Defaults to `false`.
+    pub fn transaction_pool_mode(&mut self, enable: bool) -> &mut Config {
+        self.transaction_pool_mode = enable;
+        self
+    }
+
+    /// Gets the transaction pool mode status.
+    pub fn get_transaction_pool_mode(&self) -> bool {
+        self.transaction_pool_mode
+    }
+
     /// Sets the host load balancing behavior.
     ///
     /// Defaults to `disable`.

tokio-postgres/src/connect_raw.rs

@@ -112,7 +112,15 @@ where
     let (process_id, secret_key, parameters) = read_info(&mut stream).await?;
 
     let (sender, receiver) = mpsc::unbounded();
-    let client = Client::new(sender, config.ssl_mode, process_id, secret_key);
+
+    let client = Client::new(
+        sender,
+        config.ssl_mode,
+        process_id,
+        secret_key,
+        config.transaction_pool_mode,
+    );
+
     let connection = Connection::new(stream.inner, stream.delayed, parameters, receiver);
 
     Ok((client, connection))