feat(engine): WriteTxn - validate schema + open each data table once per write

crates/omnigraph/src/db/mod.rs

@@ -10,6 +10,7 @@ pub use commit_graph::GraphCommit;
 pub use graph_coordinator::{GraphCoordinator, ReadTarget, ResolvedTarget, SnapshotId};
 pub use manifest::{Snapshot, SubTableEntry, SubTableUpdate};
 pub(crate) use omnigraph::ensure_public_branch_ref;
+pub(crate) use omnigraph::WriteTxn;
 pub use omnigraph::{
     CleanupPolicyOptions, InitOptions, MergeOutcome, Omnigraph, OpenMode, PendingIndex,
     RepairAction, RepairClassification, RepairOptions, RepairStats, SchemaApplyOptions,

crates/omnigraph/src/db/omnigraph.rs

@@ -41,6 +41,7 @@ pub use repair::{
 };
 pub use schema_apply::SchemaApplyOptions;
 pub use table_ops::PendingIndex;
+pub(crate) use table_ops::OpenedForMutation;
 
 use super::commit_graph::GraphCommit;
 use super::manifest::{
@@ -79,6 +80,35 @@ pub struct SchemaApplyPreview {
     pub catalog: Catalog,
 }
 
+/// A capture-once write transaction (RFC-013 step 3b). Pins the operation's read
+/// base ONCE so the per-table opens reuse the pinned version instead of
+/// re-resolving / re-validating per table. The schema contract is validated once
+/// (when `base` is captured). NOT a general "no re-resolution" handle — the
+/// commit-time OCC re-read, the live-HEAD drift probe, and the fork-authority reads
+/// stay fresh (correctness machinery). Step 5 (PublishPlan unification) makes this
+/// the non-optional publish carrier and adds session-aware base opens there, gated
+/// by an S3 cost test — the warm-session benefit on the single remaining open is an
+/// object-store phenomenon, so it earns its own gate rather than riding this PR.
+///
+/// Threaded as `Option<&WriteTxn>` through the mutate/load write chain
+/// (`open_for_mutation_on_branch`, `commit_all`, `commit_updates_on_branch_with_expected`)
+/// so a single write validates the schema contract EXACTLY ONCE — at capture. When
+/// present, the per-table resolves source the pinned `base` entry instead of calling
+/// `resolved_branch_target` / `snapshot_for_branch` / `fresh_snapshot_for_branch`
+/// (each of which re-runs `ensure_schema_state_valid`). When absent (`None` — every
+/// non-mutate/load caller), every threaded function behaves byte-identically to
+/// before. The carrier never removes a version guard or changes which dataset version
+/// the per-table open targets: strict ops keep `open_dataset_head_for_write` +
+/// `ensure_expected_version`, and the commit-time OCC re-read still opens a fresh
+/// manifest snapshot (via `fresh_snapshot_for_branch_unchecked`) — only the redundant
+/// schema re-validation is dropped.
+pub(crate) struct WriteTxn {
+    /// The resolved branch (`None` = main).
+    pub(crate) branch: Option<String>,
+    /// The pinned base snapshot (per-table location + version + e_tag), captured once.
+    pub(crate) base: Snapshot,
+}
+
 /// Top-level handle to an Omnigraph database.
 ///
 /// An Omnigraph is a Lance-native graph database with git-style branching.
@@ -736,6 +766,21 @@ impl Omnigraph {
         *self.coordinator.write().await = coordinator;
     }
 
+    /// Open a capture-once write transaction (RFC-013 step 3b): validate the schema
+    /// contract ONCE and pin the base snapshot + the shared per-graph `Session`. The
+    /// per-table opens take `Option<&WriteTxn>` and, on the bound branch for the
+    /// non-strict (Insert/Merge) path, source the pinned base entry — instead of
+    /// re-resolving (re-validating the schema) per table. Strict ops, the fork path,
+    /// and the commit-time OCC re-read keep their fresh reads (those are correctness
+    /// machinery — see the handoff doc).
+    pub(crate) async fn open_write_txn(&self, branch: Option<&str>) -> Result<WriteTxn> {
+        let resolved = self.resolved_branch_target(branch).await?;
+        Ok(WriteTxn {
+            branch: resolved.branch,
+            base: resolved.snapshot,
+        })
+    }
+
     pub(crate) async fn resolved_branch_target(
         &self,
         branch: Option<&str>,
@@ -770,12 +815,39 @@ impl Omnigraph {
 
     pub(crate) async fn fresh_snapshot_for_branch(&self, branch: Option<&str>) -> Result<Snapshot> {
         self.ensure_schema_state_valid().await?;
-        let requested = ReadTarget::Branch(branch.unwrap_or("main").to_string());
-        let coord = self.coordinator.read().await;
-        coord
-            .resolve_target(&requested)
-            .await
-            .map(|resolved| resolved.snapshot)
+        self.fresh_snapshot_for_branch_unchecked(branch).await
+    }
+
+    /// Fresh per-branch manifest snapshot WITHOUT the schema-contract
+    /// re-validation. Identical OCC freshness to [`fresh_snapshot_for_branch`]
+    /// — a fresh manifest re-read from storage, never the warm cache — only the
+    /// redundant `ensure_schema_state_valid` is dropped. Used inside a single
+    /// write once a `WriteTxn` has already validated the contract at capture: the
+    /// commit-time drift re-read needs the live manifest, not a second contract
+    /// read. Callers with no `WriteTxn` MUST use the checked variant.
+    ///
+    /// Reads the manifest directly via `ManifestCoordinator` rather than
+    /// `resolve_target`. The OCC re-read uses only the returned `Snapshot`
+    /// (per-table location + version), which `ManifestCoordinator::open().snapshot()`
+    /// produces identically to `GraphCoordinator::open(...).snapshot()` — but
+    /// `resolve_target` additionally opens the commit graph (an extra
+    /// `_graph_commits.lance` probe) the OCC read never consults. Skipping that
+    /// load is a pure read-cost reduction, not a freshness change. The checked
+    /// `fresh_snapshot_for_branch` delegates here, so its no-`txn` callers
+    /// (commit_all's None arm, optimize, repair, fork reclaim) get the same
+    /// identical `Snapshot` via this lighter manifest-only read; they consume
+    /// only the snapshot and never relied on the commit-graph side load.
+    pub(crate) async fn fresh_snapshot_for_branch_unchecked(
+        &self,
+        branch: Option<&str>,
+    ) -> Result<Snapshot> {
+        let manifest = match branch {
+            Some(branch) => {
+                crate::db::manifest::ManifestCoordinator::open_at_branch(self.uri(), branch).await?
+            }
+            None => crate::db::manifest::ManifestCoordinator::open(self.uri()).await?,
+        };
+        Ok(manifest.snapshot())
     }
 
     pub(crate) async fn version(&self) -> u64 {
@@ -1599,7 +1671,7 @@ impl Omnigraph {
         &self,
         table_key: &str,
         op_kind: crate::db::MutationOpKind,
-    ) -> Result<(SnapshotHandle, String, Option<String>)> {
+    ) -> Result<OpenedForMutation> {
         table_ops::open_for_mutation(self, table_key, op_kind).await
     }
 
@@ -1608,8 +1680,9 @@ impl Omnigraph {
         branch: Option<&str>,
         table_key: &str,
         op_kind: crate::db::MutationOpKind,
-    ) -> Result<(SnapshotHandle, String, Option<String>)> {
-        table_ops::open_for_mutation_on_branch(self, branch, table_key, op_kind).await
+        txn: Option<&crate::db::WriteTxn>,
+    ) -> Result<OpenedForMutation> {
+        table_ops::open_for_mutation_on_branch(self, branch, table_key, op_kind, txn).await
     }
 
     /// Fork `table_key` onto `active_branch` from the given source state,
@@ -1728,13 +1801,17 @@ impl Omnigraph {
         updates: &[crate::db::SubTableUpdate],
         expected_table_versions: &std::collections::HashMap<String, u64>,
         actor_id: Option<&str>,
+        txn: Option<&crate::db::WriteTxn>,
+        committed_handles: std::collections::HashMap<String, crate::storage_layer::SnapshotHandle>,
     ) -> Result<u64> {
         table_ops::commit_updates_on_branch_with_expected(
             self,
             branch,
             updates,
             expected_table_versions,
             actor_id,
+            txn,
+            committed_handles,
         )
         .await
     }
@@ -2466,10 +2543,13 @@ edge WorksAt: Person -> Company
     }
 
     async fn seed_person_row(db: &mut Omnigraph, name: &str, age: Option<i32>) {
+        // No-txn entry, so the handle is always `Some` (collapse #1's skip is
+        // gated on `txn.is_some()`).
         let (ds, full_path, table_branch) = db
             .open_for_mutation("node:Person", crate::db::MutationOpKind::Insert)
             .await
-            .unwrap();
+            .unwrap()
+            .require_handle("seed_person_row test");
         let schema: Arc<Schema> = Arc::new(ds.dataset().schema().into());
         let columns: Vec<Arc<dyn Array>> = schema
             .fields()