test(proxy): check a connection's mirrored state against a PostgreSQL model

CHANGELOG.md

@@ -55,6 +55,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
   `DISCARD ALL`, the default reset query of PgBouncer, and `DISCARD TEMP` now end that refusal once their transaction commits, so the next client of a pooled connection is not refused.
   If the transaction rolls back, the table and the refusal both come back.
 
+- **Unnamed prepared statement after a simple Query**: PostgreSQL drops a client's unnamed prepared statement whenever it runs a simple Query, but Proxy kept its copy.
+  Proxy now releases it when PostgreSQL answers the Query, so the statements Proxy tracks match PostgreSQL's.
+
+- **Statement metrics after a pipelined re-Parse**: when a client prepared a name again in the same pipeline as a SQL `DEALLOCATE` of that name, Proxy never finished the metrics of the deallocated statement.
+  They are now finished when PostgreSQL confirms the `DEALLOCATE`.
+
+- **Statement restored after `DEALLOCATE ALL` in a failed pipeline**: when a pipeline ran `DEALLOCATE ALL` and then failed before a `Close` or `Parse` of a statement it removed, Proxy brought that statement back.
+  Proxy now forgets every statement `DEALLOCATE ALL` removes, including those a pipelined `Close` or `Parse` had already replaced.
+
 - **`--db-user`**: the flag now sets the database username. It previously set an unused variable, so it had no effect.
 
 - **Unknown column error message**: the type-checking error for a column a table does not have now names the column and the table the right way round, for example `Column: email not found for table: users`. It previously swapped them.

packages/cipherstash-proxy/proptest-regressions/postgresql/middleware/tests/connection_model.txt

@@ -0,0 +1,8 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 335ad1263fb8dcd2668352a2c67efab0b56612df51dd1886b260a0bc5923ae87 # shrinks to pipelined = true, steps = [Parse { statement: 2, sql: DeallocateAll }, Parse { statement: 1, sql: Select }, Bind { statement: 2, portal: 0 }, Execute { portal: 0, fails: false }, Execute { portal: 0, fails: false }, CloseStatement(1)]
+cc 3ee759241739404886acf5e1805cdd8bce8e2b2beb4410d92c8d0d1b7039dd8b # shrinks to pipelined = true, steps = [Parse { statement: 1, sql: Select }, Parse { statement: 0, sql: Select }, CloseStatement(2), Parse { statement: 0, sql: DeallocateAll }, Parse { statement: 0, sql: DeallocateAll }, CloseStatement(2), Bind { statement: 0, portal: 1 }, Parse { statement: 2, sql: Select }, Execute { portal: 1, fails: false }, Bind { statement: 0, portal: 1 }, CloseStatement(1)]

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

@@ -120,6 +120,9 @@ struct OperationContext {
     /// SQL deallocations in a simple Query or an executed portal, applied in order as
     /// PostgreSQL confirms each one
     deallocations: VecDeque<Deallocation>,
+    /// A simple Query releases the unnamed statement once PostgreSQL runs it, which it
+    /// does not while it skips messages until Sync
+    releases_unnamed_statement: bool,
 }
 
 /// A prepared-statement removal requested with SQL rather than a protocol `Close`.
@@ -274,6 +277,23 @@ impl ConnectionProtocolState {
 }
 
 impl ConnectionProtocolState {
+    /// Forgets the statements PostgreSQL removed during `operation`, other than any a later
+    /// Parse prepared again
+    fn deallocate_statements(
+        &mut self,
+        operation: OperationId,
+        names: &[Name],
+    ) -> Vec<SessionMetricsContext> {
+        let mut candidates = Vec::new();
+        for name in names {
+            // Also when a later Parse prepared the name again, so that Parse's rollback
+            // cannot restore the statement PostgreSQL removed
+            candidates.extend(self.settings_protocol.deallocated(operation, name));
+        }
+        candidates.extend(self.prepared_statements.deallocate(names, operation));
+        self.take_unreferenced_metrics(candidates)
+    }
+
     /// Everything Proxy keeps for the named statement, for a rollback to restore
     fn statement_snapshot(&self, name: &Name) -> settings_protocol::MetadataSnapshot {
         settings_protocol::MetadataSnapshot {
@@ -803,6 +823,41 @@ where
         Ok(())
     }
 
+    /// PostgreSQL runs a simple Query in the unnamed statement, so it releases the client's
+    /// once it runs the Query.
+    pub fn query_releases_unnamed_statement(&self, operation: OperationId) -> Result<(), Error> {
+        self.protocol_state
+            .write()
+            .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?
+            .operations
+            .entry(operation)
+            .or_default()
+            .releases_unnamed_statement = true;
+        Ok(())
+    }
+
+    /// PostgreSQL answered `operation`, so a simple Query has released the unnamed statement
+    pub fn operation_answered(&self, operation: OperationId) -> Result<(), Error> {
+        let finished = {
+            let mut state = self
+                .protocol_state
+                .write()
+                .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?;
+            let releases = state
+                .operations
+                .get_mut(&operation)
+                .is_some_and(|context| std::mem::take(&mut context.releases_unnamed_statement));
+            if !releases {
+                return Ok(());
+            }
+            state.deallocate_statements(operation, &[Name::new()])
+        };
+        for metrics in finished {
+            self.record_finished_metrics_scope(Some(metrics));
+        }
+        Ok(())
+    }
+
     /// Record that the SQL of statement `name` is a `DEALLOCATE` or `DISCARD ALL`, so that
     /// executing it through the extended protocol releases what it removes.
     pub fn set_statement_deallocation(
@@ -899,14 +954,19 @@ where
 
             debug!(target: CONTEXT, client_id = self.client_id, msg = "Deallocate", ?deallocation);
 
-            let names = state.prepared_statements.deallocation_names(deallocation);
-            for name in &names {
-                // Also when a later Parse prepared the name again, so that Parse's rollback
-                // cannot restore the statement PostgreSQL removed
-                state.settings_protocol.deallocated(operation, name);
+            let mut names = state
+                .prepared_statements
+                .deallocation_names(deallocation.clone());
+            if deallocation == Deallocation::AllStatements {
+                // Also those a pipelined Close or Parse took out of view, whose rollback could
+                // otherwise restore them
+                for name in state.settings_protocol.journalled_statement_names() {
+                    if !names.contains(&name) {
+                        names.push(name);
+                    }
+                }
             }
-            let candidates = state.prepared_statements.deallocate(&names, operation);
-            state.take_unreferenced_metrics(candidates)
+            state.deallocate_statements(operation, &names)
         };
 
         for metrics in finished_metrics {

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

@@ -7,7 +7,7 @@ use crate::error::Error;
 use crate::postgresql::{Name, OperationId};
 use sqltk::parser::ast::Statement;
 use std::{
-    collections::{BTreeMap, HashMap, VecDeque},
+    collections::{BTreeMap, HashMap, HashSet, VecDeque},
     sync::Arc,
 };
 
@@ -481,9 +481,9 @@ impl SettingsProtocol {
         }
         retired_metrics
     }
-    /// Forgets a statement that PostgreSQL removed with SQL `DEALLOCATE` or `DISCARD ALL`
-    /// during `operation`, unless a later Parse prepared it again.
-    pub fn deallocated(&mut self, operation: OperationId, name: &Name) {
+    /// Forgets a statement that PostgreSQL removed during `operation`, unless a later Parse
+    /// prepared it again. Returns the metrics scopes of the rollback snapshot it forgot.
+    pub fn deallocated(&mut self, operation: OperationId, name: &Name) -> Vec<SessionId> {
         if self
             .prepared
             .get(name)
@@ -506,8 +506,20 @@ impl SettingsProtocol {
             })
         {
             mutation.previous = Previous::Statement(None);
-            mutation.metadata = MetadataSnapshot::default();
+            return std::mem::take(&mut mutation.metadata).metrics().collect();
         }
+        Vec::new()
+    }
+    /// Names of the named statements whose rollback snapshots the journal holds, which
+    /// include statements Proxy already closed or replaced
+    pub fn journalled_statement_names(&self) -> HashSet<Name> {
+        self.journal
+            .values()
+            .filter(|mutation| {
+                !mutation.name.is_empty() && matches!(mutation.previous, Previous::Statement(_))
+            })
+            .map(|mutation| mutation.name.clone())
+            .collect()
     }
     #[cfg(test)]
     pub fn statement_count(&self) -> usize {

packages/cipherstash-proxy/src/postgresql/middleware/backend.rs

@@ -160,6 +160,13 @@ impl<S: EncryptionService> Backend<S> {
     ) -> Result<BackendMiddlewareOutput, Error> {
         let mut outbound_message = protocol_message.clone();
 
+        // A row is never the first answer to an operation, so skip the lock for it
+        if let Some(operation) =
+            operation.filter(|_| !matches!(protocol_message, BackendMessage::DataRow(_)))
+        {
+            self.context.operation_answered(operation)?;
+        }
+
         if let (BackendMessage::CommandComplete(tag), Some(operation)) =
             (&protocol_message, operation)
         {

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

@@ -185,6 +185,9 @@ impl<S: EncryptionService> Frontend<S> {
             self.context.discard_operation(operation)?;
             return Ok(FrontendMiddlewareOutput::Suppress(protocol_message));
         }
+        if matches!(protocol_message, FrontendMessage::Query(_)) {
+            self.context.query_releases_unnamed_statement(operation)?;
+        }
 
         if self.context.mapping_disabled() {
             match &protocol_message {