apalis_postgres/
backend.rs

1use std::{
2    marker::PhantomData,
3    task::{Context, Poll},
4};
5
6use apalis_codec::json::JsonCodec;
7use apalis_core::{
8    backend::{
9        Backend, BackendConfig, WireFormatBackend,
10        ext::poll_strategy::{PollWith, StreamStrategy},
11        finalize::Durable,
12        persistence::{Persisted, TaskPersistLayer},
13    },
14    features_table,
15    worker::context::WorkerContext,
16};
17use serde_json::Value;
18use sqlx::PgPool;
19use ulid::Ulid;
20
21use crate::{PgTask, config::Config, error::Error, persistence::SqlxPersistence, pubsub::Pubsub};
22
23/// A backend for persisting and consuming jobs behind a postgres database
24#[doc = features_table! {
25    setup = r#"
26        # {
27        #   use apalis_postgres::PostgresStorage;
28        #   use sqlx::PgPool;
29        #   let pool = PgPool::connect(std::env::var("DATABASE_URL").unwrap().as_str()).await.unwrap();
30        #   PostgresStorage::setup(&pool).await.unwrap();
31        #   PostgresStorage::<u32>::new(&pool)
32        # };
33    "#,
34
35    Backend => supported("Supports storage and retrieval of tasks", true),
36    TaskSink => supported("Ability to push new tasks", true),
37    Serialization => supported("Serialization support for arguments", true),
38    Workflow => supported("Flexible enough to support workflows", true),
39    WebUI => supported("Expose a web interface for monitoring tasks", true),
40    FetchById => supported("Allow fetching a task by its ID", false),
41    RegisterWorker => supported("Allow registering a worker with the backend", false),
42    MakeShared => supported("Share one connection across multiple workers via [`PostgresStorageFactory`]", false),
43    WaitForCompletion => supported("Wait for tasks to complete without blocking", true),
44    ResumeById => supported("Resume a task by its ID", false),
45    ResumeAbandoned => supported("Resume abandoned tasks", false),
46    ListWorkers => supported("List all workers registered with the backend", false),
47    ListTasks => supported("List all tasks in the backend", false),
48}]
49///
50/// [`PostgresStorageFactory`]: crate::factory::PostgresStorageFactory
51#[pin_project::pin_project]
52pub struct PostgresStorage<Args> {
53    #[pin]
54    pub(crate) persistence: Persisted<SqlxPersistence>,
55    codec: JsonCodec,
56    _marker: PhantomData<Args>,
57}
58
59impl<Args> Clone for PostgresStorage<Args> {
60    fn clone(&self) -> Self {
61        Self {
62            persistence: self.persistence.clone(),
63            codec: self.codec.clone(),
64            _marker: PhantomData,
65        }
66    }
67}
68
69impl PostgresStorage<()> {
70    /// Runs the PostgreSQL storage migrations.
71    ///
72    /// ## Fresh databases
73    ///
74    /// No manual setup is required. Calling `setup()` will create the required
75    /// tables and migration history.
76    ///
77    /// ## Upgrading to `1.0`
78    ///
79    /// > **⚠️ Important:** Existing databases created by a pre-`1.0` version
80    /// > require a **one-time manual migration** before calling `setup()`.
81    ///
82    /// The `1.0` migration history is no longer relocated automatically by
83    /// `setup()`. Follow the **"Upgrading to 1.0"** section in the README to
84    /// perform the required transition.
85    ///
86    /// After the transition has been completed, `setup()` can be used normally
87    /// for subsequent migrations.
88    ///
89    /// ## Example
90    ///
91    /// ```no_run
92    /// use apalis_postgres::PostgresStorage;
93    /// use sqlx::PgPool;
94    ///
95    /// # async fn run(pool: PgPool) -> Result<(), apalis_postgres::Error> {
96    /// PostgresStorage::<()>::setup(&pool).await?;
97    /// # Ok(())
98    /// # }
99    /// ```
100    ///
101    /// ## Errors
102    ///
103    /// Returns an error if the migrations cannot be applied to the database.
104    #[cfg(feature = "migrate")]
105    pub async fn setup(pool: &PgPool) -> Result<(), Error> {
106        Self::migrations().run(pool).await.map_err(|r| {
107            use sqlx::migrate::MigrateError::VersionMismatch;
108
109            if let VersionMismatch(20220530084123) = r {
110                Error::UpgradeRequired {
111                    current: "0.7.4".into(),
112                    required: "1.0.0".into(),
113                    link: Some(
114                        "https://github.com/apalis-dev/apalis-postgres#upgrading-to-10".into(),
115                    ),
116                }
117            } else {
118                Error::Database(r.into())
119            }
120        })?;
121        Ok(())
122    }
123
124    /// Get postgres migrations without running them
125    #[cfg(feature = "migrate")]
126    pub fn migrations() -> sqlx::migrate::Migrator {
127        sqlx::migrate!("./migrations")
128    }
129}
130
131impl<Args> PostgresStorage<Args> {
132    /// Creates a new PostgresStorage instance.
133    pub fn new(pool: &PgPool) -> Self {
134        let config = Config::default().queue(std::any::type_name::<Args>());
135        let persistence = Persisted::new(SqlxPersistence {
136            config,
137            pool: pool.clone(),
138        });
139        Self {
140            _marker: PhantomData,
141            codec: JsonCodec::default(),
142            persistence,
143        }
144    }
145
146    /// Mount a standalone [`Pubsub`] which uses its own connection under the hood
147    pub fn with_pubsub(self) -> PollWith<Self, StreamStrategy<Pubsub>> {
148        let pool = self.pool().clone();
149        let config = self.config();
150        let namespace = config.queue.to_string();
151        PollWith::new(self, StreamStrategy::new(Pubsub::new(pool, namespace)))
152    }
153
154    /// Configure a new PostgresStorage instance.
155    pub fn with_config(mut self, config: Config) -> Self {
156        self.persistence.config = config;
157        self
158    }
159
160    /// Returns a reference to the pool.
161    pub fn pool(&self) -> &PgPool {
162        &self.persistence.pool
163    }
164
165    /// Returns a reference to the config.
166    pub fn config(&self) -> &Config {
167        &self.persistence.config
168    }
169}
170
171impl<Args> Backend for PostgresStorage<Args> {
172    type Task = PgTask;
173    type Error = Error;
174
175    fn poll_ready(
176        &mut self,
177        cx: &mut Context<'_>,
178        worker: &WorkerContext,
179    ) -> Poll<Result<(), Self::Error>> {
180        self.persistence
181            .poll_ready(cx, worker, self.config().heartbeat_interval)
182    }
183
184    fn poll_next(
185        &mut self,
186        cx: &mut Context<'_>,
187        worker: &WorkerContext,
188    ) -> Poll<Option<Result<PgTask, Self::Error>>> {
189        self.persistence.poll_next(cx, worker)
190    }
191
192    fn poll_close(
193        &mut self,
194        cx: &mut Context<'_>,
195        worker: &WorkerContext,
196    ) -> Poll<Result<(), Self::Error>> {
197        self.persistence.poll_close(cx, worker)
198    }
199}
200
201impl<Args> BackendConfig for PostgresStorage<Args> {
202    type Args = Args;
203
204    type Kind = Durable;
205
206    type Id = Ulid;
207
208    type Config = Config;
209
210    type Layer = TaskPersistLayer<JsonCodec<Value>, Value>;
211
212    fn config(&self) -> &Self::Config {
213        &self.persistence.config
214    }
215
216    fn middleware(&mut self, _: &mut WorkerContext) -> Self::Layer {
217        self.persistence
218            .layer(JsonCodec::<Value>::default(), self.config().batch_size)
219            .persist_results(self.config().persist_results)
220            .lock_tasks(self.config().lock_tasks)
221    }
222}
223
224impl<Args> WireFormatBackend for PostgresStorage<Args> {
225    type Codec = JsonCodec;
226
227    type Compact = Vec<u8>;
228
229    fn codec(&self) -> &Self::Codec {
230        &self.codec
231    }
232}
apalis_postgres/backend.rs

apalis_postgres/
backend.rs
