From 2bf1e7a5736fb18c0f3ea2fde7a36d3249040284 Mon Sep 17 00:00:00 2001 From: Eric Hauser Date: Sun, 13 Sep 2026 21:44:39 -0600 Subject: [PATCH 1/3] fix(prepared_statements): follow the client's resolved pooler mode `ConfigAndUsers::prepared_statements()` only looked at the `[general]` pooler mode. With `[general] pooler_mode = "session"` and a user or database overridden to `"transaction"`, prepared statement tracking was disabled for that user even though its pool multiplexes server connections. Under read/write split a named Parse (routed to a replica) and its Bind (routed to the primary when untracked) landed on different servers, so drivers with statement caches (pgx, npgsql, etc.) failed with `prepared statement "..." does not exist` (SQLSTATE 26000). Add `ConfigAndUsers::prepared_statements_for(pooler_mode)` and have the client and mirror derive the level from the pooler mode of the cluster they are actually bound to, while still reading the level itself from live config on every request as before. Admin clients have no cluster and keep the `[general]` behaviour. Fixes #1536 --- pgdog-config/src/core.rs | 60 ++++++++++++++++++- .../src/backend/pool/connection/mirror/mod.rs | 22 +++++-- pgdog/src/frontend/client/mod.rs | 15 ++++- pgdog/src/frontend/client/query_engine/mod.rs | 11 ++++ pgdog/src/frontend/client/test/mod.rs | 8 +-- pgdog/src/frontend/client/test/test_client.rs | 7 ++- 6 files changed, 108 insertions(+), 15 deletions(-) diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs index 010a4518e..93bff3d9c 100644 --- a/pgdog-config/src/core.rs +++ b/pgdog-config/src/core.rs @@ -167,10 +167,23 @@ 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`. + /// + /// Session mode pins a client to one server connection, so its prepared + /// statements are always present on the server it talks to and tracking is + /// unnecessary. Transaction mode multiplexes server connections (and, with + /// read/write split, primary and replica servers), so statements must be + /// tracked and re-prepared per server. + pub fn prepared_statements_for(&self, pooler_mode: PoolerMode) -> PreparedStatementsLevel { + if pooler_mode == PoolerMode::Session { PreparedStatementsLevel::Disabled } else { self.config.general.prepared_statements @@ -1028,6 +1041,47 @@ 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, and vice versa. + #[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::Session), + PreparedStatementsLevel::Disabled + ); + + config.config.general.pooler_mode = PoolerMode::Transaction; + assert_eq!( + config.prepared_statements(), + PreparedStatementsLevel::Extended + ); + assert_eq!( + config.prepared_statements_for(PoolerMode::Session), + PreparedStatementsLevel::Disabled, + "session-pooled client never needs tracking" + ); + + 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#" diff --git a/pgdog/src/backend/pool/connection/mirror/mod.rs b/pgdog/src/backend/pool/connection/mirror/mod.rs index b196cc5b2..fff981ada 100644 --- a/pgdog/src/backend/pool/connection/mirror/mod.rs +++ b/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(¶ms, &config); + let mut mirror = Self::new(¶ms, &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(¶ms, &config); + // Destination cluster in session mode: no tracking needed. + let cluster = Cluster::new_test_session_mode(&config); + let mirror = Mirror::new(¶ms, &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(¶ms, &config, &cluster); + assert_eq!( + mirror.prepared_statements.level(), + PreparedStatementsLevel::Extended + ); } } diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 96eb23e9a..e0e66baf8 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/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, cancellation_token: &CancellationToken, ) -> Result { 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; diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index 9196bcf8b..92b5897db 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/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 { + self.backend + .cluster() + .ok() + .map(|cluster| cluster.pooler_mode()) + } + /// Handle client request. pub(crate) async fn handle( &mut self, diff --git a/pgdog/src/frontend/client/test/mod.rs b/pgdog/src/frontend/client/test/mod.rs index febb77a5e..72e75117b 100644 --- a/pgdog/src/frontend/client/test/mod.rs +++ b/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; diff --git a/pgdog/src/frontend/client/test/test_client.rs b/pgdog/src/frontend/client/test/test_client.rs index 6adadff89..c535d0d4e 100644 --- a/pgdog/src/frontend/client/test/test_client.rs +++ b/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> { + 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?; From 8d7a8cd8c48e29c45717a46401318adc57c00d50 Mon Sep 17 00:00:00 2001 From: Eric Hauser Date: Mon, 14 Sep 2026 09:56:45 -0600 Subject: [PATCH 2/3] fix(prepared_statements): keep tracking for session-pooled sharded clients `prepared_statements_for(Session)` forced tracking off per cluster, which broke `test_session_multiple_statements_in_transaction`: a session-pooled user on a schema-sharded cluster sends a named Parse + Bind once psycopg auto-prepares, and without a tracked statement the router has no SQL to route the Bind, so it lands on the wrong shard ("cannot switch shards in a direct-to-shard transaction"). Only multiplexing pools (transaction and statement mode) override the `[general]` default; session-pooled clients keep the existing `[general]`-derived behaviour, so nothing changes for them relative to main. --- pgdog-config/src/core.rs | 43 +++++++++++++------ .../src/backend/pool/connection/mirror/mod.rs | 2 +- 2 files changed, 30 insertions(+), 15 deletions(-) diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs index 93bff3d9c..867886967 100644 --- a/pgdog-config/src/core.rs +++ b/pgdog-config/src/core.rs @@ -172,21 +172,30 @@ impl ConfigAndUsers { /// 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 { - self.prepared_statements_for(self.config.general.pooler_mode) + // Disable prepared statements automatically in session mode + if self.config.general.pooler_mode == PoolerMode::Session { + PreparedStatementsLevel::Disabled + } else { + self.config.general.prepared_statements + } } /// Prepared statements level for a client whose pool runs in `pooler_mode`. /// - /// Session mode pins a client to one server connection, so its prepared - /// statements are always present on the server it talks to and tracking is - /// unnecessary. Transaction mode multiplexes server connections (and, with - /// read/write split, primary and replica servers), so statements must be - /// tracked and re-prepared per server. + /// 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 { - PreparedStatementsLevel::Disabled - } else { - self.config.general.prepared_statements + match pooler_mode { + PoolerMode::Session => self.prepared_statements(), + PoolerMode::Transaction | PoolerMode::Statement => { + self.config.general.prepared_statements + } } } @@ -1042,7 +1051,8 @@ 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, and vice versa. + /// 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(); @@ -1058,9 +1068,14 @@ tls_server_private_key = "/certs/replica-client.key" 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 + PreparedStatementsLevel::Disabled, + "session-pooled client follows the session default" ); config.config.general.pooler_mode = PoolerMode::Transaction; @@ -1070,8 +1085,8 @@ tls_server_private_key = "/certs/replica-client.key" ); assert_eq!( config.prepared_statements_for(PoolerMode::Session), - PreparedStatementsLevel::Disabled, - "session-pooled client never needs tracking" + PreparedStatementsLevel::Extended, + "session-pooled client on a transaction default keeps tracking (needed for sharded routing)" ); config.config.general.prepared_statements = PreparedStatementsLevel::Disabled; diff --git a/pgdog/src/backend/pool/connection/mirror/mod.rs b/pgdog/src/backend/pool/connection/mirror/mod.rs index fff981ada..6b99f37fd 100644 --- a/pgdog/src/backend/pool/connection/mirror/mod.rs +++ b/pgdog/src/backend/pool/connection/mirror/mod.rs @@ -378,7 +378,7 @@ mod test { }, ]); - // Destination cluster in session mode: no tracking needed. + // Destination cluster in session mode follows the `[general]` default. let cluster = Cluster::new_test_session_mode(&config); let mirror = Mirror::new(¶ms, &config, &cluster); assert_eq!( From f67fcde2cfefe14d9e9eec408e33c9f19c97ce52 Mon Sep 17 00:00:00 2001 From: Eric Hauser Date: Thu, 17 Sep 2026 14:20:37 -0600 Subject: [PATCH 3/3] refactor: centralize prepared statement level resolution --- pgdog-config/src/core.rs | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs index 867886967..2584ef2ac 100644 --- a/pgdog-config/src/core.rs +++ b/pgdog-config/src/core.rs @@ -172,12 +172,7 @@ impl ConfigAndUsers { /// 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 { - PreparedStatementsLevel::Disabled - } else { - self.config.general.prepared_statements - } + self.prepared_statements_for(self.config.general.pooler_mode) } /// Prepared statements level for a client whose pool runs in `pooler_mode`. @@ -191,11 +186,12 @@ impl ConfigAndUsers { /// 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 { - match pooler_mode { - PoolerMode::Session => self.prepared_statements(), - PoolerMode::Transaction | PoolerMode::Statement => { - self.config.general.prepared_statements - } + if pooler_mode == PoolerMode::Session + && self.config.general.pooler_mode == PoolerMode::Session + { + PreparedStatementsLevel::Disabled + } else { + self.config.general.prepared_statements } }