Skip to main content

mixtape_tools/sqlite/query/
write.rs

1//! Write query tool
2
3use crate::prelude::*;
4use crate::sqlite::error::SqliteToolError;
5use crate::sqlite::manager::with_connection;
6use crate::sqlite::types::json_to_sql;
7
8/// Input for write query execution
9#[derive(Debug, Deserialize, JsonSchema)]
10pub struct WriteQueryInput {
11    /// SQL query to execute (INSERT, UPDATE, DELETE)
12    pub query: String,
13
14    /// Query parameters for prepared statements
15    #[serde(default)]
16    pub params: Vec<serde_json::Value>,
17
18    /// Database file path. If not specified, uses the default database.
19    #[serde(default)]
20    pub db_path: Option<String>,
21}
22
23/// Write query result
24#[derive(Debug, Serialize, JsonSchema)]
25struct WriteResult {
26    status: String,
27    rows_affected: usize,
28    last_insert_rowid: Option<i64>,
29}
30
31/// Tool for executing data modification queries (DESTRUCTIVE)
32///
33/// Executes INSERT, UPDATE, and DELETE queries.
34/// Returns the number of rows affected and last insert rowid (for INSERT).
35pub struct WriteQueryTool;
36
37impl WriteQueryTool {
38    /// Validates that a query is a write operation
39    fn is_write_query(sql: &str) -> bool {
40        let normalized = sql.trim().to_uppercase();
41        let write_prefixes = ["INSERT", "UPDATE", "DELETE", "REPLACE"];
42        write_prefixes
43            .iter()
44            .any(|prefix| normalized.starts_with(prefix))
45    }
46}
47
48impl Tool for WriteQueryTool {
49    type Input = WriteQueryInput;
50
51    fn name(&self) -> &str {
52        "sqlite_write_query"
53    }
54
55    fn description(&self) -> &str {
56        "Execute a data modification SQL query (INSERT, UPDATE, DELETE). Returns the number of rows affected."
57    }
58
59    async fn execute(&self, input: Self::Input) -> Result<ToolResult, ToolError> {
60        // Validate query is a write operation
61        if !Self::is_write_query(&input.query) {
62            return Err(SqliteToolError::InvalidQuery(
63                "Only INSERT, UPDATE, DELETE, and REPLACE queries are allowed. Use sqlite_read_query for SELECT or sqlite_schema_query for DDL.".to_string()
64            ).into());
65        }
66
67        let query = input.query;
68        let params = input.params;
69
70        let result = with_connection(input.db_path, move |conn| {
71            // Convert params to rusqlite values
72            let params_ref: Vec<Box<dyn rusqlite::ToSql>> =
73                params.iter().map(|v| json_to_sql(v)).collect();
74
75            let params_slice: Vec<&dyn rusqlite::ToSql> =
76                params_ref.iter().map(|b| b.as_ref()).collect();
77
78            let rows_affected = conn.execute(&query, params_slice.as_slice())?;
79
80            // Get last insert rowid for INSERT queries
81            let last_insert_rowid = if query.trim().to_uppercase().starts_with("INSERT") {
82                Some(conn.last_insert_rowid())
83            } else {
84                None
85            };
86
87            Ok(WriteResult {
88                status: "success".to_string(),
89                rows_affected,
90                last_insert_rowid,
91            })
92        })
93        .await?;
94
95        Ok(ToolResult::Json(serde_json::to_value(result)?))
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use crate::sqlite::test_utils::{unwrap_json, TestDatabase};
103    use mixtape_core::tool::Tool;
104
105    #[tokio::test]
106    async fn test_write_query_insert() {
107        let db =
108            TestDatabase::with_schema("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT)")
109                .await;
110
111        let tool = WriteQueryTool;
112        let result = tool
113            .execute(WriteQueryInput {
114                query: "INSERT INTO users (name) VALUES (?)".to_string(),
115                params: vec![serde_json::json!("Alice")],
116                db_path: Some(db.key()),
117            })
118            .await
119            .unwrap();
120
121        let json = unwrap_json(result);
122        assert_eq!(json["rows_affected"], 1);
123        assert!(json["last_insert_rowid"].as_i64().is_some());
124    }
125
126    #[tokio::test]
127    async fn test_write_query_update() {
128        let db = TestDatabase::with_schema(
129            "CREATE TABLE users (id INTEGER, name TEXT);
130             INSERT INTO users VALUES (1, 'Alice');
131             INSERT INTO users VALUES (2, 'Bob');",
132        )
133        .await;
134
135        let tool = WriteQueryTool;
136        let result = tool
137            .execute(WriteQueryInput {
138                query: "UPDATE users SET name = 'Updated'".to_string(),
139                params: vec![],
140                db_path: Some(db.key()),
141            })
142            .await
143            .unwrap();
144
145        let json = unwrap_json(result);
146        assert_eq!(json["rows_affected"], 2);
147    }
148
149    #[tokio::test]
150    async fn test_reject_select_query() {
151        let db = TestDatabase::new().await;
152
153        let result = WriteQueryTool
154            .execute(WriteQueryInput {
155                query: "SELECT * FROM users".to_string(),
156                params: vec![],
157                db_path: Some(db.key()),
158            })
159            .await;
160        assert!(result.is_err());
161    }
162
163    #[test]
164    fn test_is_write_query() {
165        assert!(WriteQueryTool::is_write_query(
166            "INSERT INTO users VALUES (1)"
167        ));
168        assert!(WriteQueryTool::is_write_query(
169            "UPDATE users SET name = 'x'"
170        ));
171        assert!(WriteQueryTool::is_write_query("DELETE FROM users"));
172        assert!(WriteQueryTool::is_write_query(
173            "REPLACE INTO users VALUES (1)"
174        ));
175
176        assert!(!WriteQueryTool::is_write_query("SELECT * FROM users"));
177        assert!(!WriteQueryTool::is_write_query(
178            "CREATE TABLE users (id INT)"
179        ));
180        assert!(!WriteQueryTool::is_write_query("DROP TABLE users"));
181    }
182
183    #[test]
184    fn test_tool_metadata() {
185        let tool = WriteQueryTool;
186        assert_eq!(tool.name(), "sqlite_write_query");
187        assert!(!tool.description().is_empty());
188    }
189
190    #[tokio::test]
191    async fn test_write_query_delete() {
192        let db = TestDatabase::with_schema(
193            "CREATE TABLE users (id INTEGER, name TEXT);
194             INSERT INTO users VALUES (1, 'Alice');
195             INSERT INTO users VALUES (2, 'Bob');
196             INSERT INTO users VALUES (3, 'Charlie');",
197        )
198        .await;
199
200        let result = WriteQueryTool
201            .execute(WriteQueryInput {
202                query: "DELETE FROM users WHERE id > 1".to_string(),
203                params: vec![],
204                db_path: Some(db.key()),
205            })
206            .await
207            .unwrap();
208
209        let json = unwrap_json(result);
210        assert_eq!(json["status"], "success");
211        assert_eq!(json["rows_affected"], 2);
212        assert!(json["last_insert_rowid"].is_null());
213        assert_eq!(db.count("users"), 1);
214    }
215
216    #[tokio::test]
217    async fn test_write_query_delete_all() {
218        let db = TestDatabase::with_schema(
219            "CREATE TABLE users (id INTEGER, name TEXT);
220             INSERT INTO users VALUES (1, 'Alice');
221             INSERT INTO users VALUES (2, 'Bob');",
222        )
223        .await;
224
225        let result = WriteQueryTool
226            .execute(WriteQueryInput {
227                query: "DELETE FROM users".to_string(),
228                params: vec![],
229                db_path: Some(db.key()),
230            })
231            .await
232            .unwrap();
233
234        let json = unwrap_json(result);
235        assert_eq!(json["rows_affected"], 2);
236    }
237
238    #[tokio::test]
239    async fn test_write_query_replace() {
240        let db = TestDatabase::with_schema(
241            "CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT);
242             INSERT INTO users VALUES (1, 'Alice');",
243        )
244        .await;
245
246        let result = WriteQueryTool
247            .execute(WriteQueryInput {
248                query: "REPLACE INTO users VALUES (1, 'Updated Alice')".to_string(),
249                params: vec![],
250                db_path: Some(db.key()),
251            })
252            .await
253            .unwrap();
254
255        let json = unwrap_json(result);
256        assert_eq!(json["status"], "success");
257        assert_eq!(json["rows_affected"], 1);
258
259        // Verify replacement
260        let rows = db.query("SELECT name FROM users WHERE id = 1");
261        assert_eq!(rows[0][0], "Updated Alice");
262        assert_eq!(db.count("users"), 1);
263    }
264
265    #[tokio::test]
266    async fn test_write_query_replace_new_row() {
267        let db = TestDatabase::with_schema(
268            "CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT);
269             INSERT INTO users VALUES (1, 'Alice');",
270        )
271        .await;
272
273        let result = WriteQueryTool
274            .execute(WriteQueryInput {
275                query: "REPLACE INTO users VALUES (2, 'Bob')".to_string(),
276                params: vec![],
277                db_path: Some(db.key()),
278            })
279            .await
280            .unwrap();
281
282        let json = unwrap_json(result);
283        assert_eq!(json["rows_affected"], 1);
284        assert_eq!(db.count("users"), 2);
285    }
286
287    #[tokio::test]
288    async fn test_write_query_parameterized_insert() {
289        let db = TestDatabase::with_schema(
290            "CREATE TABLE data (id INTEGER, name TEXT, score REAL, active INTEGER)",
291        )
292        .await;
293
294        let result = WriteQueryTool
295            .execute(WriteQueryInput {
296                query: "INSERT INTO data VALUES (?, ?, ?, ?)".to_string(),
297                params: vec![
298                    serde_json::json!(1),
299                    serde_json::json!("Alice"),
300                    serde_json::json!(95.5),
301                    serde_json::json!(true),
302                ],
303                db_path: Some(db.key()),
304            })
305            .await
306            .unwrap();
307
308        let json = unwrap_json(result);
309        assert_eq!(json["rows_affected"], 1);
310
311        // Verify data
312        let rows = db.query("SELECT name, score, active FROM data WHERE id = 1");
313        assert_eq!(rows[0][0], "Alice");
314        assert_eq!(rows[0][2], 1); // true -> 1
315    }
316
317    #[tokio::test]
318    async fn test_write_query_parameterized_update() {
319        let db = TestDatabase::with_schema(
320            "CREATE TABLE users (id INTEGER, name TEXT);
321             INSERT INTO users VALUES (1, 'Alice');
322             INSERT INTO users VALUES (2, 'Bob');",
323        )
324        .await;
325
326        let result = WriteQueryTool
327            .execute(WriteQueryInput {
328                query: "UPDATE users SET name = ? WHERE id = ?".to_string(),
329                params: vec![serde_json::json!("Updated"), serde_json::json!(1)],
330                db_path: Some(db.key()),
331            })
332            .await
333            .unwrap();
334
335        let json = unwrap_json(result);
336        assert_eq!(json["rows_affected"], 1);
337        assert!(json["last_insert_rowid"].is_null());
338
339        let rows = db.query("SELECT name FROM users WHERE id = 1");
340        assert_eq!(rows[0][0], "Updated");
341    }
342
343    #[tokio::test]
344    async fn test_write_query_parameterized_delete() {
345        let db = TestDatabase::with_schema(
346            "CREATE TABLE users (id INTEGER, name TEXT);
347             INSERT INTO users VALUES (1, 'Alice');
348             INSERT INTO users VALUES (2, 'Bob');
349             INSERT INTO users VALUES (3, 'Charlie');",
350        )
351        .await;
352
353        let result = WriteQueryTool
354            .execute(WriteQueryInput {
355                query: "DELETE FROM users WHERE name = ?".to_string(),
356                params: vec![serde_json::json!("Bob")],
357                db_path: Some(db.key()),
358            })
359            .await
360            .unwrap();
361
362        let json = unwrap_json(result);
363        assert_eq!(json["rows_affected"], 1);
364        assert_eq!(db.count("users"), 2);
365    }
366
367    #[tokio::test]
368    async fn test_write_query_null_parameter() {
369        let db = TestDatabase::with_schema("CREATE TABLE users (id INTEGER, name TEXT)").await;
370
371        let result = WriteQueryTool
372            .execute(WriteQueryInput {
373                query: "INSERT INTO users VALUES (?, ?)".to_string(),
374                params: vec![serde_json::json!(1), serde_json::Value::Null],
375                db_path: Some(db.key()),
376            })
377            .await
378            .unwrap();
379
380        let json = unwrap_json(result);
381        assert_eq!(json["rows_affected"], 1);
382
383        let rows = db.query("SELECT name FROM users WHERE id = 1");
384        assert!(rows[0][0].is_null());
385    }
386
387    #[tokio::test]
388    async fn test_write_query_json_object_parameter() {
389        let db = TestDatabase::with_schema("CREATE TABLE data (id INTEGER, metadata TEXT)").await;
390
391        let result = WriteQueryTool
392            .execute(WriteQueryInput {
393                query: "INSERT INTO data VALUES (?, ?)".to_string(),
394                params: vec![
395                    serde_json::json!(1),
396                    serde_json::json!({"key": "value", "count": 42}),
397                ],
398                db_path: Some(db.key()),
399            })
400            .await
401            .unwrap();
402
403        let json = unwrap_json(result);
404        assert_eq!(json["rows_affected"], 1);
405
406        let rows = db.query("SELECT metadata FROM data WHERE id = 1");
407        let parsed: serde_json::Value = serde_json::from_str(rows[0][0].as_str().unwrap()).unwrap();
408        assert_eq!(parsed["key"], "value");
409        assert_eq!(parsed["count"], 42);
410    }
411
412    #[tokio::test]
413    async fn test_write_query_no_rows_affected() {
414        let db = TestDatabase::with_schema(
415            "CREATE TABLE users (id INTEGER, name TEXT);
416             INSERT INTO users VALUES (1, 'Alice');",
417        )
418        .await;
419
420        let result = WriteQueryTool
421            .execute(WriteQueryInput {
422                query: "DELETE FROM users WHERE id = 999".to_string(),
423                params: vec![],
424                db_path: Some(db.key()),
425            })
426            .await
427            .unwrap();
428
429        let json = unwrap_json(result);
430        assert_eq!(json["status"], "success");
431        assert_eq!(json["rows_affected"], 0);
432    }
433
434    #[tokio::test]
435    async fn test_update_no_last_insert_rowid() {
436        let db = TestDatabase::with_schema(
437            "CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT);
438             INSERT INTO users VALUES (1, 'Alice');",
439        )
440        .await;
441
442        let result = WriteQueryTool
443            .execute(WriteQueryInput {
444                query: "UPDATE users SET name = 'Updated' WHERE id = 1".to_string(),
445                params: vec![],
446                db_path: Some(db.key()),
447            })
448            .await
449            .unwrap();
450
451        let json = unwrap_json(result);
452        assert!(json["last_insert_rowid"].is_null());
453    }
454}