From 451585ee598e4cbbdc4a698baa921152725c2d3c Mon Sep 17 00:00:00 2001 From: aaltshuler Date: Sat, 4 Jul 2026 19:10:16 +0300 Subject: [PATCH] test(engine): shared test_session() helper for TableStore construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 5 +-- crates/omnigraph/tests/helpers/mod.rs | 7 ++++ crates/omnigraph/tests/recovery.rs | 5 +-- crates/omnigraph/tests/staged_writes.rs | 44 ++++++++++++++----------- crates/omnigraph/tests/writes.rs | 2 +- 5 files changed, 39 insertions(+), 24 deletions(-) diff --git a/crates/omnigraph/tests/failpoints.rs b/crates/omnigraph/tests/failpoints.rs index 9155ab46..91ec9085 100644 --- a/crates/omnigraph/tests/failpoints.rs +++ b/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 diff --git a/crates/omnigraph/tests/helpers/mod.rs b/crates/omnigraph/tests/helpers/mod.rs index 597a104a..ada91266 100644 --- a/crates/omnigraph/tests/helpers/mod.rs +++ b/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 { + 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(); diff --git a/crates/omnigraph/tests/recovery.rs b/crates/omnigraph/tests/recovery.rs index 7206b7c8..e1f47b2b 100644 --- a/crates/omnigraph/tests/recovery.rs +++ b/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 diff --git a/crates/omnigraph/tests/staged_writes.rs b/crates/omnigraph/tests/staged_writes.rs index e50edfc1..579d3655 100644 --- a/crates/omnigraph/tests/staged_writes.rs +++ b/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 { + std::sync::Arc::new(lance::session::Session::default()) +} use std::sync::Arc; fn person_schema() -> Arc { @@ -123,7 +129,7 @@ fn collect_age_for_id(batches: &[RecordBatch], needle: &str) -> Option { 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 { 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 diff --git a/crates/omnigraph/tests/writes.rs b/crates/omnigraph/tests/writes.rs index 8e6f4485..aac44bf7 100644 --- a/crates/omnigraph/tests/writes.rs +++ b/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),