From 099afe66ecfb569a8a04b66425ded29665e6a37c Mon Sep 17 00:00:00 2001 From: Ellie Huxtable Date: Sat, 13 Feb 2021 19:37:00 +0000 Subject: [PATCH] Implement history import --- Cargo.lock | 62 +++++++++++++++++++++++- Cargo.toml | 4 +- src/command/history.rs | 63 ++++++++++++++++++++++++ src/command/import.rs | 107 +++++++++++++++++++++++++++++++++++++++++ src/command/mod.rs | 2 + src/local/database.rs | 31 +++++++++++- src/local/history.rs | 4 +- src/local/import.rs | 101 ++++++++++++++++++++++++++++++++------ src/local/mod.rs | 1 + src/main.rs | 68 +++----------------------- 10 files changed, 361 insertions(+), 82 deletions(-) create mode 100644 src/command/history.rs create mode 100644 src/command/import.rs create mode 100644 src/command/mod.rs diff --git a/Cargo.lock b/Cargo.lock index feb0404..fcf1175 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", @@ -131,6 +133,21 @@ dependencies = [ "vec_map", ] +[[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" @@ -189,6 +206,12 @@ dependencies = [ "winapi", ] +[[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" @@ -282,6 +305,15 @@ dependencies = [ "libc", ] +[[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" @@ -297,6 +329,18 @@ version = "0.3.2" 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" @@ -354,6 +398,12 @@ dependencies = [ "autocfg", ] +[[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" @@ -576,6 +626,16 @@ dependencies = [ "winapi-util", ] +[[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" diff --git a/Cargo.toml b/Cargo.toml index 3172a30..57117c2 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 "] 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 0000000..72f821c --- /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 }, + + #[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 0000000..5ece13a --- /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::::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::::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 0000000..c61d228 --- /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 b2c009b..2a4cc58 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) -> Result<()>; fn load(&self, id: &str) -> Result; 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) -> 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 { debug!("loading history item"); diff --git a/src/local/history.rs b/src/local/history.rs index 3c9a906..0010962 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 8db8f0e..ce141c5 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, + + pub loc: u64, } -impl ImportBash { - pub fn new(path: &str) -> Result { +// this could probably be sped up +fn count_lines(path: &str) -> Result { + let file = File::open(path)?; + let buf = BufReader::new(file); + + Ok(buf.lines().count()) +} + +impl ImportZsh { + pub fn new(path: &str) -> Result { + 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); - fn next(&mut self) -> Option { + if s.ends_with('\n') { + s.pop(); + if s.ends_with('\r') { + s.pop(); + } + } + + s +} + +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::() + .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::().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; + + fn next(&mut self) -> Option { + // 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(_) => Some(History { - cwd: "none".to_string(), - command: line, - timestamp: -1, - }), + Ok(_) => { + let extended = line.starts_with(":"); + + 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 f587d01..a11ee21 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 57688a4..19357cb 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 ", @@ -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 }, - - #[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();