Skip to main content

plexus_substrate/activations/mustache/
storage.rs

1//! Mustache template storage using SQLite
2
3use 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/// Configuration for Mustache storage
12#[derive(Debug, Clone)]
13pub struct MustacheStorageConfig {
14    /// Path to SQLite database for templates
15    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
26/// Storage layer for mustache templates
27pub struct MustacheStorage {
28    pool: SqlitePool,
29}
30
31impl MustacheStorage {
32    /// Create a new mustache storage instance
33    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    /// Run database migrations
43    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    /// Get a template by plugin_id, method, and name
69    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    /// Set (insert or update) a template
89    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        // Check if template exists
99        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            // Update existing template
114            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            // Insert new template
129            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    /// List all templates for a plugin
158    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    /// Delete a template
188    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
208/// Get current Unix timestamp in seconds
209fn 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        // Set a template
236        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        // Get the template
246        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        // Set initial template
260        let info1 = storage
261            .set_template(&plugin_id, "chat", "default", "v1")
262            .await
263            .unwrap();
264
265        // Update template
266        let info2 = storage
267            .set_template(&plugin_id, "chat", "default", "v2")
268            .await
269            .unwrap();
270
271        // ID and created_at should be preserved
272        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        // Content should be updated
277        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        // Deleting again should return false
329        let deleted_again = storage
330            .delete_template(&plugin_id, "chat", "default")
331            .await
332            .unwrap();
333        assert!(!deleted_again);
334    }
335}