use diesel::{ pg::PgConnection, prelude::* }; use dotenv::dotenv; use std::{ env, sync::{Arc, Mutex} }; use crate::{ errors::USrvResult }; use u_lib::models::*; 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) } } #[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() ) } }