a lil refactoring

Cargo.lock

@@ -151,6 +151,7 @@ dependencies = [
  "sodiumoxide",
  "sql-builder",
  "sqlx",
+ "tantivy",
  "tokio",
  "urlencoding",
  "uuid",

atuin-client/Cargo.toml

@@ -52,6 +52,7 @@ fs-err = "2.9"
 sql-builder = "3"
 lazy_static = "1"
 memchr = "2.5"
+tantivy = "0.19"
 
 # sync
 urlencoding = { version = "2.1.0", optional = true }

atuin-client/src/lib.rs

@@ -15,3 +15,4 @@ pub mod history;
 pub mod import;
 pub mod ordering;
 pub mod settings;
+pub mod tantivy;

atuin-client/src/sync.rs

@@ -58,6 +58,10 @@ async fn sync_download(
 
     let host = if force { Some(String::from("")) } else { None };
 
+    let (hs, schema) = crate::tantivy::schema();
+    let index = crate::tantivy::index(schema)?;
+    let mut writer = index.writer(3_000_000)?;
+
     while remote_count > local_count {
         let page = client
             .get_history(
@@ -90,6 +94,8 @@ async fn sync_download(
         } else {
             last_timestamp = page_last;
         }
+
+        crate::tantivy::bulk_write_history(&mut writer, &hs, page)?;
     }
 
     for i in remote_status.deleted {

atuin-client/src/tantivy.rs

@@ -0,0 +1,103 @@
+use crate::{database::Database, history::History};
+use eyre::Result;
+use tantivy::{
+    directory::MmapDirectory,
+    doc,
+    schema::{Field, Schema, FAST, STORED, STRING, TEXT},
+    DateTime, Index, IndexWriter,
+};
+
+pub fn schema() -> (HistorySchema, Schema) {
+    let mut schema_builder = Schema::builder();
+
+    (
+        HistorySchema {
+            id: schema_builder.add_text_field("id", STRING),
+            command: schema_builder.add_text_field("command", TEXT | STORED),
+            cwd: schema_builder.add_text_field("cwd", STRING | FAST),
+            session: schema_builder.add_text_field("session", STRING | FAST),
+            hostname: schema_builder.add_text_field("hostname", STRING | FAST),
+            timestamp: schema_builder.add_date_field("timestamp", STORED),
+            duration: schema_builder.add_i64_field("duration", STORED),
+            exit: schema_builder.add_i64_field("exit", STORED),
+        },
+        schema_builder.build(),
+    )
+}
+
+pub struct HistorySchema {
+    pub id: Field,
+    pub command: Field,
+    pub cwd: Field,
+    pub session: Field,
+    pub hostname: Field,
+    pub timestamp: Field,
+    pub duration: Field,
+    pub exit: Field,
+}
+
+pub fn index(schema: Schema) -> Result<Index> {
+    let data_dir = atuin_common::utils::data_dir();
+    let tantivy_dir = data_dir.join("tantivy");
+
+    fs_err::create_dir_all(&tantivy_dir)?;
+    let dir = MmapDirectory::open(tantivy_dir)?;
+
+    Ok(Index::open_or_create(dir, schema)?)
+}
+
+pub fn write_history(h: impl IntoIterator<Item = History>) -> Result<()> {
+    let (hs, schema) = schema();
+    let index = index(schema)?;
+    let mut writer = index.writer(3_000_000)?;
+
+    bulk_write_history(&mut writer, &hs, h)?;
+
+    Ok(())
+}
+
+pub fn bulk_write_history(
+    writer: &mut IndexWriter,
+    schema: &HistorySchema,
+    h: impl IntoIterator<Item = History>,
+) -> Result<()> {
+    for h in h {
+        write_single_history(writer, schema, h)?;
+    }
+    writer.commit()?;
+
+    Ok(())
+}
+
+fn write_single_history(
+    writer: &mut IndexWriter,
+    schema: &HistorySchema,
+    h: History,
+) -> Result<()> {
+    let timestamp = DateTime::from_timestamp_millis(h.timestamp.timestamp_millis());
+    writer.add_document(doc!(
+        schema.id => h.id,
+        schema.command => h.command,
+        schema.cwd => h.cwd,
+        schema.session => h.session,
+        schema.hostname => h.hostname,
+        schema.timestamp => timestamp,
+        schema.duration => h.duration,
+        schema.exit => h.exit,
+    ))?;
+
+    Ok(())
+}
+
+pub async fn refresh(db: &mut impl Database) -> Result<()> {
+    let history = db.all_with_count().await?;
+
+    // delete the index
+    let data_dir = atuin_common::utils::data_dir();
+    let tantivy_dir = dbg!(data_dir.join("tantivy"));
+    fs_err::remove_dir_all(tantivy_dir)?;
+
+    tokio::task::spawn_blocking(|| write_history(history.into_iter().map(|(h, _)| h))).await??;
+
+    Ok(())
+}

src/command/client.rs

@@ -32,7 +32,7 @@ pub enum Cmd {
     Search(search::Cmd),
 
     /// Interactive history search
-    RefreshTantivyIndex(search::tantivy_impl::Cmd),
+    RefreshTantivyIndex,
 
     #[cfg(feature = "sync")]
     #[command(flatten)]
@@ -56,8 +56,8 @@ impl Cmd {
             Self::History(history) => history.run(&settings, &mut db).await,
             Self::Import(import) => import.run(&mut db).await,
             Self::Stats(stats) => stats.run(&mut db, &settings).await,
-            Self::Search(search) => search.run(&mut db, &mut settings).await,
-            Self::RefreshTantivyIndex(refresh) => refresh.run(&mut db).await,
+            Self::Search(search) => search.run(db, &mut settings).await,
+            Self::RefreshTantivyIndex => atuin_client::tantivy::refresh(&mut db).await,
             #[cfg(feature = "sync")]
             Self::Sync(sync) => sync.run(settings, &mut db).await,
         }

src/command/client/history.rs

@@ -211,6 +211,7 @@ impl Cmd {
                 h.duration = chrono::Utc::now().timestamp_nanos() - h.timestamp.timestamp_nanos();
 
                 db.update(&h).await?;
+                atuin_client::tantivy::write_history([h])?;
 
                 if settings.should_sync()? {
                     #[cfg(feature = "sync")]

src/command/client/search.rs

@@ -14,11 +14,9 @@ use super::history::ListMode;
 
 mod cursor;
 mod duration;
+mod engines;
 mod history_list;
 mod interactive;
-mod skim_impl;
-mod db_impl;
-pub mod tantivy_impl;
 pub use duration::{format_duration, format_duration_into};
 
 #[allow(clippy::struct_excessive_bools)]
@@ -89,7 +87,7 @@ pub struct Cmd {
 }
 
 impl Cmd {
-    pub async fn run(self, db: &mut impl Database, settings: &mut Settings) -> Result<()> {
+    pub async fn run(self, mut db: impl Database, settings: &mut Settings) -> Result<()> {
         if self.search_mode.is_some() {
             settings.search_mode = self.search_mode.unwrap();
         }
@@ -114,7 +112,7 @@ impl Cmd {
                 self.after,
                 self.limit,
                 &self.query,
-                db,
+                &mut db,
             )
             .await?;
 

src/command/client/search/engines.rs

@@ -0,0 +1,68 @@
+use std::{ops::Deref, sync::Arc};
+
+use async_trait::async_trait;
+use atuin_client::{
+    database::{Context, Database},
+    history::History,
+    settings::{FilterMode, SearchMode},
+};
+use eyre::Result;
+
+use super::cursor::Cursor;
+
+pub mod db;
+pub mod skim;
+pub mod tantivy;
+
+pub fn engine(search_mode: SearchMode) -> Result<Box<dyn SearchEngine>> {
+    Ok(match search_mode {
+        SearchMode::Skim => Box::new(skim::Search::new()) as Box<_>,
+        SearchMode::Tantivy => Box::new(tantivy::Search::new()?) as Box<_>,
+        mode => Box::new(db::Search(mode)) as Box<_>,
+    })
+}
+
+pub struct SearchState {
+    pub input: Cursor,
+    pub filter_mode: FilterMode,
+    pub context: Context,
+}
+
+#[async_trait]
+pub trait SearchEngine: Send + Sync + 'static {
+    async fn full_query(
+        &mut self,
+        state: &SearchState,
+        db: &mut dyn Database,
+    ) -> Result<Vec<Arc<HistoryWrapper>>>;
+
+    async fn query(
+        &mut self,
+        state: &SearchState,
+        db: &mut dyn Database,
+    ) -> Result<Vec<Arc<HistoryWrapper>>> {
+        if state.input.as_str().is_empty() {
+            Ok(db
+                .list(state.filter_mode, &state.context, Some(200), true)
+                .await?
+                .into_iter()
+                .map(|history| HistoryWrapper { history, count: 1 })
+                .map(Arc::new)
+                .collect::<Vec<_>>())
+        } else {
+            self.full_query(state, db).await
+        }
+    }
+}
+
+pub struct HistoryWrapper {
+    pub history: History,
+    pub count: i32,
+}
+impl Deref for HistoryWrapper {
+    type Target = History;
+
+    fn deref(&self) -> &Self::Target {
+        &self.history
+    }
+}

src/command/client/search/engines/db.rs

@@ -4,13 +4,13 @@ use async_trait::async_trait;
 use atuin_client::{database::Database, settings::SearchMode};
 use eyre::Result;
 
-use super::interactive::{HistoryWrapper, SearchEngine, SearchState};
+use super::{HistoryWrapper, SearchEngine, SearchState};
 
 pub struct Search(pub SearchMode);
 
 #[async_trait]
 impl SearchEngine for Search {
-    async fn query(
+    async fn full_query(
         &mut self,
         state: &SearchState,
         db: &mut dyn Database,

src/command/client/search/engines/skim.rs

@@ -4,25 +4,26 @@ use async_trait::async_trait;
 use atuin_client::{database::Database, settings::FilterMode};
 use chrono::Utc;
 use eyre::Result;
-use skim::{prelude::ExactOrFuzzyEngineFactory, MatchEngineFactory};
+use skim::{prelude::ExactOrFuzzyEngineFactory, MatchEngineFactory, SkimItem};
 use tokio::task::yield_now;
 
-use super::interactive::{HistoryWrapper, SearchEngine, SearchState};
+use super::{HistoryWrapper, SearchEngine, SearchState};
 
-#[derive(Default)]
 pub struct Search {
     all_history: Vec<Arc<HistoryWrapper>>,
 }
 
 impl Search {
     pub fn new() -> Self {
-        Search::default()
+        Search {
+            all_history: vec![],
+        }
     }
 }
 
 #[async_trait]
 impl SearchEngine for Search {
-    async fn query(
+    async fn full_query(
         &mut self,
         state: &SearchState,
         db: &mut dyn Database,
@@ -125,3 +126,9 @@ pub async fn fuzzy_search(
 
     set
 }
+
+impl SkimItem for HistoryWrapper {
+    fn text(&self) -> std::borrow::Cow<str> {
+        std::borrow::Cow::Borrowed(self.history.command.as_str())
+    }
+}

src/command/client/search/engines/tantivy.rs

@@ -1,73 +1,22 @@
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use atuin_client::{database::Database, history::History, settings::FilterMode};
+use atuin_client::{
+    database::Database,
+    history::History,
+    settings::FilterMode,
+    tantivy::{index, schema, HistorySchema},
+};
 use chrono::{TimeZone, Utc};
-use clap::Parser;
 use eyre::Result;
 use tantivy::{
     collector::TopDocs,
-    directory::MmapDirectory,
-    doc,
     query::{BooleanQuery, ConstScoreQuery, FuzzyTermQuery, Occur, Query, TermQuery},
-    schema::{Field, Schema, Value, FAST, STORED, STRING, TEXT},
-    DateTime, Index, IndexWriter, Searcher, Term,
+    schema::Value,
+    Searcher, Term,
 };
 
-use super::interactive::{HistoryWrapper, SearchEngine, SearchState};
-
-fn schema() -> (HistorySchema, Schema) {
-    let mut schema_builder = Schema::builder();
-
-    (
-        HistorySchema {
-            id: schema_builder.add_text_field("id", STRING),
-            command: schema_builder.add_text_field("command", TEXT | STORED),
-            cwd: schema_builder.add_text_field("cwd", STRING | FAST),
-            session: schema_builder.add_text_field("session", STRING | FAST),
-            hostname: schema_builder.add_text_field("hostname", STRING | FAST),
-            timestamp: schema_builder.add_date_field("timestamp", STORED),
-            duration: schema_builder.add_i64_field("duration", STORED),
-            exit: schema_builder.add_i64_field("exit", STORED),
-        },
-        schema_builder.build(),
-    )
-}
-
-struct HistorySchema {
-    id: Field,
-    command: Field,
-    cwd: Field,
-    session: Field,
-    hostname: Field,
-    timestamp: Field,
-    duration: Field,
-    exit: Field,
-}
-
-fn index(schema: Schema) -> Result<Index> {
-    let data_dir = atuin_common::utils::data_dir();
-    let tantivy_dir = data_dir.join("tantivy");
-
-    fs_err::create_dir_all(&tantivy_dir)?;
-    let dir = MmapDirectory::open(tantivy_dir)?;
-
-    Ok(Index::open_or_create(dir, schema)?)
-}
-
-pub fn write_history(h: impl IntoIterator<Item = History>) -> Result<()> {
-    let (hs, schema) = schema();
-    let index = index(schema)?;
-    let mut writer = index.writer(3_000_000)?;
-
-    for h in h {
-        write_single_history(&mut writer, &hs, h)?;
-    }
-
-    writer.commit()?;
-
-    Ok(())
-}
+use super::{HistoryWrapper, SearchEngine, SearchState};
 
 pub struct Search {
     schema: HistorySchema,
@@ -91,7 +40,7 @@ impl Search {
 
 #[async_trait]
 impl SearchEngine for Search {
-    async fn query(
+    async fn full_query(
         &mut self,
         state: &SearchState,
         _: &mut dyn Database,
@@ -182,43 +131,3 @@ impl SearchEngine for Search {
         Ok(output)
     }
 }
-
-fn write_single_history(
-    writer: &mut IndexWriter,
-    schema: &HistorySchema,
-    h: History,
-) -> Result<()> {
-    let timestamp = DateTime::from_timestamp_millis(h.timestamp.timestamp_millis());
-    writer.add_document(doc!(
-        schema.id => h.id,
-        schema.command => h.command,
-        schema.cwd => h.cwd,
-        schema.session => h.session,
-        schema.hostname => h.hostname,
-        schema.timestamp => timestamp,
-        schema.duration => h.duration,
-        schema.exit => h.exit,
-    ))?;
-
-    Ok(())
-}
-
-#[allow(clippy::struct_excessive_bools)]
-#[derive(Parser)]
-pub struct Cmd {}
-
-impl Cmd {
-    pub async fn run(self, db: &mut impl Database) -> Result<()> {
-        let history = db.all_with_count().await?;
-
-        // delete the index
-        let data_dir = atuin_common::utils::data_dir();
-        let tantivy_dir = dbg!(data_dir.join("tantivy"));
-        fs_err::remove_dir_all(tantivy_dir)?;
-
-        tokio::task::spawn_blocking(|| write_history(history.into_iter().map(|(h, _)| h)))
-            .await??;
-
-        Ok(())
-    }
-}

src/command/client/search/history_list.rs

@@ -8,7 +8,7 @@ use ratatui::{
     widgets::{Block, StatefulWidget, Widget},
 };
 
-use super::{format_duration, interactive::HistoryWrapper};
+use super::{engines::HistoryWrapper, format_duration};
 
 pub struct HistoryList<'a> {
     history: &'a [Arc<HistoryWrapper>],

src/command/client/search/interactive.rs

@@ -1,11 +1,9 @@
 use std::{
     io::{stdout, Write},
-    ops::Deref,
     sync::Arc,
     time::Duration,
 };
 
-use async_trait::async_trait;
 use crossterm::{
     event::{self, Event, KeyCode, KeyEvent, KeyModifiers, MouseEvent},
     execute, terminal,
@@ -13,21 +11,20 @@ use crossterm::{
 use eyre::Result;
 use futures_util::FutureExt;
 use semver::Version;
-use skim::SkimItem;
 use unicode_width::UnicodeWidthStr;
 
 use atuin_client::{
+    database::current_context,
     database::Database,
-    database::{current_context, Context},
-    history::History,
     settings::{ExitMode, FilterMode, SearchMode, Settings},
 };
 
 use super::{
     cursor::Cursor,
+    engines::{HistoryWrapper, SearchEngine, SearchState},
     history_list::{HistoryList, ListState, PREFIX_LENGTH},
 };
-use crate::VERSION;
+use crate::{command::client::search::engines, VERSION};
 use ratatui::{
     backend::{Backend, CrosstermBackend},
     layout::{Alignment, Constraint, Direction, Layout},
@@ -44,50 +41,16 @@ struct State {
     history_count: i64,
     update_needed: Option<Version>,
     results_state: ListState,
-    search: SearchState,
+    switched_search_mode: bool,
+    search_mode: SearchMode,
 
+    search: SearchState,
     engine: Box<dyn SearchEngine>,
 }
 
-pub struct SearchState {
-    pub input: Cursor,
-    pub filter_mode: FilterMode,
-    pub search_mode: SearchMode,
-    /// Store if the user has _just_ changed the search mode.
-    /// If so, we change the UI to show the search mode instead
-    /// of the filter mode until user starts typing again.
-    switched_search_mode: bool,
-    pub context: Context,
-}
-
-#[async_trait]
-pub trait SearchEngine: Send + Sync + 'static {
-    async fn query(
-        &mut self,
-        state: &SearchState,
-        db: &mut dyn Database,
-    ) -> Result<Vec<Arc<HistoryWrapper>>>;
-}
-
 impl State {
-    async fn query_results(&mut self, db: &mut impl Database) -> Result<Vec<Arc<HistoryWrapper>>> {
-        let i = self.search.input.as_str();
-        let results = if i.is_empty() {
-            db.list(
-                self.search.filter_mode,
-                &self.search.context,
-                Some(200),
-                true,
-            )
-            .await?
-            .into_iter()
-            .map(|history| HistoryWrapper { history, count: 1 })
-            .map(Arc::new)
-            .collect::<Vec<_>>()
-        } else {
-            self.engine.query(&self.search, db).await?
-        };
-
+    async fn query_results(&mut self, db: &mut dyn Database) -> Result<Vec<Arc<HistoryWrapper>>> {
+        let results = self.engine.query(&self.search, db).await?;
         self.results_state.select(0);
         Ok(results)
     }
@@ -137,7 +100,7 @@ impl State {
         let ctrl = input.modifiers.contains(KeyModifiers::CONTROL);
         let alt = input.modifiers.contains(KeyModifiers::ALT);
         // reset the state, will be set to true later if user really did change it
-        self.search.switched_search_mode = false;
+        self.switched_search_mode = false;
         match input.code {
             KeyCode::Char('c' | 'd' | 'g') if ctrl => return Some(RETURN_ORIGINAL),
             KeyCode::Esc => {
@@ -211,8 +174,10 @@ impl State {
                 self.search.filter_mode = FILTER_MODES[i];
             }
             KeyCode::Char('s') if ctrl => {
-                self.search.switched_search_mode = true;
-                self.search.search_mode = self.search.search_mode.next(settings);
+                self.switched_search_mode = true;
+                self.search_mode = self.search_mode.next(settings);
+                self.engine =
+                    engines::engine(self.search_mode).expect("could not switch search engine");
             }
             KeyCode::Down if self.results_state.selected() == 0 => return Some(RETURN_ORIGINAL),
             KeyCode::Down => {
@@ -385,8 +350,8 @@ impl State {
     fn build_input(&mut self, compact: bool, chunk_width: usize) -> Paragraph {
         /// Max width of the UI box showing current mode
         const MAX_WIDTH: usize = 14;
-        let (pref, mode) = if self.search.switched_search_mode {
-            (" SRCH:", self.search.search_mode.as_str())
+        let (pref, mode) = if self.switched_search_mode {
+            (" SRCH:", self.search_mode.as_str())
         } else {
             ("", self.search.filter_mode.as_str())
         };
@@ -483,23 +448,6 @@ impl Write for Stdout {
     }
 }
 
-pub struct HistoryWrapper {
-    pub history: History,
-    pub count: i32,
-}
-impl Deref for HistoryWrapper {
-    type Target = History;
-
-    fn deref(&self) -> &Self::Target {
-        &self.history
-    }
-}
-impl SkimItem for HistoryWrapper {
-    fn text(&self) -> std::borrow::Cow<str> {
-        std::borrow::Cow::Borrowed(self.history.command.as_str())
-    }
-}
-
 // this is a big blob of horrible! clean it up!
 // for now, it works. But it'd be great if it were more easily readable, and
 // modular. I'd like to add some more stats and stuff at some point
@@ -507,7 +455,7 @@ impl SkimItem for HistoryWrapper {
 pub async fn history(
     query: &[String],
     settings: &Settings,
-    db: &mut impl Database,
+    mut db: impl Database,
 ) -> Result<String> {
     let stdout = Stdout::new()?;
     let backend = CrosstermBackend::new(stdout);
@@ -523,16 +471,14 @@ pub async fn history(
 
     let context = current_context();
 
-    let engine = match settings.search_mode {
-        SearchMode::Skim => Box::new(super::skim_impl::Search::new()) as Box<_>,
-        SearchMode::Tantivy => Box::new(super::tantivy_impl::Search::new()?) as Box<_>,
-        mode => Box::new(super::db_impl::Search(mode)) as Box<_>,
-    };
+    let history_count = db.history_count().await?;
 
     let mut app = State {
-        history_count: db.history_count().await?,
+        history_count,
         results_state: ListState::default(),
         update_needed: None,
+        switched_search_mode: false,
+        search_mode: settings.search_mode,
         search: SearchState {
             input,
             context,
@@ -543,13 +489,11 @@ pub async fn history(
             } else {
                 settings.filter_mode
             },
-            search_mode: settings.search_mode,
-            switched_search_mode: false,
         },
-        engine,
+        engine: engines::engine(settings.search_mode)?,
     };
 
-    let mut results = app.query_results(db).await?;
+    let mut results = app.query_results(&mut db).await?;
 
     let index = 'render: loop {
         let compact = match settings.style {
@@ -563,7 +507,7 @@ pub async fn history(
 
         let initial_input = app.search.input.as_str().to_owned();
         let initial_filter_mode = app.search.filter_mode;
-        let initial_search_mode = app.search.search_mode;
+        let initial_search_mode = app.search_mode;
 
         let event_ready = tokio::task::spawn_blocking(|| event::poll(Duration::from_millis(250)));
 
@@ -587,9 +531,9 @@ pub async fn history(
 
         if initial_input != app.search.input.as_str()
             || initial_filter_mode != app.search.filter_mode
-            || initial_search_mode != app.search.search_mode
+            || initial_search_mode != app.search_mode
         {
-            results = app.query_results(db).await?;
+            results = app.query_results(&mut db).await?;
         }
     };
     if index < results.len() {