agtrace_index/
db.rs

1use rusqlite::Connection;
2use std::path::Path;
3
4use crate::{Result, queries, records::*, schema};
5
6pub struct Database {
7    conn: Connection,
8}
9
10impl Database {
11    pub fn open(db_path: &Path) -> Result<Self> {
12        let conn = Connection::open(db_path)?;
13        let db = Self { conn };
14        schema::init_schema(&db.conn)?;
15        Ok(db)
16    }
17
18    pub fn open_in_memory() -> Result<Self> {
19        let conn = Connection::open_in_memory()?;
20        let db = Self { conn };
21        schema::init_schema(&db.conn)?;
22        Ok(db)
23    }
24
25    // Project operations
26    pub fn insert_or_update_project(&self, project: &ProjectRecord) -> Result<()> {
27        queries::project::insert_or_update(&self.conn, project)
28    }
29
30    pub fn get_project(&self, hash: &str) -> Result<Option<ProjectRecord>> {
31        queries::project::get(&self.conn, hash)
32    }
33
34    pub fn list_projects(&self) -> Result<Vec<ProjectRecord>> {
35        queries::project::list(&self.conn)
36    }
37
38    pub fn count_sessions_for_project(&self, project_hash: &str) -> Result<usize> {
39        queries::project::count_sessions(&self.conn, project_hash)
40    }
41
42    // Session operations
43    pub fn insert_or_update_session(&self, session: &SessionRecord) -> Result<()> {
44        queries::session::insert_or_update(&self.conn, session)
45    }
46
47    pub fn get_session_by_id(&self, session_id: &str) -> Result<Option<SessionSummary>> {
48        queries::session::get_by_id(&self.conn, session_id)
49    }
50
51    pub fn list_sessions(
52        &self,
53        project_hash: Option<&agtrace_types::ProjectHash>,
54        provider: Option<&str>,
55        limit: Option<usize>,
56    ) -> Result<Vec<SessionSummary>> {
57        queries::session::list(&self.conn, project_hash, provider, limit)
58    }
59
60    pub fn find_session_by_prefix(&self, prefix: &str) -> Result<Option<String>> {
61        queries::session::find_by_prefix(&self.conn, prefix)
62    }
63
64    // Log file operations
65    pub fn insert_or_update_log_file(&self, log_file: &LogFileRecord) -> Result<()> {
66        queries::log_file::insert_or_update(&self.conn, log_file)
67    }
68
69    pub fn get_session_files(&self, session_id: &str) -> Result<Vec<LogFileRecord>> {
70        queries::log_file::get_session_files(&self.conn, session_id)
71    }
72
73    pub fn get_all_log_files(&self) -> Result<Vec<LogFileRecord>> {
74        queries::log_file::get_all(&self.conn)
75    }
76
77    // Utility operations
78    pub fn vacuum(&self) -> Result<()> {
79        self.conn.execute("VACUUM", [])?;
80        println!("Database vacuumed successfully");
81        Ok(())
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use crate::schema::SCHEMA_VERSION;
89    use rusqlite::params;
90
91    #[test]
92    fn test_schema_initialization() {
93        let db = Database::open_in_memory().unwrap();
94
95        let projects = db.list_projects().unwrap();
96        assert_eq!(projects.len(), 0);
97    }
98
99    #[test]
100    fn test_insert_project() {
101        let db = Database::open_in_memory().unwrap();
102
103        let project = ProjectRecord {
104            hash: agtrace_types::ProjectHash::from("abc123"),
105            root_path: Some("/path/to/project".to_string()),
106            last_scanned_at: Some("2025-12-10T10:00:00Z".to_string()),
107        };
108
109        db.insert_or_update_project(&project).unwrap();
110
111        let retrieved = db.get_project("abc123").unwrap().unwrap();
112        assert_eq!(retrieved.hash, agtrace_types::ProjectHash::from("abc123"));
113        assert_eq!(retrieved.root_path, Some("/path/to/project".to_string()));
114    }
115
116    #[test]
117    fn test_insert_session_with_fk() {
118        let db = Database::open_in_memory().unwrap();
119
120        let project = ProjectRecord {
121            hash: agtrace_types::ProjectHash::from("abc123"),
122            root_path: Some("/path/to/project".to_string()),
123            last_scanned_at: Some("2025-12-10T10:00:00Z".to_string()),
124        };
125        db.insert_or_update_project(&project).unwrap();
126
127        let session = SessionRecord {
128            id: "session-001".to_string(),
129            project_hash: agtrace_types::ProjectHash::from("abc123"),
130            provider: "claude".to_string(),
131            start_ts: Some("2025-12-10T10:05:00Z".to_string()),
132            end_ts: Some("2025-12-10T10:15:00Z".to_string()),
133            snippet: Some("Test session".to_string()),
134            is_valid: true,
135        };
136
137        db.insert_or_update_session(&session).unwrap();
138
139        let sessions = db
140            .list_sessions(
141                Some(&agtrace_types::ProjectHash::from("abc123")),
142                None,
143                Some(10),
144            )
145            .unwrap();
146        assert_eq!(sessions.len(), 1);
147        assert_eq!(sessions[0].id, "session-001");
148        assert_eq!(sessions[0].provider, "claude");
149    }
150
151    #[test]
152    fn test_insert_log_file() {
153        let db = Database::open_in_memory().unwrap();
154
155        let project = ProjectRecord {
156            hash: agtrace_types::ProjectHash::from("abc123"),
157            root_path: Some("/path/to/project".to_string()),
158            last_scanned_at: Some("2025-12-10T10:00:00Z".to_string()),
159        };
160        db.insert_or_update_project(&project).unwrap();
161
162        let session = SessionRecord {
163            id: "session-001".to_string(),
164            project_hash: agtrace_types::ProjectHash::from("abc123"),
165            provider: "claude".to_string(),
166            start_ts: Some("2025-12-10T10:05:00Z".to_string()),
167            end_ts: None,
168            snippet: None,
169            is_valid: true,
170        };
171        db.insert_or_update_session(&session).unwrap();
172
173        let log_file = LogFileRecord {
174            path: "/path/to/log.jsonl".to_string(),
175            session_id: "session-001".to_string(),
176            role: "main".to_string(),
177            file_size: Some(1024),
178            mod_time: Some("2025-12-10T10:05:00Z".to_string()),
179        };
180
181        db.insert_or_update_log_file(&log_file).unwrap();
182
183        let files = db.get_session_files("session-001").unwrap();
184        assert_eq!(files.len(), 1);
185        assert_eq!(files[0].path, "/path/to/log.jsonl");
186        assert_eq!(files[0].role, "main");
187    }
188
189    #[test]
190    fn test_list_sessions_query() {
191        let db = Database::open_in_memory().unwrap();
192
193        let project = ProjectRecord {
194            hash: agtrace_types::ProjectHash::from("abc123"),
195            root_path: Some("/path/to/project".to_string()),
196            last_scanned_at: Some("2025-12-10T10:00:00Z".to_string()),
197        };
198        db.insert_or_update_project(&project).unwrap();
199
200        for i in 1..=5 {
201            let session = SessionRecord {
202                id: format!("session-{:03}", i),
203                project_hash: agtrace_types::ProjectHash::from("abc123"),
204                provider: "claude".to_string(),
205                start_ts: Some(format!("2025-12-10T10:{:02}:00Z", i)),
206                end_ts: None,
207                snippet: Some(format!("Session {}", i)),
208                is_valid: true,
209            };
210            db.insert_or_update_session(&session).unwrap();
211        }
212
213        let sessions = db
214            .list_sessions(
215                Some(&agtrace_types::ProjectHash::from("abc123")),
216                None,
217                Some(10),
218            )
219            .unwrap();
220        assert_eq!(sessions.len(), 5);
221
222        let sessions_limited = db
223            .list_sessions(
224                Some(&agtrace_types::ProjectHash::from("abc123")),
225                None,
226                Some(3),
227            )
228            .unwrap();
229        assert_eq!(sessions_limited.len(), 3);
230    }
231
232    #[test]
233    fn test_count_sessions_for_project() {
234        let db = Database::open_in_memory().unwrap();
235
236        let project = ProjectRecord {
237            hash: agtrace_types::ProjectHash::from("abc123"),
238            root_path: Some("/path/to/project".to_string()),
239            last_scanned_at: Some("2025-12-10T10:00:00Z".to_string()),
240        };
241        db.insert_or_update_project(&project).unwrap();
242
243        for i in 1..=3 {
244            let session = SessionRecord {
245                id: format!("session-{:03}", i),
246                project_hash: agtrace_types::ProjectHash::from("abc123"),
247                provider: "claude".to_string(),
248                start_ts: Some(format!("2025-12-10T10:{:02}:00Z", i)),
249                end_ts: None,
250                snippet: None,
251                is_valid: true,
252            };
253            db.insert_or_update_session(&session).unwrap();
254        }
255
256        let count = db.count_sessions_for_project("abc123").unwrap();
257        assert_eq!(count, 3);
258    }
259
260    #[test]
261    fn test_schema_version_set_on_init() {
262        let db = Database::open_in_memory().unwrap();
263
264        let version: i32 = db
265            .conn
266            .query_row("PRAGMA user_version", [], |row| row.get(0))
267            .unwrap();
268
269        assert_eq!(version, SCHEMA_VERSION);
270    }
271
272    #[test]
273    fn test_schema_rebuild_on_version_mismatch() {
274        let conn = Connection::open_in_memory().unwrap();
275
276        conn.execute_batch(
277            r#"
278            CREATE TABLE projects (hash TEXT PRIMARY KEY);
279            CREATE TABLE sessions (id TEXT PRIMARY KEY);
280            PRAGMA user_version = 999;
281            "#,
282        )
283        .unwrap();
284
285        conn.execute(
286            "INSERT INTO projects (hash) VALUES (?1)",
287            params!["old_data"],
288        )
289        .unwrap();
290
291        let db = Database { conn };
292        schema::init_schema(&db.conn).unwrap();
293
294        let version: i32 = db
295            .conn
296            .query_row("PRAGMA user_version", [], |row| row.get(0))
297            .unwrap();
298        assert_eq!(version, SCHEMA_VERSION);
299
300        let count: i64 = db
301            .conn
302            .query_row("SELECT COUNT(*) FROM projects", [], |row| row.get(0))
303            .unwrap();
304        assert_eq!(count, 0);
305    }
306
307    #[test]
308    fn test_schema_preserved_on_version_match() {
309        let db = Database::open_in_memory().unwrap();
310
311        let project = ProjectRecord {
312            hash: agtrace_types::ProjectHash::from("abc123"),
313            root_path: Some("/path/to/project".to_string()),
314            last_scanned_at: Some("2025-12-10T10:00:00Z".to_string()),
315        };
316        db.insert_or_update_project(&project).unwrap();
317
318        schema::init_schema(&db.conn).unwrap();
319
320        let retrieved = db.get_project("abc123").unwrap();
321        assert!(retrieved.is_some());
322        assert_eq!(
323            retrieved.unwrap().hash,
324            agtrace_types::ProjectHash::from("abc123")
325        );
326    }
327}
agtrace_index/db.rs

agtrace_index/
db.rs
