use anyhow::Result;
use indicatif::ProgressBar;
use regex::Regex;
use rusqlite::{params, Connection};
use std::convert::TryInto;
use crate::error::GcdError;
pub struct Database {
conn: rusqlite::Connection,
}
impl Database {
pub fn new(database_file: &str) -> Result<Self> {
let conn = Connection::open(database_file)?;
create_database(&conn)?;
Ok(Database { conn })
}
pub fn increment(&self, project: String) -> Result<()> {
increment_project_ref_count(&self.conn, project)
}
pub fn add(&self, projects: Vec<String>) -> Result<()> {
add_all_projects(&self.conn, projects)
}
pub fn add_new(&self, projects: Vec<String>) -> Result<()> {
add_new_projects(&self.conn, projects)
}
pub fn find(&self, input: &str) -> Result<Vec<String>> {
find(&self.conn, input)
}
pub fn all(&self) -> Result<Vec<String>> {
Ok(select_all_projects(&self.conn)?
.iter()
.map(|p| p.0.clone())
.collect())
}
pub fn remove(&self, project: String) -> Result<()> {
remove_project(&self.conn, project)
}
pub fn append(&self, project: String) -> Result<()> {
add_project(&self.conn, project)
}
pub fn alias(&self, project: &str, alias: &str) -> Result<()> {
add_alias(&self.conn, project, alias)
}
pub fn remove_alias(&self, project: &str) -> Result<()> {
remove_alias(&self.conn, project)
}
pub fn remove_alias_by_alias(&self, alias: &str) -> Result<()> {
remove_alias_by_alias(&self.conn, alias)
}
pub fn all_aliased(&self) -> Result<Vec<(String, String)>> {
let mut aliases: Vec<(String, String)> = select_all_projects(&self.conn)?
.iter()
.filter(|p| p.1.is_some())
.map(|p| (p.1.as_ref().unwrap().to_owned(), p.0.clone()))
.collect();
aliases.sort_by(|a, b| a.0.cmp(&b.0));
Ok(aliases)
}
pub fn move_project(&self, from_location: &str, to_location: &str) -> Result<()> {
move_project(&self.conn, from_location, to_location)
}
}
fn create_database(conn: &rusqlite::Connection) -> Result<()> {
conn.execute(
"CREATE TABLE IF NOT EXISTS projects (name TEXT PRIMARY KEY, alias TEXT NULL, ref_count INTEGER)",
[],
)?;
Ok(())
}
fn increment_project_ref_count(conn: &rusqlite::Connection, project: String) -> Result<()> {
let nr_updated = conn
.execute(
"UPDATE projects SET ref_count = ref_count + 1 WHERE name = ?",
[project.clone()],
)?;
if nr_updated != 1 {
return Err(GcdError::new(format!(
"Failed to update ref_count for project {}, project not found.",
project
))
.into());
}
Ok(())
}
fn add_alias(conn: &rusqlite::Connection, project: &str, alias: &str) -> Result<()> {
let nr_updated = conn
.execute(
"UPDATE projects SET alias = ?1 WHERE name = ?2",
[alias, project],
)?;
if nr_updated != 1 {
return Err(GcdError::new(format!(
"Failed to set alias {} for project {}, project not found.",
alias, project
))
.into());
}
Ok(())
}
fn remove_alias(conn: &rusqlite::Connection, project: &str) -> Result<()> {
let nr_updated = conn
.execute(
"UPDATE projects SET alias = null WHERE name = ?",
[project],
)?;
if nr_updated == 0 {
return Err(GcdError::new(format!(
"Failed to remove alias for project {}, project not found.",
project
))
.into());
}
Ok(())
}
fn remove_alias_by_alias(conn: &rusqlite::Connection, alias: &str) -> Result<()> {
let nr_updated = conn
.execute(
"UPDATE projects SET alias = null WHERE alias = ?",
[alias],
)?;
if nr_updated == 0 {
return Err(GcdError::new(format!(
"Failed to remove alias {}, alias not found.",
alias
))
.into());
}
Ok(())
}
fn delete_all_projects(conn: &rusqlite::Connection) -> Result<()> {
conn.execute("DELETE FROM projects", [])?;
Ok(())
}
fn select_all_projects(conn: &rusqlite::Connection) -> Result<Vec<(String, Option<String>)>> {
let mut stmt = conn
.prepare("SELECT name, alias FROM projects ORDER BY ref_count DESC, alias, name")?;
let projects_iter = stmt.query_map(params![], |row| match row.get(1) {
Ok(alias) => Ok((row.get(0)?, Some(alias))),
Err(_) => Ok((row.get(0)?, None)),
})?;
let mut projects: Vec<(String, Option<String>)> = vec![];
for project in projects_iter {
projects.push(project?);
}
Ok(projects)
}
fn add_all_projects(conn: &rusqlite::Connection, projects: Vec<String>) -> Result<()> {
delete_all_projects(conn)?;
add_new_projects(conn, projects)?;
Ok(())
}
fn add_new_projects(conn: &rusqlite::Connection, projects: Vec<String>) -> Result<()> {
let progress_bar = ProgressBar::new(projects.len().try_into()?);
progress_bar.println("Adding found projects to database");
let mut stmt = conn
.prepare("INSERT INTO projects(name, alias, ref_count) VALUES(?, null, 0)")?;
let mut added: usize = 0;
for project in projects {
match stmt.execute([project.clone()]) {
Ok(_) => {
progress_bar.set_message(format!("Added {}", project));
added += 1;
}
Err(e) => {
progress_bar.set_message(format!("Skipping {} - {}", project, e));
}
};
progress_bar.inc(1);
}
progress_bar.finish_with_message(format!("done. Added {} new projects", added));
Ok(())
}
fn add_project(conn: &rusqlite::Connection, project: String) -> Result<()> {
conn.execute(
"INSERT INTO projects(name, alias, ref_count) VALUES(?, null, 0)",
[project.clone()],
)?;
Ok(())
}
fn find(conn: &rusqlite::Connection, find: &str) -> Result<Vec<String>> {
match Regex::new(find) {
Ok(filter) => {
let mut projects: Vec<String> = vec![];
for (name, alias) in select_all_projects(conn)? {
if filter.is_match(&name) || (alias.is_some() && filter.is_match(&alias.unwrap())) {
projects.push(name);
}
}
Ok(projects)
}
Err(_) => Ok(vec![]),
}
}
fn remove_project(conn: &rusqlite::Connection, project: String) -> Result<()> {
conn.execute("DELETE FROM projects WHERE name =?", [project.clone()])?;
Ok(())
}
fn move_project(
conn: &rusqlite::Connection,
from_location: &str,
to_location: &str,
) -> Result<()> {
let nr_updated = conn
.execute(
"UPDATE projects SET name = ?1 WHERE name = ?2",
[to_location, from_location],
)?;
if nr_updated != 1 {
return Err(GcdError::new(format!(
"Failed to move project {} to {}, project not found.",
from_location, to_location
))
.into());
}
Ok(())
}
#[cfg(test)]
mod test {
use super::*;
use std::path::MAIN_SEPARATOR;
#[test]
fn open_empty_db() {
let result = Database::new("open_test.db");
assert!(result.is_ok());
assert!(std::fs::remove_file("open_test.db").is_ok());
}
#[test]
fn update_non_existing_database() {
let db = setup_db("non_err_test.db");
tear_down("non_err_test.db");
assert!(db.increment("project".to_owned()).is_err());
}
#[test]
fn increment_non_existing_project() {
let db = setup_db("increment_err_test.db");
assert!(db.increment("project".to_owned()).is_err());
tear_down("increment_err_test.db");
}
#[test]
fn increment_existing_project() {
let db = setup_db("increment_test.db");
assert!(db.increment("aproject".to_owned()).is_ok());
tear_down("increment_test.db");
}
#[test]
fn add_existing_project() {
let db = setup_db("add_exsisting_test.db");
assert!(db.add(vec!["aproject".to_owned()]).is_ok());
tear_down("add_exsisting_test.db");
}
#[test]
fn add_new_existing_project() {
let db = setup_db("add_new_exsisting_test.db");
assert!(db.add_new(vec!["cproject".to_owned()]).is_ok());
tear_down("add_new_exsisting_test.db");
}
#[test]
fn find_project() {
let db = setup_db("find_test.db");
let projects = db.find("apro").unwrap();
assert!(projects.iter().any(|p| p == "aproject"));
tear_down("find_test.db");
}
#[test]
fn find_project_by_alias() {
let db = setup_db("find_by_alias_test.db");
assert!(db.alias("aproject", "a-alias").is_ok());
let projects = db.find("a-alias").unwrap();
assert!(projects.iter().any(|p| p == "aproject"));
tear_down("find_by_alias_test.db");
}
#[test]
fn alias() {
let db = setup_db("alias_test.db");
assert!(db.alias("aproject", "a-alias").is_ok());
assert!(db.alias("bproject", "b-alias").is_ok());
assert!(db.alias("cproject", "c-alias").is_err());
assert!(db.remove_alias("bproject").is_ok());
assert!(db.remove_alias("dproject").is_err());
assert!(db.remove_alias_by_alias("a-alias").is_ok());
assert!(db.remove_alias_by_alias("d-alias").is_err());
tear_down("alias_test.db");
}
#[test]
fn move_project() {
let db = setup_db("move_test.db");
assert!(db.move_project("aproject", "cproject").is_ok());
assert!(db.move_project("aproject", "dproject").is_err());
tear_down("move_test.db");
}
fn setup_db(dbname: &str) -> Database {
let result = Database::new(format!("target{}{}", MAIN_SEPARATOR, dbname).as_str());
assert!(result.is_ok());
let db = result.unwrap();
assert!(db.add(vec!["aproject".to_owned(), "bproject".to_owned()]).is_ok());
assert_eq!(db.all().unwrap().len(), 2);
db
}
fn tear_down(dbname: &str) {
assert!(std::fs::remove_file(format!("target{}{}", MAIN_SEPARATOR, dbname).as_str()).is_ok());
}
}