hehe_store/
router.rs

1use crate::error::Result;
2use crate::traits::{CacheStore, RelationalStore, SearchStore, VectorStore};
3use std::sync::Arc;
4
5pub struct StoreRouter {
6    relational: Arc<dyn RelationalStore>,
7    vector: Arc<dyn VectorStore>,
8    cache: Arc<dyn CacheStore>,
9    search: Arc<dyn SearchStore>,
10}
11
12impl StoreRouter {
13    pub fn new(
14        relational: Arc<dyn RelationalStore>,
15        vector: Arc<dyn VectorStore>,
16        cache: Arc<dyn CacheStore>,
17        search: Arc<dyn SearchStore>,
18    ) -> Self {
19        Self {
20            relational,
21            vector,
22            cache,
23            search,
24        }
25    }
26
27    pub fn relational(&self) -> &dyn RelationalStore {
28        self.relational.as_ref()
29    }
30
31    pub fn vector(&self) -> &dyn VectorStore {
32        self.vector.as_ref()
33    }
34
35    pub fn cache(&self) -> &dyn CacheStore {
36        self.cache.as_ref()
37    }
38
39    pub fn search(&self) -> &dyn SearchStore {
40        self.search.as_ref()
41    }
42
43    pub async fn health_check(&self) -> StoreHealth {
44        let relational_ok = self.relational.ping().await.is_ok();
45
46        StoreHealth {
47            relational: relational_ok,
48            vector: true,
49            cache: true,
50            search: true,
51        }
52    }
53}
54
55#[derive(Debug, Clone)]
56pub struct StoreHealth {
57    pub relational: bool,
58    pub vector: bool,
59    pub cache: bool,
60    pub search: bool,
61}
62
63impl StoreHealth {
64    pub fn is_healthy(&self) -> bool {
65        self.relational && self.vector && self.cache && self.search
66    }
67
68    pub fn unhealthy_backends(&self) -> Vec<&'static str> {
69        let mut unhealthy = Vec::new();
70        if !self.relational {
71            unhealthy.push("relational");
72        }
73        if !self.vector {
74            unhealthy.push("vector");
75        }
76        if !self.cache {
77            unhealthy.push("cache");
78        }
79        if !self.search {
80            unhealthy.push("search");
81        }
82        unhealthy
83    }
84}
85
86#[cfg(all(feature = "sqlite", feature = "memory-cache"))]
87impl StoreRouter {
88    pub fn local_default() -> Result<Self> {
89        use crate::local::{MemoryCache, MemoryVectorStore, SqliteFtsStore, SqliteStore};
90
91        let sqlite = Arc::new(SqliteStore::memory()?);
92        let vector = Arc::new(MemoryVectorStore::new());
93        let cache = Arc::new(MemoryCache::new(10000));
94        let search = Arc::new(SqliteFtsStore::new(Arc::clone(&sqlite)));
95
96        Ok(Self {
97            relational: sqlite,
98            vector,
99            cache,
100            search,
101        })
102    }
103
104    pub fn local_persistent(data_dir: &str) -> Result<Self> {
105        use crate::local::{MemoryCache, MemoryVectorStore, SqliteFtsStore, SqliteStore};
106
107        std::fs::create_dir_all(data_dir)
108            .map_err(|e| crate::error::StoreError::connection(e.to_string()))?;
109
110        let db_path = format!("{}/hehe.db", data_dir);
111        let sqlite = Arc::new(SqliteStore::open(&db_path)?);
112        let vector = Arc::new(MemoryVectorStore::new());
113        let cache = Arc::new(MemoryCache::new(10000));
114        let search = Arc::new(SqliteFtsStore::new(Arc::clone(&sqlite)));
115
116        Ok(Self {
117            relational: sqlite,
118            vector,
119            cache,
120            search,
121        })
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[cfg(all(feature = "sqlite", feature = "memory-cache"))]
130    #[tokio::test]
131    async fn test_store_router_creation() {
132        let router = StoreRouter::local_default().unwrap();
133
134        assert_eq!(router.relational().backend_name(), "sqlite");
135        assert_eq!(router.vector().backend_name(), "memory");
136        assert_eq!(router.cache().backend_name(), "memory");
137        assert_eq!(router.search().backend_name(), "sqlite-fts5");
138    }
139
140    #[cfg(all(feature = "sqlite", feature = "memory-cache"))]
141    #[tokio::test]
142    async fn test_store_router_health_check() {
143        let router = StoreRouter::local_default().unwrap();
144        let health = router.health_check().await;
145
146        assert!(health.is_healthy());
147        assert!(health.unhealthy_backends().is_empty());
148    }
149}
hehe_store/router.rs

hehe_store/
router.rs
