refactor(context): push frontend coordination logic into Context

Add helper methods to Context that encapsulate multi-step operations
previously done inline in frontend.rs:

- with_session: update session metrics if session ID present, no-op otherwise
- set_execute_for_portal: look up session ID and set execute state
- close_statement_and_portal: close both statement and portal together
- get_statement_session_or_latest: session lookup with fallback logic
- reload_schema_if_changed: conditional schema reload

Simplifies frontend.rs by centralizing related operations in Context.

Author: tobyhede

packages/cipherstash-proxy/src/postgresql/context/mod.rs

@@ -302,6 +302,12 @@ where
         let _ = self.execute.write().map(|mut queue| queue.add(ctx));
     }
 
+    /// Set execute state for portal, looking up session ID internally.
+    pub fn set_execute_for_portal(&mut self, name: Name) {
+        let session_id = self.get_portal_session_id(&name);
+        self.set_execute(name, session_id);
+    }
+
     /// Marks the current Execution as Complete.
     ///
     /// Transfers accumulated timing data from ExecuteContext to SessionMetricsContext.phase_timing:
@@ -393,6 +399,12 @@ where
             .map(|mut guarded| guarded.remove(name));
     }
 
+    /// Close both statement and its associated portal.
+    pub fn close_statement_and_portal(&mut self, name: &Name) {
+        self.close_portal(name);
+        self.close_statement(name);
+    }
+
     pub fn add_portal(&mut self, name: Name, portal: Portal) {
         debug!(target: CONTEXT, client_id = self.client_id, name = ?name, portal = ?portal);
         let _ = self.portals.write().map(|mut portals| {
@@ -421,6 +433,24 @@ where
         sessions.get(name).copied()
     }
 
+    /// Get session for statement, falling back to latest session with warning log.
+    pub fn get_statement_session_or_latest(&self, name: &Name) -> Option<SessionId> {
+        if let Some(id) = self.get_statement_session(name) {
+            return Some(id);
+        }
+
+        let fallback = self.latest_session_id();
+        if fallback.is_some() {
+            warn!(
+                target: CONTEXT,
+                client_id = self.client_id,
+                prepared_statement = ?name,
+                msg = "Session lookup failed for prepared statement, using latest session"
+            );
+        }
+        fallback
+    }
+
     ///
     /// Close the portal identified by `name`
     /// Portal is removed from  queue
@@ -740,6 +770,13 @@ where
         debug!(target: CONTEXT, msg = "Database schema reloaded", ?response);
     }
 
+    /// Reload schema if it has changed since last check.
+    pub async fn reload_schema_if_changed(&self) {
+        if self.schema_changed() {
+            self.reload_schema().await;
+        }
+    }
+
     pub fn is_passthrough(&self) -> bool {
         self.encrypt_config.is_empty() || self.config.mapping_disabled()
     }
@@ -927,6 +964,16 @@ where
         });
     }
 
+    /// Update statement metadata if session ID is present, no-op otherwise.
+    pub fn with_session<F>(&mut self, session_id: Option<SessionId>, f: F)
+    where
+        F: FnOnce(&mut SessionMetricsContext),
+    {
+        if let Some(sid) = session_id {
+            self.with_session_metrics_mut(sid, f);
+        }
+    }
+
     /// Record server wait for first response; otherwise accumulate response time for the current execute
     pub fn record_execute_server_timing(&mut self, duration: Duration) {
         if let Ok(mut queue) = self.execute.write() {

packages/cipherstash-proxy/src/postgresql/frontend.rs

@@ -11,7 +11,7 @@ use super::parser::SqlParser;
 use super::protocol::{self};
 use crate::connect::Sender;
 use crate::error::{EncryptError, Error, MappingError};
-use crate::log::{CONTEXT, MAPPER, PROTOCOL};
+use crate::log::{MAPPER, PROTOCOL};
 use crate::postgresql::context::column::Column;
 use crate::postgresql::context::statement_metadata::{ProtocolType, StatementType};
 use crate::postgresql::context::Portal;
@@ -268,9 +268,7 @@ where
                     ?code,
                 );
 
-                if self.context.schema_changed() {
-                    self.context.reload_schema().await;
-                }
+                self.context.reload_schema_if_changed().await;
 
                 if self.error_state.is_some() {
                     debug!(target: PROTOCOL,
@@ -331,20 +329,16 @@ where
         debug!(target: PROTOCOL, client_id = self.context.client_id, ?close);
         match close.target {
             Target::Portal => self.context.close_portal(&close.name),
-            Target::Statement => {
-                self.context.close_portal(&close.name);
-                self.context.close_statement(&close.name);
-            }
+            Target::Statement => self.context.close_statement_and_portal(&close.name),
         }
         Ok(())
     }
 
     async fn execute_handler(&mut self, bytes: &BytesMut) -> Result<(), Error> {
         let execute = Execute::try_from(bytes)?;
         debug!(target: PROTOCOL, client_id = self.context.client_id, ?execute);
-        let session_id = self.context.get_portal_session_id(&execute.portal);
         self.context
-            .set_execute(execute.portal.to_owned(), session_id);
+            .set_execute_for_portal(execute.portal.to_owned());
         Ok(())
     }
 
@@ -968,28 +962,14 @@ where
 
         let mut bind = Bind::try_from(bytes)?;
 
-        let session_id = match self.context.get_statement_session(&bind.prepared_statement) {
-            Some(id) => Some(id),
-            None => {
-                let fallback = self.context.latest_session_id();
-                if fallback.is_some() {
-                    warn!(
-                        target: CONTEXT,
-                        client_id = self.context.client_id,
-                        prepared_statement = ?bind.prepared_statement,
-                        msg = "Session lookup failed for prepared statement, using latest session"
-                    );
-                }
-                fallback
-            }
-        };
+        let session_id = self
+            .context
+            .get_statement_session_or_latest(&bind.prepared_statement);
 
         // Track param bytes for diagnostics
         let param_bytes: usize = bind.param_values.iter().map(|p| p.bytes.len()).sum();
-        if let Some(session_id) = session_id {
-            self.context
-                .update_statement_metadata(session_id, |m| m.set_param_bytes(param_bytes));
-        }
+        self.context
+            .with_session(session_id, |m| m.metadata.set_param_bytes(param_bytes));
 
         debug!(target: PROTOCOL, client_id = self.context.client_id, bind = ?bind);
 
@@ -1008,10 +988,8 @@ where
                     bind.result_columns_format_codes.to_owned(),
                     session_id,
                 );
-                if let Some(session_id) = session_id {
-                    self.context
-                        .update_statement_metadata(session_id, |m| m.encrypted = true);
-                }
+                self.context
+                    .with_session(session_id, |m| m.metadata.encrypted = true);
             }
         };
 
@@ -1065,16 +1043,13 @@ where
 
         // Record timing and metadata for this encryption operation
         let encrypted_count = encrypted.iter().filter(|e| e.is_some()).count();
-        if let Some(sid) = session_id {
+        self.context.with_session(session_id, |m| {
             // Add to phase timing diagnostics (accumulate)
-            self.context.add_encrypt_duration(sid, duration);
-
+            m.phase_timing.add_encrypt(duration);
             // Always update metadata for slow-statement logging
-            self.context.update_statement_metadata(sid, |m| {
-                m.encrypted = true;
-                m.set_encrypted_values_count(encrypted_count);
-            });
-        }
+            m.metadata.encrypted = true;
+            m.metadata.set_encrypted_values_count(encrypted_count);
+        });
 
         // Prometheus metrics remain gated
         if self.context.prometheus_enabled() {