use std::{fmt::Debug, future::Future, pin::Pin}; use futures_core::{future::BoxFuture, stream::BoxStream}; use rsmq_async::{RedisBytes, RsmqConnection}; use sqlx::{Database, Postgres, Transaction}; pub enum RedisOp { SendMessage(RedisBytes, Option>, String), DeleteMessage(String, String), } unsafe impl Send for RedisOp {} impl RedisOp { pub async fn apply(self, rsmq: &mut R) -> Result<(), rsmq_async::RsmqError> { match self { RedisOp::SendMessage(bytes, time, queue) => { rsmq.send_message( &queue, bytes, time.map(|t| (t - chrono::Utc::now()).num_seconds()) .and_then(|e| e.try_into().ok()), ) .await?; } RedisOp::DeleteMessage(id, queue) => { rsmq.delete_message(&queue, &id).await?; } }; Ok(()) } } pub struct RedisTransaction { rsmq: R, queued_ops: Vec, } impl From for RedisTransaction { fn from(value: R) -> Self { Self { rsmq: value, queued_ops: Vec::new() } } } impl RedisTransaction { pub async fn commit(self) -> Result<(), rsmq_async::RsmqError> { let mut rsmq = self.rsmq; for op in self.queued_ops { op.apply(&mut rsmq).await?; } Ok(()) } pub fn send_message>( &mut self, bytes: E, delay_until: Option>, queue: String, ) { self.queued_ops .push(RedisOp::SendMessage(bytes.into(), delay_until, queue)) } pub fn delete_message(&mut self, id: String, queue: String) { self.queued_ops.push(RedisOp::DeleteMessage(id, queue)) } } pub struct QueueTransaction<'c, R: RsmqConnection> { pub rsmq: Option>, transaction: Transaction<'c, Postgres>, } impl<'c, R: RsmqConnection> From<(Option, Transaction<'c, Postgres>)> for QueueTransaction<'c, R> { fn from(value: (Option, Transaction<'c, Postgres>)) -> Self { Self { rsmq: value.0.map(|e| e.into()), transaction: value.1 } } } impl<'c, R: RsmqConnection> Debug for QueueTransaction<'c, R> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("QueueTransaction") .field("rsmq", &self.rsmq.as_ref().map(|_| ())) // do not require R: Debug .field("transaction", &self.transaction) .finish() } } impl<'c, R: RsmqConnection> QueueTransaction<'c, R> { pub async fn commit(self) -> Result<(), windmill_common::error::Error> { self.transaction.commit().await?; if let Some(rsmq) = self.rsmq { rsmq.commit().await.map_err(|e| anyhow::anyhow!(e))?; } Ok(()) } pub fn transaction_mut<'a>(&'a mut self) -> &'a mut Transaction<'c, Postgres> { &mut self.transaction } } impl<'c, 'b, R: RsmqConnection + Send> sqlx::Executor<'b> for &'b mut QueueTransaction<'c, R> { type Database = Postgres; fn fetch_many<'e, 'q: 'e, E: 'q>( self, query: E, ) -> BoxStream< 'e, Result< sqlx::Either< ::QueryResult, ::Row, >, sqlx::Error, >, > where 'b: 'e, E: sqlx::Execute<'q, Self::Database>, { self.transaction.fetch_many(query) } fn fetch_optional<'e, 'q: 'e, E: 'q>( self, query: E, ) -> BoxFuture<'e, Result::Row>, sqlx::Error>> where 'b: 'e, E: sqlx::Execute<'q, Self::Database>, { self.transaction.fetch_optional(query) } fn prepare_with<'e, 'q>( self, sql: &'q str, parameters: &'e [::TypeInfo], ) -> Pin< Box< dyn Future::Statement<'q>, sqlx::Error>> + Send + 'e, >, > where 'q: 'e, 'c: 'e, 'b: 'e, { self.transaction.prepare_with(sql, parameters) } fn describe<'e, 'q: 'e>( self, sql: &'q str, ) -> BoxFuture<'e, Result, sqlx::Error>> where 'b: 'e, { self.transaction.describe(sql) } }