fix(proxy): keep schema intent with its prepared statement

CHANGELOG.md

@@ -47,6 +47,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Fixed
 
+- **DDL in a prepared statement after a rejected re-Parse**: when PostgreSQL rejected a second Parse that reused a prepared DDL statement's name, Proxy forgot that the retained statement was DDL. Executing it then changed the schema without Proxy updating its view of the schema or publishing the change to other connections. Proxy now keeps a statement's schema effects with the statement, so they are restored, closed and deallocated together.
+
+- **Pooled connections refused after a temporary table shadowed an encrypted table**: after a temporary table with the same name as an encrypted table, Proxy refuses statements that need encryption on that connection until the table is gone. `DISCARD ALL`, the default reset query of PgBouncer, now ends that refusal, so the next client of a pooled connection is not refused. `DISCARD TEMP` still keeps it, because a rollback can bring the table back.
+
 - **`--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/CONTEXT.md

@@ -33,7 +33,9 @@ are different things; do not call either of them a Statement here.
 **Prepared statement**:
 A name the client prepared on its connection with Parse, and everything Proxy keeps for
 it: the Parse that prepared it, its Statement when it touches encrypted columns, its
-statement metrics scope, and whether its SQL is a deallocation. It lasts until a Close, a
+statement metrics scope, whether its SQL is a deallocation, and its schema intent (what
+executing it does to schema state, which Schema middleware applies once PostgreSQL confirms
+it). It lasts until a Close, a
 Parse that reuses the name, or a `DEALLOCATE` or `DISCARD ALL` that PostgreSQL confirms,
 and all of it is released together.
 
@@ -176,10 +178,11 @@ closed without forwarding readiness, and the dirty publication remains eligible
 
 **Schema middleware**:
 The owner of transactional schema state. Frontend and Backend report protocol lifecycle events;
-they do not directly change overlays or dirty flags. The middleware owns DDL detection, prepared
-DDL effects, successful-execution activation, savepoint and transaction transitions,
-effective-schema resolution, and the decision that publication is required. `Context` performs the
-authoritative reload round trip, while `SchemaManager` coalesces reloads and orders their
+they do not directly change overlays or dirty flags. The middleware owns DDL detection, the
+classification of each statement's schema intent, successful-execution activation, savepoint and
+transaction transitions, effective-schema resolution, and the decision that publication is
+required. `Context` keeps each schema intent with its prepared statement and portals, and performs
+the authoritative reload round trip, while `SchemaManager` coalesces reloads and orders their
 generations. See `docs/adr/0001-transaction-aware-schema-middleware.md`.
 
 ## Note on `session`

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

@@ -150,7 +150,7 @@ impl<S: EncryptionService> Context<S> {
             self.wait_for_schema_execution().await;
             self.ensure_schema_modelled()?;
         }
-        self.prepare_schema_statement(name.clone(), statement.clone());
+        self.prepare_schema_statement(name.clone(), statement.clone())?;
         self.update_statement_metadata(metrics, |m| {
             m.statement_type = Some(StatementType::from_statement(statement));
             m.set_query_fingerprint(sql);
@@ -235,7 +235,6 @@ impl<S: EncryptionService> Context<S> {
         statement: &Name,
     ) -> Result<BindKnowledge, Error> {
         self.bind_settings_statement(operation, portal.clone(), statement)?;
-        self.bind_schema_statement(portal.clone(), statement);
         let mapping_disabled = self.unsafe_disable_mapping()?;
         let metrics = if mapping_disabled {
             None
@@ -309,7 +308,7 @@ impl<S: EncryptionService> Context<S> {
             return Ok(false);
         }
         let settings = self.execute_settings_portal(operation, portal)?;
-        let schema = self.execute_schema_portal(portal);
+        let schema = self.execute_schema_portal(portal)?;
         self.set_execute_for_portal(operation, portal.clone())?;
         Ok(settings || schema)
     }

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

@@ -21,7 +21,7 @@ use crate::{
         STATEMENTS_SESSION_DURATION_SECONDS,
     },
     proxy::{
-        schema::{CommittedSchemaStore, SchemaMiddleware},
+        schema::{CommittedSchemaStore, SchemaIntent, SchemaMiddleware},
         EncryptConfig, EncryptionService, ReloadCommand, ReloadSender,
     },
 };
@@ -162,11 +162,13 @@ struct ConnectionProtocolState<K = OperationId> {
     portal_operations: HashMap<Name, PortalOwner<K>>,
 }
 
-/// The Bind operation that created a portal, and the SQL deallocation it runs when executed
+/// The Bind operation that created a portal, and the SQL deallocation and schema intent of
+/// the statement it executes
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub(crate) struct PortalOwner<K = OperationId> {
     operation: K,
     deallocation: Option<Deallocation>,
+    schema: Option<Arc<SchemaIntent>>,
 }
 
 impl<K> Default for ConnectionProtocolState<K> {
@@ -207,6 +209,7 @@ impl ConnectionProtocolState {
             PortalOwner {
                 operation,
                 deallocation: None,
+                schema: None,
             },
         );
         self.portals.insert(name, Arc::new(portal));
@@ -229,6 +232,10 @@ impl ConnectionProtocolState {
                     self.prepared_statements
                         .set_deallocation(restore.name.clone(), deallocation);
                 }
+                if let Some(schema) = restore.metadata.statement_schema {
+                    self.prepared_statements
+                        .set_schema(restore.name.clone(), schema);
+                }
                 if let Some(statement) = restore.metadata.statement {
                     self.prepared_statements
                         .set_statement(restore.name.clone(), statement);
@@ -274,6 +281,7 @@ impl ConnectionProtocolState {
             statement_metrics: self.prepared_statements.metrics_scope(name),
             statement_operation: self.prepared_statements.parse(name),
             statement_deallocation: self.prepared_statements.deallocation(name),
+            statement_schema: self.prepared_statements.schema(name),
             ..Default::default()
         }
     }
@@ -810,15 +818,18 @@ where
         Ok(())
     }
 
-    /// Carry the deallocation of `statement`, if any, onto the portal just bound from it.
+    /// Carry the deallocation and schema intent of `statement`, if any, onto the portal just
+    /// bound from it.
     pub fn bind_statement_effects(&self, portal: &Name, statement: &Name) -> Result<(), Error> {
         let mut state = self
             .protocol_state
             .write()
             .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?;
         let deallocation = state.prepared_statements.deallocation(statement);
+        let schema = state.prepared_statements.schema(statement);
         if let Some(owner) = state.portal_operations.get_mut(portal) {
             owner.deallocation = deallocation;
+            owner.schema = schema;
         }
         Ok(())
     }
@@ -1353,18 +1364,30 @@ where
     }
 
     /// Records schema intent for a parsed prepared statement.
-    fn prepare_schema_statement(&self, name: Name, statement: sqltk::parser::ast::Statement) {
-        self.schema_middleware.prepare(name, statement);
-    }
-
-    /// Associates a bound portal with its prepared statement's schema intent.
-    fn bind_schema_statement(&self, portal: Name, prepared_statement: &Name) {
-        self.schema_middleware.bind(portal, prepared_statement);
+    fn prepare_schema_statement(
+        &self,
+        name: Name,
+        statement: sqltk::parser::ast::Statement,
+    ) -> Result<(), Error> {
+        let intent = Arc::new(self.schema_middleware.intent(statement));
+        self.protocol_state
+            .write()
+            .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?
+            .prepared_statements
+            .set_schema(name, intent);
+        Ok(())
     }
 
     /// Records a portal execution and returns whether its DDL needs an injected flush.
-    fn execute_schema_portal(&self, portal: &Name) -> bool {
-        self.schema_middleware.execute(portal)
+    fn execute_schema_portal(&self, portal: &Name) -> Result<bool, Error> {
+        let intent = self
+            .protocol_state
+            .read()
+            .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?
+            .portal_operations
+            .get(portal)
+            .and_then(|owner| owner.schema.as_deref().cloned());
+        Ok(self.schema_middleware.execute(intent))
     }
 
     /// Records statements in one simple-query protocol message.

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

@@ -1,5 +1,6 @@
 use super::{Deallocation, SessionId, Statement};
 use crate::postgresql::rewrite::Name;
+use crate::proxy::schema::SchemaIntent;
 use std::{collections::HashMap, sync::Arc};
 
 /// What Proxy holds for one statement prepared on the connection.
@@ -12,6 +13,8 @@ struct PreparedStatement<K> {
     parse: Option<K>,
     /// Set when the statement's SQL is a `DEALLOCATE` or `DISCARD ALL`
     deallocation: Option<Deallocation>,
+    /// What executing the statement does to schema state
+    schema: Option<Arc<SchemaIntent>>,
 }
 
 impl<K> Default for PreparedStatement<K> {
@@ -21,6 +24,7 @@ impl<K> Default for PreparedStatement<K> {
             metrics_scope: None,
             parse: None,
             deallocation: None,
+            schema: None,
         }
     }
 }
@@ -82,6 +86,14 @@ where
         self.statements.entry(name).or_default().deallocation = Some(deallocation);
     }
 
+    pub fn schema(&self, name: &Name) -> Option<Arc<SchemaIntent>> {
+        self.statements.get(name)?.schema.clone()
+    }
+
+    pub fn set_schema(&mut self, name: Name, schema: Arc<SchemaIntent>) {
+        self.statements.entry(name).or_default().schema = Some(schema);
+    }
+
     /// Remove the statement, returning its metrics scope.
     pub fn remove(&mut self, name: &Name) -> Option<SessionId> {
         self.statements.remove(name)?.metrics_scope

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

@@ -90,6 +90,7 @@ pub(super) struct MetadataSnapshot {
     pub statement_metrics: Option<SessionId>,
     pub statement_operation: Option<OperationId>,
     pub statement_deallocation: Option<super::Deallocation>,
+    pub statement_schema: Option<Arc<crate::proxy::schema::SchemaIntent>>,
     pub portal: Option<Arc<super::Portal>>,
     pub portal_operation: Option<super::PortalOwner>,
     pub suspended: Option<(OperationId, super::ExecuteContext)>,

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

@@ -563,6 +563,45 @@ mod tests {
         }
     }
 
+    both_configs! {
+        async fn closed_statements_and_portals_retain_nothing(connect) {
+            let mut connection = connect();
+            for i in 0..ITERATIONS {
+                let name = Bytes::from(format!("s{i}"));
+                let portal = Bytes::from(format!("p{i}"));
+                let parse = connection.parse_named(format!("s{i}")).await;
+                let bind = connection.bind(name.clone(), portal.clone()).await;
+                let close_portal = connection
+                    .client(FrontendMessage::Close(Close {
+                        target: DescribeTarget::Portal,
+                        name: portal,
+                    }))
+                    .await;
+                let close_statement = connection
+                    .client(FrontendMessage::Close(Close {
+                        target: DescribeTarget::Statement,
+                        name,
+                    }))
+                    .await;
+                let sync = connection.client(FrontendMessage::Sync).await;
+                connection
+                    .server(parse, BackendMessage::ParseComplete)
+                    .await;
+                connection.server(bind, BackendMessage::BindComplete).await;
+                connection
+                    .server(close_portal, BackendMessage::CloseComplete)
+                    .await;
+                connection
+                    .server(close_statement, BackendMessage::CloseComplete)
+                    .await;
+                connection
+                    .server(sync, ready(TransactionStatus::Idle))
+                    .await;
+            }
+            assert_eq!(connection.sizes(), ProtocolStateSizes::default());
+        }
+    }
+
     /// Parse and Sync each named statement, with PostgreSQL's responses
     async fn prepare(connection: &mut Connection, names: &[&str]) {
         for name in names {
@@ -974,6 +1013,80 @@ mod tests {
         }
     }
 
+    /// Prepares `d` as DDL, has PostgreSQL acknowledge it, then re-prepares `d` as `sql`,
+    /// which never reaches PostgreSQL or is rejected by it
+    async fn reprepare_ddl(connection: &mut Connection, sql: &str, rejected_by_postgresql: bool) {
+        let parse = connection
+            .parse_sql(b"d", "CREATE TABLE reparsed (id int)".to_owned())
+            .await;
+        let sync = connection.client(FrontendMessage::Sync).await;
+        connection
+            .server(parse, BackendMessage::ParseComplete)
+            .await;
+        connection
+            .server(sync, ready(TransactionStatus::Idle))
+            .await;
+
+        let parse = connection.parse_sql(b"d", sql.to_owned()).await;
+        let sync = connection.client(FrontendMessage::Sync).await;
+        if rejected_by_postgresql {
+            connection.server(parse, error()).await;
+        }
+        connection
+            .server(sync, ready(TransactionStatus::Idle))
+            .await;
+    }
+
+    /// Executing `d` must still be known to run its DDL
+    async fn assert_executes_ddl(connection: &mut Connection) {
+        connection
+            .bind(Bytes::from_static(b"d"), Bytes::new())
+            .await;
+        connection.execute(Bytes::new()).await;
+        assert!(connection.context.schema_ddl_in_flight());
+    }
+
+    both_configs! {
+        async fn executing_a_prepared_ddl_statement_waits_for_its_outcome(connect) {
+            let mut connection = connect();
+            let parse = connection
+                .parse_sql(b"d", "CREATE TABLE reparsed (id int)".to_owned())
+                .await;
+            let sync = connection.client(FrontendMessage::Sync).await;
+            connection
+                .server(parse, BackendMessage::ParseComplete)
+                .await;
+            connection
+                .server(sync, ready(TransactionStatus::Idle))
+                .await;
+            assert_executes_ddl(&mut connection).await;
+        }
+    }
+
+    both_configs! {
+        async fn reparse_rejected_by_postgresql_keeps_the_schema_intent_of_the_retained_statement(connect) {
+            let mut connection = connect();
+            // PostgreSQL rejects a second Parse of `d` and keeps the DDL
+            reprepare_ddl(&mut connection, "SELECT 1", true).await;
+            assert_executes_ddl(&mut connection).await;
+        }
+    }
+
+    both_configs! {
+        async fn reparse_refused_by_proxy_refuses_the_retained_statement(connect) {
+            let mut connection = connect();
+            // Proxy refuses the second Parse of `d` before forwarding it, and forgets `d`
+            reprepare_ddl(&mut connection, "SET CIPHERSTASH.KEYSET_ID = 'not-a-uuid'", false).await;
+
+            // So it cannot run the DDL PostgreSQL still holds as `d` without knowing
+            let (_, output) = connection.client_output(bind_message(b"d")).await;
+            let FrontendMiddlewareOutput::Forward(FrontendMessage::Bind(forwarded)) = output else {
+                panic!("expected a Bind, got {output:?}");
+            };
+            assert_ne!(forwarded.statement, Bytes::from_static(b"d"));
+        }
+    }
+
     fn bind_message(statement: &'static [u8]) -> FrontendMessage {
         FrontendMessage::Bind(Bind {
             portal: Bytes::new(),