Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 26 additions & 1 deletion crates/lance-context-core/src/rollout_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,12 +176,37 @@ impl RolloutStore {
}

/// Open a rollout dataset with explicit storage and shard configuration.
/// Creates the dataset if it does not exist.
pub async fn open_with_options(uri: &str, options: RolloutStoreOptions) -> LanceResult<Self> {
Self::open_inner(uri, options, true).await
}

/// Open an **existing** rollout dataset. Unlike [`Self::open_with_options`],
/// this does **not** create the dataset when it is absent — it returns the
/// underlying [`LanceError::DatasetNotFound`] instead.
///
/// This is the read/write path used by the server's lazy cache-fill: a
/// cache miss must load a store that genuinely exists on object storage
/// (create-on-absence would silently materialize an empty table for a
/// mistyped name, masking the 404). Store creation goes exclusively through
/// [`Self::open_with_options`] on the explicit `create` route.
pub async fn open_existing_with_options(
uri: &str,
options: RolloutStoreOptions,
) -> LanceResult<Self> {
Self::open_inner(uri, options, false).await
}

async fn open_inner(
uri: &str,
options: RolloutStoreOptions,
create_if_missing: bool,
) -> LanceResult<Self> {
let storage_options = options.storage_options.clone();
let write_shard = derive_shard_id(options.shard_id.as_deref());
let dataset = match Self::load_with_options(uri, storage_options.clone()).await {
Ok(dataset) => dataset,
Err(LanceError::DatasetNotFound { .. }) => {
Err(LanceError::DatasetNotFound { .. }) if create_if_missing => {
Self::create_with_options(uri, storage_options.clone()).await?
}
Err(err) => return Err(err),
Expand Down
26 changes: 25 additions & 1 deletion crates/lance-context-core/src/store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -337,7 +337,31 @@ impl ContextStore {
}

/// Open a dataset with explicit object store configuration (e.g. S3 credentials).
/// Creates the dataset if it does not exist.
pub async fn open_with_options(uri: &str, options: ContextStoreOptions) -> LanceResult<Self> {
Self::open_inner(uri, options, true).await
}

/// Open an **existing** context dataset. Unlike [`Self::open_with_options`],
/// this does **not** create the dataset when it is absent — it returns the
/// underlying [`LanceError::DatasetNotFound`] instead.
///
/// Used by the server's lazy cache-fill path: a cache miss must load a store
/// that genuinely exists on object storage. Create-on-absence would silently
/// materialize an empty context for a mistyped name, masking the 404. Store
/// creation goes exclusively through [`Self::open_with_options`].
pub async fn open_existing_with_options(
uri: &str,
options: ContextStoreOptions,
) -> LanceResult<Self> {
Self::open_inner(uri, options, false).await
}

async fn open_inner(
uri: &str,
options: ContextStoreOptions,
create_if_missing: bool,
) -> LanceResult<Self> {
// Validate blob_columns
for col in &options.blob_columns {
if !VALID_BLOB_COLUMNS.contains(&col.as_str()) {
Expand All @@ -359,7 +383,7 @@ impl ContextStore {
let blob_columns = options.blob_columns.clone();
let (dataset, created) = match Self::load_with_options(uri, storage_options.clone()).await {
Ok(dataset) => (dataset, false),
Err(LanceError::DatasetNotFound { .. }) => {
Err(LanceError::DatasetNotFound { .. }) if create_if_missing => {
let dataset = Self::create_with_options(
uri,
storage_options.clone(),
Expand Down
14 changes: 2 additions & 12 deletions crates/lance-context-server/src/routes/compact.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,7 @@ pub async fn compact(
Path(name): Path<String>,
Json(req): Json<CompactRequest>,
) -> Result<Json<CompactResponse>, AppError> {
let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let config = if req.target_rows_per_fragment.is_some() || req.materialize_deletions.is_some() {
let mut c = CompactionConfig::default();
Expand Down Expand Up @@ -48,12 +43,7 @@ pub async fn compact_stats(
State(state): State<Arc<AppState>>,
Path(name): Path<String>,
) -> Result<Json<CompactStatsResponse>, AppError> {
let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let store = store_lock.read().await;
let stats = store
Expand Down
5 changes: 1 addition & 4 deletions crates/lance-context-server/src/routes/contexts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -93,10 +93,7 @@ pub async fn get_context(
State(state): State<Arc<AppState>>,
Path(name): Path<String>,
) -> Result<Json<ContextInfo>, AppError> {
let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?;
let store_lock = state.get_or_open_context_store(&name).await?;
let store = store_lock.read().await;

Ok(Json(ContextInfo {
Expand Down
126 changes: 60 additions & 66 deletions crates/lance-context-server/src/routes/records.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,7 @@ pub async fn add_records(
));
}

let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let run_id = Uuid::new_v4().to_string();
let mut ids = Vec::with_capacity(req.records.len());
Expand Down Expand Up @@ -83,12 +78,7 @@ pub async fn upsert_record(
));
}

let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let record = record_from_add_request(
&req.record,
Expand Down Expand Up @@ -141,12 +131,7 @@ pub async fn upsert_records(
}
}

let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let core_records: Vec<ContextRecord> = req
.records
Expand Down Expand Up @@ -191,12 +176,7 @@ pub async fn update_record(
));
}

let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let patch = patch_from_dto(&req.patch);
let mut store = store_lock.write().await;
Expand Down Expand Up @@ -236,12 +216,7 @@ pub async fn get_record(
State(state): State<Arc<AppState>>,
Path((name, id)): Path<(String, String)>,
) -> Result<Json<GetRecordResponse>, AppError> {
let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let store = store_lock.read().await;
let record = store.get(&id).await.map_err(AppError::from_lance)?;
Expand All @@ -260,12 +235,7 @@ pub async fn fetch_payload(
State(state): State<Arc<AppState>>,
Path((name, id)): Path<(String, String)>,
) -> Result<Response, AppError> {
let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let store = store_lock.read().await;
let record = store
Expand Down Expand Up @@ -306,12 +276,7 @@ pub async fn get_record_by_external_id(
Path(name): Path<String>,
Query(params): Query<ExternalIdParams>,
) -> Result<Json<GetRecordResponse>, AppError> {
let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let store = store_lock.read().await;
let record = store
Expand All @@ -328,12 +293,7 @@ pub async fn delete_record(
State(state): State<Arc<AppState>>,
Path((name, id)): Path<(String, String)>,
) -> Result<Json<DeleteRecordResponse>, AppError> {
let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let mut store = store_lock.write().await;
let deleted = store
Expand All @@ -350,12 +310,7 @@ pub async fn delete_record_by_external_id(
Path(name): Path<String>,
Query(params): Query<ExternalIdParams>,
) -> Result<Json<DeleteRecordResponse>, AppError> {
let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let mut store = store_lock.write().await;
let deleted = store
Expand Down Expand Up @@ -396,12 +351,7 @@ pub async fn list_records(
})
.transpose()?;

let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let store = store_lock.read().await;
let records = store
Expand Down Expand Up @@ -435,12 +385,7 @@ pub async fn related_records(
Path(name): Path<String>,
Query(params): Query<RelatedParams>,
) -> Result<Json<ListRecordsResponse>, AppError> {
let stores = state.stores.read().await;
let store_lock = stores
.get(&name)
.ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?
.clone();
drop(stores);
let store_lock = state.get_or_open_context_store(&name).await?;

let store = store_lock.read().await;
let records = store
Expand Down Expand Up @@ -1255,4 +1200,53 @@ mod tests {
.await;
assert!(matches!(result, Err(AppError::InvalidRequest(_))));
}

/// A second server instance (empty cache, shared data dir) must lazily open
/// a context created elsewhere instead of 404ing — the multi-replica bug.
#[tokio::test]
async fn second_instance_lazily_opens_context_created_elsewhere() {
let context_name = "ctx";
let (state_a, dir) = test_state(context_name).await;

// Instance A writes a record.
let _ = add_records(
State(state_a.clone()),
Path(context_name.to_string()),
Json(AddRecordsRequest {
records: vec![text_record("hello")],
}),
)
.await
.expect("write on instance A");

// Instance B: fresh AppState over the same data dir, empty cache.
let state_b = Arc::new(AppState {
stores: RwLock::new(HashMap::new()),
rollout_stores: RwLock::new(HashMap::new()),
base_path: dir.path().to_path_buf(),
instance_id: None,
rollout_merge_after_generations: 0,
rollout_cleanup_interval_secs: 0,
rollout_cleanup_min_generations: 1,
});

let Json(list) = list_records(
State(state_b.clone()),
Path(context_name.to_string()),
Query(ListParams::default()),
)
.await
.expect("instance B lazily opens the context instead of 404");
assert_eq!(list.records.len(), 1);

// Reading a context that was never created anywhere still 404s.
let err = list_records(
State(state_b),
Path("no-such-context".to_string()),
Query(ListParams::default()),
)
.await
.unwrap_err();
assert!(matches!(err, AppError::NotFound(_)));
}
}
Loading
Loading