apalis_postgres/
sink.rs

1use futures::{FutureExt, Sink, TryFutureExt};
2use sqlx::{Executor, postgres::types::PgHstore};
3use std::{
4    pin::Pin,
5    task::{Context, Poll},
6};
7use ulid::Ulid;
8
9use crate::{PgTask, backend::PostgresStorage, error::Error, timestamp::Timestamp};
10
11/// Push a batch of tasks to the database
12pub fn push_tasks<E>(
13    conn: &mut E,
14    queue: &str,
15    buffer: Vec<PgTask>,
16) -> impl futures::Future<Output = Result<(), Error>> + Send
17where
18    for<'e> &'e mut E: Executor<'e, Database = sqlx::Postgres> + Send,
19{
20    // Build the multi-row INSERT with UNNEST
21    let mut ids = Vec::new();
22    let mut job_data = Vec::new();
23    let mut run_ats = Vec::new();
24    let mut priorities = Vec::new();
25    let mut max_attempts_vec = Vec::new();
26    let mut metadata = Vec::new();
27    let mut idempotency_key: Vec<Option<String>> = Vec::new();
28    let now = Timestamp::now();
29    for task in buffer {
30        ids.push(
31            task.task_id()
32                .map(|id| id.to_string())
33                .unwrap_or(Ulid::generate().to_string()),
34        );
35
36        run_ats.push(task.run_at().map(|f| f as i64).unwrap_or(now.0 as i64));
37        priorities.push(task.priority().map(|f| f as i32).unwrap_or_default());
38        max_attempts_vec.push(task.max_attempts().map(|f| f as i32).unwrap_or(25));
39        metadata.push(PgHstore(
40            task.metadata()
41                .clone()
42                .into_inner()
43                .into_iter()
44                .map(|(k, v)| (k, Some(v)))
45                .collect(),
46        ));
47        idempotency_key.push(task.idempotency_key().map(|a| a.to_owned()));
48        job_data.push(task.args);
49    }
50
51    sqlx::query_file!(
52        "queries/task/sink.sql",
53        &ids,
54        &queue,
55        &job_data,
56        &max_attempts_vec,
57        &run_ats,
58        &priorities,
59        &metadata,
60        &idempotency_key as &[Option<String>]
61    )
62    .execute(conn)
63    .map_ok(|_| ())
64    .map_err(|e| e.into())
65    .boxed()
66}
67
68impl<Args> Sink<PgTask> for PostgresStorage<Args>
69where
70    Args: Unpin + Send + Sync + 'static,
71{
72    type Error = Error;
73
74    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
75        self.project().persistence.poll_ready(cx)
76    }
77
78    fn start_send(self: Pin<&mut Self>, item: PgTask) -> Result<(), Self::Error> {
79        self.project().persistence.start_send(item)
80    }
81
82    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
83        self.project().persistence.poll_flush(cx)
84    }
85
86    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
87        Sink::poll_close(self.project().persistence, cx)
88    }
89}
apalis_postgres/sink.rs

apalis_postgres/
sink.rs
