From 5642a3d45fe63493b7cdfef497d65b08661ff302 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Fri, 18 Sep 2026 14:16:26 +0530 Subject: [PATCH 01/12] fix --- pgdog/src/backend/pool/cleanup.rs | 32 +++++++++++++++- pgdog/src/backend/pool/guard.rs | 63 ++++++++++++++++++++++++++++++- 2 files changed, 92 insertions(+), 3 deletions(-) diff --git a/pgdog/src/backend/pool/cleanup.rs b/pgdog/src/backend/pool/cleanup.rs index df077a4ca..57e5bd343 100644 --- a/pgdog/src/backend/pool/cleanup.rs +++ b/pgdog/src/backend/pool/cleanup.rs @@ -13,9 +13,12 @@ static PREPARED: Lazy> = Lazy::new(|| vec![Query::new("DEALLOCATE ALL /// static DIRTY: Lazy> = Lazy::new(|| { vec![ - Query::new("RESET ALL"), // Reset all parameters. + // RESET ALL deliberately leaves role and session_authorization alone. + // Resetting session authorization also clears the active role. + Query::new("SET SESSION AUTHORIZATION DEFAULT"), + Query::new("RESET ALL"), // Reset all other parameters. Query::new("SELECT pg_advisory_unlock_all()"), // Remove all advisory locks. - Query::new("DISCARD TEMP"), // Drop all temporary tables. + Query::new("DISCARD TEMP"), // Drop all temporary tables. ] }); @@ -136,3 +139,28 @@ impl Cleanup { self.deallocate } } + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn dirty_cleanup_resets_session_identity_before_other_state() { + let cleanup = Cleanup::parameters(); + let queries = cleanup + .queries() + .iter() + .map(Query::query) + .collect::>(); + + assert_eq!( + queries, + [ + "SET SESSION AUTHORIZATION DEFAULT", + "RESET ALL", + "SELECT pg_advisory_unlock_all()", + "DISCARD TEMP", + ] + ); + } +} diff --git a/pgdog/src/backend/pool/guard.rs b/pgdog/src/backend/pool/guard.rs index a5eeb39ab..8432e86fc 100644 --- a/pgdog/src/backend/pool/guard.rs +++ b/pgdog/src/backend/pool/guard.rs @@ -114,9 +114,27 @@ mod test { }, server::test::test_server, }, - net::{Describe, Flush, Parse, Protocol, ProtocolMessage, Query, Sync}, + net::{DataRow, Describe, Flush, Parse, Protocol, ProtocolMessage, Query, Sync}, }; + async fn identity(guard: &mut Guard) -> (String, String) { + let messages = guard + .execute("SELECT session_user, current_user") + .await + .unwrap(); + let row: DataRow = messages + .into_iter() + .find(|message| message.code() == 'D') + .expect("identity query should return one row") + .try_into() + .unwrap(); + + ( + row.get_text(0).expect("session_user"), + row.get_text(1).expect("current_user"), + ) + } + #[tokio::test] async fn test_cleanup_dirty() { crate::logger(); @@ -177,6 +195,49 @@ mod test { drop(guard); } + #[tokio::test] + async fn test_cleanup_dirty_resets_session_identity() { + crate::logger(); + let pool = pool(); + let mut guard = pool.get(&Request::default()).await.unwrap(); + let server_id = guard.id(); + let role = format!("pgdog_cleanup_role_{}", std::process::id()); + + guard + .execute_checked(format!("CREATE ROLE {role}")) + .await + .unwrap(); + guard + .execute_checked(format!("SET ROLE {role}")) + .await + .unwrap(); + assert_eq!(identity(&mut guard).await, ("pgdog".into(), role.clone())); + + guard.mark_dirty(true); + drop(guard); + + let mut guard = pool.get(&Request::default()).await.unwrap(); + assert_eq!(guard.id(), server_id); + assert_eq!(identity(&mut guard).await, ("pgdog".into(), "pgdog".into())); + + guard + .execute_checked(format!("SET SESSION AUTHORIZATION {role}")) + .await + .unwrap(); + assert_eq!(identity(&mut guard).await, (role.clone(), role.clone())); + + guard.mark_dirty(true); + drop(guard); + + let mut guard = pool.get(&Request::default()).await.unwrap(); + assert_eq!(guard.id(), server_id); + assert_eq!(identity(&mut guard).await, ("pgdog".into(), "pgdog".into())); + guard + .execute_checked(format!("DROP ROLE {role}")) + .await + .unwrap(); + } + #[tokio::test] async fn test_cleanup_prepared_statements() { crate::logger(); From 86235623954fbbaf7dfc2237bb46349bae09122a Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 21 Sep 2026 10:59:16 +0530 Subject: [PATCH 02/12] mark backend dirty on set role --- pgdog/src/backend/pool/connection/mod.rs | 5 ++ pgdog/src/frontend/client/query_engine/set.rs | 8 +++ .../client/query_engine/test/advisory_lock.rs | 71 +++++++++++++++++++ 3 files changed, 84 insertions(+) diff --git a/pgdog/src/backend/pool/connection/mod.rs b/pgdog/src/backend/pool/connection/mod.rs index dbc1a01a5..2c1bd71b8 100644 --- a/pgdog/src/backend/pool/connection/mod.rs +++ b/pgdog/src/backend/pool/connection/mod.rs @@ -456,6 +456,11 @@ impl Connection { } } + /// Mark connected servers as containing client-specific session state. + pub(crate) fn mark_dirty(&mut self) { + self.binding.dirty(); + } + /// Check if any held server connection is currently locked to a client. #[cfg(test)] pub(crate) fn locked(&self) -> bool { diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index 9202c5452..adf03f8dd 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -10,6 +10,7 @@ use super::*; /// queries route to a different shard than the one we're pinned to, so they may /// only be set before any query connects to a backend. const SHARD_TARGETING_PARAMS: [&str; 2] = [PGDOG_SHARD, PGDOG_SHARDING_KEY]; +const SESSION_IDENTITY_PARAMS: [&str; 2] = ["role", "session_authorization"]; impl QueryEngine { /// Handle a `SET` statement or equivalent `SELECT set_config([...])` query. @@ -58,6 +59,13 @@ impl QueryEngine { } if self.backend.connected() { + if params.iter().any(|param| { + SESSION_IDENTITY_PARAMS + .iter() + .any(|name| param.name.eq_ignore_ascii_case(name)) + }) { + self.backend.mark_dirty(); + } self.execute(context, None).await?; } else { let fake_response = set_config diff --git a/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs b/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs index 3d73de1c6..bf0ea8830 100644 --- a/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs +++ b/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs @@ -1,5 +1,40 @@ +use crate::{ + backend::databases::reload_from_existing, + config::{config, load_test, set}, + expect_message, + net::{CommandComplete, DataRow, ReadyForQuery, RowDescription}, +}; + use super::prelude::*; +fn load_single_connection_test_pool() { + load_test(); + + let mut config = (*config()).clone(); + config.config.general.default_pool_size = 1; + config.config.general.min_pool_size = 0; + set(config).unwrap(); + reload_from_existing().unwrap(); +} + +async fn identity(client: &mut TestClient) -> (i64, String, String) { + client + .send_simple(Query::new( + "SELECT pg_backend_pid(), session_user, current_user", + )) + .await; + expect_message!(client.read().await, RowDescription); + let row = expect_message!(client.read().await, DataRow); + expect_message!(client.read().await, CommandComplete); + expect_message!(client.read().await, ReadyForQuery); + + ( + row.get_int(0, true).expect("backend pid"), + row.get_text(1).expect("session_user"), + row.get_text(2).expect("current_user"), + ) +} + #[tokio::test] async fn test_session_lock_tracked_outside_transaction() { let mut client = TestClient::new_sharded(Parameters::default()).await; @@ -28,6 +63,42 @@ async fn test_session_lock_tracked_outside_transaction() { assert!(client.engine.advisory_locks().contains(101)); } +#[tokio::test] +async fn test_session_lock_cleanup_resets_role_before_backend_reuse() { + load_single_connection_test_pool(); + let mut source = TestClient::new(Parameters::default()).await; + let role = format!("pgdog_cleanup_role_{}", std::process::id()); + + source + .send_simple(Query::new(format!("CREATE ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + source + .send_simple(Query::new("SELECT pg_advisory_lock(707519)")) + .await; + source.read_until('Z').await.unwrap(); + source + .send_simple(Query::new(format!("SET ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + + let before = identity(&mut source).await; + assert_eq!( + (before.1.as_str(), before.2.as_str()), + ("pgdog", role.as_str()) + ); + drop(source.leak_pool()); + + let mut peer = TestClient::new(Parameters::default()).await; + let after = identity(&mut peer).await; + assert_eq!(after.0, before.0, "physical backend should be reused"); + assert_eq!((after.1.as_str(), after.2.as_str()), ("pgdog", "pgdog")); + + peer.send_simple(Query::new(format!("DROP ROLE {role}"))) + .await; + peer.read_until('Z').await.unwrap(); +} + #[tokio::test] async fn test_session_lock_inside_transaction_survives_commit() { // A plain pg_advisory_lock taken inside a transaction lives past COMMIT From 535521ea2c50c28b7deeeb9c829db3b9cfbe651b Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 21 Sep 2026 23:58:26 +0530 Subject: [PATCH 03/12] mark deferred set role backend dirty --- pgdog/src/backend/server.rs | 24 +++++++++++++++++++ .../router/parser/query/test/test_set.rs | 23 ++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index ac277eb80..78ac22db8 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -718,6 +718,7 @@ impl Server { let mut executed = if !params.identical(&self.client_params) { // Construct client parameter SET queries. let tracked = params.tracked_and_different(&self.client_params); + let sets_role = tracked.get("role").is_some(); // Construct RESET queries to reset any current params // to their default values. let mut queries = self.client_params.reset_queries(params); @@ -732,6 +733,9 @@ impl Server { debug!("syncing {} params", queries.len()); self.execute_batch(&queries).await?; + if sets_role { + self.mark_dirty(true); + } clear_params = true; } @@ -2477,6 +2481,7 @@ pub(crate) mod test { .link_client(FrontendPid::new(), ¶ms, None) .await?; assert_eq!(changed, 1); + assert!(!server.dirty()); let changed = server .link_client(FrontendPid::new(), ¶ms, None) @@ -2501,6 +2506,25 @@ pub(crate) mod test { Ok(()) } + #[tokio::test] + async fn test_link_client_marks_server_dirty_when_setting_role() + -> Result<(), Box> { + let mut params = Parameters::default(); + params.insert("role", "pgdog"); + + let mut server = test_server().await; + assert!(!server.dirty()); + + let changed = server + .link_client(FrontendPid::new(), ¶ms, None) + .await?; + + assert_eq!(changed, 1); + assert!(server.dirty()); + + Ok(()) + } + #[tokio::test] async fn test_copy_protocol() { let mut server = test_server().await; diff --git a/pgdog/src/frontend/router/parser/query/test/test_set.rs b/pgdog/src/frontend/router/parser/query/test/test_set.rs index 8917f2993..261b67322 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_set.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_set.rs @@ -79,6 +79,29 @@ fn test_set_comment() { ); } +#[test] +fn test_set_role_is_tracked_as_parameter() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![Query::new("SET ROLE other_user").into()]); + + match command { + Command::Set { + params, set_config, .. + } => { + assert_eq!(params.len(), 1); + assert_eq!(params[0].name, "role"); + assert_eq!( + params[0].value, + Some(ParameterValue::String("other_user".into())) + ); + assert!(!params[0].local); + assert!(!set_config); + } + _ => panic!("expected Command::Set, got {command:#?}"), + } +} + #[test] fn test_set_config_null_value() { let mut test = QueryParserTest::new(); From 8c81b97d32c0ea5432b69f2488e8ec120c335936 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 22 Sep 2026 00:15:24 +0530 Subject: [PATCH 04/12] role sync cleanup --- pgdog/src/backend/server.rs | 39 ++++++++++++++++++++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 78ac22db8..56d3ea500 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -732,10 +732,10 @@ impl Server { if !queries.is_empty() { debug!("syncing {} params", queries.len()); - self.execute_batch(&queries).await?; if sets_role { self.mark_dirty(true); } + self.execute_batch(&queries).await?; clear_params = true; } @@ -756,6 +756,9 @@ impl Server { if !transaction_sets.is_empty() { debug!("syncing {} in-transaction params", transaction_sets.len()); + if params.get("role").is_some() { + self.mark_dirty(true); + } self.execute_batch(&transaction_sets).await?; clear_params = true; @@ -2525,6 +2528,40 @@ pub(crate) mod test { Ok(()) } + #[tokio::test] + async fn test_link_client_marks_server_dirty_before_role_batch_failure() { + let mut params = Parameters::default(); + params.insert("role", "pgdog"); + params.insert("work_mem", "not-a-size"); + + let mut server = test_server().await; + assert!( + server + .link_client(FrontendPid::new(), ¶ms, None) + .await + .is_err() + ); + assert!(server.dirty()); + } + + #[tokio::test] + async fn test_link_client_marks_server_dirty_for_transaction_role() + -> Result<(), Box> { + let mut params = Parameters::default(); + params.insert_transaction("role", "pgdog", false); + + let mut server = test_server().await; + let changed = server + .link_client(FrontendPid::new(), ¶ms, Some("BEGIN")) + .await?; + + assert_eq!(changed, 1); + assert!(server.dirty()); + server.rollback().await?; + + Ok(()) + } + #[tokio::test] async fn test_copy_protocol() { let mut server = test_server().await; From 3136c845e80d3fd54a8ad83cff2c68fa3310bb77 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 22 Sep 2026 00:21:06 +0530 Subject: [PATCH 05/12] e2e test --- .../client/query_engine/test/advisory_lock.rs | 82 ++++++++++++++++++- 1 file changed, 81 insertions(+), 1 deletion(-) diff --git a/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs b/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs index bf0ea8830..d6cb2a203 100644 --- a/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs +++ b/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs @@ -67,7 +67,7 @@ async fn test_session_lock_tracked_outside_transaction() { async fn test_session_lock_cleanup_resets_role_before_backend_reuse() { load_single_connection_test_pool(); let mut source = TestClient::new(Parameters::default()).await; - let role = format!("pgdog_cleanup_role_{}", std::process::id()); + let role = format!("pgdog_cleanup_role_{}_advisory", std::process::id()); source .send_simple(Query::new(format!("CREATE ROLE {role}"))) @@ -99,6 +99,86 @@ async fn test_session_lock_cleanup_resets_role_before_backend_reuse() { peer.read_until('Z').await.unwrap(); } +#[tokio::test] +async fn test_deferred_role_cleanup_before_backend_reuse() { + load_single_connection_test_pool(); + let mut source = TestClient::new(Parameters::default()).await; + let role = format!("pgdog_cleanup_role_{}_deferred", std::process::id()); + + source + .send_simple(Query::new(format!("CREATE ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + source + .send_simple(Query::new(format!("SET ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + assert!( + !source.backend_connected(), + "SET ROLE should be deferred until a query checks out a backend" + ); + + let before = identity(&mut source).await; + assert_eq!( + (before.1.as_str(), before.2.as_str()), + ("pgdog", role.as_str()) + ); + assert!(!source.backend_connected()); + drop(source.leak_pool()); + + let mut peer = TestClient::new(Parameters::default()).await; + let after = identity(&mut peer).await; + assert_eq!(after.0, before.0, "physical backend should be reused"); + assert_eq!((after.1.as_str(), after.2.as_str()), ("pgdog", "pgdog")); + + peer.send_simple(Query::new(format!("DROP ROLE {role}"))) + .await; + peer.read_until('Z').await.unwrap(); +} + +#[tokio::test] +async fn test_deferred_transaction_role_cleanup_before_backend_reuse() { + load_single_connection_test_pool(); + let mut source = TestClient::new(Parameters::default()).await; + let role = format!("pgdog_cleanup_role_{}_transaction", std::process::id()); + + source + .send_simple(Query::new(format!("CREATE ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + source.send_simple(Query::new("BEGIN")).await; + source.read_until('Z').await.unwrap(); + source + .send_simple(Query::new(format!("SET ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + assert!( + !source.backend_connected(), + "BEGIN and SET ROLE should remain deferred until the first query" + ); + + let before = identity(&mut source).await; + assert_eq!( + (before.1.as_str(), before.2.as_str()), + ("pgdog", role.as_str()) + ); + assert!(source.backend_connected()); + + source.send_simple(Query::new("COMMIT")).await; + source.read_until('Z').await.unwrap(); + assert!(!source.backend_connected()); + drop(source.leak_pool()); + + let mut peer = TestClient::new(Parameters::default()).await; + let after = identity(&mut peer).await; + assert_eq!(after.0, before.0, "physical backend should be reused"); + assert_eq!((after.1.as_str(), after.2.as_str()), ("pgdog", "pgdog")); + + peer.send_simple(Query::new(format!("DROP ROLE {role}"))) + .await; + peer.read_until('Z').await.unwrap(); +} + #[tokio::test] async fn test_session_lock_inside_transaction_survives_commit() { // A plain pg_advisory_lock taken inside a transaction lives past COMMIT From 4a9b2a7d03d5ad770eca13b9e4ca336a1f4eb968 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 28 Sep 2026 13:13:52 +0530 Subject: [PATCH 06/12] session params --- pgdog/src/backend/pool/connection/binding.rs | 16 +- pgdog/src/backend/server.rs | 69 ++- pgdog/src/frontend/client/query_engine/mod.rs | 3 + .../src/frontend/client/query_engine/query.rs | 1 + pgdog/src/frontend/client/query_engine/set.rs | 45 +- .../client/query_engine/test/advisory_lock.rs | 151 ------- .../frontend/client/query_engine/test/mod.rs | 1 + .../query_engine/test/session_identity.rs | 290 ++++++++++++ pgdog/src/frontend/router/parser/query/set.rs | 28 +- .../router/parser/query/test/test_set.rs | 51 +++ pgdog/src/net/parameter.rs | 416 +++++++++++++++++- 11 files changed, 867 insertions(+), 204 deletions(-) create mode 100644 pgdog/src/frontend/client/query_engine/test/session_identity.rs diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index 4c32e8868..adb4187b3 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -430,13 +430,21 @@ impl Binding { let results = join_all(futures).await; let mut max = 0; + let mut error = None; for result in results { - let synced = result?; - if max < synced { - max = synced; + match result { + Ok(synced) => max = max.max(synced), + Err(err) => error = Some(err), } } - Ok(max) + if let Some(error) = error { + for server in servers { + server.force_close(); + } + Err(error) + } else { + Ok(max) + } } _ => Ok(0), diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 56d3ea500..fbf829937 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -674,6 +674,7 @@ impl Server { "DISCARD ALL" => { self.prepared_statements.clear(); self.client_params.clear(); + self.client_params.clear_session_identity(); } "RESET" => self.client_params.clear(), // Someone reset params, we're gonna need to re-sync. _ => (), @@ -718,10 +719,15 @@ impl Server { let mut executed = if !params.identical(&self.client_params) { // Construct client parameter SET queries. let tracked = params.tracked_and_different(&self.client_params); - let sets_role = tracked.get("role").is_some(); + let identity_queries = self + .client_params + .session_identity(false) + .reconcile(¶ms.session_identity(false), false); + let identity_changed = !identity_queries.is_empty(); + let mut queries = identity_queries; // Construct RESET queries to reset any current params // to their default values. - let mut queries = self.client_params.reset_queries(params); + queries.extend(self.client_params.reset_queries(params)); // Combine both to create a new, fresh session state // on this connection. @@ -732,10 +738,12 @@ impl Server { if !queries.is_empty() { debug!("syncing {} params", queries.len()); - if sets_role { - self.mark_dirty(true); + if let Err(err) = self.execute_batch(&queries).await { + if identity_changed { + self.force_close(); + } + return Err(err); } - self.execute_batch(&queries).await?; clear_params = true; } @@ -751,15 +759,28 @@ impl Server { // need to be revered on rollback or commited on commit. if let Some(start_transaction) = start_transaction { self.execute(start_transaction).await?; - let transaction_sets = params.set_queries(true); + let mut current_identity = self.client_params.session_identity(false); + let mut identity_queries = vec![]; + if let Some(identity) = params.transaction_session_identity() { + identity_queries.extend(current_identity.reconcile(identity, false)); + current_identity.clone_from(identity); + } + if let Some(identity) = params.local_session_identity() { + identity_queries.extend(current_identity.reconcile(identity, true)); + } + let identity_changed = !identity_queries.is_empty(); + let mut transaction_sets = identity_queries; + transaction_sets.extend(params.set_queries(true)); if !transaction_sets.is_empty() { debug!("syncing {} in-transaction params", transaction_sets.len()); - if params.get("role").is_some() { - self.mark_dirty(true); + if let Err(err) = self.execute_batch(&transaction_sets).await { + if identity_changed { + self.force_close(); + } + return Err(err); } - self.execute_batch(&transaction_sets).await?; clear_params = true; self.client_params.copy_in_transaction(params); @@ -1128,6 +1149,7 @@ impl Server { #[inline] pub(crate) fn reset_params(&mut self) { self.client_params.clear(); + self.client_params.clear_session_identity(); } #[inline] @@ -2510,10 +2532,10 @@ pub(crate) mod test { } #[tokio::test] - async fn test_link_client_marks_server_dirty_when_setting_role() - -> Result<(), Box> { + async fn test_link_client_reconciles_session_identity() -> Result<(), Box> + { let mut params = Parameters::default(); - params.insert("role", "pgdog"); + params.insert_identity("role", &"pgdog".into(), false, false); let mut server = test_server().await; assert!(!server.dirty()); @@ -2522,16 +2544,21 @@ pub(crate) mod test { .link_client(FrontendPid::new(), ¶ms, None) .await?; - assert_eq!(changed, 1); - assert!(server.dirty()); + assert_eq!(changed, 2); + assert!(!server.dirty()); + + let changed = server + .link_client(FrontendPid::new(), ¶ms, None) + .await?; + assert_eq!(changed, 0); Ok(()) } #[tokio::test] - async fn test_link_client_marks_server_dirty_before_role_batch_failure() { + async fn test_link_client_force_closes_after_identity_batch_failure() { let mut params = Parameters::default(); - params.insert("role", "pgdog"); + params.insert_identity("role", &"pgdog".into(), false, false); params.insert("work_mem", "not-a-size"); let mut server = test_server().await; @@ -2541,22 +2568,22 @@ pub(crate) mod test { .await .is_err() ); - assert!(server.dirty()); + assert!(server.is_force_close()); } #[tokio::test] - async fn test_link_client_marks_server_dirty_for_transaction_role() + async fn test_link_client_reconciles_transaction_identity() -> Result<(), Box> { let mut params = Parameters::default(); - params.insert_transaction("role", "pgdog", false); + params.insert_identity("role", &"pgdog".into(), true, false); let mut server = test_server().await; let changed = server .link_client(FrontendPid::new(), ¶ms, Some("BEGIN")) .await?; - assert_eq!(changed, 1); - assert!(server.dirty()); + assert_eq!(changed, 2); + assert!(!server.dirty()); server.rollback().await?; Ok(()) diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index 9196bcf8b..8dabc3c4c 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -75,6 +75,7 @@ pub(crate) struct QueryEngine { // or disconnect. manual_lock: bool, temp_tables: TempTables, + last_server_error: bool, } impl QueryEngine { @@ -103,6 +104,7 @@ impl QueryEngine { advisory_locks: AdvisoryLocks::default(), manual_lock: false, temp_tables: Default::default(), + last_server_error: false, }) } @@ -139,6 +141,7 @@ impl QueryEngine { return Ok(result); } + self.last_server_error = false; self.stats .received(context.client_request.total_message_len()); self.set_state(State::Active); // Client is active. diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 0e23a1e34..3ee019ef8 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -174,6 +174,7 @@ impl QueryEngine { } if code == 'E' { + self.last_server_error = true; if let Some(state) = self.pending_explain.as_mut() { state.annotated = true; } diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index adf03f8dd..2fba50c3c 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -1,6 +1,7 @@ use crate::frontend::SetParam; use crate::frontend::client::query_engine::fake::FakeResponse; use crate::frontend::router::parameter_hints::{PGDOG_PIN, PGDOG_SHARD, PGDOG_SHARDING_KEY}; +use crate::frontend::router::parser::ShardWithPriority; use crate::net::messages::ErrorResponse; use super::*; @@ -10,7 +11,6 @@ use super::*; /// queries route to a different shard than the one we're pinned to, so they may /// only be set before any query connects to a backend. const SHARD_TARGETING_PARAMS: [&str; 2] = [PGDOG_SHARD, PGDOG_SHARDING_KEY]; -const SESSION_IDENTITY_PARAMS: [&str; 2] = ["role", "session_authorization"]; impl QueryEngine { /// Handle a `SET` statement or equivalent `SELECT set_config([...])` query. @@ -25,12 +25,24 @@ impl QueryEngine { return Ok(()); } + let identity_changed = params + .iter() + .any(|param| Parameters::is_session_identity(¶m.name)); + let identity_before = (identity_changed && self.backend.connected()) + .then(|| Box::new(context.params.clone())); let mut fake_command = "SET"; for param in params { let is_pin = param.name == PGDOG_PIN; if let Some(value) = param.value.clone() { - if context.in_transaction() { + if context.params.insert_identity( + ¶m.name, + &value, + context.in_transaction(), + param.local, + ) { + continue; + } else if context.in_transaction() { context .params .insert_transaction(¶m.name, value, param.local); @@ -47,7 +59,13 @@ impl QueryEngine { } } else { fake_command = "RESET"; - context.params.reset(¶m.name); + if !context.params.reset_identity( + ¶m.name, + context.in_transaction(), + param.local, + ) { + context.params.reset(¶m.name); + } if is_pin { self.manual_lock = false; } @@ -59,14 +77,23 @@ impl QueryEngine { } if self.backend.connected() { - if params.iter().any(|param| { - SESSION_IDENTITY_PARAMS - .iter() - .any(|name| param.name.eq_ignore_ascii_case(name)) - }) { + if identity_changed { + context.client_request.route = Some(Route::write( + ShardWithPriority::new_override_transaction(Shard::All), + )); self.backend.mark_dirty(); } - self.execute(context, None).await?; + let result = self.execute(context, None).await; + if result.is_err() || self.last_server_error { + if let Some(params) = identity_before { + context.params.clone_from(¶ms); + self.comms.update_params(context.params); + } + if self.backend.connected() { + self.backend.force_close(); + } + } + result?; } else { let fake_response = set_config .then(|| params.iter().map(|p| p.value.as_ref())) diff --git a/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs b/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs index d6cb2a203..3d73de1c6 100644 --- a/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs +++ b/pgdog/src/frontend/client/query_engine/test/advisory_lock.rs @@ -1,40 +1,5 @@ -use crate::{ - backend::databases::reload_from_existing, - config::{config, load_test, set}, - expect_message, - net::{CommandComplete, DataRow, ReadyForQuery, RowDescription}, -}; - use super::prelude::*; -fn load_single_connection_test_pool() { - load_test(); - - let mut config = (*config()).clone(); - config.config.general.default_pool_size = 1; - config.config.general.min_pool_size = 0; - set(config).unwrap(); - reload_from_existing().unwrap(); -} - -async fn identity(client: &mut TestClient) -> (i64, String, String) { - client - .send_simple(Query::new( - "SELECT pg_backend_pid(), session_user, current_user", - )) - .await; - expect_message!(client.read().await, RowDescription); - let row = expect_message!(client.read().await, DataRow); - expect_message!(client.read().await, CommandComplete); - expect_message!(client.read().await, ReadyForQuery); - - ( - row.get_int(0, true).expect("backend pid"), - row.get_text(1).expect("session_user"), - row.get_text(2).expect("current_user"), - ) -} - #[tokio::test] async fn test_session_lock_tracked_outside_transaction() { let mut client = TestClient::new_sharded(Parameters::default()).await; @@ -63,122 +28,6 @@ async fn test_session_lock_tracked_outside_transaction() { assert!(client.engine.advisory_locks().contains(101)); } -#[tokio::test] -async fn test_session_lock_cleanup_resets_role_before_backend_reuse() { - load_single_connection_test_pool(); - let mut source = TestClient::new(Parameters::default()).await; - let role = format!("pgdog_cleanup_role_{}_advisory", std::process::id()); - - source - .send_simple(Query::new(format!("CREATE ROLE {role}"))) - .await; - source.read_until('Z').await.unwrap(); - source - .send_simple(Query::new("SELECT pg_advisory_lock(707519)")) - .await; - source.read_until('Z').await.unwrap(); - source - .send_simple(Query::new(format!("SET ROLE {role}"))) - .await; - source.read_until('Z').await.unwrap(); - - let before = identity(&mut source).await; - assert_eq!( - (before.1.as_str(), before.2.as_str()), - ("pgdog", role.as_str()) - ); - drop(source.leak_pool()); - - let mut peer = TestClient::new(Parameters::default()).await; - let after = identity(&mut peer).await; - assert_eq!(after.0, before.0, "physical backend should be reused"); - assert_eq!((after.1.as_str(), after.2.as_str()), ("pgdog", "pgdog")); - - peer.send_simple(Query::new(format!("DROP ROLE {role}"))) - .await; - peer.read_until('Z').await.unwrap(); -} - -#[tokio::test] -async fn test_deferred_role_cleanup_before_backend_reuse() { - load_single_connection_test_pool(); - let mut source = TestClient::new(Parameters::default()).await; - let role = format!("pgdog_cleanup_role_{}_deferred", std::process::id()); - - source - .send_simple(Query::new(format!("CREATE ROLE {role}"))) - .await; - source.read_until('Z').await.unwrap(); - source - .send_simple(Query::new(format!("SET ROLE {role}"))) - .await; - source.read_until('Z').await.unwrap(); - assert!( - !source.backend_connected(), - "SET ROLE should be deferred until a query checks out a backend" - ); - - let before = identity(&mut source).await; - assert_eq!( - (before.1.as_str(), before.2.as_str()), - ("pgdog", role.as_str()) - ); - assert!(!source.backend_connected()); - drop(source.leak_pool()); - - let mut peer = TestClient::new(Parameters::default()).await; - let after = identity(&mut peer).await; - assert_eq!(after.0, before.0, "physical backend should be reused"); - assert_eq!((after.1.as_str(), after.2.as_str()), ("pgdog", "pgdog")); - - peer.send_simple(Query::new(format!("DROP ROLE {role}"))) - .await; - peer.read_until('Z').await.unwrap(); -} - -#[tokio::test] -async fn test_deferred_transaction_role_cleanup_before_backend_reuse() { - load_single_connection_test_pool(); - let mut source = TestClient::new(Parameters::default()).await; - let role = format!("pgdog_cleanup_role_{}_transaction", std::process::id()); - - source - .send_simple(Query::new(format!("CREATE ROLE {role}"))) - .await; - source.read_until('Z').await.unwrap(); - source.send_simple(Query::new("BEGIN")).await; - source.read_until('Z').await.unwrap(); - source - .send_simple(Query::new(format!("SET ROLE {role}"))) - .await; - source.read_until('Z').await.unwrap(); - assert!( - !source.backend_connected(), - "BEGIN and SET ROLE should remain deferred until the first query" - ); - - let before = identity(&mut source).await; - assert_eq!( - (before.1.as_str(), before.2.as_str()), - ("pgdog", role.as_str()) - ); - assert!(source.backend_connected()); - - source.send_simple(Query::new("COMMIT")).await; - source.read_until('Z').await.unwrap(); - assert!(!source.backend_connected()); - drop(source.leak_pool()); - - let mut peer = TestClient::new(Parameters::default()).await; - let after = identity(&mut peer).await; - assert_eq!(after.0, before.0, "physical backend should be reused"); - assert_eq!((after.1.as_str(), after.2.as_str()), ("pgdog", "pgdog")); - - peer.send_simple(Query::new(format!("DROP ROLE {role}"))) - .await; - peer.read_until('Z').await.unwrap(); -} - #[tokio::test] async fn test_session_lock_inside_transaction_survives_commit() { // A plain pg_advisory_lock taken inside a transaction lives past COMMIT diff --git a/pgdog/src/frontend/client/query_engine/test/mod.rs b/pgdog/src/frontend/client/query_engine/test/mod.rs index 83cf247e7..99a837189 100644 --- a/pgdog/src/frontend/client/query_engine/test/mod.rs +++ b/pgdog/src/frontend/client/query_engine/test/mod.rs @@ -34,6 +34,7 @@ mod rewrite_insert_split; mod rewrite_offset; mod rewrite_simple_prepared; mod schema_changed; +mod session_identity; mod set; mod set_schema_sharding; mod sharded; diff --git a/pgdog/src/frontend/client/query_engine/test/session_identity.rs b/pgdog/src/frontend/client/query_engine/test/session_identity.rs new file mode 100644 index 000000000..494b4a348 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/session_identity.rs @@ -0,0 +1,290 @@ +use crate::{ + backend::databases::reload_from_existing, + config::{config, load_test, set}, + expect_message, + net::{CommandComplete, DataRow, ErrorResponse, ReadyForQuery, RowDescription}, +}; + +use super::prelude::*; + +fn load_single_connection_test_pool() { + load_test(); + + let mut config = (*config()).clone(); + config.config.general.default_pool_size = 1; + config.config.general.min_pool_size = 0; + set(config).unwrap(); + reload_from_existing().unwrap(); +} + +async fn identity(client: &mut TestClient) -> (i64, String, String) { + client + .send_simple(Query::new( + "SELECT pg_backend_pid(), session_user, current_user", + )) + .await; + expect_message!(client.read().await, RowDescription); + let row = expect_message!(client.read().await, DataRow); + expect_message!(client.read().await, CommandComplete); + expect_message!(client.read().await, ReadyForQuery); + + ( + row.get_int(0, true).expect("backend pid"), + row.get_text(1).expect("session_user"), + row.get_text(2).expect("current_user"), + ) +} + +async fn create_role(client: &mut TestClient, suffix: &str) -> String { + let role = format!("pgdog_cleanup_role_{}_{suffix}", std::process::id()); + client + .send_simple(Query::new(format!("CREATE ROLE {role}"))) + .await; + client.read_until('Z').await.unwrap(); + role +} + +async fn drop_role(client: &mut TestClient, role: &str) { + client + .send_simple(Query::new(format!("DROP ROLE {role}"))) + .await; + client.read_until('Z').await.unwrap(); +} + +fn assert_default_identity(identity: &(i64, String, String)) { + assert_eq!( + (identity.1.as_str(), identity.2.as_str()), + ("pgdog", "pgdog") + ); +} + +#[tokio::test] +async fn test_session_lock_cleanup_resets_role_before_backend_reuse() { + load_single_connection_test_pool(); + let mut source = TestClient::new(Parameters::default()).await; + let role = create_role(&mut source, "advisory").await; + + source + .send_simple(Query::new("SELECT pg_advisory_lock(707519)")) + .await; + source.read_until('Z').await.unwrap(); + source + .send_simple(Query::new(format!("SET ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + + let before = identity(&mut source).await; + assert_eq!( + (before.1.as_str(), before.2.as_str()), + ("pgdog", role.as_str()) + ); + drop(source.leak_pool()); + + let mut peer = TestClient::new(Parameters::default()).await; + let after = identity(&mut peer).await; + assert_eq!(after.0, before.0, "physical backend should be reused"); + assert_default_identity(&after); + drop_role(&mut peer, &role).await; +} + +#[tokio::test] +async fn test_deferred_role_reconciled_before_backend_reuse() { + load_single_connection_test_pool(); + let mut source = TestClient::new(Parameters::default()).await; + let role = create_role(&mut source, "deferred").await; + + source + .send_simple(Query::new(format!("SET ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + assert!( + !source.backend_connected(), + "SET ROLE should be deferred until a query checks out a backend" + ); + + let before = identity(&mut source).await; + assert_eq!( + (before.1.as_str(), before.2.as_str()), + ("pgdog", role.as_str()) + ); + assert!(!source.backend_connected()); + + source.send_simple(Query::new("RESET ALL")).await; + source.read_until('Z').await.unwrap(); + let after_reset_all = identity(&mut source).await; + assert_eq!( + after_reset_all, before, + "PostgreSQL RESET ALL must preserve SET ROLE" + ); + assert!(!source.backend_connected()); + drop(source.leak_pool()); + + let mut peer = TestClient::new(Parameters::default()).await; + let after = identity(&mut peer).await; + assert_eq!(after.0, before.0, "physical backend should be reused"); + assert_default_identity(&after); + drop_role(&mut peer, &role).await; +} + +#[tokio::test] +async fn test_session_authorization_reapplied_and_reconciled_before_backend_reuse() { + load_single_connection_test_pool(); + let mut source = TestClient::new(Parameters::default()).await; + let role = create_role(&mut source, "authorization").await; + + source + .send_simple(Query::new(format!("SET SESSION AUTHORIZATION {role}"))) + .await; + source.read_until('Z').await.unwrap(); + assert!(!source.backend_connected()); + + let before = identity(&mut source).await; + assert_eq!( + (before.1.as_str(), before.2.as_str()), + (role.as_str(), role.as_str()) + ); + assert!(!source.backend_connected()); + + let reapplied = identity(&mut source).await; + assert_eq!(reapplied, before); + assert!(!source.backend_connected()); + drop(source.leak_pool()); + + let mut peer = TestClient::new(Parameters::default()).await; + let after = identity(&mut peer).await; + assert_eq!(after.0, before.0, "physical backend should be reused"); + assert_default_identity(&after); + drop_role(&mut peer, &role).await; +} + +#[tokio::test] +async fn test_deferred_transaction_role_reconciled_before_backend_reuse() { + load_single_connection_test_pool(); + let mut source = TestClient::new(Parameters::default()).await; + let role = create_role(&mut source, "transaction").await; + + source.send_simple(Query::new("BEGIN")).await; + source.read_until('Z').await.unwrap(); + source + .send_simple(Query::new(format!("SET ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + assert!( + !source.backend_connected(), + "BEGIN and SET ROLE should remain deferred until the first query" + ); + + let before = identity(&mut source).await; + assert_eq!( + (before.1.as_str(), before.2.as_str()), + ("pgdog", role.as_str()) + ); + assert!(source.backend_connected()); + + source.send_simple(Query::new("COMMIT")).await; + source.read_until('Z').await.unwrap(); + assert!(!source.backend_connected()); + + let reapplied = identity(&mut source).await; + assert_eq!( + (reapplied.1.as_str(), reapplied.2.as_str()), + ("pgdog", role.as_str()) + ); + assert!(!source.backend_connected()); + drop(source.leak_pool()); + + let mut peer = TestClient::new(Parameters::default()).await; + let after = identity(&mut peer).await; + assert_eq!(after.0, before.0, "physical backend should be reused"); + assert_default_identity(&after); + drop_role(&mut peer, &role).await; +} + +#[tokio::test] +async fn test_transaction_role_rollback_restores_identity() { + load_single_connection_test_pool(); + let mut client = TestClient::new(Parameters::default()).await; + let role = create_role(&mut client, "rollback").await; + + client.send_simple(Query::new("BEGIN")).await; + client.read_until('Z').await.unwrap(); + let original = identity(&mut client).await; + + client + .send_simple(Query::new(format!("SET ROLE {role}"))) + .await; + client.read_until('Z').await.unwrap(); + let changed = identity(&mut client).await; + assert_eq!(changed.0, original.0); + assert_eq!( + (changed.1.as_str(), changed.2.as_str()), + ("pgdog", role.as_str()) + ); + + client.send_simple(Query::new("ROLLBACK")).await; + client.read_until('Z').await.unwrap(); + let restored = identity(&mut client).await; + assert_eq!(restored.0, original.0); + assert_default_identity(&restored); + drop_role(&mut client, &role).await; +} + +#[tokio::test] +async fn test_rejected_connected_role_restores_state_and_discards_backend() { + load_single_connection_test_pool(); + let mut client = TestClient::new(Parameters::default()).await; + + client.send_simple(Query::new("BEGIN")).await; + client.read_until('Z').await.unwrap(); + let before = identity(&mut client).await; + + client + .send_simple(Query::new("SET ROLE pgdog_role_does_not_exist")) + .await; + expect_message!(client.read().await, ErrorResponse); + expect_message!(client.read().await, ReadyForQuery); + + client.send_simple(Query::new("ROLLBACK")).await; + client.read_until('Z').await.unwrap(); + let after = identity(&mut client).await; + assert_default_identity(&after); + assert_ne!( + after.0, before.0, + "identity error should discard the ambiguous physical backend" + ); +} + +#[tokio::test] +async fn test_connected_transaction_role_reconciled_before_backend_reuse() { + load_single_connection_test_pool(); + let mut source = TestClient::new(Parameters::default()).await; + let role = create_role(&mut source, "connected").await; + + source.send_simple(Query::new("BEGIN")).await; + source.read_until('Z').await.unwrap(); + + let connected = identity(&mut source).await; + assert!(source.backend_connected()); + + source + .send_simple(Query::new(format!("SET ROLE {role}"))) + .await; + source.read_until('Z').await.unwrap(); + let changed = identity(&mut source).await; + assert_eq!(changed.0, connected.0); + assert_eq!( + (changed.1.as_str(), changed.2.as_str()), + ("pgdog", role.as_str()) + ); + + source.send_simple(Query::new("COMMIT")).await; + source.read_until('Z').await.unwrap(); + assert!(!source.backend_connected()); + drop(source.leak_pool()); + + let mut peer = TestClient::new(Parameters::default()).await; + let after = identity(&mut peer).await; + assert_eq!(after.0, connected.0, "physical backend should be reused"); + assert_default_identity(&after); + drop_role(&mut peer, &role).await; +} diff --git a/pgdog/src/frontend/router/parser/query/set.rs b/pgdog/src/frontend/router/parser/query/set.rs index 923da72a7..83637d1cf 100644 --- a/pgdog/src/frontend/router/parser/query/set.rs +++ b/pgdog/src/frontend/router/parser/query/set.rs @@ -23,7 +23,7 @@ impl QueryParser { .with_read(context.read_only), )) } else { - let param = Self::parse_set_param(stmt)?; + let param = Self::parse_set_param(stmt, context.query()?.query())?; Ok(Command::Set { params: vec![param], route: Route::write(context.shards_calculator.shard()), @@ -33,25 +33,38 @@ impl QueryParser { } /// Parse a single SET statement into a SetParam - fn parse_set_param(stmt: &nodes::VariableSetStmt) -> Result { - let value = if stmt.kind == VAR_SET_VALUE { + fn parse_set_param(stmt: &nodes::VariableSetStmt, query: &str) -> Result { + let mut value = if stmt.kind == VAR_SET_VALUE { Some(Self::parse_set_values(stmt)?) } else if stmt.kind == VAR_RESET || stmt.kind == VAR_SET_DEFAULT { None } else { panic!("parse_set_param called on invalid kind {}", stmt.kind); }; + let name = stmt.name().expect("SET always has name"); + + // PostgreSQL's raw parse tree normalizes both NONE and a quoted role + // named "none" to the same string. Preserve the keyword form by + // checking the original token at its parser-provided location. + if name == "role" + && value.as_ref().and_then(ParameterValue::as_str) == Some("none") + && let Some(Node::A_Const(constant)) = stmt.args().first() + && let Some(first) = query.as_bytes().get(constant.location as usize) + && !matches!(first, b'\'' | b'"') + { + value = None; + } match value { value @ Some(_) => Ok(SetParam { - name: stmt.name().expect("SET always has name").to_string(), + name: name.to_string(), value, local: stmt.is_local, }), None => Ok(SetParam { - name: stmt.name().expect("SET always has name").to_string(), + name: name.to_string(), value: None, - local: false, + local: stmt.is_local, }), } } @@ -70,12 +83,13 @@ impl QueryParser { context: &QueryParserContext, ) -> Result, Error> { let mut has_other = false; + let query = context.query()?.query(); let params = stmts .into_iter() .filter_map(|stmt| match stmt.stmt() { Node::VariableSetStmt(stmt) if stmt.kind != VAR_SET_MULTI => { - Some(Self::parse_set_param(stmt)) + Some(Self::parse_set_param(stmt, query)) } _ => { has_other = true; diff --git a/pgdog/src/frontend/router/parser/query/test/test_set.rs b/pgdog/src/frontend/router/parser/query/test/test_set.rs index 261b67322..eb936423c 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_set.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_set.rs @@ -102,6 +102,57 @@ fn test_set_role_is_tracked_as_parameter() { } } +#[test] +fn test_session_identity_commands_are_parameters() { + let mut test = QueryParserTest::new(); + + for (query, name, value) in [ + ( + "SET SESSION AUTHORIZATION other_user", + "session_authorization", + Some("other_user"), + ), + ( + "SET SESSION AUTHORIZATION DEFAULT", + "session_authorization", + None, + ), + ("RESET SESSION AUTHORIZATION", "session_authorization", None), + ("SET ROLE NONE", "role", None), + (r#"SET ROLE "none""#, "role", Some("none")), + ("RESET ROLE", "role", None), + ] { + let command = test.execute(vec![Query::new(query).into()]); + let Command::Set { params, .. } = command else { + panic!("expected Command::Set for {query}, got {command:#?}"); + }; + assert_eq!(params.len(), 1); + assert_eq!(params[0].name, name); + assert_eq!( + params[0].value.as_ref().and_then(ParameterValue::as_str), + value + ); + } +} + +#[test] +fn test_local_identity_reset_preserves_scope() { + let mut test = QueryParserTest::new(); + + for query in [ + "SET LOCAL ROLE NONE", + "SET LOCAL SESSION AUTHORIZATION DEFAULT", + ] { + let command = test.execute(vec![Query::new(query).into()]); + let Command::Set { params, .. } = command else { + panic!("expected Command::Set for {query}, got {command:#?}"); + }; + assert_eq!(params.len(), 1); + assert_eq!(params[0].value, None); + assert!(params[0].local); + } +} + #[test] fn test_set_config_null_value() { let mut test = QueryParserTest::new(); diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index 1edec4bcb..d75e67e24 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -32,6 +32,8 @@ static UNTRACKED_PARAMS: Lazy> = Lazy::new(|| { String::from("server_version"), String::from("server_encoding"), String::from("integer_datetimes"), + // Client SETs are tracked separately in SessionIdentity; don't replay + // the server-reported ParameterStatus as an ordinary GUC. String::from("session_authorization"), String::from("in_hot_standby"), String::from("pgdog.role"), @@ -162,6 +164,85 @@ impl ParameterValue { } } +const ROLE: &str = "role"; +const SESSION_AUTHORIZATION: &str = "session_authorization"; + +/// PostgreSQL session identity overrides. +/// +/// `None` means the authenticated connection user. +#[derive(Default, Debug, Clone, PartialEq, Eq)] +pub(crate) struct SessionIdentity { + session_authorization: Option, + role: Option, +} + +impl SessionIdentity { + fn parameter(name: &str) -> bool { + matches!(name, ROLE | SESSION_AUTHORIZATION) + } + + fn set(&mut self, name: &str, value: String) { + match name { + ROLE => self.role = Some(value), + SESSION_AUTHORIZATION => { + self.session_authorization = Some(value); + self.role = None; + } + _ => unreachable!("not a session identity parameter"), + } + } + + fn reset(&mut self, name: &str) { + match name { + ROLE => self.role = None, + SESSION_AUTHORIZATION => { + self.role = None; + self.session_authorization = None; + } + _ => unreachable!("not a session identity parameter"), + } + } + + fn set_parameter(&mut self, name: &str, value: &str) { + self.set(name, value.to_owned()); + } + + /// Queries that transform this identity into `target`. + pub(crate) fn reconcile(&self, target: &Self, local: bool) -> Vec { + if self == target { + return vec![]; + } + + let set = if local { "SET LOCAL" } else { "SET" }; + let mut queries = vec![Query::new(format!("{set} SESSION AUTHORIZATION DEFAULT"))]; + + if let Some(session_authorization) = &target.session_authorization { + queries.push(Query::new(format!( + "{set} SESSION AUTHORIZATION {}", + ParameterValue::String(session_authorization.clone()) + ))); + } + + if let Some(role) = &target.role { + queries.push(Query::new(format!( + "{set} ROLE {}", + ParameterValue::String(role.clone()) + ))); + } + + queries + } +} + +impl MemoryUsage for SessionIdentity { + fn memory_usage(&self) -> usize { + self.session_authorization + .as_ref() + .map_or(0, MemoryUsage::memory_usage) + + self.role.as_ref().map_or(0, MemoryUsage::memory_usage) + } +} + /// List of parameters. #[derive(Default, Debug, Clone, PartialEq)] pub(crate) struct Parameters { @@ -176,10 +257,16 @@ pub(crate) struct Parameters { /// what but we need to intercept them for databases that have cross shard /// queries disabled. transaction_local_params: BTreeMap, - /// Hash of `params` to avoid syncing params between clients and servers - /// when they are the same. /// Reset params. Stored here to support ROLLBACK. reset_params: BTreeMap, + /// Session identity committed outside a transaction. + identity: SessionIdentity, + /// Session identity changed by SET inside a transaction. + transaction_identity: Option>, + /// Session identity changed by SET LOCAL inside a transaction. + transaction_local_identity: Option>, + /// Hash of `params` to avoid syncing params between clients and servers + /// when they are the same. hash: u64, } @@ -198,7 +285,18 @@ impl Display for Parameters { impl MemoryUsage for Parameters { #[inline] fn memory_usage(&self) -> usize { - self.params.memory_usage() + self.hash.memory_usage() + self.params.memory_usage() + + self.identity.memory_usage() + + self.transaction_identity.as_ref().map_or(0, |identity| { + std::mem::size_of::() + identity.memory_usage() + }) + + self + .transaction_local_identity + .as_ref() + .map_or(0, |identity| { + std::mem::size_of::() + identity.memory_usage() + }) + + self.hash.memory_usage() } } @@ -223,6 +321,18 @@ impl Parameters { self.hash = Self::compute_hash(&self.params); } + /// Restore the authenticated PostgreSQL identity. + pub(crate) fn clear_session_identity(&mut self) { + self.identity = SessionIdentity::default(); + self.transaction_identity = None; + self.transaction_local_identity = None; + } + + /// Whether a parameter changes PostgreSQL session identity. + pub(crate) fn is_session_identity(name: &str) -> bool { + SessionIdentity::parameter(&name.to_lowercase()) + } + /// Get parameter. pub(crate) fn get(&self, name: &str) -> Option<&ParameterValue> { if let Some(param) = self.transaction_local_params.get(name) { @@ -234,6 +344,111 @@ impl Parameters { } } + /// Store a PostgreSQL session identity parameter outside or inside a transaction. + /// + /// Returns true when `name` is a session identity parameter. + pub(crate) fn insert_identity( + &mut self, + name: &str, + value: &ParameterValue, + transaction: bool, + local: bool, + ) -> bool { + let name = name.to_lowercase(); + if !SessionIdentity::parameter(&name) { + return false; + } + if local && !transaction { + return true; + } + + let value = value + .as_str() + .expect("session identity has exactly one string value"); + + if transaction { + let base = self + .transaction_identity + .as_deref() + .unwrap_or(&self.identity) + .clone(); + if local { + self.transaction_local_identity + .get_or_insert_with(|| Box::new(base)) + .set_parameter(&name, value); + } else { + self.transaction_identity + .get_or_insert_with(|| Box::new(base)) + .set_parameter(&name, value); + if let Some(identity) = self.transaction_local_identity.as_mut() { + identity.set_parameter(&name, value); + } + } + } else { + self.identity.set_parameter(&name, value); + } + + true + } + + /// Reset a PostgreSQL session identity parameter. + /// + /// Returns true when `name` is a session identity parameter. + pub(crate) fn reset_identity(&mut self, name: &str, transaction: bool, local: bool) -> bool { + let name = name.to_lowercase(); + if !SessionIdentity::parameter(&name) { + return false; + } + if local && !transaction { + return true; + } + + if transaction { + let base = self + .transaction_identity + .as_deref() + .unwrap_or(&self.identity) + .clone(); + if local { + self.transaction_local_identity + .get_or_insert_with(|| Box::new(base)) + .reset(&name); + } else { + self.transaction_identity + .get_or_insert_with(|| Box::new(base)) + .reset(&name); + if let Some(identity) = self.transaction_local_identity.as_mut() { + identity.reset(&name); + } + } + } else { + self.identity.reset(&name); + } + + true + } + + /// Current session identity, including transaction overrides when requested. + pub(crate) fn session_identity(&self, transaction: bool) -> SessionIdentity { + if !transaction { + return self.identity.clone(); + } + + self.transaction_local_identity + .as_deref() + .or(self.transaction_identity.as_deref()) + .unwrap_or(&self.identity) + .clone() + } + + pub(crate) fn transaction_session_identity(&self) -> Option<&SessionIdentity> { + self.transaction_identity.as_deref() + } + + pub(crate) fn local_session_identity(&self) -> Option<&SessionIdentity> { + self.transaction_local_identity.as_deref() + } + /// Insert a parameter, but only for the duration of the transaction. pub(crate) fn insert_transaction( &mut self, @@ -267,6 +482,9 @@ impl Parameters { /// dropping everything changed since with `SET`. pub(crate) fn restore_startup(&mut self, startup: &Parameters) { self.params.clone_from(&startup.params); + self.identity.clone_from(&startup.identity); + self.transaction_identity = None; + self.transaction_local_identity = None; self.reset_params.clear(); self.hash = Self::compute_hash(&self.params); } @@ -280,7 +498,7 @@ impl Parameters { keys.dedup(); for key in keys { - if !UNTRACKED_PARAMS.contains(&key) { + if !UNTRACKED_PARAMS.contains(&key) && !SessionIdentity::parameter(&key) { self.reset(&key); } } @@ -292,11 +510,17 @@ impl Parameters { "saved {} in-transaction params", self.transaction_params.len() ); - let changed = !self.transaction_params.is_empty() || !self.reset_params.is_empty(); + let changed = !self.transaction_params.is_empty() + || !self.reset_params.is_empty() + || self.transaction_identity.is_some(); + if let Some(identity) = self.transaction_identity.take() { + self.identity = *identity; + } self.params .extend(std::mem::take(&mut self.transaction_params)); self.transaction_local_params.clear(); + self.transaction_local_identity = None; self.reset_params.clear(); if changed { @@ -310,12 +534,12 @@ impl Parameters { pub(crate) fn rollback(&mut self) { self.transaction_params.clear(); self.transaction_local_params.clear(); + self.transaction_identity = None; + self.transaction_local_identity = None; - let mut reset = false; - for (name, value) in std::mem::take(&mut self.reset_params) { - self.params.insert(name, value); - reset = true; - } + let reset_params = std::mem::take(&mut self.reset_params); + let reset = !reset_params.is_empty(); + self.params.extend(reset_params); if reset { self.hash = Self::compute_hash(&self.params); @@ -357,6 +581,7 @@ impl Parameters { Self { params, + identity: self.identity.clone(), hash, ..Default::default() } @@ -390,7 +615,7 @@ impl Parameters { /// Merge params from self into other, generating the queries /// needed to sync that state on the server. pub(crate) fn identical(&self, other: &Self) -> bool { - self.hash == other.hash + self.hash == other.hash && self.identity == other.identity } /// Generate SET queries to change server state. @@ -409,12 +634,14 @@ impl Parameters { let mut sets = self .transaction_params .iter() + .filter(|(key, _)| !SessionIdentity::parameter(key)) .map(|(key, value)| query(key, value, false)) .collect::>(); sets.extend( self.transaction_local_params .iter() + .filter(|(key, _)| !SessionIdentity::parameter(key)) .map(|(key, value)| query(key, value, true)), ); @@ -458,6 +685,10 @@ impl Parameters { /// Copy params set inside the transaction. pub(crate) fn copy_in_transaction(&mut self, other: &Self) { + self.transaction_identity + .clone_from(&other.transaction_identity); + self.transaction_local_identity + .clone_from(&other.transaction_local_identity); self.transaction_params.extend( other .transaction_params @@ -505,6 +736,9 @@ impl From> for Parameters { transaction_params: BTreeMap::new(), transaction_local_params: BTreeMap::new(), reset_params: BTreeMap::new(), + identity: SessionIdentity::default(), + transaction_identity: None, + transaction_local_identity: None, } } } @@ -529,7 +763,7 @@ mod test { use crate::net::ToBytes; use crate::net::parameter::ParameterValue; - use super::Parameters; + use super::{Parameters, SessionIdentity}; #[test] fn test_identical() { @@ -550,6 +784,164 @@ mod test { assert!(Parameters::default().identical(&Parameters::default())); } + #[test] + fn test_session_identity_reconcile_order() { + let current = SessionIdentity { + session_authorization: Some("old_auth".into()), + role: Some("old_role".into()), + }; + let target = SessionIdentity { + session_authorization: Some("new_auth".into()), + role: Some("new_role".into()), + }; + + let queries = current + .reconcile(&target, false) + .into_iter() + .map(|query| query.query().to_owned()) + .collect::>(); + + assert_eq!( + queries, + [ + "SET SESSION AUTHORIZATION DEFAULT", + r#"SET SESSION AUTHORIZATION "new_auth""#, + r#"SET ROLE "new_role""#, + ] + ); + + let queries = target + .reconcile(&SessionIdentity::default(), true) + .into_iter() + .map(|query| query.query().to_owned()) + .collect::>(); + assert_eq!(queries, ["SET LOCAL SESSION AUTHORIZATION DEFAULT"]); + } + + #[test] + fn test_session_identity_commit_and_rollback() { + let mut params = Parameters::default(); + params.insert_identity("role", &"base_role".into(), false, false); + params.insert_identity("role", &"transaction_role".into(), true, false); + + assert_eq!( + params.session_identity(true).role.as_deref(), + Some("transaction_role") + ); + params.rollback(); + assert_eq!( + params.session_identity(false).role.as_deref(), + Some("base_role") + ); + + params.insert_identity("role", &"committed_role".into(), true, false); + params.insert_identity("role", &"local_role".into(), true, true); + assert_eq!( + params.session_identity(true).role.as_deref(), + Some("local_role") + ); + params.commit(); + assert_eq!( + params.session_identity(false).role.as_deref(), + Some("committed_role") + ); + } + + #[test] + fn test_session_identity_respects_set_local_order() { + let mut params = Parameters::default(); + params.insert_identity("role", &"local_first".into(), true, true); + params.insert_identity("role", &"session_last".into(), true, false); + assert_eq!( + params.session_identity(true).role.as_deref(), + Some("session_last") + ); + params.commit(); + assert_eq!( + params.session_identity(false).role.as_deref(), + Some("session_last") + ); + + params.insert_identity("role", &"session_first".into(), true, false); + params.insert_identity("role", &"local_last".into(), true, true); + assert_eq!( + params.session_identity(true).role.as_deref(), + Some("local_last") + ); + params.commit(); + assert_eq!( + params.session_identity(false).role.as_deref(), + Some("session_first") + ); + + params.reset_identity("role", true, true); + assert_eq!(params.session_identity(true).role, None); + params.commit(); + assert_eq!( + params.session_identity(false).role.as_deref(), + Some("session_first") + ); + } + + #[test] + fn test_set_local_identity_outside_transaction_is_not_persisted() { + let mut params = Parameters::default(); + params.insert_identity("role", &"ignored".into(), false, true); + assert_eq!(params.session_identity(false), SessionIdentity::default()); + } + + #[test] + fn test_reset_session_authorization_is_transactional_and_clears_role() { + let mut params = Parameters::default(); + params.insert_identity( + "session_authorization", + &"delegated_user".into(), + false, + false, + ); + params.insert_identity("role", &"reporting".into(), false, false); + + params.reset_identity("session_authorization", true, false); + assert_eq!(params.session_identity(true), SessionIdentity::default()); + params.rollback(); + assert_eq!( + params.session_identity(false), + SessionIdentity { + session_authorization: Some("delegated_user".into()), + role: Some("reporting".into()), + } + ); + + params.reset_identity("session_authorization", true, false); + params.commit(); + assert_eq!(params.session_identity(false), SessionIdentity::default()); + } + + #[test] + fn test_reset_all_preserves_session_identity() { + let mut params = Parameters::default(); + params.insert("work_mem", "1MB"); + params.insert_identity("role", &"reporting".into(), false, false); + + params.reset_all(); + + assert!(params.get("work_mem").is_none()); + assert_eq!( + params.session_identity(false).role.as_deref(), + Some("reporting") + ); + } + + #[test] + fn test_role_named_none_is_distinct_from_reset() { + let mut params = Parameters::default(); + params.insert_identity("role", &"none".into(), false, false); + + assert_eq!(params.session_identity(false).role.as_deref(), Some("none")); + params.reset_identity("role", false, false); + assert_eq!(params.session_identity(false).role, None); + } + #[test] fn test_tracked_and_different() { let mut client = Parameters::default(); From 612bb1e3ff1f2adab161aff52c4f95a73d581555 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 28 Sep 2026 13:35:47 +0530 Subject: [PATCH 07/12] fix test --- pgdog/src/backend/server.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index fbf829937..5263e5c73 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -738,6 +738,9 @@ impl Server { if !queries.is_empty() { debug!("syncing {} params", queries.len()); + if identity_changed { + self.mark_dirty(true); + } if let Err(err) = self.execute_batch(&queries).await { if identity_changed { self.force_close(); @@ -775,6 +778,9 @@ impl Server { if !transaction_sets.is_empty() { debug!("syncing {} in-transaction params", transaction_sets.len()); + if identity_changed { + self.mark_dirty(true); + } if let Err(err) = self.execute_batch(&transaction_sets).await { if identity_changed { self.force_close(); @@ -2545,7 +2551,7 @@ pub(crate) mod test { .await?; assert_eq!(changed, 2); - assert!(!server.dirty()); + assert!(server.dirty()); let changed = server .link_client(FrontendPid::new(), ¶ms, None) @@ -2583,7 +2589,7 @@ pub(crate) mod test { .await?; assert_eq!(changed, 2); - assert!(!server.dirty()); + assert!(server.dirty()); server.rollback().await?; Ok(()) From f062ea9a03cef091490b307cb82a7ac30e0612eb Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 28 Sep 2026 14:39:42 +0530 Subject: [PATCH 08/12] fix --- pgdog/src/backend/server.rs | 5 ++++- pgdog/src/net/parameter.rs | 21 +++++++++++++++++++-- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 5263e5c73..bfbf41128 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -2517,7 +2517,10 @@ pub(crate) mod test { let changed = server .link_client(FrontendPid::new(), ¶ms, None) .await?; - assert_eq!(changed, 0); + assert_eq!( + changed, 2, + "identity must be reapplied from the authenticated baseline" + ); for i in 0..25 { let value = format!("apples_{}", i); diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index d75e67e24..57cb4891a 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -570,7 +570,11 @@ impl Parameters { .filter(|(k, _)| !UNTRACKED_PARAMS.contains(k)) } - /// Filter our parameters that we would track with SET queries. + /// Filter parameters tracked on a pooled server. + /// + /// Session identity is deliberately omitted: any backend that receives an + /// identity override is marked dirty and restored to authenticated defaults + /// before it returns to the pool. pub(crate) fn tracked(&self) -> Parameters { let params = self .tracked_iter() @@ -581,7 +585,6 @@ impl Parameters { Self { params, - identity: self.identity.clone(), hash, ..Default::default() } @@ -932,6 +935,20 @@ mod test { ); } + #[test] + fn test_server_snapshot_omits_session_identity() { + let mut params = Parameters::default(); + params.insert_identity("role", &"reporting".into(), false, false); + + let snapshot = params.tracked(); + + assert_eq!( + params.session_identity(false).role.as_deref(), + Some("reporting") + ); + assert_eq!(snapshot.session_identity(false), SessionIdentity::default()); + } + #[test] fn test_role_named_none_is_distinct_from_reset() { let mut params = Parameters::default(); From 615d85828d80a92da234ff3ccadac4735b3ae2ea Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 28 Sep 2026 14:58:30 +0530 Subject: [PATCH 09/12] fix --- pgdog/src/backend/server.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index bfbf41128..2c2510926 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -2517,10 +2517,7 @@ pub(crate) mod test { let changed = server .link_client(FrontendPid::new(), ¶ms, None) .await?; - assert_eq!( - changed, 2, - "identity must be reapplied from the authenticated baseline" - ); + assert_eq!(changed, 0); for i in 0..25 { let value = format!("apples_{}", i); @@ -2559,7 +2556,10 @@ pub(crate) mod test { let changed = server .link_client(FrontendPid::new(), ¶ms, None) .await?; - assert_eq!(changed, 0); + assert_eq!( + changed, 2, + "identity must be reapplied from the authenticated baseline" + ); Ok(()) } From d53f4fd2721e9313f2cb4e62cb8dbfdcbe267fe8 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 28 Sep 2026 15:34:05 +0530 Subject: [PATCH 10/12] preserve session identity on reset --- pgdog/src/frontend/client/query_engine/set.rs | 4 ++- .../query_engine/test/session_identity.rs | 5 ++++ pgdog/src/net/parameter.rs | 29 ++++++++++++++++++- 3 files changed, 36 insertions(+), 2 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index f4ac42846..e696186c6 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -137,7 +137,9 @@ impl QueryEngine { if context.in_transaction() || self.backend.connected() { context.params.reset_all(); } else { - context.params.restore_startup(context.startup_params); + context + .params + .restore_startup_parameters(context.startup_params); self.comms.update_params(context.params); } diff --git a/pgdog/src/frontend/client/query_engine/test/session_identity.rs b/pgdog/src/frontend/client/query_engine/test/session_identity.rs index 494b4a348..fe637a1a2 100644 --- a/pgdog/src/frontend/client/query_engine/test/session_identity.rs +++ b/pgdog/src/frontend/client/query_engine/test/session_identity.rs @@ -111,6 +111,11 @@ async fn test_deferred_role_reconciled_before_backend_reuse() { source.send_simple(Query::new("RESET ALL")).await; source.read_until('Z').await.unwrap(); + assert_eq!( + source.client().params.session_identity(false).role(), + Some(role.as_str()), + "RESET ALL must preserve the frontend's logical role" + ); let after_reset_all = identity(&mut source).await; assert_eq!( after_reset_all, before, diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index 3af15716a..b9817c458 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -206,6 +206,11 @@ impl SessionIdentity { self.set(name, value.to_owned()); } + #[cfg(test)] + pub(crate) fn role(&self) -> Option<&str> { + self.role.as_deref() + } + /// Queries that transform this identity into `target`. pub(crate) fn reconcile(&self, target: &Self, local: bool) -> Vec { if self == target { @@ -479,10 +484,15 @@ impl Parameters { /// Restore parameters to the values supplied in the startup message, /// dropping everything changed since with `SET`. pub(crate) fn restore_startup(&mut self, startup: &Parameters) { - self.params.clone_from(&startup.params); + self.restore_startup_parameters(startup); self.identity.clone_from(&startup.identity); self.transaction_identity = None; self.transaction_local_identity = None; + } + + /// Restore ordinary startup parameters while preserving session identity. + pub(crate) fn restore_startup_parameters(&mut self, startup: &Parameters) { + self.params.clone_from(&startup.params); self.reset_params.clear(); self.hash = Self::compute_hash(&self.params); } @@ -947,6 +957,23 @@ mod test { assert_eq!(snapshot.session_identity(false), SessionIdentity::default()); } + #[test] + fn test_restore_startup_parameters_preserves_session_identity() { + let mut startup = Parameters::default(); + startup.insert("search_path", "public"); + + let mut params = Parameters::default(); + params.insert("search_path", "private"); + params.insert_identity("role", &"reporting".into(), false, false); + params.restore_startup_parameters(&startup); + + assert_eq!(params.get("search_path"), startup.get("search_path")); + assert_eq!( + params.session_identity(false).role.as_deref(), + Some("reporting") + ); + } + #[test] fn test_role_named_none_is_distinct_from_reset() { let mut params = Parameters::default(); From 5fb3c980da578ecad2c9cc3f7d962f76cb86b8d0 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 28 Sep 2026 23:47:00 +0530 Subject: [PATCH 11/12] cleanup --- pgdog/src/backend/pool/connection/binding.rs | 28 ++--- pgdog/src/backend/pool/connection/mod.rs | 5 - pgdog/src/backend/server.rs | 81 ++++--------- pgdog/src/frontend/client/query_engine/set.rs | 25 +---- .../query_engine/test/session_identity.rs | 7 +- pgdog/src/net/parameter.rs | 106 ++++++++++++------ 6 files changed, 115 insertions(+), 137 deletions(-) diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index adb4187b3..edc52963f 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -430,21 +430,11 @@ impl Binding { let results = join_all(futures).await; let mut max = 0; - let mut error = None; for result in results { - match result { - Ok(synced) => max = max.max(synced), - Err(err) => error = Some(err), - } - } - if let Some(error) = error { - for server in servers { - server.force_close(); - } - Err(error) - } else { - Ok(max) + let synced = result?; + max = max.max(synced); } + Ok(max) } _ => Ok(0), @@ -462,6 +452,18 @@ impl Binding { } } + pub(crate) fn sync_client_params(&mut self, params: &Parameters) { + match self { + Binding::Direct(server, ..) => server.sync_client_params(params), + Binding::MultiShard(servers, _) => { + for server in servers { + server.sync_client_params(params); + } + } + _ => (), + } + } + pub(crate) fn changed_params(&mut self) -> Parameters { match self { Binding::Direct(server, ..) => server.changed_params().clone(), diff --git a/pgdog/src/backend/pool/connection/mod.rs b/pgdog/src/backend/pool/connection/mod.rs index 34c69b941..2ed8ccb8b 100644 --- a/pgdog/src/backend/pool/connection/mod.rs +++ b/pgdog/src/backend/pool/connection/mod.rs @@ -456,11 +456,6 @@ impl Connection { } } - /// Mark connected servers as containing client-specific session state. - pub(crate) fn mark_dirty(&mut self) { - self.binding.dirty(); - } - /// Check if any held server connection is currently locked to a client. #[cfg(test)] pub(crate) fn locked(&self) -> bool { diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index ff3fc5ef5..746675f64 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -721,15 +721,9 @@ impl Server { let mut executed = if !params.identical(&self.client_params) { // Construct client parameter SET queries. let tracked = params.tracked_and_different(&self.client_params); - let identity_queries = self - .client_params - .session_identity(false) - .reconcile(¶ms.session_identity(false), false); - let identity_changed = !identity_queries.is_empty(); - let mut queries = identity_queries; // Construct RESET queries to reset any current params - // to their default values. - queries.extend(self.client_params.reset_queries(params)); + // to their default values. Session identity is included here. + let mut queries = self.client_params.reset_queries(params); // Combine both to create a new, fresh session state // on this connection. @@ -740,15 +734,7 @@ impl Server { if !queries.is_empty() { debug!("syncing {} params", queries.len()); - if identity_changed { - self.mark_dirty(true); - } - if let Err(err) = self.execute_batch(&queries).await { - if identity_changed { - self.force_close(); - } - return Err(err); - } + self.execute_batch(&queries).await?; clear_params = true; } @@ -764,31 +750,12 @@ impl Server { // need to be revered on rollback or commited on commit. if let Some(start_transaction) = start_transaction { self.execute(start_transaction).await?; - let mut current_identity = self.client_params.session_identity(false); - let mut identity_queries = vec![]; - if let Some(identity) = params.transaction_session_identity() { - identity_queries.extend(current_identity.reconcile(identity, false)); - current_identity.clone_from(identity); - } - if let Some(identity) = params.local_session_identity() { - identity_queries.extend(current_identity.reconcile(identity, true)); - } - let identity_changed = !identity_queries.is_empty(); - let mut transaction_sets = identity_queries; - transaction_sets.extend(params.set_queries(true)); + let transaction_sets = params.set_queries(true); if !transaction_sets.is_empty() { debug!("syncing {} in-transaction params", transaction_sets.len()); - if identity_changed { - self.mark_dirty(true); - } - if let Err(err) = self.execute_batch(&transaction_sets).await { - if identity_changed { - self.force_close(); - } - return Err(err); - } + self.execute_batch(&transaction_sets).await?; clear_params = true; self.client_params.copy_in_transaction(params); @@ -1165,6 +1132,15 @@ impl Server { self.client_params.clear_session_identity(); } + /// Record the client's current parameters on this backend. + /// + /// Used after a connected `SET` so the next checkout can diff against + /// the identity and GUCs this connection actually has. + pub(crate) fn sync_client_params(&mut self, params: &Parameters) { + self.client_params = params.tracked(); + self.client_params.copy_in_transaction(params); + } + pub(crate) fn reset_re_synced(&mut self) { self.re_synced = false; } @@ -2536,33 +2512,18 @@ pub(crate) mod test { .await?; assert_eq!(changed, 2); - assert!(server.dirty()); + assert!(!server.dirty()); let changed = server .link_client(FrontendPid::new(), ¶ms, None) .await?; - assert_eq!( - changed, 2, - "identity must be reapplied from the authenticated baseline" - ); + assert_eq!(changed, 0, "identity stays on the server snapshot"); - Ok(()) - } + let peer = Parameters::default(); + let changed = server.link_client(FrontendPid::new(), &peer, None).await?; + assert_eq!(changed, 1, "next client resets identity from the snapshot"); - #[tokio::test] - async fn test_link_client_force_closes_after_identity_batch_failure() { - let mut params = Parameters::default(); - params.insert_identity("role", &"pgdog".into(), false, false); - params.insert("work_mem", "not-a-size"); - - let mut server = test_server().await; - assert!( - server - .link_client(FrontendPid::new(), ¶ms, None) - .await - .is_err() - ); - assert!(server.is_force_close()); + Ok(()) } #[tokio::test] @@ -2577,7 +2538,7 @@ pub(crate) mod test { .await?; assert_eq!(changed, 2); - assert!(server.dirty()); + assert!(!server.dirty()); server.rollback().await?; Ok(()) diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index e696186c6..f48fa3215 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -1,7 +1,6 @@ use crate::frontend::SetParam; use crate::frontend::client::query_engine::fake::FakeResponse; use crate::frontend::router::parameter_hints::{PGDOG_PIN, PGDOG_SHARD, PGDOG_SHARDING_KEY}; -use crate::frontend::router::parser::ShardWithPriority; use crate::net::messages::ErrorResponse; use super::*; @@ -25,11 +24,6 @@ impl QueryEngine { return Ok(()); } - let identity_changed = params - .iter() - .any(|param| Parameters::is_session_identity(¶m.name)); - let identity_before = (identity_changed && self.backend.connected()) - .then(|| Box::new(context.params.clone())); let mut fake_command = "SET"; for param in params { let is_pin = param.name == PGDOG_PIN; @@ -77,23 +71,10 @@ impl QueryEngine { } if self.backend.connected() { - if identity_changed { - context.client_request.route = Some(Route::write( - ShardWithPriority::new_override_transaction(Shard::All), - )); - self.backend.mark_dirty(); - } - let result = self.execute(context, None).await; - if result.is_err() || self.last_server_error { - if let Some(params) = identity_before { - context.params.clone_from(¶ms); - self.comms.update_params(context.params); - } - if self.backend.connected() { - self.backend.force_close(); - } + self.execute(context, None).await?; + if !self.last_server_error { + self.backend.sync_client_params(context.params); } - result?; } else { let fake_response = set_config .then(|| params.iter().map(|p| p.value.as_ref())) diff --git a/pgdog/src/frontend/client/query_engine/test/session_identity.rs b/pgdog/src/frontend/client/query_engine/test/session_identity.rs index fe637a1a2..be28b40c8 100644 --- a/pgdog/src/frontend/client/query_engine/test/session_identity.rs +++ b/pgdog/src/frontend/client/query_engine/test/session_identity.rs @@ -235,7 +235,7 @@ async fn test_transaction_role_rollback_restores_identity() { } #[tokio::test] -async fn test_rejected_connected_role_restores_state_and_discards_backend() { +async fn test_rejected_connected_role_restores_identity_on_rollback() { load_single_connection_test_pool(); let mut client = TestClient::new(Parameters::default()).await; @@ -252,11 +252,8 @@ async fn test_rejected_connected_role_restores_state_and_discards_backend() { client.send_simple(Query::new("ROLLBACK")).await; client.read_until('Z').await.unwrap(); let after = identity(&mut client).await; + assert_eq!(after.0, before.0); assert_default_identity(&after); - assert_ne!( - after.0, before.0, - "identity error should discard the ambiguous physical backend" - ); } #[tokio::test] diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index b9817c458..6b782f25a 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -34,6 +34,7 @@ static UNTRACKED_PARAMS: Lazy> = Lazy::new(|| { String::from("integer_datetimes"), // Client SETs are tracked separately in SessionIdentity; don't replay // the server-reported ParameterStatus as an ordinary GUC. + String::from("role"), String::from("session_authorization"), String::from("in_hot_standby"), String::from("pgdog.role"), @@ -212,7 +213,7 @@ impl SessionIdentity { } /// Queries that transform this identity into `target`. - pub(crate) fn reconcile(&self, target: &Self, local: bool) -> Vec { + fn reconcile(&self, target: &Self, local: bool) -> Vec { if self == target { return vec![]; } @@ -331,11 +332,6 @@ impl Parameters { self.transaction_local_identity = None; } - /// Whether a parameter changes PostgreSQL session identity. - pub(crate) fn is_session_identity(name: &str) -> bool { - SessionIdentity::parameter(&name.to_lowercase()) - } - /// Get parameter. pub(crate) fn get(&self, name: &str) -> Option<&ParameterValue> { if let Some(param) = self.transaction_local_params.get(name) { @@ -432,6 +428,7 @@ impl Parameters { } /// Current session identity, including transaction overrides when requested. + #[cfg(test)] pub(crate) fn session_identity(&self, transaction: bool) -> SessionIdentity { if !transaction { return self.identity.clone(); @@ -444,14 +441,6 @@ impl Parameters { .clone() } - pub(crate) fn transaction_session_identity(&self) -> Option<&SessionIdentity> { - self.transaction_identity.as_deref() - } - - pub(crate) fn local_session_identity(&self) -> Option<&SessionIdentity> { - self.transaction_local_identity.as_deref() - } - /// Insert a parameter, but only for the duration of the transaction. pub(crate) fn insert_transaction( &mut self, @@ -579,10 +568,6 @@ impl Parameters { } /// Filter parameters tracked on a pooled server. - /// - /// Session identity is deliberately omitted: any backend that receives an - /// identity override is marked dirty and restored to authenticated defaults - /// before it returns to the pool. pub(crate) fn tracked(&self) -> Parameters { let params = self .tracked_iter() @@ -594,6 +579,7 @@ impl Parameters { Self { params, hash, + identity: self.identity.clone(), ..Default::default() } } @@ -642,12 +628,22 @@ impl Parameters { } if transaction_only { - let mut sets = self - .transaction_params - .iter() - .filter(|(key, _)| !SessionIdentity::parameter(key)) - .map(|(key, value)| query(key, value, false)) - .collect::>(); + let mut current = self.identity.clone(); + let mut sets = vec![]; + if let Some(identity) = self.transaction_identity.as_deref() { + sets.extend(current.reconcile(identity, false)); + current.clone_from(identity); + } + if let Some(identity) = self.transaction_local_identity.as_deref() { + sets.extend(current.reconcile(identity, true)); + } + + sets.extend( + self.transaction_params + .iter() + .filter(|(key, _)| !SessionIdentity::parameter(key)) + .map(|(key, value)| query(key, value, false)), + ); sets.extend( self.transaction_local_params @@ -660,6 +656,7 @@ impl Parameters { } else { self.params .iter() + .filter(|(key, _)| !SessionIdentity::parameter(key)) .map(|(key, value)| query(key, value, false)) .collect() } @@ -674,11 +671,16 @@ impl Parameters { /// have a value on the incoming client. /// pub(crate) fn reset_queries(&self, other: &Self) -> Vec { - self.params - .keys() - .filter(|name| !other.contains_key(*name)) - .map(|name| Query::new(format!(r#"RESET "{}""#, name))) - .collect() + let mut queries = self.identity.reconcile(&other.identity, false); + queries.extend( + self.params + .keys() + .filter(|name| { + !other.contains_key(*name) && !SessionIdentity::parameter(name.as_str()) + }) + .map(|name| Query::new(format!(r#"RESET "{}""#, name))), + ); + queries } /// Get parameter value or returned an error. @@ -944,17 +946,57 @@ mod test { } #[test] - fn test_server_snapshot_omits_session_identity() { + fn test_server_snapshot_keeps_session_identity() { let mut params = Parameters::default(); params.insert_identity("role", &"reporting".into(), false, false); let snapshot = params.tracked(); assert_eq!( - params.session_identity(false).role.as_deref(), + snapshot.session_identity(false).role.as_deref(), Some("reporting") ); - assert_eq!(snapshot.session_identity(false), SessionIdentity::default()); + } + + #[test] + fn test_reset_queries_include_session_identity() { + let mut server = Parameters::default(); + server.insert_identity("role", &"reporting".into(), false, false); + server.insert("work_mem", "1MB"); + + let client = Parameters::default(); + let queries = server + .reset_queries(&client) + .into_iter() + .map(|query| query.query().to_owned()) + .collect::>(); + + assert_eq!( + queries, + ["SET SESSION AUTHORIZATION DEFAULT", r#"RESET "work_mem""#,] + ); + } + + #[test] + fn test_set_queries_include_transaction_identity() { + let mut params = Parameters::default(); + params.insert_identity("role", &"reporting".into(), true, false); + params.insert_transaction("work_mem", "1MB", false); + + let queries = params + .set_queries(true) + .into_iter() + .map(|query| query.query().to_owned()) + .collect::>(); + + assert_eq!( + queries, + [ + "SET SESSION AUTHORIZATION DEFAULT", + r#"SET ROLE "reporting""#, + r#"SET "work_mem" TO "1MB""#, + ] + ); } #[test] From b9b19dcfa37eab1c9300ba0d2a114c5df9032010 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 29 Sep 2026 00:12:40 +0530 Subject: [PATCH 12/12] clippy --- .../client/query_engine/test/schema_changed.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/pgdog/src/frontend/client/query_engine/test/schema_changed.rs b/pgdog/src/frontend/client/query_engine/test/schema_changed.rs index e77cb1254..3aaf50ceb 100644 --- a/pgdog/src/frontend/client/query_engine/test/schema_changed.rs +++ b/pgdog/src/frontend/client/query_engine/test/schema_changed.rs @@ -14,6 +14,19 @@ fn get_pool_ids(engine: &mut QueryEngine) -> Vec { /// Helper to run DDL in a transaction and verify schema_changed is set and reload occurs. async fn assert_ddl_sets_schema_changed(ddl: &str, expected_cmd: &str, cleanup: Option<&str>) { + Box::pin(assert_ddl_sets_schema_changed_inner( + ddl, + expected_cmd, + cleanup, + )) + .await +} + +async fn assert_ddl_sets_schema_changed_inner( + ddl: &str, + expected_cmd: &str, + cleanup: Option<&str>, +) { let mut test_client = TestClient::new_sharded(Parameters::default()).await; // Capture pool IDs before DDL