Skip to main content

apalis_sqlite/
sink.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4    time::{SystemTime, UNIX_EPOCH},
5};
6
7use futures::Sink;
8use sqlx::Executor;
9use ulid::Ulid;
10
11use crate::{SqliteStorage, SqliteTask, error::Error};
12
13/// Push a batch of tasks into the database
14pub async fn push_tasks<'a, E>(
15    executor: &'a mut E,
16    queue: &str,
17    tasks: &[SqliteTask],
18) -> Result<(), Error>
19where
20    for<'e> &'e mut E: Executor<'e, Database = sqlx::Sqlite>,
21{
22    for task in tasks {
23        let id = task
24            .task_id()
25            .as_ref()
26            .map(|id| id.to_string())
27            .unwrap_or(Ulid::generate().to_string());
28        let run_at = task.run_at().unwrap_or_else(|| {
29            SystemTime::now()
30                .duration_since(UNIX_EPOCH)
31                .unwrap()
32                .as_secs()
33        }) as i64;
34        let max_attempts = task.max_attempts().unwrap_or(25) as i64;
35        let priority = task.priority().unwrap_or_default() as i64;
36        let args = &task.args;
37        let idempotency_key = task.idempotency_key();
38        let meta = serde_json::to_string(task.metadata()).unwrap_or_default();
39        sqlx::query_file!(
40            "queries/task/sink.sql",
41            args,
42            id,
43            queue,
44            max_attempts,
45            run_at,
46            priority,
47            meta,
48            idempotency_key
49        )
50        .execute(&mut *executor)
51        .await?;
52    }
53    Ok(())
54}
55
56impl<Args> Sink<SqliteTask> for SqliteStorage<Args> {
57    type Error = Error;
58
59    fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
60        Poll::Ready(Ok(()))
61    }
62
63    fn start_send(self: Pin<&mut Self>, item: SqliteTask) -> Result<(), Self::Error> {
64        self.project().persistence.start_send(item)
65    }
66
67    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
68        Sink::poll_flush(self.project().persistence, cx)
69    }
70
71    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
72        Sink::poll_close(self.project().persistence, cx)
73    }
74}