This commit is contained in:
despiegk 2025-04-21 06:26:56 +02:00
parent 3b5f9c6012
commit 6757a62fec
7 changed files with 299 additions and 23 deletions

View File

@ -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"

View File

@ -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<dyn std::error::Error>> {
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::<Customer>()
.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::<Customer>()?;
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::<Customer>("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::<Customer>("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::<Customer>(2)?;
// List all customers again
println!("\nListing all customers after deletion:");
let customers = db.list::<Customer>()?;
for customer in &customers {
println!("- {} (ID: {})", customer.name, customer.id);
}
println!("Total: {} customers", customers.len());
println!("\nExample completed successfully!");
Ok(())
}

View File

@ -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::<T>(),
@ -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<T: Model>(&self, id: u32) -> DbResult<()> {
// First, get the model to extract its index keys
let model = self.get::<T>(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<T: Model>(&self, index_name: &str, index_value: &str) -> DbResult<Vec<T>> {
// 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::<T>(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<T: Model>(&self, index_name: &str, index_value_prefix: &str) -> DbResult<Vec<T>> {
// 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::<T>(id) {
Ok(model) => result.push(model),
Err(DbError::NotFound(_)) => continue, // Skip if not found
Err(e) => return Err(e),
}
}
Ok(result)
}
}

View File

@ -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;

View File

@ -1,4 +1,4 @@
...use crate::db::error::{DbError, DbResult};
use crate::db::error::{DbError, DbResult};
use serde::{Deserialize, Serialize};
use std::fmt::Debug;

View File

@ -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<u8>) -> 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<u8>, 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<Vec<(u32, Vec<u8>)>> {
// 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::<u32>().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<Vec<u32>> {
// 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<Vec<u32>> {
// 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)]

View File

@ -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<IndexKey> {
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
}
}