Record command exit code and duration

3 files changed, 100 insertions, 18 deletions

diff --git a/src/local/database.rs b/src/local/database.rs
index b94a6445..b2c009b6 100644
--- a/src/local/database.rs
+++ b/src/local/database.rs
@@ -1,6 +1,6 @@
 use std::path::Path;
 
-use eyre::Result;
+use eyre::{eyre, Result};
 
 use rusqlite::NO_PARAMS;
 use rusqlite::{params, Connection};
@@ -9,7 +9,9 @@ use crate::History;
 
 pub trait Database {
     fn save(&self, h: History) -> Result<()>;
+    fn load(&self, id: &str) -> Result<History>;
     fn list(&self) -> Result<()>;
+    fn update(&self, h: History) -> Result<()>;
 }
 
 // Intended for use on a developer machine and not a sync server.
@@ -44,8 +46,10 @@ impl SqliteDatabase {
 
         conn.execute(
             "create table if not exists history (
-                id integer primary key,
+                id text primary key,
                 timestamp integer not null,
+                duration integer not null,
+                exit integer not null,
                 command text not null,
                 cwd text not null
             )",
@@ -62,11 +66,53 @@ impl Database for SqliteDatabase {
 
         self.conn.execute(
             "insert into history (
+                id,
                 timestamp,
+                duration,
+                exit,
                 command,
                 cwd
-            ) values (?1, ?2, ?3)",
-            params![h.timestamp, h.command, h.cwd],
+            ) values (?1, ?2, ?3, ?4, ?5, ?6)",
+            params![h.id, h.timestamp, h.duration, h.exit, h.command, h.cwd],
+        )?;
+
+        Ok(())
+    }
+
+    fn load(&self, id: &str) -> Result<History> {
+        debug!("loading history item");
+
+        let mut stmt = self.conn.prepare(
+            "select id, timestamp, duration, exit, command, cwd from history
+                where id = ?1",
+        )?;
+
+        let iter = stmt.query_map(params![id], |row| {
+            Ok(History {
+                id: String::from(id),
+                timestamp: row.get(1)?,
+                duration: row.get(2)?,
+                exit: row.get(3)?,
+                command: row.get(4)?,
+                cwd: row.get(5)?,
+            })
+        })?;
+
+        for i in iter {
+            return Ok(i.unwrap());
+        }
+
+        return Err(eyre!("Failed to fetch history: {}", id));
+    }
+
+    fn update(&self, h: History) -> Result<()> {
+        debug!("updating sqlite history");
+
+        self.conn.execute(
+            "update history
+                set timestamp = ?2, duration = ?3, exit = ?4, command = ?5, cwd = ?6
+                where id = ?1",
+            params![h.id, h.timestamp, h.duration, h.exit, h.command, h.cwd],
         )?;
 
         Ok(())
@@ -77,19 +123,26 @@ impl Database for SqliteDatabase {
 
         let mut stmt = self
             .conn
-            .prepare("SELECT timestamp, command, cwd FROM history")?;
+            .prepare("SELECT id, timestamp, duration, exit, command, cwd FROM history")?;
+
         let history_iter = stmt.query_map(params![], |row| {
             Ok(History {
-                timestamp: row.get(0)?,
-                command: row.get(1)?,
-                cwd: row.get(2)?,
+                id: row.get(0)?,
+                timestamp: row.get(1)?,
+                duration: row.get(2)?,
+                exit: row.get(3)?,
+                command: row.get(4)?,
+                cwd: row.get(5)?,
             })
         })?;
 
         for h in history_iter {
             let h = h.unwrap();
 
-            println!("{}:{}:{}", h.timestamp, h.cwd, h.command);
+            println!(
+                "{} | {} | {} | {} | {}",
+                h.timestamp, h.cwd, h.duration, h.exit, h.command
+            );
         }
 
         Ok(())
diff --git a/src/local/history.rs b/src/local/history.rs
index e84d718c..3c9a9069 100644
--- a/src/local/history.rs
+++ b/src/local/history.rs
@@ -1,18 +1,25 @@
 use chrono;
+use uuid::Uuid;
 
 #[derive(Debug)]
 pub struct History {
+    pub id: String,
     pub timestamp: i64,
+    pub duration: i64,
+    pub exit: i64,
     pub command: String,
     pub cwd: String,
 }
 
 impl History {
-    pub fn new(command: String, cwd: String) -> History {
+    pub fn new(command: String, cwd: String, exit: i64, duration: i64) -> History {
         History {
+            id: Uuid::new_v4().to_simple().to_string(),
             timestamp: chrono::Utc::now().timestamp_millis(),
             command,
             cwd,
+            exit,
+            duration,
         }
     }
 }
diff --git a/src/main.rs b/src/main.rs
index a9b4b8af..57688a4a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -74,10 +74,20 @@ impl Atuin {
 #[derive(StructOpt)]
 enum HistoryCmd {
     #[structopt(
-        about="add a new command to the history",
-        aliases=&["a", "ad"],
+        about="begins a new command in the history",
+        aliases=&["s", "st", "sta", "star"],
     )]
-    Add { command: Vec<String> },
+    Start { command: Vec<String> },
+
+    #[structopt(
+        about="finishes a new command in the history (adds time, exit code)",
+        aliases=&["e", "en"],
+    )]
+    End {
+        id: String,
+        #[structopt(long, short)]
+        exit: i64,
+    },
 
     #[structopt(
         about="list all items in history",
@@ -87,16 +97,28 @@ enum HistoryCmd {
 }
 
 impl HistoryCmd {
-    fn run(self, db: SqliteDatabase) -> Result<()> {
+    fn run(&self, db: SqliteDatabase) -> Result<()> {
         match self {
-            HistoryCmd::Add { command: words } => {
+            HistoryCmd::Start { command: words } => {
                 let command = words.join(" ");
                 let cwd = env::current_dir()?.display().to_string();
-                let h = History::new(command, cwd);
 
-                debug!("adding history: {:?}", h);
+                let h = History::new(command, cwd, -1, -1);
+
+                // print the ID
+                // we use this as the key for calling end
+                println!("{}", h.id);
                 db.save(h)?;
-                debug!("saved history to sqlite");
+                Ok(())
+            }
+
+            HistoryCmd::End { id, exit } => {
+                let mut h = db.load(id)?;
+                h.exit = *exit;
+                h.duration = chrono::Utc::now().timestamp_millis() - h.timestamp;
+
+                db.update(h)?;
+
                 Ok(())
             }