apalis_sqlite/
lib.rs

1#![warn(
2    missing_debug_implementations,
3    missing_docs,
4    rust_2018_idioms,
5    unreachable_pub,
6    bad_style,
7    dead_code,
8    improper_ctypes,
9    non_shorthand_field_patterns,
10    overflowing_literals,
11    path_statements,
12    patterns_in_fns_without_body,
13    unconditional_recursion,
14    unused,
15    unused_allocation,
16    unused_comparisons,
17    unused_parens,
18    while_true
19)]
20#![doc = include_str!("../README.md")]
21
22use std::{
23    fmt::Debug,
24    marker::PhantomData,
25    task::{Context, Poll},
26};
27
28pub use apalis_codec::json::JsonCodec;
29
30use apalis_core::{
31    backend::{
32        Backend, BackendConfig, TryNewBackend, WireFormatBackend,
33        ext::{
34            BackendExt,
35            lifecycle::BeforeStart,
36            poll_strategy::{PollWith, StreamStrategy},
37        },
38        finalize::Durable,
39        persistence::{Persisted, TaskPersistLayer},
40    },
41    features_table,
42    task::Task,
43    worker::context::WorkerContext,
44};
45use futures::{
46    FutureExt,
47    channel::mpsc::{self},
48};
49use serde_json::Value;
50use sqlx::{Sqlite, pool::PoolOptions};
51pub use sqlx::{
52    SqliteConnection, SqlitePool,
53    error::Error as SqlxError,
54    sqlite::{SqliteConnectOptions, SqlitePoolOptions},
55};
56use ulid::Ulid;
57
58pub mod callback;
59mod from_row;
60mod persist;
61pub mod queries;
62pub mod shared;
63/// Sink module for pushing tasks to sqlite backend
64mod sink;
65
66use persist::SqlxPersistence;
67const JOBS_TABLE: &str = "Jobs";
68
69/// An alias for [`Task`] specialized for Sqlite
70pub type SqliteTask<Args = Vec<u8>> = Task<Args>;
71
72mod error;
73
74mod config;
75
76pub use config::Config;
77pub use error::Error;
78
79use crate::callback::{DbEvent, HookCallbackListener, update_hook_callback};
80/// SqliteStorage is a storage backend for apalis using sqlite as the database.
81///
82/// It supports both standard polling and event-driven (hooked) storage mechanisms.
83///
84#[doc = features_table! {
85    setup = r#"
86        # {
87        #   use apalis_sqlite::SqliteStorage;
88        #   use sqlx::SqlitePool;
89        #   let pool = SqlitePool::connect(":memory:").await.unwrap();
90        #   SqliteStorage::setup(&pool).await.unwrap();
91        #   SqliteStorage::<u32>::new(&pool)
92        # };
93    "#,
94
95    Backend => supported("Supports storage and retrieval of tasks", true),
96    TaskSink => supported("Ability to push new tasks", true),
97    Serialization => supported("Serialization support for arguments", true),
98    Workflow => supported("Flexible enough to support workflows", true),
99    WebUI => supported("Expose a web interface for monitoring tasks", true),
100    FetchById => supported("Allow fetching a task by its ID", false),
101    RegisterWorker => supported("Allow registering a worker with the backend", false),
102    BackendFactory => supported("Share one connection across multiple workers via [`SqliteStorageFactory`]", false),
103    WaitForCompletion => supported("Wait for tasks to complete without blocking", true),
104    ResumeById => supported("Resume a task by its ID", false),
105    ResumeAbandoned => supported("Resume abandoned tasks", false),
106    ListWorkers => supported("List all workers registered with the backend", false),
107    ListTasks => supported("List all tasks in the backend", false),
108}]
109///
110/// [`SqliteStorageFactory`]: crate::shared::SqliteStorageFactory
111#[pin_project::pin_project]
112#[derive(Debug)]
113pub struct SqliteStorage<Args> {
114    #[pin]
115    persistence: Persisted<SqlxPersistence>,
116    job_type: PhantomData<Args>,
117    codec: JsonCodec,
118}
119
120impl<T> Clone for SqliteStorage<T> {
121    fn clone(&self) -> Self {
122        Self {
123            persistence: self.persistence.clone(),
124            job_type: PhantomData,
125            codec: self.codec.clone(),
126        }
127    }
128}
129
130impl SqliteStorage<()> {
131    /// Connects to a database returning a pool and listener
132    pub fn connect_with_callback(url: &str) -> Result<(SqlitePool, HookCallbackListener), Error> {
133        let (tx, rx) = mpsc::unbounded::<DbEvent>();
134        let listener = HookCallbackListener::new(rx);
135        let pool = PoolOptions::<Sqlite>::new()
136            .after_connect(move |conn, _| {
137                let mut tx = tx.clone();
138                Box::pin(async move {
139                    let mut lock_handle = conn.lock_handle().await?;
140                    lock_handle.set_update_hook(move |ev| update_hook_callback(ev, &mut tx));
141                    Ok(())
142                })
143            })
144            .connect_lazy(url)?;
145        Ok((pool, listener))
146    }
147    /// Perform migrations for storage
148    #[cfg(feature = "migrate")]
149    pub async fn setup(pool: &SqlitePool) -> Result<(), Error> {
150        sqlx::query("PRAGMA journal_mode = 'WAL';")
151            .execute(pool)
152            .await?;
153        sqlx::query("PRAGMA temp_store = MEMORY;")
154            .execute(pool)
155            .await?;
156        sqlx::query("PRAGMA synchronous = OFF;")
157            .execute(pool)
158            .await?;
159        sqlx::query("PRAGMA cache_size = 64000;")
160            .execute(pool)
161            .await?;
162        sqlx::query("PRAGMA journal_size_limit = 67108864;")
163            .execute(pool)
164            .await?;
165        sqlx::query("PRAGMA optimize;").execute(pool).await?;
166        Self::migrations()
167            .run(pool)
168            .await
169            .map_err(sqlx::Error::from)?;
170        Ok(())
171    }
172
173    /// Get sqlite migrations without running them
174    #[cfg(feature = "migrate")]
175    #[must_use]
176    pub fn migrations() -> sqlx::migrate::Migrator {
177        sqlx::migrate!("./migrations")
178    }
179}
180
181impl<T> SqliteStorage<T> {
182    /// Create a new SqliteStorage
183    #[must_use]
184    pub fn new(pool: &SqlitePool) -> SqliteStorage<T> {
185        let config = Config::default().queue(std::any::type_name::<T>());
186        SqliteStorage {
187            job_type: PhantomData,
188            codec: JsonCodec::default(),
189            persistence: Persisted::new(SqlxPersistence {
190                pool: pool.clone(),
191                config,
192            }),
193        }
194    }
195
196    /// Create a new SqliteStorage with a custom configuration
197    #[must_use]
198    pub fn with_config(mut self, config: Config) -> SqliteStorage<T> {
199        self.persistence.config = config;
200        self
201    }
202
203    /// Attach a callback to an instance
204    pub fn with_callback(
205        self,
206        callback: HookCallbackListener,
207    ) -> PollWith<Self, StreamStrategy<HookCallbackListener>> {
208        self.poll_with_stream(callback)
209    }
210
211    /// Get the underlying pool
212    pub fn pool(&self) -> &SqlitePool {
213        &self.persistence.pool
214    }
215}
216
217impl<Args> Backend for SqliteStorage<Args> {
218    type Task = Task<Vec<u8>>;
219
220    type Error = Error;
221
222    fn poll_ready(
223        &mut self,
224        cx: &mut Context<'_>,
225        worker: &WorkerContext,
226    ) -> Poll<Result<(), Self::Error>> {
227        self.persistence
228            .poll_ready(cx, worker, self.persistence.config.heartbeat_interval)
229    }
230
231    fn poll_next(
232        &mut self,
233        cx: &mut Context<'_>,
234        worker: &WorkerContext,
235    ) -> Poll<Option<Result<SqliteTask, Self::Error>>> {
236        self.persistence.poll_next(cx, worker)
237    }
238
239    fn poll_close(
240        &mut self,
241        cx: &mut Context<'_>,
242        worker: &WorkerContext,
243    ) -> Poll<Result<(), Self::Error>> {
244        self.persistence.poll_close(cx, worker)
245    }
246}
247
248impl<Args> BackendConfig for SqliteStorage<Args> {
249    type Args = Args;
250
251    type Id = Ulid;
252
253    type Kind = Durable;
254
255    type Config = Config;
256
257    type Layer = TaskPersistLayer<JsonCodec<Value>, Value>;
258
259    fn config(&self) -> &Self::Config {
260        &self.persistence.config
261    }
262
263    fn middleware(&mut self, _: &mut WorkerContext) -> Self::Layer {
264        self.persistence
265            .layer(JsonCodec::default(), self.config().batch_size)
266            .persist_results(self.config().persist_results)
267            .lock_tasks(self.config().lock_tasks)
268    }
269}
270
271impl<Args> WireFormatBackend for SqliteStorage<Args> {
272    type Codec = JsonCodec<Vec<u8>>;
273
274    type Compact = Vec<u8>;
275    fn codec(&self) -> &Self::Codec {
276        &self.codec
277    }
278}
279
280impl<T> TryNewBackend for SqliteStorage<T> {
281    type Backend = BeforeStart<Self, Self::Error>;
282    fn try_new(config: Self::Config) -> Result<Self::Backend, Self::Error> {
283        let pool = SqlitePoolOptions::new()
284            .connect_lazy(config.database_url.as_deref().unwrap_or(":memory:"))?;
285
286        Ok(SqliteStorage::new(&pool).before_start(|s| {
287            let pool = s.persistence.pool.clone();
288            async move {
289                SqliteStorage::setup(&pool).await?;
290                Ok(())
291            }
292            .boxed()
293        }))
294    }
295}
296#[cfg(test)]
297mod tests {
298    use apalis::prelude::*;
299    use apalis_codec::bincode::BincodeCodec;
300    use apalis_core::backend::ext::BackendExt;
301    use apalis_workflow::*;
302    use futures::{StreamExt, future::ready, stream};
303    use serde::{Deserialize, Serialize};
304    use sqlx::SqlitePool;
305    use std::time::{Duration, Instant};
306
307    use super::*;
308
309    #[tokio::test]
310    async fn basic_worker() {
311        const ITEMS: usize = 3;
312        let url = std::env::var("DATABASE_URL").unwrap_or(":memory:".to_owned());
313        let pool = SqlitePool::connect(&url).await.unwrap();
314
315        let mut backend = SqliteStorage::new(&pool).before_start(|b| {
316            let pool = b.pool().clone();
317            async move { SqliteStorage::setup(&pool).await }
318        });
319
320        let mut start: usize = 0;
321
322        let mut items = stream::repeat_with(move || {
323            start += 1;
324            start
325        })
326        .take(ITEMS);
327
328        backend.push_stream(&mut items).await.unwrap();
329
330        async fn send_reminder(item: usize, wrk: WorkerContext) -> Result<(), BoxDynError> {
331            if ITEMS == item {
332                wrk.stop().unwrap();
333            }
334            Ok(())
335        }
336        let inst = Instant::now();
337        let worker = WorkerBuilder::new("rango-tango")
338            .backend(backend)
339            .retry(RetryPolicy::retries(3))
340            .build(send_reminder);
341        worker.run().await.unwrap();
342
343        println!("Done in {:?}", inst.elapsed());
344    }
345
346    #[tokio::test]
347    async fn hooked_worker() {
348        const ITEMS: usize = 10;
349        let url = &std::env::var("DATABASE_URL").unwrap_or(":memory:".to_owned());
350        let (pool, callback) = SqliteStorage::connect_with_callback(url).unwrap();
351        SqliteStorage::setup(&pool).await.unwrap();
352
353        let backend = SqliteStorage::new(&pool).with_callback(callback);
354        let queue = backend.config().queue.to_string();
355
356        tokio::spawn(async move {
357            let mut start = 0;
358            tokio::time::sleep(Duration::from_secs(5)).await;
359            loop {
360                start += 1;
361                tokio::time::sleep(Duration::from_secs(1)).await;
362
363                let items = stream::repeat_with(move || {
364                    TaskBuilder::new(serde_json::to_vec(&start).unwrap()).build()
365                })
366                .take(1)
367                .collect::<Vec<_>>()
368                .await;
369                let mut tx = pool.begin().await.unwrap();
370                let _ = crate::sink::push_tasks(&mut tx, &queue, &items).await;
371                tx.commit().await.unwrap();
372            }
373        });
374
375        async fn send_reminder(
376            item: usize,
377            wrk: WorkerContext,
378            token: TaskContext,
379        ) -> Result<(), BoxDynError> {
380            let _ctx = token.execution_context().unwrap();
381            if item == ITEMS {
382                wrk.emit(format!("Processed {} tasks", ITEMS));
383                wrk.stop().unwrap();
384            }
385            Ok(())
386        }
387
388        let worker = WorkerBuilder::new("rango-tango-hooked")
389            .backend(backend)
390            .on_event(|_, ev| println!("{ev}"))
391            .build(send_reminder);
392        worker.run().await.unwrap();
393    }
394
395    #[tokio::test]
396    async fn test_workflow() {
397        let workflow = SteppedFlow::new("odd-numbers-workflow")
398            .and_then(|a: usize| async move { Ok::<_, BoxDynError>((0..=a).collect::<Vec<_>>()) })
399            .delay_for(Duration::from_millis(5000))
400            .filter_map(|x| ready(if x % 2 != 0 { Some(x) } else { None }))
401            .and_then(|a: Vec<usize>| async move {
402                println!("Sum: {}", a.iter().sum::<usize>());
403                Err::<(), BoxDynError>("Intentional Error".into())
404            });
405
406        let pool =
407            SqlitePool::connect(&std::env::var("DATABASE_URL").unwrap_or(":memory:".to_owned()))
408                .await
409                .unwrap();
410        SqliteStorage::setup(&pool).await.unwrap();
411
412        let mut sqlite = SqliteStorage::new(&pool)
413            .with_codec(BincodeCodec)
414            // Our worker may sleep at the delay, (our heartbeat is 30s, so we wake the worker every second)
415            .poll_with_interval(Duration::from_secs(1));
416
417        sqlite.push_start(42).await.unwrap();
418
419        let worker = WorkerBuilder::new("rango-tango-workflow")
420            .backend(sqlite)
421            .on_event(|ctx, ev| {
422                println!("On Event = {:?}", ev);
423                if matches!(ev, Event::Error(_)) {
424                    ctx.stop().unwrap();
425                }
426            })
427            .build(workflow);
428        worker.run().await.unwrap();
429    }
430
431    #[tokio::test]
432    async fn test_workflow_complete() {
433        #[derive(Debug, Serialize, Deserialize, Clone)]
434        struct PipelineConfig {
435            min_confidence: f32,
436            enable_sentiment: bool,
437        }
438
439        #[derive(Debug, Serialize, Deserialize)]
440        struct UserInput {
441            text: String,
442        }
443
444        #[derive(Debug, Serialize, Deserialize)]
445        struct Classified {
446            text: String,
447            label: String,
448            confidence: f32,
449        }
450
451        #[derive(Debug, Serialize, Deserialize)]
452        struct Summary {
453            text: String,
454            sentiment: Option<String>,
455        }
456
457        let workflow = SteppedFlow::new("text-pipeline")
458            // Step 1: Preprocess input (e.g., tokenize, lowercase)
459            .and_then(|input: UserInput, worker: WorkerContext| async move {
460                worker.emit(format!("Preprocessing input: {}", input.text));
461                let processed = input.text.to_lowercase();
462                Ok::<_, BoxDynError>(processed)
463            })
464            // Step 2: Classify text
465            .and_then(|text: String| async move {
466                let confidence = 0.85; // pretend model confidence
467                let items = text.split_whitespace().collect::<Vec<_>>();
468                let results = items
469                    .into_iter()
470                    .map(|x| Classified {
471                        text: x.to_string(),
472                        label: if x.contains("rust") {
473                            "Tech"
474                        } else {
475                            "General"
476                        }
477                        .to_string(),
478                        confidence,
479                    })
480                    .collect::<Vec<_>>();
481                Ok::<_, BoxDynError>(results)
482            })
483            .delay_for(Duration::from_millis(5000))
484            // Step 3: Filter out low-confidence predictions
485            .filter_map(
486                |c: Classified| async move { if c.confidence >= 0.6 { Some(c) } else { None } },
487            )
488            .filter_map(move |c: Classified, config: Data<PipelineConfig>| {
489                let cfg = config.enable_sentiment;
490                async move {
491                    if !cfg {
492                        return Some(Summary {
493                            text: c.text,
494                            sentiment: None,
495                        });
496                    }
497
498                    // pretend we run a sentiment model
499                    let sentiment = if c.text.contains("delightful") {
500                        "positive"
501                    } else {
502                        "neutral"
503                    };
504                    Some(Summary {
505                        text: c.text,
506                        sentiment: Some(sentiment.to_string()),
507                    })
508                }
509            })
510            .and_then(|a: Vec<Summary>, worker: WorkerContext| async move {
511                worker.emit(format!("Generated {} summaries", a.len()));
512                worker.stop()
513            });
514
515        let pool = SqlitePool::connect(":memory:").await.unwrap();
516
517        SqliteStorage::setup(&pool).await.unwrap();
518
519        let backoff = BackoffConfig::new(Duration::from_millis(5000));
520        let mut sqlite =
521            SqliteStorage::new(&pool).poll_with_backoff(Duration::from_millis(200), backoff);
522
523        let input = UserInput {
524            text: "Rust makes systems programming delightful!".to_string(),
525        };
526        sqlite.push_start(input).await.unwrap();
527
528        let worker = WorkerBuilder::new("rango-tango")
529            .backend(sqlite)
530            .data(PipelineConfig {
531                min_confidence: 0.8,
532                enable_sentiment: true,
533            })
534            .on_event(|ctx, ev| match ev {
535                Event::Custom(msg) => {
536                    if let Some(m) = msg.downcast_ref::<String>() {
537                        println!("Custom Message: {}", m);
538                    }
539                }
540                Event::Error(_) => {
541                    println!("On Error = {:?}", ev);
542                    ctx.stop().unwrap();
543                }
544                _ => {
545                    println!("On Event = {:?}", ev);
546                }
547            })
548            .build(workflow);
549        worker.run().await.unwrap();
550    }
551}
apalis_sqlite/lib.rs

apalis_sqlite/
lib.rs
