use redb::{Database, ReadableDatabase, ReadableTableMetadata, TableDefinition};
use std::path::Path;
use std::sync::Arc;
pub const MEMORY_LOG: TableDefinition<u128, &[u8]> = TableDefinition::new("memory_log");
pub const MEMORY_KV: TableDefinition<u128, &[u8]> = TableDefinition::new("memory_kv");
pub const ENTITY_INDEX: TableDefinition<u64, &[u8]> = TableDefinition::new("entity_index");
pub const TOPIC_INDEX: TableDefinition<u64, &[u8]> = TableDefinition::new("topic_index");
pub const GOAL_INDEX: TableDefinition<u64, &[u8]> = TableDefinition::new("goal_index");
pub const EVENT_INDEX: TableDefinition<u64, &[u8]> = TableDefinition::new("event_index");
pub const TEMPORAL_INDEX: TableDefinition<u32, &[u8]> = TableDefinition::new("temporal_index");
pub const CAUSAL_INDEX: TableDefinition<u64, &[u8]> = TableDefinition::new("causal_index");
pub const LINK_OVERLAY: TableDefinition<u128, &[u8]> = TableDefinition::new("link_overlay");
pub const SUMMARY_OVERLAY: TableDefinition<u128, &[u8]> = TableDefinition::new("summary_overlay");
pub const CORRECTION_OVERLAY: TableDefinition<u128, &[u8]> =
TableDefinition::new("correction_overlay");
pub const ACTIVATION_LOG: TableDefinition<u128, &[u8]> = TableDefinition::new("activation_log");
pub const CONSOLIDATION_QUEUE: TableDefinition<u128, &[u8]> =
TableDefinition::new("consolidation_queue");
const U128_TABLES: &[TableDefinition<u128, &[u8]>] = &[
MEMORY_LOG,
MEMORY_KV,
LINK_OVERLAY,
SUMMARY_OVERLAY,
CORRECTION_OVERLAY,
ACTIVATION_LOG,
CONSOLIDATION_QUEUE,
];
const U64_TABLES: &[TableDefinition<u64, &[u8]>] = &[
ENTITY_INDEX,
TOPIC_INDEX,
GOAL_INDEX,
EVENT_INDEX,
CAUSAL_INDEX,
];
const U32_TABLES: &[TableDefinition<u32, &[u8]>] = &[TEMPORAL_INDEX];
pub trait Store: Sized {
fn open(path: impl AsRef<Path>) -> Result<Self, StoreError>;
fn close(self) -> Result<(), StoreError>;
}
pub struct RedbStore {
db: Arc<Database>,
}
#[derive(Debug, thiserror::Error)]
pub enum StoreError {
#[error("Database error: {0}")]
Database(#[from] redb::DatabaseError),
#[error("Transaction error: {0}")]
Transaction(#[from] redb::TransactionError),
#[error("Table error: {0}")]
Table(#[from] redb::TableError),
#[error("Commit error: {0}")]
Commit(#[from] redb::CommitError),
#[error("Storage error: {0}")]
Storage(#[from] redb::StorageError),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("Record already exists, cannot be overwritten: id={0}")]
RecordExists(u128),
}
pub type StoreResult<T> = Result<T, StoreError>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoreStats {
pub memory_log_count: u64,
pub memory_kv_count: u64,
pub entity_index_size: u64,
pub topic_index_size: u64,
pub goal_index_size: u64,
pub event_index_size: u64,
pub temporal_index_size: u64,
pub causal_index_size: u64,
}
impl Store for RedbStore {
fn open(path: impl AsRef<Path>) -> StoreResult<Self> {
let path = path.as_ref();
if let Some(parent) = path.parent() {
if !parent.exists() {
std::fs::create_dir_all(parent)?;
}
}
let db = Database::create(path)?;
{
let txn = db.begin_write()?;
for def in U128_TABLES {
let _ = txn.open_table(*def)?;
}
for def in U64_TABLES {
let _ = txn.open_table(*def)?;
}
for def in U32_TABLES {
let _ = txn.open_table(*def)?;
}
txn.commit()?;
}
Ok(Self { db: Arc::new(db) })
}
fn close(self) -> StoreResult<()> {
drop(self.db);
Ok(())
}
}
impl RedbStore {
pub fn db(&self) -> &Database {
&self.db
}
pub fn db_arc(&self) -> Arc<Database> {
Arc::clone(&self.db)
}
pub fn stats(&self) -> StoreResult<StoreStats> {
let txn = self.db.begin_read()?;
let memory_log_count = txn.open_table(MEMORY_LOG)?.len()?;
let memory_kv_count = txn.open_table(MEMORY_KV)?.len()?;
let entity_index_size = txn.open_table(ENTITY_INDEX)?.len()?;
let topic_index_size = txn.open_table(TOPIC_INDEX)?.len()?;
let goal_index_size = txn.open_table(GOAL_INDEX)?.len()?;
let event_index_size = txn.open_table(EVENT_INDEX)?.len()?;
let temporal_index_size = txn.open_table(TEMPORAL_INDEX)?.len()?;
let causal_index_size = txn.open_table(CAUSAL_INDEX)?.len()?;
Ok(StoreStats {
memory_log_count,
memory_kv_count,
entity_index_size,
topic_index_size,
goal_index_size,
event_index_size,
temporal_index_size,
causal_index_size,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use redb::TableHandle;
#[test]
fn table_definitions_have_names() {
assert_eq!(MEMORY_LOG.name(), "memory_log");
assert_eq!(MEMORY_KV.name(), "memory_kv");
assert_eq!(ENTITY_INDEX.name(), "entity_index");
assert_eq!(TOPIC_INDEX.name(), "topic_index");
assert_eq!(GOAL_INDEX.name(), "goal_index");
assert_eq!(EVENT_INDEX.name(), "event_index");
assert_eq!(TEMPORAL_INDEX.name(), "temporal_index");
assert_eq!(CAUSAL_INDEX.name(), "causal_index");
assert_eq!(LINK_OVERLAY.name(), "link_overlay");
assert_eq!(SUMMARY_OVERLAY.name(), "summary_overlay");
assert_eq!(CORRECTION_OVERLAY.name(), "correction_overlay");
assert_eq!(ACTIVATION_LOG.name(), "activation_log");
assert_eq!(CONSOLIDATION_QUEUE.name(), "consolidation_queue");
}
#[test]
fn all_tables_count_is_thirteen() {
let total = U128_TABLES.len() + U64_TABLES.len() + U32_TABLES.len();
assert_eq!(total, 13);
}
}