use indicatif::ProgressBar;
use regex::Regex;
use rusqlite::{params, Connection, NO_PARAMS};
use std::convert::TryInto;
pub struct Database {
conn: rusqlite::Connection,
}
impl Database {
pub fn new(database_file: String) -> Self {
let conn = Connection::open(&database_file).expect(&format!("Failed to open database {}", database_file));
create_database(&conn);
Database { conn }
}
pub fn increment(&self, project: String) {
increment_project_ref_count(&self.conn, project);
}
pub fn add(&self, projects: Vec<String>) {
add_all_projects(&self.conn, projects);
}
pub fn add_new(&self, projects: Vec<String>) {
add_new_projects(&self.conn, projects);
}
pub fn find(&self, input: String) -> Vec<String> {
find(&self.conn, input)
}
pub fn all(&self) -> Vec<String> {
select_all_projects(&self.conn)
.iter()
.map(|p| p.0.clone())
.collect()
}
pub fn remove(&self, project: String) {
remove_project(&self.conn, project)
}
pub fn append(&self, project: String) {
add_project(&self.conn, project)
}
pub fn alias(&self, project: String, alias: String) {
add_alias(&self.conn, project, alias)
}
pub fn remove_alias(&self, project: String) {
remove_alias(&self.conn, project)
}
pub fn all_aliased(&self) -> 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_string(), p.0.clone()))
.collect();
aliases.sort_by(|a,b| a.0.cmp(&b.0));
aliases
}
}
fn create_database(conn: &rusqlite::Connection) {
match conn.execute(
"CREATE TABLE IF NOT EXISTS projects (name TEXT PRIMARY KEY, alias TEXT NULL, ref_count INTEGER)",
NO_PARAMS,
) {
Ok(created) => {
if created >= 1 {
println!("Projects table created.");
}
}
Err(e) => {
println!("Failed to create projects table, {}", e);
}
}
}
fn increment_project_ref_count(conn: &rusqlite::Connection, project: String) {
match conn.execute(
"UPDATE projects SET ref_count = ref_count + 1 WHERE name = ?",
&[project.clone()],
) {
Ok(nr_updated) => {
if nr_updated != 1 {
println!(
"Failed to update ref_count for project {}, project not found.",
project
);
}
}
Err(e) => {
println!("Failed to update ref_count for project {}, {}", project, e);
}
}
}
fn add_alias(conn: &rusqlite::Connection, project: String, alias: String) {
match conn.execute(
"UPDATE projects SET alias = ?1 WHERE name = ?2",
&[alias.clone(), project.clone()],
) {
Ok(nr_updated) => {
if nr_updated != 1 {
println!(
"Failed to set alias {} for project {}, project not found.",
alias,
project
);
}
}
Err(e) => println!("Failed to set alias {} for project {}, {}", alias, project, e),
}
}
fn remove_alias(conn: &rusqlite::Connection, project: String) {
match conn.execute(
"UPDATE projects SET alias = null WHERE name = ?",
&[project.clone()],
) {
Ok(nr_updated) => {
if nr_updated != 1 {
println!(
"Failed to remove alias for project {}, project not found.",
project
);
}
}
Err(e) => println!("Failed to remove alias for project {}, {}", project, e),
}
}
fn delete_all_projects(conn: &rusqlite::Connection) {
if let Err(e) = conn.execute("DELETE FROM projects", NO_PARAMS) {
println!("Failed to clean database, {}", e);
}
}
fn select_all_projects(conn: &rusqlite::Connection) -> Vec<(String, Option<String>)> {
let mut stmt = conn
.prepare("SELECT name, alias FROM projects ORDER BY ref_count DESC, alias, name")
.expect("Failed to query projects database");
let projects_iter = stmt
.query_map(params![], |row| match row.get(1) {
Ok(alias) => Ok((row.get(0).unwrap(), Some(alias))),
Err(_) => Ok((row.get(0).unwrap(), None)),
})
.unwrap();
let mut projects: Vec<(String, Option<String>)> = vec![];
for project in projects_iter {
projects.push(project.unwrap());
}
projects
}
fn add_all_projects(conn: &rusqlite::Connection, projects: Vec<String>) {
let progress_bar = ProgressBar::new(projects.len().try_into().unwrap());
progress_bar.println("Adding found projects to database");
delete_all_projects(conn);
let mut stmt = conn
.prepare("INSERT INTO projects(name, alias, ref_count) VALUES(?, null, 0)")
.expect("Failed to create insert statement.");
for project in projects {
match stmt.execute(&[project.clone()]) {
Ok(_) => {
progress_bar.set_message(&format!("Added {}", project));
}
Err(e) => {
progress_bar.set_message(&format!("Failed to add {}, {}", project, e));
}
};
progress_bar.inc(1);
}
progress_bar.finish_with_message("done.");
}
fn add_new_projects(conn: &rusqlite::Connection, projects: Vec<String>) {
let progress_bar = ProgressBar::new(projects.len().try_into().unwrap());
progress_bar.println("Adding found projects to database");
let mut stmt = conn
.prepare("INSERT INTO projects(name, alias, ref_count) VALUES(?, null, 0)")
.expect("Failed to create insert statement.");
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(_) => {
progress_bar.set_message(&format!("Skipping {}", project));
}
};
progress_bar.inc(1);
}
progress_bar.finish_with_message(&format!("done. Added {} new projects", added));
}
fn add_project(conn: &rusqlite::Connection, project: String) {
if let Err(e) = conn.execute(
"INSERT INTO projects(name, alias, ref_count) VALUES(?, null, 0)",
&[project],
) {
println!("Failed to add project, {}", e);
}
}
fn find(conn: &rusqlite::Connection, find: String) -> Vec<String> {
match Regex::new(&find) {
Ok(filter) => {
let mut projects: Vec<String> = vec![];
for project in select_all_projects(conn) {
if filter.is_match(&project.0)
|| (project.1.is_some() && filter.is_match(&project.1.unwrap()))
{
projects.push(project.0);
}
}
projects
}
Err(_) => vec![],
}
}
fn remove_project(conn: &rusqlite::Connection, project: String) {
if let Err(e) = conn.execute("DELETE FROM projects WHERE name =?", &[project]) {
println!("Failed to delete project, {}", e);
}
}