intercom_rs/jetstream/

context.rs

//! JetStream context for managing streams and publishing messages.

use std::marker::PhantomData;

use serde::{de::DeserializeOwned, Serialize};

use crate::{codec::CodecType, error::Result};

use super::stream::{Stream, StreamBuilder};

/// A typed JetStream context with configurable codec.
///
/// Provides methods for creating streams, publishing messages, and managing
/// JetStream resources.
///
/// # Type Parameters
///
/// * `C` - The codec type used for message serialization
#[derive(Clone)]
pub struct JetStreamContext<C: CodecType> {
    inner: async_nats::jetstream::Context,
    _codec: PhantomData<C>,
}

impl<C: CodecType> JetStreamContext<C> {
    /// Create a new JetStream context.
    pub(crate) fn new(inner: async_nats::jetstream::Context) -> Self {
        Self {
            inner,
            _codec: PhantomData,
        }
    }

    /// Get the underlying async-nats JetStream context.
    pub fn inner(&self) -> &async_nats::jetstream::Context {
        &self.inner
    }

    /// Create a stream builder for configuring and creating a stream.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use intercom::{Client, MsgPackCodec};
    ///
    /// # async fn example() -> intercom::Result<()> {
    /// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
    /// let jetstream = client.jetstream();
    ///
    /// let stream = jetstream
    ///     .stream_builder("my-stream")
    ///     .subjects(vec!["events.>".to_string()])
    ///     .max_messages(1_000_000)
    ///     .create()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn stream_builder(&self, name: &str) -> StreamBuilder<C> {
        StreamBuilder::new(self.inner.clone(), name.to_string())
    }

    /// Get an existing stream by name.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use intercom::{Client, MsgPackCodec};
    ///
    /// # async fn example() -> intercom::Result<()> {
    /// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
    /// let jetstream = client.jetstream();
    ///
    /// let stream = jetstream.get_stream("my-stream").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn get_stream(&self, name: &str) -> Result<Stream<C>> {
        let inner = self
            .inner
            .get_stream(name)
            .await
            .map_err(|e| crate::error::Error::JetStreamStream(e.to_string()))?;
        Ok(Stream::new(inner))
    }

    /// Delete a stream by name.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use intercom::{Client, MsgPackCodec};
    ///
    /// # async fn example() -> intercom::Result<()> {
    /// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
    /// let jetstream = client.jetstream();
    ///
    /// jetstream.delete_stream("my-stream").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn delete_stream(&self, name: &str) -> Result<()> {
        self.inner
            .delete_stream(name)
            .await
            .map_err(|e| crate::error::Error::JetStreamStream(e.to_string()))?;
        Ok(())
    }

    /// Publish a typed message to a JetStream subject.
    ///
    /// # Type Parameters
    ///
    /// * `T` - The message type (must implement Serialize)
    ///
    /// # Example
    ///
    /// ```no_run
    /// use intercom::{Client, MsgPackCodec};
    /// use serde::{Deserialize, Serialize};
    ///
    /// #[derive(Serialize, Deserialize)]
    /// struct Event { id: u64, data: String }
    ///
    /// # async fn example() -> intercom::Result<()> {
    /// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
    /// let jetstream = client.jetstream();
    ///
    /// let ack = jetstream.publish::<Event>("events.user", &Event {
    ///     id: 1,
    ///     data: "created".to_string(),
    /// }).await?;
    ///
    /// println!("Published to stream: {}, seq: {}", ack.stream, ack.sequence);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn publish<T: Serialize>(&self, subject: &str, message: &T) -> Result<PublishAck> {
        let data = C::encode(message)?;
        let ack = self
            .inner
            .publish(subject.to_string(), data.into())
            .await
            .map_err(|e| crate::error::Error::JetStream(e.to_string()))?
            .await
            .map_err(|e| crate::error::Error::JetStream(e.to_string()))?;
        Ok(PublishAck {
            stream: ack.stream.to_string(),
            sequence: ack.sequence,
            duplicate: ack.duplicate,
        })
    }

    /// Publish a typed message and get the ack future immediately.
    ///
    /// This allows for more efficient batching by not waiting for the ack.
    ///
    /// # Type Parameters
    ///
    /// * `T` - The message type (must implement Serialize)
    ///
    /// # Example
    ///
    /// ```no_run
    /// use intercom::{Client, MsgPackCodec};
    /// use serde::{Deserialize, Serialize};
    ///
    /// #[derive(Serialize, Deserialize)]
    /// struct Event { id: u64 }
    ///
    /// # async fn example() -> intercom::Result<()> {
    /// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
    /// let jetstream = client.jetstream();
    ///
    /// // Publish without waiting for ack
    /// let ack_future = jetstream.publish_async::<Event>("events.user", &Event { id: 1 }).await?;
    ///
    /// // Do other work...
    ///
    /// // Now wait for the ack
    /// let ack = ack_future.await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn publish_async<T: Serialize>(
        &self,
        subject: &str,
        message: &T,
    ) -> Result<PublishAckFuture> {
        let data = C::encode(message)?;
        let future = self
            .inner
            .publish(subject.to_string(), data.into())
            .await
            .map_err(|e| crate::error::Error::JetStream(e.to_string()))?;
        Ok(PublishAckFuture { inner: future })
    }

    /// Request-reply pattern over JetStream.
    ///
    /// # Type Parameters
    ///
    /// * `T` - The request message type
    /// * `R` - The response message type
    pub async fn request<T: Serialize, R: DeserializeOwned>(
        &self,
        subject: &str,
        message: &T,
    ) -> Result<R> {
        let data = C::encode(message)?;
        let response: async_nats::Message = self
            .inner
            .request(subject.to_string(), &bytes::Bytes::from(data))
            .await
            .map_err(|e| crate::error::Error::JetStream(e.to_string()))?;
        C::decode(&response.payload)
    }
}

/// Acknowledgment from a JetStream publish.
#[derive(Debug, Clone)]
pub struct PublishAck {
    /// The stream name.
    pub stream: String,
    /// The sequence number.
    pub sequence: u64,
    /// Whether this was a duplicate.
    pub duplicate: bool,
}

/// A future that resolves to a publish acknowledgment.
pub struct PublishAckFuture {
    inner: async_nats::jetstream::context::PublishAckFuture,
}

impl PublishAckFuture {
    /// Wait for the acknowledgment.
    pub async fn await_ack(self) -> Result<PublishAck> {
        let ack = self
            .inner
            .await
            .map_err(|e| crate::error::Error::JetStream(e.to_string()))?;
        Ok(PublishAck {
            stream: ack.stream.to_string(),
            sequence: ack.sequence,
            duplicate: ack.duplicate,
        })
    }
}

impl std::future::IntoFuture for PublishAckFuture {
    type Output = Result<PublishAck>;
    type IntoFuture = std::pin::Pin<Box<dyn std::future::Future<Output = Self::Output> + Send>>;

    fn into_future(self) -> Self::IntoFuture {
        Box::pin(self.await_ack())
    }
}