//! Collection manager for multi-collection operations use dashmap::DashMap; use parking_lot::RwLock; use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; use crate::collection::{Collection, CollectionConfig, CollectionStats}; use crate::error::{CollectionError, Result}; /// Metadata for persisting collections #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] struct CollectionMetadata { name: String, config: CollectionConfig, created_at: i64, updated_at: i64, } /// Manages multiple vector collections with alias support #[derive(Debug)] pub struct CollectionManager { /// Active collections collections: DashMap>>, /// Alias mappings (alias -> collection_name) aliases: DashMap, /// Base path for storing collections base_path: PathBuf, } impl CollectionManager { /// Create a new collection manager /// /// # Arguments /// /// * `base_path` - Directory where collections will be stored /// /// # Example /// /// ```no_run /// use ruvector_collections::CollectionManager; /// use std::path::PathBuf; /// /// let manager = CollectionManager::new(PathBuf::from("./collections")).unwrap(); /// ``` pub fn new(base_path: PathBuf) -> Result { // Create base directory if it doesn't exist std::fs::create_dir_all(&base_path)?; let manager = Self { collections: DashMap::new(), aliases: DashMap::new(), base_path, }; // Load existing collections manager.load_collections()?; Ok(manager) } /// Create a new collection /// /// # Arguments /// /// * `name` - Collection name (must be unique) /// * `config` - Collection configuration /// /// # Errors /// /// Returns `CollectionAlreadyExists` if a collection with the same name exists pub fn create_collection(&self, name: &str, config: CollectionConfig) -> Result<()> { // Validate collection name Self::validate_name(name)?; // Check if collection already exists if self.collections.contains_key(name) { return Err(CollectionError::CollectionAlreadyExists { name: name.to_string(), }); } // Check if an alias with this name exists if self.aliases.contains_key(name) { return Err(CollectionError::InvalidName { name: name.to_string(), reason: "An alias with this name already exists".to_string(), }); } // Create storage path for this collection let storage_path = self.base_path.join(name); std::fs::create_dir_all(&storage_path)?; let db_path = storage_path .join("vectors.db") .to_string_lossy() .to_string(); // Create collection let collection = Collection::new(name.to_string(), config, db_path)?; // Save metadata self.save_collection_metadata(&collection)?; // Add to collections map self.collections .insert(name.to_string(), Arc::new(RwLock::new(collection))); Ok(()) } /// Delete a collection /// /// # Arguments /// /// * `name` - Collection name to delete /// /// # Errors /// /// Returns `CollectionNotFound` if collection doesn't exist /// Returns `CollectionHasAliases` if collection has active aliases pub fn delete_collection(&self, name: &str) -> Result<()> { // Check if collection exists if !self.collections.contains_key(name) { return Err(CollectionError::CollectionNotFound { name: name.to_string(), }); } // Check for active aliases let active_aliases: Vec = self .aliases .iter() .filter(|entry| entry.value() == name) .map(|entry| entry.key().clone()) .collect(); if !active_aliases.is_empty() { return Err(CollectionError::CollectionHasAliases { collection: name.to_string(), aliases: active_aliases, }); } // Remove from collections map self.collections.remove(name); // Delete from disk let collection_path = self.base_path.join(name); if collection_path.exists() { std::fs::remove_dir_all(&collection_path)?; } Ok(()) } /// Get a collection by name or alias /// /// # Arguments /// /// * `name` - Collection name or alias pub fn get_collection(&self, name: &str) -> Option>> { // Try to resolve as alias first let collection_name = self.resolve_alias(name).unwrap_or_else(|| name.to_string()); self.collections .get(&collection_name) .map(|entry| entry.value().clone()) } /// List all collection names pub fn list_collections(&self) -> Vec { self.collections .iter() .map(|entry| entry.key().clone()) .collect() } /// Check if a collection exists /// /// # Arguments /// /// * `name` - Collection name (not alias) pub fn collection_exists(&self, name: &str) -> bool { self.collections.contains_key(name) } /// Get statistics for a collection pub fn collection_stats(&self, name: &str) -> Result { let collection = self.get_collection(name) .ok_or_else(|| CollectionError::CollectionNotFound { name: name.to_string(), })?; let guard = collection.read(); guard.stats() } // ===== Alias Management ===== /// Create an alias for a collection /// /// # Arguments /// /// * `alias` - Alias name (must be unique) /// * `collection` - Target collection name /// /// # Errors /// /// Returns `AliasAlreadyExists` if alias already exists /// Returns `CollectionNotFound` if target collection doesn't exist pub fn create_alias(&self, alias: &str, collection: &str) -> Result<()> { // Validate alias name Self::validate_name(alias)?; // Check if alias already exists if self.aliases.contains_key(alias) { return Err(CollectionError::AliasAlreadyExists { alias: alias.to_string(), }); } // Check if a collection with this name exists if self.collections.contains_key(alias) { return Err(CollectionError::InvalidName { name: alias.to_string(), reason: "A collection with this name already exists".to_string(), }); } // Verify target collection exists if !self.collections.contains_key(collection) { return Err(CollectionError::CollectionNotFound { name: collection.to_string(), }); } // Create alias self.aliases .insert(alias.to_string(), collection.to_string()); // Save aliases self.save_aliases()?; Ok(()) } /// Delete an alias /// /// # Arguments /// /// * `alias` - Alias name to delete /// /// # Errors /// /// Returns `AliasNotFound` if alias doesn't exist pub fn delete_alias(&self, alias: &str) -> Result<()> { if self.aliases.remove(alias).is_none() { return Err(CollectionError::AliasNotFound { alias: alias.to_string(), }); } // Save aliases self.save_aliases()?; Ok(()) } /// Switch an alias to point to a different collection /// /// # Arguments /// /// * `alias` - Alias name /// * `new_collection` - New target collection name /// /// # Errors /// /// Returns `AliasNotFound` if alias doesn't exist /// Returns `CollectionNotFound` if new collection doesn't exist pub fn switch_alias(&self, alias: &str, new_collection: &str) -> Result<()> { // Verify alias exists if !self.aliases.contains_key(alias) { return Err(CollectionError::AliasNotFound { alias: alias.to_string(), }); } // Verify new collection exists if !self.collections.contains_key(new_collection) { return Err(CollectionError::CollectionNotFound { name: new_collection.to_string(), }); } // Update alias self.aliases .insert(alias.to_string(), new_collection.to_string()); // Save aliases self.save_aliases()?; Ok(()) } /// Resolve an alias to a collection name /// /// # Arguments /// /// * `name_or_alias` - Collection name or alias /// /// # Returns /// /// `Some(collection_name)` if it's an alias, `None` if it's not an alias pub fn resolve_alias(&self, name_or_alias: &str) -> Option { self.aliases .get(name_or_alias) .map(|entry| entry.value().clone()) } /// List all aliases with their target collections pub fn list_aliases(&self) -> Vec<(String, String)> { self.aliases .iter() .map(|entry| (entry.key().clone(), entry.value().clone())) .collect() } /// Check if a name is an alias pub fn is_alias(&self, name: &str) -> bool { self.aliases.contains_key(name) } // ===== Internal Methods ===== /// Validate a collection or alias name fn validate_name(name: &str) -> Result<()> { if name.is_empty() { return Err(CollectionError::InvalidName { name: name.to_string(), reason: "Name cannot be empty".to_string(), }); } if name.len() > 255 { return Err(CollectionError::InvalidName { name: name.to_string(), reason: "Name too long (max 255 characters)".to_string(), }); } // Only allow alphanumeric, hyphens, underscores if !name .chars() .all(|c| c.is_alphanumeric() || c == '-' || c == '_') { return Err(CollectionError::InvalidName { name: name.to_string(), reason: "Name can only contain letters, numbers, hyphens, and underscores" .to_string(), }); } Ok(()) } /// Load existing collections from disk fn load_collections(&self) -> Result<()> { if !self.base_path.exists() { return Ok(()); } // Load aliases self.load_aliases()?; // Scan for collection directories for entry in std::fs::read_dir(&self.base_path)? { let entry = entry?; let path = entry.path(); if path.is_dir() { let name = path .file_name() .and_then(|n| n.to_str()) .unwrap_or("") .to_string(); // Skip special directories if name.starts_with('.') || name == "aliases.json" { continue; } // Try to load collection metadata if let Ok(metadata) = self.load_collection_metadata(&name) { let db_path = path.join("vectors.db").to_string_lossy().to_string(); // Recreate collection if let Ok(mut collection) = Collection::new(metadata.name.clone(), metadata.config, db_path) { collection.created_at = metadata.created_at; collection.updated_at = metadata.updated_at; self.collections .insert(name.clone(), Arc::new(RwLock::new(collection))); } } } } Ok(()) } /// Save collection metadata to disk fn save_collection_metadata(&self, collection: &Collection) -> Result<()> { let metadata = CollectionMetadata { name: collection.name.clone(), config: collection.config.clone(), created_at: collection.created_at, updated_at: collection.updated_at, }; let metadata_path = self.base_path.join(&collection.name).join("metadata.json"); let json = serde_json::to_string_pretty(&metadata)?; std::fs::write(metadata_path, json)?; Ok(()) } /// Load collection metadata from disk fn load_collection_metadata(&self, name: &str) -> Result { let metadata_path = self.base_path.join(name).join("metadata.json"); let json = std::fs::read_to_string(metadata_path)?; let metadata: CollectionMetadata = serde_json::from_str(&json)?; Ok(metadata) } /// Save aliases to disk fn save_aliases(&self) -> Result<()> { let aliases: HashMap = self .aliases .iter() .map(|entry| (entry.key().clone(), entry.value().clone())) .collect(); let aliases_path = self.base_path.join("aliases.json"); let json = serde_json::to_string_pretty(&aliases)?; std::fs::write(aliases_path, json)?; Ok(()) } /// Load aliases from disk fn load_aliases(&self) -> Result<()> { let aliases_path = self.base_path.join("aliases.json"); if !aliases_path.exists() { return Ok(()); } let json = std::fs::read_to_string(aliases_path)?; let aliases: HashMap = serde_json::from_str(&json)?; for (alias, collection) in aliases { self.aliases.insert(alias, collection); } Ok(()) } } #[cfg(test)] mod tests { use super::*; #[test] fn test_validate_name() { assert!(CollectionManager::validate_name("valid-name_123").is_ok()); assert!(CollectionManager::validate_name("").is_err()); assert!(CollectionManager::validate_name("invalid name").is_err()); assert!(CollectionManager::validate_name("invalid/name").is_err()); } #[test] fn test_collection_manager() -> Result<()> { let temp_dir = std::env::temp_dir().join("ruvector_test_collections"); let _ = std::fs::remove_dir_all(&temp_dir); let manager = CollectionManager::new(temp_dir.clone())?; // Create collection let config = CollectionConfig::with_dimensions(128); manager.create_collection("test", config)?; assert!(manager.collection_exists("test")); assert_eq!(manager.list_collections().len(), 1); // Create alias manager.create_alias("test_alias", "test")?; assert!(manager.is_alias("test_alias")); assert_eq!( manager.resolve_alias("test_alias"), Some("test".to_string()) ); // Get collection by alias assert!(manager.get_collection("test_alias").is_some()); // Cleanup manager.delete_alias("test_alias")?; manager.delete_collection("test")?; let _ = std::fs::remove_dir_all(&temp_dir); Ok(()) } }