fix incremental DDL dry run: save/restore tx snapshot between statements

The incremental dry run created a fresh durable transaction for each
statement but processed ops against an accumulated CatalogState from
previous dry runs. For the 2nd+ statement, the tx (reflecting durable
storage) and the state (reflecting accumulated changes) were out of
sync, causing 'retraction does not match existing value' panics when
applying diffs.

Fix: after each dry run, export the transaction's current state as a
Snapshot (Transaction::current_snapshot). On the next dry run, initialize
the transaction from this saved snapshot via
DurableCatalogState::transaction_from_snapshot, keeping tx and state in
sync. This preserves the O(N)-per-statement optimization.

Author: aljoscha

src/adapter/src/catalog/transact.rs

@@ -27,7 +27,7 @@ use mz_audit_log::{
 };
 use mz_catalog::SYSTEM_CONN_ID;
 use mz_catalog::builtin::BuiltinLog;
-use mz_catalog::durable::{NetworkPolicy, Transaction};
+use mz_catalog::durable::{NetworkPolicy, Snapshot, Transaction};
 use mz_catalog::expr_cache::LocalExpressions;
 use mz_catalog::memory::error::{AmbiguousRename, Error, ErrorKind};
 use mz_catalog::memory::objects::{
@@ -504,17 +504,26 @@ impl Catalog {
 
     /// Performs an incremental dry-run catalog transaction: processes only the
     /// NEW ops against an accumulated `CatalogState` from previous dry runs.
-    /// This avoids the O(N^2) replay cost of `catalog_transact_with_ddl_transaction`.
+    /// This avoids the O(N^2) replay cost of replaying all accumulated ops.
     ///
-    /// Returns the new accumulated state and the OID allocator position.
+    /// The durable transaction is intentionally never committed and no storage
+    /// controller prepare-state side effects are run.
+    ///
+    /// If `prev_snapshot` is `Some`, the transaction is initialized from that
+    /// snapshot (which represents the tx state after the previous dry run),
+    /// ensuring it starts in sync with `base_state`. If `None` (first
+    /// statement), a fresh transaction is loaded from durable storage.
+    ///
+    /// Returns the new accumulated state and a snapshot of the transaction's
+    /// state for use in subsequent incremental dry runs.
     pub async fn transact_incremental_dry_run(
         &self,
         base_state: &CatalogState,
         ops: Vec<Op>,
         session: Option<&ConnMeta>,
-        prev_next_oid: Option<u64>,
+        prev_snapshot: Option<Snapshot>,
         oracle_write_ts: mz_repr::Timestamp,
-    ) -> Result<(CatalogState, u64), AdapterError> {
+    ) -> Result<(CatalogState, Snapshot), AdapterError> {
         // For DDL transactions, items are not temporary (CREATE TABLE FROM SOURCE, etc.)
         // but we still need to check for collisions.
         let temporary_ids = self.temporary_ids(&ops, BTreeSet::new())?;
@@ -523,15 +532,20 @@ impl Catalog {
         let mut catalog_updates = vec![];
         let mut audit_events = vec![];
         let mut storage = self.storage().await;
-        let mut tx = storage
-            .transaction()
-            .await
-            .unwrap_or_terminate("starting catalog transaction");
-
-        // Advance OID counter past previously-allocated OIDs from earlier dry runs.
-        if let Some(next_oid) = prev_next_oid {
-            tx.set_next_oid(next_oid)?;
-        }
+        let mut tx = if let Some(snapshot) = prev_snapshot {
+            // Restore transaction from saved snapshot so it starts in sync
+            // with the accumulated CatalogState from previous dry runs.
+            storage
+                .transaction_from_snapshot(snapshot)
+                .unwrap_or_terminate("starting catalog transaction from snapshot")
+        } else {
+            // First statement: fresh transaction from durable storage, which
+            // is in sync with the real catalog state.
+            storage
+                .transaction()
+                .await
+                .unwrap_or_terminate("starting catalog transaction")
+        };
 
         // Process only the new ops against the accumulated state in dry-run mode.
         let new_state = Self::transact_inner(
@@ -549,14 +563,16 @@ impl Catalog {
         )
         .await?;
 
-        let new_next_oid = tx.get_next_oid();
+        // Save the transaction's current state as a snapshot for the next
+        // incremental dry run.
+        let new_snapshot = tx.current_snapshot();
 
         // Transaction is NOT committed — drop it.
         drop(storage);
 
         // transact_inner returns Some(state) when ops produced changes.
         let state = new_state.unwrap_or_else(|| base_state.clone());
-        Ok((state, new_next_oid))
+        Ok((state, new_snapshot))
     }
 
     /// Extracts optimized expressions from `Op::CreateItem` operations for views

src/adapter/src/coord/command_handler.rs

@@ -1151,7 +1151,7 @@ impl Coordinator {
                             state,
                             revision,
                             side_effects: vec![],
-                            next_oid: None,
+                            snapshot: None,
                         }) {
                             return ctx.retire(Err(err));
                         }

src/adapter/src/coord/ddl.rs

@@ -226,7 +226,7 @@ impl Coordinator {
                     ops: txn_ops,
                     revision: txn_revision,
                     state: txn_state,
-                    next_oid: txn_next_oid,
+                    snapshot: txn_snapshot,
                     side_effects: _,
                 },
             ..
@@ -254,7 +254,7 @@ impl Coordinator {
         // Clone what we need from the session before taking &mut below.
         let txn_ops_clone = txn_ops.clone();
         let txn_state_clone = txn_state.clone();
-        let prev_next_oid = *txn_next_oid;
+        let prev_snapshot = txn_snapshot.clone();
 
         // Validate resource limits with all accumulated + new ops (cheap O(N) counting).
         let mut combined_ops = txn_ops_clone;
@@ -268,14 +268,18 @@ impl Coordinator {
         // Get ConnMeta for the session.
         let conn = self.active_conns.get(ctx.session().conn_id());
 
-        // Incremental dry run: only process NEW ops against accumulated state.
-        let (new_state, new_next_oid) = self
+        // Incremental dry run: process only NEW ops against accumulated state.
+        // If we have a saved snapshot from a previous dry run, use it to
+        // initialize the transaction so it starts in sync with the accumulated
+        // state. Otherwise (first statement), the fresh durable transaction is
+        // already in sync with the real catalog state.
+        let (new_state, new_snapshot) = self
             .catalog()
             .transact_incremental_dry_run(
                 &txn_state_clone,
                 ops.clone(),
                 conn,
-                prev_next_oid,
+                prev_snapshot,
                 oracle_write_ts,
             )
             .await?;
@@ -289,7 +293,7 @@ impl Coordinator {
                 state: new_state,
                 side_effects: vec![Box::new(side_effect)],
                 revision: self.catalog().transient_revision(),
-                next_oid: Some(new_next_oid),
+                snapshot: Some(new_snapshot),
             });
 
         self.metrics

src/adapter/src/coord/sequencer/inner.rs

@@ -2163,7 +2163,7 @@ impl Coordinator {
                         state: _,
                         side_effects,
                         revision,
-                        next_oid: _,
+                        snapshot: _,
                     } => {
                         // Make sure our catalog hasn't changed.
                         if *revision != self.catalog().transient_revision() {

src/adapter/src/session.rs

@@ -53,6 +53,7 @@ use uuid::Uuid;
 
 use crate::catalog::CatalogState;
 use crate::client::RecordFirstRowStream;
+use mz_catalog::durable::Snapshot;
 use crate::coord::appends::BuiltinTableAppendNotify;
 use crate::coord::in_memory_oracle::InMemoryTimestampOracle;
 use crate::coord::peek::PeekResponseUnary;
@@ -1397,14 +1398,14 @@ impl<T: TimestampManipulation> TransactionStatus<T> {
                         revision: og_revision,
                         state: og_state,
                         side_effects,
-                        next_oid: og_next_oid,
+                        snapshot: og_snapshot,
                     } => match add_ops {
                         TransactionOps::DDL {
                             ops: new_ops,
                             revision: new_revision,
                             side_effects: mut net_new_side_effects,
                             state: new_state,
-                            next_oid: new_next_oid,
+                            snapshot: new_snapshot,
                         } => {
                             if *og_revision != new_revision {
                                 return Err(AdapterError::DDLTransactionRace);
@@ -1413,7 +1414,7 @@ impl<T: TimestampManipulation> TransactionStatus<T> {
                             if !new_ops.is_empty() {
                                 *og_ops = new_ops;
                                 *og_state = new_state;
-                                *og_next_oid = new_next_oid;
+                                *og_snapshot = new_snapshot;
                             }
                             side_effects.append(&mut net_new_side_effects);
                         }
@@ -1604,10 +1605,11 @@ pub enum TransactionOps<T> {
         >,
         /// Transient revision of the `Catalog` when this transaction started.
         revision: u64,
-        /// Tracks the OID allocator position after the last dry run, so that
-        /// incremental dry runs can advance past previously-allocated OIDs.
-        /// `None` for the first statement in the transaction (before any dry run).
-        next_oid: Option<u64>,
+        /// Snapshot of the durable transaction state after the last dry run.
+        /// Used to initialize the next dry run's transaction so it starts
+        /// in sync with the accumulated `state`. `None` for the first
+        /// statement in the transaction (before any dry run).
+        snapshot: Option<Snapshot>,
     },
 }
 

src/catalog/src/durable.rs

@@ -31,7 +31,8 @@ pub use crate::durable::error::{CatalogError, DurableCatalogError, FenceError};
 pub use crate::durable::metrics::Metrics;
 pub use crate::durable::objects::state_update::StateUpdate;
 use crate::durable::objects::state_update::{StateUpdateKindJson, TryIntoStateUpdateKind};
-use crate::durable::objects::{AuditLog, Snapshot};
+use crate::durable::objects::AuditLog;
+pub use crate::durable::objects::Snapshot;
 pub use crate::durable::objects::{
     Cluster, ClusterConfig, ClusterReplica, ClusterVariant, ClusterVariantManaged, Comment,
     Database, DefaultPrivilege, IntrospectionSourceIndex, Item, NetworkPolicy, ReplicaConfig,
@@ -298,6 +299,16 @@ pub trait DurableCatalogState: ReadOnlyDurableCatalogState {
     /// Creates a new durable catalog state transaction.
     async fn transaction(&mut self) -> Result<Transaction, CatalogError>;
 
+    /// Creates a new transaction initialized from the given [`Snapshot`]
+    /// instead of reading from durable storage. Used for incremental DDL
+    /// dry runs where the transaction state from a previous dry run has been
+    /// saved and needs to be restored so it stays in sync with the accumulated
+    /// `CatalogState`.
+    fn transaction_from_snapshot(
+        &mut self,
+        snapshot: Snapshot,
+    ) -> Result<Transaction, CatalogError>;
+
     /// Commits a durable catalog state transaction. The transaction will be committed at
     /// `commit_ts`.
     ///

src/catalog/src/durable/persist.rs

@@ -1669,6 +1669,14 @@ impl DurableCatalogState for PersistCatalogState {
         Transaction::new(self, snapshot, commit_ts)
     }
 
+    fn transaction_from_snapshot(
+        &mut self,
+        snapshot: Snapshot,
+    ) -> Result<Transaction, CatalogError> {
+        let commit_ts = self.upper.clone();
+        Transaction::new(self, snapshot, commit_ts)
+    }
+
     #[mz_ore::instrument(level = "debug")]
     async fn commit_transaction(
         &mut self,

src/catalog/src/durable/transaction.rs

@@ -1022,31 +1022,37 @@ impl<'a> Transaction<'a> {
             .map(|oids| oids.into_element())
     }
 
-    /// Returns the current next-OID value from the id allocator.
-    pub fn get_next_oid(&self) -> u64 {
-        self.id_allocator
-            .items()
-            .get(&IdAllocKey {
-                name: OID_ALLOC_KEY.to_string(),
-            })
-            .unwrap_or_else(|| panic!("{OID_ALLOC_KEY} id allocator missing"))
-            .next_id
-    }
-
-    /// Advances the OID allocator to `next_oid`. Used for incremental DDL
-    /// transaction dry runs where a fresh Transaction needs to skip past OIDs
-    /// already allocated in previous dry runs.
-    pub fn set_next_oid(&mut self, next_oid: u64) -> Result<(), CatalogError> {
-        self.id_allocator.set(
-            IdAllocKey {
-                name: OID_ALLOC_KEY.to_string(),
-            },
-            Some(IdAllocValue {
-                next_id: next_oid,
-            }),
-            self.op_id,
-        )?;
-        Ok(())
+    /// Exports the current state of this transaction as a [`Snapshot`].
+    ///
+    /// This merges each `TableTransaction`'s initial data with its pending
+    /// changes to produce the current view, then converts back to proto types.
+    /// Used to persist transaction state between incremental DDL dry runs so
+    /// the next dry run's fresh `Transaction` starts in sync with the
+    /// accumulated `CatalogState`.
+    pub fn current_snapshot(&self) -> Snapshot {
+        Snapshot {
+            databases: self.databases.current_items_proto(),
+            schemas: self.schemas.current_items_proto(),
+            roles: self.roles.current_items_proto(),
+            role_auth: self.role_auth.current_items_proto(),
+            items: self.items.current_items_proto(),
+            comments: self.comments.current_items_proto(),
+            clusters: self.clusters.current_items_proto(),
+            network_policies: self.network_policies.current_items_proto(),
+            cluster_replicas: self.cluster_replicas.current_items_proto(),
+            introspection_sources: self.introspection_sources.current_items_proto(),
+            id_allocator: self.id_allocator.current_items_proto(),
+            configs: self.configs.current_items_proto(),
+            settings: self.settings.current_items_proto(),
+            system_object_mappings: self.system_gid_mapping.current_items_proto(),
+            system_configurations: self.system_configurations.current_items_proto(),
+            default_privileges: self.default_privileges.current_items_proto(),
+            source_references: self.source_references.current_items_proto(),
+            system_privileges: self.system_privileges.current_items_proto(),
+            storage_collection_metadata: self.storage_collection_metadata.current_items_proto(),
+            unfinalized_shards: self.unfinalized_shards.current_items_proto(),
+            txn_wal_shard: self.txn_wal_shard.current_items_proto(),
+        }
     }
 
     pub(crate) fn insert_id_allocator(
@@ -3130,6 +3136,22 @@ where
         items
     }
 
+    /// Returns the current items as proto-typed key-value pairs, suitable for
+    /// constructing a [`Snapshot`]. This merges `initial` and `pending` to
+    /// produce the current view and converts back to proto types.
+    fn current_items_proto<KP, VP>(&self) -> BTreeMap<KP, VP>
+    where
+        K: RustType<KP>,
+        V: RustType<VP>,
+        KP: Ord,
+    {
+        let mut items = BTreeMap::new();
+        self.for_values(|k, v| {
+            items.insert(k.into_proto(), v.into_proto());
+        });
+        items
+    }
+
     /// Returns the items viewable in the current transaction as references. Returns a map
     /// of references.
     fn items(&self) -> BTreeMap<&K, &V> {