1 files changed, 69 insertions, 1 deletions

diff --git a/src/database.rs b/src/database.rs
index 8752255..da514ef 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -2,6 +2,8 @@ use rusqlite::{Connection, Result, Rows, Statement};
 use rusqlite::config::DbConfig::{*};
 use chrono::Utc;
 use std::path::PathBuf;
+use text_io::scan;
+use std::io::{stdout, Write};
 
 pub struct Database {
     date: i64,
@@ -23,6 +25,12 @@ impl Database {
     pub fn add_register(&self, filepath: &PathBuf) -> Result<()> {
         let conn = Connection::open(filepath)?;
         conn.set_db_config(SQLITE_DBCONFIG_ENABLE_FKEY, true)?;
+
+        if !&self.check_agent_exists(&conn)? {
+            println!("The Agent '{}' is not registered in the database", &self.person);
+            return Ok(());
+        }
+
         conn.execute("
             INSERT INTO Registers (agent_id, register_date, amount, note)
             VALUES (
@@ -45,6 +53,38 @@ impl Database {
         Ok(())
     }
 
+    pub fn update_agent(&self, filepath: &PathBuf, new_name: String) -> Result <()> {
+        let conn = Connection::open(filepath)?;
+
+        if !&self.check_agent_exists(&conn)? {
+            println!("The Agent '{}' is not registered in the database", &self.person);
+            return Ok(());
+        }
+
+        conn.execute("UPDATE Agents SET name=?1 WHERE name=?2", (&new_name, &self.person))?;
+
+        println!("Agent '{}' was updated to '{}'", &self.person, &new_name);
+        Ok(())
+    }
+
+    pub fn delete_agent(&self, filepath: &PathBuf) -> Result<()>{
+        let conn = Connection::open(filepath)?;
+        let prompt = format!("The Agent '{}' will be removed are you sure?", &self.person);
+
+        if !&self.check_agent_exists(&conn)? {
+            println!("The Agent '{}' is not registered in the database", &self.person);
+            return Ok(());
+        }
+        if !Database::ask_user_confirmaton(prompt.as_str()) {
+            println!("Operation aborted");
+            return Ok(());
+        }
+        conn.set_db_config(SQLITE_DBCONFIG_ENABLE_FKEY, true)?;
+        conn.execute("DELETE FROM Agents WHERE name=?1", (&self.person,))?;
+        println!("agent '{}' deleted successfully", &self.person);
+        Ok(())
+    }
+
     pub fn view_history(filepath: &PathBuf, filter: Option<String>) -> Result<()> {
         let conn = Connection::open(filepath)?;
         let mut hist_query: String = "
@@ -123,9 +163,37 @@ impl Database {
                 register_date INTEGER NOT NULL,
                 amount INTEGER NOT NULL,
                 note TEXT,
-                FOREIGN KEY(agent_id) REFERENCES Agents(id))", ())
+                FOREIGN KEY(agent_id) REFERENCES Agents(id)
+                    ON DELETE CASCADE)", ())
             .expect("SQL initialization error");
         println!("'{}' database created", &filepath.display());
         Ok(())
     }
+
+    fn check_agent_exists(&self, conn: &Connection) -> Result<bool> {
+        let mut stmt = conn.prepare("SELECT name FROM Agents WHERE name=?1")?;
+        let mut reg = stmt.query([&self.person])?;
+
+        if let Some(_) = reg.next()? {
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+
+    fn ask_user_confirmaton(prompt: &str) -> bool {
+        let mut input: String;
+        let complete_prompt = format!("{} (Y/n)", prompt);
+        loop {
+            print!("{}: ", complete_prompt);
+            stdout().flush().unwrap();
+            scan!("{}\n", input);
+
+            match input.to_lowercase().as_str() {
+                "y"|"" => return true,
+                "n" => return false,
+                _ => println!("Invalid input")
+            }
+        }
+    }
 }