From 057ed96c0a7770fe9dd06b15d120fa1a3bdb6f91 Mon Sep 17 00:00:00 2001 From: Ajeet D'Souza <98ajeet@gmail.com> Date: Sat, 28 Mar 2020 23:51:46 +0530 Subject: [PATCH] Refactor DB architecture --- src/config.rs | 7 ++- src/db.rs | 130 ++++++++++++++++++++++++++++----------- src/dir.rs | 5 +- src/main.rs | 1 - src/subcommand/import.rs | 4 +- src/subcommand/query.rs | 3 +- src/subcommand/remove.rs | 4 +- src/types.rs | 2 - src/util.rs | 4 +- 9 files changed, 113 insertions(+), 47 deletions(-) delete mode 100644 src/types.rs diff --git a/src/config.rs b/src/config.rs index 39a8ceb..6e41b1b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,4 +1,5 @@ -use crate::types::Rank; +use crate::db::DBVersion; +use crate::dir::Rank; use anyhow::{bail, Context, Result}; @@ -6,6 +7,10 @@ use std::env; use std::fs; use std::path::PathBuf; +pub const DB_MAX_SIZE: u64 = 8 * 1024 * 1024; // 8 MiB + +pub const DB_VERSION: DBVersion = 3; + pub fn zo_data() -> Result { let path = match env::var_os("_ZO_DATA") { Some(data_osstr) => PathBuf::from(data_osstr), diff --git a/src/db.rs b/src/db.rs index dbb9d3e..32d3f2a 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1,40 +1,75 @@ -use crate::dir::Dir; -use crate::types::{Epoch, Rank}; +use crate::config; +use crate::dir::{Dir, Epoch, Rank}; use anyhow::{anyhow, bail, Context, Result}; use indoc::indoc; +use serde::{Deserialize, Serialize}; + +use std::cmp::Ordering; use std::fs::{self, File}; use std::io::{self, BufRead, BufReader, BufWriter}; use std::path::{Path, PathBuf}; +pub use i32 as DBVersion; + pub struct DB { - path: PathBuf, - dirs: Vec, + data: DBData, modified: bool, + path: PathBuf, + path_old: Option, } impl DB { pub fn open>(path: P) -> Result { - let path = path.as_ref().to_path_buf(); - - let dirs = match File::open(&path) { + let data = match File::open(&path) { Ok(file) => { let reader = BufReader::new(&file); bincode::config() - .limit(8 * 1024 * 1024) // only databases upto 8 MiB are supported + .limit(config::DB_MAX_SIZE) .deserialize_from(reader) .context("could not deserialize database")? } Err(err) => match err.kind() { - io::ErrorKind::NotFound => Vec::::new(), + io::ErrorKind::NotFound => DBData::default(), _ => return Err(err).context("could not open database file"), }, }; + if data.version != config::DB_VERSION { + bail!("this database version ({}) is unsupported", data.version); + } + Ok(DB { - path, - dirs, + data, modified: false, + path: path.as_ref().to_path_buf(), + path_old: None, + }) + } + + pub fn open_old(path_old: P1, path: P2) -> Result + where + P1: AsRef, + P2: AsRef, + { + let file = File::open(&path_old).context("could not open old database file")?; + let reader = BufReader::new(&file); + + let dirs = bincode::config() + .limit(config::DB_MAX_SIZE) + .deserialize_from(reader) + .context("could not deserialize old database")?; + + let data = DBData { + version: config::DB_VERSION, + dirs, + }; + + Ok(DB { + data, + modified: true, + path: path.as_ref().to_path_buf(), + path_old: Some(path_old.as_ref().to_path_buf()), }) } @@ -46,7 +81,8 @@ impl DB { File::create(&path_tmp).context("could not open temporary database file")?; let writer = BufWriter::new(&file_tmp); - bincode::serialize_into(writer, &self.dirs).context("could not serialize database")?; + bincode::serialize_into(writer, &self.data) + .context("could not serialize database")?; fs::rename(&path_tmp, &self.path).context("could not move temporary database file")?; } @@ -55,7 +91,7 @@ impl DB { } pub fn import>(&mut self, path: P, merge: bool) -> Result<()> { - if !self.dirs.is_empty() && !merge { + if !self.data.dirs.is_empty() && !merge { bail!(indoc!( "To prevent conflicts, you can only import from z with an empty zoxide database! If you wish to merge the two, specify the `--merge` flag." @@ -109,7 +145,9 @@ impl DB { if merge { // If the path exists in the database, add the ranks and set the epoch to // the largest of the parsed epoch and the already present epoch. - if let Some(dir) = self.dirs.iter_mut().find(|dir| dir.path == path_abs) { + if let Some(dir) = + self.data.dirs.iter_mut().find(|dir| dir.path == path_abs) + { dir.rank += rank; dir.last_accessed = Epoch::max(epoch, dir.last_accessed); @@ -117,7 +155,7 @@ impl DB { }; } - self.dirs.push(Dir { + self.data.dirs.push(Dir { path: path_abs, rank, last_accessed: epoch, @@ -142,8 +180,8 @@ impl DB { .canonicalize() .with_context(|| anyhow!("could not access directory: {}", path.as_ref().display()))?; - match self.dirs.iter_mut().find(|dir| dir.path == path_abs) { - None => self.dirs.push(Dir { + match self.data.dirs.iter_mut().find(|dir| dir.path == path_abs) { + None => self.data.dirs.push(Dir { path: path_abs, last_accessed: now, rank: 1.0, @@ -154,15 +192,15 @@ impl DB { } }; - let sum_age = self.dirs.iter().map(|dir| dir.rank).sum::(); + let sum_age = self.data.dirs.iter().map(|dir| dir.rank).sum::(); if sum_age > max_age { let factor = 0.9 * max_age / sum_age; - for dir in &mut self.dirs { + for dir in &mut self.data.dirs { dir.rank *= factor; } - self.dirs.retain(|dir| dir.rank >= 1.0); + self.data.dirs.retain(|dir| dir.rank >= 1.0); } self.modified = true; @@ -170,26 +208,36 @@ impl DB { } pub fn query(&mut self, keywords: &[String], now: Epoch) -> Option { - let (idx, dir) = self + let (idx, dir, _) = self + .data .dirs .iter() .enumerate() .filter(|(_, dir)| dir.is_match(&keywords)) - .max_by_key(|(_, dir)| dir.get_frecency(now) as i64)?; + .map(|(idx, dir)| (idx, dir, dir.get_frecency(now))) + .max_by(|(_, _, frecency1), (_, _, frecency2)| { + frecency1.partial_cmp(frecency2).unwrap_or(Ordering::Equal) + })?; if dir.is_dir() { Some(dir.to_owned()) } else { - self.dirs.swap_remove(idx); + self.data.dirs.swap_remove(idx); self.modified = true; self.query(keywords, now) } } pub fn query_all(&mut self, keywords: &[String]) -> Vec { - self.remove_invalid(); + let orig_len = self.data.dirs.len(); + self.data.dirs.retain(Dir::is_dir); - self.dirs + if orig_len != self.data.dirs.len() { + self.modified = true; + } + + self.data + .dirs .iter() .filter(|dir| dir.is_match(&keywords)) .cloned() @@ -202,8 +250,8 @@ impl DB { Err(_) => path.as_ref().to_path_buf(), }; - if let Some(idx) = self.dirs.iter().position(|dir| dir.path == path_abs) { - self.dirs.swap_remove(idx); + if let Some(idx) = self.data.dirs.iter().position(|dir| dir.path == path_abs) { + self.data.dirs.swap_remove(idx); self.modified = true; } @@ -215,21 +263,31 @@ impl DB { path_tmp.set_file_name("db.zo.tmp"); path_tmp } - - fn remove_invalid(&mut self) { - let orig_len = self.dirs.len(); - self.dirs.retain(Dir::is_dir); - - if orig_len != self.dirs.len() { - self.modified = true; - } - } } impl Drop for DB { fn drop(&mut self) { if let Err(e) = self.save() { eprintln!("{:#}", e); + } else if let Some(path_old) = &self.path_old { + if let Err(e) = fs::remove_file(path_old).context("could not remove old database") { + eprintln!("{:#}", e); + } + } + } +} + +#[derive(Debug, Deserialize, Serialize)] +struct DBData { + version: DBVersion, + dirs: Vec, +} + +impl Default for DBData { + fn default() -> DBData { + DBData { + version: config::DB_VERSION, + dirs: Vec::new(), } } } diff --git a/src/dir.rs b/src/dir.rs index c24d74a..aa7f67d 100644 --- a/src/dir.rs +++ b/src/dir.rs @@ -1,8 +1,9 @@ -use crate::types::{Epoch, Rank}; - use serde::{Deserialize, Serialize}; use std::path::PathBuf; +pub use f64 as Rank; +pub use i64 as Epoch; // use a signed integer so subtraction can be performed on it + #[derive(Clone, Debug, Default, Deserialize, Serialize)] pub struct Dir { pub path: PathBuf, diff --git a/src/main.rs b/src/main.rs index f1c2766..a480b45 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,7 +2,6 @@ mod config; mod db; mod dir; mod subcommand; -mod types; mod util; use anyhow::Result; diff --git a/src/subcommand/import.rs b/src/subcommand/import.rs index f691afd..3e10191 100644 --- a/src/subcommand/import.rs +++ b/src/subcommand/import.rs @@ -3,10 +3,12 @@ use crate::util; use anyhow::Result; use structopt::StructOpt; +use std::path::PathBuf; + #[derive(Debug, StructOpt)] #[structopt(about = "Import from z database")] pub struct Import { - path: String, + path: PathBuf, #[structopt(long, help = "Merge entries into existing database")] merge: bool, diff --git a/src/subcommand/query.rs b/src/subcommand/query.rs index 4bae377..cba97d1 100644 --- a/src/subcommand/query.rs +++ b/src/subcommand/query.rs @@ -1,9 +1,10 @@ use crate::util; use anyhow::{bail, Result}; +use structopt::StructOpt; + use std::io::{self, Write}; use std::path::Path; -use structopt::StructOpt; #[derive(Debug, StructOpt)] #[structopt(about = "Search for a directory")] diff --git a/src/subcommand/remove.rs b/src/subcommand/remove.rs index b115c28..5ba579d 100644 --- a/src/subcommand/remove.rs +++ b/src/subcommand/remove.rs @@ -3,10 +3,12 @@ use crate::util; use anyhow::Result; use structopt::StructOpt; +use std::path::PathBuf; + #[derive(Debug, StructOpt)] #[structopt(about = "Remove a directory")] pub struct Remove { - path: String, + path: PathBuf, } impl Remove { diff --git a/src/types.rs b/src/types.rs deleted file mode 100644 index 1ddffd1..0000000 --- a/src/types.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub use f64 as Rank; -pub use i64 as Epoch; // use a signed integer so subtraction can be performed on it diff --git a/src/util.rs b/src/util.rs index b6615ce..c9e4057 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,9 +1,9 @@ use crate::config; use crate::db::DB; -use crate::dir::Dir; -use crate::types::Epoch; +use crate::dir::{Dir, Epoch}; use anyhow::{anyhow, bail, Context, Result}; + use std::cmp::{Ordering, PartialOrd}; use std::io::{Read, Write}; use std::path::Path;