apunkt/omnigraph
1
0
Fork 0
mirror of https://github.com/ModernRelay/omnigraph.git synced 2026-07-12 03:12:11 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

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

Browse source 
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.
This commit is contained in:
aaltshuler 2026-07-04 19:10:16 +03:00 • committed by Andrew Altshuler
parent d502603a52
commit 451585ee59
5 changed files with 39 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

5
crates/omnigraph/tests/failpoints.rs
View file

@ -15,7 +15,8 @@ use helpers::recovery::{
branch_head_commit_id, single_sidecar_operation_id, branch_head_commit_id, single_sidecar_operation_id,
}; };
use helpers::{ 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"; 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; // publisher deliberately skips the normal index-rebuild preparation;
// the failed writer below is still the real `ensure_indices_on`. // the failed writer below is still the real `ensure_indices_on`.
let person_uri = node_table_uri(&uri, "Person"); 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 let mut ds = store
.open_dataset_head(&person_uri, Some("feature")) .open_dataset_head(&person_uri, Some("feature"))
.await .await

7
crates/omnigraph/tests/helpers/mod.rs
View file

@ -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. /// Init a graph and load the standard test data.
pub async fn init_and_load(dir: &tempfile::TempDir) -> Omnigraph { pub async fn init_and_load(dir: &tempfile::TempDir) -> Omnigraph {
let uri = dir.path().to_str().unwrap(); let uri = dir.path().to_str().unwrap();

5
crates/omnigraph/tests/recovery.rs
View file

@ -18,6 +18,7 @@ use lance::Dataset;
use omnigraph::db::Omnigraph; use omnigraph::db::Omnigraph;
mod helpers; mod helpers;
use helpers::test_session;
use helpers::recovery::{RecoveryExpectation, TableExpectation, assert_post_recovery_invariants}; use helpers::recovery::{RecoveryExpectation, TableExpectation, assert_post_recovery_invariants};
const TEST_SCHEMA: &str = include_str!("fixtures/test.pg"); 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 // Bypass the manifest: append directly to Person's Lance HEAD on the
// feature branch ref to advance HEAD past v_pin. // feature branch ref to advance HEAD past v_pin.
let person_uri = node_table_uri(uri, "Person"); 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 let mut ds = store
.open_dataset_head(&person_uri, feature_branch_name.as_deref()) .open_dataset_head(&person_uri, feature_branch_name.as_deref())
.await .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 // Bypass the manifest: append on the feature ref to advance HEAD past
// the manifest pin. // the manifest pin.
let person_uri = node_table_uri(uri, "Person"); 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 let mut ds = store
.open_dataset_head(&person_uri, feature_branch_name.as_deref()) .open_dataset_head(&person_uri, feature_branch_name.as_deref())
.await .await

44
crates/omnigraph/tests/staged_writes.rs
View file

@ -28,6 +28,12 @@ use lance_index::IndexType;
use lance_linalg::distance::MetricType; use lance_linalg::distance::MetricType;
use lance_table::format::Fragment; use lance_table::format::Fragment;
use omnigraph::table_store::{StagedWrite, TableStore}; 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; use std::sync::Arc;
fn person_schema() -> Arc<Schema> { 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() { async fn stage_append_is_visible_via_scan_with_staged() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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. // Seed: one committed row.
let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) 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() { async fn stage_merge_insert_dedupes_superseded_committed_fragment() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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. // Seed: alice age 30 in one committed fragment.
let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) 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() { async fn count_rows_with_staged_matches_scan() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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))])) let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
.await .await
@ -250,7 +256,7 @@ async fn count_rows_with_staged_matches_scan() {
async fn chained_stage_appends_have_distinct_row_ids() { async fn chained_stage_appends_have_distinct_row_ids() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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))])) let ds = TableStore::write_dataset(&uri, person_batch(&[("seed", Some(0))]))
.await .await
@ -336,7 +342,7 @@ fn combine_for_scan(ds: &Dataset, staged: &[StagedWrite]) -> Vec<Fragment> {
async fn stage_append_then_commit_persists_data() { async fn stage_append_then_commit_persists_data() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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))])) let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
.await .await
@ -369,7 +375,7 @@ async fn stage_append_then_commit_persists_data() {
async fn stage_merge_insert_then_commit_persists_merged_view() { async fn stage_merge_insert_then_commit_persists_merged_view() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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))])) let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
.await .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() { async fn stage_merge_insert_commit_rebases_over_disjoint_committed_delete() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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)) let ds = TableStore::write_dataset(&uri, numbered_person_batch(0..100))
.await .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() { async fn scan_with_staged_with_filter_silently_drops_staged_rows() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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 // Committed: alice=30, carol=40
let ds = TableStore::write_dataset( 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() { async fn chained_stage_merge_insert_with_shared_key_documents_duplicate_behavior() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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). // Seed empty (an unrelated row keeps the schema unambiguous).
let ds = TableStore::write_dataset(&uri, person_batch(&[("seed", Some(0))])) 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() { async fn stage_overwrite_does_not_advance_head_until_commit() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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))])) let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))]))
.await .await
@ -638,7 +644,7 @@ async fn stage_overwrite_does_not_advance_head_until_commit() {
async fn stage_overwrite_preserves_stable_row_ids() { async fn stage_overwrite_preserves_stable_row_ids() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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 // `write_dataset` creates with `enable_stable_row_ids: true` — see
// ADR 0001. We verify that as a precondition so a future change to // 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() { async fn stage_overwrite_replaces_all_fragments() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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( let ds = TableStore::write_dataset(
&uri, &uri,
@ -719,7 +725,7 @@ async fn stage_overwrite_replaces_all_fragments() {
async fn stage_overwrite_empty_batch_replaces_all_rows() { async fn stage_overwrite_empty_batch_replaces_all_rows() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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( let ds = TableStore::write_dataset(
&uri, &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() { async fn stage_create_btree_index_does_not_advance_head_until_commit() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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( let ds = TableStore::write_dataset(
&uri, &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() { async fn stage_create_inverted_index_does_not_advance_head_until_commit() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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( let ds = TableStore::write_dataset(
&uri, &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() { async fn stage_delete_does_not_advance_head_and_reads_through_staged() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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( let ds = TableStore::write_dataset(
&uri, &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() { async fn stage_delete_commit_rebases_over_disjoint_committed_delete() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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)) let ds = TableStore::write_dataset(&uri, numbered_person_batch(0..100))
.await .await
@ -955,7 +961,7 @@ async fn create_vector_index_advances_head_inline_documents_residual() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/vec.lance", dir.path().to_str().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 // Build a small dataset with a fixed-size vector column. Vector index
// training requires multiple rows; provide enough. // 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 dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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))])) let mut ds = TableStore::write_dataset(&uri, person_batch(&[("seed", Some(0))]))
.await .await

2
crates/omnigraph/tests/writes.rs
View file

@ -1411,7 +1411,7 @@ async fn scan_with_pending_rejects_key_column_missing_from_projection() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let uri = format!("{}/people.lance", dir.path().to_str().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![ let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Utf8, false), Field::new("id", DataType::Utf8, false),