use diesel::{ pg::PgConnection, prelude::*, result::Error as DslError }; use dotenv::dotenv; use std::{ env, sync::{Arc, Mutex} }; use crate::{ errors::* }; use u_lib::{ models::* }; use uuid::Uuid; pub type Storage = Arc>; pub struct UDB { pub conn: PgConnection } impl UDB { pub fn new() -> USrvResult { dotenv().unwrap(); let db_path = env::var("DATABASE_URL").unwrap(); let conn = PgConnection::establish(&db_path).unwrap(); let instance = UDB { conn }; Ok(Arc::new(Mutex::new(instance))) } pub fn new_agent(&self, agent: &IAgent) -> USrvResult<()> { use schema::agents; diesel::insert_into(agents::table) .values(agent) .execute(&self.conn)?; Ok(()) } pub fn get_agents(&self, uid: Option) -> USrvResult> { use schema::agents; let result = if uid.is_some() { agents::table .filter(agents::id.eq(uid.unwrap())) .load::(&self.conn)? } else { agents::table .load::(&self.conn)? }; Ok(result) } pub fn get_jobs(&self, uid: Option) -> USrvResult> { use schema::jobs; let result = if uid.is_some() { jobs::table .filter(jobs::id.eq(uid.unwrap())) .get_results::(&self.conn)? } else { jobs::table .load::(&self.conn)? }; Ok(result) } //TODO: belongs_to pub fn get_agent_jobs(&self, uid: Option, personal: bool) -> USrvResult> { use schema::{results, jobs}; let mut q = results::table .inner_join(jobs::table) .into_boxed(); if uid.is_some() { q = q.filter(results::agent_id.eq(uid.unwrap())) } if personal { q = q.filter(results::state.eq(JobState::Queued)) } let result = q.select( (jobs::alias, jobs::id, jobs::exec_type, jobs::platform, jobs::payload) ) .get_results::(&self.conn)?; Ok(result) } pub fn get_results(&self, uid: Option) -> USrvResult> { use schema::results; let result = if uid.is_some() { results::table .filter(results::agent_id.eq(uid.unwrap())) .or_filter(results::job_id.eq(uid.unwrap())) .or_filter(results::id.eq(uid.unwrap())) .load::(&self.conn)? } else { results::table .load::(&self.conn)? }; Ok(result) } pub fn add_jobs(&self, job_metas: &Vec) -> USrvResult<()> { use schema::jobs; diesel::insert_into(jobs::table) .values(job_metas) .execute(&self.conn)?; Ok(()) } pub fn set_jobs_for_agent(&self, agent_uid: &Uuid, job_uids: &Vec) -> USrvResult<()> { use schema::{agents::dsl::agents, jobs::dsl::jobs, results}; if let Err(DslError::NotFound) = agents.find(agent_uid).first::(&self.conn) { return Err(USrvError::NotFound(agent_uid.to_string())) } let not_found_jobs = job_uids.iter().filter_map(|job_uid| { if let Err(DslError::NotFound) = jobs.find(job_uid).first::(&self.conn) { Some(job_uid.to_string()) } else { None } }).collect::>(); if not_found_jobs.len() > 0 { return Err(USrvError::NotFound(not_found_jobs.join(", "))); } let job_requests = job_uids.iter().map(|job_uid| { JobResult { job_id: *job_uid, agent_id: *agent_uid, ..Default::default() } }).collect::>(); diesel::insert_into(results::table) .values(&job_requests) .execute(&self.conn)?; Ok(()) } pub fn del_jobs(&self, uids: &Vec) -> USrvResult { use schema::jobs; let mut affected = 0; for &uid in uids { let deleted = diesel::delete(jobs::table) .filter(jobs::id.eq(uid)) .execute(&self.conn)?; affected += deleted; } Ok(affected) } pub fn del_agents(&self, uids: &Vec) -> USrvResult { use schema::agents; let mut affected = 0; for &uid in uids { let deleted = diesel::delete(agents::table) .filter(agents::id.eq(uid)) .execute(&self.conn)?; affected += deleted; } Ok(affected) } pub fn del_results(&self, uids: &Vec) -> USrvResult { use schema::results; let mut affected = 0; for &uid in uids { let deleted = diesel::delete(results::table) .filter(results::id.eq(uid)) .execute(&self.conn)?; affected += deleted; } Ok(affected) } } /* #[cfg(test)] mod tests { use super::*; fn setup_db() -> Storage { return UDB::new().unwrap(); } #[tokio::test] async fn test_add_agent() { let db = setup_db(); let agent = IAgent { alias: None, id: "000-000".to_string(), hostname: "test".to_string(), is_root: false, is_root_allowed: false, platform: "linux".to_string(), status: None, token: None, username: "test".to_string() }; db.lock().unwrap().new_agent(agent).unwrap(); let result = db.lock().unwrap().get_agents().unwrap(); assert_eq!( result[0].username, "test".to_string() ) } } */