mirror of
https://github.com/ModernRelay/omnigraph.git
synced 2026-07-06 02:52:11 +02:00
* feat(engine): unify constraint validation across all write surfaces
Constraint enforcement (value/range/check, enum, uniqueness, edge
referential integrity, cardinality) was implemented three times — once
each in the bulk loader, the mutation executor, and the branch-merge
path — and had drifted: merge validated @range/@check but not enum, and
neither the mutation nor the load path enforced cross-version uniqueness
against already-committed rows.
Introduce one catalog-derived evaluator (`crate::validate`) that all
three surfaces route through. It is delta-scoped (checks only the change
set, not the whole graph) and index-backed (probes committed state
through the @key/@unique/src/dst BTREEs instead of full-scanning every
catalog table), reusing the existing leaf checks
(validate_value_constraints, validate_enum_constraints,
composite_unique_key) so the surfaces cannot drift again. A one-row-delta
merge now opens ~3 data tables instead of ~6+, and validation cost is
flat in graph size rather than O(V+E).
Behavior changes (all stricter, none relaxed):
- Enum constraints are now enforced on the merge path (was a gap).
- A write or load whose @unique value collides with an already-committed
different row is now rejected (cross-version uniqueness); re-upserting
an existing @key still upserts.
- Uniqueness distinguishes a duplicate key WITHIN one input batch (two
distinct records -> rejected, e.g. a bulk load listing a @key twice)
from the SAME id reappearing ACROSS batches (ordered supersession of
one logical row -> coalesced, e.g. a mutation insert-then-update).
- Overwrite loads validate per-table: a touched table's committed view is
its replacement image (empty), but a table absent from the batch keeps
its committed rows, so an edges-only overwrite still resolves
referential integrity against retained nodes.
Remove the per-surface validation orchestration the evaluator supersedes,
and the now-orphaned version-pinned dataset opener from the sealed
storage trait (reads route through the snapshot path). Docs (invariants,
testing) updated; full engine suite green.
* test(engine): pin orphan-edge validation on adopt-by-pointer merge
Regression for a gap in the unified merge validation: when a table is
adopted by pointer switch (AdoptSourceState) — source on main, target on a
branch — build_merge_changeset skips it, so referential integrity is never
checked for it. Merging main into a branch that deleted a node while main
added an edge to that node silently publishes the orphan edge.
This test merges main -> feature where feature deleted Bob and main added
Knows Alice->Bob, and asserts an OrphanEdge conflict. Red against HEAD
(merge returns Merged); turns green with the AdoptSourceState validation fix.
* fix(engine): validate adopt-by-pointer merge tables (AdoptSourceState)
The unified merge validator skipped any table classified AdoptSourceState
(a pointer switch / fork), so referential integrity, uniqueness, and
cardinality were never checked for it. Merging main into a branch that
deleted a node while main added an edge to that node silently published the
orphan edge — the prior full-scan validation caught it.
Root cause: classify_adopt keyed AdoptSourceState on the publish mechanism
("does it advance Lance HEAD") and returned before computing any delta, and
build_merge_changeset then skipped the table. Fix decouples the validation
input from the publish mechanism: classify_adopt now always computes the
source-vs-target delta (base == target on this path, so it is the right
validation delta) and carries it as AdoptSourceState { validation_delta };
build_merge_changeset validates it exactly like AdoptWithDelta. The publish
stays a pointer/fork (delta ignored) and remains excluded from recovery
pins, so publish/recovery semantics are unchanged — only validation is
restored. Closes the class: no publish optimization can bypass validation.
Turns the orphan-edge regression test green.
* test(engine): pin typed committed-uniqueness probe on non-String columns
The cross-version @unique check pushes a committed-state filter built from
the stringified key. On a non-String @unique column (e.g. Date) this compares
a Date32 column to a Utf8 literal — and the stringified key is the raw day
count, so the probe raises "Cannot cast string '20633' to Date32" for ANY
second write to the table (colliding or not).
Two regressions: a colliding Date value must surface a proper "@unique
violation" (not a coercion error), and a non-colliding write must succeed.
Both red against HEAD; green with the typed-literal probe fix.
* fix(engine): build committed uniqueness probe from typed column values
The cross-version @unique check pushed a Lance filter built with a
stringified key (lit(String)) against the real, typed column. On a
non-String @unique column this compared a Date32/numeric/bool column to a
Utf8 literal: a coercion error on Date/Bool (failing every write to the
table) or a silent miss on Float. For Date the stringified key was even the
raw day count, so the literal could never parse.
unique_holders now takes typed ScalarValues, built at the call site via
ScalarValue::try_from_array(group_column, row), so the pushed-down predicate
compares like-typed for any scalar @unique. The in-memory intra-delta dedup
keeps the stringified key (a type-agnostic equality grouping, unaffected).
Turns the Date @unique cross-version regression tests green.
* test(engine): pin id-keyed cardinality on merge-load edge moves/dups
Two cardinality drifts between validation and what commit persists:
- Move (B): a Merge-load that moves an edge to a new src only recounts the
new src, so vacating a src and dropping it below @card min is missed —
moving Alice's only WorksAt to Bob silently succeeds under @card(1..).
- Dup (A): a Merge-load batch listing one edge id under two srcs counts it
under both, but commit dedupes by id (last-wins). Alice gets a phantom
second edge and a spurious "has 2 edges (max 1)" violation under @card(0..1).
Both red against HEAD; green with the id-keyed last-wins cardinality model.
* fix(engine): key merge/load cardinality by edge id, last-wins
@card validation diverged from what commit persists in two ways: (1) it only
recounted the new src of a delta edge, so a Merge-load that moves an edge to a
new src never rechecked the vacated src and missed a drop below @card min; (2)
it counted raw delta rows, so the same edge id under two srcs in one batch was
counted under both, while commit dedupes by id (last-wins) — a phantom edge
and a spurious max violation.
evaluate_cardinality now coalesces the delta by edge id (last-wins, matching
dedupe_merge_batches_by_id) and builds the affected-src set from both the new
src of each delta edge AND the old committed src of each changed/deleted edge
id; a committed edge is dropped from its src when the delta deletes or
re-places it. The validated edge set per src now equals the committed image.
Turns the edge-move and duplicate-id cardinality regression tests green.
* docs(rfcs): add RFC 0001 — branch merge by fragment adoption
Proposed design for the by-design fix to merge cost/OOM: adopt the source
branch's Lance fragments by reference (base_paths) instead of re-materializing
rows, with a re-home reconciler + branch-delete reference guard closing the
dangling-reference lifecycle, and a reachability-complete cleanup sweep. Grounded
in the public Lance 7.0.0 multi-base APIs and the prior art (Delta shallow/deep
clone, Iceberg/lakeFS reachability GC). Status: Proposed.
* test(engine): pin @card validation on direct edge delete
Deletes stage as predicates, not constructive batches, so a delete-only
mutation produces an empty change-set and validate_changeset no-ops — a
`delete WorksAt where from = X` that removes a source's only edge commits
below @card(1..), while the merge path (which carries deleted_ids) rejects it.
Red against HEAD (the delete commits); green once the delete path resolves
its predicates into the validation change-set.
* fix(engine): validate edge cardinality on delete via resolved predicates
A delete-only mutation produced an empty change-set (deletes stage as
predicates, not constructive batches), so validate_changeset no-op'd and a
`delete Edge` that dropped a source below @card min committed silently — while
the merge path, which carries deleted_ids, rejects it.
validate_staged_mutation now resolves each staged delete predicate against the
live committed table (CommittedState::deleted_ids_matching, a SQL-filter scan
projecting id) and folds the matched ids into the change-set's deleted_ids for
that table. The existing evaluator then recounts the srcs a delete empties
(@card min) and sees removed rows for RI/node-delete — the same faithful
change-set the merge path already builds, so validation matches what commits.
Covers direct edge deletes, node deletes, and node-delete edge cascades
uniformly (all are staged predicates).
Turns the direct-edge-delete @card regression test green.
* refactor(engine): capture deleted ids at delete time, drop validation re-scan
The delete-cardinality fix resolved staged delete predicates a second time at
validation. Instead, capture the removed ids during the delete op's own scan:
execute_delete_edge and the node-delete edge cascade now scan id (not
count_rows), record the ids via MutationStaging::record_deleted_ids, and
to_changeset() folds them into the change-set's deleted_ids. validate_staged_
mutation reverts to plain to_changeset(); CommittedState::deleted_ids_matching
and scan_filtered_sql are removed.
Behavior-preserving (the @card-on-delete test stays green) and strictly fewer
scans — one scan at delete time replaces count-here + resolve-at-validation.
Node deletes already scanned their ids; this reuses that via a shared
ids_from_batches helper. Full engine suite green; workspace builds clean.
* test(engine): pin overwrite-removal RI + coalesced-unique final image
Two reviewer findings, both red against HEAD:
- F1 (High): overwriting a node table removes nodes without expressing them as
deleted_ids, so a retained edge in a non-overwritten table that references a
removed node is published as an orphan (edge-RI path-b never runs).
overwrite_node_removal_rejects_retained_orphan_edge.
- F2 (Medium): evaluate_unique accumulates superseded keys across batches, so a
mutation that frees a @unique value (Alice.email temp -> final) and reuses it
(insert Carol.email = temp) false-rejects a valid final image.
chained_unique_update_then_reuse_freed_value_is_not_a_violation.
* fix(engine): validate overwrite removals (orphan edges, emptied srcs)
An Overwrite load replaces each touched table, but to_changeset() only recorded
the new batch, never the committed rows the overwrite removes. So overwriting
node:Person to drop Bob while a retained edge:Knows(Alice->Bob) referenced him
published an orphan edge unchecked — edge-RI path-b is gated on the node's
deleted_ids, which were empty.
The loader now computes per overwritten table the removed ids (committed ids in
the pinned base minus the replacement batch's ids, via validate::
overwrite_removed_ids) and folds them into the change-set's deleted_ids. The
evaluator then runs RI path-b and cardinality against them — the same faithful
change-set the merge path builds. Overwrite is per-table, so a table absent from
the batch is untouched; a removed node referenced by a retained edge is now a
loud OrphanEdge.
Updates two tests that asserted the old silent-orphan behavior to
self-consistent overwrites (per-table Overwrite can't drop edge endpoints
without also overwriting the edge tables): end_to_end::overwrite_replaces_data
and writes::load_overwrite_with_bad_edge_reference_unblocks_next_load. The
orphan-rejection case itself is pinned by the new validators test.
* fix(engine): evaluate @unique against the coalesced final delta image
evaluate_unique iterated the raw delta batches and accumulated every key it saw
into one cross-batch map, so a coalesced write that frees then reuses a @unique
value within a query — update a row's email to 'temp', update the same row to
'final', insert a new row with 'temp' — false-rejected: 'temp' lingered in the
seen-set from the superseded first write though it no longer holds in the final
image that commits.
Restructure to validate the final coalesced image — the bytes that actually
publish:
- Pass 1 coalesces the delta by id (last-wins) into each id's final key, and
flags genuine within-ONE-batch duplicates (two distinct input records — the
bulk-load contract) before coalescing, so an unordered load batch with a real
dup still rejects.
- Pass 2 checks two distinct final ids holding the same key.
- Pass 3 does the committed cross-version lookup, excluding the delta's own ids.
Entries are sorted by id before the cross-row/committed passes so violation
order never depends on HashMap iteration. Coalescing first also drops the
redundant committed probes a superseded key used to issue.
Pinned by the chained-update red test; preserves intra-batch dup rejection
(consistency::loader_rejects_intra_batch_duplicate_keys) and cross-version
uniqueness (validators).
* style(engine): drop trailing blank line at staging.rs EOF
Left by a block-delete in an earlier refactor; flagged by git diff --check.
* docs(engine): refresh validate.rs module doc to current consumers
The module doc still said the merge path was the only consumer and the write
path a later, mechanical migration, and listed cardinality as a later
increment. Mutation and bulk load have since migrated onto the evaluator and
cardinality ships — correct both so the doc reflects that all three write
surfaces route through one evaluator.
2832 lines
117 KiB
Rust
2832 lines
117 KiB
Rust
use std::collections::{BTreeSet, HashMap, HashSet};
|
||
use std::io::Write;
|
||
use std::sync::Arc;
|
||
|
||
use arc_swap::ArcSwap;
|
||
use arrow_array::{
|
||
Array, BinaryArray, BooleanArray, Date32Array, FixedSizeListArray, Float32Array, Float64Array,
|
||
Int32Array, Int64Array, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray,
|
||
RecordBatch, StringArray, StructArray, UInt32Array, UInt64Array, new_null_array,
|
||
};
|
||
use arrow_schema::{DataType, Field, Schema};
|
||
use lance::Dataset;
|
||
use lance::blob::{BlobArrayBuilder, blob_field};
|
||
use lance::dataset::BlobFile;
|
||
use lance::dataset::scanner::ColumnOrdering;
|
||
use lance::datatypes::BlobKind;
|
||
use omnigraph_compiler::catalog::{Catalog, EdgeType, NodeType};
|
||
use omnigraph_compiler::schema::parser::parse_schema;
|
||
use omnigraph_compiler::types::{PropType, ScalarType};
|
||
use omnigraph_compiler::{
|
||
DropMode, SchemaIR, SchemaMigrationPlan, SchemaMigrationStep, SchemaTypeKind,
|
||
build_catalog_from_ir, build_schema_ir, plan_schema_migration,
|
||
};
|
||
|
||
use crate::db::graph_coordinator::{GraphCoordinator, PublishedSnapshot};
|
||
use crate::error::{OmniError, Result};
|
||
use crate::runtime_cache::RuntimeCache;
|
||
use crate::storage::{StorageAdapter, join_uri, normalize_root_uri, storage_for_uri};
|
||
use crate::storage_layer::SnapshotHandle;
|
||
use crate::table_store::TableStore;
|
||
|
||
mod export;
|
||
mod optimize;
|
||
mod repair;
|
||
mod schema_apply;
|
||
mod table_ops;
|
||
|
||
pub use optimize::{CleanupPolicyOptions, SkipReason, TableCleanupStats, TableOptimizeStats};
|
||
pub use repair::{
|
||
RepairAction, RepairClassification, RepairOptions, RepairStats, TableRepairStats,
|
||
};
|
||
pub use schema_apply::SchemaApplyOptions;
|
||
pub use table_ops::PendingIndex;
|
||
pub(crate) use table_ops::OpenedForMutation;
|
||
|
||
use super::commit_graph::GraphCommit;
|
||
use super::manifest::{
|
||
ManifestChange, Snapshot, SubTableEntry, TableRegistration, TableTombstone,
|
||
table_path_for_table_key,
|
||
};
|
||
use super::schema_state::{
|
||
SCHEMA_SOURCE_FILENAME, load_or_bootstrap_schema_contract, read_accepted_schema_ir,
|
||
recover_schema_state_files, schema_ir_staging_uri, schema_ir_uri, schema_source_staging_uri,
|
||
schema_source_uri, schema_state_staging_uri, schema_state_uri, validate_schema_contract,
|
||
write_schema_contract, write_schema_contract_staging,
|
||
};
|
||
use super::{
|
||
ReadTarget, ResolvedTarget, SCHEMA_APPLY_LOCK_BRANCH, SnapshotId, is_internal_system_branch,
|
||
is_schema_apply_lock_branch,
|
||
};
|
||
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum MergeOutcome {
|
||
AlreadyUpToDate,
|
||
FastForward,
|
||
Merged,
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct SchemaApplyResult {
|
||
pub supported: bool,
|
||
pub applied: bool,
|
||
pub manifest_version: u64,
|
||
pub steps: Vec<SchemaMigrationStep>,
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct SchemaApplyPreview {
|
||
pub plan: SchemaMigrationPlan,
|
||
pub catalog: Catalog,
|
||
}
|
||
|
||
/// A capture-once write transaction (RFC-013 step 3b). Pins the operation's read
|
||
/// base ONCE so the per-table opens reuse the pinned version instead of
|
||
/// re-resolving / re-validating per table. The schema contract is validated once
|
||
/// (when `base` is captured). NOT a general "no re-resolution" handle — the
|
||
/// commit-time OCC re-read, the live-HEAD drift probe, and the fork-authority reads
|
||
/// stay fresh (correctness machinery). Step 5 (PublishPlan unification) makes this
|
||
/// the non-optional publish carrier and adds session-aware base opens there, gated
|
||
/// by an S3 cost test — the warm-session benefit on the single remaining open is an
|
||
/// object-store phenomenon, so it earns its own gate rather than riding this PR.
|
||
///
|
||
/// Threaded as `Option<&WriteTxn>` through the mutate/load write chain
|
||
/// (`open_for_mutation_on_branch`, `commit_all`, `commit_updates_on_branch_with_expected`)
|
||
/// so a single write validates the schema contract EXACTLY ONCE — at capture. When
|
||
/// present, the per-table resolves source the pinned `base` entry instead of calling
|
||
/// `resolved_branch_target` / `snapshot_for_branch` / `fresh_snapshot_for_branch`
|
||
/// (each of which re-runs `ensure_schema_state_valid`). When absent (`None` — every
|
||
/// non-mutate/load caller), every threaded function behaves byte-identically to
|
||
/// before. The carrier never removes a version guard or changes which dataset version
|
||
/// the per-table open targets: strict ops keep `open_dataset_head_for_write` +
|
||
/// `ensure_expected_version`, and the commit-time OCC re-read still opens a fresh
|
||
/// manifest snapshot (via `fresh_snapshot_for_branch_unchecked`) — only the redundant
|
||
/// schema re-validation is dropped.
|
||
pub(crate) struct WriteTxn {
|
||
/// The resolved branch (`None` = main).
|
||
pub(crate) branch: Option<String>,
|
||
/// The pinned base snapshot (per-table location + version + e_tag), captured once.
|
||
pub(crate) base: Snapshot,
|
||
}
|
||
|
||
/// Top-level handle to an Omnigraph database.
|
||
///
|
||
/// An Omnigraph is a Lance-native graph database with git-style branching.
|
||
/// It stores typed property graphs as per-type Lance datasets coordinated
|
||
/// through a Lance manifest table.
|
||
pub struct Omnigraph {
|
||
root_uri: String,
|
||
storage: Arc<dyn StorageAdapter>,
|
||
/// Coordinator state behind a tokio `RwLock`. PR 2 (MR-686) wraps
|
||
/// this so engine write APIs can be `&self` (the HTTP server's
|
||
/// `AppState` holds `Arc<Omnigraph>` and dispatches concurrent
|
||
/// calls without a global write lock). Reads (`snapshot`, `version`,
|
||
/// `current_branch`, `branch_list`, `resolve_*`, `head_commit_id`,
|
||
/// `list_commits`, …) acquire `.read().await` and parallelize.
|
||
/// Writes (`refresh`, `branch_create`, `branch_delete`, `commit_*`)
|
||
/// acquire `.write().await` and serialize. The atomic commit invariant —
|
||
/// table-version rows and the graph commit are one unit — holds by
|
||
/// construction since RFC-013 Phase 7: both ride a SINGLE manifest publish
|
||
/// CAS (`commit_changes_with_lineage`), so there is no two-write window to
|
||
/// keep atomic. PR 2 Phase 2
|
||
/// converted from `Mutex` to `RwLock` because the bench showed
|
||
/// the Mutex was the dominant serializer for disjoint-table
|
||
/// workloads. Lock acquisition order: always before `runtime_cache`
|
||
/// (when both are needed in one scope).
|
||
coordinator: Arc<tokio::sync::RwLock<GraphCoordinator>>,
|
||
table_store: TableStore,
|
||
runtime_cache: RuntimeCache,
|
||
/// Per-graph read caches: one shared Lance `Session` plus the held-`Dataset`
|
||
/// handle cache, handed to live-Branch-read snapshots (via
|
||
/// `resolved_target`) so table opens reuse handles (0 IO on a warm repeat)
|
||
/// and one session. Invalidated alongside `runtime_cache` on branch switch /
|
||
/// refresh — hygiene only; version-in-key carries correctness.
|
||
read_caches: Arc<crate::runtime_cache::ReadCaches>,
|
||
/// Read-heavy on every query, written only by `apply_schema`. ArcSwap
|
||
/// gives atomic pointer swap with zero-cost reads (`load()` returns a
|
||
/// `Guard<Arc<Catalog>>`), so concurrent queries on different actors
|
||
/// don't contend on a lock to read the catalog.
|
||
catalog: Arc<ArcSwap<Catalog>>,
|
||
/// Read-heavy on schema introspection paths, written only by
|
||
/// `apply_schema`. Same ArcSwap rationale as `catalog`.
|
||
schema_source: Arc<ArcSwap<String>>,
|
||
/// Per-`(table_key, branch)` writer queues — the engine's
|
||
/// write-serialization mechanism (the server holds the engine as a
|
||
/// lockless `Arc<Omnigraph>`). Reachable from engine internals
|
||
/// (mutation finalize, schema_apply, branch_merge, ensure_indices,
|
||
/// the fork path, recovery reconciler).
|
||
write_queue: Arc<crate::db::write_queue::WriteQueueManager>,
|
||
/// Process-wide mutex held across the swap → operate → restore window
|
||
/// in `branch_merge_impl`. Two concurrent merges with distinct targets
|
||
/// would otherwise interleave their three separate
|
||
/// `coordinator.write().await` acquisitions, leaving each merge's
|
||
/// inner body running against the other's swapped coord. Pinned by
|
||
/// `concurrent_branch_merges_distinct_targets_do_not_swap_into_each_other`
|
||
/// in `crates/omnigraph-server/tests/server.rs`.
|
||
///
|
||
/// Cost: serializes ALL concurrent branch merges process-wide.
|
||
/// Acceptable because branch merges are heavy (table rewrites, index
|
||
/// rebuilds), per-(table, branch) queues inside `commit_all` already
|
||
/// serialize the data path, and merges are rare relative to /change
|
||
/// or /ingest. A finer-grained per-target-branch mutex is a follow-up
|
||
/// if telemetry shows merge concurrency matters.
|
||
///
|
||
/// The deeper fix — refactor `branch_merge_on_current_target` to take
|
||
/// an explicit target coord parameter so `self.coordinator` is never
|
||
/// used as scratch space — is the round-1 shape applied to
|
||
/// `branch_create_from_impl`. Deferred because it requires unwinding
|
||
/// every `self.snapshot()` call inside the merge body.
|
||
merge_exclusive: Arc<tokio::sync::Mutex<()>>,
|
||
/// Optional policy checker for engine-layer enforcement (MR-722).
|
||
/// `None` = no enforcement; mutating methods are unconditionally
|
||
/// allowed (this is the embedded/dev default). `Some` = every
|
||
/// mutating method calls `self.enforce(action, scope, actor)` at
|
||
/// entry; denial returns `OmniError::Policy`.
|
||
///
|
||
/// Per chassis design (see `omnigraph_policy::PolicyChecker`), the
|
||
/// trait surface is deliberately coarse — action × scope × actor.
|
||
/// Per-row / per-type / per-column scope lives at the query layer
|
||
/// (MR-725), which extends the same trait with a different method.
|
||
/// Don't be tempted to add per-row enforcement here.
|
||
///
|
||
/// Set via `with_policy(checker)` after construction. Today only
|
||
/// `apply_schema_as` consults this field (PR #2 proof-of-concept);
|
||
/// PR #3 fans the `enforce()` call out to the remaining writers.
|
||
policy: Option<Arc<dyn omnigraph_policy::PolicyChecker>>,
|
||
/// Lazily-built, reused-across-queries embedding client. Built on the first
|
||
/// `nearest($v, "string")` that needs server-side embedding (so a graph that
|
||
/// never embeds needs no provider key), then shared by every later query —
|
||
/// avoids the per-query `from_env()` rebuild and keeps the provider HTTP
|
||
/// connection pool warm. `OnceCell` guarantees a single initialization.
|
||
embedding: Arc<tokio::sync::OnceCell<crate::embedding::EmbeddingClient>>,
|
||
/// Optional pre-resolved embedding config (RFC-012 Phase 5), injected from an
|
||
/// applied cluster `providers.embedding` profile via [`Omnigraph::with_embedding_config`].
|
||
/// When set, the embedding cell builds its client from this instead of
|
||
/// `EmbeddingClient::from_env()`; `None` keeps the env fallback.
|
||
embedding_config: Option<Arc<crate::embedding::EmbeddingConfig>>,
|
||
}
|
||
|
||
/// Whether [`Omnigraph::open`] runs the open-time recovery sweep.
|
||
///
|
||
/// Recovery requires Lance writes (`Dataset::restore`, `ManifestBatchPublisher::publish`).
|
||
/// Read-only consumers — NDJSON export, `commit list`, `read`, schema
|
||
/// inspection — should not trigger writes (they may run with read-only
|
||
/// object-store credentials, and silent open-time mutations are
|
||
/// surprising). They also don't need recovery: reads always resolve
|
||
/// through the manifest pin, which is the consistent snapshot regardless
|
||
/// of any Phase B → Phase C drift on the per-table side.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum OpenMode {
|
||
/// Run the recovery sweep on open. Default for `Omnigraph::open`.
|
||
ReadWrite,
|
||
/// Skip the recovery sweep. Use for read-only consumers via
|
||
/// [`Omnigraph::open_read_only`].
|
||
ReadOnly,
|
||
}
|
||
|
||
/// Options for [`Omnigraph::init_with_options`].
|
||
///
|
||
/// `force` controls the safety preflight that prevents an
|
||
/// accidental re-init from overwriting an existing graph's schema
|
||
/// metadata. Default behavior (`force: false`) fails fast with
|
||
/// [`OmniError::AlreadyInitialized`] if any of `_schema.pg`,
|
||
/// `_schema.ir.json`, or `__schema_state.json` already exists at
|
||
/// the target URI. With `force: true` the preflight is skipped —
|
||
/// existing schema files are overwritten in place. Force does NOT
|
||
/// purge old Lance datasets or `__manifest/`; reclaiming those
|
||
/// still requires deleting the graph directory by hand (or via a
|
||
/// future `DELETE /graphs/{id}`).
|
||
#[derive(Debug, Clone, Copy, Default)]
|
||
pub struct InitOptions {
|
||
/// Skip the existing-graph preflight. Operators set this when
|
||
/// they actually mean to overwrite — e.g. `omnigraph init --force`.
|
||
pub force: bool,
|
||
}
|
||
|
||
impl Omnigraph {
|
||
/// Create a new graph at `uri` from schema source.
|
||
///
|
||
/// Strict mode: errors with [`OmniError::AlreadyInitialized`] if
|
||
/// `uri` already holds any of the three schema artifacts. To
|
||
/// overwrite an existing graph deliberately, call
|
||
/// [`Self::init_with_options`] with `InitOptions { force: true }`.
|
||
pub async fn init(uri: &str, schema_source: &str) -> Result<Self> {
|
||
Self::init_with_options(uri, schema_source, InitOptions::default()).await
|
||
}
|
||
|
||
/// Create a new graph at `uri`, with explicit init-time options.
|
||
///
|
||
/// See [`InitOptions`] for the safety contract — by default this
|
||
/// behaves identically to [`Self::init`].
|
||
pub async fn init_with_options(
|
||
uri: &str,
|
||
schema_source: &str,
|
||
options: InitOptions,
|
||
) -> Result<Self> {
|
||
Self::init_with_storage(uri, schema_source, storage_for_uri(uri)?, options).await
|
||
}
|
||
|
||
pub(crate) async fn init_with_storage(
|
||
uri: &str,
|
||
schema_source: &str,
|
||
storage: Arc<dyn StorageAdapter>,
|
||
options: InitOptions,
|
||
) -> Result<Self> {
|
||
let root = normalize_root_uri(uri)?;
|
||
|
||
// Preflight: refuse to clobber an existing graph unless the
|
||
// operator passed `force`. This runs BEFORE any parse or
|
||
// write so a misdirected `init` against an existing graph
|
||
// URI cannot reach a code path that overwrites or, on a
|
||
// later cleanup, deletes the schema files.
|
||
//
|
||
// Closes the "init is destructive against existing state"
|
||
// class: there is no longer a code path where strict-mode
|
||
// `init` can mutate a populated graph root.
|
||
if !options.force {
|
||
for candidate in [
|
||
schema_source_uri(&root),
|
||
schema_ir_uri(&root),
|
||
schema_state_uri(&root),
|
||
] {
|
||
if storage.exists(&candidate).await? {
|
||
return Err(OmniError::AlreadyInitialized { uri: root.clone() });
|
||
}
|
||
}
|
||
}
|
||
|
||
let schema_ir = read_schema_ir_from_source(schema_source)?;
|
||
let mut catalog = build_catalog_from_ir(&schema_ir)?;
|
||
fixup_blob_schemas(&mut catalog);
|
||
|
||
// Establish an atomic ownership claim on `_schema.pg` before
|
||
// writing the remaining init artifacts. A check-then-write preflight
|
||
// is not enough under concurrent `init` calls: two callers can both
|
||
// observe an empty root, one can successfully initialize, and the
|
||
// loser can then fail in Lance `WriteMode::Create`. Only the caller
|
||
// that atomically created `_schema.pg` may clean up schema artifacts
|
||
// on later failure.
|
||
let schema_pg_claimed = if options.force {
|
||
false
|
||
} else {
|
||
let schema_path = join_uri(&root, SCHEMA_SOURCE_FILENAME);
|
||
if !storage
|
||
.write_text_if_absent(&schema_path, schema_source)
|
||
.await?
|
||
{
|
||
return Err(OmniError::AlreadyInitialized { uri: root.clone() });
|
||
}
|
||
if let Err(err) = crate::failpoints::maybe_fail(crate::failpoints::names::INIT_AFTER_SCHEMA_PG_WRITTEN) {
|
||
best_effort_cleanup_init_artifacts(&root, storage.as_ref()).await;
|
||
return Err(err);
|
||
}
|
||
true
|
||
};
|
||
|
||
// Run the I/O phase. On any error, best-effort-clean schema
|
||
// artifacts only when this invocation owns them: strict mode owns
|
||
// them after the atomic `_schema.pg` claim above; force mode owns
|
||
// destructive overwrite semantics by explicit operator request.
|
||
//
|
||
// Coverage gap: Lance per-type datasets and `__manifest/`
|
||
// directory created by `GraphCoordinator::init` are NOT cleaned
|
||
// up here — fully recursive directory deletion requires a
|
||
// `StorageAdapter::delete_prefix` primitive that's deferred
|
||
// along with `DELETE /graphs/{id}` (PR 2b in the MR-668 plan
|
||
// is currently deferred). If `init` fails after coordinator
|
||
// init succeeds, operators may need to remove the graph
|
||
// directory manually before retrying `init` on the same URI.
|
||
// Documented in the PR 2a commit message and `init` rustdoc.
|
||
let coordinator = match init_storage_phase(
|
||
&root,
|
||
schema_source,
|
||
&schema_ir,
|
||
&catalog,
|
||
&storage,
|
||
!schema_pg_claimed,
|
||
)
|
||
.await
|
||
{
|
||
Ok(coordinator) => coordinator,
|
||
Err(err) => {
|
||
if schema_pg_claimed || options.force {
|
||
best_effort_cleanup_init_artifacts(&root, storage.as_ref()).await;
|
||
}
|
||
return Err(err);
|
||
}
|
||
};
|
||
|
||
Ok(Self {
|
||
root_uri: root.clone(),
|
||
storage,
|
||
coordinator: Arc::new(tokio::sync::RwLock::new(coordinator)),
|
||
table_store: TableStore::new(&root),
|
||
runtime_cache: RuntimeCache::default(),
|
||
// One shared Session per graph (LanceDB's one-session-per-connection
|
||
// model) plus the held-handle cache, created once and reused across
|
||
// reads. Session::default() caps are lazy (6 GiB index / 1 GiB
|
||
// metadata); multi-graph cap/sharing is a deferred follow-up.
|
||
read_caches: Arc::new(crate::runtime_cache::ReadCaches {
|
||
session: Arc::new(lance::session::Session::default()),
|
||
handles: Arc::new(crate::runtime_cache::TableHandleCache::default()),
|
||
}),
|
||
catalog: Arc::new(ArcSwap::from_pointee(catalog)),
|
||
schema_source: Arc::new(ArcSwap::from_pointee(schema_source.to_string())),
|
||
write_queue: Arc::new(crate::db::write_queue::WriteQueueManager::new()),
|
||
merge_exclusive: Arc::new(tokio::sync::Mutex::new(())),
|
||
policy: None,
|
||
embedding: Arc::new(tokio::sync::OnceCell::new()),
|
||
embedding_config: None,
|
||
})
|
||
}
|
||
|
||
/// Open an existing graph (read-write).
|
||
///
|
||
/// Reads `_schema.pg`, parses it, builds the catalog, and opens `__manifest`.
|
||
/// Runs the open-time recovery sweep before returning — see [`OpenMode`].
|
||
pub async fn open(uri: &str) -> Result<Self> {
|
||
Self::open_with_storage_and_mode(uri, storage_for_uri(uri)?, OpenMode::ReadWrite).await
|
||
}
|
||
|
||
/// Open an existing graph for read-only consumers (NDJSON export,
|
||
/// `commit list`, etc.). Skips the recovery sweep — see [`OpenMode`].
|
||
pub async fn open_read_only(uri: &str) -> Result<Self> {
|
||
Self::open_with_storage_and_mode(uri, storage_for_uri(uri)?, OpenMode::ReadOnly).await
|
||
}
|
||
|
||
/// Open with a caller-supplied [`StorageAdapter`]. Used by init/test paths
|
||
/// and by embedding/test consumers that wrap storage (e.g. a counting
|
||
/// decorator for IO-budget tests). Defaults to `OpenMode::ReadWrite`.
|
||
pub async fn open_with_storage(uri: &str, storage: Arc<dyn StorageAdapter>) -> Result<Self> {
|
||
Self::open_with_storage_and_mode(uri, storage, OpenMode::ReadWrite).await
|
||
}
|
||
|
||
pub(crate) async fn open_with_storage_and_mode(
|
||
uri: &str,
|
||
storage: Arc<dyn StorageAdapter>,
|
||
mode: OpenMode,
|
||
) -> Result<Self> {
|
||
let root = normalize_root_uri(uri)?;
|
||
// Refuse a `__manifest` this binary cannot serve before the coordinator
|
||
// reads any branch state — newer than CURRENT (an old binary must not
|
||
// silently misread a newer graph) or below MIN_SUPPORTED (an older
|
||
// storage format this binary does not read — rebuild via export/import).
|
||
// Both open modes refuse: there is no in-place migration, and the check is
|
||
// a stamp read with no object-store writes, so it is safe under ReadOnly.
|
||
crate::db::manifest::refuse_if_internal_schema_unsupported(&root).await?;
|
||
// Open the coordinator first so the schema-staging recovery sweep can
|
||
// compare its snapshot against any leftover staging files.
|
||
let mut coordinator = GraphCoordinator::open(&root, Arc::clone(&storage)).await?;
|
||
// Both the schema-state recovery sweep AND the manifest-drift
|
||
// recovery sweep are gated on `OpenMode::ReadWrite`. Read-only
|
||
// consumers (NDJSON export, `commit list`, schema show) shouldn't
|
||
// trigger object-store mutations: they may run with read-only
|
||
// credentials, and silent open-time writes are surprising. Both
|
||
// sweeps' work is recoverable on the next ReadWrite open, so
|
||
// skipping under ReadOnly doesn't lose any safety guarantees —
|
||
// the manifest pin is the consistent snapshot regardless of
|
||
// drift on the per-table side or leftover schema-staging files.
|
||
if matches!(mode, OpenMode::ReadWrite) {
|
||
let schema_state_recovery =
|
||
recover_schema_state_files(&root, Arc::clone(&storage), &coordinator.snapshot())
|
||
.await?;
|
||
// Recovery sweep: close the Phase B → Phase C residual on
|
||
// any sidecar left over from a crashed writer. Long-running
|
||
// processes additionally converge in-process: the staged-
|
||
// write entry points and `refresh` run the roll-forward-only
|
||
// heal (`heal_pending_sidecars_roll_forward`); only
|
||
// rollback-eligible sidecars wait for this open-time sweep.
|
||
crate::db::manifest::recover_manifest_drift(
|
||
&root,
|
||
Arc::clone(&storage),
|
||
&mut coordinator,
|
||
crate::db::manifest::RecoveryMode::Full,
|
||
schema_state_recovery,
|
||
)
|
||
.await?;
|
||
}
|
||
// Read _schema.pg (post-recovery — may have just been renamed in).
|
||
let schema_path = schema_source_uri(&root);
|
||
let schema_source = storage.read_text(&schema_path).await?;
|
||
let current_source_ir = read_schema_ir_from_source(&schema_source)?;
|
||
let branches = coordinator.branch_list().await?;
|
||
let (accepted_ir, _) = load_or_bootstrap_schema_contract(
|
||
&root,
|
||
Arc::clone(&storage),
|
||
&branches,
|
||
¤t_source_ir,
|
||
)
|
||
.await?;
|
||
let mut catalog = build_catalog_from_ir(&accepted_ir)?;
|
||
fixup_blob_schemas(&mut catalog);
|
||
|
||
Ok(Self {
|
||
root_uri: root.clone(),
|
||
storage,
|
||
coordinator: Arc::new(tokio::sync::RwLock::new(coordinator)),
|
||
table_store: TableStore::new(&root),
|
||
runtime_cache: RuntimeCache::default(),
|
||
// One shared Session per graph (LanceDB's one-session-per-connection
|
||
// model) plus the held-handle cache, created once and reused across
|
||
// reads. Session::default() caps are lazy (6 GiB index / 1 GiB
|
||
// metadata); multi-graph cap/sharing is a deferred follow-up.
|
||
read_caches: Arc::new(crate::runtime_cache::ReadCaches {
|
||
session: Arc::new(lance::session::Session::default()),
|
||
handles: Arc::new(crate::runtime_cache::TableHandleCache::default()),
|
||
}),
|
||
catalog: Arc::new(ArcSwap::from_pointee(catalog)),
|
||
schema_source: Arc::new(ArcSwap::from_pointee(schema_source)),
|
||
write_queue: Arc::new(crate::db::write_queue::WriteQueueManager::new()),
|
||
merge_exclusive: Arc::new(tokio::sync::Mutex::new(())),
|
||
policy: None,
|
||
embedding: Arc::new(tokio::sync::OnceCell::new()),
|
||
embedding_config: None,
|
||
})
|
||
}
|
||
|
||
/// Returns an `Arc<Catalog>` snapshot. Cheap clone of the current
|
||
/// catalog pointer; callers can hold the returned `Arc` across awaits
|
||
/// without blocking concurrent `apply_schema`.
|
||
pub fn catalog(&self) -> Arc<Catalog> {
|
||
self.catalog.load_full()
|
||
}
|
||
|
||
/// Returns an `Arc<String>` snapshot of the schema source.
|
||
pub fn schema_source(&self) -> Arc<String> {
|
||
self.schema_source.load_full()
|
||
}
|
||
|
||
/// Atomically swap the in-memory catalog. Concurrent readers see
|
||
/// either the old or the new pointer; never a torn state. Used by
|
||
/// `apply_schema` and `reload_schema_if_source_changed`.
|
||
pub(crate) fn store_catalog(&self, catalog: Catalog) {
|
||
self.catalog.store(Arc::new(catalog));
|
||
}
|
||
|
||
/// Atomically swap the in-memory schema source. Same rationale as
|
||
/// [`store_catalog`](Self::store_catalog).
|
||
pub(crate) fn store_schema_source(&self, schema_source: String) {
|
||
self.schema_source.store(Arc::new(schema_source));
|
||
}
|
||
|
||
pub fn uri(&self) -> &str {
|
||
&self.root_uri
|
||
}
|
||
|
||
/// Install a policy checker for engine-layer enforcement (MR-722).
|
||
/// Builder-style setter — consumes `self`, returns `Self`. Calling
|
||
/// this on a `Omnigraph` previously without policy enables
|
||
/// `enforce()` to fire at every mutating engine method that's been
|
||
/// wired to call it (currently `apply_schema_as`; PR #3 fans out to
|
||
/// the remaining writers).
|
||
///
|
||
/// Embedded callers that don't care about authorization should
|
||
/// just not call this. Server / CLI callers that have loaded a
|
||
/// `PolicyEngine` from `policy.yaml` pass it here.
|
||
pub fn with_policy(mut self, checker: Arc<dyn omnigraph_policy::PolicyChecker>) -> Self {
|
||
self.policy = Some(checker);
|
||
self
|
||
}
|
||
|
||
/// The lazily-initialized, reused-across-queries embedding client cell
|
||
/// (see the `embedding` field doc). The query executor resolves the client
|
||
/// through this on the first `nearest($v, "string")` that needs embedding.
|
||
pub(crate) fn embedding_cell(
|
||
&self,
|
||
) -> &tokio::sync::OnceCell<crate::embedding::EmbeddingClient> {
|
||
&self.embedding
|
||
}
|
||
|
||
/// Install a pre-resolved embedding config (RFC-012 Phase 5). Builder-style,
|
||
/// mirroring [`Omnigraph::with_policy`]: a graph served from a cluster
|
||
/// embedding provider profile injects it here; an embedded/CLI caller that doesn't
|
||
/// call this keeps the `EmbeddingClient::from_env()` fallback.
|
||
pub fn with_embedding_config(mut self, config: Arc<crate::embedding::EmbeddingConfig>) -> Self {
|
||
self.embedding_config = Some(config);
|
||
self
|
||
}
|
||
|
||
/// The injected embedding config, if any (see the `embedding_config` field).
|
||
pub(crate) fn embedding_config_ref(&self) -> Option<&crate::embedding::EmbeddingConfig> {
|
||
self.embedding_config.as_deref()
|
||
}
|
||
|
||
/// Engine-layer policy enforcement gate (MR-722 chassis core).
|
||
///
|
||
/// * If no policy is installed → no-op (returns `Ok(())`).
|
||
/// * If policy is installed AND actor is None → denial with a
|
||
/// clear "no actor for engine-layer policy check" message.
|
||
/// Forces server / CLI / SDK callers to thread an actor through
|
||
/// when policy is configured — silent bypass via "I forgot the
|
||
/// actor" is exactly the footgun this gate is here to prevent.
|
||
/// * If policy is installed AND actor is Some → call
|
||
/// `PolicyChecker::check(action, scope, actor)`; map denial /
|
||
/// internal failure to `OmniError::Policy(...)`.
|
||
pub(crate) fn enforce(
|
||
&self,
|
||
action: omnigraph_policy::PolicyAction,
|
||
scope: &omnigraph_policy::ResourceScope,
|
||
actor: Option<&str>,
|
||
) -> Result<()> {
|
||
let Some(checker) = self.policy.as_ref() else {
|
||
return Ok(());
|
||
};
|
||
let Some(actor) = actor else {
|
||
return Err(OmniError::Policy(
|
||
"no actor for engine-layer policy check (policy is configured but the call site \
|
||
didn't thread an actor through — this is almost certainly a bug, not an \
|
||
intended bypass)"
|
||
.to_string(),
|
||
));
|
||
};
|
||
checker
|
||
.check(action, scope, actor)
|
||
.map_err(|err| OmniError::Policy(err.to_string()))
|
||
}
|
||
|
||
pub(crate) async fn ensure_schema_state_valid(&self) -> Result<()> {
|
||
// Full per-call validation is intentional: a long-lived handle must
|
||
// detect external drift of the schema source, IR, OR state on its next
|
||
// operation (see lifecycle::long_lived_handle_rejects_schema_* tests). A
|
||
// source-only fast path would miss IR/state drift when _schema.pg is
|
||
// unchanged, so the only safe latency win is not calling this twice per
|
||
// query (finding A removes the redundant caller in exec/query.rs).
|
||
validate_schema_contract(self.uri(), Arc::clone(&self.storage)).await
|
||
}
|
||
|
||
pub async fn plan_schema(&self, desired_schema_source: &str) -> Result<SchemaMigrationPlan> {
|
||
self.plan_schema_with_options(desired_schema_source, SchemaApplyOptions::default())
|
||
.await
|
||
}
|
||
|
||
pub async fn plan_schema_with_options(
|
||
&self,
|
||
desired_schema_source: &str,
|
||
options: SchemaApplyOptions,
|
||
) -> Result<SchemaMigrationPlan> {
|
||
schema_apply::plan_schema(self, desired_schema_source, options).await
|
||
}
|
||
|
||
pub async fn preview_schema_apply_with_options(
|
||
&self,
|
||
desired_schema_source: &str,
|
||
options: SchemaApplyOptions,
|
||
) -> Result<SchemaApplyPreview> {
|
||
schema_apply::preview_schema_apply(self, desired_schema_source, options).await
|
||
}
|
||
|
||
pub async fn apply_schema(&self, desired_schema_source: &str) -> Result<SchemaApplyResult> {
|
||
self.apply_schema_as(desired_schema_source, SchemaApplyOptions::default(), None)
|
||
.await
|
||
}
|
||
|
||
pub async fn apply_schema_with_options(
|
||
&self,
|
||
desired_schema_source: &str,
|
||
options: SchemaApplyOptions,
|
||
) -> Result<SchemaApplyResult> {
|
||
self.apply_schema_as(desired_schema_source, options, None)
|
||
.await
|
||
}
|
||
|
||
/// Apply a schema migration with an explicit actor for engine-layer
|
||
/// policy enforcement (MR-722). When a `PolicyChecker` is installed
|
||
/// via [`Self::with_policy`], this method calls `enforce(SchemaApply,
|
||
/// Branch("main"), actor)` before any apply work happens. Denial
|
||
/// returns `OmniError::Policy` and leaves the manifest untouched.
|
||
///
|
||
/// The no-actor variants (`apply_schema`, `apply_schema_with_options`)
|
||
/// pass `None` here. They work fine without a policy; if a policy IS
|
||
/// installed and actor is None, enforcement intentionally fails to
|
||
/// prevent silent-bypass-via-forgetting-the-actor footguns.
|
||
pub async fn apply_schema_as(
|
||
&self,
|
||
desired_schema_source: &str,
|
||
options: SchemaApplyOptions,
|
||
actor: Option<&str>,
|
||
) -> Result<SchemaApplyResult> {
|
||
self.apply_schema_as_with_catalog_check(desired_schema_source, options, actor, |_| Ok(()))
|
||
.await
|
||
}
|
||
|
||
pub async fn apply_schema_as_with_catalog_check<F>(
|
||
&self,
|
||
desired_schema_source: &str,
|
||
options: SchemaApplyOptions,
|
||
actor: Option<&str>,
|
||
validate_catalog: F,
|
||
) -> Result<SchemaApplyResult>
|
||
where
|
||
F: FnOnce(&Catalog) -> Result<()>,
|
||
{
|
||
schema_apply::apply_schema(
|
||
self,
|
||
desired_schema_source,
|
||
options,
|
||
actor,
|
||
validate_catalog,
|
||
)
|
||
.await
|
||
}
|
||
|
||
pub(crate) async fn ensure_schema_apply_idle(&self, operation: &str) -> Result<()> {
|
||
schema_apply::ensure_schema_apply_idle(self, operation).await
|
||
}
|
||
|
||
async fn ensure_schema_apply_not_locked(&self, operation: &str) -> Result<()> {
|
||
schema_apply::ensure_schema_apply_not_locked(self, operation).await
|
||
}
|
||
|
||
/// Engine-facing trait surface around `TableStore`.
|
||
///
|
||
/// This is the **only** accessor for engine code reaching into the
|
||
/// storage layer. The trait's signatures use opaque `SnapshotHandle`
|
||
/// / `StagedHandle` instead of leaking `lance::Dataset` /
|
||
/// `lance::dataset::transaction::Transaction`, so newly-added engine
|
||
/// call sites cannot drift the staged-write invariant by mistake
|
||
/// (the trait's `stage_*` + `commit_staged` pair is the only way to
|
||
/// land a write).
|
||
pub(crate) fn storage(&self) -> &dyn crate::storage_layer::TableStorage {
|
||
&self.table_store
|
||
}
|
||
|
||
/// Inline-commit residual surface (`create_vector_index`) — the sole
|
||
/// write Lance cannot yet express as a stage-then-commit pair (segment
|
||
/// commit needs `build_index_metadata_from_segments`, Lance #6666).
|
||
/// Deliberately separate from [`Self::storage`] so the default storage
|
||
/// surface is staged-only and a new writer cannot couple "write bytes" with
|
||
/// "advance HEAD" by reaching for `db.storage()`. Only the vector-index
|
||
/// build uses this accessor — delete migrated to the staged path
|
||
/// (`stage_delete`) in MR-A. See
|
||
/// `crate::storage_layer::InlineCommitResidual` for the per-method blocker.
|
||
pub(crate) fn storage_inline_residual(
|
||
&self,
|
||
) -> &dyn crate::storage_layer::InlineCommitResidual {
|
||
&self.table_store
|
||
}
|
||
|
||
/// Engine-level access to the object-store adapter (S3 / local fs).
|
||
/// Used by the recovery sidecar protocol — writers in the engine
|
||
/// call this to write/delete sidecars at `__recovery/{ulid}.json`.
|
||
pub(crate) fn storage_adapter(&self) -> &dyn crate::storage::StorageAdapter {
|
||
self.storage.as_ref()
|
||
}
|
||
|
||
/// Per-`(table_key, branch)` writer queues.
|
||
///
|
||
/// Engine-internal writers (mutation finalize, schema_apply,
|
||
/// branch_merge, ensure_indices) and the future MR-870
|
||
/// recovery reconciler reach the queue manager via this accessor.
|
||
/// Returns an `Arc` clone so callers can hold the manager across
|
||
/// `&mut self` engine API boundaries.
|
||
pub(crate) fn write_queue(&self) -> Arc<crate::db::write_queue::WriteQueueManager> {
|
||
Arc::clone(&self.write_queue)
|
||
}
|
||
|
||
/// Engine-internal access to the merge-exclusive mutex. Held across
|
||
/// the swap → operate → restore window in `branch_merge_impl` so
|
||
/// concurrent merges with distinct targets don't corrupt
|
||
/// `self.coordinator` mid-operation. See the field doc on
|
||
/// `Omnigraph::merge_exclusive` for the full design rationale.
|
||
pub(crate) fn merge_exclusive(&self) -> Arc<tokio::sync::Mutex<()>> {
|
||
Arc::clone(&self.merge_exclusive)
|
||
}
|
||
|
||
/// Engine-level access to the graph's normalized root URI. Used by
|
||
/// the recovery sidecar protocol to compute `__recovery/` paths.
|
||
pub(crate) fn root_uri(&self) -> &str {
|
||
&self.root_uri
|
||
}
|
||
|
||
pub(crate) async fn open_coordinator_for_branch(
|
||
&self,
|
||
branch: Option<&str>,
|
||
) -> Result<GraphCoordinator> {
|
||
match branch {
|
||
Some(branch) => {
|
||
GraphCoordinator::open_branch(self.uri(), branch, Arc::clone(&self.storage)).await
|
||
}
|
||
None => GraphCoordinator::open(self.uri(), Arc::clone(&self.storage)).await,
|
||
}
|
||
}
|
||
|
||
pub(crate) async fn swap_coordinator_for_branch(
|
||
&self,
|
||
branch: Option<&str>,
|
||
) -> Result<GraphCoordinator> {
|
||
let next = self.open_coordinator_for_branch(branch).await?;
|
||
let mut coord = self.coordinator.write().await;
|
||
Ok(std::mem::replace(&mut *coord, next))
|
||
}
|
||
|
||
pub(crate) async fn restore_coordinator(&self, coordinator: GraphCoordinator) {
|
||
*self.coordinator.write().await = coordinator;
|
||
}
|
||
|
||
/// Open a capture-once write transaction (RFC-013 step 3b): validate the schema
|
||
/// contract ONCE and pin the base snapshot. The per-table opens take
|
||
/// `Option<&WriteTxn>` and, on the bound branch for the non-strict (Insert/Merge)
|
||
/// path, source the pinned base entry — instead of re-resolving (re-validating the
|
||
/// schema) per table. Strict ops, the fork path, and the commit-time OCC re-read
|
||
/// keep their fresh reads (those are correctness machinery — see the handoff doc).
|
||
///
|
||
/// "Once" covers the table-touch hot path captured here (proven by the node-insert
|
||
/// gate `write_validates_schema_contract_once`); it does NOT yet cover edge endpoint
|
||
/// / cardinality RI validation (`ensure_node_id_exists`, the loader's RI/cardinality),
|
||
/// which still resolve through `snapshot_for_branch` and re-validate. Those reads must
|
||
/// observe LIVE committed state, so unifying them (validate-once + pinned + re-checked
|
||
/// read-set) is step 4's §7.1 work — threading `txn.base` there would re-introduce the
|
||
/// stale-read class the #298 cardinality fix removed. A session-aware base open is
|
||
/// likewise deferred to step 5 (handoff §1d).
|
||
pub(crate) async fn open_write_txn(&self, branch: Option<&str>) -> Result<WriteTxn> {
|
||
let resolved = self.resolved_branch_target(branch).await?;
|
||
Ok(WriteTxn {
|
||
branch: resolved.branch,
|
||
base: resolved.snapshot,
|
||
})
|
||
}
|
||
|
||
pub(crate) async fn resolved_branch_target(
|
||
&self,
|
||
branch: Option<&str>,
|
||
) -> Result<ResolvedTarget> {
|
||
self.ensure_schema_state_valid().await?;
|
||
let requested = ReadTarget::Branch(branch.unwrap_or("main").to_string());
|
||
let normalized = normalize_branch_name(branch.unwrap_or("main"))?;
|
||
let coord = self.coordinator.read().await;
|
||
if normalized.as_deref() == coord.current_branch() {
|
||
let snapshot_id = coord.head_commit_id().await?.unwrap_or_else(|| {
|
||
SnapshotId::synthetic(
|
||
coord.current_branch(),
|
||
coord.version(),
|
||
coord.manifest_incarnation().e_tag.as_deref(),
|
||
)
|
||
});
|
||
return Ok(ResolvedTarget {
|
||
requested,
|
||
branch: coord.current_branch().map(str::to_string),
|
||
snapshot_id,
|
||
snapshot: coord.snapshot(),
|
||
});
|
||
}
|
||
coord.resolve_target(&requested).await
|
||
}
|
||
|
||
pub(crate) async fn snapshot_for_branch(&self, branch: Option<&str>) -> Result<Snapshot> {
|
||
self.resolved_branch_target(branch)
|
||
.await
|
||
.map(|resolved| resolved.snapshot)
|
||
}
|
||
|
||
pub(crate) async fn fresh_snapshot_for_branch(&self, branch: Option<&str>) -> Result<Snapshot> {
|
||
self.ensure_schema_state_valid().await?;
|
||
self.fresh_snapshot_for_branch_unchecked(branch).await
|
||
}
|
||
|
||
/// Fresh per-branch manifest snapshot WITHOUT the schema-contract
|
||
/// re-validation. Identical OCC freshness to [`fresh_snapshot_for_branch`]
|
||
/// — a fresh manifest re-read from storage, never the warm cache — only the
|
||
/// redundant `ensure_schema_state_valid` is dropped. Used inside a single
|
||
/// write once a `WriteTxn` has already validated the contract at capture: the
|
||
/// commit-time drift re-read needs the live manifest, not a second contract
|
||
/// read. Callers with no `WriteTxn` MUST use the checked variant.
|
||
///
|
||
/// Reads the manifest directly via `ManifestCoordinator` rather than
|
||
/// `resolve_target`. The OCC re-read uses only the returned `Snapshot`
|
||
/// (per-table location + version), which `ManifestCoordinator::open().snapshot()`
|
||
/// produces identically to `GraphCoordinator::open(...).snapshot()` — but
|
||
/// `resolve_target` additionally opens the commit graph (an extra
|
||
/// `_graph_commits.lance` probe) the OCC read never consults. Skipping that
|
||
/// load is a pure read-cost reduction, not a freshness change. The checked
|
||
/// `fresh_snapshot_for_branch` delegates here, so its no-`txn` callers
|
||
/// (commit_all's None arm, optimize, repair, fork reclaim) get the same
|
||
/// identical `Snapshot` via this lighter manifest-only read; they consume
|
||
/// only the snapshot and never relied on the commit-graph side load.
|
||
pub(crate) async fn fresh_snapshot_for_branch_unchecked(
|
||
&self,
|
||
branch: Option<&str>,
|
||
) -> Result<Snapshot> {
|
||
let manifest = match branch {
|
||
Some(branch) => {
|
||
crate::db::manifest::ManifestCoordinator::open_at_branch(self.uri(), branch).await?
|
||
}
|
||
None => crate::db::manifest::ManifestCoordinator::open(self.uri()).await?,
|
||
};
|
||
Ok(manifest.snapshot())
|
||
}
|
||
|
||
/// Probe-gated OCC re-capture snapshot (RFC-013 PR2 #1a). The commit-time OCC
|
||
/// re-read MUST be as fresh as a cold re-scan (see `commit_all`), but a cold
|
||
/// `__manifest` full scan per write is O(fragments) and a dominant write-path
|
||
/// cost. When the warm coordinator is already bound to `branch` AND a cheap
|
||
/// incarnation probe (one object-store op, no row scan) proves it is current,
|
||
/// the warm snapshot IS byte-identical to a fresh re-read, so reuse it with
|
||
/// zero scans. On ANY mismatch — a concurrent in- or cross-process advance, or
|
||
/// a different bound branch — fall through to the cold
|
||
/// `fresh_snapshot_for_branch_unchecked` read, preserving the "must be fresh"
|
||
/// contract and cross-process drift detection. Mirrors the read path's
|
||
/// `resolve_target_inner` probe-and-reuse idiom; same-branch sequential writes
|
||
/// keep `coordinator` current (each commit refreshes its `known_state`), so the
|
||
/// common case reuses and skips the scan. The publish CAS
|
||
/// (`expected_table_versions`) remains the final arbiter.
|
||
pub(crate) async fn occ_snapshot_for_branch(&self, branch: Option<&str>) -> Result<Snapshot> {
|
||
{
|
||
let coord = self.coordinator.read().await;
|
||
if branch == coord.current_branch() {
|
||
let held = coord.manifest_incarnation();
|
||
if coord.probe_latest_incarnation().await?.matches(&held) {
|
||
return Ok(coord.snapshot());
|
||
}
|
||
}
|
||
}
|
||
self.fresh_snapshot_for_branch_unchecked(branch).await
|
||
}
|
||
|
||
pub(crate) async fn version(&self) -> u64 {
|
||
self.coordinator.read().await.version()
|
||
}
|
||
|
||
/// Return an immutable Snapshot from the known manifest state. No storage I/O.
|
||
pub(crate) async fn snapshot(&self) -> Snapshot {
|
||
self.coordinator.read().await.snapshot()
|
||
}
|
||
|
||
pub async fn snapshot_of(&self, target: impl Into<ReadTarget>) -> Result<Snapshot> {
|
||
self.resolved_target(target)
|
||
.await
|
||
.map(|resolved| resolved.snapshot)
|
||
}
|
||
|
||
pub async fn version_of(&self, target: impl Into<ReadTarget>) -> Result<u64> {
|
||
self.snapshot_of(target)
|
||
.await
|
||
.map(|snapshot| snapshot.version())
|
||
}
|
||
|
||
/// The on-disk internal-schema version of `target`'s branch (the storage-format
|
||
/// version this graph is stamped at). Surfaced via `omnigraph snapshot`.
|
||
pub async fn internal_schema_version_of(
|
||
&self,
|
||
target: impl Into<ReadTarget>,
|
||
) -> Result<u32> {
|
||
let branch = self.resolved_branch_of(target).await?;
|
||
crate::db::manifest::internal_schema_stamp_at(self.uri(), branch.as_deref()).await
|
||
}
|
||
|
||
pub async fn resolved_branch_of(
|
||
&self,
|
||
target: impl Into<ReadTarget>,
|
||
) -> Result<Option<String>> {
|
||
self.resolved_target(target)
|
||
.await
|
||
.map(|resolved| resolved.branch)
|
||
}
|
||
|
||
/// Synchronize this handle's write base to the latest head of the named branch.
|
||
pub async fn sync_branch(&self, branch: &str) -> Result<()> {
|
||
self.ensure_schema_state_valid().await?;
|
||
let branch = normalize_branch_name(branch)?;
|
||
let next = self.open_coordinator_for_branch(branch.as_deref()).await?;
|
||
*self.coordinator.write().await = next;
|
||
self.invalidate_read_caches().await;
|
||
Ok(())
|
||
}
|
||
|
||
async fn invalidate_read_caches(&self) {
|
||
self.runtime_cache.invalidate_all().await;
|
||
self.read_caches.handles.invalidate_all().await;
|
||
}
|
||
|
||
/// Re-read the handle-local coordinator state from storage AND run
|
||
/// in-process recovery. Closes the Phase B → Phase C residual (e.g.
|
||
/// `MutationStaging::finalize` crash mid-publish in a long-running
|
||
/// server) without restart.
|
||
///
|
||
/// Composition mirrors `Omnigraph::open_with_storage_and_mode`'s
|
||
/// recovery sequence, in the same order, with one restriction: the
|
||
/// manifest-drift heal runs in `RollForwardOnly` mode (rollback /
|
||
/// abort cases defer to the next ReadWrite open because
|
||
/// `Dataset::restore` is unsafe under concurrency). Each step:
|
||
///
|
||
/// 1. `coordinator.refresh()` — re-read manifest.
|
||
/// 2. `recover_schema_state_files` — complete an in-flight
|
||
/// schema_apply's staging→final rename if a SchemaApply sidecar
|
||
/// is on disk; idempotent + early-returns when no staging files
|
||
/// exist. Required BEFORE manifest-drift recovery so a
|
||
/// SchemaApply roll-forward doesn't publish the manifest while
|
||
/// the staging files remain unrenamed (which would corrupt the
|
||
/// graph: data on new schema, catalog on old).
|
||
/// 3. `heal_pending_sidecars_roll_forward` — close the
|
||
/// finalize→publisher residual via roll-forward; defer rollback
|
||
/// work to next ReadWrite open. Serializes against live writers
|
||
/// by acquiring each sidecar's per-(table_key, branch) write
|
||
/// queues, so refresh never rolls forward an in-flight writer's
|
||
/// sidecar from under it.
|
||
/// 4. `runtime_cache.invalidate_all` — drop stale per-snapshot caches.
|
||
///
|
||
/// Steady state cost: one `list_dir` of `__recovery/` (typically
|
||
/// returns empty → early return for both passes). No additional
|
||
/// Lance reads.
|
||
///
|
||
/// The staged-write entry points (`load_as`, `mutate_as`) run the
|
||
/// same heal via
|
||
/// [`heal_pending_recovery_sidecars`](Self::heal_pending_recovery_sidecars),
|
||
/// so a long-lived server converges on the next write without an
|
||
/// explicit refresh. Engine-internal callers that already hold an
|
||
/// in-flight sidecar (e.g. `schema_apply` mid-write) MUST use
|
||
/// [`refresh_coordinator_only`](Self::refresh_coordinator_only) to
|
||
/// avoid the recovery sweep racing their own sidecar.
|
||
pub async fn refresh(&self) -> Result<()> {
|
||
// Standalone schema-staging reconcile ONLY when no recovery
|
||
// sidecar exists (legacy/manual staging residue). When sidecars
|
||
// exist, the heal below owns the reconcile — per SchemaApply
|
||
// sidecar, under that sidecar's queue guards — because an
|
||
// unserialized reconcile can promote a LIVE schema apply's
|
||
// staging files from under it, and a pre-promoted result would
|
||
// make the heal's own guarded reconcile see clean staging and
|
||
// wrongly defer the sidecar. The no-sidecar case cannot race a
|
||
// live apply: its sidecar is on disk before its staging files.
|
||
//
|
||
// Scope the coord write guard to the schema-state section only.
|
||
// `reload_schema_if_source_changed` (below) acquires
|
||
// `self.coordinator.read().await` when the on-disk schema source
|
||
// has drifted from the cached `schema_source`. Tokio's RwLock is
|
||
// not reentrant, so holding the write across that call deadlocks.
|
||
// Pinned by `composite_flow_schema_apply_then_branch_ops_no_deadlock_in_refresh`.
|
||
// The heal also takes the lock itself (queues → coordinator
|
||
// order), so it must run after this guard is released.
|
||
{
|
||
// Hold the schema-apply serialization key across the
|
||
// list-then-reconcile pair: without it, a live apply can
|
||
// write its sidecar + staging between the empty check and
|
||
// the reconcile (the same race, through a smaller window).
|
||
// Queue before coordinator — the documented lock order.
|
||
//
|
||
// Liveness note: with a pending NON-SchemaApply sidecar
|
||
// (e.g. a Mutation residual), this gate skips the standalone
|
||
// reconcile and the heal below reconciles only per
|
||
// SchemaApply sidecar — so pre-sidecar-era orphaned staging
|
||
// residue waits for the NEXT refresh after the sidecars are
|
||
// consumed. Convergence holds, one pass late. Do not "fix"
|
||
// by re-running the reconcile unserialized here: that is
|
||
// exactly the live-apply race this block exists to close.
|
||
let _serial = self
|
||
.write_queue
|
||
.acquire(&crate::db::manifest::schema_apply_serial_queue_key())
|
||
.await;
|
||
if crate::db::manifest::list_sidecars(&self.root_uri, self.storage.as_ref())
|
||
.await?
|
||
.is_empty()
|
||
{
|
||
let mut coord = self.coordinator.write().await;
|
||
coord.refresh().await?;
|
||
recover_schema_state_files(
|
||
&self.root_uri,
|
||
Arc::clone(&self.storage),
|
||
&coord.snapshot(),
|
||
)
|
||
.await?;
|
||
}
|
||
} // ← guards released before the heal's queue acquisition
|
||
crate::db::manifest::heal_pending_sidecars_roll_forward(
|
||
&self.root_uri,
|
||
Arc::clone(&self.storage),
|
||
&self.coordinator,
|
||
&self.write_queue,
|
||
)
|
||
.await?;
|
||
self.reload_schema_if_source_changed().await?;
|
||
self.invalidate_read_caches().await;
|
||
Ok(())
|
||
}
|
||
|
||
/// Write-entry heal: converge any pending recovery sidecars (a
|
||
/// previously failed writer's Phase B → Phase C residual) before
|
||
/// starting a new staged write, so a long-lived process (the HTTP
|
||
/// server, an embedded handle) recovers on its next write instead
|
||
/// of wedging every write on the commit-time drift guard until
|
||
/// restart. Roll-forward only; rollback-eligible sidecars defer to
|
||
/// the next ReadWrite open exactly as [`refresh`](Self::refresh)
|
||
/// does.
|
||
///
|
||
/// Steady-state cost: one `list_dir` of `__recovery/` (typically
|
||
/// empty → immediate return). See
|
||
/// `recovery::heal_pending_sidecars_roll_forward` for the
|
||
/// concurrency contract (per-table write-queue acquisition).
|
||
pub(crate) async fn heal_pending_recovery_sidecars(&self) -> Result<()> {
|
||
let processed = crate::db::manifest::heal_pending_sidecars_roll_forward(
|
||
&self.root_uri,
|
||
Arc::clone(&self.storage),
|
||
&self.coordinator,
|
||
&self.write_queue,
|
||
)
|
||
.await?;
|
||
if processed {
|
||
// A rolled-forward SchemaApply sidecar moved disk + manifest
|
||
// to the new schema (staging promoted, registrations
|
||
// published); the in-memory catalog must follow or the very
|
||
// write that triggered the heal validates against the stale
|
||
// schema. Same post-heal step as `refresh`.
|
||
self.reload_schema_if_source_changed().await?;
|
||
self.invalidate_read_caches().await;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
async fn reload_schema_if_source_changed(&self) -> Result<()> {
|
||
let schema_path = schema_source_uri(&self.root_uri);
|
||
let schema_source = self.storage.read_text(&schema_path).await?;
|
||
if schema_source == *self.schema_source.load_full() {
|
||
return Ok(());
|
||
}
|
||
let current_source_ir = read_schema_ir_from_source(&schema_source)?;
|
||
let branches = self.coordinator.read().await.branch_list().await?;
|
||
let (accepted_ir, _) = load_or_bootstrap_schema_contract(
|
||
&self.root_uri,
|
||
Arc::clone(&self.storage),
|
||
&branches,
|
||
¤t_source_ir,
|
||
)
|
||
.await?;
|
||
let mut catalog = build_catalog_from_ir(&accepted_ir)?;
|
||
fixup_blob_schemas(&mut catalog);
|
||
self.store_schema_source(schema_source);
|
||
self.store_catalog(catalog);
|
||
Ok(())
|
||
}
|
||
|
||
/// Refresh coordinator state and invalidate the runtime cache WITHOUT
|
||
/// running the recovery sweep. Engine-internal callers that hold an
|
||
/// in-flight sidecar (e.g. `schema_apply::apply_schema_with_lock`'s
|
||
/// internal lease-check refresh) need this variant: running recovery
|
||
/// here would observe the caller's own sidecar, classify it as
|
||
/// RolledPastExpected, and roll it forward — racing the caller's
|
||
/// own publish path.
|
||
pub(crate) async fn refresh_coordinator_only(&self) -> Result<()> {
|
||
self.coordinator.write().await.refresh().await?;
|
||
self.invalidate_read_caches().await;
|
||
Ok(())
|
||
}
|
||
|
||
pub async fn resolve_snapshot(&self, branch: &str) -> Result<SnapshotId> {
|
||
self.ensure_schema_state_valid().await?;
|
||
self.coordinator
|
||
.read()
|
||
.await
|
||
.resolve_snapshot_id(branch)
|
||
.await
|
||
}
|
||
|
||
pub(crate) async fn resolved_target(
|
||
&self,
|
||
target: impl Into<ReadTarget>,
|
||
) -> Result<ResolvedTarget> {
|
||
self.ensure_schema_state_valid().await?;
|
||
let target = target.into();
|
||
let mut resolved = self.resolve_target_inner(&target).await?;
|
||
// Attach the read caches (shared Session + held-handle cache) for live
|
||
// Branch reads so table opens reuse handles (0 IO on a warm repeat).
|
||
// Snapshot-id reads are deliberately NOT cached: they pin a historical
|
||
// version `cleanup` may GC, so bypassing the cache sidesteps the
|
||
// cleanup-vs-cached-handle edge. Writes never reach here (they use
|
||
// `resolved_branch_target`), so they never receive a pinned handle.
|
||
if matches!(target, ReadTarget::Branch(_)) {
|
||
resolved
|
||
.snapshot
|
||
.set_read_caches(Arc::clone(&self.read_caches));
|
||
}
|
||
Ok(resolved)
|
||
}
|
||
|
||
/// Resolve a read target to its snapshot, without attaching read caches.
|
||
/// Same-branch reads reuse the warm coordinator, gated by a cheap version
|
||
/// probe (invariant 6: strong consistency, never a blind warm read). Reads do
|
||
/// not need the commit graph (the manifest version is the visibility
|
||
/// authority, invariant 2), so the id is synthetic and no commit-graph scan
|
||
/// happens on this path.
|
||
async fn resolve_target_inner(&self, target: &ReadTarget) -> Result<ResolvedTarget> {
|
||
if let ReadTarget::Branch(branch) = target {
|
||
let normalized = normalize_branch_name(branch)?;
|
||
{
|
||
let coord = self.coordinator.read().await;
|
||
if normalized.as_deref() != coord.current_branch() {
|
||
// Different branch: cold resolve (opens that branch).
|
||
return coord.resolve_target(target).await;
|
||
}
|
||
let held = coord.manifest_incarnation();
|
||
if coord.probe_latest_incarnation().await?.matches(&held) {
|
||
return Ok(warm_resolved_target(&coord, target));
|
||
}
|
||
// Stale: refresh under the write lock below.
|
||
}
|
||
let mut coord = self.coordinator.write().await;
|
||
if normalized.as_deref() == coord.current_branch() {
|
||
// Re-check after taking the write lock; another writer may have
|
||
// refreshed (tokio RwLock has no read->write upgrade).
|
||
let held = coord.manifest_incarnation();
|
||
let mut refreshed = false;
|
||
if !coord.probe_latest_incarnation().await?.matches(&held) {
|
||
coord.refresh_manifest_only().await?;
|
||
refreshed = true;
|
||
}
|
||
let resolved = warm_resolved_target(&coord, target);
|
||
drop(coord);
|
||
if refreshed {
|
||
self.invalidate_read_caches().await;
|
||
}
|
||
return Ok(resolved);
|
||
}
|
||
// Branch changed while waiting for the write lock: cold resolve.
|
||
return coord.resolve_target(target).await;
|
||
}
|
||
|
||
// Snapshot target: resolve through the commit graph as before.
|
||
self.coordinator.read().await.resolve_target(target).await
|
||
}
|
||
|
||
// ─── Change detection ────────────────────────────────────────────────
|
||
|
||
pub async fn diff_between(
|
||
&self,
|
||
from: impl Into<ReadTarget>,
|
||
to: impl Into<ReadTarget>,
|
||
filter: &crate::changes::ChangeFilter,
|
||
) -> Result<crate::changes::ChangeSet> {
|
||
let from_resolved = self.resolved_target(from).await?;
|
||
let to_resolved = self.resolved_target(to).await?;
|
||
crate::changes::diff_snapshots(
|
||
self.uri(),
|
||
&from_resolved.snapshot,
|
||
&to_resolved.snapshot,
|
||
filter,
|
||
to_resolved.branch.clone().or(from_resolved.branch.clone()),
|
||
)
|
||
.await
|
||
}
|
||
|
||
/// Diff two graph commits. Resolves each commit to `(manifest_branch, manifest_version)`
|
||
/// and creates branch-aware snapshots. Supports cross-branch comparison.
|
||
pub async fn diff_commits(
|
||
&self,
|
||
from_commit_id: &str,
|
||
to_commit_id: &str,
|
||
filter: &crate::changes::ChangeFilter,
|
||
) -> Result<crate::changes::ChangeSet> {
|
||
let coord = self.coordinator.read().await;
|
||
let from_commit = coord
|
||
.resolve_commit(&SnapshotId::new(from_commit_id))
|
||
.await?;
|
||
let to_commit = coord.resolve_commit(&SnapshotId::new(to_commit_id)).await?;
|
||
let from_snap = coord
|
||
.resolve_target(&ReadTarget::Snapshot(SnapshotId::new(
|
||
from_commit.graph_commit_id.clone(),
|
||
)))
|
||
.await?;
|
||
let to_snap = coord
|
||
.resolve_target(&ReadTarget::Snapshot(SnapshotId::new(
|
||
to_commit.graph_commit_id.clone(),
|
||
)))
|
||
.await?;
|
||
drop(coord);
|
||
crate::changes::diff_snapshots(
|
||
self.uri(),
|
||
&from_snap.snapshot,
|
||
&to_snap.snapshot,
|
||
filter,
|
||
to_snap.branch.clone().or(from_snap.branch.clone()),
|
||
)
|
||
.await
|
||
}
|
||
|
||
pub async fn entity_at_target(
|
||
&self,
|
||
target: impl Into<ReadTarget>,
|
||
table_key: &str,
|
||
id: &str,
|
||
) -> Result<Option<serde_json::Value>> {
|
||
export::entity_at_target(self, target, table_key, id).await
|
||
}
|
||
|
||
/// Read one entity at a specific manifest version via time travel (on-demand enrichment).
|
||
pub async fn entity_at(
|
||
&self,
|
||
table_key: &str,
|
||
id: &str,
|
||
version: u64,
|
||
) -> Result<Option<serde_json::Value>> {
|
||
export::entity_at(self, table_key, id, version).await
|
||
}
|
||
|
||
/// Create a Snapshot at any historical manifest version.
|
||
pub async fn snapshot_at_version(&self, version: u64) -> Result<Snapshot> {
|
||
self.ensure_schema_state_valid().await?;
|
||
self.coordinator
|
||
.read()
|
||
.await
|
||
.snapshot_at_version(version)
|
||
.await
|
||
}
|
||
|
||
pub async fn export_jsonl(
|
||
&self,
|
||
branch: &str,
|
||
type_names: &[String],
|
||
table_keys: &[String],
|
||
) -> Result<String> {
|
||
export::export_jsonl(self, branch, type_names, table_keys).await
|
||
}
|
||
|
||
pub async fn export_jsonl_to_writer<W: Write>(
|
||
&self,
|
||
branch: &str,
|
||
type_names: &[String],
|
||
table_keys: &[String],
|
||
writer: &mut W,
|
||
) -> Result<()> {
|
||
export::export_jsonl_to_writer(self, branch, type_names, table_keys, writer).await
|
||
}
|
||
|
||
// ─── Graph index ──────────────────────────────────────────────────────
|
||
|
||
/// Get or build the graph index for the current snapshot.
|
||
pub async fn graph_index(&self) -> Result<Arc<crate::graph_index::GraphIndex>> {
|
||
table_ops::graph_index(self).await
|
||
}
|
||
|
||
pub(crate) async fn graph_index_for_resolved(
|
||
&self,
|
||
resolved: &ResolvedTarget,
|
||
edge_types: &std::collections::HashMap<String, (String, String)>,
|
||
) -> Result<Arc<crate::graph_index::GraphIndex>> {
|
||
table_ops::graph_index_for_resolved(self, resolved, edge_types).await
|
||
}
|
||
|
||
/// Ensure BTree scalar indices exist on key columns.
|
||
/// Idempotent — Lance skips if index already exists.
|
||
///
|
||
/// Opens sub-tables at their latest version (not snapshot-pinned) because
|
||
/// indices must be created on the current head. Any version drift from the
|
||
/// snapshot is expected and logged. The resulting versions are committed
|
||
/// back to the manifest.
|
||
///
|
||
/// On named branches, indexing preserves lazy branching:
|
||
/// unbranched subtables keep inheriting `main`, while subtables inherited
|
||
/// from an ancestor branch are first forked into the active branch before
|
||
/// their index metadata is updated.
|
||
/// Returns the declared indexes that could not be materialized on this
|
||
/// pass (today: vector columns with no trainable vectors yet). They are
|
||
/// deferred, not errors; a later `ensure_indices`/`optimize` builds them
|
||
/// once the column is trainable. Reads stay correct (brute-force) meanwhile.
|
||
pub async fn ensure_indices(&self) -> Result<Vec<PendingIndex>> {
|
||
table_ops::ensure_indices(self).await
|
||
}
|
||
|
||
pub async fn ensure_indices_on(&self, branch: &str) -> Result<Vec<PendingIndex>> {
|
||
table_ops::ensure_indices_on(self, branch).await
|
||
}
|
||
|
||
#[cfg(feature = "failpoints")]
|
||
#[doc(hidden)]
|
||
pub async fn failpoint_publish_table_head_without_index_rebuild_for_test(
|
||
&mut self,
|
||
branch: &str,
|
||
table_key: &str,
|
||
table_branch: Option<&str>,
|
||
) -> Result<u64> {
|
||
table_ops::failpoint_publish_table_head_without_index_rebuild_for_test(
|
||
self,
|
||
branch,
|
||
table_key,
|
||
table_branch,
|
||
)
|
||
.await
|
||
}
|
||
|
||
/// Compact small Lance fragments into fewer larger ones across every
|
||
/// node + edge table on `main`. See [`optimize`] for details.
|
||
pub async fn optimize(&self) -> Result<Vec<optimize::TableOptimizeStats>> {
|
||
optimize::optimize_all_tables(self).await
|
||
}
|
||
|
||
/// Classify and explicitly repair uncovered manifest/head drift. See
|
||
/// [`repair`] for the distinction between safe maintenance drift and
|
||
/// suspicious/unverifiable drift.
|
||
pub async fn repair(&self, options: repair::RepairOptions) -> Result<repair::RepairStats> {
|
||
repair::repair_all_tables(self, options).await
|
||
}
|
||
|
||
/// Remove Lance manifests (and the fragments they uniquely own) per the
|
||
/// given [`optimize::CleanupPolicyOptions`]. Destructive to version
|
||
/// history. See [`optimize`] for details.
|
||
pub async fn cleanup(
|
||
&mut self,
|
||
options: optimize::CleanupPolicyOptions,
|
||
) -> Result<Vec<optimize::TableCleanupStats>> {
|
||
optimize::cleanup_all_tables(self, options).await
|
||
}
|
||
|
||
/// Read a blob from a node by its string ID and property name.
|
||
///
|
||
/// Returns a `BlobFile` handle with async `read()`, `seek()`, `tell()`,
|
||
/// and metadata accessors (`size()`, `kind()`, `uri()`).
|
||
///
|
||
/// ```ignore
|
||
/// let blob = db.read_blob("Document", "readme", "content").await?;
|
||
/// let bytes = blob.read().await?;
|
||
/// ```
|
||
pub async fn read_blob(&self, type_name: &str, id: &str, property: &str) -> Result<BlobFile> {
|
||
self.ensure_schema_state_valid().await?;
|
||
let catalog = self.catalog();
|
||
let node_type = catalog
|
||
.node_types
|
||
.get(type_name)
|
||
.ok_or_else(|| OmniError::manifest(format!("unknown node type '{}'", type_name)))?;
|
||
if !node_type.blob_properties.contains(property) {
|
||
return Err(OmniError::manifest(format!(
|
||
"property '{}' on type '{}' is not a Blob",
|
||
property, type_name
|
||
)));
|
||
}
|
||
|
||
let snapshot = self.snapshot().await;
|
||
let table_key = format!("node:{}", type_name);
|
||
let handle = self
|
||
.storage()
|
||
.open_snapshot_at_table(&snapshot, &table_key)
|
||
.await?;
|
||
|
||
let filter_sql = format!("id = '{}'", id.replace('\'', "''"));
|
||
let row_id = self
|
||
.storage()
|
||
.first_row_id_for_filter(&handle, &filter_sql)
|
||
.await?
|
||
.ok_or_else(|| {
|
||
OmniError::manifest(format!("no {} with id '{}' found", type_name, id))
|
||
})?;
|
||
|
||
// `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
|
||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||
|
||
blobs.pop().ok_or_else(|| {
|
||
OmniError::manifest(format!(
|
||
"blob '{}' on {} '{}' returned no data",
|
||
property, type_name, id
|
||
))
|
||
})
|
||
}
|
||
|
||
pub(crate) async fn active_branch(&self) -> Option<String> {
|
||
self.coordinator
|
||
.read()
|
||
.await
|
||
.current_branch()
|
||
.map(str::to_string)
|
||
}
|
||
|
||
async fn ensure_branch_delete_safe(&self, branch: &str, branches: &[String]) -> Result<()> {
|
||
let descendants = self
|
||
.coordinator
|
||
.read()
|
||
.await
|
||
.branch_descendants(branch)
|
||
.await?;
|
||
if let Some(descendant) = descendants.first() {
|
||
return Err(OmniError::manifest_conflict(format!(
|
||
"cannot delete branch '{}' because descendant branch '{}' still depends on it",
|
||
branch, descendant
|
||
)));
|
||
}
|
||
|
||
for other_branch in branches
|
||
.iter()
|
||
.filter(|candidate| candidate.as_str() != branch)
|
||
{
|
||
let snapshot = self
|
||
.snapshot_of(ReadTarget::branch(other_branch.as_str()))
|
||
.await?;
|
||
if snapshot
|
||
.entries()
|
||
.any(|entry| entry.table_branch.as_deref() == Some(branch))
|
||
{
|
||
return Err(OmniError::manifest_conflict(format!(
|
||
"cannot delete branch '{}' because branch '{}' still depends on it",
|
||
branch, other_branch
|
||
)));
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Best-effort reclaim of the per-table Lance forks a just-deleted branch
|
||
/// owned. Runs AFTER the manifest authority flip, so the branch is already
|
||
/// gone and these forks are unreachable orphans. A failure here (transient
|
||
/// object-store error, the `branch_delete.before_table_cleanup` failpoint)
|
||
/// is logged and swallowed: the `cleanup` reconciler is the guaranteed
|
||
/// backstop that converges any leftover orphan. Uses `force_delete_branch`
|
||
/// so a partially-reclaimed retry is idempotent.
|
||
async fn cleanup_deleted_branch_tables(&self, branch: &str, owned_tables: &[(String, String)]) {
|
||
let mut seen_paths = HashSet::new();
|
||
let mut cleanup_targets = owned_tables
|
||
.iter()
|
||
.filter(|(_, table_path)| seen_paths.insert(table_path.clone()))
|
||
.cloned()
|
||
.collect::<Vec<_>>();
|
||
cleanup_targets.sort_by(|left, right| left.0.cmp(&right.0));
|
||
|
||
for (table_key, table_path) in cleanup_targets {
|
||
let dataset_uri = self.storage().dataset_uri(&table_path);
|
||
let outcome = match crate::failpoints::maybe_fail(crate::failpoints::names::BRANCH_DELETE_BEFORE_TABLE_CLEANUP)
|
||
{
|
||
Ok(()) => {
|
||
self.storage()
|
||
.force_delete_branch(&dataset_uri, branch)
|
||
.await
|
||
}
|
||
Err(injected) => Err(injected),
|
||
};
|
||
if let Err(err) = outcome {
|
||
tracing::warn!(
|
||
target: "omnigraph::branch_delete::cleanup",
|
||
branch = %branch,
|
||
table = %table_key,
|
||
error = %err,
|
||
"best-effort fork reclaim failed; cleanup will reconcile the orphan",
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn delete_branch_storage_only(&self, branch: &str) -> Result<()> {
|
||
let active = self
|
||
.coordinator
|
||
.read()
|
||
.await
|
||
.current_branch()
|
||
.map(str::to_string);
|
||
if active.as_deref() == Some(branch) {
|
||
return Err(OmniError::manifest_conflict(format!(
|
||
"cannot delete currently active branch '{}'",
|
||
branch
|
||
)));
|
||
}
|
||
|
||
let branch_snapshot = self.snapshot_of(ReadTarget::branch(branch)).await?;
|
||
let owned_tables = branch_snapshot
|
||
.entries()
|
||
.filter(|entry| entry.table_branch.as_deref() == Some(branch))
|
||
.map(|entry| (entry.table_key.clone(), entry.table_path.clone()))
|
||
.collect::<Vec<_>>();
|
||
|
||
// Authority flip (+ best-effort commit-graph reclaim) — must succeed.
|
||
self.coordinator.write().await.branch_delete(branch).await?;
|
||
// Best-effort per-table fork reclaim; cleanup reconciles any leftover.
|
||
self.cleanup_deleted_branch_tables(branch, &owned_tables)
|
||
.await;
|
||
Ok(())
|
||
}
|
||
|
||
pub(crate) fn normalize_branch_name(branch: &str) -> Result<Option<String>> {
|
||
normalize_branch_name(branch)
|
||
}
|
||
|
||
pub(crate) async fn head_commit_id_for_branch(
|
||
&self,
|
||
branch: Option<&str>,
|
||
) -> Result<Option<String>> {
|
||
let coordinator = self.open_coordinator_for_branch(branch).await?;
|
||
coordinator
|
||
.head_commit_id()
|
||
.await
|
||
.map(|id| id.map(|snapshot_id| snapshot_id.as_str().to_string()))
|
||
}
|
||
|
||
pub async fn branch_create(&self, name: &str) -> Result<()> {
|
||
self.branch_create_as(name, None).await
|
||
}
|
||
|
||
/// Create a branch from the coordinator's currently-open snapshot,
|
||
/// with an explicit actor for engine-layer policy enforcement
|
||
/// (MR-722 fan-out). Scope is `TargetBranch(name)` — symmetric with
|
||
/// `branch_delete_as`: the branch being acted upon is the target.
|
||
/// Cedar rules using `target_branch_scope: protected` therefore see
|
||
/// the new-branch name and can deny e.g. creating any branch named
|
||
/// `main` from a non-privileged actor.
|
||
pub async fn branch_create_as(&self, name: &str, actor: Option<&str>) -> Result<()> {
|
||
self.enforce(
|
||
omnigraph_policy::PolicyAction::BranchCreate,
|
||
&omnigraph_policy::ResourceScope::TargetBranch(name.to_string()),
|
||
actor,
|
||
)?;
|
||
self.ensure_schema_state_valid().await?;
|
||
self.ensure_schema_apply_idle("branch_create").await?;
|
||
ensure_public_branch_ref(name, "branch_create")?;
|
||
self.coordinator.write().await.branch_create(name).await
|
||
}
|
||
|
||
pub async fn branch_create_from(&self, from: impl Into<ReadTarget>, name: &str) -> Result<()> {
|
||
self.branch_create_from_as(from, name, None).await
|
||
}
|
||
|
||
/// Create a branch from a specific source branch with an explicit
|
||
/// actor for engine-layer policy enforcement (MR-722 fan-out).
|
||
///
|
||
/// Scope is `BranchTransition { source, target }` — matches the
|
||
/// HTTP-layer convention at `server_branch_create`
|
||
/// (branch=Some(from), target_branch=Some(name)), so engine and
|
||
/// HTTP fire the same Cedar decision. Pinned-snapshot sources
|
||
/// (which aren't a branch ref) materialize as the sentinel
|
||
/// `<snapshot>` for the policy check; Cedar rules using
|
||
/// `branch_scope: any` still match, rules pinning a specific
|
||
/// source branch correctly do not.
|
||
pub async fn branch_create_from_as(
|
||
&self,
|
||
from: impl Into<ReadTarget>,
|
||
name: &str,
|
||
actor: Option<&str>,
|
||
) -> Result<()> {
|
||
let target = from.into();
|
||
let source_branch = match &target {
|
||
ReadTarget::Branch(b) => b.clone(),
|
||
_ => "<snapshot>".to_string(),
|
||
};
|
||
self.enforce(
|
||
omnigraph_policy::PolicyAction::BranchCreate,
|
||
&omnigraph_policy::ResourceScope::BranchTransition {
|
||
source: source_branch,
|
||
target: name.to_string(),
|
||
},
|
||
actor,
|
||
)?;
|
||
self.ensure_schema_apply_idle("branch_create_from").await?;
|
||
self.branch_create_from_impl(target, name, false).await
|
||
}
|
||
|
||
async fn branch_create_from_impl(
|
||
&self,
|
||
from: impl Into<ReadTarget>,
|
||
name: &str,
|
||
allow_internal_refs: bool,
|
||
) -> Result<()> {
|
||
let target = from.into();
|
||
let ReadTarget::Branch(branch_name) = target else {
|
||
return Err(OmniError::manifest(
|
||
"branch creation from pinned snapshots is not supported yet".to_string(),
|
||
));
|
||
};
|
||
if !allow_internal_refs {
|
||
ensure_public_branch_ref(&branch_name, "branch_create_from")?;
|
||
ensure_public_branch_ref(name, "branch_create_from")?;
|
||
}
|
||
let branch = normalize_branch_name(&branch_name)?;
|
||
// Operate on a freshly-opened source coordinator that's owned locally
|
||
// — never touch `self.coordinator`. The pre-fix implementation used
|
||
// `swap_coordinator_for_branch` + operate + `restore_coordinator` as
|
||
// three separate `coordinator.write().await` acquisitions; under
|
||
// `&self` concurrency, a second `branch_create_from` could swap
|
||
// self.coordinator between this caller's swap and operate steps,
|
||
// making the operate run against the wrong source branch and
|
||
// forking off the wrong HEAD. Pinned by
|
||
// `concurrent_branch_create_from_distinct_parents_does_not_corrupt_coordinator`
|
||
// in `crates/omnigraph-server/tests/server.rs`.
|
||
//
|
||
// `branch_create` mutates only the local coord's commit-graph cache;
|
||
// the manifest write is durable on disk regardless of which
|
||
// coord-handle issued it. Discarding `source_coord` after the call
|
||
// is the right shape — the new branch is reachable from any
|
||
// subsequent open of any coord.
|
||
let mut source_coord = self.open_coordinator_for_branch(branch.as_deref()).await?;
|
||
source_coord.branch_create(name).await
|
||
}
|
||
|
||
pub async fn branch_list(&self) -> Result<Vec<String>> {
|
||
self.ensure_schema_state_valid().await?;
|
||
self.coordinator.read().await.branch_list().await
|
||
}
|
||
|
||
pub async fn branch_delete(&self, name: &str) -> Result<()> {
|
||
self.branch_delete_as(name, None).await
|
||
}
|
||
|
||
/// Delete a branch with an explicit actor for engine-layer policy
|
||
/// enforcement (MR-722 fan-out). Scope is `TargetBranch(name)` —
|
||
/// matches the HTTP-layer convention at `server_branch_delete`
|
||
/// (branch=None, target_branch=Some(name)). Cedar rules using
|
||
/// `target_branch_scope: protected` therefore correctly gate
|
||
/// deletion of protected branches (e.g. deny BranchDelete against
|
||
/// `main`).
|
||
pub async fn branch_delete_as(&self, name: &str, actor: Option<&str>) -> Result<()> {
|
||
self.enforce(
|
||
omnigraph_policy::PolicyAction::BranchDelete,
|
||
&omnigraph_policy::ResourceScope::TargetBranch(name.to_string()),
|
||
actor,
|
||
)?;
|
||
self.ensure_schema_state_valid().await?;
|
||
self.ensure_schema_apply_idle("branch_delete").await?;
|
||
ensure_public_branch_ref(name, "branch_delete")?;
|
||
self.refresh().await?;
|
||
let branch = normalize_branch_name(name)?
|
||
.ok_or_else(|| OmniError::manifest("cannot delete branch 'main'".to_string()))?;
|
||
let branches = self.coordinator.read().await.branch_list().await?;
|
||
if !branches.iter().any(|candidate| candidate == &branch) {
|
||
return Err(OmniError::manifest_not_found(format!(
|
||
"branch '{}' not found",
|
||
branch
|
||
)));
|
||
}
|
||
|
||
self.ensure_branch_delete_safe(&branch, &branches).await?;
|
||
self.delete_branch_storage_only(&branch).await
|
||
}
|
||
|
||
pub async fn get_commit(&self, commit_id: &str) -> Result<GraphCommit> {
|
||
self.ensure_schema_state_valid().await?;
|
||
self.coordinator
|
||
.read()
|
||
.await
|
||
.resolve_commit(&SnapshotId::new(commit_id))
|
||
.await
|
||
}
|
||
|
||
pub async fn list_commits(&self, branch: Option<&str>) -> Result<Vec<GraphCommit>> {
|
||
self.ensure_schema_state_valid().await?;
|
||
let branch = match branch {
|
||
Some(branch) => normalize_branch_name(branch)?,
|
||
None => None,
|
||
};
|
||
let coordinator = self.open_coordinator_for_branch(branch.as_deref()).await?;
|
||
coordinator.list_commits().await
|
||
}
|
||
|
||
/// Open a sub-table for mutation with version-drift guard.
|
||
///
|
||
/// Checks that the dataset's current version matches the snapshot-pinned
|
||
/// version. If another writer has advanced the version, returns an error
|
||
/// prompting the caller to refresh and retry (optimistic concurrency).
|
||
pub(crate) async fn open_for_mutation(
|
||
&self,
|
||
table_key: &str,
|
||
op_kind: crate::db::MutationOpKind,
|
||
) -> Result<OpenedForMutation> {
|
||
table_ops::open_for_mutation(self, table_key, op_kind).await
|
||
}
|
||
|
||
pub(crate) async fn open_for_mutation_on_branch(
|
||
&self,
|
||
branch: Option<&str>,
|
||
table_key: &str,
|
||
op_kind: crate::db::MutationOpKind,
|
||
txn: Option<&crate::db::WriteTxn>,
|
||
) -> Result<OpenedForMutation> {
|
||
table_ops::open_for_mutation_on_branch(self, branch, table_key, op_kind, txn).await
|
||
}
|
||
|
||
/// Fork `table_key` onto `active_branch` from the given source state,
|
||
/// self-healing a manifest-unreferenced leftover fork if one is in the
|
||
/// way. Callers that reach this MUST already hold the per-`(table_key,
|
||
/// active_branch)` write queue (so the reclaim cannot race an in-process
|
||
/// fork) and must have confirmed via the live manifest that the table is
|
||
/// not yet on `active_branch`. Both the first-write fork path
|
||
/// (`open_owned_dataset_for_branch_write`) and `branch_merge` satisfy this.
|
||
pub(crate) async fn fork_dataset_from_entry_state(
|
||
&self,
|
||
table_key: &str,
|
||
full_path: &str,
|
||
source_branch: Option<&str>,
|
||
source_version: u64,
|
||
active_branch: &str,
|
||
) -> Result<SnapshotHandle> {
|
||
match table_ops::fork_dataset_from_entry_state(
|
||
self,
|
||
table_key,
|
||
full_path,
|
||
source_branch,
|
||
source_version,
|
||
active_branch,
|
||
)
|
||
.await?
|
||
{
|
||
crate::storage_layer::ForkOutcome::Created(ds) => Ok(ds),
|
||
crate::storage_layer::ForkOutcome::RefAlreadyExists => {
|
||
table_ops::reclaim_orphaned_fork_and_refork(
|
||
self,
|
||
table_key,
|
||
full_path,
|
||
source_branch,
|
||
source_version,
|
||
active_branch,
|
||
)
|
||
.await
|
||
}
|
||
}
|
||
}
|
||
|
||
pub(crate) async fn reopen_for_mutation(
|
||
&self,
|
||
table_key: &str,
|
||
full_path: &str,
|
||
table_branch: Option<&str>,
|
||
expected_version: u64,
|
||
op_kind: crate::db::MutationOpKind,
|
||
) -> Result<SnapshotHandle> {
|
||
table_ops::reopen_for_mutation(
|
||
self,
|
||
table_key,
|
||
full_path,
|
||
table_branch,
|
||
expected_version,
|
||
op_kind,
|
||
)
|
||
.await
|
||
}
|
||
|
||
pub(crate) async fn build_indices_on_dataset(
|
||
&self,
|
||
table_key: &str,
|
||
ds: &mut SnapshotHandle,
|
||
) -> Result<Vec<PendingIndex>> {
|
||
table_ops::build_indices_on_dataset(self, table_key, ds).await
|
||
}
|
||
|
||
// Used only by in-tree tests (`#[cfg(test)]`); the runtime path now
|
||
// uses `commit_updates_on_branch_with_expected` exclusively.
|
||
#[cfg(test)]
|
||
pub(crate) async fn commit_updates(
|
||
&mut self,
|
||
updates: &[crate::db::SubTableUpdate],
|
||
) -> Result<u64> {
|
||
table_ops::commit_updates(self, updates).await
|
||
}
|
||
|
||
/// Publish a branch merge: the merged table `updates` and the merge commit
|
||
/// in one manifest CAS (RFC-013 Phase 7). The merge commit's merged-in parent
|
||
/// is `merged_parent_commit_id` (the source head); its first parent is the
|
||
/// live target-branch head, resolved by the publisher.
|
||
pub(crate) async fn commit_merge_with_actor(
|
||
&self,
|
||
updates: &[crate::db::SubTableUpdate],
|
||
merged_parent_commit_id: &str,
|
||
actor_id: Option<&str>,
|
||
) -> Result<String> {
|
||
table_ops::commit_merge_with_actor(self, updates, merged_parent_commit_id, actor_id).await
|
||
}
|
||
|
||
pub(crate) async fn commit_updates_on_branch_with_expected(
|
||
&self,
|
||
branch: Option<&str>,
|
||
updates: &[crate::db::SubTableUpdate],
|
||
expected_table_versions: &std::collections::HashMap<String, u64>,
|
||
actor_id: Option<&str>,
|
||
txn: Option<&crate::db::WriteTxn>,
|
||
committed_handles: std::collections::HashMap<String, crate::storage_layer::SnapshotHandle>,
|
||
) -> Result<u64> {
|
||
table_ops::commit_updates_on_branch_with_expected(
|
||
self,
|
||
branch,
|
||
updates,
|
||
expected_table_versions,
|
||
actor_id,
|
||
txn,
|
||
committed_handles,
|
||
)
|
||
.await
|
||
}
|
||
|
||
/// Invalidate the cached graph index. Called after edge mutations.
|
||
pub(crate) async fn invalidate_graph_index(&self) {
|
||
table_ops::invalidate_graph_index(self).await
|
||
}
|
||
}
|
||
|
||
pub(crate) fn normalize_branch_name(branch: &str) -> Result<Option<String>> {
|
||
let branch = branch.trim();
|
||
if branch.is_empty() {
|
||
return Err(OmniError::manifest(
|
||
"branch name cannot be empty".to_string(),
|
||
));
|
||
}
|
||
if branch == "main" {
|
||
return Ok(None);
|
||
}
|
||
Ok(Some(branch.to_string()))
|
||
}
|
||
|
||
/// Build a `ResolvedTarget` from the warm coordinator without opening the commit
|
||
/// graph. The live branch snapshot is pinned by the manifest incarnation, so the
|
||
/// id is synthetic `(branch, version, e_tag when available)`; nothing on the read
|
||
/// path needs a real commit ULID (only `RuntimeCache` keys on the id, where
|
||
/// synthetic is consistent).
|
||
fn warm_resolved_target(coord: &GraphCoordinator, requested: &ReadTarget) -> ResolvedTarget {
|
||
ResolvedTarget {
|
||
requested: requested.clone(),
|
||
branch: coord.current_branch().map(str::to_string),
|
||
snapshot_id: SnapshotId::synthetic(
|
||
coord.current_branch(),
|
||
coord.version(),
|
||
coord.manifest_incarnation().e_tag.as_deref(),
|
||
),
|
||
snapshot: coord.snapshot(),
|
||
}
|
||
}
|
||
|
||
pub(crate) fn ensure_public_branch_ref(branch: &str, operation: &str) -> Result<()> {
|
||
if is_internal_system_branch(branch) {
|
||
return Err(OmniError::manifest(format!(
|
||
"{} does not allow internal system ref '{}'",
|
||
operation, branch
|
||
)));
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn concat_or_empty_batches(schema: Arc<Schema>, batches: Vec<RecordBatch>) -> Result<RecordBatch> {
|
||
if batches.is_empty() {
|
||
return Ok(RecordBatch::new_empty(schema));
|
||
}
|
||
if batches.len() == 1 {
|
||
return Ok(batches.into_iter().next().unwrap());
|
||
}
|
||
let batch_schema = batches[0].schema();
|
||
arrow_select::concat::concat_batches(&batch_schema, &batches)
|
||
.map_err(|e| OmniError::Lance(e.to_string()))
|
||
}
|
||
|
||
fn blob_properties_for_table_key<'a>(
|
||
catalog: &'a Catalog,
|
||
table_key: &str,
|
||
) -> Result<&'a std::collections::HashSet<String>> {
|
||
if let Some(type_name) = table_key.strip_prefix("node:") {
|
||
return catalog
|
||
.node_types
|
||
.get(type_name)
|
||
.map(|node_type| &node_type.blob_properties)
|
||
.ok_or_else(|| OmniError::manifest(format!("unknown node type '{}'", type_name)));
|
||
}
|
||
if let Some(type_name) = table_key.strip_prefix("edge:") {
|
||
return catalog
|
||
.edge_types
|
||
.get(type_name)
|
||
.map(|edge_type| &edge_type.blob_properties)
|
||
.ok_or_else(|| OmniError::manifest(format!("unknown edge type '{}'", type_name)));
|
||
}
|
||
Err(OmniError::manifest(format!(
|
||
"invalid table key '{}'",
|
||
table_key
|
||
)))
|
||
}
|
||
|
||
fn blob_description_is_null(descriptions: &StructArray, row: usize) -> Result<bool> {
|
||
if descriptions.is_null(row) {
|
||
return Ok(true);
|
||
}
|
||
|
||
let kind = descriptions
|
||
.column_by_name("kind")
|
||
.and_then(|col| col.as_any().downcast_ref::<UInt32Array>())
|
||
.and_then(|arr| (!arr.is_null(row)).then(|| arr.value(row) as u8))
|
||
.or_else(|| {
|
||
descriptions
|
||
.column_by_name("kind")
|
||
.and_then(|col| col.as_any().downcast_ref::<arrow_array::UInt8Array>())
|
||
.and_then(|arr| (!arr.is_null(row)).then(|| arr.value(row)))
|
||
});
|
||
let position = descriptions
|
||
.column_by_name("position")
|
||
.and_then(|col| col.as_any().downcast_ref::<UInt64Array>())
|
||
.and_then(|arr| (!arr.is_null(row)).then(|| arr.value(row)));
|
||
let size = descriptions
|
||
.column_by_name("size")
|
||
.and_then(|col| col.as_any().downcast_ref::<UInt64Array>())
|
||
.and_then(|arr| (!arr.is_null(row)).then(|| arr.value(row)));
|
||
let blob_uri = descriptions
|
||
.column_by_name("blob_uri")
|
||
.and_then(|col| col.as_any().downcast_ref::<StringArray>())
|
||
.and_then(|arr| (!arr.is_null(row)).then(|| arr.value(row)));
|
||
|
||
let Some(kind) = kind else {
|
||
return Ok(true);
|
||
};
|
||
let kind = BlobKind::try_from(kind).map_err(|e| OmniError::Lance(e.to_string()))?;
|
||
if kind != BlobKind::Inline {
|
||
return Ok(false);
|
||
}
|
||
|
||
Ok(position.unwrap_or(0) == 0 && size.unwrap_or(0) == 0 && blob_uri.unwrap_or("").is_empty())
|
||
}
|
||
|
||
/// Replace placeholder `LargeBinary` fields with Lance blob v2 fields.
|
||
///
|
||
/// The compiler crate has no Lance dependency, so `ScalarType::Blob` maps to
|
||
/// `DataType::LargeBinary` as a placeholder. This function replaces those
|
||
/// fields with the real blob v2 struct type via `lance::blob::blob_field()`.
|
||
fn fixup_blob_schemas(catalog: &mut Catalog) {
|
||
for node_type in catalog.node_types.values_mut() {
|
||
if node_type.blob_properties.is_empty() {
|
||
continue;
|
||
}
|
||
let fields: Vec<Field> = node_type
|
||
.arrow_schema
|
||
.fields()
|
||
.iter()
|
||
.map(|f| {
|
||
if node_type.blob_properties.contains(f.name()) {
|
||
blob_field(f.name(), f.is_nullable())
|
||
} else {
|
||
f.as_ref().clone()
|
||
}
|
||
})
|
||
.collect();
|
||
node_type.arrow_schema = Arc::new(Schema::new(fields));
|
||
}
|
||
for edge_type in catalog.edge_types.values_mut() {
|
||
if edge_type.blob_properties.is_empty() {
|
||
continue;
|
||
}
|
||
let fields: Vec<Field> = edge_type
|
||
.arrow_schema
|
||
.fields()
|
||
.iter()
|
||
.map(|f| {
|
||
if edge_type.blob_properties.contains(f.name()) {
|
||
blob_field(f.name(), f.is_nullable())
|
||
} else {
|
||
f.as_ref().clone()
|
||
}
|
||
})
|
||
.collect();
|
||
edge_type.arrow_schema = Arc::new(Schema::new(fields));
|
||
}
|
||
}
|
||
|
||
fn read_schema_ir_from_source(schema_source: &str) -> Result<SchemaIR> {
|
||
let schema_ast = parse_schema(schema_source)?;
|
||
build_schema_ir(&schema_ast).map_err(|err| OmniError::manifest(err.to_string()))
|
||
}
|
||
|
||
/// I/O phase of `Omnigraph::init_with_storage`. Split out so the caller
|
||
/// can pattern-match on the result and run cleanup on error before
|
||
/// returning the original error.
|
||
///
|
||
/// Failpoints fire at the phase boundaries:
|
||
/// * `init.after_schema_pg_written` — `_schema.pg` is on disk. In strict mode
|
||
/// this fires in the caller immediately after the atomic ownership claim; in
|
||
/// force mode it fires here after the explicit overwrite.
|
||
/// * `init.after_schema_contract_written` — `_schema.pg` + `_schema.ir.json`
|
||
/// + `__schema_state.json` are on disk.
|
||
/// * `init.after_coordinator_init` — all schema files plus Lance per-type
|
||
/// datasets and `__manifest/` are on disk. (The cleanup wrapper can only
|
||
/// remove the schema files; Lance directories need `delete_prefix` —
|
||
/// deferred along with `DELETE /graphs/{id}`.)
|
||
async fn init_storage_phase(
|
||
root: &str,
|
||
schema_source: &str,
|
||
schema_ir: &SchemaIR,
|
||
catalog: &Catalog,
|
||
storage: &Arc<dyn StorageAdapter>,
|
||
write_schema_pg: bool,
|
||
) -> Result<GraphCoordinator> {
|
||
if write_schema_pg {
|
||
let schema_path = join_uri(root, SCHEMA_SOURCE_FILENAME);
|
||
storage.write_text(&schema_path, schema_source).await?;
|
||
crate::failpoints::maybe_fail(crate::failpoints::names::INIT_AFTER_SCHEMA_PG_WRITTEN)?;
|
||
}
|
||
|
||
write_schema_contract(root, storage.as_ref(), schema_ir).await?;
|
||
crate::failpoints::maybe_fail(crate::failpoints::names::INIT_AFTER_SCHEMA_CONTRACT_WRITTEN)?;
|
||
|
||
let coordinator = GraphCoordinator::init(root, catalog, Arc::clone(storage)).await?;
|
||
crate::failpoints::maybe_fail(crate::failpoints::names::INIT_AFTER_COORDINATOR_INIT)?;
|
||
|
||
Ok(coordinator)
|
||
}
|
||
|
||
/// Best-effort cleanup of init-phase artifacts. Called from
|
||
/// `init_with_storage` on any error returned by `init_storage_phase`.
|
||
///
|
||
/// Removes the three schema files: `_schema.pg`, `_schema.ir.json`,
|
||
/// `__schema_state.json`. Lance datasets and `__manifest/` are not
|
||
/// touched here — recursive directory deletion requires a
|
||
/// `StorageAdapter::delete_prefix` primitive that's deferred along
|
||
/// with `DELETE /graphs/{id}` (MR-668 PR 2b).
|
||
///
|
||
/// Failures to delete are logged via `tracing::warn` and do not mask
|
||
/// the original init error.
|
||
async fn best_effort_cleanup_init_artifacts(root: &str, storage: &dyn StorageAdapter) {
|
||
for uri in [
|
||
schema_source_uri(root),
|
||
schema_ir_uri(root),
|
||
schema_state_uri(root),
|
||
] {
|
||
if let Err(err) = storage.delete(&uri).await {
|
||
tracing::warn!(
|
||
target: "omnigraph::init::cleanup",
|
||
uri = %uri,
|
||
error = %err,
|
||
"init failed; best-effort cleanup could not delete artifact",
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
fn schema_table_key(type_kind: SchemaTypeKind, name: &str) -> String {
|
||
match type_kind {
|
||
SchemaTypeKind::Node => format!("node:{}", name),
|
||
SchemaTypeKind::Edge => format!("edge:{}", name),
|
||
SchemaTypeKind::Interface => unreachable!("interfaces do not map to tables"),
|
||
}
|
||
}
|
||
|
||
fn schema_for_table_key(catalog: &Catalog, table_key: &str) -> Result<Arc<Schema>> {
|
||
if let Some(type_name) = table_key.strip_prefix("node:") {
|
||
let node_type: &NodeType = catalog
|
||
.node_types
|
||
.get(type_name)
|
||
.ok_or_else(|| OmniError::manifest(format!("unknown node type '{}'", type_name)))?;
|
||
return Ok(node_type.arrow_schema.clone());
|
||
}
|
||
if let Some(type_name) = table_key.strip_prefix("edge:") {
|
||
let edge_type: &EdgeType = catalog
|
||
.edge_types
|
||
.get(type_name)
|
||
.ok_or_else(|| OmniError::manifest(format!("unknown edge type '{}'", type_name)))?;
|
||
return Ok(edge_type.arrow_schema.clone());
|
||
}
|
||
Err(OmniError::manifest(format!(
|
||
"invalid table key '{}'",
|
||
table_key
|
||
)))
|
||
}
|
||
|
||
fn record_batch_row_to_json(batch: &RecordBatch, row: usize) -> Result<serde_json::Value> {
|
||
let mut obj = serde_json::Map::new();
|
||
for (i, field) in batch.schema().fields().iter().enumerate() {
|
||
obj.insert(
|
||
field.name().clone(),
|
||
json_value_from_array(batch.column(i).as_ref(), row)?,
|
||
);
|
||
}
|
||
Ok(serde_json::Value::Object(obj))
|
||
}
|
||
|
||
fn json_value_from_array(array: &dyn Array, row: usize) -> Result<serde_json::Value> {
|
||
if array.is_null(row) {
|
||
return Ok(serde_json::Value::Null);
|
||
}
|
||
|
||
match array.data_type() {
|
||
DataType::Utf8 => Ok(serde_json::Value::String(
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<StringArray>()
|
||
.ok_or_else(|| OmniError::Lance("expected StringArray".to_string()))?
|
||
.value(row)
|
||
.to_string(),
|
||
)),
|
||
DataType::LargeUtf8 => Ok(serde_json::Value::String(
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<LargeStringArray>()
|
||
.ok_or_else(|| OmniError::Lance("expected LargeStringArray".to_string()))?
|
||
.value(row)
|
||
.to_string(),
|
||
)),
|
||
DataType::Boolean => Ok(serde_json::Value::Bool(
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<BooleanArray>()
|
||
.ok_or_else(|| OmniError::Lance("expected BooleanArray".to_string()))?
|
||
.value(row),
|
||
)),
|
||
DataType::Int32 => Ok(serde_json::Value::Number(serde_json::Number::from(
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<Int32Array>()
|
||
.ok_or_else(|| OmniError::Lance("expected Int32Array".to_string()))?
|
||
.value(row),
|
||
))),
|
||
DataType::Int64 => Ok(serde_json::Value::Number(serde_json::Number::from(
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<Int64Array>()
|
||
.ok_or_else(|| OmniError::Lance("expected Int64Array".to_string()))?
|
||
.value(row),
|
||
))),
|
||
DataType::UInt32 => Ok(serde_json::Value::Number(serde_json::Number::from(
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<UInt32Array>()
|
||
.ok_or_else(|| OmniError::Lance("expected UInt32Array".to_string()))?
|
||
.value(row),
|
||
))),
|
||
DataType::UInt64 => Ok(serde_json::Value::Number(serde_json::Number::from(
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<UInt64Array>()
|
||
.ok_or_else(|| OmniError::Lance("expected UInt64Array".to_string()))?
|
||
.value(row),
|
||
))),
|
||
DataType::Float32 => {
|
||
let value = array
|
||
.as_any()
|
||
.downcast_ref::<Float32Array>()
|
||
.ok_or_else(|| OmniError::Lance("expected Float32Array".to_string()))?
|
||
.value(row) as f64;
|
||
Ok(serde_json::Value::Number(
|
||
serde_json::Number::from_f64(value).ok_or_else(|| {
|
||
OmniError::Lance(format!("cannot encode f32 value '{}' as JSON", value))
|
||
})?,
|
||
))
|
||
}
|
||
DataType::Float64 => {
|
||
let value = array
|
||
.as_any()
|
||
.downcast_ref::<Float64Array>()
|
||
.ok_or_else(|| OmniError::Lance("expected Float64Array".to_string()))?
|
||
.value(row);
|
||
Ok(serde_json::Value::Number(
|
||
serde_json::Number::from_f64(value).ok_or_else(|| {
|
||
OmniError::Lance(format!("cannot encode f64 value '{}' as JSON", value))
|
||
})?,
|
||
))
|
||
}
|
||
DataType::Date32 => Ok(serde_json::Value::Number(serde_json::Number::from(
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<Date32Array>()
|
||
.ok_or_else(|| OmniError::Lance("expected Date32Array".to_string()))?
|
||
.value(row),
|
||
))),
|
||
DataType::Binary => Ok(serde_json::Value::String(base64::Engine::encode(
|
||
&base64::engine::general_purpose::STANDARD,
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<BinaryArray>()
|
||
.ok_or_else(|| OmniError::Lance("expected BinaryArray".to_string()))?
|
||
.value(row),
|
||
))),
|
||
DataType::LargeBinary => Ok(serde_json::Value::String(base64::Engine::encode(
|
||
&base64::engine::general_purpose::STANDARD,
|
||
array
|
||
.as_any()
|
||
.downcast_ref::<LargeBinaryArray>()
|
||
.ok_or_else(|| OmniError::Lance("expected LargeBinaryArray".to_string()))?
|
||
.value(row),
|
||
))),
|
||
DataType::List(_) => {
|
||
let list = array
|
||
.as_any()
|
||
.downcast_ref::<ListArray>()
|
||
.ok_or_else(|| OmniError::Lance("expected ListArray".to_string()))?;
|
||
let values = list.value(row);
|
||
let mut out = Vec::with_capacity(values.len());
|
||
for idx in 0..values.len() {
|
||
out.push(json_value_from_array(values.as_ref(), idx)?);
|
||
}
|
||
Ok(serde_json::Value::Array(out))
|
||
}
|
||
DataType::LargeList(_) => {
|
||
let list = array
|
||
.as_any()
|
||
.downcast_ref::<LargeListArray>()
|
||
.ok_or_else(|| OmniError::Lance("expected LargeListArray".to_string()))?;
|
||
let values = list.value(row);
|
||
let mut out = Vec::with_capacity(values.len());
|
||
for idx in 0..values.len() {
|
||
out.push(json_value_from_array(values.as_ref(), idx)?);
|
||
}
|
||
Ok(serde_json::Value::Array(out))
|
||
}
|
||
DataType::FixedSizeList(_, _) => {
|
||
let list = array
|
||
.as_any()
|
||
.downcast_ref::<FixedSizeListArray>()
|
||
.ok_or_else(|| OmniError::Lance("expected FixedSizeListArray".to_string()))?;
|
||
let values = list.value(row);
|
||
let mut out = Vec::with_capacity(values.len());
|
||
for idx in 0..values.len() {
|
||
out.push(json_value_from_array(values.as_ref(), idx)?);
|
||
}
|
||
Ok(serde_json::Value::Array(out))
|
||
}
|
||
DataType::Struct(fields) => {
|
||
let struct_array = array
|
||
.as_any()
|
||
.downcast_ref::<StructArray>()
|
||
.ok_or_else(|| OmniError::Lance("expected StructArray".to_string()))?;
|
||
let mut obj = serde_json::Map::new();
|
||
for (field_idx, field) in fields.iter().enumerate() {
|
||
obj.insert(
|
||
field.name().clone(),
|
||
json_value_from_array(struct_array.column(field_idx).as_ref(), row)?,
|
||
);
|
||
}
|
||
Ok(serde_json::Value::Object(obj))
|
||
}
|
||
_ => {
|
||
let value = arrow_cast::display::array_value_to_string(array, row)
|
||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||
Ok(serde_json::Value::String(value))
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::db::manifest::ManifestCoordinator;
|
||
use async_trait::async_trait;
|
||
use serde_json::Value;
|
||
use std::sync::{Arc, Mutex};
|
||
|
||
use crate::storage::{ObjectStorageAdapter, StorageAdapter, join_uri};
|
||
|
||
const TEST_SCHEMA: &str = r#"
|
||
node Person {
|
||
name: String @key
|
||
age: I32?
|
||
}
|
||
node Company {
|
||
name: String @key
|
||
}
|
||
edge Knows: Person -> Person {
|
||
since: Date?
|
||
}
|
||
edge WorksAt: Person -> Company
|
||
"#;
|
||
|
||
#[derive(Debug)]
|
||
struct RecordingStorageAdapter {
|
||
inner: ObjectStorageAdapter,
|
||
reads: Mutex<Vec<String>>,
|
||
writes: Mutex<Vec<String>>,
|
||
exists_checks: Mutex<Vec<String>>,
|
||
renames: Mutex<Vec<(String, String)>>,
|
||
deletes: Mutex<Vec<String>>,
|
||
}
|
||
|
||
impl Default for RecordingStorageAdapter {
|
||
fn default() -> Self {
|
||
Self {
|
||
inner: ObjectStorageAdapter::local(),
|
||
reads: Mutex::default(),
|
||
writes: Mutex::default(),
|
||
exists_checks: Mutex::default(),
|
||
renames: Mutex::default(),
|
||
deletes: Mutex::default(),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl RecordingStorageAdapter {
|
||
fn reads(&self) -> Vec<String> {
|
||
self.reads.lock().unwrap().clone()
|
||
}
|
||
|
||
fn writes(&self) -> Vec<String> {
|
||
self.writes.lock().unwrap().clone()
|
||
}
|
||
|
||
fn exists_checks(&self) -> Vec<String> {
|
||
self.exists_checks.lock().unwrap().clone()
|
||
}
|
||
}
|
||
|
||
#[async_trait]
|
||
impl StorageAdapter for RecordingStorageAdapter {
|
||
async fn read_text(&self, uri: &str) -> Result<String> {
|
||
self.reads.lock().unwrap().push(uri.to_string());
|
||
self.inner.read_text(uri).await
|
||
}
|
||
|
||
async fn write_text(&self, uri: &str, contents: &str) -> Result<()> {
|
||
self.writes.lock().unwrap().push(uri.to_string());
|
||
self.inner.write_text(uri, contents).await
|
||
}
|
||
|
||
async fn write_text_if_absent(&self, uri: &str, contents: &str) -> Result<bool> {
|
||
self.writes.lock().unwrap().push(uri.to_string());
|
||
self.inner.write_text_if_absent(uri, contents).await
|
||
}
|
||
|
||
async fn exists(&self, uri: &str) -> Result<bool> {
|
||
self.exists_checks.lock().unwrap().push(uri.to_string());
|
||
self.inner.exists(uri).await
|
||
}
|
||
|
||
async fn rename_text(&self, from_uri: &str, to_uri: &str) -> Result<()> {
|
||
self.renames
|
||
.lock()
|
||
.unwrap()
|
||
.push((from_uri.to_string(), to_uri.to_string()));
|
||
self.inner.rename_text(from_uri, to_uri).await
|
||
}
|
||
|
||
async fn delete(&self, uri: &str) -> Result<()> {
|
||
self.deletes.lock().unwrap().push(uri.to_string());
|
||
self.inner.delete(uri).await
|
||
}
|
||
|
||
async fn list_dir(&self, dir_uri: &str) -> Result<Vec<String>> {
|
||
self.inner.list_dir(dir_uri).await
|
||
}
|
||
|
||
async fn read_text_versioned(&self, uri: &str) -> Result<(String, String)> {
|
||
self.inner.read_text_versioned(uri).await
|
||
}
|
||
|
||
async fn write_text_if_match(
|
||
&self,
|
||
uri: &str,
|
||
contents: &str,
|
||
expected_version: &str,
|
||
) -> Result<Option<String>> {
|
||
self.inner
|
||
.write_text_if_match(uri, contents, expected_version)
|
||
.await
|
||
}
|
||
|
||
async fn delete_prefix(&self, prefix_uri: &str) -> Result<()> {
|
||
self.inner.delete_prefix(prefix_uri).await
|
||
}
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
struct InitRaceStorageAdapter {
|
||
inner: ObjectStorageAdapter,
|
||
root: String,
|
||
barrier: Arc<tokio::sync::Barrier>,
|
||
}
|
||
|
||
#[async_trait]
|
||
impl StorageAdapter for InitRaceStorageAdapter {
|
||
async fn read_text(&self, uri: &str) -> Result<String> {
|
||
self.inner.read_text(uri).await
|
||
}
|
||
|
||
async fn write_text(&self, uri: &str, contents: &str) -> Result<()> {
|
||
self.inner.write_text(uri, contents).await
|
||
}
|
||
|
||
async fn write_text_if_absent(&self, uri: &str, contents: &str) -> Result<bool> {
|
||
self.inner.write_text_if_absent(uri, contents).await
|
||
}
|
||
|
||
async fn exists(&self, uri: &str) -> Result<bool> {
|
||
let exists = self.inner.exists(uri).await?;
|
||
if uri == schema_state_uri(&self.root) {
|
||
self.barrier.wait().await;
|
||
}
|
||
Ok(exists)
|
||
}
|
||
|
||
async fn rename_text(&self, from_uri: &str, to_uri: &str) -> Result<()> {
|
||
self.inner.rename_text(from_uri, to_uri).await
|
||
}
|
||
|
||
async fn delete(&self, uri: &str) -> Result<()> {
|
||
self.inner.delete(uri).await
|
||
}
|
||
|
||
async fn list_dir(&self, dir_uri: &str) -> Result<Vec<String>> {
|
||
self.inner.list_dir(dir_uri).await
|
||
}
|
||
|
||
async fn read_text_versioned(&self, uri: &str) -> Result<(String, String)> {
|
||
self.inner.read_text_versioned(uri).await
|
||
}
|
||
|
||
async fn write_text_if_match(
|
||
&self,
|
||
uri: &str,
|
||
contents: &str,
|
||
expected_version: &str,
|
||
) -> Result<Option<String>> {
|
||
self.inner
|
||
.write_text_if_match(uri, contents, expected_version)
|
||
.await
|
||
}
|
||
|
||
async fn delete_prefix(&self, prefix_uri: &str) -> Result<()> {
|
||
self.inner.delete_prefix(prefix_uri).await
|
||
}
|
||
}
|
||
|
||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||
async fn concurrent_strict_init_does_not_delete_winning_schema_files() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap().to_string();
|
||
let root = normalize_root_uri(&uri).unwrap();
|
||
let storage: Arc<dyn StorageAdapter> = Arc::new(InitRaceStorageAdapter {
|
||
inner: ObjectStorageAdapter::local(),
|
||
root,
|
||
barrier: Arc::new(tokio::sync::Barrier::new(2)),
|
||
});
|
||
|
||
let left = Omnigraph::init_with_storage(
|
||
&uri,
|
||
TEST_SCHEMA,
|
||
Arc::clone(&storage),
|
||
InitOptions::default(),
|
||
);
|
||
let right = Omnigraph::init_with_storage(
|
||
&uri,
|
||
TEST_SCHEMA,
|
||
Arc::clone(&storage),
|
||
InitOptions::default(),
|
||
);
|
||
let (left, right) = tokio::join!(left, right);
|
||
let ok_count = usize::from(left.is_ok()) + usize::from(right.is_ok());
|
||
assert_eq!(ok_count, 1, "exactly one concurrent init should win");
|
||
|
||
assert!(
|
||
dir.path().join("_schema.pg").exists(),
|
||
"winning init must leave _schema.pg in place"
|
||
);
|
||
assert!(
|
||
dir.path().join("_schema.ir.json").exists(),
|
||
"winning init must leave _schema.ir.json in place"
|
||
);
|
||
assert!(
|
||
dir.path().join("__schema_state.json").exists(),
|
||
"winning init must leave __schema_state.json in place"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_init_and_open_route_graph_metadata_through_storage_adapter() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let adapter = Arc::new(RecordingStorageAdapter::default());
|
||
|
||
Omnigraph::init_with_storage(uri, TEST_SCHEMA, adapter.clone(), InitOptions::default())
|
||
.await
|
||
.unwrap();
|
||
assert!(adapter.writes().contains(&join_uri(uri, "_schema.pg")));
|
||
assert!(adapter.writes().contains(&join_uri(uri, "_schema.ir.json")));
|
||
assert!(
|
||
adapter
|
||
.writes()
|
||
.contains(&join_uri(uri, "__schema_state.json"))
|
||
);
|
||
|
||
Omnigraph::open_with_storage(uri, adapter.clone())
|
||
.await
|
||
.unwrap();
|
||
assert!(adapter.reads().contains(&join_uri(uri, "_schema.pg")));
|
||
assert!(adapter.reads().contains(&join_uri(uri, "_schema.ir.json")));
|
||
assert!(
|
||
adapter
|
||
.reads()
|
||
.contains(&join_uri(uri, "__schema_state.json"))
|
||
);
|
||
assert!(
|
||
adapter
|
||
.exists_checks()
|
||
.contains(&join_uri(uri, "_schema.ir.json"))
|
||
);
|
||
assert!(
|
||
adapter
|
||
.exists_checks()
|
||
.contains(&join_uri(uri, "__schema_state.json"))
|
||
);
|
||
// (Phase B retired `_graph_commits.lance`: open no longer probes for it.)
|
||
}
|
||
|
||
async fn table_rows_json(db: &Omnigraph, table_key: &str) -> Vec<Value> {
|
||
let snapshot = db.snapshot().await;
|
||
let ds = db
|
||
.storage()
|
||
.open_snapshot_at_table(&snapshot, table_key)
|
||
.await
|
||
.unwrap();
|
||
let batches = db.storage().scan_batches(&ds).await.unwrap();
|
||
batches
|
||
.into_iter()
|
||
.flat_map(|batch| {
|
||
(0..batch.num_rows())
|
||
.map(|row| record_batch_row_to_json(&batch, row).unwrap())
|
||
.collect::<Vec<_>>()
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
async fn seed_person_row(db: &mut Omnigraph, name: &str, age: Option<i32>) {
|
||
// No-txn entry, so the handle is always `Some` (collapse #1's skip is
|
||
// gated on `txn.is_some()`).
|
||
let (ds, full_path, table_branch) = db
|
||
.open_for_mutation("node:Person", crate::db::MutationOpKind::Insert)
|
||
.await
|
||
.unwrap()
|
||
.require_handle("seed_person_row test");
|
||
let schema: Arc<Schema> = Arc::new(ds.dataset().schema().into());
|
||
let columns: Vec<Arc<dyn Array>> = schema
|
||
.fields()
|
||
.iter()
|
||
.map(|field| match field.name().as_str() {
|
||
"id" => Arc::new(StringArray::from(vec![name])) as Arc<dyn Array>,
|
||
"name" => Arc::new(StringArray::from(vec![name])) as Arc<dyn Array>,
|
||
"age" => Arc::new(Int32Array::from(vec![age])) as Arc<dyn Array>,
|
||
_ => new_null_array(field.data_type(), 1),
|
||
})
|
||
.collect();
|
||
let batch = RecordBatch::try_new(Arc::clone(&schema), columns).unwrap();
|
||
let staged = db.storage().stage_append(&ds, batch, &[]).await.unwrap();
|
||
let committed = db.storage().commit_staged(ds, staged).await.unwrap();
|
||
let state = db
|
||
.storage()
|
||
.table_state(&full_path, &committed)
|
||
.await
|
||
.unwrap();
|
||
db.commit_updates(&[crate::db::SubTableUpdate {
|
||
table_key: "node:Person".to_string(),
|
||
table_version: state.version,
|
||
table_branch,
|
||
row_count: state.row_count,
|
||
version_metadata: state.version_metadata,
|
||
}])
|
||
.await
|
||
.unwrap();
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_apply_schema_adds_nullable_property_and_preserves_rows() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let mut db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||
seed_person_row(&mut db, "Alice", Some(30)).await;
|
||
|
||
let desired = TEST_SCHEMA.replace(
|
||
" age: I32?\n}",
|
||
" age: I32?\n nickname: String?\n}",
|
||
);
|
||
let result = db.apply_schema(&desired).await.unwrap();
|
||
assert!(result.applied);
|
||
|
||
let reopened = Omnigraph::open(uri).await.unwrap();
|
||
let rows = table_rows_json(&reopened, "node:Person").await;
|
||
assert_eq!(rows.len(), 1);
|
||
assert_eq!(rows[0]["name"], "Alice");
|
||
assert_eq!(rows[0]["age"], 30);
|
||
assert!(rows[0]["nickname"].is_null());
|
||
assert!(
|
||
reopened.catalog().node_types["Person"]
|
||
.properties
|
||
.contains_key("nickname")
|
||
);
|
||
assert!(dir.path().join("_schema.pg").exists());
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_apply_schema_renames_property_and_preserves_values() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let mut db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||
seed_person_row(&mut db, "Alice", Some(30)).await;
|
||
|
||
let desired = TEST_SCHEMA.replace(
|
||
" age: I32?\n}",
|
||
" years: I32? @rename_from(\"age\")\n}",
|
||
);
|
||
db.apply_schema(&desired).await.unwrap();
|
||
|
||
let reopened = Omnigraph::open(uri).await.unwrap();
|
||
let rows = table_rows_json(&reopened, "node:Person").await;
|
||
assert_eq!(rows[0]["name"], "Alice");
|
||
assert_eq!(rows[0]["years"], 30);
|
||
assert!(rows[0].get("age").is_none());
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_apply_schema_renames_type_and_preserves_historical_snapshot() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let mut db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||
seed_person_row(&mut db, "Alice", Some(30)).await;
|
||
let before_version = db.snapshot().await.version();
|
||
|
||
let desired = TEST_SCHEMA
|
||
.replace("node Person {\n", "node Human @rename_from(\"Person\") {\n")
|
||
.replace("edge Knows: Person -> Person", "edge Knows: Human -> Human")
|
||
.replace(
|
||
"edge WorksAt: Person -> Company",
|
||
"edge WorksAt: Human -> Company",
|
||
);
|
||
db.apply_schema(&desired).await.unwrap();
|
||
|
||
let head = db.snapshot().await;
|
||
assert!(head.entry("node:Person").is_none());
|
||
assert!(head.entry("node:Human").is_some());
|
||
let historical = ManifestCoordinator::snapshot_at(uri, None, before_version)
|
||
.await
|
||
.unwrap();
|
||
assert!(historical.entry("node:Person").is_some());
|
||
assert!(historical.entry("node:Human").is_none());
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_apply_schema_succeeds_after_load() {
|
||
// Historical: schema apply used to be blocked by leftover
|
||
// `__run__` branches. The Run state machine was removed in
|
||
// MR-771, so a fresh graph never creates a `__run__` branch;
|
||
// legacy ones are swept by the v2→v3 manifest migration. This
|
||
// asserts the invariant a current graph upholds: publish leaves
|
||
// no `__run__` branch behind, so schema apply proceeds.
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let mut db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||
|
||
crate::loader::load_jsonl(
|
||
&mut db,
|
||
r#"{"type": "Person", "data": {"name": "Alice", "age": 30}}"#,
|
||
crate::loader::LoadMode::Overwrite,
|
||
)
|
||
.await
|
||
.unwrap();
|
||
|
||
let all_branches = db.coordinator.read().await.all_branches().await.unwrap();
|
||
assert!(
|
||
!all_branches.iter().any(|b| b.starts_with("__run__")),
|
||
"no __run__ branch should exist after publish, got: {:?}",
|
||
all_branches
|
||
);
|
||
|
||
let desired = TEST_SCHEMA.replace(
|
||
" age: I32?\n}",
|
||
" age: I32?\n nickname: String?\n}",
|
||
);
|
||
let result = db.apply_schema(&desired).await.unwrap();
|
||
assert!(result.applied, "schema apply should have applied");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_apply_schema_defers_index_then_reconciler_builds_it() {
|
||
// iss-848: schema apply records the @index intent but builds nothing
|
||
// inline; a later ensure_indices materializes it once the table has
|
||
// rows. (Use `age`, which is unindexed in TEST_SCHEMA — `name @key` is
|
||
// already FTS-indexed at seed, so it can't show the deferral.)
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let mut db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||
seed_person_row(&mut db, "Alice", Some(30)).await;
|
||
|
||
let desired = TEST_SCHEMA.replace("age: I32?", "age: I32? @index");
|
||
db.apply_schema(&desired).await.unwrap();
|
||
|
||
// Apply built nothing — the BTREE on `age` is deferred.
|
||
let snapshot = db.snapshot().await;
|
||
let ds = db
|
||
.storage()
|
||
.open_snapshot_at_table(&snapshot, "node:Person")
|
||
.await
|
||
.unwrap();
|
||
assert!(
|
||
!db.storage().has_btree_index(&ds, "age").await.unwrap(),
|
||
"apply must not build the index inline (deferred to the reconciler)"
|
||
);
|
||
|
||
// The reconciler materializes it (Person has a row).
|
||
db.ensure_indices().await.unwrap();
|
||
let snapshot = db.snapshot().await;
|
||
let ds = db
|
||
.storage()
|
||
.open_snapshot_at_table(&snapshot, "node:Person")
|
||
.await
|
||
.unwrap();
|
||
assert!(
|
||
db.storage().has_btree_index(&ds, "age").await.unwrap(),
|
||
"ensure_indices must build the deferred index"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_apply_schema_rewrite_defers_index_then_reconciler_restores() {
|
||
// iss-848: an AddProperty rewrite writes a new dataset version without
|
||
// rebuilding indexes inline (deferred); ensure_indices restores them.
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let initial_schema = TEST_SCHEMA.replace("name: String @key", "name: String @key @index");
|
||
let mut db = Omnigraph::init(uri, &initial_schema).await.unwrap();
|
||
seed_person_row(&mut db, "Alice", Some(30)).await;
|
||
|
||
let desired = initial_schema.replace(
|
||
" age: I32?\n}",
|
||
" age: I32?\n nickname: String?\n}",
|
||
);
|
||
db.apply_schema(&desired).await.unwrap();
|
||
|
||
// After the rewrite the reconciler restores index coverage.
|
||
db.ensure_indices().await.unwrap();
|
||
let snapshot = db.snapshot().await;
|
||
let ds = db
|
||
.storage()
|
||
.open_snapshot_at_table(&snapshot, "node:Person")
|
||
.await
|
||
.unwrap();
|
||
assert!(db.storage().has_btree_index(&ds, "id").await.unwrap());
|
||
assert!(db.storage().has_fts_index(&ds, "name").await.unwrap());
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_open_for_mutation_rejects_while_schema_apply_locked() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||
let mut db = db;
|
||
db.coordinator
|
||
.write()
|
||
.await
|
||
.branch_create(SCHEMA_APPLY_LOCK_BRANCH)
|
||
.await
|
||
.unwrap();
|
||
|
||
let err = db
|
||
.open_for_mutation("node:Person", crate::db::MutationOpKind::Insert)
|
||
.await
|
||
.unwrap_err();
|
||
assert!(
|
||
err.to_string()
|
||
.contains("write is unavailable while schema apply is in progress")
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_commit_updates_rejects_while_schema_apply_locked() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let mut db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||
db.coordinator
|
||
.write()
|
||
.await
|
||
.branch_create(SCHEMA_APPLY_LOCK_BRANCH)
|
||
.await
|
||
.unwrap();
|
||
|
||
let err = db.commit_updates(&[]).await.unwrap_err();
|
||
assert!(
|
||
err.to_string()
|
||
.contains("write commit is unavailable while schema apply is in progress")
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_branch_list_hides_schema_apply_lock_branch() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let uri = dir.path().to_str().unwrap();
|
||
let mut db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||
db.coordinator
|
||
.write()
|
||
.await
|
||
.branch_create(SCHEMA_APPLY_LOCK_BRANCH)
|
||
.await
|
||
.unwrap();
|
||
|
||
let branches = db.branch_list().await.unwrap();
|
||
assert_eq!(branches, vec!["main".to_string()]);
|
||
}
|
||
}
|