Importer V3 (#395)

* start of importer refactor

* fish

* resh

* zsh

8 files changed, 61 insertions, 132 deletions

diff --git a/src/command/client.rs b/src/command/client.rs
index c75872aa..b9d43b39 100644
--- a/src/command/client.rs
+++ b/src/command/client.rs
@@ -58,6 +58,7 @@ pub enum Cmd {
 }
 
 impl Cmd {
+    #[tokio::main(flavor = "current_thread")]
     pub async fn run(self) -> Result<()> {
         pretty_env_logger::init();
 
diff --git a/src/command/client/history.rs b/src/command/client/history.rs
index d001658c..805fe4ca 100644
--- a/src/command/client/history.rs
+++ b/src/command/client/history.rs
@@ -128,11 +128,7 @@ pub fn print_cmd_only(w: &mut StdoutLock, h: &[History]) {
 }
 
 impl Cmd {
-    pub async fn run(
-        &self,
-        settings: &Settings,
-        db: &mut (impl Database + Send + Sync),
-    ) -> Result<()> {
+    pub async fn run(&self, settings: &Settings, db: &mut impl Database) -> Result<()> {
         let context = current_context();
 
         match self {
diff --git a/src/command/client/import.rs b/src/command/client/import.rs
index c70446d5..580e4b0e 100644
--- a/src/command/client/import.rs
+++ b/src/command/client/import.rs
@@ -1,13 +1,14 @@
-use std::{env, path::PathBuf};
+use std::env;
 
+use async_trait::async_trait;
 use clap::Parser;
-use eyre::{eyre, Result};
+use eyre::Result;
 use indicatif::ProgressBar;
 
 use atuin_client::{
     database::Database,
     history::History,
-    import::{bash::Bash, fish::Fish, resh::Resh, zsh::Zsh, Importer},
+    import::{bash::Bash, fish::Fish, resh::Resh, zsh::Zsh, Importer, Loader},
 };
 
 #[derive(Parser)]
@@ -18,13 +19,10 @@ pub enum Cmd {
 
     /// Import history from the zsh history file
     Zsh,
-
     /// Import history from the bash history file
     Bash,
-
     /// Import history from the resh history file
     Resh,
-
     /// Import history from the fish history file
     Fish,
 }
@@ -32,7 +30,7 @@ pub enum Cmd {
 const BATCH_SIZE: usize = 100;
 
 impl Cmd {
-    pub async fn run(&self, db: &mut (impl Database + Send + Sync)) -> Result<()> {
+    pub async fn run<DB: Database>(&self, db: &mut DB) -> Result<()> {
         println!("        Atuin         ");
         println!("======================");
         println!("          \u{1f30d}          ");
@@ -47,124 +45,73 @@ impl Cmd {
 
                 if shell.ends_with("/zsh") {
                     println!("Detected ZSH");
-                    import::<Zsh<_>, _>(db, BATCH_SIZE).await
+                    import::<Zsh, DB>(db).await
                 } else if shell.ends_with("/fish") {
                     println!("Detected Fish");
-                    import::<Fish<_>, _>(db, BATCH_SIZE).await
+                    import::<Fish, DB>(db).await
                 } else if shell.ends_with("/bash") {
                     println!("Detected Bash");
-                    import::<Bash<_>, _>(db, BATCH_SIZE).await
+                    import::<Bash, DB>(db).await
                 } else {
                     println!("cannot import {} history", shell);
                     Ok(())
                 }
             }
 
-            Self::Zsh => import::<Zsh<_>, _>(db, BATCH_SIZE).await,
-            Self::Bash => import::<Bash<_>, _>(db, BATCH_SIZE).await,
-            Self::Resh => import::<Resh, _>(db, BATCH_SIZE).await,
-            Self::Fish => import::<Fish<_>, _>(db, BATCH_SIZE).await,
+            Self::Zsh => import::<Zsh, DB>(db).await,
+            Self::Bash => import::<Bash, DB>(db).await,
+            Self::Resh => import::<Resh, DB>(db).await,
+            Self::Fish => import::<Fish, DB>(db).await,
         }
     }
 }
 
-async fn import<I: Importer + Send, DB: Database + Send + Sync>(
-    db: &mut DB,
-    buf_size: usize,
-) -> Result<()>
-where
-    I::IntoIter: Send,
-{
-    println!("Importing history from {}", I::NAME);
-
-    let histpath = get_histpath::<I>()?;
-    let contents = I::parse(histpath)?;
-
-    let iter = contents.into_iter();
-    let progress = if let (_, Some(upper_bound)) = iter.size_hint() {
-        ProgressBar::new(upper_bound as u64)
-    } else {
-        ProgressBar::new_spinner()
-    };
-
-    let mut buf = Vec::<History>::with_capacity(buf_size);
-    let mut iter = progress.wrap_iter(iter);
-    loop {
-        // fill until either no more entries
-        // or until the buffer is full
-        let done = fill_buf(&mut buf, &mut iter);
-
-        // flush
-        db.save_bulk(&buf).await?;
-
-        if done {
-            break;
-        }
-    }
-
-    println!("Import complete!");
-
-    Ok(())
+pub struct HistoryImporter<'db, DB: Database> {
+    pb: ProgressBar,
+    buf: Vec<History>,
+    db: &'db mut DB,
 }
 
-fn get_histpath<I: Importer>() -> Result<PathBuf> {
-    if let Ok(p) = env::var("HISTFILE") {
-        is_file(PathBuf::from(p))
-    } else {
-        is_file(I::histpath()?)
+impl<'db, DB: Database> HistoryImporter<'db, DB> {
+    fn new(db: &'db mut DB, len: usize) -> Self {
+        Self {
+            pb: ProgressBar::new(len as u64),
+            buf: Vec::with_capacity(BATCH_SIZE),
+            db,
+        }
     }
-}
 
-fn is_file(p: PathBuf) -> Result<PathBuf> {
-    if p.is_file() {
-        Ok(p)
-    } else {
-        Err(eyre!(
-            "Could not find history file {:?}. Try setting $HISTFILE",
-            p
-        ))
+    async fn flush(self) -> Result<()> {
+        if !self.buf.is_empty() {
+            self.db.save_bulk(&self.buf).await?;
+        }
+        self.pb.finish();
+        Ok(())
     }
 }
 
-fn fill_buf<T, E>(buf: &mut Vec<T>, iter: &mut impl Iterator<Item = Result<T, E>>) -> bool {
-    buf.clear();
-    loop {
-        match iter.next() {
-            Some(Ok(t)) => buf.push(t),
-            Some(Err(_)) => (),
-            None => break true,
-        }
-
-        if buf.len() == buf.capacity() {
-            break false;
+#[async_trait]
+impl<'db, DB: Database> Loader for HistoryImporter<'db, DB> {
+    async fn push(&mut self, hist: History) -> Result<()> {
+        self.pb.inc(1);
+        self.buf.push(hist);
+        if self.buf.len() == self.buf.capacity() {
+            self.db.save_bulk(&self.buf).await?;
+            self.buf.clear();
         }
+        Ok(())
     }
 }
 
-#[cfg(test)]
-mod tests {
-    use super::fill_buf;
-
-    #[test]
-    fn test_fill_buf() {
-        let mut buf = Vec::with_capacity(4);
-        let mut iter = vec![
-            Ok(1),
-            Err(2),
-            Ok(3),
-            Ok(4),
-            Err(5),
-            Ok(6),
-            Ok(7),
-            Err(8),
-            Ok(9),
-        ]
-        .into_iter();
+async fn import<I: Importer + Send, DB: Database>(db: &mut DB) -> Result<()> {
+    println!("Importing history from {}", I::NAME);
 
-        assert!(!fill_buf(&mut buf, &mut iter));
-        assert_eq!(buf, vec![1, 3, 4, 6]);
+    let mut importer = I::new().await?;
+    let len = importer.entries().await.unwrap();
+    let mut loader = HistoryImporter::new(db, len);
+    importer.load(&mut loader).await?;
+    loader.flush().await?;
 
-        assert!(fill_buf(&mut buf, &mut iter));
-        assert_eq!(buf, vec![7, 9]);
-    }
+    println!("Import complete!");
+    Ok(())
 }
diff --git a/src/command/client/search.rs b/src/command/client/search.rs
index 8c60bd36..de6e7961 100644
--- a/src/command/client/search.rs
+++ b/src/command/client/search.rs
@@ -75,11 +75,7 @@ pub struct Cmd {
 }
 
 impl Cmd {
-    pub async fn run(
-        self,
-        db: &mut (impl Database + Send + Sync),
-        settings: &Settings,
-    ) -> Result<()> {
+    pub async fn run(self, db: &mut impl Database, settings: &Settings) -> Result<()> {
         if self.interactive {
             let item = select_history(
                 &self.query,
@@ -257,7 +253,7 @@ impl State {
 async fn query_results(
     app: &mut State,
     search_mode: SearchMode,
-    db: &mut (impl Database + Send + Sync),
+    db: &mut impl Database,
 ) -> Result<()> {
     let results = match app.input.as_str() {
         "" => {
@@ -284,7 +280,7 @@ async fn query_results(
 async fn key_handler(
     input: Key,
     search_mode: SearchMode,
-    db: &mut (impl Database + Send + Sync),
+    db: &mut impl Database,
     app: &mut State,
 ) -> Option<String> {
     match input {
@@ -537,7 +533,7 @@ async fn select_history(
     search_mode: SearchMode,
     filter_mode: FilterMode,
     style: atuin_client::settings::Style,
-    db: &mut (impl Database + Send + Sync),
+    db: &mut impl Database,
 ) -> Result<String> {
     let stdout = stdout().into_raw_mode()?;
     let stdout = MouseTerminal::from(stdout);
@@ -596,7 +592,7 @@ async fn run_non_interactive(
     after: Option<String>,
     limit: Option<i64>,
     query: &[String],
-    db: &mut (impl Database + Send + Sync),
+    db: &mut impl Database,
 ) -> Result<()> {
     let dir = if cwd.as_deref() == Some(".") {
         let current = std::env::current_dir()?;
diff --git a/src/command/client/stats.rs b/src/command/client/stats.rs
index 80450989..157496de 100644
--- a/src/command/client/stats.rs
+++ b/src/command/client/stats.rs
@@ -62,11 +62,7 @@ fn compute_stats(history: &[History]) -> Result<()> {
 }
 
 impl Cmd {
-    pub async fn run(
-        &self,
-        db: &mut (impl Database + Send + Sync),
-        settings: &Settings,
-    ) -> Result<()> {
+    pub async fn run(&self, db: &mut impl Database, settings: &Settings) -> Result<()> {
         let context = current_context();
         let words = if self.period.is_empty() {
             String::from("all")
diff --git a/src/command/client/sync.rs b/src/command/client/sync.rs
index 6fbf8dfc..af809f3e 100644
--- a/src/command/client/sync.rs
+++ b/src/command/client/sync.rs
@@ -31,11 +31,7 @@ pub enum Cmd {
 }
 
 impl Cmd {
-    pub async fn run(
-        self,
-        settings: Settings,
-        db: &mut (impl Database + Send + Sync),
-    ) -> Result<()> {
+    pub async fn run(self, settings: Settings, db: &mut impl Database) -> Result<()> {
         match self {
             Self::Sync { force } => run(&settings, force, db).await,
             Self::Login(l) => l.run(&settings).await,
@@ -52,11 +48,7 @@ impl Cmd {
     }
 }
 
-async fn run(
-    settings: &Settings,
-    force: bool,
-    db: &mut (impl Database + Send + Sync),
-) -> Result<()> {
+async fn run(settings: &Settings, force: bool, db: &mut impl Database) -> Result<()> {
     atuin_client::sync::sync(settings, force, db).await?;
     println!(
         "Sync complete! {} items in database, force: {}",
diff --git a/src/command/mod.rs b/src/command/mod.rs
index 953b76ba..c86e76f5 100644
--- a/src/command/mod.rs
+++ b/src/command/mod.rs
@@ -19,11 +19,11 @@ pub enum AtuinCmd {
 }
 
 impl AtuinCmd {
-    pub async fn run(self) -> Result<()> {
+    pub fn run(self) -> Result<()> {
         match self {
-            Self::Client(client) => client.run().await,
+            Self::Client(client) => client.run(),
             #[cfg(feature = "server")]
-            Self::Server(server) => server.run().await,
+            Self::Server(server) => server.run(),
         }
     }
 }
diff --git a/src/command/server.rs b/src/command/server.rs
index 1d514bb2..495f85d0 100644
--- a/src/command/server.rs
+++ b/src/command/server.rs
@@ -21,6 +21,7 @@ pub enum Cmd {
 }
 
 impl Cmd {
+    #[tokio::main]
     pub async fn run(self) -> Result<()> {
         tracing_subscriber::registry()
             .with(fmt::layer())