test(engine): shared test_session() helper for TableStore construction

Greptile review follow-up on the session-threading commit: the
Arc::new(Session::default()) boilerplate at every test TableStore
construction collapses into helpers::test_session() (and a local twin in
staged_writes.rs, which is primitive-level and deliberately does not
include the helpers module) — a future session-constructor change is a
single-point edit.

crates/omnigraph/tests/failpoints.rs

@@ -15,7 +15,8 @@ use helpers::recovery::{
     branch_head_commit_id, single_sidecar_operation_id,
 };
 use helpers::{
-    MUTATION_QUERIES, collect_column_strings, mixed_params, mutate_main, read_table, version_main,
+    MUTATION_QUERIES, collect_column_strings, mixed_params, mutate_main, read_table, test_session,
+    version_main,
 };
 
 const SCHEMA_V1: &str = "node Person { name: String @key }\n";
@@ -1131,7 +1132,7 @@ async fn recovery_rolls_forward_ensure_indices_on_feature_branch() {
     // publisher deliberately skips the normal index-rebuild preparation;
     // the failed writer below is still the real `ensure_indices_on`.
     let person_uri = node_table_uri(&uri, "Person");
-    let store = TableStore::new(&uri, std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(&uri, test_session());
     let mut ds = store
         .open_dataset_head(&person_uri, Some("feature"))
         .await

crates/omnigraph/tests/helpers/mod.rs

@@ -47,6 +47,13 @@ query insert_person_and_friend($name: String, $age: I32, $friend: String) {
 }
 "#;
 
+/// A standalone Lance `Session` for tests that construct a `TableStore`
+/// directly (production stores share the graph's per-connection session;
+/// tests get a fresh one — the cache scope is the test).
+pub fn test_session() -> std::sync::Arc<lance::session::Session> {
+    std::sync::Arc::new(lance::session::Session::default())
+}
+
 /// Init a graph and load the standard test data.
 pub async fn init_and_load(dir: &tempfile::TempDir) -> Omnigraph {
     let uri = dir.path().to_str().unwrap();

crates/omnigraph/tests/recovery.rs

@@ -18,6 +18,7 @@ use lance::Dataset;
 use omnigraph::db::Omnigraph;
 
 mod helpers;
+use helpers::test_session;
 use helpers::recovery::{RecoveryExpectation, TableExpectation, assert_post_recovery_invariants};
 
 const TEST_SCHEMA: &str = include_str!("fixtures/test.pg");
@@ -1441,7 +1442,7 @@ async fn recovery_classifies_feature_branch_sidecar_against_feature_branch() {
     // Bypass the manifest: append directly to Person's Lance HEAD on the
     // feature branch ref to advance HEAD past v_pin.
     let person_uri = node_table_uri(uri, "Person");
-    let store = TableStore::new(uri, std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(uri, test_session());
     let mut ds = store
         .open_dataset_head(&person_uri, feature_branch_name.as_deref())
         .await
@@ -1556,7 +1557,7 @@ async fn recovery_rolls_back_feature_branch_sidecar_against_feature_branch() {
     // Bypass the manifest: append on the feature ref to advance HEAD past
     // the manifest pin.
     let person_uri = node_table_uri(uri, "Person");
-    let store = TableStore::new(uri, std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(uri, test_session());
     let mut ds = store
         .open_dataset_head(&person_uri, feature_branch_name.as_deref())
         .await

crates/omnigraph/tests/staged_writes.rs

@@ -28,6 +28,12 @@ use lance_index::IndexType;
 use lance_linalg::distance::MetricType;
 use lance_table::format::Fragment;
 use omnigraph::table_store::{StagedWrite, TableStore};
+
+/// A standalone Lance `Session` per test store (this binary is primitive-level
+/// and deliberately does not include the shared `helpers` module).
+fn test_session() -> std::sync::Arc<lance::session::Session> {
+    std::sync::Arc::new(lance::session::Session::default())
+}
 use std::sync::Arc;
 
 fn person_schema() -> Arc<Schema> {
@@ -123,7 +129,7 @@ fn collect_age_for_id(batches: &[RecordBatch], needle: &str) -> Option<i32> {
 async fn stage_append_is_visible_via_scan_with_staged() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     // Seed: one committed row.
     let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
@@ -153,7 +159,7 @@ async fn stage_append_is_visible_via_scan_with_staged() {
 async fn stage_merge_insert_dedupes_superseded_committed_fragment() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     // Seed: alice age 30 in one committed fragment.
     let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
@@ -215,7 +221,7 @@ async fn stage_merge_insert_dedupes_superseded_committed_fragment() {
 async fn count_rows_with_staged_matches_scan() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
         .await
@@ -250,7 +256,7 @@ async fn count_rows_with_staged_matches_scan() {
 async fn chained_stage_appends_have_distinct_row_ids() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(&uri, person_batch(&[("seed", Some(0))]))
         .await
@@ -336,7 +342,7 @@ fn combine_for_scan(ds: &Dataset, staged: &[StagedWrite]) -> Vec<Fragment> {
 async fn stage_append_then_commit_persists_data() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
         .await
@@ -369,7 +375,7 @@ async fn stage_append_then_commit_persists_data() {
 async fn stage_merge_insert_then_commit_persists_merged_view() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
         .await
@@ -401,7 +407,7 @@ async fn stage_merge_insert_then_commit_persists_merged_view() {
 async fn stage_merge_insert_commit_rebases_over_disjoint_committed_delete() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(&uri, numbered_person_batch(0..100))
         .await
@@ -461,7 +467,7 @@ async fn stage_merge_insert_commit_rebases_over_disjoint_committed_delete() {
 async fn scan_with_staged_with_filter_silently_drops_staged_rows() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     // Committed: alice=30, carol=40
     let ds = TableStore::write_dataset(
@@ -528,7 +534,7 @@ async fn scan_with_staged_with_filter_silently_drops_staged_rows() {
 async fn chained_stage_merge_insert_with_shared_key_documents_duplicate_behavior() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     // Seed empty (an unrelated row keeps the schema unambiguous).
     let ds = TableStore::write_dataset(&uri, person_batch(&[("seed", Some(0))]))
@@ -594,7 +600,7 @@ async fn chained_stage_merge_insert_with_shared_key_documents_duplicate_behavior
 async fn stage_overwrite_does_not_advance_head_until_commit() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
         .await
@@ -638,7 +644,7 @@ async fn stage_overwrite_does_not_advance_head_until_commit() {
 async fn stage_overwrite_preserves_stable_row_ids() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     // `write_dataset` creates with `enable_stable_row_ids: true` — see
     // ADR 0001. We verify that as a precondition so a future change to
@@ -680,7 +686,7 @@ async fn stage_overwrite_preserves_stable_row_ids() {
 async fn stage_overwrite_replaces_all_fragments() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(
         &uri,
@@ -719,7 +725,7 @@ async fn stage_overwrite_replaces_all_fragments() {
 async fn stage_overwrite_empty_batch_replaces_all_rows() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(
         &uri,
@@ -774,7 +780,7 @@ async fn stage_overwrite_empty_batch_replaces_all_rows() {
 async fn stage_create_btree_index_does_not_advance_head_until_commit() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(
         &uri,
@@ -821,7 +827,7 @@ async fn stage_create_btree_index_does_not_advance_head_until_commit() {
 async fn stage_create_inverted_index_does_not_advance_head_until_commit() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(
         &uri,
@@ -862,7 +868,7 @@ async fn stage_create_inverted_index_does_not_advance_head_until_commit() {
 async fn stage_delete_does_not_advance_head_and_reads_through_staged() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(
         &uri,
@@ -917,7 +923,7 @@ async fn stage_delete_does_not_advance_head_and_reads_through_staged() {
 async fn stage_delete_commit_rebases_over_disjoint_committed_delete() {
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let ds = TableStore::write_dataset(&uri, numbered_person_batch(0..100))
         .await
@@ -955,7 +961,7 @@ async fn create_vector_index_advances_head_inline_documents_residual() {
 
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/vec.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     // Build a small dataset with a fixed-size vector column. Vector index
     // training requires multiple rows; provide enough.
@@ -1196,7 +1202,7 @@ async fn commit_staged_skips_auto_cleanup_so_pinned_versions_survive() {
 
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let mut ds = TableStore::write_dataset(&uri, person_batch(&[("seed", Some(0))]))
         .await

crates/omnigraph/tests/writes.rs

@@ -1411,7 +1411,7 @@ async fn scan_with_pending_rejects_key_column_missing_from_projection() {
 
     let dir = tempfile::tempdir().unwrap();
     let uri = format!("{}/people.lance", dir.path().to_str().unwrap());
-    let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default()));
+    let store = TableStore::new(dir.path().to_str().unwrap(), test_session());
 
     let schema = Arc::new(Schema::new(vec![
         Field::new("id", DataType::Utf8, false),