herocode/db
13
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

...

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

74
herodb/src/db/db.rs
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)
}
}

2
herodb/src/db/mod.rs
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;

2
herodb/src/db/model.rs
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;

127
herodb/src/db/tst_index.rs
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)]

20
herodb/src/models/biz/customer.rs
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
}
}