1use crate::core::error::{Result, SearchError};
2use crate::core::types::{ContentPreview, ExclusionRule, ExclusionRuleType, FileEntry, IndexStats};
3use crate::storage::migrations::MigrationManager;
4use chrono::{TimeZone, Utc};
5use r2d2::Pool;
6use r2d2_sqlite::SqliteConnectionManager;
7use rusqlite::{params, OptionalExtension};
8use std::path::{Path, PathBuf};
9
10pub type DbPool = Pool<SqliteConnectionManager>;
11
12pub struct Database {
13 pool: DbPool,
14}
15
16impl Database {
17 pub fn new<P: AsRef<Path>>(path: P, pool_size: u32) -> Result<Self> {
18 let manager = SqliteConnectionManager::file(path.as_ref());
19 let pool = Pool::builder()
20 .max_size(pool_size)
21 .build(manager)?;
22
23 {
24 let conn = pool.get()?;
25 MigrationManager::initialize_schema(&conn)?;
26 }
27
28 Ok(Self { pool })
29 }
30
31 pub fn in_memory(pool_size: u32) -> Result<Self> {
32 let manager = SqliteConnectionManager::memory();
33 let pool = Pool::builder()
34 .max_size(pool_size)
35 .build(manager)?;
36
37 {
38 let conn = pool.get()?;
39 MigrationManager::initialize_schema(&conn)?;
40 }
41
42 Ok(Self { pool })
43 }
44
45 pub fn insert_file(&self, file: &FileEntry) -> Result<i64> {
46 let conn = self.pool.get()?;
47
48 let created_at = file.created_at.map(|dt| dt.timestamp());
49 let modified_at = file.modified_at.map(|dt| dt.timestamp());
50 let accessed_at = file.accessed_at.map(|dt| dt.timestamp());
51 let indexed_at = file.indexed_at.timestamp();
52 let last_verified = file.last_verified.timestamp();
53
54 conn.execute(
55 r#"
56 INSERT INTO files (
57 path, name, extension, size, created_at, modified_at, accessed_at,
58 is_directory, is_hidden, is_symlink, parent_path, mime_type, file_hash,
59 indexed_at, last_verified
60 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)
61 ON CONFLICT(path) DO UPDATE SET
62 name = excluded.name,
63 extension = excluded.extension,
64 size = excluded.size,
65 modified_at = excluded.modified_at,
66 accessed_at = excluded.accessed_at,
67 is_directory = excluded.is_directory,
68 is_hidden = excluded.is_hidden,
69 is_symlink = excluded.is_symlink,
70 mime_type = excluded.mime_type,
71 file_hash = excluded.file_hash,
72 last_verified = excluded.last_verified
73 "#,
74 params![
75 file.path.to_string_lossy().to_string(),
76 file.name,
77 file.extension,
78 file.size as i64,
79 created_at,
80 modified_at,
81 accessed_at,
82 file.is_directory as i32,
83 file.is_hidden as i32,
84 file.is_symlink as i32,
85 file.parent_path.as_ref().map(|p| p.to_string_lossy().to_string()),
86 file.mime_type,
87 file.file_hash,
88 indexed_at,
89 last_verified,
90 ],
91 )?;
92
93 Ok(conn.last_insert_rowid())
94 }
95
96 pub fn insert_files_batch(&self, files: &[FileEntry]) -> Result<()> {
97 let mut conn = self.pool.get()?;
98 let tx = conn.transaction()?;
99
100 for file in files {
101 let created_at = file.created_at.map(|dt| dt.timestamp());
102 let modified_at = file.modified_at.map(|dt| dt.timestamp());
103 let accessed_at = file.accessed_at.map(|dt| dt.timestamp());
104 let indexed_at = file.indexed_at.timestamp();
105 let last_verified = file.last_verified.timestamp();
106
107 tx.execute(
108 r#"
109 INSERT INTO files (
110 path, name, extension, size, created_at, modified_at, accessed_at,
111 is_directory, is_hidden, is_symlink, parent_path, mime_type, file_hash,
112 indexed_at, last_verified
113 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)
114 ON CONFLICT(path) DO UPDATE SET
115 name = excluded.name,
116 extension = excluded.extension,
117 size = excluded.size,
118 modified_at = excluded.modified_at,
119 accessed_at = excluded.accessed_at,
120 is_directory = excluded.is_directory,
121 is_hidden = excluded.is_hidden,
122 is_symlink = excluded.is_symlink,
123 mime_type = excluded.mime_type,
124 file_hash = excluded.file_hash,
125 last_verified = excluded.last_verified
126 "#,
127 params![
128 file.path.to_string_lossy().to_string(),
129 file.name,
130 file.extension,
131 file.size as i64,
132 created_at,
133 modified_at,
134 accessed_at,
135 file.is_directory as i32,
136 file.is_hidden as i32,
137 file.is_symlink as i32,
138 file.parent_path.as_ref().map(|p| p.to_string_lossy().to_string()),
139 file.mime_type,
140 file.file_hash,
141 indexed_at,
142 last_verified,
143 ],
144 )?;
145 }
146
147 tx.commit()?;
148 Ok(())
149 }
150
151 pub fn find_by_path(&self, path: &Path) -> Result<Option<FileEntry>> {
152 let conn = self.pool.get()?;
153
154 let result = conn
155 .query_row(
156 r#"
157 SELECT id, path, name, extension, size, created_at, modified_at, accessed_at,
158 is_directory, is_hidden, is_symlink, parent_path, mime_type, file_hash,
159 indexed_at, last_verified
160 FROM files WHERE path = ?1
161 "#,
162 params![path.to_string_lossy().to_string()],
163 |row| Self::row_to_file_entry(row),
164 )
165 .optional()?;
166
167 Ok(result)
168 }
169
170 pub fn find_by_id(&self, id: i64) -> Result<Option<FileEntry>> {
171 let conn = self.pool.get()?;
172
173 let result = conn
174 .query_row(
175 r#"
176 SELECT id, path, name, extension, size, created_at, modified_at, accessed_at,
177 is_directory, is_hidden, is_symlink, parent_path, mime_type, file_hash,
178 indexed_at, last_verified
179 FROM files WHERE id = ?1
180 "#,
181 params![id],
182 |row| Self::row_to_file_entry(row),
183 )
184 .optional()?;
185
186 Ok(result)
187 }
188
189 pub fn delete_by_path(&self, path: &Path) -> Result<()> {
190 let conn = self.pool.get()?;
191 conn.execute(
192 "DELETE FROM files WHERE path = ?1",
193 params![path.to_string_lossy().to_string()],
194 )?;
195 Ok(())
196 }
197
198 pub fn search_by_name(&self, pattern: &str, limit: usize) -> Result<Vec<FileEntry>> {
199 let conn = self.pool.get()?;
200 let mut stmt = conn.prepare(
201 r#"
202 SELECT id, path, name, extension, size, created_at, modified_at, accessed_at,
203 is_directory, is_hidden, is_symlink, parent_path, mime_type, file_hash,
204 indexed_at, last_verified
205 FROM files WHERE name LIKE ?1 LIMIT ?2
206 "#,
207 )?;
208
209 let files = stmt
210 .query_map(params![format!("%{}%", pattern), limit], |row| {
211 Self::row_to_file_entry(row)
212 })?
213 .collect::<rusqlite::Result<Vec<_>>>()?;
214
215 Ok(files)
216 }
217
218 pub fn search_by_extension(&self, extension: &str, limit: usize) -> Result<Vec<FileEntry>> {
219 let conn = self.pool.get()?;
220 let mut stmt = conn.prepare(
221 r#"
222 SELECT id, path, name, extension, size, created_at, modified_at, accessed_at,
223 is_directory, is_hidden, is_symlink, parent_path, mime_type, file_hash,
224 indexed_at, last_verified
225 FROM files WHERE extension = ?1 LIMIT ?2
226 "#,
227 )?;
228
229 let files = stmt
230 .query_map(params![extension, limit], |row| {
231 Self::row_to_file_entry(row)
232 })?
233 .collect::<rusqlite::Result<Vec<_>>>()?;
234
235 Ok(files)
236 }
237
238 pub fn get_all_files(&self, limit: usize, offset: usize) -> Result<Vec<FileEntry>> {
239 let conn = self.pool.get()?;
240 let mut stmt = conn.prepare(
241 r#"
242 SELECT id, path, name, extension, size, created_at, modified_at, accessed_at,
243 is_directory, is_hidden, is_symlink, parent_path, mime_type, file_hash,
244 indexed_at, last_verified
245 FROM files LIMIT ?1 OFFSET ?2
246 "#,
247 )?;
248
249 let files = stmt
250 .query_map(params![limit, offset], |row| Self::row_to_file_entry(row))?
251 .collect::<rusqlite::Result<Vec<_>>>()?;
252
253 Ok(files)
254 }
255
256 pub fn insert_content(&self, file_id: i64, preview: &ContentPreview) -> Result<()> {
257 let conn = self.pool.get()?;
258
259 conn.execute(
260 r#"
261 INSERT INTO file_contents (file_id, content_preview, word_count, line_count, encoding)
262 VALUES (?1, ?2, ?3, ?4, ?5)
263 ON CONFLICT(file_id) DO UPDATE SET
264 content_preview = excluded.content_preview,
265 word_count = excluded.word_count,
266 line_count = excluded.line_count,
267 encoding = excluded.encoding
268 "#,
269 params![
270 file_id,
271 preview.preview,
272 preview.word_count as i64,
273 preview.line_count as i64,
274 preview.encoding
275 ],
276 )?;
277
278 Ok(())
279 }
280
281 pub fn insert_fts_entry(&self, file_id: i64, name: &str, path: &str, content: &str) -> Result<()> {
282 let conn = self.pool.get()?;
283
284 conn.execute(
285 "INSERT INTO files_fts (file_id, name, path, content) VALUES (?1, ?2, ?3, ?4)",
286 params![file_id, name, path, content],
287 )?;
288
289 Ok(())
290 }
291
292 pub fn search_content(&self, query: &str, limit: usize) -> Result<Vec<i64>> {
293 let conn = self.pool.get()?;
294 let mut stmt = conn.prepare(
295 "SELECT file_id FROM files_fts WHERE files_fts MATCH ?1 LIMIT ?2"
296 )?;
297
298 let file_ids = stmt
299 .query_map(params![query, limit], |row| row.get(0))?
300 .collect::<rusqlite::Result<Vec<_>>>()?;
301
302 Ok(file_ids)
303 }
304
305 pub fn add_exclusion_rule(&self, rule: &ExclusionRule) -> Result<i64> {
306 let conn = self.pool.get()?;
307
308 let rule_type = match rule.rule_type {
309 ExclusionRuleType::Glob => "glob",
310 ExclusionRuleType::Regex => "regex",
311 ExclusionRuleType::Path => "path",
312 };
313
314 conn.execute(
315 "INSERT INTO exclusion_rules (pattern, rule_type, created_at) VALUES (?1, ?2, ?3)",
316 params![rule.pattern, rule_type, Utc::now().timestamp()],
317 )?;
318
319 Ok(conn.last_insert_rowid())
320 }
321
322 pub fn get_exclusion_rules(&self) -> Result<Vec<ExclusionRule>> {
323 let conn = self.pool.get()?;
324 let mut stmt = conn.prepare("SELECT pattern, rule_type FROM exclusion_rules")?;
325
326 let rules = stmt
327 .query_map([], |row| {
328 let pattern: String = row.get(0)?;
329 let rule_type_str: String = row.get(1)?;
330 let rule_type = match rule_type_str.as_str() {
331 "glob" => ExclusionRuleType::Glob,
332 "regex" => ExclusionRuleType::Regex,
333 "path" => ExclusionRuleType::Path,
334 _ => ExclusionRuleType::Glob,
335 };
336
337 Ok(ExclusionRule { pattern, rule_type })
338 })?
339 .collect::<rusqlite::Result<Vec<_>>>()?;
340
341 Ok(rules)
342 }
343
344 pub fn log_access(&self, file_id: i64) -> Result<()> {
345 let conn = self.pool.get()?;
346 conn.execute(
347 "INSERT INTO access_log (file_id, accessed_at) VALUES (?1, ?2)",
348 params![file_id, Utc::now().timestamp()],
349 )?;
350 Ok(())
351 }
352
353 pub fn get_stats(&self) -> Result<IndexStats> {
354 let conn = self.pool.get()?;
355
356 let total_files: i64 = conn.query_row(
357 "SELECT COUNT(*) FROM files WHERE is_directory = 0",
358 [],
359 |row| row.get(0),
360 )?;
361
362 let total_directories: i64 = conn.query_row(
363 "SELECT COUNT(*) FROM files WHERE is_directory = 1",
364 [],
365 |row| row.get(0),
366 )?;
367
368 let total_size: i64 = conn.query_row(
369 "SELECT COALESCE(SUM(size), 0) FROM files WHERE is_directory = 0",
370 [],
371 |row| row.get(0),
372 )?;
373
374 let indexed_files: i64 = conn.query_row(
375 "SELECT COUNT(*) FROM file_contents",
376 [],
377 |row| row.get(0),
378 )?;
379
380 let last_update_ts: Option<i64> = conn
381 .query_row(
382 "SELECT MAX(indexed_at) FROM files",
383 [],
384 |row| row.get(0),
385 )
386 .optional()?
387 .flatten();
388
389 let last_update = last_update_ts
390 .and_then(|ts| Utc.timestamp_opt(ts, 0).single())
391 .unwrap_or_else(Utc::now);
392
393 let index_size = std::fs::metadata(
394 conn.path().ok_or_else(|| {
395 SearchError::Configuration("Cannot get database path".to_string())
396 })?,
397 )
398 .map(|m| m.len())
399 .unwrap_or(0);
400
401 Ok(IndexStats {
402 total_files: total_files as usize,
403 total_directories: total_directories as usize,
404 total_size: total_size as u64,
405 indexed_files: indexed_files as usize,
406 last_update,
407 index_size,
408 })
409 }
410
411 pub fn clear_all(&self) -> Result<()> {
412 let conn = self.pool.get()?;
413 let tx = conn.unchecked_transaction()?;
414
415 tx.execute("DELETE FROM files", [])?;
416 tx.execute("DELETE FROM file_contents", [])?;
417 tx.execute("DELETE FROM files_fts", [])?;
418 tx.execute("DELETE FROM access_log", [])?;
419 tx.execute("DELETE FROM search_history", [])?;
420
421 tx.commit()?;
422 Ok(())
423 }
424
425 pub fn vacuum(&self) -> Result<()> {
426 let conn = self.pool.get()?;
427 conn.execute("VACUUM", [])?;
428 Ok(())
429 }
430
431 fn row_to_file_entry(row: &rusqlite::Row) -> rusqlite::Result<FileEntry> {
432 let id: i64 = row.get(0)?;
433 let path: String = row.get(1)?;
434 let name: String = row.get(2)?;
435 let extension: Option<String> = row.get(3)?;
436 let size: i64 = row.get(4)?;
437 let created_at: Option<i64> = row.get(5)?;
438 let modified_at: Option<i64> = row.get(6)?;
439 let accessed_at: Option<i64> = row.get(7)?;
440 let is_directory: i32 = row.get(8)?;
441 let is_hidden: i32 = row.get(9)?;
442 let is_symlink: i32 = row.get(10)?;
443 let parent_path: Option<String> = row.get(11)?;
444 let mime_type: Option<String> = row.get(12)?;
445 let file_hash: Option<String> = row.get(13)?;
446 let indexed_at: i64 = row.get(14)?;
447 let last_verified: i64 = row.get(15)?;
448
449 Ok(FileEntry {
450 id: Some(id),
451 path: PathBuf::from(path),
452 name,
453 extension,
454 size: size as u64,
455 created_at: created_at.and_then(|ts| Utc.timestamp_opt(ts, 0).single()),
456 modified_at: modified_at.and_then(|ts| Utc.timestamp_opt(ts, 0).single()),
457 accessed_at: accessed_at.and_then(|ts| Utc.timestamp_opt(ts, 0).single()),
458 is_directory: is_directory != 0,
459 is_hidden: is_hidden != 0,
460 is_symlink: is_symlink != 0,
461 parent_path: parent_path.map(PathBuf::from),
462 mime_type,
463 file_hash,
464 indexed_at: Utc.timestamp_opt(indexed_at, 0).single().unwrap_or_else(Utc::now),
465 last_verified: Utc.timestamp_opt(last_verified, 0).single().unwrap_or_else(Utc::now),
466 })
467 }
468}