Add statement mode (#612)

levkk · web-flow · commit c46f1dd0c045 · 2025-11-17T13:13:31.000-08:00
diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs
@@ -302,6 +302,21 @@ impl Config {
             }
         }
 
+        // Check pooler mode.
+        let mut pooler_mode = HashMap::<String, Option<PoolerMode>>::new();
+        for database in &self.databases {
+            if let Some(mode) = pooler_mode.get(&database.name) {
+                if mode != &database.pooler_mode {
+                    warn!(
+                        "database \"{}\" (shard={}, role={}) has a different \"pooler_mode\" setting, ignoring",
+                        database.name, database.shard, database.role,
+                    );
+                }
+            } else {
+                pooler_mode.insert(database.name.clone(), database.pooler_mode.clone());
+            }
+        }
+
         // Check that idle_healthcheck_interval is shorter than ban_timeout.
         if self.general.ban_timeout > 0
             && self.general.idle_healthcheck_interval >= self.general.ban_timeout
@@ -316,12 +331,12 @@ impl Config {
         match self.general.passthrough_auth {
             PassthoughAuth::Enabled if !self.general.tls_client_required => {
                 warn!(
-                    "consider setting tls_client_required while passthrough_auth is enabled to prevent clients from exposing plaintext passwords"
+                    "consider setting \"tls_client_required\" while \"passthrough_auth\" is enabled to prevent clients from exposing plaintext passwords"
                 );
             }
             PassthoughAuth::EnabledPlain => {
                 warn!(
-                    "passthrough_auth plain is enabled - network traffic may expose plaintext passwords"
+                    "\"passthrough_auth\" is set to \"plain\", network traffic may expose plaintext passwords"
                 )
             }
             _ => (),
diff --git a/pgdog-config/src/pooling.rs b/pgdog-config/src/pooling.rs
@@ -49,13 +49,15 @@ pub enum PoolerMode {
     #[default]
     Transaction,
     Session,
+    Statement,
 }
 
 impl std::fmt::Display for PoolerMode {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Transaction => write!(f, "transaction"),
             Self::Session => write!(f, "session"),
+            Self::Statement => write!(f, "statement"),
         }
     }
 }
diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs
@@ -85,6 +85,21 @@ pub struct ClusterShardConfig {
     pub replicas: Vec<PoolConfig>,
 }
 
+impl ClusterShardConfig {
+    pub fn pooler_mode(&self) -> PoolerMode {
+        // One of these will exist.
+
+        if let Some(ref primary) = self.primary {
+            return primary.config.pooler_mode;
+        }
+
+        self.replicas
+            .first()
+            .map(|replica| replica.config.pooler_mode)
+            .unwrap_or_default()
+    }
+}
+
 /// Cluster creation config.
 pub struct ClusterConfig<'a> {
     pub name: &'a str,
@@ -122,12 +137,17 @@ impl<'a> ClusterConfig<'a> {
         sharded_schemas: ShardedSchemas,
         rewrite: &'a Rewrite,
     ) -> Self {
+        let pooler_mode = shards
+            .first()
+            .map(|shard| shard.pooler_mode())
+            .unwrap_or(user.pooler_mode.unwrap_or(general.pooler_mode));
+
         Self {
             name: &user.database,
             password: user.password(),
             user: &user.name,
             replication_sharding: user.replication_sharding.clone(),
-            pooler_mode: user.pooler_mode.unwrap_or(general.pooler_mode),
+            pooler_mode,
             lb_strategy: general.load_balancing_strategy,
             shards,
             sharded_tables,
diff --git a/pgdog/src/backend/pool/connection/mod.rs b/pgdog/src/backend/pool/connection/mod.rs
@@ -402,18 +402,17 @@ impl Connection {
         self.cluster.as_ref().ok_or(Error::NotConnected)
     }
 
-    /// Transaction mode pooling.
+    /// Pooler is in session mode.
     #[inline]
-    pub(crate) fn transaction_mode(&self) -> bool {
+    pub(crate) fn session_mode(&self) -> bool {
         self.cluster()
-            .map(|c| c.pooler_mode() == PoolerMode::Transaction)
+            .map(|c| c.pooler_mode() == PoolerMode::Session)
             .unwrap_or(true)
     }
 
-    /// Pooler is in session mod
     #[inline]
-    pub(crate) fn session_mode(&self) -> bool {
-        !self.transaction_mode()
+    pub(crate) fn pooler_mode(&self) -> PoolerMode {
+        self.cluster().map(|c| c.pooler_mode()).unwrap_or_default()
     }
 
     /// This is an admin DB connection.
diff --git a/pgdog/src/config/mod.rs b/pgdog/src/config/mod.rs
@@ -144,19 +144,26 @@ pub fn overrides(overrides: Overrides) {
 // Test helper functions
 #[cfg(test)]
 pub fn load_test() {
+    load_test_with_pooler_mode(PoolerMode::Transaction)
+}
+
+#[cfg(test)]
+pub fn load_test_with_pooler_mode(pooler_mode: PoolerMode) {
     use crate::backend::databases::init;
 
     let mut config = ConfigAndUsers::default();
     config.config.databases = vec![Database {
         name: "pgdog".into(),
         host: "127.0.0.1".into(),
         port: 5432,
+        pooler_mode: Some(pooler_mode),
         ..Default::default()
     }];
     config.users.users = vec![User {
         name: "pgdog".into(),
         database: "pgdog".into(),
         password: Some("pgdog".into()),
+        pooler_mode: Some(pooler_mode),
         ..Default::default()
     }];
 
diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs
@@ -9,7 +9,7 @@ use crate::{
     state::State,
 };
 
-use tracing::debug;
+use tracing::{debug, error};
 
 use super::*;
 
@@ -264,7 +264,7 @@ impl QueryEngine {
 
             // Release the connection back into the pool before flushing data to client.
             // Flushing can take a minute and we don't want to block the connection from being reused.
-            if self.backend.transaction_mode() && context.requests_left == 0 {
+            if !self.backend.session_mode() && context.requests_left == 0 {
                 self.backend.disconnect();
             }
 
@@ -312,18 +312,14 @@ impl QueryEngine {
             && !context.admin
             && context.client_request.executable()
         {
-            let bytes_sent = context
-                .stream
-                .error(
-                    ErrorResponse::cross_shard_disabled(),
-                    context.in_transaction(),
-                )
-                .await?;
-            self.stats.sent(bytes_sent);
+            let error = ErrorResponse::cross_shard_disabled();
+
+            self.error_response(context, error).await?;
 
             if self.backend.connected() && self.backend.done() {
                 self.backend.disconnect();
             }
+
             Ok(false)
         } else {
             Ok(true)
@@ -367,19 +363,31 @@ impl QueryEngine {
         {
             let error = ErrorResponse::in_failed_transaction();
 
-            self.hooks.on_engine_error(context, &error)?;
-
-            let bytes_sent = context
-                .stream
-                .error(error, context.in_transaction())
-                .await?;
-            self.stats.sent(bytes_sent);
+            self.error_response(context, error).await?;
 
             Ok(false)
         } else {
             Ok(true)
         }
     }
+
+    pub(super) async fn error_response(
+        &mut self,
+        context: &mut QueryEngineContext<'_>,
+        error: ErrorResponse,
+    ) -> Result<(), Error> {
+        error!("{:?} [{:?}]", error.message, context.stream.peer_addr());
+
+        self.hooks.on_engine_error(context, &error)?;
+
+        let bytes_sent = context
+            .stream
+            .error(error, context.in_transaction())
+            .await?;
+        self.stats.sent(bytes_sent);
+
+        return Ok(());
+    }
 }
 
 #[derive(Debug, Default, Clone)]
diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs
@@ -1,4 +1,5 @@
-use tracing::{error, trace};
+use pgdog_config::PoolerMode;
+use tracing::trace;
 
 use super::*;
 
@@ -7,6 +8,14 @@ impl QueryEngine {
         &mut self,
         context: &mut QueryEngineContext<'_>,
     ) -> Result<bool, Error> {
+        // Check that we can route this transaction at all.
+        if self.backend.pooler_mode() == PoolerMode::Statement && context.client_request.is_begin()
+        {
+            self.error_response(context, ErrorResponse::transaction_statement_mode())
+                .await?;
+            return Ok(false);
+        }
+
         // Admin doesn't have a cluster.
         let cluster = if let Ok(cluster) = self.backend.cluster() {
             if !context.in_transaction() && !cluster.online() {
@@ -17,18 +26,14 @@ impl QueryEngine {
 
                 match self.backend.cluster() {
                     Ok(cluster) => cluster,
-                    Err(err) => {
+                    Err(_) => {
                         // Cluster is gone.
-                        error!("{:?} [{:?}]", err, context.stream.peer_addr());
-                        let error =
-                            ErrorResponse::connection(&identifier.user, &identifier.database);
-                        self.hooks.on_engine_error(context, &error)?;
+                        self.error_response(
+                            context,
+                            ErrorResponse::connection(&identifier.user, &identifier.database),
+                        )
+                        .await?;
 
-                        let bytes_sent = context
-                            .stream
-                            .error(error, context.in_transaction())
-                            .await?;
-                        self.stats.sent(bytes_sent);
                         return Ok(false);
                     }
                 }
@@ -55,17 +60,9 @@ impl QueryEngine {
                 );
             }
             Err(err) => {
-                error!("{:?} [{:?}]", err, context.stream.peer_addr());
-
-                let error = ErrorResponse::syntax(err.to_string().as_str());
-
-                self.hooks.on_engine_error(context, &error)?;
-
-                let bytes_sent = context
-                    .stream
-                    .error(error, context.in_transaction())
+                self.error_response(context, ErrorResponse::syntax(err.to_string().as_str()))
                     .await?;
-                self.stats.sent(bytes_sent);
+
                 return Ok(false);
             }
         }
diff --git a/pgdog/src/frontend/client/test/mod.rs b/pgdog/src/frontend/client/test/mod.rs
@@ -1,5 +1,6 @@
 use std::time::{Duration, Instant};
 
+use pgdog_config::PoolerMode;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     net::{TcpListener, TcpStream},
@@ -10,7 +11,10 @@ use bytes::{Buf, BufMut, BytesMut};
 
 use crate::{
     backend::databases::databases,
-    config::{config, load_test, load_test_replicas, set, PreparedStatements, Role},
+    config::{
+        config, load_test, load_test_replicas, load_test_with_pooler_mode, set, PreparedStatements,
+        Role,
+    },
     frontend::{
         client::{BufferEvent, QueryEngine},
         prepared_statements, Client,
@@ -739,6 +743,25 @@ async fn test_client_login_timeout() {
     handle.await.unwrap().unwrap();
 }
 
+#[tokio::test]
+async fn test_statement_mode() {
+    crate::logger();
+
+    load_test_with_pooler_mode(PoolerMode::Statement);
+    let (mut conn, mut client) = parallel_test_client().await;
+
+    let _ = tokio::spawn(async move {
+        client.run().await.unwrap();
+    });
+
+    let req = buffer!({ Query::new("BEGIN") });
+    conn.write_all(&req).await.unwrap();
+
+    let msgs = read!(conn, ['E', 'Z']);
+    let error = ErrorResponse::from_bytes(msgs[0].clone().freeze()).unwrap();
+    assert_eq!(error.code, "58000");
+}
+
 #[tokio::test]
 async fn test_client_login_timeout_does_not_affect_queries() {
     crate::logger();
diff --git a/pgdog/src/frontend/client_request.rs b/pgdog/src/frontend/client_request.rs
@@ -2,6 +2,7 @@
 use std::ops::{Deref, DerefMut};
 
 use lazy_static::lazy_static;
+use regex::Regex;
 
 use crate::{
     frontend::router::parser::RewritePlan,
@@ -108,6 +109,26 @@ impl ClientRequest {
         Ok(None)
     }
 
+    pub fn is_begin(&self) -> bool {
+        lazy_static! {
+            static ref BEGIN: Regex = Regex::new("(?i)^BEGIN").unwrap();
+        }
+
+        for message in &self.messages {
+            let query = match message {
+                ProtocolMessage::Parse(parse) => parse.query(),
+                ProtocolMessage::Query(query) => query.query(),
+                _ => continue,
+            };
+
+            if BEGIN.is_match(query) {
+                return true;
+            }
+        }
+
+        false
+    }
+
     /// If this buffer contains bound parameters, retrieve them.
     pub fn parameters(&self) -> Result<Option<&Bind>, Error> {
         for message in &self.messages {
@@ -462,4 +483,16 @@ mod test {
         assert_eq!(second_slice[2].code(), 'H'); // Flush
         assert_eq!(second_slice[3].code(), 'S'); // Sync
     }
+
+    #[test]
+    fn test_detect_begin() {
+        for query in [
+            ProtocolMessage::Query(Query::new("begin")),
+            ProtocolMessage::Query(Query::new("BEGIN WORK REPEATABLE READ")),
+            ProtocolMessage::Parse(Parse::new_anonymous("BEGIN")),
+        ] {
+            let req = ClientRequest::from(vec![query]);
+            assert!(req.is_begin());
+        }
+    }
 }
diff --git a/pgdog/src/net/messages/error_response.rs b/pgdog/src/net/messages/error_response.rs

Original file line number	Diff line number	Diff line change
`@@ -49,13 +49,15 @@ pub enum PoolerMode {`
`49`	`49`	`#[default]`
`50`	`50`	`Transaction,`
`51`	`51`	`Session,`
	`52`	`+ Statement,`
`52`	`53`	`}`
`53`	`54`
`54`	`55`	`impl std::fmt::Display for PoolerMode {`
`55`	`56`	`fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {`
`56`	`57`	`match self {`
`57`	`58`	`Self::Transaction => write!(f, "transaction"),`
`58`	`59`	`Self::Session => write!(f, "session"),`
	`60`	`+ Self::Statement => write!(f, "statement"),`
`59`	`61`	`}`
`60`	`62`	`}`
`61`	`63`	`}`