use kevy_embedded::{PubsubFrame, Store};
use sqlx::PgPool;
use crate::queue;
use crate::queue::QueuedMessage;
use crate::store::{Notifier, QueueStore, StoreError};
#[derive(Debug, Clone)]
pub struct PgQueueStore {
pool: PgPool,
}
impl PgQueueStore {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
}
#[async_trait::async_trait]
impl QueueStore for PgQueueStore {
async fn enqueue(
&self,
sender: &str,
recipient: &str,
domain: &str,
message_data: &[u8],
message_id: Option<&str>,
now: i64,
is_forwarded: bool,
) -> Result<i64, StoreError> {
Ok(queue::enqueue_ex(
&self.pool,
sender,
recipient,
domain,
message_data,
message_id,
now,
is_forwarded,
)
.await?)
}
async fn enqueue_scheduled(
&self,
sender: &str,
recipient: &str,
domain: &str,
message_data: &[u8],
message_id: Option<&str>,
created_at: i64,
scheduled_at: i64,
) -> Result<i64, StoreError> {
Ok(queue::enqueue_scheduled(
&self.pool,
sender,
recipient,
domain,
message_data,
message_id,
created_at,
scheduled_at,
)
.await?)
}
async fn dequeue(&self, now: i64, limit: u32) -> Result<Vec<QueuedMessage>, StoreError> {
Ok(queue::dequeue(&self.pool, now, limit).await?)
}
async fn recover_stale_inflight(&self, now: i64) -> Result<u64, StoreError> {
Ok(queue::recover_stale_inflight(&self.pool, now).await?)
}
async fn mark_inflight(&self, id: i64, now: i64) -> Result<(), StoreError> {
Ok(queue::mark_inflight(&self.pool, id, now).await?)
}
async fn mark_delivered(&self, id: i64, now: i64) -> Result<(), StoreError> {
Ok(queue::mark_delivered(&self.pool, id, now).await?)
}
async fn mark_failed(
&self,
id: i64,
error: &str,
next_retry: i64,
now: i64,
) -> Result<(), StoreError> {
Ok(queue::mark_failed(&self.pool, id, error, next_retry, now).await?)
}
async fn mark_bounced(&self, id: i64, error: &str, now: i64) -> Result<(), StoreError> {
Ok(queue::mark_bounced(&self.pool, id, error, now).await?)
}
async fn get_message(&self, id: i64) -> Result<Option<QueuedMessage>, StoreError> {
Ok(queue::get_message(&self.pool, id).await?)
}
async fn queue_stats(&self) -> Result<Vec<(String, i64)>, StoreError> {
Ok(queue::queue_stats(&self.pool).await?)
}
async fn list_recent(&self, limit: i32) -> Result<Vec<QueuedMessage>, StoreError> {
Ok(queue::list_recent(&self.pool, limit).await?)
}
async fn cancel_pending(&self, id: i64) -> Result<bool, StoreError> {
Ok(queue::cancel_pending(&self.pool, id).await?)
}
async fn cancel_pending_by_message_id(
&self,
message_id: &str,
sender: &str,
) -> Result<bool, StoreError> {
Ok(queue::cancel_pending_by_message_id(&self.pool, message_id, sender).await?)
}
async fn retry_message(&self, id: i64, now: i64) -> Result<bool, StoreError> {
Ok(queue::retry_message(&self.pool, id, now).await?)
}
async fn is_suppressed(&self, email: &str) -> bool {
queue::is_suppressed(&self.pool, email).await
}
async fn add_suppression(
&self,
email: &str,
reason: &str,
smtp_code: Option<i32>,
) -> Result<(), StoreError> {
Ok(queue::add_suppression(&self.pool, email, reason, smtp_code).await?)
}
async fn remove_suppression(&self, email: &str) -> Result<bool, StoreError> {
Ok(queue::remove_suppression(&self.pool, email).await?)
}
async fn list_suppressions(
&self,
limit: i64,
) -> Result<Vec<(String, String, Option<i32>, i64)>, StoreError> {
Ok(queue::list_suppressions(&self.pool, limit).await?)
}
}
#[derive(Clone)]
pub struct KevyNotifier {
store: Store,
}
impl std::fmt::Debug for KevyNotifier {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("KevyNotifier").finish_non_exhaustive()
}
}
impl KevyNotifier {
pub fn new(store: Store) -> Self {
Self { store }
}
}
#[async_trait::async_trait]
impl Notifier for KevyNotifier {
async fn notify(&self) {
let _ = self.store.publish(b"queue:notify", b"1");
}
async fn wait(&self) {
let store = self.store.clone();
let _ = tokio::task::spawn_blocking(move || {
let sub = store.subscribe(&[b"queue:notify"]);
loop {
match sub.recv() {
Ok(PubsubFrame::Message { .. } | PubsubFrame::Pmessage { .. }) => break,
Ok(_) => continue, Err(_) => break, }
}
})
.await;
}
}