use super::TaskExecutionDAL;
use crate::dal::unified::models::UnifiedTaskExecution;
use crate::database::schema::unified::task_executions;
use crate::database::universal_types::UniversalUuid;
use crate::error::ValidationError;
use crate::models::task_execution::TaskExecution;
use diesel::prelude::*;
impl<'a> TaskExecutionDAL<'a> {
pub async fn get_pending_tasks(
&self,
workflow_execution_id: UniversalUuid,
) -> Result<Vec<TaskExecution>, ValidationError> {
let tasks: Vec<UnifiedTaskExecution> = crate::interact_on_backend!(self.dal, |conn| {
task_executions::table
.filter(task_executions::workflow_execution_id.eq(workflow_execution_id))
.filter(task_executions::status.eq("NotStarted"))
.load(conn)
})?;
Ok(tasks.into_iter().map(Into::into).collect())
}
pub async fn get_pending_tasks_batch(
&self,
workflow_execution_ids: Vec<UniversalUuid>,
) -> Result<Vec<TaskExecution>, ValidationError> {
if workflow_execution_ids.is_empty() {
return Ok(Vec::new());
}
let tasks: Vec<UnifiedTaskExecution> = crate::interact_on_backend!(self.dal, |conn| {
task_executions::table
.filter(task_executions::workflow_execution_id.eq_any(&workflow_execution_ids))
.filter(task_executions::status.eq_any(vec!["NotStarted", "Pending"]))
.load(conn)
})?;
Ok(tasks.into_iter().map(Into::into).collect())
}
pub async fn count_running_tasks(&self) -> Result<i64, ValidationError> {
let n: i64 = crate::interact_on_backend!(self.dal, |conn| {
task_executions::table
.filter(task_executions::status.eq("Running"))
.count()
.get_result(conn)
})?;
Ok(n)
}
pub async fn check_workflow_completion(
&self,
workflow_execution_id: UniversalUuid,
) -> Result<bool, ValidationError> {
let incomplete_count: i64 = crate::interact_on_backend!(self.dal, |conn| {
task_executions::table
.filter(task_executions::workflow_execution_id.eq(workflow_execution_id))
.filter(task_executions::status.ne_all(vec!["Completed", "Failed", "Skipped"]))
.count()
.get_result(conn)
})?;
Ok(incomplete_count == 0)
}
pub async fn get_task_status(
&self,
workflow_execution_id: UniversalUuid,
task_name: &str,
) -> Result<String, ValidationError> {
let task_name_owned = task_name.to_string();
let status: String = crate::interact_on_backend!(self.dal, |conn| {
task_executions::table
.filter(task_executions::workflow_execution_id.eq(workflow_execution_id))
.filter(task_executions::task_name.eq(&task_name_owned))
.select(task_executions::status)
.first(conn)
})?;
Ok(status)
}
pub async fn get_task_statuses_batch(
&self,
workflow_execution_id: UniversalUuid,
task_names: Vec<String>,
) -> Result<std::collections::HashMap<String, String>, ValidationError> {
use std::collections::HashMap;
if task_names.is_empty() {
return Ok(HashMap::new());
}
let results: Vec<(String, String)> = crate::interact_on_backend!(self.dal, |conn| {
task_executions::table
.filter(task_executions::workflow_execution_id.eq(workflow_execution_id))
.filter(task_executions::task_name.eq_any(&task_names))
.select((task_executions::task_name, task_executions::status))
.load(conn)
})?;
Ok(results.into_iter().collect())
}
pub async fn get_all_task_statuses_for_executions(
&self,
workflow_execution_ids: Vec<UniversalUuid>,
) -> Result<
std::collections::HashMap<UniversalUuid, std::collections::HashMap<String, String>>,
ValidationError,
> {
use std::collections::HashMap;
if workflow_execution_ids.is_empty() {
return Ok(HashMap::new());
}
let rows: Vec<(UniversalUuid, String, String)> =
crate::interact_on_backend!(self.dal, |conn| {
task_executions::table
.filter(task_executions::workflow_execution_id.eq_any(&workflow_execution_ids))
.select((
task_executions::workflow_execution_id,
task_executions::task_name,
task_executions::status,
))
.load(conn)
})?;
let mut grouped: HashMap<UniversalUuid, HashMap<String, String>> = HashMap::new();
for (exec_id, name, status) in rows {
grouped.entry(exec_id).or_default().insert(name, status);
}
Ok(grouped)
}
}