Skip to main content

the_code_graph_storage/
lib.rs

1pub mod embedding_store;
2pub mod graph_store;
3pub mod mapping;
4pub mod schema;
5pub mod search_index;
6
7use domain::error::{CodeGraphError, Result};
8use r2d2::Pool;
9use r2d2_sqlite::SqliteConnectionManager;
10use std::path::Path;
11
12#[derive(Debug, Clone)]
13pub struct SqliteStore {
14    pool: Pool<SqliteConnectionManager>,
15}
16
17impl SqliteStore {
18    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
19        let manager = SqliteConnectionManager::file(path.as_ref()).with_init(|c| {
20            c.execute_batch(
21                "PRAGMA journal_mode = WAL;
22                 PRAGMA foreign_keys = ON;
23                 PRAGMA busy_timeout = 5000;",
24            )
25        });
26        let pool = Pool::builder()
27            .build(manager)
28            .map_err(|e| CodeGraphError::Storage(e.to_string()))?;
29        let conn = pool
30            .get()
31            .map_err(|e| CodeGraphError::Storage(e.to_string()))?;
32        schema::ensure_schema(&conn)?;
33        Ok(Self { pool })
34    }
35
36    pub fn open_in_memory() -> Result<Self> {
37        let manager = SqliteConnectionManager::memory().with_init(|c| {
38            c.execute_batch(
39                "PRAGMA foreign_keys = ON;
40                 PRAGMA busy_timeout = 5000;",
41            )
42        });
43        let pool = Pool::builder()
44            .max_size(1)
45            .build(manager)
46            .map_err(|e| CodeGraphError::Storage(e.to_string()))?;
47        let conn = pool
48            .get()
49            .map_err(|e| CodeGraphError::Storage(e.to_string()))?;
50        schema::ensure_schema(&conn)?;
51        Ok(Self { pool })
52    }
53
54    pub(crate) fn conn(&self) -> Result<r2d2::PooledConnection<SqliteConnectionManager>> {
55        self.pool
56            .get()
57            .map_err(|e| CodeGraphError::Storage(e.to_string()))
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64    use domain::model::*;
65    use domain::ports::{GraphStore, SearchIndex};
66
67    fn assert_send_sync<T: Send + Sync>() {}
68
69    #[test]
70    fn sqlite_store_is_send_sync() {
71        assert_send_sync::<SqliteStore>();
72    }
73
74    #[test]
75    fn open_in_memory_creates_schema() {
76        let store = SqliteStore::open_in_memory().unwrap();
77        let conn = store.conn().unwrap();
78        let version: i32 = conn
79            .query_row("PRAGMA user_version", [], |r| r.get(0))
80            .unwrap();
81        assert_eq!(version, 2);
82    }
83
84    #[test]
85    fn open_in_memory_creates_all_tables() {
86        let store = SqliteStore::open_in_memory().unwrap();
87        let conn = store.conn().unwrap();
88        let tables: Vec<String> = conn
89            .prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
90            .unwrap()
91            .query_map([], |r| r.get(0))
92            .unwrap()
93            .filter_map(|r| r.ok())
94            .collect();
95        assert!(tables.contains(&"files".to_string()));
96        assert!(tables.contains(&"non_parsed_files".to_string()));
97        assert!(tables.contains(&"symbols".to_string()));
98        assert!(tables.contains(&"edges".to_string()));
99        assert!(tables.contains(&"metadata".to_string()));
100    }
101
102    #[test]
103    fn pragmas_are_set() {
104        let store = SqliteStore::open_in_memory().unwrap();
105        let conn = store.conn().unwrap();
106        let fk: i32 = conn
107            .query_row("PRAGMA foreign_keys", [], |r| r.get(0))
108            .unwrap();
109        assert_eq!(fk, 1);
110    }
111
112    #[test]
113    fn unsupported_schema_version_errors() {
114        let store = SqliteStore::open_in_memory().unwrap();
115        {
116            let conn = store.conn().unwrap();
117            conn.pragma_update(None, "user_version", 99).unwrap();
118        }
119        let conn = store.conn().unwrap();
120        let result = schema::ensure_schema(&conn);
121        assert!(result.is_err());
122        let err_msg = format!("{}", result.unwrap_err());
123        assert!(err_msg.contains("unsupported schema version"));
124    }
125
126    // --- Integration tests (T07) ---
127
128    #[test]
129    fn fts5_trigger_sync_on_store_file_data() {
130        let store = SqliteStore::open_in_memory().unwrap();
131        let file = FileNode {
132            path: "a.rs".into(),
133            language: Language::Rust,
134            hash: "h".into(),
135        };
136        let sym = SymbolNode {
137            name: "BatchSymbol".into(),
138            qualified_name: "a.rs::BatchSymbol".into(),
139            kind: SymbolKind::Function,
140            location: Location {
141                file: "a.rs".into(),
142                line_start: 1,
143                line_end: 5,
144                col_start: 0,
145                col_end: 0,
146            },
147            visibility: Visibility::Public,
148            is_exported: true,
149            is_async: false,
150            is_test: false,
151            decorators: vec![],
152            signature: None,
153        };
154        store.store_file_data(&file, &[sym], &[]).unwrap();
155        let results = store.search("BatchSymbol", 10).unwrap();
156        assert_eq!(results.len(), 1);
157    }
158
159    #[test]
160    fn cascade_delete_removes_from_fts() {
161        let store = SqliteStore::open_in_memory().unwrap();
162        let file = FileNode {
163            path: "a.rs".into(),
164            language: Language::Rust,
165            hash: "h".into(),
166        };
167        let sym = SymbolNode {
168            name: "Doomed".into(),
169            qualified_name: "a.rs::Doomed".into(),
170            kind: SymbolKind::Class,
171            location: Location {
172                file: "a.rs".into(),
173                line_start: 1,
174                line_end: 5,
175                col_start: 0,
176                col_end: 0,
177            },
178            visibility: Visibility::Public,
179            is_exported: true,
180            is_async: false,
181            is_test: false,
182            decorators: vec![],
183            signature: None,
184        };
185        store.store_file_data(&file, &[sym], &[]).unwrap();
186        store.remove_file("a.rs".as_ref()).unwrap();
187        assert!(store.search("Doomed", 10).unwrap().is_empty());
188    }
189
190    #[test]
191    fn concurrent_reads_do_not_deadlock() {
192        use std::thread;
193        let store = SqliteStore::open_in_memory().unwrap();
194        let file = FileNode {
195            path: "a.rs".into(),
196            language: Language::Rust,
197            hash: "h".into(),
198        };
199        store.upsert_file(&file).unwrap();
200
201        let s1 = store.clone();
202        let s2 = store.clone();
203        let t1 = thread::spawn(move || s1.all_files().unwrap());
204        let t2 = thread::spawn(move || s2.stats().unwrap());
205        t1.join().unwrap();
206        t2.join().unwrap();
207    }
208
209    #[test]
210    fn store_file_data_atomicity() {
211        let store = SqliteStore::open_in_memory().unwrap();
212        let file = FileNode {
213            path: "a.rs".into(),
214            language: Language::Rust,
215            hash: "h".into(),
216        };
217        let sym = SymbolNode {
218            name: "X".into(),
219            qualified_name: "a.rs::X".into(),
220            kind: SymbolKind::Function,
221            location: Location {
222                file: "a.rs".into(),
223                line_start: 1,
224                line_end: 2,
225                col_start: 0,
226                col_end: 0,
227            },
228            visibility: Visibility::Public,
229            is_exported: false,
230            is_async: false,
231            is_test: false,
232            decorators: vec![],
233            signature: None,
234        };
235        store.store_file_data(&file, &[sym], &[]).unwrap();
236        assert!(store.get_file("a.rs".as_ref()).unwrap().is_some());
237        assert!(store.get_symbol("a.rs::X").unwrap().is_some());
238    }
239}