plexus_substrate/activations/mustache/
storage.rs1use super::types::{MustacheError, TemplateInfo};
4use sqlx::{sqlite::{SqliteConnectOptions, SqlitePool}, ConnectOptions, Row};
5use std::path::PathBuf;
6use std::time::{SystemTime, UNIX_EPOCH};
7use uuid::Uuid;
8
9#[derive(Debug, Clone)]
11pub struct MustacheStorageConfig {
12 pub db_path: PathBuf,
14}
15
16impl Default for MustacheStorageConfig {
17 fn default() -> Self {
18 Self {
19 db_path: PathBuf::from("templates.db"),
20 }
21 }
22}
23
24pub struct MustacheStorage {
26 pool: SqlitePool,
27}
28
29impl MustacheStorage {
30 pub async fn new(config: MustacheStorageConfig) -> Result<Self, MustacheError> {
32 let db_url = format!("sqlite:{}?mode=rwc", config.db_path.display());
33 let mut connect_options: SqliteConnectOptions = db_url.parse()
34 .map_err(|e| format!("Failed to parse database URL: {}", e))?;
35 connect_options.disable_statement_logging();
36 let pool = SqlitePool::connect_with(connect_options.clone())
37 .await
38 .map_err(|e| format!("Failed to connect to templates database: {}", e))?;
39
40 let storage = Self { pool };
41 storage.run_migrations().await?;
42
43 Ok(storage)
44 }
45
46 async fn run_migrations(&self) -> Result<(), MustacheError> {
48 sqlx::query(
49 r#"
50 CREATE TABLE IF NOT EXISTS templates (
51 id TEXT PRIMARY KEY,
52 plugin_id TEXT NOT NULL,
53 method TEXT NOT NULL,
54 name TEXT NOT NULL,
55 template TEXT NOT NULL,
56 created_at INTEGER NOT NULL,
57 updated_at INTEGER NOT NULL,
58 UNIQUE(plugin_id, method, name)
59 );
60
61 CREATE INDEX IF NOT EXISTS idx_templates_plugin ON templates(plugin_id);
62 CREATE INDEX IF NOT EXISTS idx_templates_lookup ON templates(plugin_id, method, name);
63 "#,
64 )
65 .execute(&self.pool)
66 .await
67 .map_err(|e| format!("Failed to run mustache migrations: {}", e))?;
68
69 Ok(())
70 }
71
72 pub async fn get_template(
74 &self,
75 plugin_id: &Uuid,
76 method: &str,
77 name: &str,
78 ) -> Result<Option<String>, MustacheError> {
79 let row = sqlx::query(
80 "SELECT template FROM templates WHERE plugin_id = ? AND method = ? AND name = ?",
81 )
82 .bind(plugin_id.to_string())
83 .bind(method)
84 .bind(name)
85 .fetch_optional(&self.pool)
86 .await
87 .map_err(|e| format!("Failed to fetch template: {}", e))?;
88
89 Ok(row.map(|r| r.get("template")))
90 }
91
92 pub async fn set_template(
94 &self,
95 plugin_id: &Uuid,
96 method: &str,
97 name: &str,
98 template: &str,
99 ) -> Result<TemplateInfo, MustacheError> {
100 let now = current_timestamp();
101
102 let existing = sqlx::query(
104 "SELECT id, created_at FROM templates WHERE plugin_id = ? AND method = ? AND name = ?",
105 )
106 .bind(plugin_id.to_string())
107 .bind(method)
108 .bind(name)
109 .fetch_optional(&self.pool)
110 .await
111 .map_err(|e| format!("Failed to check existing template: {}", e))?;
112
113 let (id, created_at) = if let Some(row) = existing {
114 let id: String = row.get("id");
115 let created_at: i64 = row.get("created_at");
116
117 sqlx::query(
119 "UPDATE templates SET template = ?, updated_at = ? WHERE id = ?",
120 )
121 .bind(template)
122 .bind(now)
123 .bind(&id)
124 .execute(&self.pool)
125 .await
126 .map_err(|e| format!("Failed to update template: {}", e))?;
127
128 (id, created_at)
129 } else {
130 let id = Uuid::new_v4().to_string();
131
132 sqlx::query(
134 "INSERT INTO templates (id, plugin_id, method, name, template, created_at, updated_at)
135 VALUES (?, ?, ?, ?, ?, ?, ?)",
136 )
137 .bind(&id)
138 .bind(plugin_id.to_string())
139 .bind(method)
140 .bind(name)
141 .bind(template)
142 .bind(now)
143 .bind(now)
144 .execute(&self.pool)
145 .await
146 .map_err(|e| format!("Failed to insert template: {}", e))?;
147
148 (id, now)
149 };
150
151 Ok(TemplateInfo {
152 id,
153 plugin_id: *plugin_id,
154 method: method.to_string(),
155 name: name.to_string(),
156 created_at,
157 updated_at: now,
158 })
159 }
160
161 pub async fn list_templates(&self, plugin_id: &Uuid) -> Result<Vec<TemplateInfo>, MustacheError> {
163 let rows = sqlx::query(
164 "SELECT id, plugin_id, method, name, created_at, updated_at
165 FROM templates WHERE plugin_id = ? ORDER BY method, name",
166 )
167 .bind(plugin_id.to_string())
168 .fetch_all(&self.pool)
169 .await
170 .map_err(|e| format!("Failed to list templates: {}", e))?;
171
172 let templates: Result<Vec<TemplateInfo>, MustacheError> = rows
173 .iter()
174 .map(|row| {
175 let plugin_id_str: String = row.get("plugin_id");
176 Ok(TemplateInfo {
177 id: row.get("id"),
178 plugin_id: Uuid::parse_str(&plugin_id_str)
179 .map_err(|e| format!("Invalid plugin ID: {}", e))?,
180 method: row.get("method"),
181 name: row.get("name"),
182 created_at: row.get("created_at"),
183 updated_at: row.get("updated_at"),
184 })
185 })
186 .collect();
187
188 templates
189 }
190
191 pub async fn delete_template(
193 &self,
194 plugin_id: &Uuid,
195 method: &str,
196 name: &str,
197 ) -> Result<bool, MustacheError> {
198 let result = sqlx::query(
199 "DELETE FROM templates WHERE plugin_id = ? AND method = ? AND name = ?",
200 )
201 .bind(plugin_id.to_string())
202 .bind(method)
203 .bind(name)
204 .execute(&self.pool)
205 .await
206 .map_err(|e| format!("Failed to delete template: {}", e))?;
207
208 Ok(result.rows_affected() > 0)
209 }
210}
211
212fn current_timestamp() -> i64 {
214 SystemTime::now()
215 .duration_since(UNIX_EPOCH)
216 .unwrap()
217 .as_secs() as i64
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223 use tempfile::{tempdir, TempDir};
224
225 async fn create_test_storage() -> (MustacheStorage, TempDir) {
226 let dir = tempdir().unwrap();
227 let config = MustacheStorageConfig {
228 db_path: dir.path().join("test_templates.db"),
229 };
230 let storage = MustacheStorage::new(config).await.unwrap();
231 (storage, dir)
232 }
233
234 #[tokio::test]
235 async fn test_set_and_get_template() {
236 let (storage, _dir) = create_test_storage().await;
237 let plugin_id = Uuid::new_v4();
238
239 let info = storage
241 .set_template(&plugin_id, "chat", "default", "[{{role}}]: {{content}}")
242 .await
243 .unwrap();
244
245 assert_eq!(info.plugin_id, plugin_id);
246 assert_eq!(info.method, "chat");
247 assert_eq!(info.name, "default");
248
249 let template = storage
251 .get_template(&plugin_id, "chat", "default")
252 .await
253 .unwrap();
254
255 assert_eq!(template, Some("[{{role}}]: {{content}}".to_string()));
256 }
257
258 #[tokio::test]
259 async fn test_update_template() {
260 let (storage, _dir) = create_test_storage().await;
261 let plugin_id = Uuid::new_v4();
262
263 let info1 = storage
265 .set_template(&plugin_id, "chat", "default", "v1")
266 .await
267 .unwrap();
268
269 let info2 = storage
271 .set_template(&plugin_id, "chat", "default", "v2")
272 .await
273 .unwrap();
274
275 assert_eq!(info1.id, info2.id);
277 assert_eq!(info1.created_at, info2.created_at);
278 assert!(info2.updated_at >= info1.updated_at);
279
280 let template = storage
282 .get_template(&plugin_id, "chat", "default")
283 .await
284 .unwrap();
285 assert_eq!(template, Some("v2".to_string()));
286 }
287
288 #[tokio::test]
289 async fn test_list_templates() {
290 let (storage, _dir) = create_test_storage().await;
291 let plugin_id = Uuid::new_v4();
292
293 storage
294 .set_template(&plugin_id, "chat", "default", "t1")
295 .await
296 .unwrap();
297 storage
298 .set_template(&plugin_id, "chat", "compact", "t2")
299 .await
300 .unwrap();
301 storage
302 .set_template(&plugin_id, "execute", "default", "t3")
303 .await
304 .unwrap();
305
306 let templates = storage.list_templates(&plugin_id).await.unwrap();
307 assert_eq!(templates.len(), 3);
308 }
309
310 #[tokio::test]
311 async fn test_delete_template() {
312 let (storage, _dir) = create_test_storage().await;
313 let plugin_id = Uuid::new_v4();
314
315 storage
316 .set_template(&plugin_id, "chat", "default", "content")
317 .await
318 .unwrap();
319
320 let deleted = storage
321 .delete_template(&plugin_id, "chat", "default")
322 .await
323 .unwrap();
324 assert!(deleted);
325
326 let template = storage
327 .get_template(&plugin_id, "chat", "default")
328 .await
329 .unwrap();
330 assert!(template.is_none());
331
332 let deleted_again = storage
334 .delete_template(&plugin_id, "chat", "default")
335 .await
336 .unwrap();
337 assert!(!deleted_again);
338 }
339}