unki/bin/u_server/src/db.rs

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<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)
    }

    pub fn get_jobs(&self, uid: Option<Uuid>) -> USrvResult<Vec<JobMeta>> {
        use schema::jobs;
        let result = if uid.is_some() {
            jobs::table
                .filter(jobs::id.like(uid.unwrap()))
        } else {
            jobs::table
        }
            .load::<JobMeta>(&self.conn)?;
        Ok(result)
    }

    pub fn get_agent_jobs(&self, uid: Option<Uuid>) -> USrvResult<Vec<JobMeta>> {
        use schema::{results, jobs};
        let result = if uid.is_some() {
            jobs::table
                .filter(jobs::id.like(uid.unwrap()))
        } else {
            jobs::table
        }
            .load::<Agent>(&self.conn)?;
        Ok(result)
    }

    pub fn add_jobs(&self, jobs: &Vec<JobMeta>) -> 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()
        )
    }
}