use diesel::{pg::PgConnection, prelude::*, result::Error as DslError}; use dotenv::dotenv; use once_cell::sync::OnceCell; use std::{ env, sync::{Arc, Mutex, MutexGuard}, }; use u_lib::{ models::{schema, Agent, ExactJob, IAgent, JobMeta, JobState}, ULocalError, ULocalResult, }; use uuid::Uuid; pub struct UDB { pub conn: PgConnection, } static DB: OnceCell>> = OnceCell::new(); pub fn lock_db() -> MutexGuard<'static, UDB> { DB.get_or_init(|| { dotenv().unwrap(); let db_path = env::var("DATABASE_URL").unwrap(); let conn = PgConnection::establish(&db_path).unwrap(); let instance = UDB { conn }; Arc::new(Mutex::new(instance)) }) .lock() .unwrap() } impl UDB { pub fn insert_jobs(&self, job_metas: &Vec) -> ULocalResult<()> { use schema::jobs; diesel::insert_into(jobs::table) .values(job_metas) .execute(&self.conn)?; Ok(()) } pub fn get_jobs(&self, uid: Option) -> ULocalResult> { 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) } pub fn insert_agents(&self, agent: &IAgent) -> ULocalResult<()> { use schema::agents; diesel::insert_into(agents::table) .values(agent) .execute(&self.conn)?; Ok(()) } pub fn get_agents(&self, uid: Option) -> ULocalResult> { 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 update_job_status(&self, uid: Uuid, status: JobState) -> ULocalResult<()> { use schema::results; diesel::update(results::table) .filter(results::id.eq(uid)) .set(results::state.eq(status)) .execute(&self.conn)?; Ok(()) } //TODO: filters possibly could work in a wrong way, check pub fn get_exact_jobs(&self, uid: Option, personal: bool) -> ULocalResult> { use schema::results; let mut q = results::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) .and(results::agent_id.eq(uid.unwrap())), ) } else if uid.is_some() { q = q .filter(results::agent_id.eq(uid.unwrap())) .or_filter(results::job_id.eq(uid.unwrap())) .or_filter(results::id.eq(uid.unwrap())) } let result = q.load::(&self.conn)?; Ok(result) } pub fn set_jobs_for_agent(&self, agent_uid: &Uuid, job_uids: &Vec) -> ULocalResult<()> { use schema::{agents::dsl::agents, jobs::dsl::jobs, results}; if let Err(DslError::NotFound) = agents.find(agent_uid).first::(&self.conn) { return Err(ULocalError::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(ULocalError::NotFound(not_found_jobs.join(", "))); } let job_requests = job_uids .iter() .map(|job_uid| ExactJob { 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) -> ULocalResult { 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_results(&self, uids: &Vec) -> ULocalResult { 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) } pub fn del_agents(&self, uids: &Vec) -> ULocalResult { 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) } } /* #[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() ) } } */