Skip to main content

agent_ask/
store.rs

1//! SQLite-backed artifact store (mirror of `src/store.ts`).
2
3use std::sync::{Arc, Mutex};
4
5use rusqlite::{params, Connection};
6use serde_json::Value;
7
8use crate::artifact::cid_of;
9use crate::error::Error;
10
11const SCHEMA: &str = r#"
12CREATE TABLE IF NOT EXISTS artifacts (
13  cid TEXT PRIMARY KEY,
14  kind TEXT NOT NULL,
15  created_at TEXT NOT NULL,
16  body TEXT NOT NULL
17);
18CREATE INDEX IF NOT EXISTS idx_artifacts_kind_created ON artifacts(kind, created_at);
19CREATE INDEX IF NOT EXISTS idx_artifacts_created ON artifacts(created_at);
20CREATE TABLE IF NOT EXISTS question_tags (
21  cid TEXT NOT NULL,
22  tag TEXT NOT NULL,
23  PRIMARY KEY (cid, tag),
24  FOREIGN KEY (cid) REFERENCES artifacts(cid) ON DELETE CASCADE
25);
26CREATE INDEX IF NOT EXISTS idx_question_tags_tag ON question_tags(tag);
27"#;
28
29#[derive(Clone)]
30pub struct Store {
31    conn: Arc<Mutex<Connection>>,
32}
33
34impl Store {
35    pub fn open(path: &str) -> Result<Self, Error> {
36        let conn = Connection::open(path)?;
37        conn.execute_batch("PRAGMA journal_mode = WAL; PRAGMA foreign_keys = ON;")?;
38        conn.execute_batch(SCHEMA)?;
39        Ok(Self { conn: Arc::new(Mutex::new(conn)) })
40    }
41
42    pub fn insert_artifact(&self, artifact: &Value) -> Result<String, Error> {
43        let cid = cid_of(artifact)?;
44        let kind = artifact.get("kind").and_then(|v| v.as_str())
45            .ok_or_else(|| Error::Invalid("missing kind".into()))?
46            .to_string();
47        let created_at = artifact.get("created_at").and_then(|v| v.as_str())
48            .ok_or_else(|| Error::Invalid("missing created_at".into()))?
49            .to_string();
50        let body = serde_json::to_string(artifact)?;
51        let conn = self.conn.lock().unwrap();
52        conn.execute(
53            "INSERT OR IGNORE INTO artifacts(cid, kind, created_at, body) VALUES (?1, ?2, ?3, ?4)",
54            params![&cid, &kind, &created_at, &body],
55        )?;
56        if kind == "question" {
57            if let Some(tags) = artifact.get("tags").and_then(|v| v.as_array()) {
58                for t in tags {
59                    if let Some(tag) = t.as_str() {
60                        conn.execute(
61                            "INSERT OR IGNORE INTO question_tags(cid, tag) VALUES (?1, ?2)",
62                            params![&cid, tag],
63                        )?;
64                    }
65                }
66            }
67        }
68        Ok(cid)
69    }
70
71    pub fn get_artifact(&self, cid: &str) -> Result<Option<Value>, Error> {
72        let conn = self.conn.lock().unwrap();
73        let mut stmt = conn.prepare("SELECT body FROM artifacts WHERE cid = ?1")?;
74        let mut rows = stmt.query(params![cid])?;
75        if let Some(row) = rows.next()? {
76            let body: String = row.get(0)?;
77            Ok(Some(serde_json::from_str(&body)?))
78        } else {
79            Ok(None)
80        }
81    }
82
83    pub fn has_artifact(&self, cid: &str) -> Result<bool, Error> {
84        let conn = self.conn.lock().unwrap();
85        let mut stmt = conn.prepare("SELECT 1 FROM artifacts WHERE cid = ?1")?;
86        Ok(stmt.exists(params![cid])?)
87    }
88
89    pub fn list_questions(
90        &self,
91        tag: Option<&str>,
92        since: Option<&str>,
93        limit: i64,
94    ) -> Result<Vec<Value>, Error> {
95        let conn = self.conn.lock().unwrap();
96        let mut clauses: Vec<&str> = vec!["a.kind = 'question'"];
97        let mut params_vec: Vec<rusqlite::types::Value> = Vec::new();
98        if let Some(t) = tag {
99            clauses.push("qt.tag = ?");
100            params_vec.push(t.to_string().into());
101        }
102        if let Some(s) = since {
103            clauses.push("a.created_at > ?");
104            params_vec.push(s.to_string().into());
105        }
106        let join = if tag.is_some() { "JOIN question_tags qt ON qt.cid = a.cid" } else { "" };
107        let sql = format!(
108            "SELECT DISTINCT a.body FROM artifacts a {join} WHERE {} ORDER BY a.created_at DESC LIMIT ?",
109            clauses.join(" AND ")
110        );
111        params_vec.push(limit.into());
112
113        let mut stmt = conn.prepare(&sql)?;
114        let rows = stmt.query_map(
115            rusqlite::params_from_iter(params_vec.iter()),
116            |row| row.get::<_, String>(0),
117        )?;
118        let mut out = Vec::new();
119        for row in rows {
120            out.push(serde_json::from_str(&row?)?);
121        }
122        Ok(out)
123    }
124
125    pub fn stream_feed(&self, since: Option<&str>, limit: i64) -> Result<Vec<Value>, Error> {
126        let conn = self.conn.lock().unwrap();
127        let mut params_vec: Vec<rusqlite::types::Value> = Vec::new();
128        let where_clause = if let Some(s) = since {
129            params_vec.push(s.to_string().into());
130            "WHERE created_at > ?"
131        } else {
132            ""
133        };
134        let sql = format!(
135            "SELECT body FROM artifacts {where_clause} ORDER BY created_at DESC LIMIT ?"
136        );
137        params_vec.push(limit.into());
138
139        let mut stmt = conn.prepare(&sql)?;
140        let rows = stmt.query_map(
141            rusqlite::params_from_iter(params_vec.iter()),
142            |row| row.get::<_, String>(0),
143        )?;
144        let mut out = Vec::new();
145        for row in rows {
146            out.push(serde_json::from_str(&row?)?);
147        }
148        Ok(out)
149    }
150}