diff options
Diffstat (limited to 'atuin-client/src/database.rs')
-rw-r--r-- | atuin-client/src/database.rs | 113 |
1 files changed, 96 insertions, 17 deletions
diff --git a/atuin-client/src/database.rs b/atuin-client/src/database.rs index abc22bb8..0855359b 100644 --- a/atuin-client/src/database.rs +++ b/atuin-client/src/database.rs @@ -2,7 +2,7 @@ use chrono::prelude::*; use chrono::Utc; use std::path::Path; -use eyre::Result; +use eyre::{eyre, Result}; use rusqlite::{params, Connection}; use rusqlite::{Params, Transaction}; @@ -14,7 +14,7 @@ pub trait Database { fn save_bulk(&mut self, h: &[History]) -> Result<()>; fn load(&self, id: &str) -> Result<History>; - fn list(&self) -> Result<Vec<History>>; + fn list(&self, max: Option<usize>, unique: bool) -> Result<Vec<History>>; fn range(&self, from: chrono::DateTime<Utc>, to: chrono::DateTime<Utc>) -> Result<Vec<History>>; @@ -27,6 +27,8 @@ pub trait Database { fn before(&self, timestamp: chrono::DateTime<Utc>, count: i64) -> Result<Vec<History>>; fn prefix_search(&self, query: &str) -> Result<Vec<History>>; + + fn search(&self, cwd: Option<String>, exit: Option<i64>, query: &str) -> Result<Vec<History>>; } // Intended for use on a developer machine and not a sync server. @@ -81,6 +83,16 @@ impl Sqlite { [], )?; + conn.execute( + "create index if not exists idx_history_timestamp on history(timestamp)", + [], + )?; + + conn.execute( + "create index if not exists idx_history_command on history(command)", + [], + )?; + Ok(()) } @@ -136,16 +148,19 @@ impl Database for Sqlite { } fn load(&self, id: &str) -> Result<History> { - debug!("loading history item"); + debug!("loading history item {}", id); - let mut stmt = self.conn.prepare( + let history = self.query( "select id, timestamp, duration, exit, command, cwd, session, hostname from history - where id = ?1", + where id = ?1 limit 1", + &[id], )?; - let history = stmt.query_row(params![id], |row| { - history_from_sqlite_row(Some(id.to_string()), row) - })?; + if history.is_empty() { + return Err(eyre!("could not find history with id {}", id)); + } + + let history = history[0].clone(); Ok(history) } @@ -163,16 +178,39 @@ impl Database for Sqlite { Ok(()) } - fn list(&self) -> Result<Vec<History>> { + // make a unique list, that only shows the *newest* version of things + fn list(&self, max: Option<usize>, unique: bool) -> Result<Vec<History>> { debug!("listing history"); - let mut stmt = self - .conn - .prepare("SELECT * FROM history order by timestamp asc")?; + // very likely vulnerable to SQL injection + // however, this is client side, and only used by the client, on their + // own data. They can just open the db file... + // otherwise building the query is awkward + let query = format!( + "select * from history h + {} + order by timestamp desc + {}", + // inject the unique check + if unique { + "where timestamp = ( + select max(timestamp) from history + where h.command = history.command + )" + } else { + "" + }, + // inject the limit + if let Some(max) = max { + format!("limit {}", max) + } else { + "".to_string() + } + ); - let history_iter = stmt.query_map(params![], |row| history_from_sqlite_row(None, row))?; + let history = self.query(query.as_str(), params![])?; - Ok(history_iter.filter_map(Result::ok).collect()) + Ok(history) } fn range( @@ -207,7 +245,7 @@ impl Database for Sqlite { fn last(&self) -> Result<History> { let mut stmt = self .conn - .prepare("SELECT * FROM history order by timestamp desc limit 1")?; + .prepare("SELECT * FROM history where duration >= 0 order by timestamp desc limit 1")?; let history = stmt.query_row(params![], |row| history_from_sqlite_row(None, row))?; @@ -235,9 +273,17 @@ impl Database for Sqlite { } fn prefix_search(&self, query: &str) -> Result<Vec<History>> { + let query = query.to_string().replace("*", "%"); // allow wildcard char + self.query( - "select * from history where command like ?1 || '%' order by timestamp asc limit 1000", - &[query], + "select * from history h + where command like ?1 || '%' + and timestamp = ( + select max(timestamp) from history + where h.command = history.command + ) + order by timestamp desc limit 200", + &[query.as_str()], ) } @@ -248,6 +294,39 @@ impl Database for Sqlite { Ok(res) } + + fn search(&self, cwd: Option<String>, exit: Option<i64>, query: &str) -> Result<Vec<History>> { + match (cwd, exit) { + (Some(cwd), Some(exit)) => self.query( + "select * from history + where command like ?1 || '%' + and cwd = ?2 + and exit = ?3 + order by timestamp asc limit 1000", + &[query, cwd.as_str(), exit.to_string().as_str()], + ), + (Some(cwd), None) => self.query( + "select * from history + where command like ?1 || '%' + and cwd = ?2 + order by timestamp asc limit 1000", + &[query, cwd.as_str()], + ), + (None, Some(exit)) => self.query( + "select * from history + where command like ?1 || '%' + and exit = ?2 + order by timestamp asc limit 1000", + &[query, exit.to_string().as_str()], + ), + (None, None) => self.query( + "select * from history + where command like ?1 || '%' + order by timestamp asc limit 1000", + &[query], + ), + } + } } fn history_from_sqlite_row( |