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