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

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

@@ -56,6 +56,7 @@ impl<S: EncryptionService> Context<S> {
     }
 
     pub fn unmapped_query(&mut self, operation: OperationId) -> Result<(), Error> {
+        self.release_unnamed_statement()?;
         let metrics = self.start_metrics_scope()?;
         self.add_portal(operation, Name::new(), Portal::passthrough(Some(metrics)))?;
         self.set_simple_query_execute_until_ready(operation, Name::new(), Some(metrics))?;
@@ -373,6 +374,7 @@ impl<S: EncryptionService> Context<S> {
         portal: Portal,
         statements: &[ast::Statement],
     ) -> Result<(), Error> {
+        self.release_unnamed_statement()?;
         self.add_portal(operation, Name::new(), portal)?;
         self.set_simple_query_execute_until_ready(operation, Name::new(), Some(metrics))?;
         self.execute_simple_schema_statements(statements);
@@ -386,6 +388,7 @@ impl<S: EncryptionService> Context<S> {
         metrics: SessionId,
         response: DiagnosticResponse,
     ) -> Result<(), Error> {
+        self.release_unnamed_statement()?;
         self.finish_metrics_scope(Some(metrics))?;
         self.set_operation_error(operation, response)?;
         self.set_execute(operation, Name::new(), None)?;

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

@@ -803,6 +803,24 @@ where
         Ok(())
     }
 
+    /// PostgreSQL runs a simple Query in the unnamed statement, releasing the client's.
+    fn release_unnamed_statement(&self) -> Result<(), Error> {
+        let finished = {
+            let mut state = self
+                .protocol_state
+                .write()
+                .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?;
+            let name = Name::new();
+            let mut candidates = state.settings_protocol.released(&name);
+            candidates.extend(state.prepared_statements.remove(&name));
+            state.take_unreferenced_metrics(candidates)
+        };
+        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(

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

@@ -509,6 +509,22 @@ impl SettingsProtocol {
             mutation.metadata = MetadataSnapshot::default();
         }
     }
+    /// Forgets a statement that PostgreSQL releases unconditionally, as a simple Query
+    /// releases the unnamed statement. A rollback of an earlier operation must not restore
+    /// it, so returns the metrics scopes its rollback snapshots referenced.
+    pub fn released(&mut self, name: &Name) -> Vec<SessionId> {
+        self.prepared.remove(name);
+        self.journal
+            .values_mut()
+            .filter(|mutation| {
+                mutation.name == *name && matches!(mutation.previous, Previous::Statement(_))
+            })
+            .flat_map(|mutation| {
+                mutation.previous = Previous::Statement(None);
+                std::mem::take(&mut mutation.metadata).metrics()
+            })
+            .collect()
+    }
     #[cfg(test)]
     pub fn statement_count(&self) -> usize {
         self.prepared.len()

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

@@ -103,6 +103,8 @@ mod tests {
 
     const ITERATIONS: usize = 200;
 
+    mod connection_model;
+
     #[derive(Clone)]
     struct TestService;
 
@@ -643,6 +645,53 @@ mod tests {
         }
     }
 
+    both_configs! {
+        async fn simple_query_releases_the_unnamed_statement(connect) {
+            let mut connection = connect();
+            prepare(&mut connection, &["", "a1"]).await;
+            let select = connection.select;
+
+            // PostgreSQL runs a simple Query in the unnamed statement, dropping the client's
+            let query = connection.query(select).await;
+            connection
+                .server(query, command_complete(b"SELECT 1"))
+                .await;
+            connection
+                .server(query, ready(TransactionStatus::Idle))
+                .await;
+
+            assert!(!connection.holds_statement(b""));
+            assert_eq!(statement_names(&connection), vec!["a1"]);
+        }
+    }
+
+    both_configs! {
+        async fn earlier_pipelined_failure_does_not_restore_the_unnamed_statement_a_query_released(connect) {
+            let mut connection = connect();
+            prepare(&mut connection, &[""]).await;
+            let select = connection.select;
+
+            let failed = connection.parse_named("a1".to_owned()).await;
+            // PostgreSQL skips this Parse after the failure, so it never answers it
+            connection.parse(b"").await;
+            let sync = connection.client(FrontendMessage::Sync).await;
+            let query = connection.query(select).await;
+            connection.server(failed, error()).await;
+            connection
+                .server(sync, ready(TransactionStatus::Idle))
+                .await;
+            connection
+                .server(query, command_complete(b"SELECT 1"))
+                .await;
+            connection
+                .server(query, ready(TransactionStatus::Idle))
+                .await;
+
+            assert!(!connection.holds_statement(b""));
+            assert!(statement_names(&connection).is_empty());
+        }
+    }
+
     both_configs! {
         async fn query_failing_part_way_keeps_the_statements_it_did_not_deallocate(connect) {
             let mut connection = connect();
@@ -835,16 +884,30 @@ mod tests {
                 .server(sync, ready(TransactionStatus::Idle))
                 .await;
 
-            let query = connection.query("DEALLOCATE ALL").await;
+            // Through the extended protocol, because a simple Query releases the unnamed
+            // statement itself
+            let parse = connection
+                .parse_sql(b"deallocate", "DEALLOCATE ALL".to_owned())
+                .await;
+            let bind = connection
+                .bind(Bytes::from_static(b"deallocate"), Bytes::new())
+                .await;
+            let execute = connection.execute(Bytes::new()).await;
+            let sync = connection.client(FrontendMessage::Sync).await;
             connection
-                .server(query, command_complete(b"DEALLOCATE ALL"))
+                .server(parse, BackendMessage::ParseComplete)
                 .await;
+            connection.server(bind, BackendMessage::BindComplete).await;
             connection
-                .server(query, ready(TransactionStatus::Idle))
+                .server(execute, command_complete(b"DEALLOCATE ALL"))
+                .await;
+            connection
+                .server(sync, ready(TransactionStatus::Idle))
                 .await;
 
             assert!(connection.holds_statement(b""));
             assert!(!connection.holds_statement(b"a1"));
+            assert!(!connection.holds_statement(b"deallocate"));
         }
     }