1use anyhow::Result;
2use indicatif::ProgressBar;
3use regex::Regex;
4use rusqlite::{params, Connection};
5use std::convert::TryInto;
6
7use crate::error::GcdError;
8
9pub struct Database {
10 conn: rusqlite::Connection,
11}
12
13impl Database {
14 pub fn new(database_file: &str) -> Result<Self> {
15 let conn = Connection::open(database_file)?;
16 create_database(&conn)?;
17 Ok(Database { conn })
18 }
19
20 pub fn increment(&self, project: String) -> Result<()> {
21 increment_project_ref_count(&self.conn, project)
22 }
23
24 pub fn add(&self, projects: Vec<String>) -> Result<()> {
25 add_all_projects(&self.conn, projects)
26 }
27 pub fn add_new(&self, projects: Vec<String>) -> Result<()> {
28 add_new_projects(&self.conn, projects)
29 }
30
31 pub fn find(&self, input: &str) -> Result<Vec<String>> {
32 find(&self.conn, input)
33 }
34
35 pub fn all(&self) -> Result<Vec<String>> {
36 Ok(select_all_projects(&self.conn)?
37 .iter()
38 .map(|p| p.0.clone())
39 .collect())
40 }
41 pub fn remove(&self, project: String) -> Result<()> {
42 remove_project(&self.conn, project)
43 }
44
45 pub fn append(&self, project: String) -> Result<()> {
46 add_project(&self.conn, project)
47 }
48
49 pub fn alias(&self, project: &str, alias: &str) -> Result<()> {
50 add_alias(&self.conn, project, alias)
51 }
52
53 pub fn remove_alias(&self, project: &str) -> Result<()> {
54 remove_alias(&self.conn, project)
55 }
56
57 pub fn remove_alias_by_alias(&self, alias: &str) -> Result<()> {
58 remove_alias_by_alias(&self.conn, alias)
59 }
60
61 pub fn all_aliased(&self) -> Result<Vec<(String, String)>> {
62 let mut aliases: Vec<(String, String)> = select_all_projects(&self.conn)?
63 .iter()
64 .filter(|p| p.1.is_some())
65 .map(|p| (p.1.as_ref().unwrap().to_owned(), p.0.clone()))
66 .collect();
67
68 aliases.sort_by(|a, b| a.0.cmp(&b.0));
69
70 Ok(aliases)
71 }
72
73 pub fn move_project(&self, from_location: &str, to_location: &str) -> Result<()> {
74 move_project(&self.conn, from_location, to_location)
75 }
76}
77
78fn create_database(conn: &rusqlite::Connection) -> Result<()> {
79 conn.execute(
80 "CREATE TABLE IF NOT EXISTS projects (name TEXT PRIMARY KEY, alias TEXT NULL, ref_count INTEGER)",
81 [],
82 )?;
83 Ok(())
84}
85
86fn increment_project_ref_count(conn: &rusqlite::Connection, project: String) -> Result<()> {
87 let nr_updated = conn
88 .execute(
89 "UPDATE projects SET ref_count = ref_count + 1 WHERE name = ?",
90 [project.clone()],
91 )?;
92
93 if nr_updated != 1 {
94 return Err(GcdError::new(format!(
95 "Failed to update ref_count for project {}, project not found.",
96 project
97 ))
98 .into());
99 }
100 Ok(())
101}
102
103fn add_alias(conn: &rusqlite::Connection, project: &str, alias: &str) -> Result<()> {
104 let nr_updated = conn
105 .execute(
106 "UPDATE projects SET alias = ?1 WHERE name = ?2",
107 [alias, project],
108 )?;
109
110 if nr_updated != 1 {
111 return Err(GcdError::new(format!(
112 "Failed to set alias {} for project {}, project not found.",
113 alias, project
114 ))
115 .into());
116 }
117
118 Ok(())
119}
120
121fn remove_alias(conn: &rusqlite::Connection, project: &str) -> Result<()> {
122 let nr_updated = conn
123 .execute(
124 "UPDATE projects SET alias = null WHERE name = ?",
125 [project],
126 )?;
127
128 if nr_updated == 0 {
129 return Err(GcdError::new(format!(
130 "Failed to remove alias for project {}, project not found.",
131 project
132 ))
133 .into());
134 }
135 Ok(())
136}
137
138fn remove_alias_by_alias(conn: &rusqlite::Connection, alias: &str) -> Result<()> {
139 let nr_updated = conn
140 .execute(
141 "UPDATE projects SET alias = null WHERE alias = ?",
142 [alias],
143 )?;
144 if nr_updated == 0 {
145 return Err(GcdError::new(format!(
146 "Failed to remove alias {}, alias not found.",
147 alias
148 ))
149 .into());
150 }
151 Ok(())
152}
153
154fn delete_all_projects(conn: &rusqlite::Connection) -> Result<()> {
155 conn.execute("DELETE FROM projects", [])?;
156 Ok(())
157}
158
159fn select_all_projects(conn: &rusqlite::Connection) -> Result<Vec<(String, Option<String>)>> {
160 let mut stmt = conn
161 .prepare("SELECT name, alias FROM projects ORDER BY ref_count DESC, alias, name")?;
162 let projects_iter = stmt.query_map(params![], |row| match row.get(1) {
163 Ok(alias) => Ok((row.get(0)?, Some(alias))),
164 Err(_) => Ok((row.get(0)?, None)),
165 })?;
166 let mut projects: Vec<(String, Option<String>)> = vec![];
167 for project in projects_iter {
168 projects.push(project?);
169 }
170 Ok(projects)
171}
172
173fn add_all_projects(conn: &rusqlite::Connection, projects: Vec<String>) -> Result<()> {
174 delete_all_projects(conn)?;
175 add_new_projects(conn, projects)?;
176 Ok(())
177}
178
179fn add_new_projects(conn: &rusqlite::Connection, projects: Vec<String>) -> Result<()> {
180 let progress_bar = ProgressBar::new(projects.len().try_into()?);
181 progress_bar.println("Adding found projects to database");
182 let mut stmt = conn
183 .prepare("INSERT INTO projects(name, alias, ref_count) VALUES(?, null, 0)")?;
184 let mut added: usize = 0;
185 for project in projects {
186 match stmt.execute([project.clone()]) {
187 Ok(_) => {
188 progress_bar.set_message(format!("Added {}", project));
189 added += 1;
190 }
191 Err(e) => {
192 progress_bar.set_message(format!("Skipping {} - {}", project, e));
193 }
194 };
195 progress_bar.inc(1);
196 }
197 progress_bar.finish_with_message(format!("done. Added {} new projects", added));
198 Ok(())
199}
200
201fn add_project(conn: &rusqlite::Connection, project: String) -> Result<()> {
202 conn.execute(
203 "INSERT INTO projects(name, alias, ref_count) VALUES(?, null, 0)",
204 [project.clone()],
205 )?;
206
207 Ok(())
208}
209
210fn find(conn: &rusqlite::Connection, find: &str) -> Result<Vec<String>> {
211 match Regex::new(find) {
212 Ok(filter) => {
213 let mut projects: Vec<String> = vec![];
214 for (name, alias) in select_all_projects(conn)? {
215 if filter.is_match(&name) || (alias.is_some() && filter.is_match(&alias.unwrap())) {
216 projects.push(name);
217 }
218 }
219 Ok(projects)
220 }
221 Err(_) => Ok(vec![]),
222 }
223}
224
225fn remove_project(conn: &rusqlite::Connection, project: String) -> Result<()> {
226 conn.execute("DELETE FROM projects WHERE name =?", [project.clone()])?;
227
228 Ok(())
229}
230
231fn move_project(
232 conn: &rusqlite::Connection,
233 from_location: &str,
234 to_location: &str,
235) -> Result<()> {
236 let nr_updated = conn
237 .execute(
238 "UPDATE projects SET name = ?1 WHERE name = ?2",
239 [to_location, from_location],
240 )?;
241
242 if nr_updated != 1 {
243 return Err(GcdError::new(format!(
244 "Failed to move project {} to {}, project not found.",
245 from_location, to_location
246 ))
247 .into());
248 }
249 Ok(())
250}
251
252#[cfg(test)]
253mod test {
254 use super::*;
255 use std::path::MAIN_SEPARATOR;
256
257 #[test]
258 fn open_empty_db() {
259 let result = Database::new("open_test.db");
260 assert!(result.is_ok());
261 assert!(std::fs::remove_file("open_test.db").is_ok());
262 }
263 #[test]
264 fn update_non_existing_database() {
265 let db = setup_db("non_err_test.db");
266 tear_down("non_err_test.db");
267 assert!(db.increment("project".to_owned()).is_err());
268 }
269
270
271 #[test]
272 fn increment_non_existing_project() {
273 let db = setup_db("increment_err_test.db");
274 assert!(db.increment("project".to_owned()).is_err());
275 tear_down("increment_err_test.db");
276
277 }
278
279 #[test]
280 fn increment_existing_project() {
281 let db = setup_db("increment_test.db");
282 assert!(db.increment("aproject".to_owned()).is_ok());
283 tear_down("increment_test.db");
284
285 }
286
287 #[test]
288 fn add_existing_project() {
289 let db = setup_db("add_exsisting_test.db");
290 assert!(db.add(vec!["aproject".to_owned()]).is_ok());
291 tear_down("add_exsisting_test.db");
292
293 }
294
295 #[test]
296 fn add_new_existing_project() {
297 let db = setup_db("add_new_exsisting_test.db");
298 assert!(db.add_new(vec!["cproject".to_owned()]).is_ok());
299 tear_down("add_new_exsisting_test.db");
300
301 }
302
303 #[test]
304 fn find_project() {
305 let db = setup_db("find_test.db");
306 let projects = db.find("apro").unwrap();
307 assert!(projects.iter().any(|p| p == "aproject"));
308 tear_down("find_test.db");
309 }
310
311 #[test]
312 fn find_project_by_alias() {
313 let db = setup_db("find_by_alias_test.db");
314 assert!(db.alias("aproject", "a-alias").is_ok());
315 let projects = db.find("a-alias").unwrap();
316 assert!(projects.iter().any(|p| p == "aproject"));
317 tear_down("find_by_alias_test.db");
318 }
319
320
321 #[test]
322 fn alias() {
323 let db = setup_db("alias_test.db");
324
325 assert!(db.alias("aproject", "a-alias").is_ok());
326 assert!(db.alias("bproject", "b-alias").is_ok());
327 assert!(db.alias("cproject", "c-alias").is_err());
328
329 assert!(db.remove_alias("bproject").is_ok());
330 assert!(db.remove_alias("dproject").is_err());
331
332 assert!(db.remove_alias_by_alias("a-alias").is_ok());
333 assert!(db.remove_alias_by_alias("d-alias").is_err());
334
335 tear_down("alias_test.db");
336
337 }
338 #[test]
339 fn move_project() {
340 let db = setup_db("move_test.db");
341
342 assert!(db.move_project("aproject", "cproject").is_ok());
343 assert!(db.move_project("aproject", "dproject").is_err());
344
345 tear_down("move_test.db");
346
347 }
348
349 fn setup_db(dbname: &str) -> Database {
350 let result = Database::new(format!("target{}{}", MAIN_SEPARATOR, dbname).as_str());
351 assert!(result.is_ok());
352 let db = result.unwrap();
353 assert!(db.add(vec!["aproject".to_owned(), "bproject".to_owned()]).is_ok());
354 assert_eq!(db.all().unwrap().len(), 2);
355
356 db
357 }
358
359 fn tear_down(dbname: &str) {
360 assert!(std::fs::remove_file(format!("target{}{}", MAIN_SEPARATOR, dbname).as_str()).is_ok());
361 }
362}