vein_database/sql/
sqlite.rs

1use anyhow::Result;
2use rusqlite::Connection;
3use std::fs;
4
5use crate::paths::Paths;
6
7struct SQLEntry {
8    id: String,
9    role: String,
10    summary: String,
11    code_style: bool,
12    agent: String,
13    timestamp: String,
14}
15
16pub struct SqliteDb {
17    conn: Connection,
18}
19
20impl SqliteDb {
21    pub fn new() -> Result<Self> {
22        let path_ref = Paths::get_dialogues_db();
23        if !path_ref.exists() {
24            fs::create_dir_all(&path_ref)?;
25        }
26        let db_path = path_ref.join("dialogues.db");
27        let conn = Connection::open(&db_path)?;
28        conn.execute(
29            "CREATE TABLE IF NOT EXISTS medium (
30                id TEXT PRIMARY KEY,
31                role TEXT NOT NULL,
32                summary TEXT NOT NULL,
33                code_style INTEGER NOT NULL,
34                agent TEXT NOT NULL,
35                timestamp TEXT NOT NULL
36            )",
37            [],
38        )?;
39        conn.execute(
40            "CREATE UNIQUE INDEX IF NOT EXISTS idx_medium_summary ON medium(summary)",
41            [],
42        )?;
43        Ok(Self { conn })
44    }
45
46    pub fn create(
47        &self,
48        id: &str,
49        role: &str,
50        summary: &str,
51        code_style: bool,
52        agent: &str,
53        timestamp: &str,
54    ) -> Result<String> {
55        self.conn.execute(
56            "INSERT INTO medium (id, role, summary, code_style, agent, timestamp) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
57            (id, role, summary, code_style as i32, agent, timestamp),
58        )?;
59        Ok(id.to_string())
60    }
61
62    pub fn upsert(
63        &self,
64        id: &str,
65        role: &str,
66        summary: &str,
67        code_style: bool,
68        agent: &str,
69        timestamp: &str,
70    ) -> Result<String> {
71        self.conn.execute(
72            "INSERT OR IGNORE INTO medium (id, role, summary, code_style, agent, timestamp) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
73            (id, role, summary, code_style as i32, agent, timestamp),
74        )?;
75        Ok(id.to_string())
76    }
77
78    pub fn update(
79        &self,
80        id: &str,
81        role: &str,
82        summary: &str,
83        code_style: bool,
84        agent: &str,
85        timestamp: &str,
86    ) -> Result<bool> {
87        let rows = self.conn.execute(
88            "UPDATE medium SET role = ?1, summary = ?2, code_style = ?3, agent = ?4, timestamp = ?5 WHERE id = ?6",
89            (role, summary, code_style as i32, agent, timestamp, id),
90        )?;
91        Ok(rows > 0)
92    }
93
94    pub fn delete(&self, id: &str) -> Result<bool> {
95        let rows = self.conn.execute("DELETE FROM medium WHERE id = ?1", [id])?;
96        Ok(rows > 0)
97    }
98
99    pub fn query(&self, query_str: &str) -> Result<String> {
100        let mut stmt = self.conn.prepare(query_str)?;
101        let rows = stmt.query_map([], |row| {
102            Ok(SQLEntry {
103                id: row.get(0)?,
104                role: row.get(1)?,
105                summary: row.get(2)?,
106                code_style: row.get::<_, i32>(3)? != 0,
107                agent: row.get(4)?,
108                timestamp: row.get(5)?,
109            })
110        })?;
111
112        let results: Vec<SQLEntry> = rows.filter_map(|r| r.ok()).collect();
113        if results.is_empty() {
114            return Ok("No results found".to_string());
115        }
116
117        let mut output = String::new();
118        for entry in results {
119            output.push_str(&format!(
120                "id: {}\nrole: {}\nsummary: {}\ncode_style: {}\nagent: {}\ntimestamp: {}\n---\n",
121                entry.id, entry.role, entry.summary, entry.code_style, entry.agent, entry.timestamp
122            ));
123        }
124        Ok(output)
125    }
126
127    pub fn count_all(&self) -> Result<usize> {
128        let count: i64 = self
129            .conn
130            .query_row("SELECT COUNT(*) FROM medium", [], |row| row.get(0))?;
131        Ok(count as usize)
132    }
133}
vein_database/sql/sqlite.rs

vein_database/sql/
sqlite.rs
