agentdb/
db.rs

1use crate::conversations::ConversationStore;
2use crate::error::Result;
3use crate::fts::FullTextStore;
4use crate::hybrid::{HybridQuery, HybridResult, HybridStore};
5use crate::memory::MemoryGraph;
6use crate::schema;
7use crate::traces::TraceStore;
8use crate::vectors::VectorStore;
9use crate::workflows::WorkflowStore;
10use rusqlite::Connection;
11use std::sync::atomic::{AtomicBool, Ordering};
12use std::sync::{Arc, Mutex};
13
14/// The main AgentDB connection — your single-file AI database.
15#[derive(Clone)]
16pub struct AgentDB {
17    conn: Arc<Mutex<Connection>>,
18    closed: Arc<AtomicBool>,
19}
20
21impl AgentDB {
22    /// Open or create an AgentDB database. Use `":memory:"` for tests.
23    pub fn open(path: &str) -> Result<Self> {
24        let conn = Connection::open(path)?;
25        schema::bootstrap(&conn)?;
26        schema::check_version(&conn)?;
27        Ok(Self {
28            conn: Arc::new(Mutex::new(conn)),
29            closed: Arc::new(AtomicBool::new(false)),
30        })
31    }
32
33    /// Access the vector store layer
34    pub fn vectors(&self) -> VectorStore {
35        VectorStore::new(Arc::clone(&self.conn))
36    }
37
38    /// Access the memory graph layer
39    pub fn memory(&self) -> MemoryGraph {
40        MemoryGraph::new(Arc::clone(&self.conn))
41    }
42
43    /// Access the full-text search layer
44    pub fn fts(&self) -> FullTextStore {
45        FullTextStore::new(Arc::clone(&self.conn))
46    }
47
48    /// Access the conversation / message-threading layer
49    pub fn conversations(&self) -> ConversationStore {
50        ConversationStore::new(Arc::clone(&self.conn))
51    }
52
53    /// Access the workflow persistence layer
54    pub fn workflows(&self) -> WorkflowStore {
55        WorkflowStore::new(Arc::clone(&self.conn))
56    }
57
58    /// Access the reasoning-trace layer
59    pub fn traces(&self) -> TraceStore {
60        TraceStore::new(Arc::clone(&self.conn))
61    }
62
63    /// Run a hybrid graph + vector query
64    pub fn hybrid_query(&self, q: HybridQuery) -> Result<Vec<HybridResult>> {
65        let dim: usize = {
66            let conn = self.conn.lock().unwrap();
67            conn.query_row(
68                "SELECT dim FROM _adb_collections WHERE name = ?1",
69                rusqlite::params![q.collection],
70                |r| r.get::<_, i64>(0).map(|v| v as usize),
71            )
72            .unwrap_or(q.embedding.len())
73        };
74        let col = self.vectors().collection(q.collection, dim)?;
75        let store = HybridStore::new(Arc::clone(&self.conn));
76        store.query(q, &col)
77    }
78
79    /// Execute a raw SQL statement
80    pub fn execute(&self, sql: &str) -> Result<usize> {
81        let conn = self.conn.lock().unwrap();
82        Ok(conn.execute(sql, [])?)
83    }
84
85    /// Execute a parameterized SQL statement
86    pub fn execute_params(&self, sql: &str, params: &[&dyn rusqlite::ToSql]) -> Result<usize> {
87        let conn = self.conn.lock().unwrap();
88        Ok(conn.execute(sql, params)?)
89    }
90
91    /// Run multiple operations atomically inside a single SQLite transaction.
92    ///
93    /// The closure receives a [`rusqlite::Transaction`] and may perform any
94    /// number of reads or writes.  If the closure returns `Ok`, the transaction
95    /// is committed; if it returns `Err` (or panics), the transaction is rolled
96    /// back automatically.
97    ///
98    /// # Example
99    /// ```rust,no_run
100    /// # use agentdb::AgentDB;
101    /// let db = AgentDB::open(":memory:").unwrap();
102    /// db.transaction(|tx| {
103    ///     tx.execute("INSERT INTO _adb_nodes (id, kind, data, created_at, updated_at) VALUES ('x','tag','{}',0,0)", [])?;
104    ///     tx.execute("INSERT INTO _adb_nodes (id, kind, data, created_at, updated_at) VALUES ('y','tag','{}',0,0)", [])?;
105    ///     Ok(())
106    /// }).unwrap();
107    /// ```
108    pub fn transaction<F, T>(&self, f: F) -> Result<T>
109    where
110        F: FnOnce(&rusqlite::Transaction) -> Result<T>,
111    {
112        let mut conn = self.conn.lock().unwrap();
113        let tx = conn.transaction()?;
114        let result = f(&tx)?;
115        tx.commit()?;
116        Ok(result)
117    }
118
119    /// Execute one or more semicolon-separated SQL statements as a single
120    /// atomic batch.  This is a convenience wrapper around
121    /// [`execute_batch`](rusqlite::Connection::execute_batch) that wraps the
122    /// statements in an explicit transaction so partial execution is never
123    /// visible to other threads.
124    ///
125    /// # Example
126    /// ```rust,no_run
127    /// # use agentdb::AgentDB;
128    /// let db = AgentDB::open(":memory:").unwrap();
129    /// db.execute_batch(
130    ///     "INSERT INTO _adb_nodes (id,kind,data,created_at,updated_at) VALUES ('a','t','{}',0,0);
131    ///      INSERT INTO _adb_nodes (id,kind,data,created_at,updated_at) VALUES ('b','t','{}',0,0);"
132    /// ).unwrap();
133    /// ```
134    pub fn execute_batch(&self, sql: &str) -> Result<()> {
135        let mut conn = self.conn.lock().unwrap();
136        let tx = conn.transaction()?;
137        tx.execute_batch(sql)?;
138        tx.commit()?;
139        Ok(())
140    }
141
142    /// Query and return rows as JSON values
143    pub fn query_json(&self, sql: &str) -> Result<Vec<serde_json::Value>> {
144        let conn = self.conn.lock().unwrap();
145        let mut stmt = conn.prepare(sql)?;
146        let col_names: Vec<String> = stmt.column_names().iter().map(|s| s.to_string()).collect();
147        let rows = stmt.query_map([], |row| {
148            let mut map = serde_json::Map::new();
149            for (i, name) in col_names.iter().enumerate() {
150                let val: rusqlite::types::Value = row.get(i)?;
151                map.insert(name.clone(), rusqlite_value_to_json(val));
152            }
153            Ok(serde_json::Value::Object(map))
154        })?;
155        rows.map(|r| r.map_err(crate::error::AgentDbError::Sqlite))
156            .collect()
157    }
158
159    /// Query with parameters and return rows as JSON values.
160    ///
161    /// Use this for any query involving user-supplied values to prevent SQL injection.
162    ///
163    /// # Example
164    /// ```rust,no_run
165    /// # use agentdb::AgentDB;
166    /// let db = AgentDB::open(":memory:").unwrap();
167    /// let rows = db.query_json_params(
168    ///     "SELECT * FROM _adb_nodes WHERE kind = ?1",
169    ///     &[&"session" as &dyn rusqlite::ToSql],
170    /// ).unwrap();
171    /// ```
172    pub fn query_json_params(
173        &self,
174        sql: &str,
175        params: &[&dyn rusqlite::ToSql],
176    ) -> Result<Vec<serde_json::Value>> {
177        let conn = self.conn.lock().unwrap();
178        let mut stmt = conn.prepare(sql)?;
179        let col_names: Vec<String> = stmt.column_names().iter().map(|s| s.to_string()).collect();
180        let rows = stmt.query_map(params, |row| {
181            let mut map = serde_json::Map::new();
182            for (i, name) in col_names.iter().enumerate() {
183                let val: rusqlite::types::Value = row.get(i)?;
184                map.insert(name.clone(), rusqlite_value_to_json(val));
185            }
186            Ok(serde_json::Value::Object(map))
187        })?;
188        rows.map(|r| r.map_err(crate::error::AgentDbError::Sqlite))
189            .collect()
190    }
191
192    /// Flush dirty HNSW indexes and close gracefully.
193    ///
194    /// After this returns, the subsequent `Drop` is a no-op (no double flush).
195    pub fn close(self) -> Result<()> {
196        let collections = self.vectors().list_collections()?;
197        for (name, dim, _) in collections {
198            let col = self.vectors().collection(&name, dim)?;
199            let is_dirty: i64 = {
200                let conn = self.conn.lock().unwrap();
201                conn.query_row(
202                    "SELECT COALESCE(
203                        (SELECT is_dirty FROM _adb_hnsw_index
204                         WHERE collection_id =
205                           (SELECT id FROM _adb_collections WHERE name = ?1)
206                        ), 0)",
207                    rusqlite::params![name],
208                    |r| r.get(0),
209                )
210                .unwrap_or(0)
211            };
212            if is_dirty == 1 {
213                col.reindex()?;
214            }
215        }
216        self.closed.store(true, Ordering::Release);
217        Ok(())
218    }
219
220    /// Return database-wide statistics (single-query implementation).
221    pub fn stats(&self) -> Result<DbStats> {
222        let conn = self.conn.lock().unwrap();
223        conn.query_row(
224            "SELECT
225                 (SELECT COUNT(*)           FROM _adb_collections)          AS collections,
226                 (SELECT COALESCE(SUM(count),0) FROM _adb_collections)      AS vectors,
227                 (SELECT COUNT(*)           FROM _adb_nodes)                 AS nodes,
228                 (SELECT COUNT(*)           FROM _adb_edges)                 AS edges,
229                 (SELECT COUNT(*)           FROM _adb_conversations)         AS conversations,
230                 (SELECT COUNT(*)           FROM _adb_messages)              AS messages,
231                 (SELECT COUNT(*)           FROM _adb_workflows)             AS workflows,
232                 (SELECT COUNT(*)           FROM _adb_workflow_steps)        AS workflow_steps,
233                 (SELECT COUNT(*)           FROM _adb_traces)                AS traces",
234            [],
235            |r| {
236                Ok(DbStats {
237                    collections:    r.get(0)?,
238                    vectors:        r.get(1)?,
239                    nodes:          r.get(2)?,
240                    edges:          r.get(3)?,
241                    conversations:  r.get(4)?,
242                    messages:       r.get(5)?,
243                    workflows:      r.get(6)?,
244                    workflow_steps: r.get(7)?,
245                    traces:         r.get(8)?,
246                })
247            },
248        )
249        .map_err(crate::error::AgentDbError::Sqlite)
250    }
251}
252
253/// Database-wide statistics returned by [`AgentDB::stats`].
254#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
255pub struct DbStats {
256    /// Number of named vector collections.
257    pub collections: i64,
258    /// Total number of vectors across all collections.
259    pub vectors: i64,
260    /// Number of nodes in the memory graph.
261    pub nodes: i64,
262    /// Number of directed edges in the memory graph.
263    pub edges: i64,
264    /// Number of conversation threads.
265    pub conversations: i64,
266    /// Total number of messages across all conversations.
267    pub messages: i64,
268    /// Number of workflow records.
269    pub workflows: i64,
270    /// Total number of workflow steps across all workflows.
271    pub workflow_steps: i64,
272    /// Total number of reasoning trace entries.
273    pub traces: i64,
274}
275
276impl Drop for AgentDB {
277    fn drop(&mut self) {
278        if self.closed.load(Ordering::Acquire) {
279            return;
280        }
281        if let Ok(collections) = self.vectors().list_collections() {
282            for (name, dim, _) in collections {
283                if let Ok(col) = self.vectors().collection(&name, dim) {
284                    let is_dirty: bool = {
285                        let conn = self.conn.lock().unwrap();
286                        conn.query_row(
287                            "SELECT COALESCE(
288                                (SELECT is_dirty FROM _adb_hnsw_index
289                                 WHERE collection_id =
290                                   (SELECT id FROM _adb_collections WHERE name = ?1)
291                                ), 0)",
292                            rusqlite::params![name],
293                            |r| r.get::<_, i64>(0),
294                        )
295                        .unwrap_or(0) == 1
296                    };
297                    if is_dirty {
298                        let _ = col.reindex();
299                    }
300                }
301            }
302        }
303    }
304}
305
306fn rusqlite_value_to_json(val: rusqlite::types::Value) -> serde_json::Value {
307    match val {
308        rusqlite::types::Value::Null => serde_json::Value::Null,
309        rusqlite::types::Value::Integer(i) => serde_json::Value::Number(i.into()),
310        rusqlite::types::Value::Real(f) => serde_json::Number::from_f64(f)
311            .map(serde_json::Value::Number)
312            .unwrap_or(serde_json::Value::Null),
313        rusqlite::types::Value::Text(s) => serde_json::Value::String(s),
314        rusqlite::types::Value::Blob(b) => {
315            serde_json::Value::String(format!("<blob {} bytes>", b.len()))
316        }
317    }
318}
agentdb/db.rs

agentdb/
db.rs
