tidy some stuff (#6)

6 files changed, 119 insertions, 143 deletions

diff --git a/src/command/history.rs b/src/command/history.rs
index 5959fc55..e40af4d6 100644
--- a/src/command/history.rs
+++ b/src/command/history.rs
@@ -34,7 +34,7 @@ pub enum Cmd {
     },
 }
 
-fn print_list(h: &Vec<History>) {
+fn print_list(h: &[History]) {
     for i in h {
         println!("{}", i.command);
     }
@@ -60,7 +60,7 @@ impl Cmd {
                 // print the ID
                 // we use this as the key for calling end
                 println!("{}", h.id);
-                db.save(h)?;
+                db.save(&h)?;
                 Ok(())
             }
 
@@ -69,7 +69,7 @@ impl Cmd {
                 h.exit = *exit;
                 h.duration = chrono::Utc::now().timestamp_nanos() - h.timestamp;
 
-                db.update(h)?;
+                db.update(&h)?;
 
                 Ok(())
             }
diff --git a/src/command/mod.rs b/src/command/mod.rs
index 5c36146a..8d463bd6 100644
--- a/src/command/mod.rs
+++ b/src/command/mod.rs
@@ -1,3 +1,23 @@
-pub mod history;
-pub mod import;
-pub mod server;
+use structopt::StructOpt;
+
+mod history;
+mod import;
+mod server;
+
+#[derive(StructOpt)]
+pub enum AtuinCmd {
+    #[structopt(
+        about="manipulate shell history",
+        aliases=&["h", "hi", "his", "hist", "histo", "histor"],
+    )]
+    History(history::Cmd),
+
+    #[structopt(about = "import shell history from file")]
+    Import(import::Cmd),
+
+    #[structopt(about = "start an atuin server")]
+    Server(server::Cmd),
+
+    #[structopt(about = "generates a UUID")]
+    Uuid,
+}
diff --git a/src/local/database.rs b/src/local/database.rs
index 43b17e50..e2df9ba5 100644
--- a/src/local/database.rs
+++ b/src/local/database.rs
@@ -3,18 +3,18 @@ use std::path::Path;
 
 use eyre::{eyre, Result};
 
-use rusqlite::NO_PARAMS;
 use rusqlite::{params, Connection};
+use rusqlite::{Transaction, NO_PARAMS};
 
-use crate::History;
+use super::history::History;
 
 pub trait Database {
-    fn save(&mut self, h: History) -> Result<()>;
+    fn save(&mut self, h: &History) -> Result<()>;
     fn save_bulk(&mut self, h: &[History]) -> Result<()>;
     fn load(&self, id: &str) -> Result<History>;
     fn list(&self) -> Result<Vec<History>>;
     fn since(&self, date: chrono::DateTime<Utc>) -> Result<Vec<History>>;
-    fn update(&self, h: History) -> Result<()>;
+    fn update(&self, h: &History) -> Result<()>;
 }
 
 // Intended for use on a developer machine and not a sync server.
@@ -65,46 +65,53 @@ impl Sqlite {
 
         Ok(())
     }
+
+    fn save_raw(tx: &Transaction, h: &History) -> Result<()> {
+        tx.execute(
+            "insert or ignore into history (
+            id,
+            timestamp,
+            duration,
+            exit,
+            command,
+            cwd,
+            session,
+            hostname
+        ) values (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
+            params![
+                h.id,
+                h.timestamp,
+                h.duration,
+                h.exit,
+                h.command,
+                h.cwd,
+                h.session,
+                h.hostname
+            ],
+        )?;
+
+        Ok(())
+    }
 }
 
 impl Database for Sqlite {
-    fn save(&mut self, h: History) -> Result<()> {
+    fn save(&mut self, h: &History) -> Result<()> {
         debug!("saving history to sqlite");
-        let v = vec![h];
 
-        self.save_bulk(&v)
+        let tx = self.conn.transaction()?;
+        Self::save_raw(&tx, h)?;
+        tx.commit()?;
+
+        Ok(())
     }
 
     fn save_bulk(&mut self, h: &[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,
-                session,
-                hostname
-            ) values (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
-                params![
-                    i.id,
-                    i.timestamp,
-                    i.duration,
-                    i.exit,
-                    i.command,
-                    i.cwd,
-                    i.session,
-                    i.hostname
-                ],
-            )?;
+            Self::save_raw(&tx, i)?
         }
-
         tx.commit()?;
 
         Ok(())
@@ -119,16 +126,7 @@ impl Database for Sqlite {
         )?;
 
         let mut 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)?,
-                session: row.get(6)?,
-                hostname: row.get(7)?,
-            })
+            history_from_sqlite_row(Some(id.to_string()), row)
         })?;
 
         let history = iter.next().unwrap();
@@ -139,7 +137,7 @@ impl Database for Sqlite {
         }
     }
 
-    fn update(&self, h: History) -> Result<()> {
+    fn update(&self, h: &History) -> Result<()> {
         debug!("updating sqlite history");
 
         self.conn.execute(
@@ -159,42 +157,43 @@ impl Database for Sqlite {
             .conn
             .prepare("SELECT * FROM history order by timestamp asc")?;
 
-        let history_iter = stmt.query_map(params![], |row| {
-            Ok(History {
-                id: row.get(0)?,
-                timestamp: row.get(1)?,
-                duration: row.get(2)?,
-                exit: row.get(3)?,
-                command: row.get(4)?,
-                cwd: row.get(5)?,
-                session: row.get(6)?,
-                hostname: row.get(7)?,
-            })
-        })?;
+        let history_iter = stmt.query_map(params![], |row| history_from_sqlite_row(None, row))?;
 
-        Ok(history_iter.filter_map(|x| x.ok()).collect())
+        Ok(history_iter.filter_map(Result::ok).collect())
     }
 
     fn since(&self, date: chrono::DateTime<Utc>) -> Result<Vec<History>> {
-        debug!("listing history");
+        debug!("listing history since {:?}", date);
 
         let mut stmt = self.conn.prepare(
             "SELECT distinct command FROM history where timestamp > ?1 order by timestamp asc",
         )?;
 
         let history_iter = stmt.query_map(params![date.timestamp_nanos()], |row| {
-            Ok(History {
-                id: row.get(0)?,
-                timestamp: row.get(1)?,
-                duration: row.get(2)?,
-                exit: row.get(3)?,
-                command: row.get(4)?,
-                cwd: row.get(5)?,
-                session: row.get(6)?,
-                hostname: row.get(7)?,
-            })
+            history_from_sqlite_row(None, row)
         })?;
 
-        Ok(history_iter.filter_map(|x| x.ok()).collect())
+        Ok(history_iter.filter_map(Result::ok).collect())
     }
 }
+
+fn history_from_sqlite_row(
+    id: Option<String>,
+    row: &rusqlite::Row,
+) -> Result<History, rusqlite::Error> {
+    let id = match id {
+        Some(id) => id,
+        None => row.get(0)?,
+    };
+
+    Ok(History {
+        id,
+        timestamp: row.get(1)?,
+        duration: row.get(2)?,
+        exit: row.get(3)?,
+        command: row.get(4)?,
+        cwd: row.get(5)?,
+        session: row.get(6)?,
+        hostname: row.get(7)?,
+    })
+}
diff --git a/src/local/history.rs b/src/local/history.rs
index e46180b8..06a350fc 100644
--- a/src/local/history.rs
+++ b/src/local/history.rs
@@ -24,17 +24,11 @@ impl History {
         session: Option<String>,
         hostname: Option<String>,
     ) -> Self {
-        // get the current session or just generate a random string
-        let env_session =
-            env::var("ATUIN_SESSION").unwrap_or_else(|_| Uuid::new_v4().to_simple().to_string());
-
-        // best attempt at getting the current hostname, or just unknown
-        let os_hostname = hostname::get().unwrap();
-        let os_hostname = os_hostname.to_str().unwrap();
-        let os_hostname = String::from(os_hostname);
-
-        let session = session.unwrap_or(env_session);
-        let hostname = hostname.unwrap_or(os_hostname);
+        let session = session.unwrap_or_else(|| {
+            env::var("ATUIN_SESSION").unwrap_or_else(|_| Uuid::new_v4().to_simple().to_string())
+        });
+        let hostname =
+            hostname.unwrap_or_else(|| hostname::get().unwrap().to_str().unwrap().to_string());
 
         Self {
             id: Uuid::new_v4().to_simple().to_string(),
diff --git a/src/local/import.rs b/src/local/import.rs
index 046c74bf..858e5786 100644
--- a/src/local/import.rs
+++ b/src/local/import.rs
@@ -1,8 +1,8 @@
 // import old shell history!
 // automatically hoover up all that we can find
 
-use std::fs::File;
-use std::io::{BufRead, BufReader};
+use std::io::{BufRead, BufReader, Seek, SeekFrom};
+use std::{fs::File, path::Path};
 
 use eyre::{eyre, Result};
 
@@ -16,19 +16,18 @@ pub struct Zsh {
 }
 
 // this could probably be sped up
-fn count_lines(path: &str) -> Result<usize> {
-    let file = File::open(path)?;
-    let buf = BufReader::new(file);
+fn count_lines(buf: &mut BufReader<File>) -> Result<usize> {
+    let lines = buf.lines().count();
+    buf.seek(SeekFrom::Start(0))?;
 
-    Ok(buf.lines().count())
+    Ok(lines)
 }
 
 impl Zsh {
-    pub fn new(path: &str) -> Result<Self> {
-        let loc = count_lines(path)?;
-
+    pub fn new(path: impl AsRef<Path>) -> Result<Self> {
         let file = File::open(path)?;
-        let buf = BufReader::new(file);
+        let mut buf = BufReader::new(file);
+        let loc = count_lines(&mut buf)?;
 
         Ok(Self {
             file: buf,
@@ -37,43 +36,25 @@ impl Zsh {
     }
 }
 
-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 parse_extended(line: &str) -> 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_else(|_| chrono::Utc::now().timestamp_nanos());
-
-    let duration_command = split.next().unwrap(); // might be 0;the command
-    let mut split = duration_command.split(';');
+    let (time, duration) = line.split_once(':').unwrap();
+    let (duration, command) = duration.split_once(';').unwrap();
 
-    let duration = split.next().unwrap_or("-1"); // should just be the 0
-    let duration = duration.parse::<i64>().unwrap_or(-1);
+    let time = time.parse::<i64>().map_or_else(
+        |_| chrono::Utc::now().timestamp_nanos(),
+        |t| t * 1_000_000_000,
+    );
 
-    let command = split.next().unwrap();
+    let duration = duration.parse::<i64>().map_or(-1, |t| t * 1_000_000_000);
 
     // use nanos, because why the hell not? we won't display them.
     History::new(
-        time * 1_000_000_000,
-        trim_newline(command),
+        time,
+        command.trim_end().to_string(),
         String::from("unknown"),
         -1,
-        duration * 1_000_000_000,
+        duration,
         None,
         None,
     )
@@ -101,7 +82,7 @@ impl Iterator for Zsh {
                 } else {
                     Some(Ok(History::new(
                         chrono::Utc::now().timestamp_nanos(), // what else? :/
-                        trim_newline(line.as_str()),
+                        line.trim_end().to_string(),
                         String::from("unknown"),
                         -1,
                         -1,
diff --git a/src/main.rs b/src/main.rs
index 022703b9..2dbeabfe 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,4 @@
+#![feature(str_split_once)]
 #![feature(proc_macro_hygiene)]
 #![feature(decl_macro)]
 #![warn(clippy::pedantic, clippy::nursery)]
@@ -15,9 +16,8 @@ extern crate log;
 #[macro_use]
 extern crate rocket;
 
-use command::{history, import, server};
+use command::AtuinCmd;
 use local::database::Sqlite;
-use local::history::History;
 
 mod command;
 mod local;
@@ -37,24 +37,6 @@ struct Atuin {
     atuin: AtuinCmd,
 }
 
-#[derive(StructOpt)]
-enum AtuinCmd {
-    #[structopt(
-        about="manipulate shell history",
-        aliases=&["h", "hi", "his", "hist", "histo", "histor"],
-    )]
-    History(history::Cmd),
-
-    #[structopt(about = "import shell history from file")]
-    Import(import::Cmd),
-
-    #[structopt(about = "start an atuin server")]
-    Server(server::Cmd),
-
-    #[structopt(about = "generates a UUID")]
-    Uuid,
-}
-
 impl Atuin {
     fn run(self) -> Result<()> {
         let db_path = if let Some(db_path) = self.db {