unki/bin/u_server/src/db.rs

use diesel::{
    sqlite::SqliteConnection,
    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: SqliteConnection
}

impl UDB {
    pub fn new(path: Option<String>) -> USrvResult<Storage> {
        dotenv()?;
        let db_path = match path {
            Some(p) => p,
            None => env::var("DATABASE_URL").unwrap_or(":memory:".to_string())
        };
        let conn = SqliteConnection::establish(&db_path)?;
        conn.execute("PRAGMA foreign_keys = ON;")?;
        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(Some(":memory:".to_string())).unwrap();
    }

    #[tokio::test]
    async fn test_add_agent() {
        let db = setup_db();
        let agent = IAgent {
            alias: None,
            agent_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()
        )
    }
}