|
|
|
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 update_job_status(&self, uid: Uuid, status: JobState) -> USrvResult<()> {
|
|
|
|
use schema::results;
|
|
|
|
diesel::update(results::table)
|
|
|
|
.filter(results::id.eq(uid))
|
|
|
|
.set(results::state.eq(status))
|
|
|
|
.execute(&self.conn)?;
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn get_agent_jobs(&self, uid: Option<Uuid>, personal: bool) -> USrvResult<Vec<ExactJob>> {
|
|
|
|
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::<ExactJob>(&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| {
|
|
|
|
ExactJob {
|
|
|
|
job_id: *job_uid,
|
|
|
|
agent_id: *agent_uid,
|
|
|
|
..Default::default()
|
|
|
|
}
|
|
|
|
}).collect::<Vec<ExactJob>>();
|
|
|
|
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()
|
|
|
|
)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
*/
|