unki/bin/u_server/src/db.rs

use diesel::{
    pg::PgConnection,
    prelude::*,
    result::Error as DslError
};
use dotenv::dotenv;
use std::{
    env,
    sync::{Arc, Mutex}
};
use crate::{
    errors::*
};
use u_lib::{
    models::*
};
use uuid::Uuid;

pub type Storage = Arc<Mutex<UDB>>;

pub struct UDB {
    pub conn: PgConnection
}

impl UDB {
    pub fn new() -> USrvResult<Storage> {
        dotenv().unwrap();
        let db_path = env::var("DATABASE_URL").unwrap();
        let conn = PgConnection::establish(&db_path).unwrap();
        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, uid: Option<Uuid>) -> USrvResult<Vec<Agent>> {
        use schema::agents;
        let result = if uid.is_some() {
            agents::table
                .filter(agents::id.eq(uid.unwrap()))
                .load::<Agent>(&self.conn)?
        } else {
            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.eq(uid.unwrap()))
                .get_results::<JobMeta>(&self.conn)?
        } else {
            jobs::table
                .load::<JobMeta>(&self.conn)?
        };
        Ok(result)
    }

    pub fn get_agent_jobs(&self, uid: Option<Uuid>, personal: bool) -> USrvResult<Vec<JobResult>> {
        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::<JobResult>(&self.conn)?;
        Ok(result)
    }

    pub fn add_jobs(&self, job_metas: &Vec<JobMeta>) -> USrvResult<()> {
        use schema::jobs;
        diesel::insert_into(jobs::table)
            .values(job_metas)
            .execute(&self.conn)?;
        Ok(())
    }

    pub fn set_jobs_for_agent(&self, agent_uid: &Uuid, job_uids: &Vec<Uuid>) -> USrvResult<()> {
        use schema::{agents::dsl::agents, jobs::dsl::jobs, results};
        if let Err(DslError::NotFound) = agents.find(agent_uid).first::<Agent>(&self.conn) {
            return Err(USrvError::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::<JobMeta>(&self.conn) {
                Some(job_uid.to_string())
            } else { None }
        }).collect::<Vec<String>>();
        if not_found_jobs.len() > 0 {
            return Err(USrvError::NotFound(not_found_jobs.join(", ")));
        }
        let job_requests = job_uids.iter().map(|job_uid| {
            JobResult {
                job_id: *job_uid,
                agent_id: *agent_uid,
                ..Default::default()
            }
        }).collect::<Vec<JobResult>>();
        diesel::insert_into(results::table)
            .values(&job_requests)
            .execute(&self.conn)?;
        Ok(())
    }

    pub fn del_jobs(&self, uids: &Vec<Uuid>) -> USrvResult<usize> {
        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_agents(&self, uids: &Vec<Uuid>) -> USrvResult<usize> {
        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)
    }

    pub fn del_results(&self, uids: &Vec<Uuid>) -> USrvResult<usize> {
        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)
    }
}
/*
#[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()
        )
    }
}
*/