fix: handle attempts in storages

packages/apalis-core/src/layers.rs

@@ -1,7 +1,9 @@
+use crate::task::attempt::Attempt;
 use crate::{request::Request, worker::WorkerId};
 use futures::channel::mpsc::{SendError, Sender};
 use futures::SinkExt;
 use futures::{future::BoxFuture, Future, FutureExt};
+use serde::{Deserialize, Serialize};
 use std::marker::PhantomData;
 use std::{fmt, sync::Arc};
 pub use tower::{
@@ -168,24 +170,16 @@ pub trait Ack<Task> {
 }
 
 /// ACK response
-#[derive(Debug)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct AckResponse<A> {
     /// The worker id
     pub worker: WorkerId,
     /// The acknowledger
     pub acknowledger: A,
     /// The stringified result
     pub result: Result<String, String>,
-}
-
-impl<A: fmt::Display> AckResponse<A> {
-    /// Output a json for the response
-    pub fn to_json(&self) -> String {
-        format!(
-            r#"{{"worker": "{}", "acknowledger": "{}", "result": "{:?}"}}"#,
-            self.worker, self.acknowledger, self.result
-        )
-    }
+    /// The number of attempts made by the request
+    pub attempts: Attempt,
 }
 
 /// A generic stream that emits (worker_id, task_id)
@@ -286,6 +280,8 @@ where
         let mut ack = self.ack.clone();
         let worker_id = self.worker_id.clone();
         let data = request.get::<<A as Ack<T>>::Acknowledger>().cloned();
+        let attempts = request.get::<Attempt>().cloned().unwrap_or_default();
+
         let fut = self.service.call(request);
         let fut_with_ack = async move {
             let res = fut.await;
@@ -299,6 +295,7 @@ where
                         worker: worker_id,
                         acknowledger: task_id,
                         result,
+                        attempts,
                     })
                     .await
                 {

packages/apalis-redis/src/storage.rs

@@ -1082,6 +1082,7 @@ mod tests {
                 acknowledger: job_id.clone(),
                 result: Ok("Success".to_string()),
                 worker: worker_id.clone(),
+                attempts: Attempt::new_with_value(0)
             })
             .await
             .expect("failed to acknowledge the job");

packages/apalis-sql/Cargo.toml

@@ -35,6 +35,8 @@ async-stream = "0.3.5"
 tokio = { version = "1", features = ["rt", "net"], optional = true }
 futures-lite = "2.3.0"
 async-std = { version = "1.12.0", optional = true }
+chrono = { version = "0.4", features = ["serde"] }
+
 
 [dev-dependencies]
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }

packages/apalis-sql/src/context.rs

@@ -2,7 +2,7 @@ use apalis_core::error::Error;
 use apalis_core::task::{attempt::Attempt, task_id::TaskId};
 use apalis_core::worker::WorkerId;
 use serde::{Deserialize, Serialize};
-use sqlx::types::chrono::{DateTime, Utc};
+use chrono::{DateTime, Utc};
 use std::{fmt, str::FromStr};
 
 /// The context for a job is represented here

packages/apalis-sql/src/from_row.rs

@@ -70,7 +70,7 @@ impl<'r, T: Decode<'r, sqlx::Sqlite> + Type<sqlx::Sqlite>>
     sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> for SqlRequest<T>
 {
     fn from_row(row: &'r sqlx::sqlite::SqliteRow) -> Result<Self, sqlx::Error> {
-        use sqlx::types::chrono::DateTime;
+        use chrono::DateTime;
         use sqlx::Row;
         use std::str::FromStr;
 

packages/apalis-sql/src/mysql.rs

@@ -17,7 +17,7 @@ use log::error;
 use serde::{de::DeserializeOwned, Serialize};
 use serde_json::Value;
 use sqlx::mysql::MySqlRow;
-use sqlx::types::chrono::{DateTime, Utc};
+use chrono::{DateTime, Utc};
 use sqlx::{MySql, Pool, Row};
 use std::any::type_name;
 use std::convert::TryInto;
@@ -396,11 +396,12 @@ impl<T: Serialize + DeserializeOwned + Sync + Send + Unpin + 'static> Backend<Re
                 .await
             {
                 for id in ids {
-                    let query = "UPDATE jobs SET status = ?, done_at = now(), last_error = ? WHERE id = ? AND lock_by = ?";
+                    let query = "UPDATE jobs SET status = ?, done_at = now(), last_error = ?, attempts = ? WHERE id = ? AND lock_by = ?";
                     let query = sqlx::query(query);
                     let query = query
                         .bind(calculate_status(&id.result).to_string())
                         .bind(serde_json::to_string(&id.result).unwrap())
+                        .bind(id.attempts.current() as u64)
                         .bind(id.acknowledger.to_string())
                         .bind(id.worker.to_string());
                     if let Err(e) = query.execute(&pool).await {
@@ -512,6 +513,7 @@ mod tests {
     use crate::context::State;
 
     use super::*;
+    use apalis_core::task::attempt::Attempt;
     use email_service::Email;
     use futures::StreamExt;
 
@@ -645,6 +647,7 @@ mod tests {
                 acknowledger: job_id.clone(),
                 result: Ok("Success".to_string()),
                 worker: worker_id.clone(),
+                attempts: Attempt::new_with_value(0)
             })
             .await
             .expect("failed to acknowledge the job");

packages/apalis-sql/src/postgres.rs

@@ -59,7 +59,7 @@ use futures::{select, stream, SinkExt};
 use log::error;
 use serde::{de::DeserializeOwned, Serialize};
 use sqlx::postgres::PgListener;
-use sqlx::types::chrono::{DateTime, Utc};
+use chrono::{DateTime, Utc};
 use sqlx::{Pool, Postgres, Row};
 use std::any::type_name;
 use std::convert::TryInto;
@@ -174,12 +174,12 @@ impl<T: Serialize + DeserializeOwned + Sync + Send + Unpin + 'static> Backend<Re
                     }
                     ids = ack_stream.next() => {
                         if let Some(ids) = ids {
-                            let ack_ids: Vec<(String, String, String, String)> = ids.iter().map(|c| {
-                                (c.acknowledger.to_string(), c.worker.to_string(), serde_json::to_string(&c.result).unwrap(), calculate_status(&c.result).to_string())
+                            let ack_ids: Vec<(String, String, String, String, u64)> = ids.iter().map(|c| {
+                                (c.acknowledger.to_string(), c.worker.to_string(), serde_json::to_string(&c.result).unwrap(), calculate_status(&c.result).to_string(), c.attempts.current() as u64)
                             }).collect();
                             let query =
-                                "UPDATE apalis.jobs SET status = Q.status, done_at = now(), lock_by = Q.lock_by, last_error = Q.result FROM (
-                                                SELECT(value-->0)::text as id, (value->>1)::text as worker_id, (value->>2)::text as result, (value->>3)::text as status FROM json_array_elements($1)
+                                "UPDATE apalis.jobs SET status = Q.status, done_at = now(), lock_by = Q.lock_by, last_error = Q.result, attempts = Q.attempts FROM (
+                                                SELECT(value-->0)::text as id, (value->>1)::text as worker_id, (value->>2)::text as result, (value->>3)::text as status, (value->>4)::int as attempts FROM json_array_elements($1)
                                             ) Q
                                             WHERE id = Q.id";
                             if let Err(e) = sqlx::query(query)
@@ -611,8 +611,9 @@ mod tests {
     use crate::context::State;
 
     use super::*;
+    use apalis_core::task::attempt::Attempt;
     use email_service::Email;
-    use sqlx::types::chrono::Utc;
+    use chrono::Utc;
 
     /// migrate DB and return a storage instance.
     async fn setup() -> PostgresStorage<Email> {
@@ -729,6 +730,8 @@ mod tests {
                 acknowledger: job_id.clone(),
                 result: Ok("Success".to_string()),
                 worker: worker_id.clone(),
+                attempts: Attempt::new_with_value(0)
+
             })
             .await
             .expect("failed to acknowledge the job");

packages/apalis-sql/src/sqlite.rs

@@ -16,7 +16,7 @@ use apalis_core::{Backend, BoxCodec};
 use async_stream::try_stream;
 use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt};
 use serde::{de::DeserializeOwned, Serialize};
-use sqlx::types::chrono::Utc;
+use chrono::Utc;
 use sqlx::{Pool, Row, Sqlite};
 use std::any::type_name;
 use std::convert::TryInto;
@@ -484,14 +484,15 @@ impl<T: Sync + Send> Ack<T> for SqliteStorage<T> {
     async fn ack(&mut self, res: AckResponse<Self::Acknowledger>) -> Result<(), sqlx::Error> {
         let pool = self.pool.clone();
         let query =
-                "UPDATE Jobs SET status = ?4, done_at = strftime('%s','now'), last_error = ?3 WHERE id = ?1 AND lock_by = ?2";
+                "UPDATE Jobs SET status = ?4, done_at = strftime('%s','now'), last_error = ?3, attempts =?5 WHERE id = ?1 AND lock_by = ?2";
         let result = serde_json::to_string(&res.result)
             .map_err(|e| sqlx::Error::Io(io::Error::new(io::ErrorKind::InvalidData, e)))?;
         sqlx::query(query)
             .bind(res.acknowledger.to_string())
             .bind(res.worker.to_string())
             .bind(result)
             .bind(calculate_status(&res.result).to_string())
+            .bind(res.attempts.current() as i64)
             .execute(&pool)
             .await?;
         Ok(())
@@ -504,9 +505,10 @@ mod tests {
     use crate::context::State;
 
     use super::*;
+    use apalis_core::task::attempt::Attempt;
     use email_service::Email;
     use futures::StreamExt;
-    use sqlx::types::chrono::Utc;
+    use chrono::Utc;
 
     /// migrate DB and return a storage instance.
     async fn setup() -> SqliteStorage<Email> {
@@ -618,6 +620,8 @@ mod tests {
                 acknowledger: job_id.clone(),
                 result: Ok("Success".to_string()),
                 worker: worker_id.clone(),
+                attempts: Attempt::new_with_value(0)
+
             })
             .await
             .expect("failed to acknowledge the job");