use diesel::{ pg::PgConnection, prelude::* }; use dotenv::dotenv; use std::{ env, sync::{Arc, Mutex} }; use crate::{ errors::USrvResult }; use u_lib::models::*; use uuid::Uuid; pub type Storage = Arc>; pub struct UDB { conn: PgConnection } impl UDB { pub fn new() -> USrvResult { dotenv()?; let db_path = env::var("DATABASE_URL").unwrap(); let conn = PgConnection::establish(&db_path)?; 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) -> USrvResult> { use schema::agents; let result = 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.like(uid.unwrap())) } else { jobs::table } .load::(&self.conn)?; Ok(result) } pub fn get_agent_jobs(&self, uid: Option) -> USrvResult> { use schema::{results, jobs}; let result = if uid.is_some() { jobs::table .filter(jobs::id.like(uid.unwrap())) } else { jobs::table } .load::(&self.conn)?; Ok(result) } pub fn add_jobs(&self, jobs: &Vec) -> USrvResult<()> { use schema::jobs; diesel::insert_into(jobs::table) .values(jobs) .execute(&self.conn)?; Ok(()) } } #[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() ) } }