oxidite_queue/
postgres.rs

1use async_trait::async_trait;
2use sqlx::{PgPool, Row};
3use crate::{QueueBackend, job::JobWrapper, Result, QueueError};
4
5/// PostgreSQL queue backend
6pub struct PostgresBackend {
7    pool: PgPool,
8    table_name: String,
9    dlq_table_name: String,
10}
11
12impl PostgresBackend {
13    /// Create a new PostgreSQL queue backend
14    pub async fn new(pool: PgPool, table_name: &str) -> Result<Self> {
15        let backend = Self {
16            pool,
17            table_name: table_name.to_string(),
18            dlq_table_name: format!("{}_dlq", table_name),
19        };
20        
21        // Initialize tables if they don't exist
22        backend.init_tables().await?;
23        
24        Ok(backend)
25    }
26
27    async fn init_tables(&self) -> Result<()> {
28        // Create main queue table
29        sqlx::query(&format!(
30            r#"CREATE TABLE IF NOT EXISTS {} (
31                id TEXT PRIMARY KEY,
32                payload JSONB NOT NULL,
33                priority INTEGER DEFAULT 0,
34                attempts INTEGER DEFAULT 0,
35                max_attempts INTEGER DEFAULT 3,
36                scheduled_at TIMESTAMP WITH TIME ZONE,
37                created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
38                updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
39                status VARCHAR(20) DEFAULT 'pending'
40            )"#,
41            self.table_name
42        ))
43        .execute(&self.pool)
44        .await
45        .map_err(|e| QueueError::BackendError(e.to_string()))?;
46
47        // Create dead letter queue table
48        sqlx::query(&format!(
49            r#"CREATE TABLE IF NOT EXISTS {} (
50                id TEXT PRIMARY KEY,
51                payload JSONB NOT NULL,
52                priority INTEGER DEFAULT 0,
53                attempts INTEGER DEFAULT 0,
54                error TEXT,
55                created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
56                updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
57                status VARCHAR(20) DEFAULT 'dead_letter'
58            )"#,
59            self.dlq_table_name
60        ))
61        .execute(&self.pool)
62        .await
63        .map_err(|e| QueueError::BackendError(e.to_string()))?;
64
65        Ok(())
66    }
67}
68
69#[async_trait]
70impl QueueBackend for PostgresBackend {
71    async fn enqueue(&self, mut job: JobWrapper) -> Result<()> {
72        job.status = crate::job::JobStatus::Pending;
73        
74        sqlx::query(&format!(
75            r#"INSERT INTO {} (id, payload, priority, attempts, max_attempts, scheduled_at, status)
76               VALUES ($1, $2, $3, $4, $5, $6, $7)"#,
77            self.table_name
78        ))
79        .bind(&job.id)
80        .bind(serde_json::to_value(&job)?)
81        .bind(job.priority)
82        .bind(job.attempts as i32)
83        .bind(job.max_retries as i32)
84        .bind(job.scheduled_at.map(|t| chrono::DateTime::<chrono::Utc>::from_timestamp(t, 0)))
85        .bind("pending")
86        .execute(&self.pool)
87        .await
88        .map_err(|e| QueueError::BackendError(e.to_string()))?;
89
90        Ok(())
91    }
92
93    async fn dequeue(&self) -> Result<Option<JobWrapper>> {
94        // Get the next job that is pending and ready to be processed
95        let row = sqlx::query(&format!(
96            r#"UPDATE {}
97               SET status = 'running', attempts = attempts + 1, updated_at = NOW()
98               WHERE id = (
99                   SELECT id FROM {}
100                   WHERE status = 'pending'
101                   AND (scheduled_at IS NULL OR scheduled_at <= NOW())
102                   ORDER BY priority DESC, created_at ASC
103                   LIMIT 1
104                   FOR UPDATE SKIP LOCKED
105               )
106               RETURNING payload"#,
107            self.table_name, self.table_name
108        ))
109        .fetch_optional(&self.pool)
110        .await
111        .map_err(|e| QueueError::BackendError(e.to_string()))?;
112
113        if let Some(row) = row {
114            let payload: serde_json::Value = row.try_get("payload")
115                .map_err(|e| QueueError::BackendError(e.to_string()))?;
116            let mut job: JobWrapper = serde_json::from_value(payload)
117                .map_err(|e| QueueError::SerializationError(e))?;
118            job.status = crate::job::JobStatus::Running;
119            job.attempts += 1;
120            Ok(Some(job))
121        } else {
122            Ok(None)
123        }
124    }
125
126    async fn complete(&self, job_id: &str) -> Result<()> {
127        sqlx::query(&format!(
128            r#"DELETE FROM {} WHERE id = $1"#,
129            self.table_name
130        ))
131        .bind(job_id)
132        .execute(&self.pool)
133        .await
134        .map_err(|e| QueueError::BackendError(e.to_string()))?;
135
136        Ok(())
137    }
138
139    async fn fail(&self, job_id: &str, error: String) -> Result<()> {
140        // Check if max attempts reached
141        let row = sqlx::query(&format!(
142            r#"SELECT attempts, max_retries FROM {} WHERE id = $1"#,
143            self.table_name
144        ))
145        .bind(job_id)
146        .fetch_optional(&self.pool)
147        .await
148        .map_err(|e| QueueError::BackendError(e.to_string()))?;
149
150        if let Some(row) = row {
151            let attempts: i32 = row.try_get("attempts")
152                .map_err(|e| QueueError::BackendError(e.to_string()))?;
153            let max_retries: i32 = row.try_get("max_retries")
154                .map_err(|e| QueueError::BackendError(e.to_string()))?;
155
156            if attempts >= max_retries {
157                // Move to dead letter queue
158                self.move_to_dead_letter_with_error(job_id, error).await?;
159            } else {
160                // Update status back to pending for retry
161                sqlx::query(&format!(
162                    r#"UPDATE {} SET status = 'pending', updated_at = NOW() WHERE id = $1"#,
163                    self.table_name
164                ))
165                .bind(job_id)
166                .execute(&self.pool)
167                .await
168                .map_err(|e| QueueError::BackendError(e.to_string()))?;
169            }
170        }
171
172        Ok(())
173    }
174
175    async fn retry(&self, job: JobWrapper) -> Result<()> {
176        // First, delete the existing job
177        sqlx::query(&format!(
178            r#"DELETE FROM {} WHERE id = $1"#,
179            self.table_name
180        ))
181        .bind(&job.id)
182        .execute(&self.pool)
183        .await
184        .map_err(|e| QueueError::BackendError(e.to_string()))?;
185
186        // Then re-enqueue with reset attempts
187        let mut job = job;
188        job.attempts = 0;
189        job.status = crate::job::JobStatus::Pending;
190        
191        self.enqueue(job).await
192    }
193
194    async fn move_to_dead_letter(&self, job: JobWrapper) -> Result<()> {
195        self.move_to_dead_letter_with_error(&job.id, job.error.unwrap_or_else(|| "Unknown error".to_string())).await
196    }
197
198    async fn list_dead_letter(&self) -> Result<Vec<JobWrapper>> {
199        let rows = sqlx::query(&format!(
200            r#"SELECT payload FROM {} ORDER BY created_at DESC"#,
201            self.dlq_table_name
202        ))
203        .fetch_all(&self.pool)
204        .await
205        .map_err(|e| QueueError::BackendError(e.to_string()))?;
206
207        let mut jobs = Vec::new();
208        for row in rows {
209            let payload: serde_json::Value = row.try_get("payload")
210                .map_err(|e| QueueError::BackendError(e.to_string()))?;
211            if let Ok(job) = serde_json::from_value::<JobWrapper>(payload) {
212                jobs.push(job);
213            }
214        }
215
216        Ok(jobs)
217    }
218
219    async fn retry_from_dead_letter(&self, job_id: &str) -> Result<()> {
220        // Get job from DLQ
221        let row = sqlx::query(&format!(
222            r#"SELECT payload FROM {} WHERE id = $1"#,
223            self.dlq_table_name
224        ))
225        .bind(job_id)
226        .fetch_optional(&self.pool)
227        .await
228        .map_err(|e| QueueError::BackendError(e.to_string()))?;
229
230        if let Some(row) = row {
231            let payload: serde_json::Value = row.try_get("payload")
232                .map_err(|e| QueueError::BackendError(e.to_string()))?;
233            let mut job: JobWrapper = serde_json::from_value(payload)
234                .map_err(|e| QueueError::SerializationError(e))?;
235
236            // Delete from DLQ
237            sqlx::query(&format!(
238                r#"DELETE FROM {} WHERE id = $1"#,
239                self.dlq_table_name
240            ))
241            .bind(job_id)
242            .execute(&self.pool)
243            .await
244            .map_err(|e| QueueError::BackendError(e.to_string()))?;
245
246            // Reset and enqueue to main queue
247            job.status = crate::job::JobStatus::Pending;
248            job.attempts = 0;
249            job.error = None;
250
251            self.enqueue(job).await?;
252        }
253
254        Ok(())
255    }
256
257    async fn clear(&self) -> Result<()> {
258        sqlx::query(&format!("DELETE FROM {}", self.table_name))
259            .execute(&self.pool)
260            .await
261            .map_err(|e| QueueError::BackendError(e.to_string()))?;
262        Ok(())
263    }
264}
265
266impl PostgresBackend {
267    async fn move_to_dead_letter_with_error(&self, job_id: &str, error: String) -> Result<()> {
268        // Get the job from the main queue
269        let row = sqlx::query(&format!(
270            r#"SELECT payload, priority, attempts FROM {} WHERE id = $1"#,
271            self.table_name
272        ))
273        .bind(job_id)
274        .fetch_optional(&self.pool)
275        .await
276        .map_err(|e| QueueError::BackendError(e.to_string()))?;
277
278        if let Some(row) = row {
279            let payload: serde_json::Value = row.try_get("payload")
280                .map_err(|e| QueueError::BackendError(e.to_string()))?;
281            let priority: i32 = row.try_get("priority")
282                .map_err(|e| QueueError::BackendError(e.to_string()))?;
283            let attempts: i32 = row.try_get("attempts")
284                .map_err(|e| QueueError::BackendError(e.to_string()))?;
285
286            // Insert into dead letter queue
287            sqlx::query(&format!(
288                r#"INSERT INTO {} (id, payload, priority, attempts, error)
289                   VALUES ($1, $2, $3, $4, $5)"#,
290                self.dlq_table_name
291            ))
292            .bind(job_id)
293            .bind(&payload)
294            .bind(priority)
295            .bind(attempts)
296            .bind(error)
297            .execute(&self.pool)
298            .await
299            .map_err(|e| QueueError::BackendError(e.to_string()))?;
300
301            // Remove from main queue
302            sqlx::query(&format!(
303                r#"DELETE FROM {} WHERE id = $1"#,
304                self.table_name
305            ))
306            .bind(job_id)
307            .execute(&self.pool)
308            .await
309            .map_err(|e| QueueError::BackendError(e.to_string()))?;
310        }
311
312        Ok(())
313    }
314}
oxidite_queue/postgres.rs

oxidite_queue/
postgres.rs
