Skip to main content

rag_rat_core/
storage.rs

1use std::{
2    fs,
3    path::{Path, PathBuf},
4};
5
6use rusqlite::Connection;
7use serde::Serialize;
8
9#[derive(Debug, Clone, Serialize)]
10pub struct StorageStatus {
11    pub backend: &'static str,
12    pub sqlite_version: String,
13    pub fts5_available: bool,
14}
15
16#[derive(Debug)]
17pub struct IndexConnection {
18    conn: Connection,
19    database_path: PathBuf,
20    source_root: Option<PathBuf>,
21}
22
23impl IndexConnection {
24    pub fn open(path: &Path) -> anyhow::Result<Self> {
25        if let Some(parent) = path.parent() {
26            fs::create_dir_all(parent)?;
27        }
28        let conn = Connection::open(path)?;
29        let storage = Self { conn, database_path: path.to_path_buf(), source_root: None };
30        storage.setup()?;
31        Ok(storage)
32    }
33
34    pub fn database_path(&self) -> &Path {
35        &self.database_path
36    }
37
38    pub fn connection(&self) -> &Connection {
39        &self.conn
40    }
41
42    pub fn source_root(&self) -> Option<&Path> {
43        self.source_root.as_deref()
44    }
45
46    pub fn set_source_root(&mut self, source_root: PathBuf) {
47        self.source_root = Some(source_root);
48    }
49
50    pub fn execute_batch(&self, sql: &str) -> anyhow::Result<()> {
51        self.conn.execute_batch(sql)?;
52        Ok(())
53    }
54
55    pub fn status(&self) -> anyhow::Result<StorageStatus> {
56        let sqlite_version =
57            self.conn.query_row("SELECT sqlite_version()", [], |row| row.get::<_, String>(0))?;
58        Ok(StorageStatus {
59            backend: "sqlite",
60            sqlite_version,
61            fts5_available: self.fts5_available(),
62        })
63    }
64
65    fn setup(&self) -> anyhow::Result<()> {
66        self.conn.execute_batch(
67            "
68            PRAGMA foreign_keys = ON;
69            PRAGMA journal_mode = WAL;
70            PRAGMA synchronous = NORMAL;
71            ",
72        )?;
73        Ok(())
74    }
75
76    fn fts5_available(&self) -> bool {
77        self.conn
78            .execute_batch(
79                "
80                CREATE VIRTUAL TABLE temp.rag_rat_fts_probe USING fts5(text);
81                DROP TABLE temp.rag_rat_fts_probe;
82                ",
83            )
84            .is_ok()
85    }
86}