feat(pgwire): implement transactional DDL with AST-based dispatch

farhan-syah · farhan-syah · commit a451bc5587bd · 2026-04-16T13:47:50.000+08:00
Wire transactional DDL semantics through the pgwire session layer:

- `ddl_buffer`: thread-local buffer that intercepts `propose_catalog_entry`
  calls while a BEGIN block is active, accumulating encoded payloads
  instead of proposing them immediately.
- `transaction_cmds`: COMMIT flushes the buffer as a single
  `MetadataEntry::Batch` proposal; ROLLBACK discards without proposing.
- `metadata_proposer`: checks the thread-local buffer before proposing,
  returning early when a transaction is active.
- `ast` + router fast path: parse DDL once into `NodedbStatement`,
  handle `IF [NOT] EXISTS` at dispatch before falling through to
  legacy string-prefix handlers.
diff --git a/nodedb/src/control/metadata_proposer.rs b/nodedb/src/control/metadata_proposer.rs
@@ -176,6 +176,15 @@ pub fn propose_catalog_entry_with_timeout(
     }
 
     let payload = catalog_entry::encode(entry)?;
+
+    // DDL transaction buffer: if a transactional DDL session is
+    // active on this thread (BEGIN ... COMMIT), buffer the payload
+    // instead of proposing immediately. The buffered entries will
+    // be proposed as a single MetadataEntry::Batch at COMMIT time.
+    if crate::control::server::pgwire::session::ddl_buffer::try_buffer(payload.clone()) {
+        return Ok(0);
+    }
+
     let metadata_entry = MetadataEntry::CatalogDdl { payload };
     let raw = encode_entry(&metadata_entry).map_err(|e| Error::Config {
         detail: format!("metadata entry encode: {e}"),
diff --git a/nodedb/src/control/server/pgwire/ddl/router/ast.rs b/nodedb/src/control/server/pgwire/ddl/router/ast.rs
@@ -0,0 +1,231 @@
+//! AST-based DDL dispatch — typed fast path.
+//!
+//! Runs before the legacy string-prefix routers. Handles
+//! `IF [NOT] EXISTS` at the dispatch level so individual handlers
+//! don't need to check. Falls through to legacy dispatch for
+//! `Other` variants and for statements where the typed path
+//! delegates to the existing handler (via `raw_sql`).
+
+use pgwire::api::results::{Response, Tag};
+use pgwire::error::PgWireResult;
+
+use nodedb_sql::ddl_ast::NodedbStatement;
+
+use crate::control::security::identity::AuthenticatedIdentity;
+use crate::control::state::SharedState;
+
+/// Try to dispatch a parsed `NodedbStatement`. Returns `Some` if
+/// fully handled, `None` if the statement should fall through to
+/// the legacy dispatch.
+pub(super) fn try_dispatch(
+    state: &SharedState,
+    identity: &AuthenticatedIdentity,
+    stmt: &NodedbStatement,
+) -> Option<PgWireResult<Vec<Response>>> {
+    match stmt {
+        // ── IF NOT EXISTS: swallow duplicate-creation errors ──────
+        NodedbStatement::CreateCollection {
+            name,
+            if_not_exists: true,
+            ..
+        } => {
+            if collection_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new("CREATE COLLECTION"))]));
+            }
+            None // fall through to legacy CREATE handler
+        }
+
+        NodedbStatement::CreateSequence {
+            name,
+            if_not_exists: true,
+            ..
+        } => {
+            if sequence_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new("CREATE SEQUENCE"))]));
+            }
+            None
+        }
+
+        // ── IF EXISTS: swallow not-found errors on DROP ──────────
+        NodedbStatement::DropCollection {
+            name,
+            if_exists: true,
+        } => {
+            if !collection_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new("DROP COLLECTION"))]));
+            }
+            None
+        }
+
+        NodedbStatement::DropIndex {
+            if_exists: true, ..
+        } => None, // legacy handler has its own check
+
+        NodedbStatement::DropTrigger {
+            name,
+            if_exists: true,
+            ..
+        } => {
+            if !trigger_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new("DROP TRIGGER"))]));
+            }
+            None
+        }
+
+        NodedbStatement::DropSchedule {
+            name,
+            if_exists: true,
+        } => {
+            if !schedule_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new("DROP SCHEDULE"))]));
+            }
+            None
+        }
+
+        NodedbStatement::DropSequence {
+            name,
+            if_exists: true,
+        } => {
+            if !sequence_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new("DROP SEQUENCE"))]));
+            }
+            None
+        }
+
+        NodedbStatement::DropAlert {
+            name,
+            if_exists: true,
+        } => {
+            if !alert_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new("DROP ALERT"))]));
+            }
+            None
+        }
+
+        NodedbStatement::DropRetentionPolicy {
+            name,
+            if_exists: true,
+        } => {
+            if !retention_policy_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new(
+                    "DROP RETENTION POLICY",
+                ))]));
+            }
+            None
+        }
+
+        NodedbStatement::DropChangeStream {
+            name,
+            if_exists: true,
+        } => {
+            if !change_stream_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new(
+                    "DROP CHANGE STREAM",
+                ))]));
+            }
+            None
+        }
+
+        NodedbStatement::DropMaterializedView {
+            name,
+            if_exists: true,
+        } => {
+            if !materialized_view_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new(
+                    "DROP MATERIALIZED VIEW",
+                ))]));
+            }
+            None
+        }
+
+        NodedbStatement::DropContinuousAggregate {
+            name,
+            if_exists: true,
+        } => {
+            if !continuous_aggregate_exists(state, identity, name) {
+                return Some(Ok(vec![Response::Execution(Tag::new(
+                    "DROP CONTINUOUS AGGREGATE",
+                ))]));
+            }
+            None
+        }
+
+        NodedbStatement::DropRlsPolicy {
+            if_exists: true, ..
+        } => {
+            // RLS policy existence check would need collection context;
+            // fall through to legacy handler which already handles this.
+            None
+        }
+
+        NodedbStatement::DropConsumerGroup {
+            if_exists: true, ..
+        } => None, // legacy handler
+
+        // All other variants fall through to legacy dispatch.
+        _ => None,
+    }
+}
+
+fn collection_exists(state: &SharedState, identity: &AuthenticatedIdentity, name: &str) -> bool {
+    let Some(catalog) = state.credentials.catalog() else {
+        return false;
+    };
+    let tid = identity.tenant_id.as_u32();
+    matches!(catalog.get_collection(tid, name), Ok(Some(_)))
+}
+
+fn trigger_exists(state: &SharedState, identity: &AuthenticatedIdentity, name: &str) -> bool {
+    let Some(catalog) = state.credentials.catalog() else {
+        return false;
+    };
+    let tid = identity.tenant_id.as_u32();
+    matches!(catalog.get_trigger(tid, name), Ok(Some(_)))
+}
+
+fn schedule_exists(state: &SharedState, identity: &AuthenticatedIdentity, name: &str) -> bool {
+    let tid = identity.tenant_id.as_u32();
+    state.schedule_registry.get(tid, name).is_some()
+}
+
+fn sequence_exists(state: &SharedState, identity: &AuthenticatedIdentity, name: &str) -> bool {
+    let tid = identity.tenant_id.as_u32();
+    state.sequence_registry.exists(tid, name)
+}
+
+fn alert_exists(state: &SharedState, identity: &AuthenticatedIdentity, name: &str) -> bool {
+    let tid = identity.tenant_id.as_u32();
+    state.alert_registry.get(tid, name).is_some()
+}
+
+fn retention_policy_exists(
+    state: &SharedState,
+    identity: &AuthenticatedIdentity,
+    name: &str,
+) -> bool {
+    let tid = identity.tenant_id.as_u32();
+    state.retention_policy_registry.get(tid, name).is_some()
+}
+
+fn change_stream_exists(state: &SharedState, identity: &AuthenticatedIdentity, name: &str) -> bool {
+    let tid = identity.tenant_id.as_u32();
+    state.stream_registry.get(tid, name).is_some()
+}
+
+fn materialized_view_exists(
+    state: &SharedState,
+    identity: &AuthenticatedIdentity,
+    name: &str,
+) -> bool {
+    let tid = identity.tenant_id.as_u32();
+    state.mv_registry.get_def(tid, name).is_some()
+}
+
+fn continuous_aggregate_exists(
+    state: &SharedState,
+    identity: &AuthenticatedIdentity,
+    name: &str,
+) -> bool {
+    let tid = identity.tenant_id.as_u32();
+    state.mv_registry.get_def(tid, name).is_some()
+}
diff --git a/nodedb/src/control/server/pgwire/ddl/router/mod.rs b/nodedb/src/control/server/pgwire/ddl/router/mod.rs
@@ -1,4 +1,5 @@
 mod admin;
+mod ast;
 mod auth;
 mod collaborative;
 mod dsl;
@@ -26,6 +27,18 @@ pub async fn dispatch(
     identity: &AuthenticatedIdentity,
     sql: &str,
 ) -> Option<PgWireResult<Vec<Response>>> {
+    // AST-typed fast path: parse once, handle IF [NOT] EXISTS at the
+    // dispatch level, then fall through to legacy handlers for the
+    // actual execution. This is the incremental migration path —
+    // once every legacy handler has been ported to accept a typed
+    // NodedbStatement, the string-prefix routers below can be
+    // removed entirely.
+    if let Some(stmt) = nodedb_sql::ddl_ast::parse(sql)
+        && let Some(r) = ast::try_dispatch(state, identity, &stmt)
+    {
+        return Some(r);
+    }
+
     let upper = sql.to_uppercase();
     let parts: Vec<&str> = sql.split_whitespace().collect();
 
diff --git a/nodedb/src/control/server/pgwire/handler/transaction_cmds.rs b/nodedb/src/control/server/pgwire/handler/transaction_cmds.rs
@@ -18,6 +18,7 @@ impl NodeDbPgHandler {
             let next = self.state.wal.next_lsn();
             crate::types::Lsn::new(next.as_u64().saturating_sub(1))
         };
+        crate::control::server::pgwire::session::ddl_buffer::activate();
         self.sessions.begin(addr, snapshot_lsn).map_err(|msg| {
             PgWireError::UserError(Box::new(ErrorInfo::new(
                 "ERROR".to_owned(),
@@ -171,6 +172,35 @@ impl NodeDbPgHandler {
             }
         }
 
+        // Flush any buffered DDL entries as a single atomic batch.
+        if let Some(payloads) = crate::control::server::pgwire::session::ddl_buffer::take()
+            && !payloads.is_empty()
+        {
+            use nodedb_cluster::{MetadataEntry, encode_entry};
+            let sub_entries: Vec<MetadataEntry> = payloads
+                .into_iter()
+                .map(|p| MetadataEntry::CatalogDdl { payload: p })
+                .collect();
+            let batch = MetadataEntry::Batch {
+                entries: sub_entries,
+            };
+            if let Some(handle) = self.state.metadata_raft.get() {
+                let raw = encode_entry(&batch).map_err(|e| {
+                    PgWireError::UserError(Box::new(ErrorInfo::new(
+                        "ERROR".to_owned(),
+                        "XX000".to_owned(),
+                        format!("DDL batch encode: {e}"),
+                    )))
+                })?;
+                handle.propose(raw).map_err(|e| {
+                    PgWireError::UserError(Box::new(ErrorInfo::new(
+                        "ERROR".to_owned(),
+                        "XX000".to_owned(),
+                        format!("DDL batch propose: {e}"),
+                    )))
+                })?;
+            }
+        }
         // Close non-WITH-HOLD cursors on transaction end.
         self.sessions.close_non_hold_cursors(addr);
         Ok(vec![Response::Execution(Tag::new("COMMIT"))])
@@ -182,6 +212,7 @@ impl NodeDbPgHandler {
         identity: &AuthenticatedIdentity,
         addr: &std::net::SocketAddr,
     ) -> PgWireResult<Vec<Response>> {
+        crate::control::server::pgwire::session::ddl_buffer::discard();
         let reservations = self.sessions.rollback(addr).unwrap_or_default();
         for handle in &reservations {
             let key = &handle.sequence_key;
diff --git a/nodedb/src/control/server/pgwire/session/ddl_buffer.rs b/nodedb/src/control/server/pgwire/session/ddl_buffer.rs
diff --git a/nodedb/src/control/server/pgwire/session/mod.rs b/nodedb/src/control/server/pgwire/session/mod.rs
