Implement history import

10 files changed, 361 insertions, 82 deletions

diff --git a/Cargo.lock b/Cargo.lock
index feb04044..fcf1175d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -49,11 +49,13 @@ dependencies = [
 
 [[package]]
 name = "atuin"
-version = "0.1.1"
+version = "0.2.0"
 dependencies = [
  "chrono",
  "directories",
  "eyre",
+ "home",
+ "indicatif",
  "log",
  "pretty_env_logger",
  "rusqlite",
@@ -132,6 +134,21 @@ dependencies = [
 ]
 
 [[package]]
+name = "console"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7cc80946b3480f421c2f17ed1cb841753a371c7c5104f51d507e13f532c856aa"
+dependencies = [
+ "encode_unicode",
+ "lazy_static",
+ "libc",
+ "regex",
+ "terminal_size",
+ "unicode-width",
+ "winapi",
+]
+
+[[package]]
 name = "constant_time_eq"
 version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -190,6 +207,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "encode_unicode"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f"
+
+[[package]]
 name = "env_logger"
 version = "0.7.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -283,6 +306,15 @@ dependencies = [
 ]
 
 [[package]]
+name = "home"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2456aef2e6b6a9784192ae780c0f15bc57df0e918585282325e8c8ac27737654"
+dependencies = [
+ "winapi",
+]
+
+[[package]]
 name = "humantime"
 version = "1.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -298,6 +330,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f4d5eb2e114fec2b7fe0fadc22888ad2658789bb7acac4dbee9cf8389f971ec8"
 
 [[package]]
+name = "indicatif"
+version = "0.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7baab56125e25686df467fe470785512329883aab42696d661247aca2a2896e4"
+dependencies = [
+ "console",
+ "lazy_static",
+ "number_prefix",
+ "regex",
+]
+
+[[package]]
 name = "lazy_static"
 version = "1.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -355,6 +399,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "number_prefix"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "17b02fc0ff9a9e4b35b3342880f48e896ebf69f2967921fe8646bf5b7125956a"
+
+[[package]]
 name = "once_cell"
 version = "1.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -577,6 +627,16 @@ dependencies = [
 ]
 
 [[package]]
+name = "terminal_size"
+version = "0.1.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "86ca8ced750734db02076f44132d802af0b33b09942331f4459dde8636fd2406"
+dependencies = [
+ "libc",
+ "winapi",
+]
+
+[[package]]
 name = "textwrap"
 version = "0.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index 3172a308..57117c29 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "atuin"
-version = "0.1.1"
+version = "0.2.0"
 authors = ["Ellie Huxtable <e@elm.sh>"]
 edition = "2018"
 license = "MIT"
@@ -15,6 +15,8 @@ shellexpand = "2.*"
 structopt = "0.3.*"
 directories = "3.*"
 uuid = { version = "0.8", features = ["serde", "v4"] }
+home = "0.5.3"
+indicatif = "0.15.0"
 
 [dependencies.rusqlite]
 version = "0.24.*"
diff --git a/src/command/history.rs b/src/command/history.rs
new file mode 100644
index 00000000..72f821c5
--- /dev/null
+++ b/src/command/history.rs
@@ -0,0 +1,63 @@
+use std::env;
+
+use eyre::Result;
+use structopt::StructOpt;
+
+use crate::local::database::{Database, SqliteDatabase};
+use crate::local::history::History;
+
+#[derive(StructOpt)]
+pub enum HistoryCmd {
+    #[structopt(
+        about="begins a new command in the history",
+        aliases=&["s", "st", "sta", "star"],
+    )]
+    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",
+        aliases=&["l", "li", "lis"],
+    )]
+    List,
+}
+
+impl HistoryCmd {
+    pub fn run(&self, db: SqliteDatabase) -> Result<()> {
+        match self {
+            HistoryCmd::Start { command: words } => {
+                let command = words.join(" ");
+                let cwd = env::current_dir()?.display().to_string();
+
+                let h = History::new(chrono::Utc::now().timestamp_nanos(), command, cwd, -1, -1);
+
+                // print the ID
+                // we use this as the key for calling end
+                println!("{}", h.id);
+                db.save(h)?;
+                Ok(())
+            }
+
+            HistoryCmd::End { id, exit } => {
+                let mut h = db.load(id)?;
+                h.exit = *exit;
+                h.duration = chrono::Utc::now().timestamp_nanos() - h.timestamp;
+
+                db.update(h)?;
+
+                Ok(())
+            }
+
+            HistoryCmd::List => db.list(),
+        }
+    }
+}
diff --git a/src/command/import.rs b/src/command/import.rs
new file mode 100644
index 00000000..5ece13a8
--- /dev/null
+++ b/src/command/import.rs
@@ -0,0 +1,107 @@
+use std::env;
+use std::path::PathBuf;
+
+use eyre::{eyre, Result};
+use home::home_dir;
+use structopt::StructOpt;
+
+use crate::local::database::{Database, SqliteDatabase};
+use crate::local::history::History;
+use crate::local::import::ImportZsh;
+use indicatif::ProgressBar;
+
+#[derive(StructOpt)]
+pub enum ImportCmd {
+    #[structopt(
+        about="import history for the current shell",
+        aliases=&["a", "au", "aut"],
+    )]
+    Auto,
+
+    #[structopt(
+        about="import history from the zsh history file",
+        aliases=&["z", "zs"],
+    )]
+    Zsh,
+}
+
+impl ImportCmd {
+    fn import_zsh(&self, db: &mut SqliteDatabase) -> Result<()> {
+        // oh-my-zsh sets HISTFILE=~/.zhistory
+        // zsh has no default value for this var, but uses ~/.zhistory.
+        // we could maybe be smarter about this in the future :)
+
+        let histpath = env::var("HISTFILE");
+
+        let histpath = match histpath {
+            Ok(p) => PathBuf::from(p),
+            Err(_) => {
+                let mut home = home_dir().unwrap();
+                home.push(".zhistory");
+
+                home
+            }
+        };
+
+        if !histpath.exists() {
+            return Err(eyre!(
+                "Could not find history file at {}, try setting $HISTFILE",
+                histpath.to_str().unwrap()
+            ));
+        }
+
+        let zsh = ImportZsh::new(histpath.to_str().unwrap())?;
+
+        let progress = ProgressBar::new(zsh.loc);
+
+        let buf_size = 100;
+        let mut buf = Vec::<History>::with_capacity(buf_size);
+
+        for i in zsh {
+            match i {
+                Ok(h) => {
+                    buf.push(h);
+                }
+                Err(e) => {
+                    error!("{}", e);
+                    continue;
+                }
+            }
+
+            if buf.len() == buf_size {
+                db.save_bulk(&buf)?;
+                progress.inc(buf.len() as u64);
+
+                buf = Vec::<History>::with_capacity(buf_size);
+            }
+        }
+
+        if buf.len() > 0 {
+            db.save_bulk(&buf)?;
+            progress.inc(buf.len() as u64);
+        }
+
+        progress.finish_with_message("Imported history!");
+
+        Ok(())
+    }
+
+    pub fn run(&self, db: &mut SqliteDatabase) -> Result<()> {
+        match self {
+            ImportCmd::Auto => {
+                let shell = env::var("SHELL").unwrap_or(String::from("NO_SHELL"));
+
+                match shell.as_str() {
+                    "/bin/zsh" => self.import_zsh(db),
+
+                    _ => {
+                        println!("cannot import {} history", shell);
+                        Ok(())
+                    }
+                }
+            }
+
+            ImportCmd::Zsh => Ok(()),
+        }
+    }
+}
diff --git a/src/command/mod.rs b/src/command/mod.rs
new file mode 100644
index 00000000..c61d2280
--- /dev/null
+++ b/src/command/mod.rs
@@ -0,0 +1,2 @@
+pub mod history;
+pub mod import;
diff --git a/src/local/database.rs b/src/local/database.rs
index b2c009b6..2a4cc582 100644
--- a/src/local/database.rs
+++ b/src/local/database.rs
@@ -9,6 +9,7 @@ use crate::History;
 
 pub trait Database {
     fn save(&self, h: History) -> Result<()>;
+    fn save_bulk(&mut self, h: &Vec<History>) -> Result<()>;
     fn load(&self, id: &str) -> Result<History>;
     fn list(&self) -> Result<()>;
     fn update(&self, h: History) -> Result<()>;
@@ -51,7 +52,9 @@ impl SqliteDatabase {
                 duration integer not null,
                 exit integer not null,
                 command text not null,
-                cwd text not null
+                cwd text not null,
+
+                unique(timestamp, cwd, command)
             )",
             NO_PARAMS,
         )?;
@@ -65,7 +68,7 @@ impl Database for SqliteDatabase {
         debug!("saving history to sqlite");
 
         self.conn.execute(
-            "insert into history (
+            "insert or ignore into history (
                 id,
                 timestamp,
                 duration,
@@ -79,6 +82,30 @@ impl Database for SqliteDatabase {
         Ok(())
     }
 
+    fn save_bulk(&mut self, h: &Vec<History>) -> Result<()> {
+        debug!("saving history to sqlite");
+
+        let tx = self.conn.transaction()?;
+
+        for i in h {
+            tx.execute(
+                "insert or ignore into history (
+                id,
+                timestamp,
+                duration,
+                exit,
+                command,
+                cwd
+            ) values (?1, ?2, ?3, ?4, ?5, ?6)",
+                params![i.id, i.timestamp, i.duration, i.exit, i.command, i.cwd],
+            )?;
+        }
+
+        tx.commit()?;
+
+        Ok(())
+    }
+
     fn load(&self, id: &str) -> Result<History> {
         debug!("loading history item");
 
diff --git a/src/local/history.rs b/src/local/history.rs
index 3c9a9069..00109621 100644
--- a/src/local/history.rs
+++ b/src/local/history.rs
@@ -12,10 +12,10 @@ pub struct History {
 }
 
 impl History {
-    pub fn new(command: String, cwd: String, exit: i64, duration: i64) -> History {
+    pub fn new(timestamp: i64, command: String, cwd: String, exit: i64, duration: i64) -> History {
         History {
             id: Uuid::new_v4().to_simple().to_string(),
-            timestamp: chrono::Utc::now().timestamp_millis(),
+            timestamp,
             command,
             cwd,
             exit,
diff --git a/src/local/import.rs b/src/local/import.rs
index 8db8f0e3..ce141c52 100644
--- a/src/local/import.rs
+++ b/src/local/import.rs
@@ -4,38 +4,109 @@
 use std::fs::File;
 use std::io::{BufRead, BufReader};
 
-use eyre::Result;
+use chrono::{TimeZone, Utc};
+use eyre::{eyre, Result};
 
-use crate::models::history::History;
+use crate::local::history::History;
 
-pub struct ImportBash {
+#[derive(Debug)]
+pub struct ImportZsh {
     file: BufReader<File>,
+
+    pub loc: u64,
 }
 
-impl ImportBash {
-    pub fn new(path: &str) -> Result<ImportBash> {
+// this could probably be sped up
+fn count_lines(path: &str) -> Result<usize> {
+    let file = File::open(path)?;
+    let buf = BufReader::new(file);
+
+    Ok(buf.lines().count())
+}
+
+impl ImportZsh {
+    pub fn new(path: &str) -> Result<ImportZsh> {
+        let loc = count_lines(path)?;
+
         let file = File::open(path)?;
         let buf = BufReader::new(file);
 
-        Ok(ImportBash { file: buf })
+        Ok(ImportZsh {
+            file: buf,
+            loc: loc as u64,
+        })
     }
 }
 
-impl Iterator for ImportBash {
-    type Item = History;
+fn trim_newline(s: &str) -> String {
+    let mut s = String::from(s);
+
+    if s.ends_with('\n') {
+        s.pop();
+        if s.ends_with('\r') {
+            s.pop();
+        }
+    }
+
+    s
+}
 
-    fn next(&mut self) -> Option<History> {
+fn parse_extended(line: String) -> History {
+    let line = line.replacen(": ", "", 2);
+    let mut split = line.splitn(2, ":");
+
+    let time = split.next().unwrap_or("-1");
+    let time = time
+        .parse::<i64>()
+        .unwrap_or(chrono::Utc::now().timestamp_nanos());
+
+    let duration = split.next().unwrap(); // might be 0;the command
+    let mut split = duration.split(";");
+
+    let duration = split.next().unwrap_or("-1"); // should just be the 0
+    let duration = duration.parse::<i64>().unwrap_or(-1);
+
+    let command = split.next().unwrap();
+
+    // use nanos, because why the hell not? we won't display them.
+    History::new(
+        Utc.timestamp(time, 0).timestamp_nanos(),
+        trim_newline(command),
+        String::from("unknown"),
+        -1,
+        duration * 1_000_000_000,
+    )
+}
+
+impl Iterator for ImportZsh {
+    type Item = Result<History>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // ZSH extended history records the timestamp + command duration
+        // These lines begin with :
+        // So, if the line begins with :, parse it. Otherwise it's just
+        // the command
         let mut line = String::new();
 
         match self.file.read_line(&mut line) {
             Ok(0) => None,
-            Err(_) => None,
+            Err(e) => Some(Err(eyre!("failed to parse line: {}", e))),
+
+            Ok(_) => {
+                let extended = line.starts_with(":");
 
-            Ok(_) => Some(History {
-                cwd: "none".to_string(),
-                command: line,
-                timestamp: -1,
-            }),
+                if extended {
+                    Some(Ok(parse_extended(line)))
+                } else {
+                    Some(Ok(History::new(
+                        chrono::Utc::now().timestamp_nanos(), // what else? :/
+                        trim_newline(line.as_str()),
+                        String::from("unknown"),
+                        -1,
+                        -1,
+                    )))
+                }
+            }
         }
     }
 }
diff --git a/src/local/mod.rs b/src/local/mod.rs
index f587d016..a11ee213 100644
--- a/src/local/mod.rs
+++ b/src/local/mod.rs
@@ -1,2 +1,3 @@
 pub mod database;
 pub mod history;
+pub mod import;
diff --git a/src/main.rs b/src/main.rs
index 57688a4a..19357cbe 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,3 @@
-use std::env;
 use std::path::PathBuf;
 
 use directories::ProjectDirs;
@@ -9,11 +8,13 @@ use structopt::StructOpt;
 extern crate log;
 use pretty_env_logger;
 
-mod local;
-
+use command::{history::HistoryCmd, import::ImportCmd};
 use local::database::{Database, SqliteDatabase};
 use local::history::History;
 
+mod command;
+mod local;
+
 #[derive(StructOpt)]
 #[structopt(
     author = "Ellie Huxtable <e@elm.sh>",
@@ -37,7 +38,7 @@ enum AtuinCmd {
     History(HistoryCmd),
 
     #[structopt(about = "import shell history from file")]
-    Import,
+    Import(ImportCmd),
 
     #[structopt(about = "start a atuin server")]
     Server,
@@ -62,71 +63,16 @@ impl Atuin {
             }
         };
 
-        let db = SqliteDatabase::new(db_path)?;
+        let mut db = SqliteDatabase::new(db_path)?;
 
         match self.atuin {
             AtuinCmd::History(history) => history.run(db),
+            AtuinCmd::Import(import) => import.run(&mut db),
             _ => Ok(()),
         }
     }
 }
 
-#[derive(StructOpt)]
-enum HistoryCmd {
-    #[structopt(
-        about="begins a new command in the history",
-        aliases=&["s", "st", "sta", "star"],
-    )]
-    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",
-        aliases=&["l", "li", "lis"],
-    )]
-    List,
-}
-
-impl HistoryCmd {
-    fn run(&self, db: SqliteDatabase) -> Result<()> {
-        match self {
-            HistoryCmd::Start { command: words } => {
-                let command = words.join(" ");
-                let cwd = env::current_dir()?.display().to_string();
-
-                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)?;
-                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(())
-            }
-
-            HistoryCmd::List => db.list(),
-        }
-    }
-}
-
 fn main() -> Result<()> {
     pretty_env_logger::init();