use crate::runnable::AsyncRunnable; use crate::fang_task_state::FangTaskState; use crate::schema::fang_tasks; use crate::errors::CronError; use crate::Scheduled::*; use crate::task::{DEFAULT_TASK_TYPE, Task}; use async_trait::async_trait; use chrono::DateTime; use chrono::Duration; use chrono::Utc; use cron::Schedule; use diesel::prelude::*; use diesel::result::Error::QueryBuilderError; use diesel::ExpressionMethods; use diesel_async::scoped_futures::ScopedFutureExt; use diesel_async::AsyncConnection; use diesel_async::{pg::AsyncPgConnection, pooled_connection::bb8::Pool, pooled_connection::bb8::PooledConnection, RunQueryDsl}; use diesel_async::pooled_connection::PoolableConnection; use sha2::{Digest, Sha256}; use std::str::FromStr; use typed_builder::TypedBuilder; use uuid::Uuid; use crate::task::NewTask; use crate::errors::AsyncQueueError; impl Task { pub async fn remove_all_scheduled_tasks( connection: &mut AsyncPgConnection, ) -> Result { let query = fang_tasks::table.filter(fang_tasks::scheduled_at.gt(Utc::now())); Ok(diesel::delete(query).execute(connection).await? as u64) } pub async fn remove_task( connection: &mut AsyncPgConnection, id: Uuid, ) -> Result { let query = fang_tasks::table.filter(fang_tasks::id.eq(id)); Ok(diesel::delete(query).execute(connection).await? as u64) } pub async fn remove_task_by_metadata( connection: &mut AsyncPgConnection, task: &dyn AsyncRunnable, ) -> Result { let metadata = serde_json::to_value(task)?; let uniq_hash = Self::calculate_hash(metadata.to_string()); let query = fang_tasks::table.filter(fang_tasks::uniq_hash.eq(uniq_hash)); Ok(diesel::delete(query).execute(connection).await? as u64) } pub async fn remove_tasks_type( connection: &mut AsyncPgConnection, task_type: &str, ) -> Result { let query = fang_tasks::table.filter(fang_tasks::task_type.eq(task_type)); Ok(diesel::delete(query).execute(connection).await? as u64) } pub async fn find_task_by_id( connection: &mut AsyncPgConnection, id: Uuid, ) -> Result { let task = fang_tasks::table .filter(fang_tasks::id.eq(id)) .first::(connection) .await?; Ok(task) } pub async fn fail_task( connection: &mut AsyncPgConnection, task: Task, error_message: &str, ) -> Result { Ok(diesel::update(&task) .set(( fang_tasks::state.eq(FangTaskState::Failed), fang_tasks::error_message.eq(error_message), fang_tasks::updated_at.eq(Utc::now()), )) .get_result::(connection) .await?) } pub async fn schedule_retry( connection: &mut AsyncPgConnection, task: &Task, backoff_seconds: u32, error: &str, ) -> Result { let now = Utc::now(); let scheduled_at = now + Duration::seconds(backoff_seconds as i64); let task = diesel::update(task) .set(( fang_tasks::state.eq(FangTaskState::Retried), fang_tasks::error_message.eq(error), fang_tasks::retries.eq(task.retries + 1), fang_tasks::scheduled_at.eq(scheduled_at), fang_tasks::updated_at.eq(now), )) .get_result::(connection) .await?; Ok(task) } pub async fn fetch_task_of_type( connection: &mut AsyncPgConnection, task_type: Option, ) -> Option { fang_tasks::table .order(fang_tasks::created_at.asc()) .order(fang_tasks::scheduled_at.asc()) .limit(1) .filter(fang_tasks::scheduled_at.le(Utc::now())) .filter(fang_tasks::state.eq_any(vec![FangTaskState::New, FangTaskState::Retried])) .filter(fang_tasks::task_type.eq(task_type.unwrap_or_else(|| DEFAULT_TASK_TYPE.to_string()))) .for_update() .skip_locked() .get_result::(connection) .await .ok() } pub async fn update_task_state( connection: &mut AsyncPgConnection, task: Task, state: FangTaskState, ) -> Result { let updated_at = Utc::now(); Ok(diesel::update(&task) .set(( fang_tasks::state.eq(state), fang_tasks::updated_at.eq(updated_at), )) .get_result::(connection) .await?) } pub async fn insert_task( connection: &mut AsyncPgConnection, params: &dyn AsyncRunnable, scheduled_at: DateTime, ) -> Result { if !params.uniq() { let new_task = NewTask::builder() .scheduled_at(scheduled_at) .uniq_hash(None) .task_type(params.task_type()) .metadata(serde_json::to_value(params).unwrap()) .build(); Ok(diesel::insert_into(fang_tasks::table) .values(new_task) .get_result::(connection) .await?) } else { let metadata = serde_json::to_value(params).unwrap(); let uniq_hash = Self::calculate_hash(metadata.to_string()); match Self::find_task_by_uniq_hash(connection, &uniq_hash).await { Some(task) => Ok(task), None => { let new_task = NewTask::builder() .scheduled_at(scheduled_at) .uniq_hash(Some(uniq_hash)) .task_type(params.task_type()) .metadata(serde_json::to_value(params).unwrap()) .build(); Ok(diesel::insert_into(fang_tasks::table) .values(new_task) .get_result::(connection) .await?) } } } } fn calculate_hash(json: String) -> String { let mut hasher = Sha256::new(); hasher.update(json.as_bytes()); let result = hasher.finalize(); hex::encode(result) } pub async fn find_task_by_uniq_hash( connection: &mut AsyncPgConnection, uniq_hash: &str, ) -> Option { fang_tasks::table .filter(fang_tasks::uniq_hash.eq(uniq_hash)) .filter(fang_tasks::state.eq_any(vec![FangTaskState::New, FangTaskState::Retried])) .first::(connection) .await .ok() } }