apalis_sqlite/
sink.rs

1use std::{
2    pin::Pin,
3    sync::Arc,
4    task::{Context, Poll},
5};
6
7use futures::{
8    FutureExt, Sink,
9    future::{BoxFuture, Shared},
10};
11use sqlx::SqlitePool;
12use ulid::Ulid;
13
14use crate::{CompactType, SqliteStorage, SqliteTask, config::Config};
15
16type FlushFuture = BoxFuture<'static, Result<(), Arc<sqlx::Error>>>;
17
18#[pin_project::pin_project]
19pub struct SqliteSink<Args, Compact, Codec> {
20    pool: SqlitePool,
21    config: Config,
22    buffer: Vec<SqliteTask<Compact>>,
23    #[pin]
24    flush_future: Option<Shared<FlushFuture>>,
25    _marker: std::marker::PhantomData<(Args, Codec)>,
26}
27
28impl<Args, Compact, Codec> Clone for SqliteSink<Args, Compact, Codec> {
29    fn clone(&self) -> Self {
30        Self {
31            pool: self.pool.clone(),
32            config: self.config.clone(),
33            buffer: Vec::new(),
34            flush_future: None,
35            _marker: std::marker::PhantomData,
36        }
37    }
38}
39
40pub async fn push_tasks(
41    pool: SqlitePool,
42    cfg: Config,
43    buffer: Vec<SqliteTask<CompactType>>,
44) -> Result<(), Arc<sqlx::Error>> {
45    let mut tx = pool.begin().await?;
46    for task in buffer {
47        let id = task
48            .parts
49            .task_id
50            .map(|id| id.to_string())
51            .unwrap_or(Ulid::new().to_string());
52        let run_at = task.parts.run_at as i64;
53        let max_attempts = task.parts.ctx.max_attempts();
54        let priority = task.parts.ctx.priority();
55        let args = task.args;
56        // Use specified queue if specified, otherwise use default
57        let job_type = match task.parts.queue {
58            Some(ref queue) => queue.to_string(),
59            None => cfg.queue().to_string(),
60        };
61        let meta = serde_json::to_string(&task.parts.ctx.meta()).unwrap_or_default();
62        sqlx::query_file!(
63            "queries/task/sink.sql",
64            args,
65            id,
66            job_type,
67            max_attempts,
68            run_at,
69            priority,
70            meta
71        )
72        .execute(&mut *tx)
73        .await?;
74    }
75    tx.commit().await?;
76
77    Ok(())
78}
79
80impl<Args, Compact, Codec> SqliteSink<Args, Compact, Codec> {
81    pub fn new(pool: &SqlitePool, config: &Config) -> Self {
82        Self {
83            pool: pool.clone(),
84            config: config.clone(),
85            buffer: Vec::new(),
86            _marker: std::marker::PhantomData,
87            flush_future: None,
88        }
89    }
90}
91
92impl<Args, Encode, Fetcher> Sink<SqliteTask<CompactType>> for SqliteStorage<Args, Encode, Fetcher>
93where
94    Args: Send + Sync + 'static,
95{
96    type Error = sqlx::Error;
97
98    fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
99        Poll::Ready(Ok(()))
100    }
101
102    fn start_send(self: Pin<&mut Self>, item: SqliteTask<CompactType>) -> Result<(), Self::Error> {
103        // Add the item to the buffer
104        self.project().sink.buffer.push(item);
105        Ok(())
106    }
107
108    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
109        let mut this = self.project();
110
111        // If there's no existing future and buffer is empty, we're done
112        if this.sink.flush_future.is_none() && this.sink.buffer.is_empty() {
113            return Poll::Ready(Ok(()));
114        }
115
116        // Create the future only if we don't have one and there's work to do
117        if this.sink.flush_future.is_none() && !this.sink.buffer.is_empty() {
118            let pool = this.pool.clone();
119            let config = this.config.clone();
120            let buffer = std::mem::take(&mut this.sink.buffer);
121            let sink_fut = push_tasks(pool, config, buffer);
122            this.sink.flush_future = Some((Box::pin(sink_fut) as FlushFuture).shared());
123        }
124
125        // Poll the existing future
126        if let Some(mut fut) = this.sink.flush_future.take() {
127            match fut.poll_unpin(cx) {
128                Poll::Ready(Ok(())) => {
129                    // Future completed successfully, don't put it back
130                    Poll::Ready(Ok(()))
131                }
132                Poll::Ready(Err(e)) => {
133                    // Future completed with error, don't put it back
134                    Poll::Ready(Err(Arc::into_inner(e).unwrap()))
135                }
136                Poll::Pending => {
137                    // Future is still pending, put it back and return Pending
138                    this.sink.flush_future = Some(fut);
139                    Poll::Pending
140                }
141            }
142        } else {
143            // No future and no work to do
144            Poll::Ready(Ok(()))
145        }
146    }
147
148    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
149        self.poll_flush(cx)
150    }
151}
apalis_sqlite/sink.rs

apalis_sqlite/
sink.rs
