|
|
|
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<Mutex<UDB>>;
|
|
|
|
|
|
|
|
pub struct UDB {
|
|
|
|
conn: PgConnection
|
|
|
|
}
|
|
|
|
|
|
|
|
impl UDB {
|
|
|
|
pub fn new() -> USrvResult<Storage> {
|
|
|
|
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<Vec<Agent>> {
|
|
|
|
use schema::agents;
|
|
|
|
let result = agents::table
|
|
|
|
.load::<Agent>(&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()
|
|
|
|
)
|
|
|
|
}
|
|
|
|
}
|