hehe-store 0.0.1

Unified storage abstraction layer for hehe AI Agent framework
Documentation
use crate::error::Result;
use crate::traits::{CacheStore, RelationalStore, SearchStore, VectorStore};
use std::sync::Arc;

pub struct StoreRouter {
    relational: Arc<dyn RelationalStore>,
    vector: Arc<dyn VectorStore>,
    cache: Arc<dyn CacheStore>,
    search: Arc<dyn SearchStore>,
}

impl StoreRouter {
    pub fn new(
        relational: Arc<dyn RelationalStore>,
        vector: Arc<dyn VectorStore>,
        cache: Arc<dyn CacheStore>,
        search: Arc<dyn SearchStore>,
    ) -> Self {
        Self {
            relational,
            vector,
            cache,
            search,
        }
    }

    pub fn relational(&self) -> &dyn RelationalStore {
        self.relational.as_ref()
    }

    pub fn vector(&self) -> &dyn VectorStore {
        self.vector.as_ref()
    }

    pub fn cache(&self) -> &dyn CacheStore {
        self.cache.as_ref()
    }

    pub fn search(&self) -> &dyn SearchStore {
        self.search.as_ref()
    }

    pub async fn health_check(&self) -> StoreHealth {
        let relational_ok = self.relational.ping().await.is_ok();

        StoreHealth {
            relational: relational_ok,
            vector: true,
            cache: true,
            search: true,
        }
    }
}

#[derive(Debug, Clone)]
pub struct StoreHealth {
    pub relational: bool,
    pub vector: bool,
    pub cache: bool,
    pub search: bool,
}

impl StoreHealth {
    pub fn is_healthy(&self) -> bool {
        self.relational && self.vector && self.cache && self.search
    }

    pub fn unhealthy_backends(&self) -> Vec<&'static str> {
        let mut unhealthy = Vec::new();
        if !self.relational {
            unhealthy.push("relational");
        }
        if !self.vector {
            unhealthy.push("vector");
        }
        if !self.cache {
            unhealthy.push("cache");
        }
        if !self.search {
            unhealthy.push("search");
        }
        unhealthy
    }
}

#[cfg(all(feature = "sqlite", feature = "memory-cache"))]
impl StoreRouter {
    pub fn local_default() -> Result<Self> {
        use crate::local::{MemoryCache, MemoryVectorStore, SqliteFtsStore, SqliteStore};

        let sqlite = Arc::new(SqliteStore::memory()?);
        let vector = Arc::new(MemoryVectorStore::new());
        let cache = Arc::new(MemoryCache::new(10000));
        let search = Arc::new(SqliteFtsStore::new(Arc::clone(&sqlite)));

        Ok(Self {
            relational: sqlite,
            vector,
            cache,
            search,
        })
    }

    pub fn local_persistent(data_dir: &str) -> Result<Self> {
        use crate::local::{MemoryCache, MemoryVectorStore, SqliteFtsStore, SqliteStore};

        std::fs::create_dir_all(data_dir)
            .map_err(|e| crate::error::StoreError::connection(e.to_string()))?;

        let db_path = format!("{}/hehe.db", data_dir);
        let sqlite = Arc::new(SqliteStore::open(&db_path)?);
        let vector = Arc::new(MemoryVectorStore::new());
        let cache = Arc::new(MemoryCache::new(10000));
        let search = Arc::new(SqliteFtsStore::new(Arc::clone(&sqlite)));

        Ok(Self {
            relational: sqlite,
            vector,
            cache,
            search,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[cfg(all(feature = "sqlite", feature = "memory-cache"))]
    #[tokio::test]
    async fn test_store_router_creation() {
        let router = StoreRouter::local_default().unwrap();

        assert_eq!(router.relational().backend_name(), "sqlite");
        assert_eq!(router.vector().backend_name(), "memory");
        assert_eq!(router.cache().backend_name(), "memory");
        assert_eq!(router.search().backend_name(), "sqlite-fts5");
    }

    #[cfg(all(feature = "sqlite", feature = "memory-cache"))]
    #[tokio::test]
    async fn test_store_router_health_check() {
        let router = StoreRouter::local_default().unwrap();
        let health = router.health_check().await;

        assert!(health.is_healthy());
        assert!(health.unhealthy_backends().is_empty());
    }
}