From ef9e7354528e4162f7da98132964a53e60797d0f Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 13 May 2026 04:02:58 +0000
Subject: [PATCH] MR-854: convert remaining table_store call sites in export.rs
 / read_blob
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two leftover `self.table_store.X` / `db.table_store.X` call sites were
missed in the initial sweep — flagged by Devin Review on PR #86. Both
now go through the trait surface:

- `entity_from_snapshot` (db/omnigraph/export.rs): switch from
  `db.table_store.open_snapshot_table` + `db.table_store.scan` to
  `db.storage().open_snapshot_at_table` + `db.storage().scan`.

- `read_blob` (db/omnigraph.rs): replace
  `snapshot.open(table_key)` + `self.table_store.first_row_id_for_filter`
  with `self.storage().open_snapshot_at_table` +
  `self.storage().first_row_id_for_filter`. The follow-up
  `take_blobs` call still needs an `Arc<Dataset>` (it's a Lance blob
  accessor not surfaced through the trait), so we hand off via
  `SnapshotHandle::into_arc()` with a comment.

After this commit, no engine code outside `table_store.rs` reaches the
inherent `TableStore` API — the docs/runs.md and docs/invariants.md
claim is now uniformly true.

Co-Authored-By: Ragnor Comerford <ragnor.comerford@gmail.com>
---
 crates/omnigraph/src/db/omnigraph.rs        | 15 ++++++++++-----
 crates/omnigraph/src/db/omnigraph/export.rs |  6 +++---
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/crates/omnigraph/src/db/omnigraph.rs b/crates/omnigraph/src/db/omnigraph.rs
index 4eee566..f10739e 100644
--- a/crates/omnigraph/src/db/omnigraph.rs
+++ b/crates/omnigraph/src/db/omnigraph.rs
@@ -1022,19 +1022,24 @@ impl Omnigraph {
 
         let snapshot = self.snapshot().await;
         let table_key = format!("node:{}", type_name);
-        let ds = snapshot.open(&table_key).await?;
+        let handle = self
+            .storage()
+            .open_snapshot_at_table(&snapshot, &table_key)
+            .await?;
 
         let filter_sql = format!("id = '{}'", id.replace('\'', "''"));
         let row_id = self
-            .table_store
-            .first_row_id_for_filter(&ds, &filter_sql)
+            .storage()
+            .first_row_id_for_filter(&handle, &filter_sql)
             .await?
             .ok_or_else(|| {
                 OmniError::manifest(format!("no {} with id '{}' found", type_name, id))
             })?;
 
-        // Use take_blobs to get the BlobFile handle
-        let ds = Arc::new(ds);
+        // `take_blobs` is a Lance-specific blob accessor not surfaced
+        // through the `TableStorage` trait — reach the inner `Arc<Dataset>`
+        // via the `pub(crate)` accessor for this read-only call.
+        let ds = handle.into_arc();
         let mut blobs = ds
             .take_blobs(&[row_id], property)
             .await
diff --git a/crates/omnigraph/src/db/omnigraph/export.rs b/crates/omnigraph/src/db/omnigraph/export.rs
index 3043ce9..7696056 100644
--- a/crates/omnigraph/src/db/omnigraph/export.rs
+++ b/crates/omnigraph/src/db/omnigraph/export.rs
@@ -60,12 +60,12 @@ async fn entity_from_snapshot(
     }
 
     let ds = db
-        .table_store
-        .open_snapshot_table(snapshot, table_key)
+        .storage()
+        .open_snapshot_at_table(snapshot, table_key)
         .await?;
     let filter_sql = format!("id = '{}'", id.replace('\'', "''"));
     let batches = db
-        .table_store
+        .storage()
         .scan(&ds, None, Some(&filter_sql), None)
         .await?;
     let Some(batch) = batches.iter().find(|batch| batch.num_rows() > 0) else {