cloudlysis/projection/src/stream/jetstream.rs

use crate::config::Settings;
use crate::types::ProjectionError;
use async_nats::jetstream::{
    self, consumer::pull::Config as PullConfig, consumer::AckPolicy, consumer::DeliverPolicy,
    consumer::ReplayPolicy, stream::Config as StreamConfig,
};

#[derive(Debug, Clone)]
pub struct JetStreamClient {
    stream: jetstream::stream::Stream,
    consumer: jetstream::consumer::PullConsumer,
}

#[derive(Debug, Clone)]
pub struct ConsumerOptions {
    pub durable_name: String,
    pub filter_subject: String,
    pub deliver_policy: DeliverPolicy,
}

impl JetStreamClient {
    pub async fn connect(settings: &Settings) -> Result<Self, ProjectionError> {
        let filter_subject = settings
            .subject_filters
            .first()
            .cloned()
            .unwrap_or_else(|| shared::NATS_SUBJECT_AGGREGATE_EVENTS_ALL.to_string());

        let options = ConsumerOptions {
            durable_name: settings.durable_name.clone(),
            filter_subject,
            deliver_policy: DeliverPolicy::All,
        };

        Self::connect_with(settings, options).await
    }

    pub async fn connect_with(
        settings: &Settings,
        options: ConsumerOptions,
    ) -> Result<Self, ProjectionError> {
        let client = async_nats::connect(&settings.nats_url).await.map_err(|e| {
            ProjectionError::StreamError(format!("Failed to connect to NATS: {}", e))
        })?;

        let jetstream = jetstream::new(client);

        let expected = stream_policy_config(&settings.stream_name);
        let mut stream = jetstream
            .get_or_create_stream(expected.clone())
            .await
            .map_err(|e| ProjectionError::StreamError(format!("Stream error: {}", e)))?;

        let info = stream
            .info()
            .await
            .map_err(|e| ProjectionError::StreamError(format!("Stream info error: {}", e)))?;
        validate_stream_config(&expected, &info.config)?;

        let policy = shared::consumer_policy_from_parts(
            settings.ack_timeout_ms,
            settings.max_in_flight,
            settings.max_deliver,
        );
        let consumer_config = PullConfig {
            durable_name: Some(options.durable_name.clone()),
            deliver_policy: options.deliver_policy,
            ack_policy: AckPolicy::Explicit,
            ack_wait: policy.ack_wait,
            filter_subject: options.filter_subject,
            replay_policy: ReplayPolicy::Instant,
            max_ack_pending: policy.max_ack_pending,
            max_deliver: policy.max_deliver,
            ..Default::default()
        };

        let consumer = stream
            .get_or_create_consumer(&options.durable_name, consumer_config)
            .await
            .map_err(|e| {
                ProjectionError::StreamError(format!("Consumer creation failed: {}", e))
            })?;

        Ok(Self { stream, consumer })
    }

    pub async fn messages(&self) -> Result<jetstream::consumer::pull::Stream, ProjectionError> {
        self.consumer
            .messages()
            .await
            .map_err(|e| ProjectionError::StreamError(format!("Message stream error: {}", e)))
    }

    pub async fn stream_last_sequence(&self) -> Result<u64, ProjectionError> {
        let mut stream = self.stream.clone();
        let info = stream
            .info()
            .await
            .map_err(|e| ProjectionError::StreamError(e.to_string()))?;
        Ok(info.state.last_sequence)
    }
}

fn stream_policy_config(name: &str) -> StreamConfig {
    let policy = shared::stream_policy_defaults(
        name.to_string(),
        vec![shared::NATS_SUBJECT_AGGREGATE_EVENTS_ALL.to_string()],
    );
    StreamConfig {
        name: policy.name,
        subjects: policy.subjects,
        max_messages: policy.max_messages,
        max_bytes: policy.max_bytes,
        max_age: policy.max_age,
        duplicate_window: policy.duplicate_window,
        ..Default::default()
    }
}

fn validate_stream_config(
    expected: &StreamConfig,
    actual: &StreamConfig,
) -> Result<(), ProjectionError> {
    let expected = shared::stream_policy_from_parts(
        expected.name.as_str(),
        expected.subjects.clone(),
        expected.max_messages,
        expected.max_bytes,
        expected.max_age,
        expected.duplicate_window,
    );
    let actual = shared::stream_policy_from_parts(
        actual.name.as_str(),
        actual.subjects.clone(),
        actual.max_messages,
        actual.max_bytes,
        actual.max_age,
        actual.duplicate_window,
    );
    shared::validate_stream_policy(&expected, &actual)
        .map_err(|e| ProjectionError::StreamError(e.to_string()))
}