Skip to main content

apalis_sqlite/
callback.rs

1//! Callbacks allow realtime listeners for new jobs
2//!
3//! ## Example usage
4//!
5//! ```ignore
6//! let (pool, callback) = SqliteStorage::connect_with_callback(url).unwrap();
7//! SqliteStorage::setup(&pool).await.unwrap();
8//! let backend = SqliteStorage::new(&pool).with_callback(callback);
9//! ```
10use apalis_core::backend::ext::poll_strategy::{PollWith, StreamStrategy};
11use futures::channel::mpsc::{UnboundedReceiver, UnboundedSender};
12use futures::{Stream, StreamExt};
13use sqlx::sqlite::{SqliteOperation, UpdateHookResult};
14
15use std::pin::Pin;
16use std::sync::{Arc, Mutex};
17use std::task::{Context, Poll};
18
19use crate::{JOBS_TABLE, SqliteStorage};
20
21/// An [SqliteStorage] that polls when [HookCallbackListener] is invoked
22pub type SqliteStorageWithHook<Args> =
23    PollWith<SqliteStorage<Args>, StreamStrategy<HookCallbackListener>>;
24
25/// Database event emitted by SQLite update hook
26#[derive(Debug)]
27pub struct DbEvent {
28    op: SqliteOperation,
29    db_name: String,
30    table_name: String,
31    rowid: i64,
32}
33
34impl DbEvent {
35    /// Get the operation type of the database event
36    #[must_use]
37    pub fn operation(&self) -> &SqliteOperation {
38        &self.op
39    }
40
41    /// Get the database name of the database event
42    #[must_use]
43    pub fn db_name(&self) -> &str {
44        &self.db_name
45    }
46
47    /// Get the table name of the database event
48    #[must_use]
49    pub fn table_name(&self) -> &str {
50        &self.table_name
51    }
52
53    /// Get the rowid of the database event
54    #[must_use]
55    pub fn rowid(&self) -> i64 {
56        self.rowid
57    }
58}
59
60// Callback for SQLite update hook
61pub(crate) fn update_hook_callback(event: UpdateHookResult<'_>, tx: &mut UnboundedSender<DbEvent>) {
62    if event.operation == SqliteOperation::Insert && event.table == JOBS_TABLE {
63        let _ = tx.start_send(DbEvent {
64            op: event.operation,
65            db_name: event.database.to_owned(),
66            table_name: event.table.to_owned(),
67            rowid: event.rowid,
68        });
69    }
70}
71
72/// Listener for database events emitted by SQLite update hook
73#[derive(Debug, Clone)]
74pub struct HookCallbackListener {
75    rx: Arc<Mutex<UnboundedReceiver<DbEvent>>>,
76}
77
78impl HookCallbackListener {
79    /// Create a new HookCallbackListener
80    #[must_use]
81    pub fn new(rx: UnboundedReceiver<DbEvent>) -> Self {
82        Self {
83            rx: Arc::new(Mutex::new(rx)),
84        }
85    }
86}
87
88impl Stream for HookCallbackListener {
89    type Item = ();
90
91    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
92        log::trace!("HookCallbackListener: poll_next");
93        match self.rx.lock().unwrap().poll_next_unpin(cx) {
94            Poll::Ready(Some(_)) => Poll::Ready(Some(())),
95            Poll::Ready(None) => Poll::Ready(None),
96            Poll::Pending => Poll::Pending,
97        }
98        .map(|s| {
99            log::trace!(
100                "HookCallbackListener: poll_ready: {ready}",
101                ready = s.is_some()
102            );
103            s
104        })
105    }
106}