Skip to main content

obsidian_cli_inspector/
db.rs

1use rusqlite::{Connection, OptionalExtension, Result, Transaction};
2use std::path::Path;
3
4mod operations;
5mod schema;
6mod stats;
7
8pub use stats::DatabaseStats;
9
10pub const SCHEMA_VERSION: i32 = 1;
11
12#[derive(Debug, Clone)]
13pub struct NoteMetadata {
14    pub id: i64,
15    pub mtime: i64,
16    pub hash: String,
17}
18
19pub struct DatabaseTransaction<'a> {
20    tx: Transaction<'a>,
21}
22
23pub struct Database {
24    conn: Connection,
25}
26
27impl Database {
28    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
29        let conn = Connection::open(path)?;
30        Ok(Database { conn })
31    }
32
33    pub fn transaction(&mut self) -> Result<DatabaseTransaction<'_>> {
34        Ok(DatabaseTransaction {
35            tx: self.conn.transaction()?,
36        })
37    }
38
39    pub fn initialize(&self, force: bool) -> Result<()> {
40        if force {
41            schema::drop_tables(&self.conn)?;
42        }
43
44        // Create schema version table
45        self.conn.execute(
46            "CREATE TABLE IF NOT EXISTS schema_version (
47                version INTEGER PRIMARY KEY,
48                applied_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
49            )",
50            [],
51        )?;
52
53        // Check current version
54        let current_version: Option<i32> = self
55            .conn
56            .query_row(
57                "SELECT version FROM schema_version ORDER BY version DESC LIMIT 1",
58                [],
59                |row| row.get(0),
60            )
61            .ok();
62
63        if current_version.is_none() || force {
64            schema::create_schema(&self.conn)?;
65            self.conn.execute(
66                "INSERT OR REPLACE INTO schema_version (version) VALUES (?1)",
67                [SCHEMA_VERSION],
68            )?;
69        }
70
71        Ok(())
72    }
73
74    pub fn get_version(&self) -> Result<Option<i32>> {
75        self.conn
76            .query_row(
77                "SELECT version FROM schema_version ORDER BY version DESC LIMIT 1",
78                [],
79                |row| row.get(0),
80            )
81            .optional()
82    }
83
84    pub fn get_stats(&self) -> Result<DatabaseStats> {
85        stats::get_stats(&self.conn)
86    }
87
88    pub fn insert_note(
89        &self,
90        path: &str,
91        title: &str,
92        mtime: u64,
93        hash: &str,
94        frontmatter_json: Option<&str>,
95    ) -> Result<i64> {
96        operations::insert_note(&self.conn, path, title, mtime, hash, frontmatter_json)
97    }
98
99    pub fn get_note_by_path(&self, path: &str) -> Result<Option<i64>> {
100        operations::get_note_by_path(&self.conn, path)
101    }
102
103    pub fn get_note_metadata_by_path(&self, path: &str) -> Result<Option<NoteMetadata>> {
104        operations::get_note_metadata_by_path(&self.conn, path)
105    }
106
107    pub fn insert_tag(&self, note_id: i64, tag: &str) -> Result<()> {
108        operations::insert_tag(&self.conn, note_id, tag)
109    }
110
111    #[allow(clippy::too_many_arguments)]
112    pub fn insert_link(
113        &self,
114        src_note_id: i64,
115        dst_text: &str,
116        kind: &str,
117        is_embed: bool,
118        alias: Option<&str>,
119        heading_ref: Option<&str>,
120        block_ref: Option<&str>,
121    ) -> Result<()> {
122        operations::insert_link(
123            &self.conn,
124            src_note_id,
125            dst_text,
126            kind,
127            is_embed,
128            alias,
129            heading_ref,
130            block_ref,
131        )
132    }
133
134    pub fn insert_chunk(&self, note_id: i64, heading_path: Option<&str>, text: &str) -> Result<()> {
135        operations::insert_chunk(&self.conn, note_id, heading_path, text)
136    }
137
138    pub fn insert_chunk_with_offset(
139        &self,
140        note_id: i64,
141        heading_path: Option<&str>,
142        text: &str,
143        byte_offset: i32,
144        byte_length: i32,
145    ) -> Result<()> {
146        operations::insert_chunk_with_offset(
147            &self.conn,
148            note_id,
149            heading_path,
150            text,
151            byte_offset,
152            byte_length,
153        )
154    }
155
156    pub fn clear_note_data(&self, note_id: i64) -> Result<()> {
157        operations::clear_note_data(&self.conn, note_id)
158    }
159
160    /// Execute a query function with access to the database connection
161    pub fn conn(&self) -> DatabaseQueryExecutor<'_> {
162        DatabaseQueryExecutor { conn: &self.conn }
163    }
164}
165
166impl DatabaseTransaction<'_> {
167    pub fn insert_note(
168        &self,
169        path: &str,
170        title: &str,
171        mtime: u64,
172        hash: &str,
173        frontmatter_json: Option<&str>,
174    ) -> Result<i64> {
175        operations::insert_note(&self.tx, path, title, mtime, hash, frontmatter_json)
176    }
177
178    pub fn get_note_metadata_by_path(&self, path: &str) -> Result<Option<NoteMetadata>> {
179        operations::get_note_metadata_by_path(&self.tx, path)
180    }
181
182    pub fn insert_tag(&self, note_id: i64, tag: &str) -> Result<()> {
183        operations::insert_tag(&self.tx, note_id, tag)
184    }
185
186    #[allow(clippy::too_many_arguments)]
187    pub fn insert_link(
188        &self,
189        src_note_id: i64,
190        dst_text: &str,
191        kind: &str,
192        is_embed: bool,
193        alias: Option<&str>,
194        heading_ref: Option<&str>,
195        block_ref: Option<&str>,
196    ) -> Result<()> {
197        operations::insert_link(
198            &self.tx,
199            src_note_id,
200            dst_text,
201            kind,
202            is_embed,
203            alias,
204            heading_ref,
205            block_ref,
206        )
207    }
208
209    pub fn insert_chunk_with_offset(
210        &self,
211        note_id: i64,
212        heading_path: Option<&str>,
213        text: &str,
214        byte_offset: i32,
215        byte_length: i32,
216    ) -> Result<()> {
217        operations::insert_chunk_with_offset(
218            &self.tx,
219            note_id,
220            heading_path,
221            text,
222            byte_offset,
223            byte_length,
224        )
225    }
226
227    pub fn clear_note_data(&self, note_id: i64) -> Result<()> {
228        operations::clear_note_data(&self.tx, note_id)
229    }
230
231    pub fn commit(self) -> Result<()> {
232        self.tx.commit()
233    }
234}
235
236pub struct DatabaseQueryExecutor<'a> {
237    conn: &'a Connection,
238}
239
240impl DatabaseQueryExecutor<'_> {
241    pub fn execute_query<T, F>(&self, f: F) -> Result<T>
242    where
243        F: FnOnce(&Connection) -> Result<T>,
244    {
245        f(self.conn)
246    }
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252    use tempfile::TempDir;
253
254    #[test]
255    fn test_database_open() {
256        let temp_dir = TempDir::new().unwrap();
257        let db_path = temp_dir.path().join("test.db");
258
259        let _db = Database::open(&db_path).unwrap();
260        assert!(db_path.exists());
261    }
262
263    #[test]
264    fn test_database_initialize() {
265        let temp_dir = TempDir::new().unwrap();
266        let db_path = temp_dir.path().join("test.db");
267
268        let db = Database::open(&db_path).unwrap();
269        db.initialize(false).unwrap();
270
271        let version = db.get_version().unwrap();
272        assert_eq!(version, Some(1));
273    }
274
275    #[test]
276    fn test_database_initialize_force() {
277        let temp_dir = TempDir::new().unwrap();
278        let db_path = temp_dir.path().join("test.db");
279
280        let db = Database::open(&db_path).unwrap();
281        db.initialize(false).unwrap();
282        let version1 = db.get_version().unwrap();
283
284        // Force reinitialize
285        db.initialize(true).unwrap();
286        let version2 = db.get_version().unwrap();
287
288        assert_eq!(version1, version2);
289    }
290
291    #[test]
292    fn test_database_transaction() {
293        let temp_dir = TempDir::new().unwrap();
294        let db_path = temp_dir.path().join("test.db");
295
296        let mut db = Database::open(&db_path).unwrap();
297        db.initialize(false).unwrap();
298
299        let tx = db.transaction().unwrap();
300        tx.commit().unwrap();
301    }
302
303    #[test]
304    fn test_note_metadata_creation() {
305        let metadata = NoteMetadata {
306            id: 1,
307            mtime: 1234567890,
308            hash: "abc123".to_string(),
309        };
310
311        assert_eq!(metadata.id, 1);
312        assert_eq!(metadata.mtime, 1234567890);
313        assert_eq!(metadata.hash, "abc123");
314    }
315
316    #[test]
317    fn test_database_stats_creation() {
318        let stats = DatabaseStats {
319            note_count: 10,
320            link_count: 20,
321            tag_count: 5,
322            chunk_count: 100,
323            unresolved_links: 2,
324        };
325
326        assert_eq!(stats.note_count, 10);
327        assert_eq!(stats.link_count, 20);
328        assert_eq!(stats.tag_count, 5);
329        assert_eq!(stats.chunk_count, 100);
330        assert_eq!(stats.unresolved_links, 2);
331    }
332
333    #[test]
334    fn test_schema_version_constant() {
335        assert_eq!(SCHEMA_VERSION, 1);
336    }
337}