refactor(engine): one dataset-open chokepoint — open_dataset(uri, VersionResolution, session, wrapper)

Unify the two instrumented openers (open_dataset_tracked latest-only,
open_table_dataset pinned-only) into a single
instrumentation::open_dataset whose version resolution (Latest | At(v))
is an explicit typed parameter, and route every production Dataset open
in the engine through it — including the previously-raw opens in the
branch ops (delete/list/force_delete_branch), manifest branch delete,
version-metadata resolution, recovery classification/restore, the
recovery-audit table, and schema-apply version GC (whose
forbidden-api-allow sentinel is now unnecessary and removed).

One chokepoint means three properties hold uniformly instead of
per-path: record_open feeds the cost probes on every open (the raw
sites were invisible to the gates), the per-query IO wrapper is
attached with the right class (manifest vs table) everywhere, and the
shared per-graph Session parameter exists on every path — threading it
into the write-side owners lands separately.

Behavior-preserving: every call site keeps its existing version
resolution, session argument, and wrapper class; in-source tests keep
raw opens (out of scope for the chokepoint by design).

crates/omnigraph/src/db/manifest.rs

@@ -256,13 +256,19 @@ impl SubTableEntry {
         // matches the physical layout the namespace path resolved, without the
         // per-open `__manifest` scan.
         let location = table_uri_for_path(root_uri, &self.table_path, self.table_branch.as_deref());
-        // Route through the instrumented data-table opener (Fix 3). With no
-        // session this is exactly the Fix-2 `from_uri(location).with_version`.
-        // This is the uncached fallback (a snapshot with no read caches); the
+        // Route through the one opener (Fix 3). With no session this is exactly
+        // the Fix-2 `from_uri(location).with_version`. This is the uncached
+        // fallback (a snapshot detached from its graph's read caches); the
         // cached path (`Snapshot::open` → handle cache) calls the same opener on
         // a miss with the shared session, so both paths count on the per-query
         // `table_wrapper`.
-        crate::instrumentation::open_table_dataset(&location, self.table_version, None).await
+        crate::instrumentation::open_dataset(
+            &location,
+            crate::instrumentation::VersionResolution::At(self.table_version),
+            None,
+            crate::instrumentation::table_wrapper(),
+        )
+        .await
     }
 }
 
@@ -586,9 +592,13 @@ impl ManifestCoordinator {
 
     pub async fn delete_branch(&mut self, name: &str) -> Result<()> {
         let uri = manifest_uri(&self.root_uri);
-        let mut ds = Dataset::open(&uri)
-            .await
-            .map_err(|e| OmniError::Lance(e.to_string()))?;
+        let mut ds = crate::instrumentation::open_dataset(
+            &uri,
+            crate::instrumentation::VersionResolution::Latest,
+            None,
+            crate::instrumentation::manifest_wrapper(),
+        )
+        .await?;
         ds.delete_branch(name)
             .await
             .map_err(|e| OmniError::Lance(e.to_string()))?;

crates/omnigraph/src/db/manifest/layout.rs

@@ -21,8 +21,10 @@ pub(crate) fn manifest_uri(root: &str) -> String {
 
 pub(super) async fn open_manifest_dataset(root_uri: &str, branch: Option<&str>) -> Result<Dataset> {
     let uri = manifest_uri(root_uri.trim_end_matches('/'));
-    let dataset = crate::instrumentation::open_dataset_tracked(
+    let dataset = crate::instrumentation::open_dataset(
         &uri,
+        crate::instrumentation::VersionResolution::Latest,
+        None,
         crate::instrumentation::manifest_wrapper(),
     )
     .await?;

crates/omnigraph/src/db/manifest/metadata.rs

@@ -229,9 +229,13 @@ pub(super) async fn table_version_metadata_for_state(
     version: u64,
 ) -> Result<TableVersionMetadata> {
     let full_path = format!("{}/{}", root_uri.trim_end_matches('/'), table_path);
-    let ds = Dataset::open(&full_path)
-        .await
-        .map_err(|e| OmniError::Lance(e.to_string()))?;
+    let ds = crate::instrumentation::open_dataset(
+        &full_path,
+        crate::instrumentation::VersionResolution::Latest,
+        None,
+        crate::instrumentation::table_wrapper(),
+    )
+    .await?;
     let ds = match branch {
         Some(branch) => ds
             .checkout_branch(branch)

crates/omnigraph/src/db/manifest/namespace.rs

@@ -107,9 +107,13 @@ impl StagedTableNamespace {
     }
 
     async fn open_head(&self) -> Result<Dataset> {
-        Dataset::open(&self.table_uri())
-            .await
-            .map_err(|e| OmniError::Lance(e.to_string()))
+        crate::instrumentation::open_dataset(
+            &self.table_uri(),
+            crate::instrumentation::VersionResolution::Latest,
+            None,
+            crate::instrumentation::manifest_wrapper(),
+        )
+        .await
     }
 
     async fn open_version(&self, version: u64) -> Result<Dataset> {

crates/omnigraph/src/db/manifest/recovery.rs

@@ -36,7 +36,6 @@
 
 use std::collections::HashMap;
 
-use lance::Dataset;
 use serde::{Deserialize, Serialize};
 use tracing::warn;
 
@@ -737,9 +736,13 @@ pub(crate) async fn restore_table_to_version(
     branch: Option<&str>,
     target_version: u64,
 ) -> Result<()> {
-    let head = Dataset::open(table_path)
-        .await
-        .map_err(|e| OmniError::Lance(e.to_string()))?;
+    let head = crate::instrumentation::open_dataset(
+        table_path,
+        crate::instrumentation::VersionResolution::Latest,
+        None,
+        crate::instrumentation::table_wrapper(),
+    )
+    .await?;
     let head = match branch {
         Some(b) if b != "main" => head
             .checkout_branch(b)
@@ -1729,9 +1732,13 @@ async fn roll_forward_all(
             continue;
         }
         let dataset_uri = format!("{}/{}", root_uri.trim_end_matches('/'), reg.table_path);
-        let head_ds = Dataset::open(&dataset_uri)
-            .await
-            .map_err(|e| OmniError::Lance(e.to_string()))?;
+        let head_ds = crate::instrumentation::open_dataset(
+            &dataset_uri,
+            crate::instrumentation::VersionResolution::Latest,
+            None,
+            crate::instrumentation::table_wrapper(),
+        )
+        .await?;
         let head_ds = match reg.table_branch.as_deref() {
             Some(b) if b != "main" => head_ds
                 .checkout_branch(b)
@@ -1826,9 +1833,13 @@ async fn push_table_update(
     updates: &mut Vec<ManifestChange>,
     expected: &mut HashMap<String, u64>,
 ) -> Result<u64> {
-    let ds = Dataset::open(table_path)
-        .await
-        .map_err(|e| OmniError::Lance(e.to_string()))?;
+    let ds = crate::instrumentation::open_dataset(
+        table_path,
+        crate::instrumentation::VersionResolution::Latest,
+        None,
+        crate::instrumentation::table_wrapper(),
+    )
+    .await?;
     let ds = match branch {
         Some(b) if b != "main" => ds
             .checkout_branch(b)
@@ -1924,9 +1935,13 @@ pub(crate) async fn has_schema_apply_sidecar(
 /// HEAD version. Recovery uses this so feature-branch sidecars classify
 /// against the feature-branch's Lance HEAD, not main's.
 async fn open_lance_head(table_path: &str, branch: Option<&str>) -> Result<u64> {
-    let ds = Dataset::open(table_path)
-        .await
-        .map_err(|e| OmniError::Lance(e.to_string()))?;
+    let ds = crate::instrumentation::open_dataset(
+        table_path,
+        crate::instrumentation::VersionResolution::Latest,
+        None,
+        crate::instrumentation::table_wrapper(),
+    )
+    .await?;
     let ds = match branch {
         Some(b) if b != "main" => ds
             .checkout_branch(b)
@@ -1977,6 +1992,7 @@ pub(crate) fn new_sidecar(
 mod tests {
     use super::*;
     use crate::storage::ObjectStorageAdapter;
+    use lance::Dataset;
     use crate::table_store::TableStore;
     use arrow_array::{Int32Array, RecordBatch, StringArray};
     use arrow_schema::{DataType, Field, Schema};

crates/omnigraph/src/db/omnigraph/schema_apply.rs

@@ -751,10 +751,13 @@ where
 async fn cleanup_dataset_old_versions(db: &Omnigraph, full_uri: &str) -> Result<()> {
     use chrono::Utc;
     use lance::dataset::cleanup::CleanupPolicy;
-    // forbidden-api-allow: maintenance (Hard-drop version GC) opens the dataset to run cleanup_old_versions.
-    let ds = lance::Dataset::open(full_uri)
-        .await
-        .map_err(|e| OmniError::Lance(e.to_string()))?;
+    let ds = crate::instrumentation::open_dataset(
+        full_uri,
+        crate::instrumentation::VersionResolution::Latest,
+        None,
+        crate::instrumentation::table_wrapper(),
+    )
+    .await?;
     let policy = CleanupPolicy {
         before_timestamp: Some(Utc::now()),
         before_version: None,

crates/omnigraph/src/db/recovery_audit.rs

@@ -101,7 +101,14 @@ impl RecoveryAudit {
     /// with no dataset yet (it is created lazily on the first append).
     pub(crate) async fn open(root_uri: &str) -> Result<Self> {
         let root = root_uri.trim_end_matches('/').to_string();
-        let dataset = Dataset::open(&recoveries_uri(&root)).await.ok();
+        let dataset = crate::instrumentation::open_dataset(
+            &recoveries_uri(&root),
+            crate::instrumentation::VersionResolution::Latest,
+            None,
+            crate::instrumentation::manifest_wrapper(),
+        )
+        .await
+        .ok();
         Ok(Self {
             root_uri: root,
             dataset,
@@ -195,9 +202,15 @@ async fn create_recoveries_dataset(root_uri: &str) -> Result<Dataset> {
         // variant, not the display string (not a Lance API contract). Same
         // discipline as `commit_graph.rs`'s create-or-open; pinned by
         // `lance_surface_guards.rs::lance_error_dataset_already_exists_variant_exists`.
-        Err(lance::Error::DatasetAlreadyExists { .. }) => Dataset::open(&uri)
+        Err(lance::Error::DatasetAlreadyExists { .. }) => {
+            crate::instrumentation::open_dataset(
+                &uri,
+                crate::instrumentation::VersionResolution::Latest,
+                None,
+                crate::instrumentation::manifest_wrapper(),
+            )
             .await
-            .map_err(|open_err| OmniError::Lance(open_err.to_string())),
+        }
         Err(err) => Err(OmniError::Lance(err.to_string())),
     }
 }

crates/omnigraph/src/instrumentation.rs

@@ -252,47 +252,49 @@ pub(crate) fn record_scan_staged_combined() {
     });
 }
 
-/// Open a Lance dataset at `uri`, attaching `wrapper` (for IO counting) when
-/// present. With no wrapper this is exactly `Dataset::open(uri)`. The wrapper is
-/// set via `ObjectStoreParams` on the builder so the open itself is counted
-/// (`Dataset::with_object_store_wrappers` only wraps an already-open store).
-pub(crate) async fn open_dataset_tracked(
+/// Which version [`open_dataset`] resolves.
+///
+/// `Latest` re-resolves the dataset's current head (the substrate's cheap
+/// latest-location probe); `At(v)` is a list-free pinned open. The choice is
+/// a correctness decision — strict read-modify-write ops need `Latest`,
+/// snapshot reads need `At(v)` — so it is an explicit parameter of the one
+/// opener rather than a property of which helper a caller happened to reach.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum VersionResolution {
+    Latest,
+    At(u64),
+}
+
+/// THE dataset-open chokepoint. Every engine `Dataset` open routes through
+/// here so three things hold uniformly, on every path:
+///
+/// 1. `record_open` feeds the per-query cost probes — an open that bypasses
+///    this function is invisible to the cost gates.
+/// 2. The per-query IO `wrapper` (manifest- or table-class) is set via
+///    `ObjectStoreParams` on the builder, so the open itself is counted
+///    (`Dataset::with_object_store_wrappers` only wraps an already-open
+///    store). No wrapper (production) adds nothing.
+/// 3. The shared per-graph `Session` (LanceDB's one-session-per-connection
+///    pattern; warms Lance's metadata/index caches across opens) is attached
+///    whenever the caller has one. `None` is for genuinely session-less
+///    contexts (a `Snapshot` detached from its graph's read caches) — owners
+///    that hold a session (`TableStore`, the handle cache) pass it
+///    unconditionally, so it cannot be silently dropped on one path.
+pub(crate) async fn open_dataset(
     uri: &str,
+    version: VersionResolution,
+    session: Option<&Arc<lance::session::Session>>,
     wrapper: Option<Arc<dyn WrappingObjectStore>>,
 ) -> Result<Dataset> {
     record_open(uri);
-    let result = match wrapper {
-        None => Dataset::open(uri).await,
-        Some(wrapper) => {
-            DatasetBuilder::from_uri(uri)
-                .with_store_params(ObjectStoreParams {
-                    object_store_wrapper: Some(wrapper),
-                    ..Default::default()
-                })
-                .load()
-                .await
-        }
-    };
-    result.map_err(|e| OmniError::Lance(e.to_string()))
-}
-
-/// Open a data-table dataset at `location` pinned to `version` — the cache-miss
-/// path of the data-read boundary (`SubTableEntry::open`). Attaches the shared
-/// per-graph `Session` (warms metadata/index caches across opens, LanceDB's
-/// one-session-per-connection pattern) and the per-query `table_wrapper` (for IO
-/// counting) when present. With neither, this is exactly the Fix-2
-/// `from_uri(location).with_version(version)` open.
-pub(crate) async fn open_table_dataset(
-    location: &str,
-    version: u64,
-    session: Option<&Arc<lance::session::Session>>,
-) -> Result<Dataset> {
-    record_open(location);
-    let mut builder = DatasetBuilder::from_uri(location).with_version(version);
+    let mut builder = DatasetBuilder::from_uri(uri);
+    if let VersionResolution::At(version) = version {
+        builder = builder.with_version(version);
+    }
     if let Some(session) = session {
         builder = builder.with_session(session.clone());
     }
-    if let Some(wrapper) = table_wrapper() {
+    if let Some(wrapper) = wrapper {
         builder = builder.with_store_params(ObjectStoreParams {
             object_store_wrapper: Some(wrapper),
             ..Default::default()

crates/omnigraph/src/runtime_cache.rs

@@ -264,7 +264,13 @@ impl TableHandleCache {
         // Miss: open without holding the lock (the open is async IO). A concurrent
         // double-miss opens twice and one wins the insert — correct (the dataset
         // at a version is immutable) and rare.
-        let ds = crate::instrumentation::open_table_dataset(location, version, session).await?;
+        let ds = crate::instrumentation::open_dataset(
+            location,
+            crate::instrumentation::VersionResolution::At(version),
+            session,
+            crate::instrumentation::table_wrapper(),
+        )
+        .await?;
         let mut inner = self.inner.lock().await;
         if let Some(existing) = inner.entries.get(&key).cloned() {
             return Ok(existing);

crates/omnigraph/src/table_store.rs

@@ -202,12 +202,14 @@ impl TableStore {
         dataset_uri: &str,
         branch: Option<&str>,
     ) -> Result<Dataset> {
-        // Direct open by URI (O(1) latest-resolution). Routed through the tracked
+        // Direct open by URI (O(1) latest-resolution). Routed through the one
         // opener so a cost test counts it via the per-query `table_wrapper`
         // (no-op in production — the task-local is unset, so this is exactly
         // `Dataset::open(uri)`).
-        let ds = crate::instrumentation::open_dataset_tracked(
+        let ds = crate::instrumentation::open_dataset(
             dataset_uri,
+            crate::instrumentation::VersionResolution::Latest,
+            None,
             crate::instrumentation::table_wrapper(),
         )
         .await?;
@@ -237,9 +239,13 @@ impl TableStore {
     }
 
     pub async fn delete_branch(&self, dataset_uri: &str, branch: &str) -> Result<()> {
-        let mut ds = Dataset::open(dataset_uri)
-            .await
-            .map_err(|e| OmniError::Lance(e.to_string()))?;
+        let mut ds = crate::instrumentation::open_dataset(
+            dataset_uri,
+            crate::instrumentation::VersionResolution::Latest,
+            None,
+            crate::instrumentation::table_wrapper(),
+        )
+        .await?;
         ds.delete_branch(branch)
             .await
             .map_err(|e| OmniError::Lance(e.to_string()))
@@ -250,9 +256,13 @@ impl TableStore {
     /// set to find orphaned per-table forks. `main`/default is not a named
     /// branch and never appears here.
     pub async fn list_branches(&self, dataset_uri: &str) -> Result<Vec<String>> {
-        let ds = Dataset::open(dataset_uri)
-            .await
-            .map_err(|e| OmniError::Lance(e.to_string()))?;
+        let ds = crate::instrumentation::open_dataset(
+            dataset_uri,
+            crate::instrumentation::VersionResolution::Latest,
+            None,
+            crate::instrumentation::table_wrapper(),
+        )
+        .await?;
         let branches = ds
             .list_branches()
             .await
@@ -274,9 +284,13 @@ impl TableStore {
     /// Used by the eager best-effort reclaim in `cleanup_deleted_branch_tables`
     /// and the `cleanup` orphan reconciler.
     pub async fn force_delete_branch(&self, dataset_uri: &str, branch: &str) -> Result<()> {
-        let mut ds = Dataset::open(dataset_uri)
-            .await
-            .map_err(|e| OmniError::Lance(e.to_string()))?;
+        let mut ds = crate::instrumentation::open_dataset(
+            dataset_uri,
+            crate::instrumentation::VersionResolution::Latest,
+            None,
+            crate::instrumentation::table_wrapper(),
+        )
+        .await?;
         match ds.force_delete_branch(branch).await {
             Ok(()) => Ok(()),
             Err(lance::Error::RefNotFound { .. }) | Err(lance::Error::NotFound { .. }) => Ok(()),