expose store_offset and query_offset in consumer (#203)

* expose store_offset

* replacing assert with better return value

* adding an OffsetNotFound unit test

* adding a non-fake offset store/query offset

* implemented code review suggestion

---------

Co-authored-by: DanielePalaia <[email protected]>

Author: DanielePalaia

src/consumer.rs

@@ -14,9 +14,12 @@ use std::{
 use rabbitmq_stream_protocol::{
     commands::subscribe::OffsetSpecification, message::Message, ResponseKind,
 };
+
 use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tracing::trace;
 
+use crate::error::ConsumerStoreOffsetError;
+
 use crate::{
     client::{MessageHandler, MessageResult},
     error::{ConsumerCloseError, ConsumerCreateError, ConsumerDeliveryError},
@@ -28,6 +31,8 @@ use rand::{seq::SliceRandom, SeedableRng};
 
 /// API for consuming RabbitMQ stream messages
 pub struct Consumer {
+    // Mandatory in case of manual offset tracking
+    name: Option<String>,
     receiver: Receiver<Result<Delivery, ConsumerDeliveryError>>,
     internal: Arc<ConsumerInternal>,
 }
@@ -50,6 +55,7 @@ impl ConsumerInternal {
 
 /// Builder for [`Consumer`]
 pub struct ConsumerBuilder {
+    pub(crate) consumer_name: Option<String>,
     pub(crate) environment: Environment,
     pub(crate) offset_specification: OffsetSpecification,
 }
@@ -110,6 +116,7 @@ impl ConsumerBuilder {
 
         if response.is_ok() {
             Ok(Consumer {
+                name: self.consumer_name,
                 receiver: rx,
                 internal: consumer,
             })
@@ -125,6 +132,11 @@ impl ConsumerBuilder {
         self.offset_specification = offset_specification;
         self
     }
+
+    pub fn name(mut self, consumer_name: &str) -> Self {
+        self.consumer_name = Some(String::from(consumer_name));
+        self
+    }
 }
 
 impl Consumer {
@@ -137,6 +149,30 @@ impl Consumer {
     pub fn is_closed(&self) -> bool {
         self.internal.is_closed()
     }
+
+    pub async fn store_offset(&self, offset: u64) -> Result<(), ConsumerStoreOffsetError> {
+        if let Some(name) = &self.name {
+            self.internal
+                .client
+                .store_offset(name.as_str(), self.internal.stream.as_str(), offset)
+                .await
+                .map(Ok)?
+        } else {
+            Err(ConsumerStoreOffsetError::NameMissing)
+        }
+    }
+
+    pub async fn query_offset(&self) -> Result<u64, ConsumerStoreOffsetError> {
+        if let Some(name) = &self.name {
+            self.internal
+                .client
+                .query_offset(name.clone(), self.internal.stream.as_str())
+                .await
+                .map(Ok)?
+        } else {
+            Err(ConsumerStoreOffsetError::NameMissing)
+        }
+    }
 }
 
 impl Stream for Consumer {

src/environment.rs

@@ -51,6 +51,7 @@ impl Environment {
     /// Returns a builder for creating a consumer
     pub fn consumer(&self) -> ConsumerBuilder {
         ConsumerBuilder {
+            consumer_name: None,
             environment: self.clone(),
             offset_specification: OffsetSpecification::Next,
         }

src/error.rs

@@ -23,6 +23,14 @@ pub enum ClientError {
     RequestError(ResponseCode),
 }
 
+#[derive(Error, Debug)]
+pub enum ConsumerStoreOffsetError {
+    #[error("Failed to store offset, missing consumer name")]
+    NameMissing,
+    #[error(transparent)]
+    Client(#[from] ClientError),
+}
+
 #[derive(Error, Debug)]
 pub enum ProtocolError {
     #[error("Encode Error {0:?}")]

tests/integration/consumer_test.rs

@@ -4,11 +4,16 @@ use crate::common::TestEnvironment;
 use fake::{Fake, Faker};
 use futures::StreamExt;
 use rabbitmq_stream_client::{
-    error::{ConsumerCloseError, ConsumerDeliveryError, ProducerCloseError},
+    error::{
+        ClientError, ConsumerCloseError, ConsumerDeliveryError, ConsumerStoreOffsetError,
+        ProducerCloseError,
+    },
     types::{Delivery, Message, OffsetSpecification},
     Consumer, NoDedup, Producer,
 };
 
+use rabbitmq_stream_protocol::ResponseCode;
+
 #[tokio::test(flavor = "multi_thread")]
 async fn consumer_test() {
     let env = TestEnvironment::create().await;
@@ -164,3 +169,130 @@ async fn consumer_create_stream_not_existing_error() {
         _ => panic!("Should be StreamNotFound error"),
     }
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn consumer_store_and_query_offset_test() {
+    let env = TestEnvironment::create().await;
+    let consumer = env
+        .env
+        .consumer()
+        .name("test-name")
+        .offset(OffsetSpecification::Next)
+        .build(&env.stream)
+        .await
+        .unwrap();
+
+    let offset: u64 = Faker.fake();
+
+    consumer.store_offset(offset).await.unwrap();
+
+    let response = consumer.query_offset().await.unwrap();
+
+    assert_eq!(offset, response);
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn consumer_query_missing_offset_test() {
+    let env = TestEnvironment::create().await;
+    let consumer = env
+        .env
+        .consumer()
+        .name("test-name-new")
+        .offset(OffsetSpecification::Next)
+        .build(&env.stream)
+        .await
+        .unwrap();
+
+    assert!(matches!(
+        consumer.query_offset().await,
+        Err(ConsumerStoreOffsetError::Client(ClientError::RequestError(
+            ResponseCode::OffsetNotFound
+        ))),
+    ))
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn consumer_store_and_query_offset_missing_name_test() {
+    let env = TestEnvironment::create().await;
+    let consumer = env
+        .env
+        .consumer()
+        .offset(OffsetSpecification::Next)
+        .build(&env.stream)
+        .await
+        .unwrap();
+
+    let offset: u64 = Faker.fake();
+
+    assert!(matches!(
+        consumer.store_offset(offset).await,
+        Err(ConsumerStoreOffsetError::NameMissing),
+    ));
+
+    assert!(matches!(
+        consumer.query_offset().await,
+        Err(ConsumerStoreOffsetError::NameMissing),
+    ));
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn consumer_test_with_store_offset() {
+    let env = TestEnvironment::create().await;
+    let reference: String = Faker.fake();
+    let offset_to_store = 4;
+
+    let message_count = 10;
+    let mut producer = env
+        .env
+        .producer()
+        .name(&reference)
+        .build(&env.stream)
+        .await
+        .unwrap();
+
+    let mut consumer_store = env
+        .env
+        .consumer()
+        .offset(OffsetSpecification::Next)
+        .name("consumer-1")
+        .build(&env.stream)
+        .await
+        .unwrap();
+
+    for n in 0..message_count {
+        let _ = producer
+            .send_with_confirm(Message::builder().body(format!("message{}", n)).build())
+            .await
+            .unwrap();
+    }
+
+    for i in 0..message_count {
+        let delivery = consumer_store.next().await.unwrap();
+
+        // Store an offset
+        if i == offset_to_store {
+            //Store the 5th element produced
+            let result = consumer_store
+                .store_offset(delivery.unwrap().offset())
+                .await;
+        }
+    }
+
+    consumer_store.handle().close().await.unwrap();
+
+    let mut consumer_query = env
+        .env
+        .consumer()
+        .offset(OffsetSpecification::First)
+        .name("consumer-1")
+        .build(&env.stream)
+        .await
+        .unwrap();
+
+    let offset = consumer_query.query_offset().await.unwrap();
+
+    assert!(offset == offset_to_store);
+
+    consumer_query.handle().close().await.unwrap();
+    producer.close().await.unwrap();
+}