diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 68c27cb..f46465c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,7 +4,8 @@ on: push: branches: [main] pull_request: - branches: [main] + # feat/** so stacked PRs (feature targeting feature) get CI too. + branches: [main, "feat/**"] env: CARGO_TERM_COLOR: always diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d87856..2a577d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ## [Unreleased] +### Added — tenant-partitioned collections (Phase 6) + +- **`config.partition_by`**: create a collection partitioned by a metadata field (e.g. `tenant_id`) and every chunk routes to an internal per-tenant partition — a full engine namespace (own LSM, indexes, attach/evict lifecycle) behind one collection API. Searches and deletes filter by the partition field (exact → one partition; `{"in": [...]}` fans out up to 16, merged by score); chunk ids are collection-unique via the parent's CAS id allocator; partitions auto-create on first ingest (writer role included), attach on demand, are hidden from listings, and cascade-delete with the parent. This moves the scale envelope from per-collection to per-tenant: RAM and refresh cost track the HOT tenant set, so one collection can hold billions of vectors across tenants while serving on bounded memory. Not yet routed on partitioned collections (clear errors): relations, facets, TAMS lookup, vector-space CRUD. + ### Added — "warm serverless" - **Stateless writer role** (`COMPASS_ROLE=writer`): durable-append-only nodes with no local indexes and instant boot. Writes validate against the bucket's collection config, mint ids from CAS-leased blocks, append one WAL fragment, and return its `seq`. Reads and delete-by-filter are refused with clear errors. Consistency contract: durable immediately, searchable on serving nodes within the refresh interval. diff --git a/crates/compass/src/collections/mod.rs b/crates/compass/src/collections/mod.rs index c84736d..7205902 100644 --- a/crates/compass/src/collections/mod.rs +++ b/crates/compass/src/collections/mod.rs @@ -10,6 +10,11 @@ // search with full scoring pipeline, vector space CRUD, background rebuild jobs. pub mod cloud; +#[cfg(all(test, feature = "object-storage"))] +mod partition_cloud_tests; +#[cfg(test)] +mod partition_tests; +pub mod partitions; pub mod rebuild; pub mod relation_store; pub mod relationships; @@ -516,6 +521,35 @@ impl CollectionManager { vector_spaces: Option>, embedding_dims: Option, config: Option, + ) -> Result> { + // The partition separator is reserved: user collections must not + // squat on internal partition namespaces. + if partitions::is_partition_ns(name) { + return Err(format!( + "Collection name '{name}' contains the reserved partition separator '{}'", + partitions::PART_SEP + ) + .into()); + } + if let Some(cfg) = &config { + if let Some(field) = &cfg.partition_by { + if field.is_empty() { + return Err("partition_by must name a metadata field".into()); + } + } + } + self.create_collection_inner(name, vector_spaces, embedding_dims, config) + .await + } + + /// Shared create path. Partition namespaces (containing [`partitions::PART_SEP`]) + /// may only be created internally by the ingest router. + async fn create_collection_inner( + &self, + name: &str, + vector_spaces: Option>, + embedding_dims: Option, + config: Option, ) -> Result> { if self.role == NodeRole::Writer { return Err( @@ -640,9 +674,13 @@ impl CollectionManager { Ok(()) => { // Fresh namespace: seed the id allocator at 0 so every // ingest path (attached or stateless) can claim blocks. - if let Err(e) = + // Partition namespaces mint from the PARENT's allocator + // (collection-unique ids) and are never seeded themselves. + if let Err(e) = if partitions::is_partition_ns(name) { + Ok(()) + } else { crate::storage::id_alloc::seed(self.storage.as_ref(), name, 0).await - { + } { let _ = crate::storage::lsm::delete_namespace(self.storage.as_ref(), name) .await; rollback_local().await; @@ -667,6 +705,16 @@ impl CollectionManager { } } + // Partitioned parents allocate chunk ids from a shared CAS allocator + // in LOCAL mode too (partitions must never mint colliding ids). This + // is a single JSON file under the collection dir — no WAL/manifest. + if !self.cloud_mode + && collection.config.partition_by.is_some() + && !partitions::is_partition_ns(name) + { + crate::storage::id_alloc::seed(self.storage.as_ref(), name, 0).await?; + } + tracing::info!("Created collection '{}'", name); Ok(collection) } @@ -683,6 +731,8 @@ impl CollectionManager { let collections = self.collections.read().await; collections.values().map(|c| c.metadata.clone()).collect() }; + // Partition namespaces are internal — the parent represents them. + out.retain(|c| !partitions::is_partition_ns(&c.name)); if self.lazy_attach { let attached: std::collections::HashSet = out.iter().map(|c| c.name.clone()).collect(); @@ -694,6 +744,9 @@ impl CollectionManager { .collect() }; for name in names { + if partitions::is_partition_ns(&name) { + continue; + } if let Ok(Some(cfg)) = cloud::read_bucket_config(self.storage.as_ref(), &name).await { out.push(Collection { @@ -726,6 +779,49 @@ impl CollectionManager { &self, name: &str, ) -> Result<(), Box> { + // Partitioned parent: cascade over every partition namespace FIRST, + // so a failure mid-cascade leaves the parent (and the retry path) + // intact. Partitions are discovered from all sources — attached map, + // lazy registry, local dirs, and the bucket (a writer node may have + // created partitions this node never saw). + if !partitions::is_partition_ns(name) { + let prefix = format!("{name}{}", partitions::PART_SEP); + let mut parts: std::collections::HashSet = std::collections::HashSet::new(); + { + let collections = self.collections.read().await; + parts.extend( + collections + .keys() + .filter(|k| k.starts_with(&prefix)) + .cloned(), + ); + } + parts.extend( + self.registered + .read() + .await + .iter() + .filter(|k| k.starts_with(&prefix)) + .cloned(), + ); + if let Ok(entries) = std::fs::read_dir(&self.data_dir) { + for e in entries.flatten() { + if let Some(n) = e.file_name().to_str() { + if n.starts_with(&prefix) { + parts.insert(n.to_string()); + } + } + } + } + if self.cloud_mode { + if let Ok(all) = crate::storage::lsm::list_namespaces(self.storage.as_ref()).await { + parts.extend(all.into_iter().filter(|n| n.starts_with(&prefix))); + } + } + for part in parts { + Box::pin(self.delete_collection(&part)).await?; + } + } // Lazy mode: the collection may be registered-but-unattached (or LRU // evicted) — deleting it must still purge the bucket. let attached = { @@ -776,6 +872,8 @@ impl CollectionManager { dims: usize, model: &str, ) -> Result<(), Box> { + self.reject_if_partitioned(collection_name, "vector-space changes") + .await?; // Validate `space_name` before it touches the filesystem. The name is // interpolated into on-disk paths (`{space_name}.bin`, `.index`, // `.keymap`), so an unconstrained value like `../../tmp/pwn` could @@ -855,6 +953,8 @@ impl CollectionManager { collection_name: &str, space_name: &str, ) -> Result<(), Box> { + self.reject_if_partitioned(collection_name, "vector-space changes") + .await?; // Same path-traversal guard as add_vector_space — the name flows into // `remove_file` calls below. validate_name_segment(space_name, "Vector space")?; @@ -914,6 +1014,8 @@ impl CollectionManager { collection_name: &str, space_name: &str, ) -> Result<(), Box> { + self.reject_if_partitioned(collection_name, "vector-space changes") + .await?; if self.role == NodeRole::Writer { return Err( "this node runs in writer role; manage vector spaces via a serving node".into(), @@ -1020,6 +1122,214 @@ impl CollectionManager { store::vectors_dir(&self.data_dir, collection_name) } + // ── Tenant partitions (Phase 6) ────────────────────────────────────── + + /// The partition field of a collection, or None for normal collections + /// and partition namespaces themselves. Attaches the parent if needed + /// (cheap: a partitioned parent holds config only, no chunk data). + async fn partition_field( + &self, + name: &str, + ) -> Result, Box> { + if partitions::is_partition_ns(name) { + return Ok(None); + } + self.ensure_attached(name).await?; + let collections = self.collections.read().await; + Ok(collections + .get(name) + .and_then(|l| l.metadata.config.partition_by.clone())) + } + + /// Make sure a partition namespace exists and is servable, creating it on + /// first sight (inheriting the parent's vector spaces + embed model). + /// Racing creators and partitions created by writer nodes resolve via the + /// create path's own already-exists handling. + async fn ensure_partition( + &self, + parent: &str, + pval: &str, + ) -> Result> { + let ns = partitions::partition_ns(parent, pval); + if self.collections.read().await.contains_key(&ns) { + return Ok(ns); + } + if self.cloud_mode && self.registered.read().await.contains(&ns) { + return Ok(ns); // lazy attach loads it at the entry point + } + let (spaces, config) = { + let collections = self.collections.read().await; + let parent_meta = collections + .get(parent) + .ok_or_else(|| not_found(format_args!("Collection '{}' not found", parent)))?; + ( + parent_meta.metadata.vector_spaces.clone(), + CollectionConfig { + embed_model: parent_meta.metadata.config.embed_model.clone(), + partition_by: None, + }, + ) + }; + match self + .create_collection_inner(&ns, Some(spaces), None, Some(config)) + .await + { + Ok(_) => Ok(ns), + // Lost a create race (local map, bucket config, or bucket data) — + // the partition exists; ensure_attached at the entry point loads it. + Err(e) if e.to_string().contains("already") => Ok(ns), + Err(e) => Err(e), + } + } + + /// Ingest into a partitioned collection: one delegated ingest per touched + /// partition. `seq` is passed through when exactly one partition was + /// touched; multi-partition batches return None (each partition has its + /// own manifest and thus its own seq domain). + async fn ingest_partitioned( + &self, + parent: &str, + field: &str, + ingest_chunks: Vec, + embed_state: &EmbedState, + ) -> Result<(usize, HashMap, Option), Box> + { + let groups = partitions::group_by_partition(field, ingest_chunks)?; + let multi = groups.len() > 1; + let mut total = 0usize; + let mut id_map = HashMap::new(); + let mut last_seq = None; + for (pval, group) in groups { + let ns = self.ensure_partition(parent, &pval).await?; + let (n, ids, seq) = Box::pin(self.ingest(&ns, group, embed_state)).await?; + total += n; + id_map.extend(ids); + last_seq = seq; + } + Ok((total, id_map, if multi { None } else { last_seq })) + } + + /// Search a partitioned collection: route to the partitions named by the + /// partition-field filter, merge by score, truncate to top_k. Partitions + /// that do not exist yet contribute zero results (a tenant with no data + /// is empty, not an error). + #[allow(clippy::type_complexity)] + async fn search_partitioned( + &self, + parent: &str, + field: &str, + req: &SearchRequest, + embed_state: &EmbedState, + ) -> Result< + ( + Vec<( + DocumentChunk, + f32, + String, + Option>, + Option>, + )>, + usize, + u64, + Option, + ), + Box, + > { + let pvals = partitions::partition_values_from_filters(field, &req.filters)?; + if req.min_seq.is_some() && pvals.len() > 1 { + return Err("min_seq applies to a single partition's write history; \ + filter to one partition value when using it" + .into()); + } + let mut merged = Vec::new(); + let mut total = 0usize; + let mut took = 0u64; + let mut explain = None; + for pval in pvals { + let ns = partitions::partition_ns(parent, &pval); + match Box::pin(self.search(&ns, req, embed_state)).await { + Ok((results, t, us, ex)) => { + merged.extend(results); + total += t; + took += us; + if explain.is_none() { + explain = ex; + } + } + Err(e) if e.downcast_ref::().is_some() => continue, + Err(e) => return Err(e), + } + } + merged.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + merged.truncate(req.top_k); + Ok((merged, total, took, explain)) + } + + /// Writer-side partition bootstrap: create-only bucket objects for a + /// partition namespace (config copy + empty manifest). Idempotent — racing + /// writers and serving nodes all converge on the first writer's objects. + /// No per-partition id allocator is seeded (ids mint from the parent's). + async fn ensure_partition_ns_cloud( + &self, + parent_cfg: &cloud::BucketConfig, + pval: &str, + ) -> Result> { + let ns = partitions::partition_ns(&parent_cfg.name, pval); + if self.bucket_configs.read().await.contains_key(&ns) { + return Ok(ns); + } + let mut part_cfg = parent_cfg.clone(); + part_cfg.name = ns.clone(); + part_cfg.config.partition_by = None; + match cloud::write_bucket_config_if_absent(self.storage.as_ref(), &ns, &part_cfg).await { + Ok(()) | Err(crate::storage::StorageError::AlreadyExists(_)) => {} + Err(e) => return Err(format!("partition config write failed: {e}").into()), + } + match crate::storage::lsm::init_namespace(self.storage.as_ref(), &ns).await { + Ok(()) | Err(crate::storage::StorageError::AlreadyExists(_)) => {} + Err(e) => return Err(format!("partition manifest init failed: {e}").into()), + } + Ok(ns) + } + + /// Typed fence for operations not yet routed on partitioned collections. + /// Writer nodes consult the bucket config (they hold no local metadata); + /// without this a writer would durably append e.g. relations into the + /// parent namespace, which no serving node ever materializes. + async fn reject_if_partitioned( + &self, + name: &str, + what: &str, + ) -> Result<(), Box> { + if self.is_partitioned_any_role(name).await? { + return Err(format!( + "{what} is not supported on a partitioned collection yet \ + (collection '{name}' is partitioned)" + ) + .into()); + } + Ok(()) + } + + /// Role-aware "is this collection partitioned?": serving nodes read local + /// metadata (attaching if needed); writers consult the bucket config. + async fn is_partitioned_any_role( + &self, + name: &str, + ) -> Result> { + if partitions::is_partition_ns(name) { + return Ok(false); + } + if self.role == NodeRole::Writer { + return Ok(self + .bucket_config(name, false) + .await + .map(|c| c.config.partition_by.is_some()) + .unwrap_or(false)); + } + Ok(self.partition_field(name).await?.is_some()) + } + // ── Ingest ─────────────────────────────────────────────────────────── /// Ingest chunks with batch parent resolution, named embeddings, and relationships. @@ -1033,6 +1343,9 @@ impl CollectionManager { count: u64, ) -> Result, Box> { use crate::storage::id_alloc; + // Partitions mint from the PARENT's allocator: chunk ids stay unique + // across the whole partitioned collection. + let ns = partitions::alloc_ns(ns); match id_alloc::claim(self.storage.as_ref(), ns, count).await { Ok(r) => Ok(r), Err(crate::storage::StorageError::NotFound(_)) => { @@ -1146,6 +1459,26 @@ impl CollectionManager { } let cfg = self.bucket_config(collection_name, false).await?; + // Partitioned collection: group by the partition field, make each + // partition's bucket objects exist (idempotent create-only writes — + // a writer may see a tenant before any serving node does), delegate. + if let Some(field) = cfg.config.partition_by.clone() { + let groups = partitions::group_by_partition(&field, ingest_chunks)?; + let multi = groups.len() > 1; + let mut total = 0usize; + let mut id_map = HashMap::new(); + let mut last_seq = None; + for (pval, group) in groups { + let ns = self.ensure_partition_ns_cloud(&cfg, &pval).await?; + let (n, ids, seq) = + Box::pin(self.ingest_stateless(&ns, group, embed_state)).await?; + total += n; + id_map.extend(ids); + last_seq = seq; + } + return Ok((total, id_map, if multi { None } else { last_seq })); + } + // Ids from the writer-side pool (same allocator as attached nodes). // The pool mutex is NEVER held across the S3 claim: drain what's // available, release, claim, push, repeat. Ids already drained are @@ -1324,6 +1657,14 @@ impl CollectionManager { .await; } + // Partitioned collection: group by the partition field and delegate + // each group to its partition namespace (partitions.rs). + if let Some(field) = self.partition_field(collection_name).await? { + return self + .ingest_partitioned(collection_name, &field, ingest_chunks, embed_state) + .await; + } + let count = ingest_chunks.len(); self.ensure_attached(collection_name).await?; @@ -1332,11 +1673,16 @@ impl CollectionManager { // BEFORE taking the write lock (its refill path does S3 round-trips). // A failed ingest after this point leaks the taken ids — gaps are fine; // the invariant is no-reuse, not density. - let cloud_ids: Option> = if self.cloud_mode && count > 0 { - Some(self.take_ids_cloud(collection_name, count).await?) - } else { - None - }; + // + // Partition namespaces use the block allocator in LOCAL mode too: all + // partitions of one collection mint from the PARENT's allocator, so a + // per-partition next_id counter would collide across siblings. + let cloud_ids: Option> = + if (self.cloud_mode || partitions::is_partition_ns(collection_name)) && count > 0 { + Some(self.take_ids_cloud(collection_name, count).await?) + } else { + None + }; let mut collections = self.collections.write().await; let loaded = collections.get_mut(collection_name).ok_or_else(|| { @@ -1901,6 +2247,13 @@ impl CollectionManager { return Err("this node runs in writer role and does not serve queries".into()); } crate::metrics::inc(&crate::metrics::SEARCH_REQUESTS_TOTAL); + + // Partitioned collection: route to the partitions named by the filter. + if let Some(field) = self.partition_field(collection_name).await? { + return self + .search_partitioned(collection_name, &field, req, embed_state) + .await; + } self.ensure_attached(collection_name).await?; // Read-your-writes: wait (bounded) until fragments up to `min_seq` are @@ -2261,6 +2614,8 @@ impl CollectionManager { collection_name: &str, new: Vec, ) -> Result, Box> { + self.reject_if_partitioned(collection_name, "creating relations") + .await?; // Writer role: build the edges without local state — target_status is // stored as "missing" and re-resolved against the live chunk set at // every read on serving nodes — and append ONE durable fragment. @@ -2404,6 +2759,8 @@ impl CollectionManager { collection_name: &str, relation_id: &str, ) -> Result> { + self.reject_if_partitioned(collection_name, "deleting relations") + .await?; // Writer role: durable relation-delete only (idempotent on replay). if self.role == NodeRole::Writer { crate::storage::lsm::append_relation_delete( @@ -2471,6 +2828,8 @@ impl CollectionManager { direction: RelationDirection, types: Option<&[String]>, ) -> Result, Box> { + self.reject_if_partitioned(collection_name, "listing relations") + .await?; if self.role == NodeRole::Writer { return Err("this node runs in writer role and does not serve queries".into()); } @@ -2664,7 +3023,11 @@ impl CollectionManager { &self, ns: &str, ) -> Result<(), Box> { - if !self.lazy_attach { + // Partition namespaces attach on demand EVEN in non-lazy cloud mode: + // partitions appear dynamically (a writer node can mint one at any + // time), so "everything attached at boot" can never hold for them. + let dynamic_partition = self.cloud_mode && partitions::is_partition_ns(ns); + if !self.lazy_attach && !dynamic_partition { return Ok(()); } if self.collections.read().await.contains_key(ns) { @@ -2982,6 +3345,18 @@ impl CollectionManager { ids: &[u64], ) -> Result<(usize, Option), Box> { crate::metrics::inc(&crate::metrics::DELETE_REQUESTS_TOTAL); + // Ids alone don't say which partition holds them (a fan-out probe of + // every partition would be unbounded) — partitioned collections + // delete via POST /delete with the partition filter. This applies to + // writer nodes too: a tombstone appended to the PARENT namespace + // would never be materialized by any serving node. + if self.is_partitioned_any_role(collection_name).await? { + return Err(format!( + "collection '{collection_name}' is partitioned: delete via filters \ + (POST .../delete with the partition field), not bare ids" + ) + .into()); + } // Writer role: durable tombstone only. Without local indexes we can't // filter to ids-that-exist; a tombstone for an absent id is an // idempotent no-op on replay, so append the deduped set as-is. @@ -2998,8 +3373,11 @@ impl CollectionManager { // Ids can never legitimately reach the allocator frontier; a bogus // huge id would otherwise poison max_id forever (rebuilds compute // next_id = max_id + 1 → overflow / id reuse). - let frontier = - crate::storage::id_alloc::frontier(self.storage.as_ref(), collection_name).await?; + let frontier = crate::storage::id_alloc::frontier( + self.storage.as_ref(), + partitions::alloc_ns(collection_name), + ) + .await?; if let Some(bad) = newly.iter().find(|id| **id >= frontier) { return Err(format!( "chunk id {bad} was never allocated in '{collection_name}' \ @@ -3125,6 +3503,26 @@ impl CollectionManager { .into(), ); } + // Partitioned collection: route to the partitions named by the filter + // (same routing rule as search; NotFound partitions delete nothing). + if let Some(field) = self.partition_field(collection_name).await? { + let pvals = partitions::partition_values_from_filters(&field, filters)?; + let multi = pvals.len() > 1; + let mut total = 0usize; + let mut last_seq = None; + for pval in pvals { + let ns = partitions::partition_ns(collection_name, &pval); + match Box::pin(self.delete_by_filter(&ns, filters)).await { + Ok((n, seq)) => { + total += n; + last_seq = seq; + } + Err(e) if e.downcast_ref::().is_some() => continue, + Err(e) => return Err(e), + } + } + return Ok((total, if multi { None } else { last_seq })); + } self.ensure_attached(collection_name).await?; // Resolve matching live ids from the roaring filter index — the same // pushdown search uses, so delete-by-filter and search can never @@ -3379,6 +3777,8 @@ impl CollectionManager { if self.role == NodeRole::Writer { return Err("this node runs in writer role and does not serve queries".into()); } + self.reject_if_partitioned(collection_name, "facet counting") + .await?; self.ensure_attached(collection_name).await?; let collections = self.collections.read().await; let loaded = collections.get(collection_name).ok_or_else(|| { @@ -3428,6 +3828,8 @@ impl CollectionManager { time_start_ms: Option, time_end_ms: Option, ) -> Result, Box> { + self.reject_if_partitioned(collection_name, "temporal segment lookup") + .await?; let collections = self.collections.read().await; let loaded = collections.get(collection_name).ok_or_else(|| { not_found(format_args!("Collection \'{}\' not found", collection_name)) diff --git a/crates/compass/src/collections/partition_cloud_tests.rs b/crates/compass/src/collections/partition_cloud_tests.rs new file mode 100644 index 0000000..8fa2bf1 --- /dev/null +++ b/crates/compass/src/collections/partition_cloud_tests.rs @@ -0,0 +1,268 @@ +// collections/partition_cloud_tests.rs — Phase 6 cloud-mode coverage: writer +// routing, cross-node partition discovery, and cold rebuild. Uses the +// in-memory object_store backend (same code path as real S3). + +use super::*; +use crate::embed::EmbedState; +use crate::storage::object_store_backend::ObjectStoreBackend; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +fn unique_data_dir() -> std::path::PathBuf { + static N: AtomicU64 = AtomicU64::new(0); + std::env::temp_dir().join(format!( + "compass-partition-cloud-{}-{}", + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )) +} + +fn embed_state() -> EmbedState { + EmbedState { + bge: None, + distilled: None, + } +} + +fn mem_storage() -> Arc { + Arc::new(ObjectStoreBackend::from_store( + std::sync::Arc::new(object_store::memory::InMemory::new()), + "object-store:memory", + )) +} + +fn tenant_chunk(tenant: &str, file: &str, text: &str, vec: [f32; 4]) -> IngestChunk { + let mut metadata = HashMap::new(); + metadata.insert( + "tenant".to_string(), + MetadataValue::String(tenant.to_string()), + ); + let mut embeddings = HashMap::new(); + embeddings.insert("default".to_string(), vec.to_vec()); + IngestChunk { + client_id: None, + file_id: file.to_string(), + chunk_index: 0, + page: None, + text: text.to_string(), + metadata, + doc_type: "chunk".to_string(), + parent_id: None, + parent_ref: None, + group_id: None, + embeddings, + embedding: None, + } +} + +fn partitioned_config() -> Option { + Some(CollectionConfig { + embed_model: "bge-small".to_string(), + partition_by: Some("tenant".to_string()), + }) +} + +fn search_req(query: &str, tenant: &str) -> SearchRequest { + let mut filters = HashMap::new(); + filters.insert( + "tenant".to_string(), + FilterValue::Exact(MetadataValue::String(tenant.to_string())), + ); + SearchRequest { + query: query.to_string(), + mode: "fts".to_string(), + vector_space: None, + top_k: 10, + query_vector: None, + filters, + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::default(), + min_seq: None, + } +} + +// A stateless writer routes partitioned ingest into per-partition namespaces +// it bootstraps itself; a serving node that booted BEFORE those partitions +// existed attaches them on demand and serves the data. +#[tokio::test] +async fn writer_partitioned_ingest_visible_on_serving_node() { + let storage = mem_storage(); + let embed = embed_state(); + + // Serving node boots first and creates the (empty) partitioned parent. + let serve_dir = unique_data_dir(); + std::fs::create_dir_all(&serve_dir).unwrap(); + let serving = CollectionManager::new_with_storage(&serve_dir, storage.clone()) + .await + .unwrap(); + serving + .create_collection("wp", None, Some(4), partitioned_config()) + .await + .unwrap(); + + // Writer node ingests for two tenants the serving node has never seen. + let writer_dir = unique_data_dir(); + std::fs::create_dir_all(&writer_dir).unwrap(); + let writer = CollectionManager::new_with_storage_opts( + &writer_dir, + storage.clone(), + NodeRole::Writer, + false, + usize::MAX, + 0, + ) + .await + .unwrap(); + let (n, id_map, _) = writer + .ingest( + "wp", + vec![ + { + let mut c = + tenant_chunk("acme", "a1", "durable acme fact", [0.9, 0.1, 0.0, 0.0]); + c.client_id = Some("a1".to_string()); + c + }, + { + let mut c = + tenant_chunk("globex", "g1", "durable globex fact", [0.1, 0.9, 0.0, 0.0]); + c.client_id = Some("g1".to_string()); + c + }, + ], + &embed, + ) + .await + .unwrap(); + assert_eq!(n, 2); + let ids: std::collections::HashSet = id_map.values().copied().collect(); + assert_eq!(ids.len(), 2, "writer ids must be collection-unique"); + + // The serving node attaches the new partitions on first query. + let (results, _, _, _) = serving + .search("wp", &search_req("durable", "acme"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1, "acme partition attaches on demand"); + assert_eq!(results[0].0.file_id, "a1"); + let (results, _, _, _) = serving + .search("wp", &search_req("durable", "globex"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.file_id, "g1"); + + // Writer refuses partitioned operations that would black-hole data. + let err = writer.delete_chunks("wp", &[0]).await.unwrap_err(); + assert!(err.to_string().contains("delete via filters"), "{err}"); + let err = writer + .create_relations( + "wp", + vec![CreateRelation { + source_chunk_id: 0, + target_chunk_id: 1, + target_document_id: None, + relation_type: "cites".to_string(), + metadata: HashMap::new(), + }], + ) + .await + .unwrap_err(); + assert!( + err.to_string().contains("not supported on a partitioned"), + "{err}" + ); + + let _ = std::fs::remove_dir_all(&serve_dir); + let _ = std::fs::remove_dir_all(&writer_dir); +} + +// A brand-new node with an empty disk recovers a partitioned collection — +// parent config, every partition's data, and the shared id allocator — from +// the bucket alone. +#[tokio::test] +async fn partitioned_collection_cold_rebuild_from_bucket() { + let storage = mem_storage(); + let embed = embed_state(); + + { + let dir_a = unique_data_dir(); + std::fs::create_dir_all(&dir_a).unwrap(); + let a = CollectionManager::new_with_storage(&dir_a, storage.clone()) + .await + .unwrap(); + a.create_collection("cold", None, Some(4), partitioned_config()) + .await + .unwrap(); + a.ingest( + "cold", + vec![ + tenant_chunk("acme", "a1", "cold acme", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("globex", "g1", "cold globex", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + let _ = std::fs::remove_dir_all(&dir_a); + } // node A gone, local disk gone; only the bucket remains. + + let dir_b = unique_data_dir(); + std::fs::create_dir_all(&dir_b).unwrap(); + let b = CollectionManager::new_with_storage(&dir_b, storage.clone()) + .await + .unwrap(); + + // Parent is listed (partitions hidden), routing metadata survived. + let listed = b.list_collections().await; + let names: Vec<&str> = listed.iter().map(|c| c.name.as_str()).collect(); + assert!(names.contains(&"cold"), "{names:?}"); + assert!(!names.iter().any(|n| n.contains(partitions::PART_SEP))); + + let (results, _, _, _) = b + .search("cold", &search_req("cold", "acme"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.file_id, "a1"); + + // New ingest keeps minting collection-unique ids from the recovered + // allocator (never reuses the pre-rebuild ids). + let (_, id_map, _) = b + .ingest( + "cold", + vec![{ + let mut c = tenant_chunk("acme", "a2", "post-rebuild", [0.7, 0.3, 0.0, 0.0]); + c.client_id = Some("a2".to_string()); + c + }], + &embed, + ) + .await + .unwrap(); + let new_id = *id_map.values().next().unwrap(); + assert!( + new_id >= 2, + "rebuilt node must not reuse ids (got {new_id})" + ); + + // Cascade delete purges parent + partitions from the bucket. + b.delete_collection("cold").await.unwrap(); + let remaining = crate::storage::lsm::list_namespaces(storage.as_ref()) + .await + .unwrap(); + assert!( + remaining.iter().all(|n| !n.starts_with("cold")), + "bucket still has: {remaining:?}" + ); + + let _ = std::fs::remove_dir_all(&dir_b); +} diff --git a/crates/compass/src/collections/partition_tests.rs b/crates/compass/src/collections/partition_tests.rs new file mode 100644 index 0000000..0b86039 --- /dev/null +++ b/crates/compass/src/collections/partition_tests.rs @@ -0,0 +1,405 @@ +// collections/partition_tests.rs — tenant-partitioned collections (Phase 6). +// Child module of `collections`: `super::*` sees the parent's private items. + +use super::*; +use crate::embed::EmbedState; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +fn unique_data_dir() -> std::path::PathBuf { + static N: AtomicU64 = AtomicU64::new(0); + std::env::temp_dir().join(format!( + "compass-partition-{}-{}", + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + )) +} + +fn embed_state() -> EmbedState { + EmbedState { + bge: None, + distilled: None, + } +} + +fn tenant_chunk(tenant: &str, file: &str, text: &str, vec: [f32; 4]) -> IngestChunk { + let mut metadata = HashMap::new(); + metadata.insert( + "tenant".to_string(), + MetadataValue::String(tenant.to_string()), + ); + let mut embeddings = HashMap::new(); + embeddings.insert("default".to_string(), vec.to_vec()); + IngestChunk { + client_id: None, + file_id: file.to_string(), + chunk_index: 0, + page: None, + text: text.to_string(), + metadata, + doc_type: "chunk".to_string(), + parent_id: None, + parent_ref: None, + group_id: None, + embeddings, + embedding: None, + } +} + +fn partitioned_config() -> Option { + Some(CollectionConfig { + embed_model: "bge-small".to_string(), + partition_by: Some("tenant".to_string()), + }) +} + +fn tenant_filter(tenant: &str) -> HashMap { + let mut f = HashMap::new(); + f.insert( + "tenant".to_string(), + FilterValue::Exact(MetadataValue::String(tenant.to_string())), + ); + f +} + +fn search_req(query: &str, vec: Option<[f32; 4]>, tenant: &str) -> SearchRequest { + SearchRequest { + query: query.to_string(), + mode: if vec.is_some() { + "semantic".to_string() + } else { + "fts".to_string() + }, + vector_space: None, + top_k: 10, + query_vector: vec.map(|v| v.to_vec()), + filters: tenant_filter(tenant), + score_weights: None, + recency: None, + recency_preset: None, + recency_field: None, + boosts: Vec::new(), + relationship_boost: None, + explain: false, + include_relations: false, + relation_types: None, + relation_direction: RelationDirection::default(), + min_seq: None, + } +} + +async fn local_manager(data_dir: &std::path::Path) -> std::sync::Arc { + std::fs::create_dir_all(data_dir).unwrap(); + CollectionManager::new(data_dir).await.unwrap() +} + +// ── Local mode ────────────────────────────────────────────────────────── + +#[tokio::test] +async fn partitioned_ingest_routes_and_isolates() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("multi", None, Some(4), partitioned_config()) + .await + .unwrap(); + + let (n, _, _) = manager + .ingest( + "multi", + vec![ + tenant_chunk("acme", "a1", "acme secret report", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("acme", "a2", "acme quarterly numbers", [0.8, 0.2, 0.0, 0.0]), + tenant_chunk("globex", "g1", "globex secret memo", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + assert_eq!(n, 3); + + // Tenant-scoped search sees ONLY that tenant's chunks — even for a query + // term both tenants share. + let (results, _, _, _) = manager + .search("multi", &search_req("secret", None, "acme"), &embed) + .await + .unwrap(); + assert_eq!( + results.len(), + 1, + "acme must see exactly its own 'secret' hit" + ); + assert_eq!(results[0].0.file_id, "a1"); + + let (results, _, _, _) = manager + .search("multi", &search_req("secret", None, "globex"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.file_id, "g1"); + + // A tenant that never ingested is empty, not an error. + let (results, total, _, _) = manager + .search("multi", &search_req("secret", None, "initech"), &embed) + .await + .unwrap(); + assert!(results.is_empty()); + assert_eq!(total, 0); + + // Unfiltered search on a partitioned collection is a clear error. + let mut req = search_req("secret", None, "acme"); + req.filters.clear(); + let err = manager.search("multi", &req, &embed).await.unwrap_err(); + assert!(err.to_string().contains("partitioned by 'tenant'"), "{err}"); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partition_ids_are_collection_unique() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("uniq", None, Some(4), partitioned_config()) + .await + .unwrap(); + + // Interleave ingests across tenants; every assigned id must be distinct + // (partitions share the parent's id allocator). + let mut all_ids = Vec::new(); + for round in 0..3 { + for tenant in ["t-a", "t-b", "t-c"] { + let (_, id_map, _) = manager + .ingest( + "uniq", + vec![{ + let mut c = tenant_chunk( + tenant, + &format!("{tenant}-{round}"), + "payload", + [0.5, 0.5, 0.0, 0.0], + ); + c.client_id = Some(format!("{tenant}-{round}")); + c + }], + &embed, + ) + .await + .unwrap(); + all_ids.extend(id_map.values().copied()); + } + } + assert_eq!(all_ids.len(), 9); + let unique: std::collections::HashSet = all_ids.iter().copied().collect(); + assert_eq!(unique.len(), 9, "ids must never collide across partitions"); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partitioned_delete_routes_by_filter_and_rejects_bare_ids() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("deltest", None, Some(4), partitioned_config()) + .await + .unwrap(); + manager + .ingest( + "deltest", + vec![ + tenant_chunk("acme", "a1", "doomed", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("globex", "g1", "doomed", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + + // Bare ids are ambiguous across partitions — rejected with guidance. + let err = manager.delete_chunks("deltest", &[0]).await.unwrap_err(); + assert!(err.to_string().contains("delete via filters"), "{err}"); + + // Filter-scoped delete removes acme's chunk only. + let (n, _) = manager + .delete_by_filter("deltest", &tenant_filter("acme")) + .await + .unwrap(); + assert_eq!(n, 1); + let (results, _, _, _) = manager + .search("deltest", &search_req("doomed", None, "acme"), &embed) + .await + .unwrap(); + assert!(results.is_empty(), "acme's chunk is gone"); + let (results, _, _, _) = manager + .search("deltest", &search_req("doomed", None, "globex"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1, "globex is untouched"); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partitions_hidden_from_listing_and_cascade_deleted() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("casc", None, Some(4), partitioned_config()) + .await + .unwrap(); + manager + .ingest( + "casc", + vec![ + tenant_chunk("acme", "a1", "x", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("globex", "g1", "y", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + + // Listing shows the parent only — partition namespaces are internal. + let listed = manager.list_collections().await; + let names: Vec<&str> = listed.iter().map(|c| c.name.as_str()).collect(); + assert!(names.contains(&"casc")); + assert!( + !names.iter().any(|n| n.contains(partitions::PART_SEP)), + "partition namespaces must not be listed: {names:?}" + ); + + // Deleting the parent removes every partition's data on disk. + manager.delete_collection("casc").await.unwrap(); + let leftovers: Vec = std::fs::read_dir(&data_dir) + .map(|rd| { + rd.flatten() + .filter_map(|e| e.file_name().to_str().map(String::from)) + .filter(|n| n.starts_with("casc")) + .collect() + }) + .unwrap_or_default(); + assert!( + leftovers.is_empty(), + "cascade left dirs behind: {leftovers:?}" + ); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partitioned_fences_and_validation() { + let data_dir = unique_data_dir(); + let manager = local_manager(&data_dir).await; + let embed = embed_state(); + manager + .create_collection("fenced", None, Some(4), partitioned_config()) + .await + .unwrap(); + manager + .ingest( + "fenced", + vec![tenant_chunk("acme", "a1", "x", [0.9, 0.1, 0.0, 0.0])], + &embed, + ) + .await + .unwrap(); + + // Reserved separator in user collection names. + let err = manager + .create_collection("evil--part--x", None, Some(4), None) + .await + .unwrap_err(); + assert!(err.to_string().contains("reserved partition separator")); + + // Chunks missing the partition field are rejected with the field name. + let mut bad = tenant_chunk("acme", "b", "x", [0.1, 0.1, 0.0, 0.0]); + bad.metadata.clear(); + let err = manager + .ingest("fenced", vec![bad], &embed) + .await + .unwrap_err(); + assert!(err.to_string().contains("missing partition field 'tenant'")); + + // Unrouted operations fail closed with a clear message. + let err = manager + .get_facets("fenced", "", &["tenant".to_string()]) + .await + .unwrap_err(); + assert!(err.to_string().contains("not supported on a partitioned")); + let err = manager + .create_relations( + "fenced", + vec![CreateRelation { + source_chunk_id: 0, + target_chunk_id: 1, + target_document_id: None, + relation_type: "cites".to_string(), + metadata: HashMap::new(), + }], + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("not supported on a partitioned")); + let err = manager + .add_vector_space("fenced", "extra", 8, "model") + .await + .unwrap_err(); + assert!(err.to_string().contains("not supported on a partitioned")); + + let _ = std::fs::remove_dir_all(&data_dir); +} + +#[tokio::test] +async fn partitioned_collection_survives_restart() { + let data_dir = unique_data_dir(); + let embed = embed_state(); + { + let manager = local_manager(&data_dir).await; + manager + .create_collection("persist", None, Some(4), partitioned_config()) + .await + .unwrap(); + manager + .ingest( + "persist", + vec![ + tenant_chunk("acme", "a1", "durable acme", [0.9, 0.1, 0.0, 0.0]), + tenant_chunk("globex", "g1", "durable globex", [0.1, 0.9, 0.0, 0.0]), + ], + &embed, + ) + .await + .unwrap(); + } + + let manager = CollectionManager::new(&data_dir).await.unwrap(); + // Routing metadata survives: tenant-scoped search still works, ids keep + // minting from the shared allocator without collision. + let (results, _, _, _) = manager + .search("persist", &search_req("durable", None, "acme"), &embed) + .await + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.file_id, "a1"); + let (_, id_map, _) = manager + .ingest( + "persist", + vec![{ + let mut c = tenant_chunk("acme", "a2", "post-restart", [0.7, 0.3, 0.0, 0.0]); + c.client_id = Some("a2".to_string()); + c + }], + &embed, + ) + .await + .unwrap(); + let new_id = *id_map.values().next().unwrap(); + assert!(new_id >= 2, "restart must not reuse ids (got {new_id})"); + + let _ = std::fs::remove_dir_all(&data_dir); +} diff --git a/crates/compass/src/collections/partitions.rs b/crates/compass/src/collections/partitions.rs new file mode 100644 index 0000000..2be819b --- /dev/null +++ b/crates/compass/src/collections/partitions.rs @@ -0,0 +1,216 @@ +// collections/partitions.rs — tenant-partitioned collections (roadmap Phase 6). +// +// Design: a partition IS a full internal collection. Each partition gets its +// own namespace — LSM manifest, WAL, segments, Tantivy dir, vector files, +// attach/evict lifecycle — by reusing the existing per-collection engine +// wholesale. Phase 6 is a ROUTER at the manager entry points, not a new +// engine: +// +// - `create` with `config.partition_by = "tenant_id"` marks the parent. +// The parent namespace holds config + the shared id allocator; chunk data +// lives only in partitions. +// - Ingest groups chunks by `metadata[partition_by]` and delegates each +// group to the partition's namespace, auto-creating it on first sight. +// - Search/deletes require a filter on the partition field and route to +// exactly the named partitions (set membership fans out, capped). +// - Chunk ids are COLLECTION-unique: every partition claims id blocks from +// the PARENT's allocator (in local mode too — the allocator is just a +// CAS-updated file; this does not create WAL/manifest objects). +// +// Why this shape: per-tenant cost isolation falls out of the existing +// machinery. A query for tenant T attaches T's partition only; LRU eviction +// and refresh scale with the HOT tenant set, not the collection. The +// per-namespace scale envelope now bounds the largest TENANT, not the +// collection, which is what makes a billion-vector multi-tenant collection +// servable on bounded RAM. +// +// Scope fences (MVP, enforced with clear errors): relations, facets, TAMS +// temporal lookup, and vector-space CRUD are not yet routed for partitioned +// collections; partition keys must be kebab-case strings; the partition +// field is immutable after create. + +use crate::models::{FilterValue, IngestChunk, MetadataValue}; +use std::collections::HashMap; + +/// Separator between parent collection name and partition value in the +/// internal namespace. User-facing collection names must not contain it +/// (enforced at create); it is otherwise valid kebab-case, so every existing +/// storage/path rule accepts partition namespaces unchanged. +pub const PART_SEP: &str = "--part--"; + +/// Max partitions a single set-membership search may fan out to. +pub const MAX_SEARCH_FANOUT: usize = 16; + +/// Internal namespace for one partition of a parent collection. +pub fn partition_ns(parent: &str, pval: &str) -> String { + format!("{parent}{PART_SEP}{pval}") +} + +/// Is this namespace a partition (vs a user-facing collection)? +pub fn is_partition_ns(ns: &str) -> bool { + ns.contains(PART_SEP) +} + +/// The parent collection of a partition namespace, or None for a normal one. +pub fn parent_of(ns: &str) -> Option<&str> { + ns.split_once(PART_SEP).map(|(parent, _)| parent) +} + +/// The namespace whose id allocator a collection mints from: partitions share +/// the PARENT's allocator so chunk ids are unique across the whole collection +/// (delete-by-id and search results would otherwise be ambiguous). +pub fn alloc_ns(ns: &str) -> &str { + parent_of(ns).unwrap_or(ns) +} + +/// Partition values become path/namespace segments — hold them to the same +/// kebab-case rule as collection names, and bound the length so a hostile +/// value can't manufacture absurd object keys. +pub fn validate_partition_value(v: &str) -> Result<(), Box> { + if v.is_empty() || v.len() > 64 || !v.chars().all(|c| c.is_ascii_alphanumeric() || c == '-') { + return Err(format!( + "partition value '{v}' is invalid: use letters, digits, and hyphens (max 64 chars)" + ) + .into()); + } + if v.contains(PART_SEP) { + return Err(format!("partition value '{v}' must not contain '{PART_SEP}'").into()); + } + Ok(()) +} + +/// Group an ingest batch by its partition value, validating that every chunk +/// carries a usable `metadata[field]` string. +pub fn group_by_partition( + field: &str, + chunks: Vec, +) -> Result>, Box> { + let mut groups: HashMap> = HashMap::new(); + for chunk in chunks { + let pval = match chunk.metadata.get(field) { + Some(MetadataValue::String(s)) => s.clone(), + Some(_) => { + return Err(format!( + "chunk '{}': partition field '{field}' must be a string", + chunk.file_id + ) + .into()); + } + None => { + return Err(format!( + "chunk '{}' is missing partition field '{field}' \ + (this collection is partitioned by it)", + chunk.file_id + ) + .into()); + } + }; + validate_partition_value(&pval)?; + groups.entry(pval).or_default().push(chunk); + } + Ok(groups) +} + +/// Resolve which partition values a filtered request targets. Exact match +/// routes to one partition; set membership (`in`) fans out (capped). Anything +/// else is an error — a partitioned collection cannot be scanned blind. +pub fn partition_values_from_filters( + field: &str, + filters: &HashMap, +) -> Result, Box> { + let missing = || -> Box { + format!( + "this collection is partitioned by '{field}': include filters.{field} \ + (exact value, or {{\"in\": [...]}} for up to {MAX_SEARCH_FANOUT} partitions)" + ) + .into() + }; + let values = match filters.get(field) { + Some(FilterValue::Exact(MetadataValue::String(s))) => vec![s.clone()], + Some(FilterValue::Condition(cond)) => match &cond.in_values { + Some(vs) if !vs.is_empty() => vs.clone(), + _ => return Err(missing()), + }, + Some(_) => return Err(missing()), + None => return Err(missing()), + }; + if values.len() > MAX_SEARCH_FANOUT { + return Err(format!( + "filters.{field} names {} partitions; max fan-out is {MAX_SEARCH_FANOUT}", + values.len() + ) + .into()); + } + for v in &values { + validate_partition_value(v)?; + } + Ok(values) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::models::FilterCondition; + + #[test] + fn namespace_roundtrip() { + let ns = partition_ns("videos", "acme"); + assert_eq!(ns, "videos--part--acme"); + assert!(is_partition_ns(&ns)); + assert_eq!(parent_of(&ns), Some("videos")); + assert_eq!(alloc_ns(&ns), "videos"); + assert!(!is_partition_ns("videos")); + assert_eq!(parent_of("videos"), None); + assert_eq!(alloc_ns("videos"), "videos"); + } + + #[test] + fn partition_value_rules() { + assert!(validate_partition_value("acme-01").is_ok()); + assert!(validate_partition_value("").is_err()); + assert!(validate_partition_value("has space").is_err()); + assert!(validate_partition_value("a--part--b").is_err()); + assert!(validate_partition_value(&"x".repeat(65)).is_err()); + } + + #[test] + fn filter_routing() { + let field = "tenant"; + let mut f = HashMap::new(); + assert!(partition_values_from_filters(field, &f).is_err()); + + f.insert( + field.to_string(), + FilterValue::Exact(MetadataValue::String("acme".into())), + ); + assert_eq!( + partition_values_from_filters(field, &f).unwrap(), + vec!["acme".to_string()] + ); + + f.insert( + field.to_string(), + FilterValue::Condition(FilterCondition { + gte: None, + lte: None, + contains: None, + in_values: Some(vec!["a".into(), "b".into()]), + }), + ); + assert_eq!(partition_values_from_filters(field, &f).unwrap().len(), 2); + + let too_many: Vec = (0..MAX_SEARCH_FANOUT + 1) + .map(|i| format!("t{i}")) + .collect(); + f.insert( + field.to_string(), + FilterValue::Condition(FilterCondition { + gte: None, + lte: None, + contains: None, + in_values: Some(too_many), + }), + ); + assert!(partition_values_from_filters(field, &f).is_err()); + } +} diff --git a/crates/compass/src/models.rs b/crates/compass/src/models.rs index 85fbf87..7f49c75 100644 --- a/crates/compass/src/models.rs +++ b/crates/compass/src/models.rs @@ -155,6 +155,11 @@ fn default_dims() -> usize { pub struct CollectionConfig { #[serde(default = "default_embed_model")] pub embed_model: String, + /// Tenant-partitioned collections: the metadata field whose (string) + /// value routes each chunk to its own internal partition namespace. + /// Immutable after create. None = normal single-namespace collection. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub partition_by: Option, } fn default_embed_model() -> String { diff --git a/crates/compass/src/storage/mod.rs b/crates/compass/src/storage/mod.rs index 674df99..0e5608b 100644 --- a/crates/compass/src/storage/mod.rs +++ b/crates/compass/src/storage/mod.rs @@ -53,8 +53,9 @@ impl Version { } /// True when the token carries no usable precondition (CAS must refuse it). - // Used by the object-store backend at runtime and by s3_integration tests. - #[cfg(any(test, feature = "object-storage"))] + // Used by the object-store backend at runtime and by s3_integration + // tests — all behind the feature; default builds never reference it. + #[cfg(feature = "object-storage")] pub fn is_empty(&self) -> bool { self.e_tag.is_empty() && self.version.as_deref().is_none_or(str::is_empty) } diff --git a/scripts/e2e.sh b/scripts/e2e.sh index e671648..44a725a 100755 --- a/scripts/e2e.sh +++ b/scripts/e2e.sh @@ -116,6 +116,30 @@ r=$(post $FULL/collections/e2e/compact '') n=$(post $FULL/collections/e2e/search '{"query":"goal","mode":"fts","top_k":5}' | jqn "len(d['results'])") [ "$n" -ge 1 ] && ok "data survives compaction" || bad post-compact "$n" +echo "── tenant partitions ──" +post $FULL/collections '{"name":"mt","embedding_dims":4,"config":{"partition_by":"tenant"}}' >/dev/null +r=$(post $FULL/collections/mt/ingest '{"chunks":[ + {"client_id":"p1","file_id":"p1","chunk_index":0,"doc_type":"chunk","text":"shared secret alpha","metadata":{"tenant":"acme"},"embeddings":{"default":[0.9,0.1,0.1,0.1]}}, + {"client_id":"p2","file_id":"p2","chunk_index":0,"doc_type":"chunk","text":"shared secret beta","metadata":{"tenant":"globex"},"embeddings":{"default":[0.1,0.9,0.1,0.1]}}]}') +[ "$(echo "$r" | jqn "d['indexed']")" = "2" ] && ok "partitioned ingest routes" || bad p-ingest x +n=$(post $FULL/collections/mt/search '{"query":"secret","mode":"fts","top_k":10,"filters":{"tenant":"acme"}}' | jqn "len(d['results'])") +f1=$(post $FULL/collections/mt/search '{"query":"secret","mode":"fts","top_k":10,"filters":{"tenant":"acme"}}' | jqn "d['results'][0]['chunk']['file_id']") +[ "$n" = "1" ] && [ "$f1" = "p1" ] && ok "tenant isolation (acme sees only its hit)" || bad p-iso "$n/$f1" +code=$(curl -s -o /dev/null -w '%{http_code}' -X POST $FULL/collections/mt/search -H 'content-type: application/json' -d '{"query":"secret","mode":"fts"}') +[ "$code" -ge 400 ] && ok "unfiltered partitioned search rejected" || bad p-nofilter "$code" +n=$(post $FULL/collections/mt/search '{"query":"secret","mode":"fts","top_k":10,"filters":{"tenant":{"in":["acme","globex"]}}}' | jqn "len(d['results'])") +[ "$n" = "2" ] && ok "set-membership fan-out merges tenants" || bad p-fanout "$n" +r=$(post $WRITER/collections/mt/ingest '{"chunks":[{"client_id":"p3","file_id":"p3","chunk_index":0,"doc_type":"chunk","text":"writer minted tenant","metadata":{"tenant":"initech"},"embeddings":{"default":[0.1,0.1,0.9,0.1]}}]}') +[ "$(echo "$r" | jqn "d['indexed']")" = "1" ] && ok "writer partitioned ingest" || bad p-writer x +sleep 1 +n=$(post $FULL/collections/mt/search '{"query":"minted","mode":"fts","top_k":5,"filters":{"tenant":"initech"}}' | jqn "len(d['results'])") +[ "$n" = "1" ] && ok "writer-minted partition attaches on serving node" || bad p-attach "$n" +r=$(post $FULL/collections/mt/delete '{"filters":{"tenant":"acme"}}') +[ "$(echo "$r" | jqn "d['deleted']")" = "1" ] && ok "partition-scoped delete" || bad p-del x +code=$(curl -s -o /dev/null -w '%{http_code}' -X DELETE $FULL/collections/mt) +[ "$code" -lt 400 ] && ok "partitioned collection cascade delete" || bad p-casc "$code" +curl -s $FULL/collections | grep -q "part--" && bad "partitions hidden from listing" leak || ok "partitions hidden from listing" + echo "── collection delete ──" code=$(curl -s -o /dev/null -w '%{http_code}' -X DELETE $FULL/collections/e2e) [ "$code" -lt 400 ] && ok "delete collection" || bad coll-del "$code"