vv_agent/runtime/stores/
sqlite.rs

1use std::io::{Error, Result};
2use std::path::Path;
3use std::sync::Mutex;
4
5use rusqlite::{params, Connection, OptionalExtension};
6
7use crate::runtime::checkpoint_codec::{
8    cycles_from_json, cycles_to_json, messages_from_json, messages_to_json,
9};
10use crate::runtime::state::{
11    checkpoint_status_from_value, checkpoint_status_value, Checkpoint, StateStore,
12};
13
14#[derive(Debug)]
15pub struct SqliteStateStore {
16    connection: Mutex<Connection>,
17}
18
19impl SqliteStateStore {
20    pub fn new(db_path: impl AsRef<Path>) -> Result<Self> {
21        let db_path = db_path.as_ref().to_string_lossy().to_string();
22        let connection = Connection::open(db_path).map_err(sqlite_to_io)?;
23        connection
24            .execute_batch(
25                r#"
26                PRAGMA journal_mode=WAL;
27                CREATE TABLE IF NOT EXISTS checkpoints (
28                    task_id TEXT PRIMARY KEY,
29                    cycle_index INTEGER NOT NULL,
30                    status TEXT NOT NULL,
31                    messages TEXT NOT NULL,
32                    cycles TEXT NOT NULL,
33                    shared_state TEXT NOT NULL
34                );
35                "#,
36            )
37            .map_err(sqlite_to_io)?;
38        Ok(Self {
39            connection: Mutex::new(connection),
40        })
41    }
42
43    pub fn close(self) -> Result<()> {
44        let connection = self
45            .connection
46            .into_inner()
47            .map_err(|_| Error::other("sqlite state store lock is poisoned"))?;
48        connection.close().map_err(|(_, error)| sqlite_to_io(error))
49    }
50}
51
52impl StateStore for SqliteStateStore {
53    fn save_checkpoint(&self, checkpoint: Checkpoint) -> Result<()> {
54        let messages_json = messages_to_json(&checkpoint.messages)?;
55        let cycles_json = cycles_to_json(&checkpoint.cycles)?;
56        let shared_state_json = serde_json::to_string(&checkpoint.shared_state)
57            .map_err(|error| Error::other(error.to_string()))?;
58        self.connection
59            .lock()
60            .map_err(|_| Error::other("sqlite state store lock is poisoned"))?
61            .execute(
62                r#"
63                INSERT OR REPLACE INTO checkpoints
64                    (task_id, cycle_index, status, messages, cycles, shared_state)
65                VALUES (?1, ?2, ?3, ?4, ?5, ?6)
66                "#,
67                params![
68                    checkpoint.task_id,
69                    checkpoint.cycle_index,
70                    checkpoint_status_value(checkpoint.status),
71                    messages_json,
72                    cycles_json,
73                    shared_state_json,
74                ],
75            )
76            .map_err(sqlite_to_io)?;
77        Ok(())
78    }
79
80    fn load_checkpoint(&self, task_id: &str) -> Result<Option<Checkpoint>> {
81        let connection = self
82            .connection
83            .lock()
84            .map_err(|_| Error::other("sqlite state store lock is poisoned"))?;
85        let row = connection
86            .query_row(
87                "SELECT task_id, cycle_index, status, messages, cycles, shared_state FROM checkpoints WHERE task_id = ?1",
88                params![task_id],
89                |row| {
90                    Ok((
91                        row.get::<_, String>(0)?,
92                        row.get::<_, u32>(1)?,
93                        row.get::<_, String>(2)?,
94                        row.get::<_, String>(3)?,
95                        row.get::<_, String>(4)?,
96                        row.get::<_, String>(5)?,
97                    ))
98                },
99            )
100            .optional()
101            .map_err(sqlite_to_io)?;
102        let Some((task_id, cycle_index, status, messages, cycles, shared_state)) = row else {
103            return Ok(None);
104        };
105        Ok(Some(Checkpoint {
106            task_id,
107            cycle_index,
108            status: checkpoint_status_from_value(&status)?,
109            messages: messages_from_json(&messages)?,
110            cycles: cycles_from_json(&cycles)?,
111            shared_state: serde_json::from_str(&shared_state)
112                .map_err(|error| Error::other(error.to_string()))?,
113        }))
114    }
115
116    fn delete_checkpoint(&self, task_id: &str) -> Result<()> {
117        self.connection
118            .lock()
119            .map_err(|_| Error::other("sqlite state store lock is poisoned"))?
120            .execute(
121                "DELETE FROM checkpoints WHERE task_id = ?1",
122                params![task_id],
123            )
124            .map_err(sqlite_to_io)?;
125        Ok(())
126    }
127
128    fn list_checkpoints(&self) -> Result<Vec<String>> {
129        let connection = self
130            .connection
131            .lock()
132            .map_err(|_| Error::other("sqlite state store lock is poisoned"))?;
133        let mut statement = connection
134            .prepare("SELECT task_id FROM checkpoints ORDER BY task_id")
135            .map_err(sqlite_to_io)?;
136        let rows = statement
137            .query_map([], |row| row.get::<_, String>(0))
138            .map_err(sqlite_to_io)?;
139        rows.collect::<rusqlite::Result<Vec<_>>>()
140            .map_err(sqlite_to_io)
141    }
142}
143
144fn sqlite_to_io(error: rusqlite::Error) -> Error {
145    Error::other(error.to_string())
146}
vv_agent/runtime/stores/sqlite.rs

vv_agent/runtime/stores/
sqlite.rs
