use crate::errors::AsyncQueueError; use crate::runnable::RunnableTask; use crate::schema::backie_tasks; use crate::task::{NewTask, TaskId, TaskType, TaskHash}; use crate::task::Task; use chrono::DateTime; use chrono::Duration; use chrono::Utc; use diesel::prelude::*; use diesel::ExpressionMethods; use diesel_async::{pg::AsyncPgConnection, RunQueryDsl}; impl Task { pub(crate) async fn remove_all( connection: &mut AsyncPgConnection, ) -> Result { Ok(diesel::delete(backie_tasks::table) .execute(connection) .await? as u64) } pub(crate) async fn remove_all_scheduled( connection: &mut AsyncPgConnection, ) -> Result { let query = backie_tasks::table.filter(backie_tasks::running_at.is_null()); Ok(diesel::delete(query).execute(connection).await? as u64) } pub(crate) async fn remove( connection: &mut AsyncPgConnection, id: TaskId, ) -> Result { let query = backie_tasks::table.filter(backie_tasks::id.eq(id)); Ok(diesel::delete(query).execute(connection).await? as u64) } pub(crate) async fn remove_by_hash( connection: &mut AsyncPgConnection, task_hash: TaskHash, ) -> Result { let query = backie_tasks::table.filter(backie_tasks::uniq_hash.eq(task_hash)); let qty = diesel::delete(query).execute(connection).await?; Ok(qty > 0) } pub(crate) async fn remove_by_type( connection: &mut AsyncPgConnection, task_type: TaskType, ) -> Result { let query = backie_tasks::table.filter(backie_tasks::task_type.eq(task_type)); Ok(diesel::delete(query).execute(connection).await? as u64) } pub(crate) async fn find_by_id( connection: &mut AsyncPgConnection, id: TaskId, ) -> Result { let task = backie_tasks::table .filter(backie_tasks::id.eq(id)) .first::(connection) .await?; Ok(task) } pub(crate) async fn fail_with_message( connection: &mut AsyncPgConnection, id: TaskId, error_message: &str, ) -> Result { let query = backie_tasks::table.filter(backie_tasks::id.eq(id)); Ok(diesel::update(query) .set(( backie_tasks::error_message.eq(error_message), backie_tasks::done_at.eq(Utc::now()), )) .get_result::(connection) .await?) } pub(crate) async fn schedule_retry( connection: &mut AsyncPgConnection, id: TaskId, backoff_seconds: u32, error: &str, ) -> Result { use crate::schema::backie_tasks::dsl; let now = Utc::now(); let scheduled_at = now + Duration::seconds(backoff_seconds as i64); let task = diesel::update(backie_tasks::table.filter(backie_tasks::id.eq(id))) .set(( backie_tasks::error_message.eq(error), backie_tasks::retries.eq(dsl::retries + 1), backie_tasks::created_at.eq(scheduled_at), backie_tasks::running_at.eq::>>(None), )) .get_result::(connection) .await?; Ok(task) } pub(crate) async fn fetch_next_pending( connection: &mut AsyncPgConnection, task_type: TaskType, ) -> Option { backie_tasks::table .filter(backie_tasks::created_at.lt(Utc::now())) // skip tasks scheduled for the future .order(backie_tasks::created_at.asc()) // get the oldest task first .filter(backie_tasks::running_at.is_null()) // that is not marked as running already .filter(backie_tasks::done_at.is_null()) // and not marked as done .filter(backie_tasks::task_type.eq(task_type)) .limit(1) .for_update() .skip_locked() .get_result::(connection) .await .ok() } pub(crate) async fn set_running( connection: &mut AsyncPgConnection, task: Task, ) -> Result { Ok(diesel::update(&task) .set(( backie_tasks::running_at.eq(Utc::now()), )) .get_result::(connection) .await?) } pub(crate) async fn set_done( connection: &mut AsyncPgConnection, id: TaskId, ) -> Result { Ok(diesel::update(backie_tasks::table.filter(backie_tasks::id.eq(id))) .set(( backie_tasks::done_at.eq(Utc::now()), )) .get_result::(connection) .await?) } pub(crate) async fn insert( connection: &mut AsyncPgConnection, runnable: &dyn RunnableTask ) -> Result { let payload = serde_json::to_value(runnable)?; match runnable.uniq() { None => { let new_task = NewTask::builder() .uniq_hash(None) .task_type(runnable.task_type()) .payload(payload) .build(); Ok(diesel::insert_into(backie_tasks::table) .values(new_task) .get_result::(connection) .await?) } Some(hash) => { match Self::find_by_uniq_hash(connection, hash.clone()).await { Some(task) => Ok(task), None => { let new_task = NewTask::builder() .uniq_hash(Some(hash)) .task_type(runnable.task_type()) .payload(payload) .build(); Ok(diesel::insert_into(backie_tasks::table) .values(new_task) .get_result::(connection) .await?) } } } } } pub(crate) async fn find_by_uniq_hash( connection: &mut AsyncPgConnection, hash: TaskHash, ) -> Option { backie_tasks::table .filter(backie_tasks::uniq_hash.eq(hash)) .first::(connection) .await .ok() } }