diff --git a/crates/api/Cargo.toml b/crates/api/Cargo.toml index 7c198a3..3fbf5d3 100644 --- a/crates/api/Cargo.toml +++ b/crates/api/Cargo.toml @@ -13,4 +13,5 @@ snafu.workspace = true snapshot.workspace = true storage.workspace = true tempfile.workspace = true +tracing.workspace = true uuid.workspace = true diff --git a/crates/api/src/collection.rs b/crates/api/src/collection.rs new file mode 100644 index 0000000..c959a20 --- /dev/null +++ b/crates/api/src/collection.rs @@ -0,0 +1,207 @@ +use std::collections::HashMap; +use std::fs; +use std::path::PathBuf; +use std::sync::{Arc, RwLock}; + +use defs::{Dimension, Similarity}; +use index::IndexType; + +use crate::{ApiError, DbConfig, Result, VectorDb, init_api}; + +pub const DEFAULT_COLLECTION: &str = "default"; +const MAX_COLLECTION_NAME_LEN: usize = 64; + +#[derive(Debug, Clone)] +pub struct CollectionSpec { + pub id: String, + pub dimension: Dimension, + pub similarity: Similarity, + pub index_type: IndexType, +} + +pub struct CollectionManager { + root: PathBuf, // DATA_PATH + defaults: DbConfig, + collections: RwLock>>, + specs: RwLock>, +} + +// TODO : rethink how to implement "default" collection (whether to have a 'default' collection or not) +impl CollectionManager { + pub fn new(root: PathBuf, defaults: DbConfig) -> Result { + let collections_root = root.join("collections"); + + fs::create_dir_all(&collections_root).map_err(|e| ApiError::InitializationFailed { + reason: format!( + "could not create collections dir {}: {e}", + collections_root.display() + ), + })?; + + let default_cfg = DbConfig { + data_path: collections_root.join(DEFAULT_COLLECTION), + ..defaults.clone() + }; + + let default_db = init_api(default_cfg)?; + + let default_spec = CollectionSpec { + id: DEFAULT_COLLECTION.to_string(), + dimension: defaults.dimension, + similarity: defaults.similarity, + index_type: defaults.index_type, + }; + + Ok(Self { + root, + defaults, + collections: RwLock::new(HashMap::from([( + DEFAULT_COLLECTION.to_string(), + Arc::new(default_db), + )])), + specs: RwLock::new(HashMap::from([( + DEFAULT_COLLECTION.to_string(), + default_spec, + )])), + }) + } + + pub fn get(&self, name: &str) -> Result> { + let key = Self::resolve(name); + + let collections = self.collections.read().map_err(|_| ApiError::LockError)?; + + collections + .get(key) + .cloned() + .ok_or_else(|| ApiError::CollectionNotFound { + name: key.to_string(), + }) + } + + fn resolve(name: &str) -> &str { + if name.is_empty() { + DEFAULT_COLLECTION + } else { + name + } + } + + pub fn create(&self, spec: CollectionSpec) -> Result> { + validate_name(&spec.id)?; + + // dimension check + if spec.dimension == 0 { + return Err(ApiError::InitializationFailed { + reason: "collection dimension must be greater than zero".to_string(), + }); + } + + // Check for existing dir/collection of the same name in root/collections dir + // (as the 'collections' HashMap may not contain the collection which have been created during previous startups) + // This may lead to opening of existing RocksDB and the new collection (acc. to the user) + // would inherit the points of the already exisiting collection + let path = self.root.join("collections").join(&spec.id); + if path.exists() { + return Err(ApiError::CollectionAlreadyExists { + name: spec.id.clone(), + }); + } + + // load/build config + let cfg = DbConfig { + data_path: path, + dimension: spec.dimension, + similarity: spec.similarity, + index_type: spec.index_type, + ..self.defaults.clone() + }; + + let db = { + let mut collections = self.collections.write().map_err(|_| ApiError::LockError)?; + + if collections.contains_key(&spec.id) { + return Err(ApiError::CollectionAlreadyExists { + name: spec.id.clone(), + }); + } + + let db = Arc::new(init_api(cfg)?); + collections.insert(spec.id.clone(), Arc::clone(&db)); + db + }; + + let mut specs = self.specs.write().map_err(|_| ApiError::LockError)?; + specs.insert(spec.id.clone(), spec); + + Ok(db) + } + + pub fn list(&self) -> Result> { + let mut specs = { + let specs = self.specs.read().map_err(|_| ApiError::LockError)?; + specs.values().cloned().collect::>() + }; + + // random thing added : sort HashMap for better visiblitiy for the usr + specs.sort_by(|a, b| a.id.cmp(&b.id)); + + Ok(specs) + } + + pub fn delete(&self, name: &str) -> Result<()> { + let key = Self::resolve(name); + + if key == DEFAULT_COLLECTION { + return Err(ApiError::ReservedCollection { + name: key.to_string(), + }); + } + + { + let collections = self.collections.read().map_err(|_| ApiError::LockError)?; + if !collections.contains_key(key) { + return Err(ApiError::CollectionNotFound { + name: key.to_string(), + }); + } + } + + // unregister first - dropping the last Arc releases the storage engine's directory handle so that the cleanup works + { + let mut collections = self.collections.write().map_err(|_| ApiError::LockError)?; + collections.remove(key); + } + { + let mut specs = self.specs.write().map_err(|_| ApiError::LockError)?; + specs.remove(key); + } + + // an ongoing request may still hold this collection open, + // so in this case the directory survives and is cleaned up during a future startup + let path = self.root.join("collections").join(key); + if path.exists() + && let Err(e) = fs::remove_dir_all(&path) + { + tracing::warn!("could not remove collection dir {}: {e}", path.display()); + } + + Ok(()) + } +} + +fn validate_name(name: &str) -> Result<()> { + let valid = !name.is_empty() + && name.len() <= MAX_COLLECTION_NAME_LEN + && name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'); + + if valid { + Ok(()) + } else { + Err(ApiError::InvalidCollectionName { + name: name.to_string(), + }) + } +} diff --git a/crates/api/src/error.rs b/crates/api/src/error.rs index 622e58b..8c98273 100644 --- a/crates/api/src/error.rs +++ b/crates/api/src/error.rs @@ -26,6 +26,18 @@ pub enum ApiError { #[snafu(display("Failed to initialize database: {reason}"))] InitializationFailed { reason: String }, + + #[snafu(display("Collection not found: {name}"))] + CollectionNotFound { name: String }, + + #[snafu(display("Collection '{name}' already exists"))] + CollectionAlreadyExists { name: String }, + + #[snafu(display("Invalid collection name: '{name}'"))] + InvalidCollectionName { name: String }, + + #[snafu(display("Collection '{name}' is reserved and cannot be deleted"))] + ReservedCollection { name: String }, } pub type Result = std::result::Result; diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index 741efb2..66ce78f 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -14,6 +14,7 @@ use storage::{StorageEngine, StorageType, VectorPage, create_storage_engine}; use uuid::Uuid; +pub mod collection; pub mod error; pub use error::{ApiError, Result}; @@ -220,7 +221,7 @@ impl SnapshottableDb for VectorDb { } } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct DbConfig { pub storage_type: StorageType, pub index_type: IndexType, diff --git a/crates/api/src/tests.rs b/crates/api/src/tests.rs index ea9504b..158565a 100644 --- a/crates/api/src/tests.rs +++ b/crates/api/src/tests.rs @@ -1,5 +1,7 @@ use super::*; +use crate::collection::{CollectionManager, CollectionSpec, DEFAULT_COLLECTION}; + // TODO: Add more exhaustive tests use std::sync::Mutex; @@ -449,3 +451,207 @@ fn test_snapshot_engine() { std::mem::drop(db); } } + +fn collection_spec(name: &str) -> CollectionSpec { + CollectionSpec { + id: name.to_string(), + dimension: 4, + similarity: Similarity::Cosine, + index_type: IndexType::Flat, + } +} + +fn create_test_manager() -> (CollectionManager, TempDir) { + let temp_dir = tempdir().unwrap(); + let config = DbConfig { + storage_type: StorageType::RocksDb, + index_type: IndexType::Flat, + data_path: temp_dir.path().to_path_buf(), + dimension: 3, + similarity: Similarity::Cosine, + hnsw_config: HnswConfig::default(), + kd_tree_config: KDTreeConfig::default(), + }; + let manager = CollectionManager::new(temp_dir.path().to_path_buf(), config).unwrap(); + (manager, temp_dir) +} + +#[test] +fn test_collection_default_resolution() { + let (manager, _temp_dir) = create_test_manager(); + + assert!(manager.get(DEFAULT_COLLECTION).is_ok()); + assert!(manager.get("").is_ok()); + + // unspecified name -> "default". + assert_eq!( + Arc::as_ptr(&manager.get("").unwrap()), + Arc::as_ptr(&manager.get(DEFAULT_COLLECTION).unwrap()) + ); + + assert!(matches!( + manager.get("nope"), + Err(ApiError::CollectionNotFound { .. }) + )); +} + +#[test] +fn test_create_and_fetch_collection() { + let (manager, temp_dir) = create_test_manager(); + + let created = manager.create(collection_spec("products")).unwrap(); + + // collection is at /collections/ + assert!( + temp_dir + .path() + .join("collections") + .join("products") + .is_dir() + ); + + assert_eq!( + Arc::as_ptr(&created), + Arc::as_ptr(&manager.get("products").unwrap()) + ); + + // check if it is a different instance from the default collection or not + assert_ne!( + Arc::as_ptr(&created), + Arc::as_ptr(&manager.get(DEFAULT_COLLECTION).unwrap()) + ); + + // check if repeated name results in error or not + assert!(matches!( + manager.create(collection_spec("products")), + Err(ApiError::CollectionAlreadyExists { .. }) + )); +} + +#[test] +fn test_create_rejects_invalid_names() { + let (manager, _temp_dir) = create_test_manager(); + + let long_name = "x".repeat(200); + for bad in ["", "../etc/passwd", "a/b", "with space", long_name.as_str()] { + assert!( + matches!( + manager.create(collection_spec(bad)), + Err(ApiError::InvalidCollectionName { .. }) + ), + "expected {bad:?} to be rejected" + ); + } +} + +fn collection_ids(manager: &CollectionManager) -> Vec { + manager.list().unwrap().into_iter().map(|s| s.id).collect() +} + +#[test] +fn test_list_collections_is_sorted() { + let (manager, _temp_dir) = create_test_manager(); + + assert_eq!(collection_ids(&manager), ["default"].map(String::from)); + + manager.create(collection_spec("products")).unwrap(); + manager.create(collection_spec("analytics")).unwrap(); + + assert_eq!( + collection_ids(&manager), + ["analytics", "default", "products"].map(String::from) + ); + + let products = manager + .list() + .unwrap() + .into_iter() + .find(|s| s.id == "products") + .unwrap(); + assert_eq!(products.dimension, 4); + assert_eq!(products.index_type, IndexType::Flat); +} + +#[test] +fn test_delete_collection() { + let (manager, temp_dir) = create_test_manager(); + + manager.create(collection_spec("products")).unwrap(); + manager.delete("products").unwrap(); + + // no longer reachable + assert!(matches!( + manager.get("products"), + Err(ApiError::CollectionNotFound { .. }) + )); + // no longer listed + assert_eq!(collection_ids(&manager), ["default"].map(String::from)); + assert!( + !temp_dir + .path() + .join("collections") + .join("products") + .exists() + ); +} + +// #[test] +// fn test_delete_default_is_refused() { +// let (manager, _temp_dir) = create_test_manager(); + +// assert!(matches!( +// manager.delete(DEFAULT_COLLECTION), +// Err(ApiError::ReservedCollection { .. }) +// )); + +// // the default collection is not deleted +// assert!(manager.get(DEFAULT_COLLECTION).is_ok()); +// assert!(manager.get("").is_ok()); +// } + +#[test] +fn test_delete_unknown_collection() { + let (manager, _temp_dir) = create_test_manager(); + + assert!(matches!( + manager.delete("nope"), + Err(ApiError::CollectionNotFound { .. }) + )); +} + +#[test] +fn test_create_refuses_existing_directory() { + let dir = tempdir().unwrap(); + let config = || DbConfig { + storage_type: StorageType::RocksDb, + index_type: IndexType::Flat, + data_path: dir.path().to_path_buf(), + dimension: 3, + similarity: Similarity::Cosine, + hnsw_config: HnswConfig::default(), + kd_tree_config: KDTreeConfig::default(), + }; + + // session A: create a collection and put a point in it + { + let manager = CollectionManager::new(dir.path().to_path_buf(), config()).unwrap(); + let products = manager.create(collection_spec("products")).unwrap(); + products + .insert(vec![1.0, 2.0, 3.0, 4.0], test_payload("old")) + .unwrap(); + } // process "ends" - the data directory outlives it + + // session B: with no catalog yet, `products` is not registered on startup... + let manager = CollectionManager::new(dir.path().to_path_buf(), config()).unwrap(); + assert!(matches!( + manager.get("products"), + Err(ApiError::CollectionNotFound { .. }) + )); + + // ...but its directory is on disk, so the name must be refused rather than + // silently reopening the old database and inheriting its points + assert!(matches!( + manager.create(collection_spec("products")), + Err(ApiError::CollectionAlreadyExists { .. }) + )); +} diff --git a/crates/grpc/proto/vector-db.proto b/crates/grpc/proto/vector-db.proto index 0b5dfe1..81fa711 100644 --- a/crates/grpc/proto/vector-db.proto +++ b/crates/grpc/proto/vector-db.proto @@ -8,35 +8,50 @@ message UUID { string value = 1; } +message CollectionID { + string name = 1; +} + service VectorDB { //Insert a vector with a payload and return the assigned PointID rpc InsertVector(InsertVectorRequest) returns (PointID) {} - //Delete a vector by its PointID - rpc DeletePoint(PointID) returns (google.protobuf.Empty) {} + //Delete a vector by its PointID and CollectionID + rpc DeletePoint(PointReference) returns (google.protobuf.Empty) {} - //Get a vector and its payload by PointID - rpc GetPoint(PointID) returns (Point) {} + //Get a vector and its payload by PointID and CollectionID + rpc GetPoint(PointReference) returns (Point) {} //Search for the k nearest vectors to a target vector given a distance function rpc SearchPoints(SearchRequest) returns (SearchResponse) {} rpc InsertVectorsBatch(InsertVectorsBatchRequest) returns (InsertVectorsBatchResponse) {} -rpc SearchPointsBatch(SearchPointsBatchRequest) returns (SearchPointsBatchResponse) {} + rpc SearchPointsBatch(SearchPointsBatchRequest) returns (SearchPointsBatchResponse) {} + + rpc CreateCollection(CollectionInfo) returns (CollectionInfo) {} + rpc GetCollection(CollectionID) returns (CollectionInfo) {} + rpc ListCollections(ListCollectionsRequest) returns (ListCollectionsResponse) {} + rpc DeleteCollection(CollectionID) returns (google.protobuf.Empty) {} } message InsertVectorRequest { DenseVector vector = 1; Payload payload = 2; + CollectionID collection_id = 3; } +message PointReference { + PointID id = 1; + CollectionID collection_id = 2; +} message SearchRequest { DenseVector query_vector = 1; Similarity similarity = 2; uint64 limit = 3; uint64 ef = 4; + CollectionID collection_id = 5; } @@ -90,3 +105,22 @@ message SearchPointsBatchRequest { message SearchPointsBatchResponse { repeated SearchResponse results = 1; } + +enum IndexType { + Flat = 0; + KDTree = 1; + HNSW = 2; +} + +message CollectionInfo { + CollectionID collection_id = 1; + uint64 dimension = 2; + Similarity similarity = 3; + IndexType index_type = 4; +} + +message ListCollectionsRequest {} + +message ListCollectionsResponse { + repeated CollectionInfo collections = 1; +} \ No newline at end of file diff --git a/crates/grpc/src/service.rs b/crates/grpc/src/service.rs index 9225077..0d46959 100644 --- a/crates/grpc/src/service.rs +++ b/crates/grpc/src/service.rs @@ -11,8 +11,10 @@ use tonic::{Request, Response, Status, service::InterceptorLayer, transport::Ser use tracing::{Level, event}; use uuid::Uuid as UuidCrate; use vectordb::{ - DenseVector, InsertVectorRequest, InsertVectorsBatchRequest, InsertVectorsBatchResponse, Point, - PointId, SearchPointsBatchRequest, SearchPointsBatchResponse, SearchRequest, SearchResponse, + CollectionId, CollectionInfo, DenseVector, InsertVectorRequest, InsertVectorsBatchRequest, + InsertVectorsBatchResponse, ListCollectionsRequest, ListCollectionsResponse, Point, PointId, + PointReference, SearchPointsBatchRequest, SearchPointsBatchResponse, SearchRequest, + SearchResponse, vector_db_server::{VectorDb, VectorDbServer}, }; @@ -58,6 +60,11 @@ impl VectorDb for VectorDBService { ContentType::Image => defs::ContentType::Image, }; + let c_id = inner_request + .collection_id + .map(|c| c.name) + .unwrap_or_default(); + let point_id = self.vector_db.insert( dense_vector.unwrap().values, defs::Payload { @@ -75,12 +82,23 @@ impl VectorDb for VectorDBService { })) } - async fn get_point(&self, request: Request) -> Result, Status> { + async fn get_point(&self, request: Request) -> Result, Status> { log_rpc("get_point", self.logging); let inner_request = request.into_inner(); - let point_id = inner_request.id.unwrap_or_default().value; + let point_id = inner_request + .id + .unwrap_or_default() + .id + .unwrap_or_default() + .value; + + let c_id = inner_request + .collection_id + .map(|c| c.name) + .unwrap_or_default(); + let point_opt = self .vector_db .get(UuidCrate::from_str(&point_id).unwrap()) @@ -131,6 +149,11 @@ impl VectorDb for VectorDBService { return Err(Status::invalid_argument("Limit must be greater than zero")); } + let c_id = search_request + .collection_id + .map(|c| c.name) + .unwrap_or_default(); + let result_point_ids = self .vector_db .search(SearchQueryInput { @@ -156,11 +179,23 @@ impl VectorDb for VectorDBService { })) } - async fn delete_point(&self, request: Request) -> Result, Status> { + async fn delete_point(&self, request: Request) -> Result, Status> { log_rpc("delete_point", self.logging); interceptors::require_write_role(&request)?; - let point_id = request.into_inner().id.unwrap_or_default().value; + let delete_request = request.into_inner(); + + let point_id = delete_request + .id + .unwrap_or_default() + .id + .unwrap_or_default() + .value; + + let c_id = delete_request + .collection_id + .map(|c| c.name) + .unwrap_or_default(); match self .vector_db @@ -185,6 +220,8 @@ impl VectorDb for VectorDBService { let req = request.into_inner(); let mut ids = Vec::with_capacity(req.vectors.len()); + // Extract CollectionID for each vector and check if all are same or not + for vec in req.vectors { let payload = vec.payload.map(|p| defs::Payload { content_type: match ContentType::try_from(p.content_type) @@ -196,6 +233,10 @@ impl VectorDb for VectorDBService { content: p.content, }); + let c_id = vec.collection_id.map(|c| c.name).unwrap_or_default(); + + // TODO : Check for same c_id throughout the all vectors + let id = self .vector_db .insert( @@ -226,6 +267,8 @@ impl VectorDb for VectorDBService { .get(query.similarity as usize) .ok_or(tonic::Status::invalid_argument("Invalid similarity"))?; + let c_id = query.collection_id.map(|c| c.name).unwrap_or_default(); + let ids = self .vector_db .search(SearchQueryInput { @@ -253,6 +296,34 @@ impl VectorDb for VectorDBService { Ok(tonic::Response::new(SearchPointsBatchResponse { results })) } + + async fn create_collection( + &self, + request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("not yet implemented")) + } + + async fn get_collection( + &self, + request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented yet")) + } + + async fn list_collections( + &self, + request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented yet")) + } + + async fn delete_collection( + &self, + request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented yet")) + } } pub async fn run_server( vector_db_service: VectorDBService, diff --git a/crates/grpc/src/tests.rs b/crates/grpc/src/tests.rs index e623e13..621f257 100644 --- a/crates/grpc/src/tests.rs +++ b/crates/grpc/src/tests.rs @@ -1,7 +1,7 @@ use crate::constants::AUTHORIZATION_HEADER_KEY; use crate::service::vectordb::vector_db_client::VectorDbClient; use crate::service::vectordb::{ - ContentType, DenseVector, InsertVectorRequest, Payload, PointId, SearchRequest, + ContentType, DenseVector, InsertVectorRequest, Payload, PointId, PointReference, SearchRequest, }; use crate::service::{VectorDBService, run_server}; use crate::utils::ServerEndpoint; @@ -92,6 +92,7 @@ async fn test_grpc_server_start() { values: test_vec.clone(), }), payload: Some(Payload::default()), + collection_id: None, }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); @@ -115,6 +116,7 @@ async fn test_insert_vector_rpc() { content_type: ContentType::Text as i32, content: "test".to_string(), }), + collection_id: None, }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); @@ -124,8 +126,10 @@ async fn test_insert_vector_rpc() { assert!(resp.is_ok()); // check if the vector is actually present in the database - let mut request = tonic::Request::new(PointId { - id: resp.unwrap().into_inner().id, + let inserted = resp.unwrap().into_inner(); + let mut request = tonic::Request::new(PointReference { + id: Some(inserted), + collection_id: None, }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); let resp = client.get_point(request).await; @@ -146,6 +150,7 @@ async fn test_insert_vector_rpc() { values: vec![1.0, 2.0], }), payload: Some(Payload::default()), + collection_id: None, }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); @@ -167,6 +172,7 @@ async fn test_delete_vector_rpc() { values: test_vec.clone(), }), payload: Some(Payload::default()), + collection_id: None, }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); @@ -177,8 +183,9 @@ async fn test_delete_vector_rpc() { let point = resp.unwrap().into_inner(); // delete the vector - let mut request = tonic::Request::new(PointId { - id: point.id.clone(), + let mut request = tonic::Request::new(PointReference { + id: Some(point.clone()), + collection_id: None, }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); @@ -188,7 +195,10 @@ async fn test_delete_vector_rpc() { assert!(resp.is_ok()); // verify that the vector is deleted - let mut request = tonic::Request::new(PointId { id: point.id }); + let mut request = tonic::Request::new(PointReference { + id: Some(point), + collection_id: None, + }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); let resp = client.get_point(request).await; @@ -209,6 +219,7 @@ async fn test_search_vector_rpc() { values: test_vec.clone(), }), payload: Some(Payload::default()), + collection_id: None, }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); @@ -228,6 +239,7 @@ async fn test_search_vector_rpc() { similarity: 0, // euclidean distance limit: 1, ef: 0, + collection_id: None, }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); @@ -256,6 +268,7 @@ async fn test_unauthorized_rpc() { values: test_vec.clone(), }), payload: Some(Payload::default()), + collection_id: None, }); append_test_auth_header(&mut request, "43121"); @@ -275,6 +288,7 @@ async fn test_readonly_key_cannot_write() { values: vec![1.0, 2.0, 3.0], }), payload: Some(Payload::default()), + collection_id: None, }); append_test_auth_header(&mut request, TEST_READONLY_BEARER_TOKEN); @@ -295,6 +309,7 @@ async fn test_readonly_key_can_search() { similarity: 0, limit: 1, ef: 0, + collection_id: None, }); append_test_auth_header(&mut request, TEST_READONLY_BEARER_TOKEN);