Skip to main content

plexus_substrate/activations/mustache/
storage.rs

1//! Mustache template storage using SQLite
2
3use 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/// Configuration for Mustache storage
10#[derive(Debug, Clone)]
11pub struct MustacheStorageConfig {
12    /// Path to SQLite database for templates
13    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
24/// Storage layer for mustache templates
25pub struct MustacheStorage {
26    pool: SqlitePool,
27}
28
29impl MustacheStorage {
30    /// Create a new mustache storage instance
31    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    /// Run database migrations
47    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    /// Get a template by plugin_id, method, and name
73    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    /// Set (insert or update) a template
93    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        // Check if template exists
103        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            // Update existing template
118            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            // Insert new template
133            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    /// List all templates for a plugin
162    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    /// Delete a template
192    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
212/// Get current Unix timestamp in seconds
213fn 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        // Set a template
240        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        // Get the template
250        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        // Set initial template
264        let info1 = storage
265            .set_template(&plugin_id, "chat", "default", "v1")
266            .await
267            .unwrap();
268
269        // Update template
270        let info2 = storage
271            .set_template(&plugin_id, "chat", "default", "v2")
272            .await
273            .unwrap();
274
275        // ID and created_at should be preserved
276        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        // Content should be updated
281        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        // Deleting again should return false
333        let deleted_again = storage
334            .delete_template(&plugin_id, "chat", "default")
335            .await
336            .unwrap();
337        assert!(!deleted_again);
338    }
339}