command_vault/db/
store.rs

1//! Database module for command-vault
2//! 
3//! This module provides SQLite-based storage for commands, tags, and parameters.
4//! It handles all database operations including CRUD operations for commands,
5//! tag management, and search functionality.
6
7use anyhow::{Result, anyhow};
8use rusqlite::Connection;
9use chrono::Utc;
10use serde_json;
11
12use super::models::Command;
13
14/// The main database interface for command-vault.
15/// 
16/// Handles all database operations including:
17/// - Command storage and retrieval
18/// - Tag management
19/// - Search functionality
20/// 
21/// # Example
22/// ```no_run
23/// use anyhow::Result;
24/// use command_vault::db::Database;
25/// 
26/// fn main() -> Result<()> {
27///     let db = Database::new("commands.db")?;
28///     db.init()?;
29///     Ok(())
30/// }
31/// ```
32pub struct Database {
33    conn: Connection,
34}
35
36impl Database {
37    /// Creates a new database connection.
38    /// 
39    /// # Arguments
40    /// * `path` - Path to the SQLite database file
41    /// 
42    /// # Returns
43    /// * `Result<Database>` - A new database instance
44    pub fn new(path: &str) -> Result<Self> {
45        let conn = Connection::open(path)?;
46        let db = Database { conn };
47        db.init()?;
48        Ok(db)
49    }
50
51    /// Initializes the database schema.
52    /// 
53    /// Creates the following tables if they don't exist:
54    /// - commands: Stores command information
55    /// - tags: Stores tag information
56    /// - command_tags: Links commands to tags
57    pub fn init(&self) -> Result<()> {
58        // Create commands table
59        self.conn.execute(
60            "CREATE TABLE IF NOT EXISTS commands (
61                id INTEGER PRIMARY KEY,
62                command TEXT NOT NULL,
63                timestamp TEXT NOT NULL,
64                directory TEXT NOT NULL,
65                tags TEXT NOT NULL DEFAULT '',
66                parameters TEXT NOT NULL DEFAULT '[]'
67            )",
68            [],
69        )?;
70        
71        // Create tags table
72        self.conn.execute(
73            "CREATE TABLE IF NOT EXISTS tags (
74                id INTEGER PRIMARY KEY,
75                name TEXT NOT NULL UNIQUE
76            )",
77            [],
78        )?;
79
80        // Create command_tags table for many-to-many relationship
81        self.conn.execute(
82            "CREATE TABLE IF NOT EXISTS command_tags (
83                command_id INTEGER NOT NULL,
84                tag_id INTEGER NOT NULL,
85                PRIMARY KEY (command_id, tag_id),
86                FOREIGN KEY (command_id) REFERENCES commands(id) ON DELETE CASCADE,
87                FOREIGN KEY (tag_id) REFERENCES tags(id) ON DELETE CASCADE
88            )",
89            [],
90        )?;
91        
92        // Create indexes
93        self.conn.execute(
94            "CREATE INDEX IF NOT EXISTS idx_commands_command ON commands(command)",
95            [],
96        )?;
97        self.conn.execute(
98            "CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(name)",
99            [],
100        )?;
101        
102        Ok(())
103    }
104
105    /// Adds a new command to the database.
106    /// 
107    /// # Arguments
108    /// * `command` - The command to add
109    /// 
110    /// # Returns
111    /// * `Result<i64>` - The ID of the newly added command
112    pub fn add_command(&mut self, command: &Command) -> Result<i64> {
113        let tx = self.conn.transaction()?;
114        
115        // Insert the command
116        tx.execute(
117            "INSERT INTO commands (command, timestamp, directory, tags, parameters)
118             VALUES (?1, ?2, ?3, ?4, ?5)",
119            (
120                &command.command,
121                &command.timestamp.to_rfc3339(),
122                &command.directory,
123                &command.tags.join(","),
124                &serde_json::to_string(&command.parameters)?,
125            ),
126        )?;
127        
128        let command_id = tx.last_insert_rowid();
129        
130        // Add tags if present
131        for tag in &command.tags {
132            // Insert or get tag
133            tx.execute(
134                "INSERT OR IGNORE INTO tags (name) VALUES (?1)",
135                [tag],
136            )?;
137            
138            let tag_id: i64 = tx.query_row(
139                "SELECT id FROM tags WHERE name = ?1",
140                [tag],
141                |row| row.get(0),
142            )?;
143            
144            // Link command to tag
145            tx.execute(
146                "INSERT OR IGNORE INTO command_tags (command_id, tag_id) VALUES (?1, ?2)",
147                rusqlite::params![command_id, tag_id],
148            )?;
149        }
150        
151        tx.commit()?;
152        Ok(command_id)
153    }
154
155    /// Adds tags to an existing command.
156    /// 
157    /// # Arguments
158    /// * `command_id` - The ID of the command to add tags to
159    /// * `tags` - The tags to add
160    /// 
161    /// # Returns
162    /// * `Result<()>` - Success or failure
163    pub fn add_tags_to_command(&mut self, command_id: i64, tags: &[String]) -> Result<()> {
164        let tx = self.conn.transaction()?;
165        
166        // Verify command exists
167        let exists: bool = tx.query_row(
168            "SELECT 1 FROM commands WHERE id = ?1",
169            [command_id],
170            |_| Ok(true),
171        ).unwrap_or(false);
172        
173        if !exists {
174            return Err(anyhow!("Command not found"));
175        }
176        
177        // Get current tags
178        let mut current_tags = Vec::new();
179        {
180            let mut stmt = tx.prepare(
181                "SELECT t.name 
182                 FROM tags t 
183                 JOIN command_tags ct ON ct.tag_id = t.id 
184                 WHERE ct.command_id = ?1"
185            )?;
186            let mut rows = stmt.query([command_id])?;
187            while let Some(row) = rows.next()? {
188                current_tags.push(row.get::<_, String>(0)?);
189            }
190        }
191        
192        for tag in tags {
193            // Skip if tag already exists
194            if current_tags.contains(tag) {
195                continue;
196            }
197            
198            // Insert or get tag
199            tx.execute(
200                "INSERT OR IGNORE INTO tags (name) VALUES (?1)",
201                [tag],
202            )?;
203            
204            let tag_id: i64 = tx.query_row(
205                "SELECT id FROM tags WHERE name = ?1",
206                [tag],
207                |row| row.get(0),
208            )?;
209            
210            // Link command to tag
211            tx.execute(
212                "INSERT OR IGNORE INTO command_tags (command_id, tag_id) VALUES (?1, ?2)",
213                rusqlite::params![command_id, tag_id],
214            )?;
215            
216            // Update tags string in commands table
217            current_tags.push(tag.clone());
218        }
219        
220        // Update the tags string in the commands table
221        tx.execute(
222            "UPDATE commands SET tags = ?1 WHERE id = ?2",
223            rusqlite::params![current_tags.join(","), command_id],
224        )?;
225        
226        tx.commit()?;
227        Ok(())
228    }
229
230    /// Removes a tag from a command.
231    /// 
232    /// # Arguments
233    /// * `command_id` - The ID of the command to remove the tag from
234    /// * `tag_name` - The name of the tag to remove
235    /// 
236    /// # Returns
237    /// * `Result<()>` - Success or failure
238    pub fn remove_tag_from_command(&mut self, command_id: i64, tag_name: &str) -> Result<()> {
239        let tx = self.conn.transaction()?;
240        
241        tx.execute(
242            "DELETE FROM command_tags 
243             WHERE command_id = ?1 
244             AND tag_id = (SELECT id FROM tags WHERE name = ?2)",
245            rusqlite::params![command_id, tag_name],
246        )?;
247        
248        tx.commit()?;
249        Ok(())
250    }
251
252    /// Searches for commands containing a given query string.
253    /// 
254    /// # Arguments
255    /// * `query` - The query string to search for
256    /// * `limit` - The maximum number of results to return
257    /// 
258    /// # Returns
259    /// * `Result<Vec<Command>>` - A list of matching commands
260    pub fn search_commands(&self, query: &str, limit: usize) -> Result<Vec<Command>> {
261        let mut stmt = self.conn.prepare(
262            "SELECT c.id, c.command, c.timestamp, c.directory, c.tags, c.parameters 
263             FROM commands c
264             WHERE c.command LIKE '%' || ?1 || '%'
265             ORDER BY c.timestamp DESC
266             LIMIT ?2"
267        )?;
268
269        let mut rows = stmt.query([query, &limit.to_string()])?;
270        let mut commands = Vec::new();
271
272        while let Some(row) = rows.next()? {
273            let id: i64 = row.get(0)?;
274            commands.push(Command {
275                id: Some(id),
276                command: row.get(1)?,
277                timestamp: chrono::DateTime::parse_from_rfc3339(&row.get::<_, String>(2)?)?
278                    .with_timezone(&Utc),
279                directory: row.get(3)?,
280                tags: row.get::<_, String>(4)?
281                    .split(',')
282                    .filter(|s| !s.is_empty())
283                    .map(|s| s.to_string())
284                    .collect(),
285                parameters: serde_json::from_str(&row.get::<_, String>(5)?)?,
286            });
287        }
288
289        Ok(commands)
290    }
291
292    /// Searches for commands with a given tag.
293    /// 
294    /// # Arguments
295    /// * `tag` - The tag to search for
296    /// * `limit` - The maximum number of results to return
297    /// 
298    /// # Returns
299    /// * `Result<Vec<Command>>` - A list of matching commands
300    pub fn search_by_tag(&self, tag: &str, limit: usize) -> Result<Vec<Command>> {
301        let mut stmt = self.conn.prepare(
302            "SELECT DISTINCT c.id, c.command, c.timestamp, c.directory, c.tags, c.parameters 
303             FROM commands c
304             JOIN command_tags ct ON ct.command_id = c.id
305             JOIN tags t ON t.id = ct.tag_id
306             WHERE t.name = ?1
307             ORDER BY c.timestamp DESC
308             LIMIT ?2"
309        )?;
310
311        let mut rows = stmt.query([tag, &limit.to_string()])?;
312        let mut commands = Vec::new();
313
314        while let Some(row) = rows.next()? {
315            let id: i64 = row.get(0)?;
316            commands.push(Command {
317                id: Some(id),
318                command: row.get(1)?,
319                timestamp: chrono::DateTime::parse_from_rfc3339(&row.get::<_, String>(2)?)?
320                    .with_timezone(&Utc),
321                directory: row.get(3)?,
322                tags: row.get::<_, String>(4)?
323                    .split(',')
324                    .filter(|s| !s.is_empty())
325                    .map(|s| s.to_string())
326                    .collect(),
327                parameters: serde_json::from_str(&row.get::<_, String>(5)?)?,
328            });
329        }
330
331        Ok(commands)
332    }
333
334    /// Lists all tags in the database.
335    /// 
336    /// # Returns
337    /// * `Result<Vec<(String, i64)>>` - A list of tags with their respective counts
338    pub fn list_tags(&self) -> Result<Vec<(String, i64)>> {
339        let mut stmt = self.conn.prepare(
340            "SELECT t.name, COUNT(ct.command_id) as count
341             FROM tags t
342             LEFT JOIN command_tags ct ON ct.tag_id = t.id
343             GROUP BY t.id, t.name
344             ORDER BY count DESC, t.name"
345        )?;
346        
347        let tags = stmt.query_map([], |row| {
348            Ok((row.get::<_, String>(0)?, row.get(1)?))
349        })?
350        .collect::<std::result::Result<Vec<_>, _>>()?;
351        
352        Ok(tags)
353    }
354
355    /// Lists all commands in the database.
356    /// 
357    /// # Arguments
358    /// * `limit` - The maximum number of results to return
359    /// * `ascending` - Whether to return results in ascending order
360    /// 
361    /// # Returns
362    /// * `Result<Vec<Command>>` - A list of commands
363    pub fn list_commands(&self, limit: usize, ascending: bool) -> Result<Vec<Command>> {
364        let query = if ascending {
365            if limit == 0 {
366                "SELECT c.id, c.command, c.timestamp, c.directory, c.tags, c.parameters 
367                 FROM commands c
368                 ORDER BY c.timestamp ASC"
369            } else {
370                "SELECT c.id, c.command, c.timestamp, c.directory, c.tags, c.parameters 
371                 FROM commands c
372                 ORDER BY c.timestamp ASC
373                 LIMIT ?1"
374            }
375        } else {
376            if limit == 0 {
377                "SELECT c.id, c.command, c.timestamp, c.directory, c.tags, c.parameters 
378                 FROM commands c
379                 ORDER BY c.timestamp DESC"
380            } else {
381                "SELECT c.id, c.command, c.timestamp, c.directory, c.tags, c.parameters 
382                 FROM commands c
383                 ORDER BY c.timestamp DESC
384                 LIMIT ?1"
385            }
386        };
387
388        let mut stmt = self.conn.prepare(query)?;
389        let mut rows = if limit == 0 {
390            stmt.query([])?
391        } else {
392            stmt.query([limit])?
393        };
394        
395        let mut commands = Vec::new();
396
397        while let Some(row) = rows.next()? {
398            let id: i64 = row.get(0)?;
399            commands.push(Command {
400                id: Some(id),
401                command: row.get(1)?,
402                timestamp: chrono::DateTime::parse_from_rfc3339(&row.get::<_, String>(2)?)?
403                    .with_timezone(&Utc),
404                directory: row.get(3)?,
405                tags: row.get::<_, String>(4)?
406                    .split(',')
407                    .filter(|s| !s.is_empty())
408                    .map(|s| s.to_string())
409                    .collect(),
410                parameters: serde_json::from_str(&row.get::<_, String>(5)?)?,
411            });
412        }
413
414        Ok(commands)
415    }
416
417    /// Gets a command by its ID.
418    /// 
419    /// # Arguments
420    /// * `id` - The ID of the command to retrieve
421    /// 
422    /// # Returns
423    /// * `Result<Option<Command>>` - The command if found
424    pub fn get_command(&self, id: i64) -> Result<Option<Command>> {
425        // First get the command details
426        let mut stmt = self.conn.prepare(
427            "SELECT command, timestamp, directory, parameters 
428             FROM commands 
429             WHERE id = ?1"
430        )?;
431
432        let command = stmt.query_row([id], |row| {
433            Ok((
434                row.get::<_, String>(0)?,
435                row.get::<_, String>(1)?,
436                row.get::<_, String>(2)?,
437                row.get::<_, String>(3)?,
438            ))
439        });
440
441        if let Ok((command, timestamp, directory, parameters)) = command {
442            // Then get the tags
443            let mut stmt = self.conn.prepare(
444                "SELECT t.name 
445                 FROM tags t 
446                 JOIN command_tags ct ON ct.tag_id = t.id 
447                 WHERE ct.command_id = ?1"
448            )?;
449
450            let mut tags = Vec::new();
451            let mut rows = stmt.query([id])?;
452            while let Some(row) = rows.next()? {
453                tags.push(row.get::<_, String>(0)?);
454            }
455
456            Ok(Some(Command {
457                id: Some(id),
458                command,
459                timestamp: chrono::DateTime::parse_from_rfc3339(&timestamp)?
460                    .with_timezone(&Utc),
461                directory,
462                tags,
463                parameters: serde_json::from_str(&parameters)?,
464            }))
465        } else {
466            Ok(None)
467        }
468    }
469
470    /// Updates an existing command.
471    /// 
472    /// # Arguments
473    /// * `command` - The updated command
474    /// 
475    /// # Returns
476    /// * `Result<()>` - Success or failure
477    pub fn update_command(&mut self, command: &Command) -> Result<()> {
478        if command.id.is_none() {
479            return Err(anyhow!("Cannot update command without id"));
480        }
481
482        let tx = self.conn.transaction()?;
483        
484        // Update command
485        tx.execute(
486            "UPDATE commands 
487             SET command = ?1, 
488                 timestamp = ?2,
489                 directory = ?3,
490                 tags = ?4,
491                 parameters = ?5
492             WHERE id = ?6",
493            rusqlite::params![
494                command.command,
495                command.timestamp.to_rfc3339(),
496                command.directory,
497                command.tags.join(","),
498                serde_json::to_string(&command.parameters)?,
499                command.id.unwrap()
500            ],
501        )?;
502
503        // Delete existing tags
504        tx.execute(
505            "DELETE FROM command_tags WHERE command_id = ?1",
506            [command.id.unwrap()],
507        )?;
508
509        // Add new tags using the same transaction
510        for tag in &command.tags {
511            // Insert or get tag
512            tx.execute(
513                "INSERT OR IGNORE INTO tags (name) VALUES (?1)",
514                [tag],
515            )?;
516            
517            let tag_id: i64 = tx.query_row(
518                "SELECT id FROM tags WHERE name = ?1",
519                [tag],
520                |row| row.get(0),
521            )?;
522            
523            // Link command to tag
524            tx.execute(
525                "INSERT OR IGNORE INTO command_tags (command_id, tag_id) VALUES (?1, ?2)",
526                rusqlite::params![command.id.unwrap(), tag_id],
527            )?;
528        }
529        
530        tx.commit()?;
531        Ok(())
532    }
533
534    /// Deletes a command by its ID.
535    /// 
536    /// # Arguments
537    /// * `command_id` - The ID of the command to delete
538    /// 
539    /// # Returns
540    /// * `Result<()>` - Success or failure
541    pub fn delete_command(&mut self, command_id: i64) -> Result<()> {
542        let tx = self.conn.transaction()?;
543        
544        // First delete from command_tags
545        tx.execute(
546            "DELETE FROM command_tags WHERE command_id = ?",
547            [command_id],
548        )?;
549
550        // Then delete from commands
551        let rows_affected = tx.execute(
552            "DELETE FROM commands WHERE id = ?",
553            [command_id],
554        )?;
555
556        if rows_affected == 0 {
557            return Err(anyhow!("Command not found"));
558        }
559
560        // Clean up unused tags
561        tx.execute(
562            "DELETE FROM tags WHERE id NOT IN (SELECT DISTINCT tag_id FROM command_tags)",
563            [],
564        )?;
565
566        tx.commit()?;
567        Ok(())
568    }
569}
command_vault/db/store.rs

command_vault/db/
store.rs
