Skip to main content

directory_indexer/storage/
sqlite.rs

1use log::info;
2use rusqlite::Connection;
3use serde_json::Value;
4use std::path::Path;
5
6use crate::{error::Result, utils::normalize_path};
7
8pub struct SqliteStore {
9    conn: Connection,
10}
11
12#[derive(Debug, Clone)]
13pub struct FileRecord {
14    pub id: i64,
15    pub path: String,
16    pub size: i64,
17    pub modified_time: i64,
18    pub hash: String,
19    pub parent_dirs: Vec<String>,
20    pub chunks_json: Option<Value>,
21    pub errors_json: Option<Value>,
22}
23
24#[derive(Debug, Clone)]
25pub struct DirectoryRecord {
26    pub id: i64,
27    pub path: String,
28    pub status: String,
29    pub indexed_at: i64,
30}
31
32impl SqliteStore {
33    pub fn new<P: AsRef<Path>>(db_path: P) -> Result<Self> {
34        let conn = Connection::open(db_path)?;
35        let store = SqliteStore { conn };
36        store.initialize_schema()?;
37        Ok(store)
38    }
39
40    fn initialize_schema(&self) -> Result<()> {
41        self.conn.execute(
42            "CREATE TABLE IF NOT EXISTS directories (
43                id INTEGER PRIMARY KEY AUTOINCREMENT,
44                path TEXT UNIQUE NOT NULL,
45                status TEXT NOT NULL DEFAULT 'pending',
46                indexed_at INTEGER NOT NULL DEFAULT 0
47            )",
48            [],
49        )?;
50
51        self.conn.execute(
52            "CREATE TABLE IF NOT EXISTS files (
53                id INTEGER PRIMARY KEY AUTOINCREMENT,
54                path TEXT UNIQUE NOT NULL,
55                size INTEGER NOT NULL,
56                modified_time INTEGER NOT NULL,
57                hash TEXT NOT NULL,
58                parent_dirs TEXT NOT NULL,
59                chunks_json TEXT,
60                errors_json TEXT
61            )",
62            [],
63        )?;
64
65        self.conn.execute(
66            "CREATE INDEX IF NOT EXISTS idx_files_path ON files(path)",
67            [],
68        )?;
69
70        self.conn.execute(
71            "CREATE INDEX IF NOT EXISTS idx_files_parent_dirs ON files(parent_dirs)",
72            [],
73        )?;
74
75        Ok(())
76    }
77
78    pub fn add_directory(&self, path: &str) -> Result<i64> {
79        let normalized_path = normalize_path(path)?;
80        let mut stmt = self.conn.prepare(
81            "INSERT OR REPLACE INTO directories (path, status, indexed_at) 
82             VALUES (?1, 'pending', strftime('%s', 'now'))",
83        )?;
84
85        stmt.execute([&normalized_path])?;
86        Ok(self.conn.last_insert_rowid())
87    }
88
89    pub fn update_directory_status(&self, path: &str, status: &str) -> Result<()> {
90        let normalized_path = normalize_path(path)?;
91        let mut stmt = self.conn.prepare(
92            "UPDATE directories SET status = ?1, indexed_at = strftime('%s', 'now') WHERE path = ?2"
93        )?;
94
95        stmt.execute([status, &normalized_path])?;
96        Ok(())
97    }
98
99    pub fn get_directories(&self) -> Result<Vec<DirectoryRecord>> {
100        let mut stmt = self
101            .conn
102            .prepare("SELECT id, path, status, indexed_at FROM directories ORDER BY path")?;
103
104        let rows = stmt.query_map([], |row| {
105            Ok(DirectoryRecord {
106                id: row.get(0)?,
107                path: row.get(1)?,
108                status: row.get(2)?,
109                indexed_at: row.get(3)?,
110            })
111        })?;
112
113        let mut directories = Vec::new();
114        for row in rows {
115            directories.push(row?);
116        }
117
118        Ok(directories)
119    }
120
121    pub fn add_file(&self, record: &FileRecord) -> Result<i64> {
122        let parent_dirs_json = serde_json::to_string(&record.parent_dirs)?;
123        let chunks_json = record
124            .chunks_json
125            .as_ref()
126            .map(serde_json::to_string)
127            .transpose()?;
128        let errors_json = record
129            .errors_json
130            .as_ref()
131            .map(serde_json::to_string)
132            .transpose()?;
133
134        let mut stmt = self.conn.prepare(
135            "INSERT OR REPLACE INTO files 
136             (path, size, modified_time, hash, parent_dirs, chunks_json, errors_json)
137             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
138        )?;
139
140        stmt.execute(rusqlite::params![
141            record.path,
142            record.size,
143            record.modified_time,
144            record.hash,
145            parent_dirs_json,
146            chunks_json,
147            errors_json
148        ])?;
149
150        Ok(self.conn.last_insert_rowid())
151    }
152
153    pub fn get_file_by_path(&self, path: &str) -> Result<Option<FileRecord>> {
154        let normalized_path = normalize_path(path)?;
155        let mut stmt = self.conn.prepare(
156            "SELECT id, path, size, modified_time, hash, parent_dirs, chunks_json, errors_json 
157             FROM files WHERE path = ?1",
158        )?;
159
160        let mut rows = stmt.query_map([&normalized_path], |row| {
161            let parent_dirs_str: String = row.get(5)?;
162            let parent_dirs: Vec<String> = serde_json::from_str(&parent_dirs_str).map_err(|e| {
163                rusqlite::Error::FromSqlConversionFailure(
164                    5,
165                    rusqlite::types::Type::Text,
166                    Box::new(e),
167                )
168            })?;
169
170            let chunks_json: Option<String> = row.get(6)?;
171            let chunks = chunks_json
172                .filter(|s| !s.is_empty())
173                .map(|s| serde_json::from_str(&s))
174                .transpose()
175                .map_err(|e| {
176                    rusqlite::Error::FromSqlConversionFailure(
177                        6,
178                        rusqlite::types::Type::Text,
179                        Box::new(e),
180                    )
181                })?;
182
183            let errors_json: Option<String> = row.get(7)?;
184            let errors = errors_json
185                .filter(|s| !s.is_empty())
186                .map(|s| serde_json::from_str(&s))
187                .transpose()
188                .map_err(|e| {
189                    rusqlite::Error::FromSqlConversionFailure(
190                        7,
191                        rusqlite::types::Type::Text,
192                        Box::new(e),
193                    )
194                })?;
195
196            Ok(FileRecord {
197                id: row.get(0)?,
198                path: row.get(1)?,
199                size: row.get(2)?,
200                modified_time: row.get(3)?,
201                hash: row.get(4)?,
202                parent_dirs,
203                chunks_json: chunks,
204                errors_json: errors,
205            })
206        })?;
207
208        match rows.next() {
209            Some(row) => Ok(Some(row?)),
210            None => Ok(None),
211        }
212    }
213
214    pub fn delete_file(&self, path: &str) -> Result<()> {
215        let normalized_path = normalize_path(path)?;
216        let mut stmt = self.conn.prepare("DELETE FROM files WHERE path = ?1")?;
217        stmt.execute([&normalized_path])?;
218        Ok(())
219    }
220
221    pub fn clear_all_files(&self) -> Result<()> {
222        info!("Clearing all file tracking data from SQLite");
223        let mut stmt = self.conn.prepare("DELETE FROM files")?;
224        stmt.execute([])?;
225        Ok(())
226    }
227
228    pub fn get_stats(&self) -> Result<(i64, i64, i64)> {
229        let directory_count: i64 =
230            self.conn
231                .query_row("SELECT COUNT(*) FROM directories", [], |row| row.get(0))?;
232
233        let file_count: i64 = self
234            .conn
235            .query_row("SELECT COUNT(*) FROM files", [], |row| row.get(0))?;
236
237        let chunk_count: i64 = self.conn.query_row(
238            "SELECT COUNT(*) FROM files WHERE chunks_json IS NOT NULL",
239            [],
240            |row| row.get(0),
241        )?;
242
243        Ok((directory_count, file_count, chunk_count))
244    }
245}