fix(prepared_statements): follow the client's resolved pooler mode

pgdog-config/src/core.rs

@@ -167,10 +167,28 @@ impl ConfigAndUsers {
         Ok(())
     }
 
-    /// Prepared statements are enabled.
+    /// Prepared statements level for clients using the `[general]` pooler mode.
+    ///
+    /// Prefer [`Self::prepared_statements_for`] with the pooler mode the client's
+    /// pool actually runs in: users and databases can override `pooler_mode`.
     pub fn prepared_statements(&self) -> PreparedStatementsLevel {
-        // Disable prepared statements automatically in session mode
-        if self.config.general.pooler_mode == PoolerMode::Session {
+        self.prepared_statements_for(self.config.general.pooler_mode)
+    }
+
+    /// Prepared statements level for a client whose pool runs in `pooler_mode`.
+    ///
+    /// A pool that multiplexes server connections (transaction and statement
+    /// mode, and with read/write split, primary and replica servers) must track
+    /// prepared statements and re-prepare them per server, even when the
+    /// `[general]` default is session mode and would otherwise disable tracking.
+    ///
+    /// Session-pooled clients keep the `[general]`-derived behaviour: tracking is
+    /// not forced off for them, because the router still needs the statement
+    /// text to route `Bind` and `Execute` on sharded clusters.
+    pub fn prepared_statements_for(&self, pooler_mode: PoolerMode) -> PreparedStatementsLevel {
+        if pooler_mode == PoolerMode::Session
+            && self.config.general.pooler_mode == PoolerMode::Session
+        {
             PreparedStatementsLevel::Disabled
         } else {
             self.config.general.prepared_statements
@@ -1028,6 +1046,53 @@ tls_server_private_key = "/certs/replica-client.key"
         );
     }
 
+    /// A `[general]` session default must not disable tracking for a user or
+    /// database configured for transaction pooling. Session-pooled clients
+    /// keep whatever the `[general]` default resolves to.
+    #[test]
+    fn test_prepared_statements_for_pooler_mode() {
+        let mut config = ConfigAndUsers::default();
+        config.config.general.pooler_mode = PoolerMode::Session;
+        config.config.general.prepared_statements = PreparedStatementsLevel::Extended;
+
+        assert_eq!(
+            config.prepared_statements(),
+            PreparedStatementsLevel::Disabled
+        );
+        assert_eq!(
+            config.prepared_statements_for(PoolerMode::Transaction),
+            PreparedStatementsLevel::Extended,
+            "transaction-pooled client must track statements despite the session default"
+        );
+        assert_eq!(
+            config.prepared_statements_for(PoolerMode::Statement),
+            PreparedStatementsLevel::Extended
+        );
+        assert_eq!(
+            config.prepared_statements_for(PoolerMode::Session),
+            PreparedStatementsLevel::Disabled,
+            "session-pooled client follows the session default"
+        );
+
+        config.config.general.pooler_mode = PoolerMode::Transaction;
+        assert_eq!(
+            config.prepared_statements(),
+            PreparedStatementsLevel::Extended
+        );
+        assert_eq!(
+            config.prepared_statements_for(PoolerMode::Session),
+            PreparedStatementsLevel::Extended,
+            "session-pooled client on a transaction default keeps tracking (needed for sharded routing)"
+        );
+
+        config.config.general.prepared_statements = PreparedStatementsLevel::Disabled;
+        assert_eq!(
+            config.prepared_statements_for(PoolerMode::Transaction),
+            PreparedStatementsLevel::Disabled,
+            "explicitly disabled stays disabled"
+        );
+    }
+
     #[test]
     fn test_mirroring_config() {
         let source = r#"

pgdog/src/backend/pool/connection/mirror/mod.rs

@@ -51,9 +51,11 @@ pub(crate) struct Mirror {
 }
 
 impl Mirror {
-    fn new(params: &Parameters, config: &ConfigAndUsers) -> Self {
+    fn new(params: &Parameters, config: &ConfigAndUsers, cluster: &Cluster) -> Self {
         let mut prepared_statements = PreparedStatements::new();
-        prepared_statements.set_level(config.prepared_statements());
+        // Follow the destination cluster's pooler mode, like a regular client
+        // of that cluster would, rather than the `[general]` default.
+        prepared_statements.set_level(config.prepared_statements_for(cluster.pooler_mode()));
 
         Self {
             id: FrontendPid::new(),
@@ -96,7 +98,7 @@ impl Mirror {
         ]);
 
         // Mirror traffic handler.
-        let mut mirror = Self::new(&params, &config);
+        let mut mirror = Self::new(&params, &config, cluster);
 
         // Same query engine as the client, except with a potentially different database config.
         // Use mirror.id so pool checkout (Request) and comms keying share one identity.
@@ -376,10 +378,22 @@ mod test {
             },
         ]);
 
-        let mirror = Mirror::new(&params, &config);
+        // Destination cluster in session mode follows the `[general]` default.
+        let cluster = Cluster::new_test_session_mode(&config);
+        let mirror = Mirror::new(&params, &config, &cluster);
         assert_eq!(
             mirror.prepared_statements.level(),
             PreparedStatementsLevel::Disabled
         );
+
+        // Destination cluster in transaction mode (the test cluster default)
+        // must track statements even though `[general]` is session mode.
+        let cluster = Cluster::new_test(&config);
+        assert_eq!(cluster.pooler_mode(), PoolerMode::Transaction);
+        let mirror = Mirror::new(&params, &config, &cluster);
+        assert_eq!(
+            mirror.prepared_statements.level(),
+            PreparedStatementsLevel::Extended
+        );
     }
 }

pgdog/src/frontend/client/mod.rs

@@ -8,6 +8,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use chrono::{DateTime, Utc};
+use pgdog_config::PoolerMode;
 use pgdog_config::users::PasswordKind;
 use timeouts::Timeouts;
 use tokio::{select, spawn};
@@ -532,6 +533,7 @@ impl Client {
             }
 
             let client_state = query_engine.client_state();
+            let pooler_mode = query_engine.pooler_mode();
 
             let cancellation_token = query_engine.cancellation_token();
 
@@ -546,7 +548,7 @@ impl Client {
                     self.server_message(&mut query_engine, message).await?;
                 }
 
-                buffer = self.buffer(client_state, &cancellation_token) => {
+                buffer = self.buffer(client_state, pooler_mode, &cancellation_token) => {
                     let event = buffer?;
 
                     // Only send requests to the backend if they are complete.
@@ -649,14 +651,21 @@ impl Client {
     async fn buffer(
         &mut self,
         state: State,
+        pooler_mode: Option<PoolerMode>,
         cancellation_token: &CancellationToken,
     ) -> Result<BufferEvent, Error> {
         self.client_request.clear();
 
         // Check config once per request.
         let config = config::config();
-        // Configure prepared statements cache.
-        self.prepared_statements.level = config.prepared_statements();
+        // Configure prepared statements cache using the pooler mode this
+        // client's pool actually runs in. The `[general]` default alone is
+        // wrong when it is session mode but this user is transaction pooled:
+        // untracked Parse and Bind could land on different servers.
+        self.prepared_statements.level = match pooler_mode {
+            Some(pooler_mode) => config.prepared_statements_for(pooler_mode),
+            None => config.prepared_statements(),
+        };
         self.timeouts = Timeouts::from_config(&config.config.general);
         self.query_log_stdout = config.config.general.query_log_stdout;
         self.query_size_limit = config.config.general.query_size_limit;

pgdog/src/frontend/client/query_engine/mod.rs

@@ -10,6 +10,7 @@ use crate::{
     net::{ErrorResponse, Message, Parameters},
     state::State,
 };
+use pgdog_config::PoolerMode;
 use tokio_util::sync::CancellationToken;
 use tracing::debug;
 
@@ -130,6 +131,16 @@ impl QueryEngine {
         self.stats.state
     }
 
+    /// Pooler mode of the cluster this client is connected to, resolved from
+    /// the user, database and `[general]` config. `None` when the client has
+    /// no cluster (e.g. admin database).
+    pub(crate) fn pooler_mode(&self) -> Option<PoolerMode> {
+        self.backend
+            .cluster()
+            .ok()
+            .map(|cluster| cluster.pooler_mode())
+    }
+
     /// Handle client request.
     pub(crate) async fn handle(
         &mut self,

pgdog/src/frontend/client/test/mod.rs

@@ -255,7 +255,7 @@ async fn test_abrupt_disconnect() {
     drop(conn);
 
     let event = client
-        .buffer(State::Idle, &CancellationToken::new())
+        .buffer(State::Idle, None, &CancellationToken::new())
         .await
         .unwrap();
     assert_eq!(event, BufferEvent::DisconnectAbrupt);
@@ -277,7 +277,7 @@ async fn test_client_idle_timeout() {
 
     let start = Instant::now();
     let res = client
-        .buffer(State::Idle, &CancellationToken::new())
+        .buffer(State::Idle, None, &CancellationToken::new())
         .await
         .unwrap();
     assert_eq!(res, BufferEvent::DisconnectAbrupt);
@@ -290,7 +290,7 @@ async fn test_client_idle_timeout() {
     assert!(
         timeout(
             Duration::from_millis(50),
-            client.buffer(State::IdleInTransaction, &CancellationToken::new())
+            client.buffer(State::IdleInTransaction, None, &CancellationToken::new())
         )
         .await
         .is_err()
@@ -643,7 +643,7 @@ async fn test_query_timeout() {
     conn.write_all(&buf).await.unwrap();
 
     client
-        .buffer(State::Idle, &CancellationToken::new())
+        .buffer(State::Idle, None, &CancellationToken::new())
         .await
         .unwrap();
     let result = client.client_messages(&mut engine).await;

pgdog/src/frontend/client/test/test_client.rs

@@ -238,8 +238,13 @@ impl TestClient {
 
     /// Process a request.
     pub(crate) async fn try_process(&mut self) -> Result<(), Box<dyn std::error::Error>> {
+        let pooler_mode = self.engine.pooler_mode();
         self.client
-            .buffer(self.engine.stats().state, &CancellationToken::new())
+            .buffer(
+                self.engine.stats().state,
+                pooler_mode,
+                &CancellationToken::new(),
+            )
             .await?;
         self.client.client_messages(&mut self.engine).await?;