From 3d934b7eb9d80ab6b24462d2f0cd8f44e4fb3c87 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 00:06:17 +0000 Subject: [PATCH 1/2] test(proxy): reproduce schema intent outliving its prepared statement Schema middleware keeps its own map of prepared statement and portal intents. These tests show three ways it goes out of step with PostgreSQL: - A re-Parse of a DDL statement name that PostgreSQL rejects replaces the intent, so executing the retained DDL is not seen as DDL. - Close, DEALLOCATE and DISCARD ALL never release the intents. - DISCARD ALL drops the temporary table that shadowed an encrypted table, but the connection stays fail-closed for the next client. Co-Authored-By: Claude Opus 5.5 --- .../src/postgresql/context/mod.rs | 4 + .../src/postgresql/middleware/mod.rs | 113 ++++++++++++++++++ .../src/proxy/schema/middleware.rs | 61 ++++++++++ 3 files changed, 178 insertions(+) diff --git a/packages/cipherstash-proxy/src/postgresql/context/mod.rs b/packages/cipherstash-proxy/src/postgresql/context/mod.rs index 44077828..cf74ab0d 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/mod.rs @@ -148,6 +148,8 @@ pub(crate) struct ProtocolStateSizes { pub settings_statements: usize, pub portals: usize, pub portal_operations: usize, + pub schema_statements: usize, + pub schema_portals: usize, } #[derive(Debug)] @@ -925,6 +927,8 @@ where settings_statements: state.settings_protocol.statement_count(), portals: state.portals.len(), portal_operations: state.portal_operations.len(), + schema_statements: self.schema_middleware.statement_count(), + schema_portals: self.schema_middleware.portal_count(), } } diff --git a/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs b/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs index 4434f8e3..f7565d66 100644 --- a/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs +++ b/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(), diff --git a/packages/cipherstash-proxy/src/proxy/schema/middleware.rs b/packages/cipherstash-proxy/src/proxy/schema/middleware.rs index 3ff448f1..9a3e655d 100644 --- a/packages/cipherstash-proxy/src/proxy/schema/middleware.rs +++ b/packages/cipherstash-proxy/src/proxy/schema/middleware.rs @@ -481,6 +481,16 @@ impl SchemaMiddleware { } } + #[cfg(test)] + pub(crate) fn statement_count(&self) -> usize { + self.prepared.read().unwrap().len() + } + + #[cfg(test)] + pub(crate) fn portal_count(&self) -> usize { + self.portals.read().unwrap().len() + } + /// Records schema intent for a named prepared statement. pub fn prepare(&self, name: Name, statement: Statement) { let intent = self.intent(statement); @@ -1837,6 +1847,57 @@ mod tests { assert!(middleware.has_unmodelled_ddl()); } + #[test] + fn discard_all_ends_fail_closed_temporary_shadowing() { + let mut encrypt_config = EncryptConfig::new(); + let column = match parse("create table users (email eql_v3_text_search)") { + Statement::CreateTable(create) => create.columns.into_iter().next().unwrap(), + _ => unreachable!(), + }; + add_column_config(&mut encrypt_config, "users", &column); + let store = CommittedSchemaStore::for_testing(Schema::new("public"), encrypt_config); + let middleware = SchemaMiddleware::from_store(store); + run(&middleware, "create temporary table users (email text)"); + middleware.ready_for_query(TransactionStatus::Idle); + assert!(middleware.has_unmodelled_ddl()); + + // A pooler's reset query drops the shadowing table, so the next client is unaffected + run(&middleware, "discard all"); + middleware.ready_for_query(TransactionStatus::Idle); + + assert!(!middleware.has_unmodelled_ddl()); + } + + #[test] + fn discard_temp_keeps_temporary_shadowing_fail_closed() { + let mut encrypt_config = EncryptConfig::new(); + let column = match parse("create table users (email eql_v3_text_search)") { + Statement::CreateTable(create) => create.columns.into_iter().next().unwrap(), + _ => unreachable!(), + }; + add_column_config(&mut encrypt_config, "users", &column); + let store = CommittedSchemaStore::for_testing(Schema::new("public"), encrypt_config); + let middleware = SchemaMiddleware::from_store(store); + run(&middleware, "create temporary table users (email text)"); + middleware.ready_for_query(TransactionStatus::Idle); + + // Unlike `DISCARD ALL`, a rollback can undo `DISCARD TEMP` + run(&middleware, "discard temp"); + + assert!(middleware.has_unmodelled_ddl()); + } + + #[test] + fn discard_all_keeps_unmodelled_catalog_ddl_fail_closed() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + run(&middleware, "create view names as select 1"); + assert!(middleware.has_unmodelled_ddl()); + + run(&middleware, "discard all"); + + assert!(middleware.has_unmodelled_ddl()); + } + #[test] fn native_temporary_table_batch_does_not_fail_closed() { let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); From 1a40c0eb04adbd1c77f4d35c7864cd1c09dd0fe5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 00:07:13 +0000 Subject: [PATCH 2/2] fix(proxy): keep schema intent with its prepared statement Schema middleware kept prepared statement and portal intents in maps of its own, outside Context's rollback journal and lifecycle. A rejected re-Parse replaced the intent of the statement PostgreSQL kept, and Close, DEALLOCATE and DISCARD ALL never released it. Context now keeps the intent in the prepared statement's entry and carries it onto each portal bound from it, as it does a deallocation. The journal restores it with the rest of the statement, and closing or deallocating the statement releases it. Schema middleware classifies a statement and receives the intent when the portal executes. DISCARD ALL now also ends the fail-closed state left by a temporary table that shadowed an encrypted table. PostgreSQL refuses DISCARD ALL in a transaction block, so no rollback can bring the table back. DISCARD TEMP can be rolled back, so it keeps the connection fail-closed. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 4 + packages/cipherstash-proxy/CONTEXT.md | 13 ++- .../src/postgresql/context/lifecycle.rs | 5 +- .../src/postgresql/context/mod.rs | 51 ++++++--- .../postgresql/context/prepared_statements.rs | 12 ++ .../postgresql/context/settings_protocol.rs | 1 + .../src/proxy/schema/middleware.rs | 107 ++++++------------ .../cipherstash-proxy/src/proxy/schema/mod.rs | 2 +- 8 files changed, 97 insertions(+), 98 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 37c446dd..d4ad7c18 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/packages/cipherstash-proxy/CONTEXT.md b/packages/cipherstash-proxy/CONTEXT.md index 3d318b1d..6906e9f0 100644 --- a/packages/cipherstash-proxy/CONTEXT.md +++ b/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` diff --git a/packages/cipherstash-proxy/src/postgresql/context/lifecycle.rs b/packages/cipherstash-proxy/src/postgresql/context/lifecycle.rs index 5b399ecf..ad0b18b9 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/lifecycle.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/lifecycle.rs @@ -150,7 +150,7 @@ impl Context { 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 Context { statement: &Name, ) -> Result { 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 Context { 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) } diff --git a/packages/cipherstash-proxy/src/postgresql/context/mod.rs b/packages/cipherstash-proxy/src/postgresql/context/mod.rs index cf74ab0d..c69951c4 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/mod.rs +++ b/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, }, }; @@ -148,8 +148,6 @@ pub(crate) struct ProtocolStateSizes { pub settings_statements: usize, pub portals: usize, pub portal_operations: usize, - pub schema_statements: usize, - pub schema_portals: usize, } #[derive(Debug)] @@ -164,11 +162,13 @@ struct ConnectionProtocolState { portal_operations: HashMap>, } -/// 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 { operation: K, deallocation: Option, + schema: Option>, } impl Default for ConnectionProtocolState { @@ -209,6 +209,7 @@ impl ConnectionProtocolState { PortalOwner { operation, deallocation: None, + schema: None, }, ); self.portals.insert(name, Arc::new(portal)); @@ -231,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); @@ -276,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() } } @@ -812,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(()) } @@ -927,8 +936,6 @@ where settings_statements: state.settings_protocol.statement_count(), portals: state.portals.len(), portal_operations: state.portal_operations.len(), - schema_statements: self.schema_middleware.statement_count(), - schema_portals: self.schema_middleware.portal_count(), } } @@ -1357,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 { + 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. diff --git a/packages/cipherstash-proxy/src/postgresql/context/prepared_statements.rs b/packages/cipherstash-proxy/src/postgresql/context/prepared_statements.rs index c9fc8609..1753e9fa 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/prepared_statements.rs +++ b/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 { parse: Option, /// Set when the statement's SQL is a `DEALLOCATE` or `DISCARD ALL` deallocation: Option, + /// What executing the statement does to schema state + schema: Option>, } impl Default for PreparedStatement { @@ -21,6 +24,7 @@ impl Default for PreparedStatement { 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> { + self.statements.get(name)?.schema.clone() + } + + pub fn set_schema(&mut self, name: Name, schema: Arc) { + self.statements.entry(name).or_default().schema = Some(schema); + } + /// Remove the statement, returning its metrics scope. pub fn remove(&mut self, name: &Name) -> Option { self.statements.remove(name)?.metrics_scope diff --git a/packages/cipherstash-proxy/src/postgresql/context/settings_protocol.rs b/packages/cipherstash-proxy/src/postgresql/context/settings_protocol.rs index 0091d56f..32d71744 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/settings_protocol.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/settings_protocol.rs @@ -90,6 +90,7 @@ pub(super) struct MetadataSnapshot { pub statement_metrics: Option, pub statement_operation: Option, pub statement_deallocation: Option, + pub statement_schema: Option>, pub portal: Option>, pub portal_operation: Option, pub suspended: Option<(OperationId, super::ExecuteContext)>, diff --git a/packages/cipherstash-proxy/src/proxy/schema/middleware.rs b/packages/cipherstash-proxy/src/proxy/schema/middleware.rs index 9a3e655d..89bd5ea3 100644 --- a/packages/cipherstash-proxy/src/proxy/schema/middleware.rs +++ b/packages/cipherstash-proxy/src/proxy/schema/middleware.rs @@ -64,7 +64,6 @@ use super::eql_domains; use super::manager::CommittedSchemaStore; -use crate::postgresql::Name; use crate::proxy::encrypt_config::from_domain::column_config_from_domain; use crate::proxy::EncryptConfig; use cipherstash_client::eql::Identifier; @@ -73,16 +72,19 @@ use sqltk::parser::ast::{ AlterColumnOperation, AlterTableOperation, ColumnDef, DataType, DiscardObject, DropBehavior, Ident, ObjectName, ObjectNamePart, ObjectType, OnCommit, Statement, }; -use std::collections::{HashMap, VecDeque}; +use std::collections::VecDeque; use std::sync::{ atomic::{AtomicBool, AtomicUsize, Ordering}, Arc, RwLock, }; use tokio::sync::Notify; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] /// The schema-relevant meaning recorded for a parsed statement. -struct Intent { +/// +/// Context keeps it with the prepared statement and its portals, so it is released and +/// restored together with them. +pub struct SchemaIntent { statement: Statement, ddl: bool, modelled: bool, @@ -90,7 +92,7 @@ struct Intent { temporary: bool, } -impl Intent { +impl SchemaIntent { /// Whether schema-dependent work must wait for this execution's outcome: /// catalog DDL, or a temporary table a later statement may use. Work after /// `DISCARD` need not wait, since a temporary table it removes can only @@ -258,7 +260,7 @@ fn drop_tables(names: Vec) -> Statement { /// One execution or synchronization boundary awaiting a backend outcome. enum PendingExecution { /// An execution, optionally carrying schema intent. - Execute(Option>), + Execute(Option>), /// The boundary terminated by the next `ReadyForQuery` message. ReadyBoundary, } @@ -316,8 +318,6 @@ pub struct SchemaMiddleware { base_encrypt_config: Arc>>, encrypt_config: Arc>>, resolver: Arc>>, - prepared: Arc>>, - portals: Arc>>, executions: Arc>>, in_flight_ddl: Arc, execution_finished: Arc, @@ -349,8 +349,6 @@ impl SchemaMiddleware { base_encrypt_config: Arc::new(RwLock::new(snapshot.encrypt_config())), encrypt_config: Arc::new(RwLock::new(snapshot.encrypt_config())), resolver: Arc::new(RwLock::new(Arc::new(TableResolver::new_editable(schema)))), - prepared: Arc::new(RwLock::new(HashMap::new())), - portals: Arc::new(RwLock::new(HashMap::new())), executions: Arc::new(RwLock::new(VecDeque::new())), in_flight_ddl: Arc::new(AtomicUsize::new(0)), execution_finished: Arc::new(Notify::new()), @@ -459,9 +457,9 @@ impl SchemaMiddleware { } /// Classifies what executing `statement` would do to schema state. - fn intent(&self, statement: Statement) -> Intent { + pub fn intent(&self, statement: Statement) -> SchemaIntent { let statement = mark_temporary(statement); - Intent { + SchemaIntent { ddl: self.is_schema_ddl(&statement), modelled: self.models_ddl(&statement), temporary: self.changes_only_temporary_tables(&statement), @@ -481,46 +479,11 @@ impl SchemaMiddleware { } } - #[cfg(test)] - pub(crate) fn statement_count(&self) -> usize { - self.prepared.read().unwrap().len() - } - - #[cfg(test)] - pub(crate) fn portal_count(&self) -> usize { - self.portals.read().unwrap().len() - } - - /// Records schema intent for a named prepared statement. - pub fn prepare(&self, name: Name, statement: Statement) { - let intent = self.intent(statement); - self.prepared.write().unwrap().insert(name, intent); - } - - /// Associates a portal with the schema intent of its prepared statement. - pub fn bind(&self, portal: Name, prepared_statement: &Name) { - let intent = self - .prepared - .read() - .unwrap() - .get(prepared_statement) - .cloned(); - let mut portals = self.portals.write().unwrap(); - match intent { - Some(intent) => { - portals.insert(portal, intent); - } - None => { - portals.remove(&portal); - } - } - } - - /// Records a portal execution awaiting a backend result and returns whether - /// the protocol adapter must flush DDL immediately to unblock dependent work. - pub fn execute(&self, portal: &Name) -> bool { - let intent = self.portals.read().unwrap().get(portal).cloned(); - let executes_ddl = intent.as_ref().is_some_and(Intent::changes_schema); + /// Records a portal execution, with the schema intent of its prepared statement, awaiting + /// a backend result. Returns whether the protocol adapter must flush DDL immediately to + /// unblock dependent work. + pub fn execute(&self, intent: Option) -> bool { + let executes_ddl = intent.as_ref().is_some_and(SchemaIntent::changes_schema); self.transaction_active.store(true, Ordering::Release); if executes_ddl { self.in_flight_ddl.fetch_add(1, Ordering::AcqRel); @@ -793,7 +756,18 @@ impl SchemaMiddleware { self.temporary.write().unwrap().apply(&statement) } SchemaChange::DiscardTemporary => { - self.temporary.write().unwrap().discard() + self.temporary.write().unwrap().discard(); + // PostgreSQL refuses `DISCARD ALL` in a transaction block, so + // no rollback can bring back the temporary tables it drops. + // `DISCARD TEMP` can be rolled back, so it keeps them unmodelled. + if matches!( + intent.statement, + Statement::Discard { + object_type: DiscardObject::ALL + } + ) { + self.local_unmodelled.store(false, Ordering::Release); + } } SchemaChange::Catalog(statement) if self.models_ddl(&statement) => { self.apply_ddl(&statement) @@ -957,7 +931,7 @@ impl SchemaMiddleware { } /// Removes the next execution without crossing a readiness boundary. - fn pop_execution(&self) -> Option> { + fn pop_execution(&self) -> Option> { let mut executions = self.executions.write().unwrap(); if matches!(executions.front(), Some(PendingExecution::Execute(_))) { match executions.pop_front().unwrap() { @@ -1251,12 +1225,8 @@ mod tests { ) } - fn execute_prepared(middleware: &SchemaMiddleware, name: String, statement: Statement) { - let statement_name = Name::from(format!("statement_{name}")); - let portal_name = Name::from(format!("portal_{name}")); - middleware.prepare(statement_name.clone(), statement); - middleware.bind(portal_name.clone(), &statement_name); - middleware.execute(&portal_name); + fn execute_prepared(middleware: &SchemaMiddleware, statement: Statement) { + middleware.execute(Some(middleware.intent(statement))); } proptest! { @@ -1271,7 +1241,6 @@ mod tests { let name = format!("generated_{id}"); execute_prepared( &middleware, - format!("create_{id}"), parse(&format!( "create table {name} (id bigint, secret eql_v3_text_search)" )), @@ -1283,7 +1252,6 @@ mod tests { let name = format!("generated_{id}"); execute_prepared( &middleware, - format!("create_{id}"), parse(&format!( "create table {name} (id bigint, secret eql_v3_text_search)" )), @@ -1291,7 +1259,7 @@ mod tests { middleware.execution_failed(); } GeneratedEvent::Rollback => { - execute_prepared(&middleware, "rollback".to_owned(), parse("rollback")); + execute_prepared(&middleware, parse("rollback")); middleware.execution_succeeded(); expected_tables.clear(); } @@ -1483,11 +1451,7 @@ mod tests { fn ordinary_extended_execution_pins_the_snapshot_until_readiness() { let store = CommittedSchemaStore::for_testing(Schema::new("public"), EncryptConfig::new()); let middleware = SchemaMiddleware::from_store(store.clone()); - let statement = Name::from("statement"); - let portal = Name::from("portal"); - middleware.prepare(statement.clone(), parse("select 1")); - middleware.bind(portal.clone(), &statement); - assert!(!middleware.execute(&portal)); + assert!(!middleware.execute(Some(middleware.intent(parse("select 1"))))); let mut published = Schema::new("public"); published.add_table(Table::new(Ident::new("reports"))); @@ -1510,12 +1474,9 @@ mod tests { #[test] fn extended_ddl_execution_requests_an_immediate_flush() { let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); - let statement = Name::from("statement"); - let portal = Name::from("portal"); - middleware.prepare(statement.clone(), parse("create table reports (id bigint)")); - middleware.bind(portal.clone(), &statement); + let intent = middleware.intent(parse("create table reports (id bigint)")); - assert!(middleware.execute(&portal)); + assert!(middleware.execute(Some(intent))); } #[test] diff --git a/packages/cipherstash-proxy/src/proxy/schema/mod.rs b/packages/cipherstash-proxy/src/proxy/schema/mod.rs index 968933cd..8a9db453 100644 --- a/packages/cipherstash-proxy/src/proxy/schema/mod.rs +++ b/packages/cipherstash-proxy/src/proxy/schema/mod.rs @@ -10,6 +10,6 @@ mod middleware; pub use manager::{CommittedSchemaStore, SchemaManager}; #[cfg(test)] pub(crate) use middleware::committed_store_from_ddl; -pub use middleware::{SchemaMiddleware, TransactionStatus}; +pub use middleware::{SchemaIntent, SchemaMiddleware, TransactionStatus}; pub(crate) use middleware::postgres_identifier;