feat: track PostgreSQL session identity across pooled connections

pgdog/src/backend/pool/cleanup.rs

@@ -13,9 +13,12 @@ static PREPARED: Lazy<Vec<Query>> = Lazy::new(|| vec![Query::new("DEALLOCATE ALL
 ///
 static DIRTY: Lazy<Vec<Query>> = 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::<Vec<_>>();
+
+        assert_eq!(
+            queries,
+            [
+                "SET SESSION AUTHORIZATION DEFAULT",
+                "RESET ALL",
+                "SELECT pg_advisory_unlock_all()",
+                "DISCARD TEMP",
+            ]
+        );
+    }
+}

pgdog/src/backend/pool/connection/binding.rs

@@ -432,9 +432,7 @@ impl Binding {
                 let mut max = 0;
                 for result in results {
                     let synced = result?;
-                    if max < synced {
-                        max = synced;
-                    }
+                    max = max.max(synced);
                 }
                 Ok(max)
             }
@@ -454,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(),

pgdog/src/backend/pool/guard.rs

@@ -113,9 +113,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();
@@ -176,6 +194,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();

pgdog/src/backend/server.rs

@@ -684,6 +684,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.
                     _ => (),
@@ -729,7 +730,7 @@ impl Server {
             // Construct client parameter SET queries.
             let tracked = params.tracked_and_different(&self.client_params);
             // Construct RESET queries to reset any current params
-            // to their default values.
+            // 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
@@ -1136,6 +1137,16 @@ impl Server {
 
     pub(crate) fn reset_params(&mut self) {
         self.client_params.clear();
+        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) {
@@ -2470,6 +2481,7 @@ pub(crate) mod test {
             .link_client(FrontendPid::new(), &params, None)
             .await?;
         assert_eq!(changed, 1);
+        assert!(!server.dirty());
 
         let changed = server
             .link_client(FrontendPid::new(), &params, None)
@@ -2494,6 +2506,52 @@ pub(crate) mod test {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_link_client_reconciles_session_identity() -> Result<(), Box<dyn std::error::Error>>
+    {
+        let mut params = Parameters::default();
+        params.insert_identity("role", &"pgdog".into(), false, false);
+
+        let mut server = test_server().await;
+        assert!(!server.dirty());
+
+        let changed = server
+            .link_client(FrontendPid::new(), &params, None)
+            .await?;
+
+        assert_eq!(changed, 2);
+        assert!(!server.dirty());
+
+        let changed = server
+            .link_client(FrontendPid::new(), &params, None)
+            .await?;
+        assert_eq!(changed, 0, "identity stays on the server snapshot");
+
+        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");
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_link_client_reconciles_transaction_identity()
+    -> Result<(), Box<dyn std::error::Error>> {
+        let mut params = Parameters::default();
+        params.insert_identity("role", &"pgdog".into(), true, false);
+
+        let mut server = test_server().await;
+        let changed = server
+            .link_client(FrontendPid::new(), &params, Some("BEGIN"))
+            .await?;
+
+        assert_eq!(changed, 2);
+        assert!(!server.dirty());
+        server.rollback().await?;
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_copy_protocol() {
         let mut server = test_server().await;

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,
         })
     }
 
@@ -140,6 +142,7 @@ impl QueryEngine {
             return Ok(result);
         }
 
+        self.last_server_error = false;
         self.stats.received(client_request.total_message_len());
         self.set_state(State::Active); // Client is active.
 

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

@@ -195,6 +195,7 @@ impl QueryEngine {
         }
 
         if code == 'E' {
+            self.last_server_error = true;
             if let Some(state) = self.pending_explain.as_mut() {
                 state.annotated = true;
             }

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

@@ -34,7 +34,14 @@ impl QueryEngine {
             let is_pin = param.name == PGDOG_PIN;
 
             if let Some(value) = param.value.clone() {
-                if context.in_transaction() {
+                if context.params.insert_identity(
+                    &param.name,
+                    &value,
+                    context.in_transaction(),
+                    param.local,
+                ) {
+                    continue;
+                } else if context.in_transaction() {
                     context
                         .params
                         .insert_transaction(&param.name, value, param.local);
@@ -51,7 +58,13 @@ impl QueryEngine {
                 }
             } else {
                 fake_command = "RESET";
-                context.params.reset(&param.name);
+                if !context.params.reset_identity(
+                    &param.name,
+                    context.in_transaction(),
+                    param.local,
+                ) {
+                    context.params.reset(&param.name);
+                }
                 if is_pin {
                     self.manual_lock = false;
                 }
@@ -64,6 +77,9 @@ impl QueryEngine {
 
         if self.backend.connected() {
             self.execute(context, client_request, None).await?;
+            if !self.last_server_error {
+                self.backend.sync_client_params(context.params);
+            }
         } else {
             let fake_response = set_config
                 .then(|| params.iter().map(|p| p.value.as_ref()))
@@ -119,7 +135,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);
         }
 

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

@@ -36,6 +36,7 @@ mod rewrite_offset;
 mod rewrite_projection;
 mod rewrite_simple_prepared;
 mod schema_changed;
+mod session_identity;
 mod set;
 mod set_schema_sharding;
 mod sharded;

pgdog/src/frontend/client/query_engine/test/schema_changed.rs

@@ -14,6 +14,19 @@ fn get_pool_ids(engine: &mut QueryEngine) -> Vec<u64> {
 
 /// 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