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