tinyflow_framework/state/
sqlite_job_state.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use tokio::sync::Mutex;
6use turso::{Builder, Connection};
7
8use tinyflow_api::error::StreamResult;
9use tinyflow_api::state::{StateStore, state_db_err};
10
11pub struct SqliteJobStateStore {
12    path: PathBuf,
13    conn: Mutex<Connection>,
14}
15
16impl SqliteJobStateStore {
17    pub async fn open(path: &Path) -> StreamResult<Arc<dyn StateStore>> {
18        if let Some(parent) = path.parent() {
19            std::fs::create_dir_all(parent)
20                .map_err(|e| state_db_err(format!("create state db parent dir: {}", e)))?;
21        }
22        let path_str = path
23            .to_str()
24            .ok_or_else(|| state_db_err("non-UTF-8 path for turso db".to_string()))?;
25        let db = Builder::new_local(path_str)
26            .build()
27            .await
28            .map_err(|e| state_db_err(format!("open turso db: {}", e)))?;
29        let conn = db
30            .connect()
31            .map_err(|e| state_db_err(format!("connect turso: {}", e)))?;
32        let mut stmt = conn
33            .prepare("PRAGMA journal_mode=WAL")
34            .await
35            .map_err(|e| state_db_err(format!("turso pragma prepare: {}", e)))?;
36        stmt.query(())
37            .await
38            .map_err(|e| state_db_err(format!("turso pragma query: {}", e)))?;
39        let mut stmt = conn
40            .prepare("PRAGMA busy_timeout=5000")
41            .await
42            .map_err(|e| state_db_err(format!("turso busy_timeout prepare: {}", e)))?;
43        stmt.query(())
44            .await
45            .map_err(|e| state_db_err(format!("turso busy_timeout query: {}", e)))?;
46        conn.execute(
47            "CREATE TABLE IF NOT EXISTS kv (
48               namespace TEXT NOT NULL,
49               key BLOB NOT NULL,
50               value BLOB NOT NULL,
51               PRIMARY KEY (namespace, key)
52             )",
53            (),
54        )
55        .await
56        .map_err(|e| state_db_err(format!("turso init schema: {}", e)))?;
57        let store: Arc<dyn StateStore> = Arc::new(Self {
58            path: path.to_path_buf(),
59            conn: Mutex::new(conn),
60        });
61        Ok(store)
62    }
63}
64
65fn prefix_upper_bound(prefix: &[u8]) -> Option<Vec<u8>> {
66    if prefix.is_empty() {
67        return None;
68    }
69    let mut upper = prefix.to_vec();
70    while let Some(last) = upper.last_mut() {
71        if *last < 255 {
72            *last += 1;
73            return Some(upper);
74        }
75        upper.pop();
76    }
77    None
78}
79
80fn hex_prefix_like(prefix: &[u8]) -> String {
81    prefix.iter().map(|b| format!("{:02x}", b)).collect()
82}
83
84#[async_trait]
85impl StateStore for SqliteJobStateStore {
86    fn persistent_db_path(&self) -> Option<&Path> {
87        Some(self.path.as_path())
88    }
89
90    async fn put(&self, namespace: &str, key: &[u8], value: &[u8]) -> StreamResult<()> {
91        let conn = self.conn.lock().await;
92        conn.execute(
93            "INSERT OR REPLACE INTO kv (namespace, key, value) VALUES (?1, ?2, ?3)",
94            (namespace, key, value),
95        )
96        .await
97        .map_err(|e| state_db_err(format!("turso: {}", e)))?;
98        Ok(())
99    }
100
101    async fn get(&self, namespace: &str, key: &[u8]) -> StreamResult<Option<Vec<u8>>> {
102        let conn = self.conn.lock().await;
103        let mut stmt = conn
104            .prepare("SELECT value FROM kv WHERE namespace = ?1 AND key = ?2")
105            .await
106            .map_err(|e| state_db_err(format!("turso: {}", e)))?;
107        let mut rows = stmt
108            .query((namespace, key))
109            .await
110            .map_err(|e| state_db_err(format!("turso: {}", e)))?;
111        if let Some(row) = rows
112            .next()
113            .await
114            .map_err(|e| state_db_err(format!("turso: {}", e)))?
115        {
116            let blob = row
117                .get_value(0)
118                .map_err(|e| state_db_err(format!("turso: {}", e)))?;
119            let v = blob
120                .as_blob()
121                .ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
122                .to_vec();
123            return Ok(Some(v));
124        }
125        Ok(None)
126    }
127
128    async fn delete(&self, namespace: &str, key: &[u8]) -> StreamResult<()> {
129        let conn = self.conn.lock().await;
130        conn.execute(
131            "DELETE FROM kv WHERE namespace = ?1 AND key = ?2",
132            (namespace, key),
133        )
134        .await
135        .map_err(|e| state_db_err(format!("turso: {}", e)))?;
136        Ok(())
137    }
138
139    async fn delete_prefix(&self, namespace: &str, key_prefix: &[u8]) -> StreamResult<()> {
140        let conn = self.conn.lock().await;
141        if let Some(upper) = prefix_upper_bound(key_prefix) {
142            conn.execute(
143                "DELETE FROM kv WHERE namespace = ?1 AND key >= ?2 AND key < ?3",
144                (namespace, key_prefix, upper.as_slice()),
145            )
146            .await
147            .map_err(|e| state_db_err(format!("turso: {}", e)))?;
148        } else if key_prefix.is_empty() {
149            conn.execute("DELETE FROM kv WHERE namespace = ?1", [namespace])
150                .await
151                .map_err(|e| state_db_err(format!("turso: {}", e)))?;
152        } else {
153            let pattern = format!("{}%", hex_prefix_like(key_prefix));
154            conn.execute(
155                "DELETE FROM kv WHERE namespace = ?1 AND hex(key) LIKE ?2",
156                (namespace, pattern.as_str()),
157            )
158            .await
159            .map_err(|e| state_db_err(format!("turso: {}", e)))?;
160        }
161        Ok(())
162    }
163
164    async fn list_prefix(
165        &self,
166        namespace: &str,
167        key_prefix: &[u8],
168    ) -> StreamResult<Vec<(Vec<u8>, Vec<u8>)>> {
169        let conn = self.conn.lock().await;
170        let mut out = Vec::new();
171        if let Some(upper) = prefix_upper_bound(key_prefix) {
172            let mut stmt = conn
173                .prepare(
174                    "SELECT key, value FROM kv WHERE namespace = ?1 AND key >= ?2 AND key < ?3 ORDER BY key",
175                )
176                .await
177                .map_err(|e| state_db_err(format!("turso: {}", e)))?;
178            let mut rows = stmt
179                .query((namespace, key_prefix, upper.as_slice()))
180                .await
181                .map_err(|e| state_db_err(format!("turso: {}", e)))?;
182            while let Some(row) = rows
183                .next()
184                .await
185                .map_err(|e| state_db_err(format!("turso: {}", e)))?
186            {
187                let k = row
188                    .get_value(0)
189                    .map_err(|e| state_db_err(format!("turso: {}", e)))?
190                    .as_blob()
191                    .ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
192                    .to_vec();
193                let v = row
194                    .get_value(1)
195                    .map_err(|e| state_db_err(format!("turso: {}", e)))?
196                    .as_blob()
197                    .ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
198                    .to_vec();
199                out.push((k, v));
200            }
201        } else if key_prefix.is_empty() {
202            let mut stmt = conn
203                .prepare("SELECT key, value FROM kv WHERE namespace = ?1 ORDER BY key")
204                .await
205                .map_err(|e| state_db_err(format!("turso: {}", e)))?;
206            let mut rows = stmt
207                .query([namespace])
208                .await
209                .map_err(|e| state_db_err(format!("turso: {}", e)))?;
210            while let Some(row) = rows
211                .next()
212                .await
213                .map_err(|e| state_db_err(format!("turso: {}", e)))?
214            {
215                let k = row
216                    .get_value(0)
217                    .map_err(|e| state_db_err(format!("turso: {}", e)))?
218                    .as_blob()
219                    .ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
220                    .to_vec();
221                let v = row
222                    .get_value(1)
223                    .map_err(|e| state_db_err(format!("turso: {}", e)))?
224                    .as_blob()
225                    .ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
226                    .to_vec();
227                out.push((k, v));
228            }
229        } else {
230            let pattern = format!("{}%", hex_prefix_like(key_prefix));
231            let mut stmt = conn
232                .prepare(
233                    "SELECT key, value FROM kv WHERE namespace = ?1 AND hex(key) LIKE ?2 ORDER BY key",
234                )
235                .await
236                .map_err(|e| state_db_err(format!("turso: {}", e)))?;
237            let mut rows = stmt
238                .query((namespace, pattern.as_str()))
239                .await
240                .map_err(|e| state_db_err(format!("turso: {}", e)))?;
241            while let Some(row) = rows
242                .next()
243                .await
244                .map_err(|e| state_db_err(format!("turso: {}", e)))?
245            {
246                let k = row
247                    .get_value(0)
248                    .map_err(|e| state_db_err(format!("turso: {}", e)))?
249                    .as_blob()
250                    .ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
251                    .to_vec();
252                let v = row
253                    .get_value(1)
254                    .map_err(|e| state_db_err(format!("turso: {}", e)))?
255                    .as_blob()
256                    .ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
257                    .to_vec();
258                out.push((k, v));
259            }
260        }
261        Ok(out)
262    }
263}
tinyflow_framework/state/sqlite_job_state.rs

tinyflow_framework/state/
sqlite_job_state.rs
