diff --git a/Cargo.lock b/Cargo.lock index 1b23b7d1..b7063297 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -151,6 +151,7 @@ dependencies = [ "sodiumoxide", "sql-builder", "sqlx", + "tantivy", "tokio", "urlencoding", "uuid", diff --git a/atuin-client/Cargo.toml b/atuin-client/Cargo.toml index 1fc5baf1..764bee23 100644 --- a/atuin-client/Cargo.toml +++ b/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 } diff --git a/atuin-client/src/lib.rs b/atuin-client/src/lib.rs index 497c5e74..b68bc47a 100644 --- a/atuin-client/src/lib.rs +++ b/atuin-client/src/lib.rs @@ -15,3 +15,4 @@ pub mod history; pub mod import; pub mod ordering; pub mod settings; +pub mod tantivy; diff --git a/atuin-client/src/sync.rs b/atuin-client/src/sync.rs index 1c0acaf8..29228682 100644 --- a/atuin-client/src/sync.rs +++ b/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 { diff --git a/atuin-client/src/tantivy.rs b/atuin-client/src/tantivy.rs new file mode 100644 index 00000000..f36e5f92 --- /dev/null +++ b/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 { + 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) -> 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, +) -> 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(()) +} diff --git a/src/command/client.rs b/src/command/client.rs index 5d7404e3..7eb8cecb 100644 --- a/src/command/client.rs +++ b/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, } diff --git a/src/command/client/history.rs b/src/command/client/history.rs index c8f6b535..e796be20 100644 --- a/src/command/client/history.rs +++ b/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")] diff --git a/src/command/client/search.rs b/src/command/client/search.rs index 4449104d..f307694b 100644 --- a/src/command/client/search.rs +++ b/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?; diff --git a/src/command/client/search/engines.rs b/src/command/client/search/engines.rs new file mode 100644 index 00000000..93d0bf40 --- /dev/null +++ b/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> { + 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>>; + + async fn query( + &mut self, + state: &SearchState, + db: &mut dyn Database, + ) -> Result>> { + 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::>()) + } 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 + } +} diff --git a/src/command/client/search/db_impl.rs b/src/command/client/search/engines/db.rs similarity index 89% rename from src/command/client/search/db_impl.rs rename to src/command/client/search/engines/db.rs index 4d197549..a6a7275b 100644 --- a/src/command/client/search/db_impl.rs +++ b/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, diff --git a/src/command/client/search/skim_impl.rs b/src/command/client/search/engines/skim.rs similarity index 93% rename from src/command/client/search/skim_impl.rs rename to src/command/client/search/engines/skim.rs index b0f4ca28..6be1c670 100644 --- a/src/command/client/search/skim_impl.rs +++ b/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>, } 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 { + std::borrow::Cow::Borrowed(self.history.command.as_str()) + } +} diff --git a/src/command/client/search/tantivy_impl.rs b/src/command/client/search/engines/tantivy.rs similarity index 58% rename from src/command/client/search/tantivy_impl.rs rename to src/command/client/search/engines/tantivy.rs index 77297e4f..9a63ed94 100644 --- a/src/command/client/search/tantivy_impl.rs +++ b/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 { - 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) -> 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(()) - } -} diff --git a/src/command/client/search/history_list.rs b/src/command/client/search/history_list.rs index 60ec15a8..8aea012c 100644 --- a/src/command/client/search/history_list.rs +++ b/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], diff --git a/src/command/client/search/interactive.rs b/src/command/client/search/interactive.rs index 30fff4a5..2c4f5863 100644 --- a/src/command/client/search/interactive.rs +++ b/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, results_state: ListState, - search: SearchState, + switched_search_mode: bool, + search_mode: SearchMode, + search: SearchState, engine: Box, } -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>>; -} - impl State { - async fn query_results(&mut self, db: &mut impl Database) -> Result>> { - 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::>() - } else { - self.engine.query(&self.search, db).await? - }; - + async fn query_results(&mut self, db: &mut dyn Database) -> Result>> { + 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 { - 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 { 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() {