1use std::path::Path;
8use std::sync::Arc;
9
10use rusqlite::OptionalExtension;
11
12use crate::error::SqliteError;
13use crate::pool::{ConnectionPool, PoolConfig};
14use crate::sql_bridge::SqlBridge;
15use crate::stores::{entity, event, graph, note, sparse, text, vectors};
16
17pub struct StorageBackend {
19 pool: Arc<ConnectionPool>,
20 is_file_backed: bool,
21 path: Option<std::path::PathBuf>,
22}
23
24impl StorageBackend {
25 pub fn sqlite(path: impl AsRef<Path>) -> Result<Self, SqliteError> {
31 crate::extension::ensure_extensions_loaded();
32 let resolved = path.as_ref().to_path_buf();
33 let config = PoolConfig {
34 path: Some(resolved.clone()),
35 ..PoolConfig::default()
36 };
37 let pool = ConnectionPool::new(config)?;
38 Ok(Self {
39 pool: Arc::new(pool),
40 is_file_backed: true,
41 path: Some(resolved),
42 })
43 }
44
45 pub fn sqlite_read_only(path: impl AsRef<Path>) -> Result<Self, SqliteError> {
55 crate::extension::ensure_extensions_loaded();
56 let resolved = path.as_ref().to_path_buf();
57 let config = PoolConfig {
58 path: Some(resolved.clone()),
59 ..PoolConfig::default()
60 };
61 let pool = ConnectionPool::new(config)?;
62 {
66 let writer = pool.try_writer()?;
67 writer
68 .conn()
69 .pragma_update(None, "query_only", "ON")
70 .map_err(SqliteError::Rusqlite)?;
71 }
72 Ok(Self {
73 pool: Arc::new(pool),
74 is_file_backed: true,
75 path: Some(resolved),
76 })
77 }
78
79 pub fn memory() -> Result<Self, SqliteError> {
85 crate::extension::ensure_extensions_loaded();
86 let config = PoolConfig {
87 path: None,
88 ..PoolConfig::default()
89 };
90 let pool = ConnectionPool::new(config)?;
91 Ok(Self {
92 pool: Arc::new(pool),
93 is_file_backed: false,
94 path: None,
95 })
96 }
97
98 pub fn sql(&self) -> Arc<dyn khive_storage::SqlAccess> {
102 Arc::new(SqlBridge::new(Arc::clone(&self.pool), self.is_file_backed))
103 }
104
105 pub fn apply_schema(
111 &self,
112 plan: &crate::migrations::ServiceSchemaPlan,
113 ) -> Result<(), SqliteError> {
114 let writer = self.pool.try_writer()?;
115 crate::migrations::apply_schema_plan(writer.conn(), plan)
116 }
117
118 pub fn apply_pack_ddl_statements(
135 &self,
136 statements: &[&'static str],
137 ) -> Result<(), SqliteError> {
138 let writer = self.pool.try_writer()?;
139 for &stmt in statements {
140 writer.conn().execute_batch(stmt)?;
141 }
142 Ok(())
143 }
144
145 pub fn entities(&self) -> Result<Arc<dyn khive_storage::EntityStore>, SqliteError> {
149 self.entities_for_namespace("local")
150 }
151
152 pub fn entities_for_namespace(
156 &self,
157 namespace: &str,
158 ) -> Result<Arc<dyn khive_storage::EntityStore>, SqliteError> {
159 if namespace.trim().is_empty() {
160 return Err(SqliteError::InvalidData(
161 "entities namespace must be non-empty".to_string(),
162 ));
163 }
164 let writer = self.pool.try_writer()?;
165 entity::ensure_entities_schema(writer.conn())?;
166
167 Ok(Arc::new(entity::SqlEntityStore::new(
168 Arc::clone(&self.pool),
169 self.is_file_backed,
170 )))
171 }
172
173 pub fn graph(&self) -> Result<Arc<dyn khive_storage::GraphStore>, SqliteError> {
178 self.graph_for_namespace("local")
179 }
180
181 pub fn graph_for_namespace(
183 &self,
184 namespace: &str,
185 ) -> Result<Arc<dyn khive_storage::GraphStore>, SqliteError> {
186 if namespace.trim().is_empty() {
187 return Err(SqliteError::InvalidData(
188 "graph namespace must be non-empty".to_string(),
189 ));
190 }
191 let writer = self.pool.try_writer()?;
192 graph::ensure_graph_schema(writer.conn())?;
193
194 Ok(Arc::new(graph::SqlGraphStore::new_scoped(
195 Arc::clone(&self.pool),
196 self.is_file_backed,
197 namespace.trim().to_string(),
198 )))
199 }
200
201 pub fn notes(&self) -> Result<Arc<dyn khive_storage::NoteStore>, SqliteError> {
205 self.notes_for_namespace("local")
206 }
207
208 pub fn notes_for_namespace(
212 &self,
213 namespace: &str,
214 ) -> Result<Arc<dyn khive_storage::NoteStore>, SqliteError> {
215 if namespace.trim().is_empty() {
216 return Err(SqliteError::InvalidData(
217 "notes namespace must be non-empty".to_string(),
218 ));
219 }
220 let writer = self.pool.try_writer()?;
221 note::ensure_notes_schema(writer.conn())?;
222
223 Ok(Arc::new(note::SqlNoteStore::new(
224 Arc::clone(&self.pool),
225 self.is_file_backed,
226 )))
227 }
228
229 pub fn events(&self) -> Result<Arc<dyn khive_storage::EventStore>, SqliteError> {
234 self.events_for_namespace("local")
235 }
236
237 pub fn events_for_namespace(
239 &self,
240 namespace: &str,
241 ) -> Result<Arc<dyn khive_storage::EventStore>, SqliteError> {
242 if namespace.trim().is_empty() {
243 return Err(SqliteError::InvalidData(
244 "events namespace must be non-empty".to_string(),
245 ));
246 }
247 let writer = self.pool.try_writer()?;
248 event::ensure_events_schema(writer.conn())?;
249
250 Ok(Arc::new(event::SqlEventStore::new_scoped(
251 Arc::clone(&self.pool),
252 self.is_file_backed,
253 namespace.trim().to_string(),
254 )))
255 }
256
257 pub fn vectors(
263 &self,
264 model_key: &str,
265 embedding_model: &str,
266 dimensions: usize,
267 ) -> Result<Arc<dyn khive_storage::VectorStore>, SqliteError> {
268 self.vectors_for_namespace(model_key, embedding_model, dimensions, "local")
269 }
270
271 pub fn vectors_for_namespace(
281 &self,
282 model_key: &str,
283 embedding_model: &str,
284 dimensions: usize,
285 namespace: &str,
286 ) -> Result<Arc<dyn khive_storage::VectorStore>, SqliteError> {
287 if model_key.is_empty()
288 || !model_key
289 .chars()
290 .all(|c| c.is_ascii_alphanumeric() || c == '_')
291 {
292 return Err(SqliteError::InvalidData(format!(
293 "invalid model_key '{}': must be non-empty and contain only \
294 alphanumeric/underscore characters",
295 model_key
296 )));
297 }
298 if namespace.trim().is_empty() {
299 return Err(SqliteError::InvalidData(
300 "vector store namespace must be non-empty".to_string(),
301 ));
302 }
303
304 crate::extension::ensure_extensions_loaded();
306
307 let table = format!("vec_{}", model_key);
308 let writer = self.pool.try_writer()?;
309
310 let table_exists: bool = writer
317 .conn()
318 .query_row(
319 "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?1",
320 rusqlite::params![&table],
321 |row| row.get::<_, i64>(0),
322 )
323 .optional()
324 .map_err(SqliteError::Rusqlite)?
325 .is_some();
326
327 if table_exists {
328 let pragma = format!("PRAGMA table_xinfo({})", table);
334 let mut stmt = writer.conn().prepare(&pragma)?;
335 let mut rows = stmt.query([])?;
336 let mut has_field = false;
337 let mut has_embedding_model = false;
338 while let Some(row) = rows.next()? {
339 let name: String = row.get(1)?;
340 if name == "field" {
341 has_field = true;
342 }
343 if name == "embedding_model" {
344 has_embedding_model = true;
345 }
346 }
347 if !has_field || !has_embedding_model {
348 return Err(SqliteError::InvalidData(format!(
349 "vec0 table '{}' is missing required column(s) (field={}, \
350 embedding_model={}); this is a pre-v0.2.8 vector schema and is \
351 not supported — recreate the database",
352 table, has_field, has_embedding_model,
353 )));
354 }
355 }
356
357 writer
366 .conn()
367 .execute_batch(crate::migrations::EMBEDDING_MODELS_DDL)?;
368
369 let ddl = format!(
372 "CREATE VIRTUAL TABLE IF NOT EXISTS vec_{} USING vec0(\
373 subject_id TEXT PRIMARY KEY, \
374 namespace TEXT NOT NULL, \
375 kind TEXT NOT NULL, \
376 field TEXT NOT NULL, \
377 embedding_model TEXT NOT NULL, \
378 embedding float[{}] distance_metric=cosine\
379 )",
380 model_key, dimensions
381 );
382 writer.conn().execute_batch(&ddl)?;
383
384 Ok(Arc::new(vectors::SqliteVecStore::new(
385 Arc::clone(&self.pool),
386 self.is_file_backed,
387 model_key.to_string(),
388 embedding_model.to_string(),
389 dimensions,
390 namespace.trim().to_string(),
391 )?))
392 }
393
394 pub fn register_embedding_model(
399 &self,
400 engine_name: &str,
401 model_id: &str,
402 key_version: &str,
403 dimensions: u32,
404 ) -> Result<(), SqliteError> {
405 let writer = self.pool.try_writer()?;
406 writer
407 .conn()
408 .execute_batch(crate::migrations::EMBEDDING_MODELS_DDL)?;
409
410 let now = chrono::Utc::now().timestamp_micros();
411 let canonical_key =
412 format!("{engine_name}:{model_id}:{key_version}:{dimensions}").into_bytes();
413 let id = uuid::Uuid::new_v4();
414 writer.conn().execute(
415 "INSERT INTO _embedding_models \
416 (id, engine_name, model_id, key_version, dim, output_dim, status, \
417 activated_at, superseded_at, superseded_by, canonical_key, created_at) \
418 VALUES (?1, ?2, ?3, ?4, ?5, NULL, 'active', ?6, NULL, NULL, ?7, ?8) \
419 ON CONFLICT(canonical_key) DO UPDATE SET \
420 status = 'active', \
421 activated_at = COALESCE(_embedding_models.activated_at, excluded.activated_at)",
422 rusqlite::params![
423 id.as_bytes().as_slice(),
424 engine_name,
425 model_id,
426 key_version,
427 dimensions as i64,
428 now,
429 canonical_key,
430 now,
431 ],
432 )?;
433 Ok(())
434 }
435
436 pub fn sparse(
440 &self,
441 model_key: &str,
442 ) -> Result<Arc<dyn khive_storage::SparseStore>, SqliteError> {
443 self.sparse_for_namespace(model_key, "local")
444 }
445
446 pub fn sparse_for_namespace(
450 &self,
451 model_key: &str,
452 namespace: &str,
453 ) -> Result<Arc<dyn khive_storage::SparseStore>, SqliteError> {
454 if model_key.is_empty()
455 || !model_key
456 .chars()
457 .all(|c| c.is_ascii_alphanumeric() || c == '_')
458 {
459 return Err(SqliteError::InvalidData(format!(
460 "invalid model_key '{}': must be non-empty and contain only alphanumeric/underscore characters",
461 model_key
462 )));
463 }
464 if namespace.trim().is_empty() {
465 return Err(SqliteError::InvalidData(
466 "sparse store namespace must be non-empty".to_string(),
467 ));
468 }
469
470 let writer = self.pool.try_writer()?;
471 sparse::ensure_sparse_schema(writer.conn(), model_key).map_err(SqliteError::Rusqlite)?;
472
473 Ok(Arc::new(sparse::SqliteSparseStore::new(
474 Arc::clone(&self.pool),
475 self.is_file_backed,
476 model_key.to_string(),
477 namespace.trim().to_string(),
478 )?))
479 }
480
481 pub fn text(&self, table_key: &str) -> Result<Arc<dyn khive_storage::TextSearch>, SqliteError> {
488 self.text_with_tokenizer(table_key, "trigram")
489 }
490
491 pub fn text_with_tokenizer(
499 &self,
500 table_key: &str,
501 tokenizer: &str,
502 ) -> Result<Arc<dyn khive_storage::TextSearch>, SqliteError> {
503 if table_key.is_empty()
504 || !table_key
505 .chars()
506 .all(|c| c.is_ascii_alphanumeric() || c == '_')
507 {
508 return Err(SqliteError::InvalidData(format!(
509 "invalid table_key '{}': must be non-empty and contain only \
510 alphanumeric/underscore characters",
511 table_key
512 )));
513 }
514 if tokenizer.is_empty()
515 || !tokenizer
516 .chars()
517 .all(|c| c.is_ascii_alphanumeric() || c == '_')
518 {
519 return Err(SqliteError::InvalidData(format!(
520 "invalid tokenizer '{}': must be non-empty and contain only \
521 alphanumeric/underscore characters",
522 tokenizer
523 )));
524 }
525
526 let ddl = format!(
527 "CREATE VIRTUAL TABLE IF NOT EXISTS fts_{} USING fts5(\
528 subject_id UNINDEXED, \
529 kind UNINDEXED, \
530 title, \
531 body, \
532 tags UNINDEXED, \
533 namespace UNINDEXED, \
534 metadata UNINDEXED, \
535 updated_at UNINDEXED, \
536 tokenize = '{}'\
537 )",
538 table_key, tokenizer
539 );
540 let writer = self.pool.try_writer()?;
541 writer.conn().execute_batch(&ddl)?;
542
543 Ok(Arc::new(text::Fts5TextSearch::new(
544 Arc::clone(&self.pool),
545 self.is_file_backed,
546 table_key.to_string(),
547 )))
548 }
549
550 pub fn is_file_backed(&self) -> bool {
552 self.is_file_backed
553 }
554
555 pub fn data_dir(&self) -> Option<std::path::PathBuf> {
558 self.path.as_ref()?.parent().map(|p| p.to_path_buf())
559 }
560
561 pub fn pool(&self) -> &ConnectionPool {
563 &self.pool
564 }
565
566 pub fn pool_arc(&self) -> Arc<ConnectionPool> {
568 Arc::clone(&self.pool)
569 }
570}
571
572#[cfg(test)]
573mod tests {
574 use super::*;
575 use khive_storage::types::{SqlStatement, SqlValue};
576
577 #[test]
578 fn memory_backend_creates_successfully() {
579 let backend = StorageBackend::memory().expect("memory backend should create");
580 assert!(!backend.is_file_backed());
581 }
582
583 #[test]
584 fn file_backend_creates_successfully() {
585 let dir = tempfile::tempdir().unwrap();
586 let path = dir.path().join("test.db");
587 let backend = StorageBackend::sqlite(&path).expect("file backend should create");
588 assert!(backend.is_file_backed());
589 assert!(path.exists());
590 }
591
592 #[test]
593 fn data_dir_returns_none_for_memory_backend() {
594 let backend = StorageBackend::memory().expect("memory backend");
595 assert!(backend.data_dir().is_none());
596 }
597
598 #[test]
599 fn data_dir_returns_parent_dir_for_file_backend() {
600 let dir = tempfile::tempdir().unwrap();
601 let path = dir.path().join("data.db");
602 let backend = StorageBackend::sqlite(&path).expect("file backend");
603 let got = backend.data_dir().expect("file backend must return Some");
604 assert_eq!(got, dir.path());
605 }
606
607 #[tokio::test]
608 async fn sql_access_memory_roundtrip() {
609 let backend = StorageBackend::memory().unwrap();
610 let sql = backend.sql();
611
612 let mut writer = sql.writer().await.unwrap();
613 writer
614 .execute_script(
615 "CREATE TABLE test_rt (id TEXT PRIMARY KEY, value INTEGER NOT NULL)".into(),
616 )
617 .await
618 .unwrap();
619
620 let affected = writer
621 .execute(SqlStatement {
622 sql: "INSERT INTO test_rt (id, value) VALUES (?1, ?2)".into(),
623 params: vec![SqlValue::Text("row1".into()), SqlValue::Integer(42)],
624 label: None,
625 })
626 .await
627 .unwrap();
628 assert_eq!(affected, 1);
629
630 let mut reader = sql.reader().await.unwrap();
631 let row = reader
632 .query_row(SqlStatement {
633 sql: "SELECT id, value FROM test_rt WHERE id = ?1".into(),
634 params: vec![SqlValue::Text("row1".into())],
635 label: None,
636 })
637 .await
638 .unwrap();
639
640 let row = row.expect("should find the inserted row");
641 assert_eq!(row.columns.len(), 2);
642 match &row.columns[0].value {
643 SqlValue::Text(s) => assert_eq!(s, "row1"),
644 other => panic!("expected Text, got {other:?}"),
645 }
646 match &row.columns[1].value {
647 SqlValue::Integer(v) => assert_eq!(*v, 42),
648 other => panic!("expected Integer, got {other:?}"),
649 }
650 }
651
652 #[tokio::test]
653 async fn sql_access_file_roundtrip() {
654 let dir = tempfile::tempdir().unwrap();
655 let path = dir.path().join("test_roundtrip.db");
656 let backend = StorageBackend::sqlite(&path).unwrap();
657 let sql = backend.sql();
658
659 let mut writer = sql.writer().await.unwrap();
660 writer
661 .execute_script("CREATE TABLE test_f (k TEXT PRIMARY KEY, v TEXT)".into())
662 .await
663 .unwrap();
664 writer
665 .execute(SqlStatement {
666 sql: "INSERT INTO test_f (k, v) VALUES (?1, ?2)".into(),
667 params: vec![
668 SqlValue::Text("hello".into()),
669 SqlValue::Text("world".into()),
670 ],
671 label: None,
672 })
673 .await
674 .unwrap();
675
676 let mut reader = sql.reader().await.unwrap();
677 let rows = reader
678 .query_all(SqlStatement {
679 sql: "SELECT k, v FROM test_f".into(),
680 params: vec![],
681 label: None,
682 })
683 .await
684 .unwrap();
685 assert_eq!(rows.len(), 1);
686 match &rows[0].columns[1].value {
687 SqlValue::Text(s) => assert_eq!(s, "world"),
688 other => panic!("expected Text, got {other:?}"),
689 }
690 }
691
692 #[tokio::test]
693 #[cfg(feature = "vectors")]
694 async fn vectors_roundtrip_via_public_api() {
695 let backend = StorageBackend::memory().unwrap();
696 let store = backend.vectors("test_api", "test_api", 3).unwrap();
697
698 let id = uuid::Uuid::new_v4();
699 store
700 .insert(
701 id,
702 khive_types::SubstrateKind::Entity,
703 "local",
704 "content",
705 vec![vec![1.0, 0.0, 0.0]],
706 )
707 .await
708 .unwrap();
709
710 let hits = store
711 .search(khive_storage::types::VectorSearchRequest {
712 query_vectors: vec![vec![1.0, 0.0, 0.0]],
713 top_k: 1,
714 namespace: None,
715 kind: None,
716 embedding_model: None,
717 filter: None,
718 backend_hints: None,
719 })
720 .await
721 .unwrap();
722
723 assert_eq!(hits.len(), 1);
724 assert_eq!(hits[0].subject_id, id);
725 assert!(hits[0].score.to_f64() > 0.99);
726 }
727
728 #[tokio::test]
729 #[cfg(feature = "vectors")]
730 async fn vectors_creates_table_idempotently() {
731 let backend = StorageBackend::memory().unwrap();
732
733 let store1 = backend.vectors("idempotent", "idempotent", 3).unwrap();
734 let store2 = backend.vectors("idempotent", "idempotent", 3).unwrap();
735
736 let id = uuid::Uuid::new_v4();
737 store1
738 .insert(
739 id,
740 khive_types::SubstrateKind::Entity,
741 "local",
742 "content",
743 vec![vec![1.0, 0.0, 0.0]],
744 )
745 .await
746 .unwrap();
747
748 let count = store2.count().await.unwrap();
749 assert_eq!(count, 1);
750 }
751
752 #[tokio::test]
753 async fn text_roundtrip_via_public_api() {
754 let backend = StorageBackend::memory().unwrap();
755 let store = backend.text("test_api").unwrap();
756
757 let id = uuid::Uuid::new_v4();
758 let doc = khive_storage::types::TextDocument {
759 subject_id: id,
760 kind: khive_types::SubstrateKind::Entity,
761 title: Some("Test Title".to_string()),
762 body: "This is a searchable document about Rust.".to_string(),
763 tags: vec!["rust".to_string()],
764 namespace: "test_ns".to_string(),
765 metadata: None,
766 updated_at: chrono::Utc::now(),
767 };
768 store.upsert_document(doc).await.unwrap();
769
770 let hits = store
771 .search(khive_storage::types::TextSearchRequest {
772 query: "Rust".to_string(),
773 mode: khive_storage::types::TextQueryMode::Plain,
774 filter: Some(khive_storage::types::TextFilter {
775 namespaces: vec!["test_ns".to_string()],
776 ..Default::default()
777 }),
778 top_k: 1,
779 snippet_chars: 64,
780 })
781 .await
782 .unwrap();
783
784 assert_eq!(hits.len(), 1);
785 assert_eq!(hits[0].subject_id, id);
786 assert!(hits[0].score.to_f64() > 0.0);
787 }
788
789 #[tokio::test]
790 async fn text_creates_table_idempotently() {
791 let backend = StorageBackend::memory().unwrap();
792
793 let store1 = backend.text("idempotent_fts").unwrap();
794 let store2 = backend.text("idempotent_fts").unwrap();
795
796 let id = uuid::Uuid::new_v4();
797 let doc = khive_storage::types::TextDocument {
798 subject_id: id,
799 kind: khive_types::SubstrateKind::Note,
800 title: None,
801 body: "Hello world.".to_string(),
802 tags: vec![],
803 namespace: "test_ns".to_string(),
804 metadata: None,
805 updated_at: chrono::Utc::now(),
806 };
807 store1.upsert_document(doc).await.unwrap();
808
809 let count = store2
810 .count(khive_storage::types::TextFilter {
811 namespaces: vec!["test_ns".to_string()],
812 ..Default::default()
813 })
814 .await
815 .unwrap();
816 assert_eq!(count, 1);
817 }
818
819 #[test]
820 fn invalid_model_key_rejected() {
821 let backend = StorageBackend::memory().unwrap();
822 assert!(backend.vectors("bad key!", "bad key!", 3).is_err());
823 assert!(backend.vectors("", "", 3).is_err());
824 }
825
826 #[test]
827 fn invalid_table_key_rejected() {
828 let backend = StorageBackend::memory().unwrap();
829 assert!(backend.text("bad key!").is_err());
830 assert!(backend.text("").is_err());
831 }
832
833 #[test]
834 fn apply_schema_runs_migrations_idempotently() {
835 static MIGRATIONS: &[crate::migrations::Migration] = &[crate::migrations::Migration {
836 id: "001_init",
837 up_sql: "CREATE TABLE IF NOT EXISTS schema_test (id TEXT PRIMARY KEY);",
838 down_sql: None,
839 is_already_applied: None,
840 }];
841 let plan = crate::migrations::ServiceSchemaPlan {
842 service: "schema_test_svc",
843 sqlite: MIGRATIONS,
844 postgres: &[],
845 };
846
847 let backend = StorageBackend::memory().unwrap();
848 backend.apply_schema(&plan).unwrap();
849 backend.apply_schema(&plan).unwrap();
850
851 let reader = backend.pool().reader().unwrap();
852 let count: i64 = reader
853 .conn()
854 .query_row(
855 "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='schema_test'",
856 [],
857 |row| row.get(0),
858 )
859 .unwrap();
860 assert_eq!(count, 1);
861 }
862}