nstreams-rabbit 0.1.0

use std::sync::Arc;

use async_trait::async_trait;
use futures::StreamExt;
use lapin::options::{
    BasicAckOptions, BasicConsumeOptions, BasicPublishOptions, BasicQosOptions,
    ExchangeDeclareOptions, QueueBindOptions, QueueDeclareOptions,
};
use lapin::types::{AMQPValue, FieldTable};
use lapin::{Connection, ConnectionProperties, Consumer, ExchangeKind};
use nstreams_core::namespace::{Namespace, NamespaceResources};
use nstreams_core::stream::WriteQueueBackend;
use nstreams_core::worker::{QueueConsumer, QueueMessage};
use thiserror::Error;
use tracing::{debug, info};

#[derive(Debug, Error)]
pub enum RabbitQueueError {
    #[error(transparent)]
    Lapin(#[from] lapin::Error),
}

impl From<RabbitQueueError> for nstreams_core::Error {
    fn from(value: RabbitQueueError) -> Self {
        Self::WriteQueue(value.to_string())
    }
}

const WRITE_EXCHANGE: &str = "nstreams.write";

#[derive(Clone)]
pub struct QueueConsumerConfig {
    pub amqp_url: String,
    pub prefetch: u16,
}

/// RabbitMQ queue backend for write-side event publishing.
#[derive(Clone)]
pub struct LapinWriteQueue {
    connection: Arc<Connection>,
}

impl LapinWriteQueue {
    pub async fn connect(amqp_url: &str) -> Result<Self, RabbitQueueError> {
        let connection = Connection::connect(amqp_url, ConnectionProperties::default()).await?;
        let queue = Self {
            connection: Arc::new(connection),
        };
        queue.declare_exchange().await?;
        Ok(queue)
    }

    async fn declare_exchange(&self) -> Result<(), RabbitQueueError> {
        let channel = self.connection.create_channel().await?;
        channel
            .exchange_declare(
                WRITE_EXCHANGE,
                ExchangeKind::Topic,
                ExchangeDeclareOptions {
                    durable: true,
                    ..Default::default()
                },
                FieldTable::default(),
            )
            .await?;
        Ok(())
    }

    async fn ensure_queue(&self, namespace: &str) -> Result<(), RabbitQueueError> {
        let resources = NamespaceResources::new(&nstreams_core::namespace::StreamNamespace::from(
            namespace,
        ));
        let channel = self.connection.create_channel().await?;
        channel
            .queue_declare(
                &resources.write_queue,
                QueueDeclareOptions {
                    durable: true,
                    ..Default::default()
                },
                FieldTable::default(),
            )
            .await?;
        channel
            .queue_bind(
                &resources.write_queue,
                WRITE_EXCHANGE,
                namespace,
                QueueBindOptions::default(),
                FieldTable::default(),
            )
            .await?;
        Ok(())
    }

    pub async fn create_consumer(
        &self,
        config: QueueConsumerConfig,
    ) -> Result<LapinQueueConsumer, RabbitQueueError> {
        let connection = Connection::connect(&config.amqp_url, ConnectionProperties::default())
            .await?;
        let channel = connection.create_channel().await?;
        channel
            .basic_qos(config.prefetch, BasicQosOptions::default())
            .await?;

        let mut args = FieldTable::default();
        args.insert("x-queue-type".into(), AMQPValue::LongString("quorum".into()));

        channel
            .queue_declare(
                "nstreams.write.worker",
                QueueDeclareOptions {
                    durable: true,
                    ..Default::default()
                },
                args,
            )
            .await?;

        channel
            .queue_bind(
                "nstreams.write.worker",
                WRITE_EXCHANGE,
                "#",
                QueueBindOptions::default(),
                FieldTable::default(),
            )
            .await?;

        let consumer = channel
            .basic_consume(
                "nstreams.write.worker",
                "nstreams-worker",
                BasicConsumeOptions::default(),
                FieldTable::default(),
            )
            .await?;

        info!("rabbitmq write queue consumer ready");
        Ok(LapinQueueConsumer { consumer })
    }
}

#[async_trait]
impl WriteQueueBackend for LapinWriteQueue {
    async fn publish<N: Namespace>(&self, namespace: &N, payload: &[u8]) -> nstreams_core::Result<()> {
        self.ensure_queue(namespace.as_str())
            .await
            .map_err(|error| nstreams_core::Error::WriteQueue(error.to_string()))?;
        let channel = self
            .connection
            .create_channel()
            .await
            .map_err(|error| nstreams_core::Error::WriteQueue(error.to_string()))?;
        channel
            .basic_publish(
                WRITE_EXCHANGE,
                namespace.as_str(),
                BasicPublishOptions::default(),
                payload,
                lapin::BasicProperties::default().with_delivery_mode(2),
            )
            .await
            .map_err(|error| nstreams_core::Error::WriteQueue(error.to_string()))?
            .await
            .map_err(|error| nstreams_core::Error::WriteQueue(error.to_string()))?;

        debug!(namespace = namespace.as_str(), "published to write queue");
        Ok(())
    }
}

pub struct LapinQueueConsumer {
    consumer: Consumer,
}

#[async_trait]
impl QueueConsumer for LapinQueueConsumer {
    async fn next(&mut self) -> Option<std::result::Result<QueueMessage, nstreams_core::Error>> {
        loop {
            match self.consumer.next().await {
                Some(Ok(delivery)) => {
                    let routing_key = delivery.routing_key.to_string();
                    let namespace = if routing_key.is_empty() {
                        String::from("unknown")
                    } else {
                        routing_key
                    };

                    if let Err(error) = delivery.ack(BasicAckOptions::default()).await {
                        return Some(Err(nstreams_core::Error::WriteQueue(error.to_string())));
                    }

                    return Some(Ok(QueueMessage {
                        namespace,
                        payload: delivery.data,
                    }));
                }
                Some(Err(error)) => {
                    return Some(Err(nstreams_core::Error::WriteQueue(error.to_string())));
                }
                None => return None,
            }
        }
    }
}