use std::process::ExitCode;
use ruststream::OutgoingMessage;
use ruststream::runtime::cli::run_main;
use ruststream::schemars::JsonSchema;
use ruststream_sqlx::dialect::Column;
use ruststream_sqlx::keys::InboxContext;
use ruststream_sqlx::spec::{Attempt, Payload};
use ruststream_sqlx::{AttemptRow, InboxSpec, InboxTable, PayloadRow};
use serde::{Deserialize, Serialize};
use sqlx::PgConnection;
const MIGRATION: &str = "
CREATE TABLE IF NOT EXISTS email_jobs (
job_id BIGSERIAL PRIMARY KEY,
name TEXT NOT NULL,
attempt SMALLINT NOT NULL DEFAULT 1,
payload BYTEA NOT NULL
);
CREATE INDEX IF NOT EXISTS email_jobs_claim ON email_jobs (name, job_id);
";
use ruststream_sqlx::prelude::*;
use sqlx::Postgres;
#[derive(sqlx::FromRow)]
struct SendEmail {
job_id: i64,
attempt: i16,
payload: Vec<u8>,
}
impl InboxTable for SendEmail {
type Id = i64;
type Table = InboxSpec<(Attempt, Payload)>;
const TABLE: Self::Table = InboxSpec::new("email_jobs", Column::new("job_id").generated())
.group(Column::new("name"))
.attempt(Column::new("attempt").generated())
.payload(Column::new("payload"));
fn id(&self) -> &i64 {
&self.job_id
}
}
impl AttemptRow for SendEmail {
type Attempt = i16;
fn attempt(&self) -> &i16 {
&self.attempt
}
}
impl PayloadRow for SendEmail {
type Column = Vec<u8>;
fn payload(&self) -> &[u8] {
&self.payload
}
}
use std::num::NonZeroUsize;
use std::time::Duration;
use sqlx::Connection;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
const WORKERS: NonZeroUsize = nonzero!(8);
fn pool() -> PgPoolOptions {
let size = u32::try_from(WORKERS.get() + 1).unwrap_or(u32::MAX);
PgPoolOptions::new()
.max_connections(size)
.acquire_timeout(Duration::from_secs(5))
}
struct Mailer;
impl Handle<Email, (), (), InboxContext> for Mailer {
async fn handle(
&self,
email: &Email,
_outs: &(),
ctx: &mut Context<'_, InboxContext>,
) -> Result<(), HandlerOutcome> {
let Ctx(attempt) = Ctx::<keys::Attempt>::from_context(ctx).await?;
tracing::info!(to = %email.to, ?attempt, "sending");
Ok(())
}
}
fn service(options: PgConnectOptions) -> RustStream {
let migrate = options.clone();
let broker = SqlxBroker::from_connect_options(options, pool()).route::<SendEmail>("emails");
RustStream::new(AppInfo::new("mailer", "0.1.0"))
.on_startup(async move |()| {
let mut conn = PgConnection::connect_with(&migrate).await?;
sqlx::raw_sql(MIGRATION).execute(&mut conn).await?;
conn.close().await
})
.with_broker_labeled("postgres", broker, |b| {
let send = subscriber("emails", Mailer)
.map_source(|_| InboxQueue::<SendEmail>::new("emails"))
.workers(WORKERS);
b.include(send.build())
.max_attempts(nonzero!(5u32))
.dead_letter("emails.dead");
})
}
fn main() -> ExitCode {
run_main(|| service(PgConnectOptions::new()))
}
impl Publish<Postgres> for SendEmail {
async fn publish(
conn: &mut PgConnection,
message: &OutgoingMessage<'_>,
) -> Result<(), sqlx::Error> {
sqlx::query("INSERT INTO email_jobs (name, payload) VALUES ($1, $2)")
.bind(message.name())
.bind(message.payload())
.execute(conn)
.await?;
Ok(())
}
}
#[derive(Debug, PartialEq, Deserialize, Serialize, JsonSchema)]
#[schemars(crate = "ruststream::schemars")]
struct Email {
to: String,
}
impl OutgoingDestination for Email {
type Form = FixedName;
const DESTINATION: &'static str = "emails";
}
impl MessageHeaders for Email {
type Contract = NoHeaders;
}
#[cfg(test)]
#[path = "../../../tests/live/mod.rs"]
mod live;
#[cfg(test)]
mod tests {
use std::error::Error;
use ruststream::asyncapi::build_spec;
use ruststream::testing::TestApp;
use ruststream_sqlx::prelude::*;
use sqlx::Postgres;
use sqlx::postgres::PgConnectOptions;
use super::{Email, live, service};
const PASSWORD: &str = "hunter2";
#[tokio::test]
async fn the_document_names_the_server_and_never_its_password() -> Result<(), Box<dyn Error>> {
let options: PgConnectOptions =
format!("postgres://mailer:{PASSWORD}@db.internal:5432/mail").parse()?;
let document = build_spec(&service(options)).to_json()?;
assert!(document.contains("db.internal:5432"), "{document}");
assert!(!document.contains(PASSWORD), "{document}");
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_table_of_an_older_release_stops_the_startup() -> Result<(), Box<dyn Error>> {
let Some(db) = live::postgres::database().await else {
return Ok(());
};
sqlx::raw_sql("ALTER TABLE email_jobs DROP COLUMN attempt")
.execute(&db.pool)
.await?;
let started = TestApp::start_live(service(db.options())).await;
let error = started
.err()
.map(|error| error.to_string())
.unwrap_or_default();
assert!(error.contains("email_jobs"), "{error}");
db.finish().await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn the_checked_statements_serve_the_queue() -> Result<(), Box<dyn Error>> {
let Some(db) = live::postgres::database().await else {
return Ok(());
};
let tb = TestApp::start_live(service(db.options())).await?;
let email = Email {
to: "ops@example.com".to_owned(),
};
tb.broker::<SqlxBroker<Postgres>>()
.message(&email)
.publish()
.await?;
tb.broker::<SqlxBroker<Postgres>>()
.subscriber("emails")
.assert_called_once()
.with(&email)
.settled(HandlerOutcome::ack());
tb.shutdown().await?;
db.finish().await;
Ok(())
}
}