|
|
|
@ -1,12 +1,14 @@ |
|
|
|
|
use crate::error::Error; |
|
|
|
|
use diesel::{pg::PgConnection, prelude::*, result::Error as DslError, Connection}; |
|
|
|
|
use std::collections::{HashMap, HashSet}; |
|
|
|
|
use std::mem::drop; |
|
|
|
|
use u_lib::{ |
|
|
|
|
db::PgAsyncPool, |
|
|
|
|
models::{schema, Agent, AssignedJob, Job, JobMeta, JobState, Payload}, |
|
|
|
|
models::{schema, Agent, AssignedJob, AssignedJobById, Job, JobMeta, JobState, Payload}, |
|
|
|
|
platform::Platform, |
|
|
|
|
types::Id, |
|
|
|
|
}; |
|
|
|
|
use uuid::Uuid; |
|
|
|
|
|
|
|
|
|
type Result<T> = std::result::Result<T, Error>; |
|
|
|
|
|
|
|
|
@ -215,39 +217,72 @@ impl UDB<'_> { |
|
|
|
|
Ok(result) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
pub fn set_jobs_for_agent(&mut self, agent_id: Id, job_ids: &[Id]) -> Result<()> { |
|
|
|
|
// todo: move to handlers
|
|
|
|
|
pub fn assign_jobs(&mut self, assigned_jobs: &[AssignedJobById]) -> Result<()> { |
|
|
|
|
use schema::{jobs, results}; |
|
|
|
|
|
|
|
|
|
let agent_platform = match self.get_agent(agent_id)? { |
|
|
|
|
Some(agent) => Platform::new(&agent.platform), |
|
|
|
|
None => { |
|
|
|
|
return Err(Error::ProcessingError(format!( |
|
|
|
|
"Agent {agent_id} not found" |
|
|
|
|
))) |
|
|
|
|
} |
|
|
|
|
}; |
|
|
|
|
struct JobBriefMeta { |
|
|
|
|
alias: Option<String>, |
|
|
|
|
target_platform: String, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
let assigned_job_ids = HashSet::<Uuid>::from_iter(assigned_jobs.iter().map(|a| a.job_id)); |
|
|
|
|
|
|
|
|
|
let jobs_meta = HashMap::<Id, JobBriefMeta>::from_iter( |
|
|
|
|
jobs::table |
|
|
|
|
.select((jobs::id, jobs::alias, jobs::target_platforms)) |
|
|
|
|
.filter(jobs::id.eq_any(&assigned_job_ids)) |
|
|
|
|
.load::<(Id, Option<String>, String)>(self.conn) |
|
|
|
|
.map_err(with_err_ctx(format!( |
|
|
|
|
"Can't find jobs {:?}", |
|
|
|
|
assigned_job_ids |
|
|
|
|
)))? |
|
|
|
|
.into_iter() |
|
|
|
|
.map(|(id, alias, target_platform)| { |
|
|
|
|
( |
|
|
|
|
id, |
|
|
|
|
JobBriefMeta { |
|
|
|
|
alias, |
|
|
|
|
target_platform, |
|
|
|
|
}, |
|
|
|
|
) |
|
|
|
|
}), |
|
|
|
|
); |
|
|
|
|
|
|
|
|
|
let existing_job_ids = HashSet::from_iter(jobs_meta.keys().copied()); |
|
|
|
|
|
|
|
|
|
let jobs_meta = jobs::table |
|
|
|
|
.select((jobs::id, jobs::alias, jobs::target_platforms)) |
|
|
|
|
.filter(jobs::id.eq_any(job_ids)) |
|
|
|
|
.load::<(Id, Option<String>, String)>(self.conn) |
|
|
|
|
.map_err(with_err_ctx(format!("Can't find jobs {job_ids:?}")))?; |
|
|
|
|
if assigned_job_ids != existing_job_ids { |
|
|
|
|
return Err(Error::ProcessingError(format!( |
|
|
|
|
"Jobs not found: {:?}", |
|
|
|
|
assigned_job_ids.difference(&existing_job_ids), |
|
|
|
|
))); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
for meta in &jobs_meta { |
|
|
|
|
if !agent_platform.matches(&meta.2) { |
|
|
|
|
for ajob in assigned_jobs { |
|
|
|
|
let meta = &jobs_meta[&ajob.job_id]; |
|
|
|
|
let agent_platform = match self.get_agent(ajob.agent_id)? { |
|
|
|
|
Some(agent) => Platform::new(&agent.platform), |
|
|
|
|
None => { |
|
|
|
|
return Err(Error::ProcessingError(format!( |
|
|
|
|
"Agent {} not found", |
|
|
|
|
ajob.agent_id |
|
|
|
|
))) |
|
|
|
|
} |
|
|
|
|
}; |
|
|
|
|
if !agent_platform.matches(&meta.target_platform) { |
|
|
|
|
return Err(Error::InsuitablePlatform( |
|
|
|
|
agent_platform.into_string(), |
|
|
|
|
meta.2.clone(), |
|
|
|
|
meta.target_platform.clone(), |
|
|
|
|
)); |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
let job_requests = jobs_meta |
|
|
|
|
let job_requests = assigned_jobs |
|
|
|
|
.into_iter() |
|
|
|
|
.map(|(job_id, alias, _)| AssignedJob { |
|
|
|
|
job_id, |
|
|
|
|
agent_id, |
|
|
|
|
alias, |
|
|
|
|
.map(|a| AssignedJob { |
|
|
|
|
job_id: a.job_id, |
|
|
|
|
agent_id: a.agent_id, |
|
|
|
|
alias: jobs_meta[&a.job_id].alias.clone(), |
|
|
|
|
..Default::default() |
|
|
|
|
}) |
|
|
|
|
.collect::<Vec<AssignedJob>>(); |
|
|
|
@ -256,9 +291,7 @@ impl UDB<'_> { |
|
|
|
|
.values(&job_requests) |
|
|
|
|
.execute(self.conn) |
|
|
|
|
.map(drop) |
|
|
|
|
.map_err(with_err_ctx(format!( |
|
|
|
|
"Can't setup jobs {job_ids:?} for agent {agent_id:?}" |
|
|
|
|
))) |
|
|
|
|
.map_err(with_err_ctx("Can't assign jobs")) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
pub fn del_jobs(&mut self, ids: &[Id]) -> Result<()> { |
|
|
|
|