Add fuzzy text search mode (#142)

6 files changed, 132 insertions, 3 deletions

diff --git a/Cargo.lock b/Cargo.lock
index b46bdbbd..abb5959f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -38,6 +38,27 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "23b62fc65de8e4e7f52534fb52b0f3ed04746ae267519eef2a83941e8085068b"
 
 [[package]]
+name = "async-stream"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "171374e7e3b2504e0e5236e3b59260560f9fe94bfe9ac39ba5e4e929c5590625"
+dependencies = [
+ "async-stream-impl",
+ "futures-core",
+]
+
+[[package]]
+name = "async-stream-impl"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "648ed8c8d2ce5409ccd57453d9d1b214b342a0d69376a6feda1fd6cae3299308"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
 name = "async-trait"
 version = "0.1.50"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -129,6 +150,7 @@ dependencies = [
  "sodiumoxide",
  "sqlx",
  "tokio",
+ "tokio-test",
  "urlencoding",
  "uuid",
  "whoami",
@@ -2324,6 +2346,19 @@ dependencies = [
 ]
 
 [[package]]
+name = "tokio-test"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "53474327ae5e166530d17f2d956afcb4f8a004de581b3cae10f12006bc8163e3"
+dependencies = [
+ "async-stream",
+ "bytes",
+ "futures-core",
+ "tokio",
+ "tokio-stream",
+]
+
+[[package]]
 name = "tokio-tungstenite"
 version = "0.13.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/atuin-client/Cargo.toml b/atuin-client/Cargo.toml
index 09bccaa0..b4c9c13e 100644
--- a/atuin-client/Cargo.toml
+++ b/atuin-client/Cargo.toml
@@ -40,3 +40,6 @@ humantime = "2.1.0"
 itertools = "0.10.0"
 shellexpand = "2"
 sqlx = { version = "0.5", features = [ "runtime-tokio-rustls", "uuid", "chrono", "sqlite" ] }
+
+[dev-dependencies]
+tokio-test = "*"
diff --git a/atuin-client/config.toml b/atuin-client/config.toml
index c8926297..bfb51149 100644
--- a/atuin-client/config.toml
+++ b/atuin-client/config.toml
@@ -24,5 +24,5 @@
 # sync_address = "https://api.atuin.sh"
 
 ## which search mode to use
-## possible values: prefix, fulltext
+## possible values: prefix, fulltext, fuzzy
 # search_mode = "prefix"
diff --git a/atuin-client/src/database.rs b/atuin-client/src/database.rs
index 160c6054..6a70ae33 100644
--- a/atuin-client/src/database.rs
+++ b/atuin-client/src/database.rs
@@ -6,6 +6,7 @@ use chrono::prelude::*;
 use chrono::Utc;
 
 use eyre::Result;
+use itertools::Itertools;
 
 use sqlx::sqlite::{
     SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions, SqliteRow,
@@ -286,6 +287,7 @@ impl Database for Sqlite {
         let query = match search_mode {
             SearchMode::Prefix => query,
             SearchMode::FullText => format!("%{}", query),
+            SearchMode::Fuzzy => query.split("").join("%"),
         };
 
         let res = sqlx::query(
@@ -318,3 +320,89 @@ impl Database for Sqlite {
         Ok(res)
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    async fn new_history_item(db: &mut impl Database, cmd: &str) -> Result<()> {
+        let history = History::new(
+            chrono::Utc::now(),
+            cmd.to_string(),
+            "/home/ellie".to_string(),
+            0,
+            1,
+            Some("beep boop".to_string()),
+            Some("booop".to_string()),
+        );
+        return db.save(&history).await;
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_search_prefix() {
+        let mut db = Sqlite::new("sqlite::memory:").await.unwrap();
+        new_history_item(&mut db, "ls /home/ellie").await.unwrap();
+
+        let mut results = db.search(None, SearchMode::Prefix, "ls").await.unwrap();
+        assert_eq!(results.len(), 1);
+
+        results = db.search(None, SearchMode::Prefix, "/home").await.unwrap();
+        assert_eq!(results.len(), 0);
+
+        results = db.search(None, SearchMode::Prefix, "ls  ").await.unwrap();
+        assert_eq!(results.len(), 0);
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_search_fulltext() {
+        let mut db = Sqlite::new("sqlite::memory:").await.unwrap();
+        new_history_item(&mut db, "ls /home/ellie").await.unwrap();
+
+        let mut results = db.search(None, SearchMode::FullText, "ls").await.unwrap();
+        assert_eq!(results.len(), 1);
+
+        results = db
+            .search(None, SearchMode::FullText, "/home")
+            .await
+            .unwrap();
+        assert_eq!(results.len(), 1);
+
+        results = db.search(None, SearchMode::FullText, "ls  ").await.unwrap();
+        assert_eq!(results.len(), 0);
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_search_fuzzy() {
+        let mut db = Sqlite::new("sqlite::memory:").await.unwrap();
+        new_history_item(&mut db, "ls /home/ellie").await.unwrap();
+        new_history_item(&mut db, "ls /home/frank").await.unwrap();
+        new_history_item(&mut db, "cd /home/ellie").await.unwrap();
+        new_history_item(&mut db, "/home/ellie/.bin/rustup")
+            .await
+            .unwrap();
+
+        let mut results = db.search(None, SearchMode::Fuzzy, "ls /").await.unwrap();
+        assert_eq!(results.len(), 2);
+
+        results = db.search(None, SearchMode::Fuzzy, "l/h/").await.unwrap();
+        assert_eq!(results.len(), 2);
+
+        results = db.search(None, SearchMode::Fuzzy, "/h/e").await.unwrap();
+        assert_eq!(results.len(), 3);
+
+        results = db.search(None, SearchMode::Fuzzy, "/hmoe/").await.unwrap();
+        assert_eq!(results.len(), 0);
+
+        results = db
+            .search(None, SearchMode::Fuzzy, "ellie/home")
+            .await
+            .unwrap();
+        assert_eq!(results.len(), 0);
+
+        results = db.search(None, SearchMode::Fuzzy, "lsellie").await.unwrap();
+        assert_eq!(results.len(), 1);
+
+        results = db.search(None, SearchMode::Fuzzy, " ").await.unwrap();
+        assert_eq!(results.len(), 3);
+    }
+}
diff --git a/atuin-client/src/settings.rs b/atuin-client/src/settings.rs
index 1d7e9a5f..0cb845c3 100644
--- a/atuin-client/src/settings.rs
+++ b/atuin-client/src/settings.rs
@@ -17,6 +17,9 @@ pub enum SearchMode {
 
     #[serde(rename = "fulltext")]
     FullText,
+
+    #[serde(rename = "fuzzy")]
+    Fuzzy,
 }
 
 // FIXME: Can use upstream Dialect enum if https://github.com/stevedonovan/chrono-english/pull/16 is merged
diff --git a/docs/config.md b/docs/config.md
index 29ce770c..414146f4 100644
--- a/docs/config.md
+++ b/docs/config.md
@@ -96,8 +96,8 @@ key = "~/.atuin-session"
 
 ### `search_mode`
 
-Which search mode to use. Atuin supports both "prefix" and full text search
-modes. The former will essentially search for "query*", and the latter "*query\*"
+Which search mode to use. Atuin supports "prefix", full text and "fuzzy" search
+modes. The prefix search for "query\*", fulltext "\*query\*", and fuzzy "\*q\*u\*e\*r\*y\*"
 
 Defaults to "prefix"