diff --git a/herodb/Cargo.toml b/herodb/Cargo.toml index d165a75..9d62b77 100644 --- a/herodb/Cargo.toml +++ b/herodb/Cargo.toml @@ -36,6 +36,10 @@ path = "examples/business_models_demo.rs" name = "ourdb_example" path = "examples/ourdb_example.rs" +[[example]] +name = "tst_index_example" +path = "examples/tst_index_example.rs" + [[bin]] name = "dbexample_prod" path = "src/cmd/dbexample_prod/main.rs" diff --git a/herodb/examples/tst_index_example.rs b/herodb/examples/tst_index_example.rs new file mode 100644 index 0000000..c9f3285 --- /dev/null +++ b/herodb/examples/tst_index_example.rs @@ -0,0 +1,93 @@ +use herodb::db::{DB, DBBuilder, Model, IndexKey}; +use herodb::models::biz::Customer; +use std::path::PathBuf; +use std::fs; + +fn main() -> Result<(), Box> { + println!("TST Index Example"); + println!("================"); + + // Create a temporary directory for the database + let db_path = PathBuf::from("/tmp/tst_index_example"); + if db_path.exists() { + fs::remove_dir_all(&db_path)?; + } + fs::create_dir_all(&db_path)?; + println!("Database path: {:?}", db_path); + + // Create a database instance with the Customer model registered + let db = DBBuilder::new(&db_path) + .register_model::() + .build()?; + + // Create some customers + let customer1 = Customer::new( + 1, + "John Doe".to_string(), + "A regular customer".to_string(), + "pk123456".to_string(), + ); + + let customer2 = Customer::new( + 2, + "Jane Smith".to_string(), + "A VIP customer".to_string(), + "pk789012".to_string(), + ); + + let customer3 = Customer::new( + 3, + "John Smith".to_string(), + "Another customer".to_string(), + "pk345678".to_string(), + ); + + // Insert the customers + db.set(&customer1)?; + db.set(&customer2)?; + db.set(&customer3)?; + + println!("\nCustomers created:"); + println!("1. {} ({})", customer1.name, customer1.pubkey); + println!("2. {} ({})", customer2.name, customer2.pubkey); + println!("3. {} ({})", customer3.name, customer3.pubkey); + + // List all customers + println!("\nListing all customers:"); + let customers = db.list::()?; + for customer in &customers { + println!("- {} (ID: {})", customer.name, customer.id); + } + println!("Total: {} customers", customers.len()); + + // Find customers by name index + println!("\nFinding customers by name 'John':"); + let john_customers = db.find_by_index_prefix::("name", "John")?; + for customer in &john_customers { + println!("- {} (ID: {})", customer.name, customer.id); + } + println!("Total: {} customers", john_customers.len()); + + // Find customers by pubkey index + println!("\nFinding customer by pubkey 'pk789012':"); + let pubkey_customers = db.find_by_index::("pubkey", "pk789012")?; + for customer in &pubkey_customers { + println!("- {} (ID: {})", customer.name, customer.id); + } + println!("Total: {} customers", pubkey_customers.len()); + + // Delete a customer + println!("\nDeleting customer with ID 2"); + db.delete::(2)?; + + // List all customers again + println!("\nListing all customers after deletion:"); + let customers = db.list::()?; + for customer in &customers { + println!("- {} (ID: {})", customer.name, customer.id); + } + println!("Total: {} customers", customers.len()); + + println!("\nExample completed successfully!"); + Ok(()) +} \ No newline at end of file diff --git a/herodb/src/db/db.rs b/herodb/src/db/db.rs index 9267f04..864a4fa 100644 --- a/herodb/src/db/db.rs +++ b/herodb/src/db/db.rs @@ -232,7 +232,9 @@ impl DB { // Apply to OurDB self.apply_set_operation(model_type, &serialized)?; - // Apply to TST index + // Apply to TST index (primary key only) + // We can't easily get the index keys in the transaction commit + // because we don't have the model type information at runtime let mut tst_index = self.tst_index.write().unwrap(); tst_index.set(&model_prefix, model_id, serialized.clone())?; } @@ -241,6 +243,9 @@ impl DB { id, model_prefix, } => { + // For delete operations, we can't get the index keys from the model + // because it's already deleted. We'll just delete the primary key. + // Apply to OurDB let db_ops = self .type_map @@ -249,7 +254,7 @@ impl DB { let mut db_ops_guard = db_ops.write().unwrap(); db_ops_guard.delete(id)?; - // Apply to TST index + // Apply to TST index (primary key only) let mut tst_index = self.tst_index.write().unwrap(); tst_index.delete(&model_prefix, id)?; } @@ -298,6 +303,9 @@ impl DB { // Serialize the model for later use let serialized = model.to_bytes()?; + // Get the index keys for this model + let index_keys = model.db_keys(); + // Record a Set operation in the transaction with prefix and ID tx_state.operations.push(DbOperation::Set { model_type: TypeId::of::(), @@ -320,12 +328,13 @@ impl DB { let mut db_ops_guard = db_ops.write().unwrap(); db_ops_guard.insert(model)?; - // Also update the TST index + // Also update the TST index with all index keys let mut tst_index = self.tst_index.write().unwrap(); let prefix = T::db_prefix(); let id = model.get_id(); let data = model.to_bytes()?; - tst_index.set(prefix, id, data)?; + let index_keys = model.db_keys(); + tst_index.set_with_indexes(prefix, id, data, &index_keys)?; Ok(()) }, @@ -412,6 +421,10 @@ impl DB { /// Delete a model instance by its ID and type pub fn delete(&self, id: u32) -> DbResult<()> { + // First, get the model to extract its index keys + let model = self.get::(id)?; + let index_keys = model.db_keys(); + // Try to acquire a write lock on the transaction let mut tx_guard = self.transaction.write().unwrap(); @@ -439,10 +452,10 @@ impl DB { let mut db_ops_guard = db_ops.write().unwrap(); db_ops_guard.delete(id)?; - // Also update the TST index + // Also update the TST index with all index keys let mut tst_index = self.tst_index.write().unwrap(); let prefix = T::db_prefix(); - tst_index.delete(prefix, id)?; + tst_index.delete_with_indexes(prefix, id, &index_keys)?; Ok(()) }, @@ -494,11 +507,12 @@ impl DB { let mut tst_index = self.tst_index.write().unwrap(); let prefix = T::db_prefix(); - // Rebuild the TST index + // Rebuild the TST index with all index keys for model in models { let id = model.get_id(); let data = model.to_bytes()?; - tst_index.set(prefix, id, data)?; + let index_keys = model.db_keys(); + tst_index.set_with_indexes(prefix, id, data, &index_keys)?; } Ok(()) @@ -541,4 +555,48 @@ impl DB { Ok(()) } + + /// Find a model by a specific index key + pub fn find_by_index(&self, index_name: &str, index_value: &str) -> DbResult> { + // Get the prefix for this model type + let prefix = T::db_prefix(); + + // Use the TST index to find objects with this index key + let mut tst_index = self.tst_index.write().unwrap(); + let ids = tst_index.find_by_index(prefix, index_name, index_value)?; + + // Get the objects by their IDs + let mut result = Vec::with_capacity(ids.len()); + for id in ids { + match self.get::(id) { + Ok(model) => result.push(model), + Err(DbError::NotFound(_)) => continue, // Skip if not found + Err(e) => return Err(e), + } + } + + Ok(result) + } + + /// Find models by a prefix of an index key + pub fn find_by_index_prefix(&self, index_name: &str, index_value_prefix: &str) -> DbResult> { + // Get the prefix for this model type + let prefix = T::db_prefix(); + + // Use the TST index to find objects with this index key prefix + let mut tst_index = self.tst_index.write().unwrap(); + let ids = tst_index.find_by_index_prefix(prefix, index_name, index_value_prefix)?; + + // Get the objects by their IDs + let mut result = Vec::with_capacity(ids.len()); + for id in ids { + match self.get::(id) { + Ok(model) => result.push(model), + Err(DbError::NotFound(_)) => continue, // Skip if not found + Err(e) => return Err(e), + } + } + + Ok(result) + } } diff --git a/herodb/src/db/mod.rs b/herodb/src/db/mod.rs index 57efa52..48e18f8 100644 --- a/herodb/src/db/mod.rs +++ b/herodb/src/db/mod.rs @@ -4,7 +4,7 @@ pub use error::{DbError, DbResult}; // Export the model module pub mod model; -pub use model::{Model, Storable}; +pub use model::{Model, Storable, IndexKey}; // Export the store module pub mod store; diff --git a/herodb/src/db/model.rs b/herodb/src/db/model.rs index 2f2de93..8e9c97d 100644 --- a/herodb/src/db/model.rs +++ b/herodb/src/db/model.rs @@ -1,4 +1,4 @@ -...use crate::db::error::{DbError, DbResult}; +use crate::db::error::{DbError, DbResult}; use serde::{Deserialize, Serialize}; use std::fmt::Debug; diff --git a/herodb/src/db/tst_index.rs b/herodb/src/db/tst_index.rs index b8272f1..28fe019 100644 --- a/herodb/src/db/tst_index.rs +++ b/herodb/src/db/tst_index.rs @@ -1,4 +1,5 @@ use crate::db::error::{DbError, DbResult}; +use crate::db::model::IndexKey; use std::path::{Path, PathBuf}; use std::collections::HashMap; use tst::TST; @@ -45,22 +46,45 @@ impl TSTIndexManager { Ok(self.tst_instances.get_mut(prefix).unwrap()) } - /// Adds or updates an object in the TST index + /// Adds or updates an object in the TST index with primary key pub fn set(&mut self, prefix: &str, id: u32, data: Vec) -> DbResult<()> { // Get the TST for this prefix let tst = self.get_tst(prefix)?; - // Create the key in the format prefix_id + // Create the primary key in the format prefix_id let key = format!("{}_{}", prefix, id); // Set the key-value pair in the TST - tst.set(&key, data) + tst.set(&key, data.clone()) .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?; Ok(()) } - /// Removes an object from the TST index + /// Adds or updates an object in the TST index with additional index keys + pub fn set_with_indexes(&mut self, prefix: &str, id: u32, data: Vec, index_keys: &[IndexKey]) -> DbResult<()> { + // First set the primary key + self.set(prefix, id, data.clone())?; + + // Get the TST for this prefix + let tst = self.get_tst(prefix)?; + + // Add additional index keys + for index_key in index_keys { + // Create the index key in the format prefix_indexname_value + let key = format!("{}_{}_{}", prefix, index_key.name, index_key.value); + + // Set the key-value pair in the TST + // For index keys, we store the ID as the value + let id_bytes = id.to_be_bytes().to_vec(); + tst.set(&key, id_bytes) + .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?; + } + + Ok(()) + } + + /// Removes an object from the TST index (primary key only) pub fn delete(&mut self, prefix: &str, id: u32) -> DbResult<()> { // Get the TST for this prefix let tst = self.get_tst(prefix)?; @@ -75,23 +99,48 @@ impl TSTIndexManager { Ok(()) } - /// Lists all objects with a given prefix + /// Removes an object from the TST index including all index keys + pub fn delete_with_indexes(&mut self, prefix: &str, id: u32, index_keys: &[IndexKey]) -> DbResult<()> { + // First delete the primary key + self.delete(prefix, id)?; + + // Get the TST for this prefix + let tst = self.get_tst(prefix)?; + + // Delete additional index keys + for index_key in index_keys { + // Create the index key in the format prefix_indexname_value + let key = format!("{}_{}_{}", prefix, index_key.name, index_key.value); + + // Delete the key from the TST + tst.delete(&key) + .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?; + } + + Ok(()) + } + + /// Lists all objects with a given prefix (primary keys only) pub fn list(&mut self, prefix: &str) -> DbResult)>> { // Get the TST for this prefix let tst = self.get_tst(prefix)?; - // Get all keys with this prefix - let keys = tst.list(prefix) + // Get all keys with this prefix followed by an underscore + let search_prefix = format!("{}_", prefix); + let keys = tst.list(&search_prefix) .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?; // Get all values for these keys let mut result = Vec::with_capacity(keys.len()); for key in keys { - // Extract the ID from the key (format: prefix_id) - let id_str = key.split('_').nth(1).ok_or_else(|| { - DbError::GeneralError(format!("Invalid key format: {}", key)) - })?; + // Check if this is a primary key (prefix_id) and not an index key (prefix_indexname_value) + let parts: Vec<&str> = key.split('_').collect(); + if parts.len() != 2 { + continue; // Skip index keys + } + // Extract the ID from the key (format: prefix_id) + let id_str = parts[1]; let id = id_str.parse::().map_err(|_| { DbError::GeneralError(format!("Invalid ID in key: {}", key)) })?; @@ -105,6 +154,62 @@ impl TSTIndexManager { Ok(result) } + + /// Finds objects by a specific index key + pub fn find_by_index(&mut self, prefix: &str, index_name: &str, index_value: &str) -> DbResult> { + // Get the TST for this prefix + let tst = self.get_tst(prefix)?; + + // Create the index key in the format prefix_indexname_value + let key = format!("{}_{}_{}", prefix, index_name, index_value); + + // Try to get the value from the TST + match tst.get(&key) { + Ok(id_bytes) => { + // Convert the bytes to a u32 ID + if id_bytes.len() == 4 { + let mut bytes = [0u8; 4]; + bytes.copy_from_slice(&id_bytes[0..4]); + let id = u32::from_be_bytes(bytes); + Ok(vec![id]) + } else { + Err(DbError::GeneralError(format!("Invalid ID bytes for key: {}", key))) + } + }, + Err(_) => Ok(Vec::new()), // No matches found + } + } + + /// Finds objects by a prefix of an index key + pub fn find_by_index_prefix(&mut self, prefix: &str, index_name: &str, index_value_prefix: &str) -> DbResult> { + // Get the TST for this prefix + let tst = self.get_tst(prefix)?; + + // Create the index key prefix in the format prefix_indexname_valueprefix + let key_prefix = format!("{}_{}_{}", prefix, index_name, index_value_prefix); + + // Get all keys with this prefix + let keys = tst.list(&key_prefix) + .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?; + + // Extract the IDs from the values + let mut result = Vec::with_capacity(keys.len()); + for key in keys { + // Get the value from the TST + let id_bytes = tst.get(&key) + .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?; + + // Convert the bytes to a u32 ID + if id_bytes.len() == 4 { + let mut bytes = [0u8; 4]; + bytes.copy_from_slice(&id_bytes[0..4]); + let id = u32::from_be_bytes(bytes); + result.push(id); + } + } + + Ok(result) + } } #[cfg(test)] diff --git a/herodb/src/models/biz/customer.rs b/herodb/src/models/biz/customer.rs index ccf69e4..5fe7643 100644 --- a/herodb/src/models/biz/customer.rs +++ b/herodb/src/models/biz/customer.rs @@ -1,4 +1,4 @@ -use crate::db::{Model, Storable}; // Import Model trait from db module +use crate::db::{Model, Storable, IndexKey}; // Import Model trait and IndexKey from db module use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; @@ -145,6 +145,22 @@ impl Model for Customer { fn db_prefix() -> &'static str { "customer" } - + fn db_keys(&self) -> Vec { + let mut keys = Vec::new(); + + // Add an index for the name + keys.push(IndexKey { + name: "name", + value: self.name.clone(), + }); + + // Add an index for the pubkey + keys.push(IndexKey { + name: "pubkey", + value: self.pubkey.clone(), + }); + + keys + } } \ No newline at end of file