Added pull consumer fetch function.

Author: s3rius

python/natsrpy/_natsrpy_rs/js/__init__.pyi

@@ -1,3 +1,5 @@
+from typing import Any
+
 from .managers import KVManager, StreamsManager
 
 class JetStream:
@@ -14,3 +16,14 @@ class JetStream:
     def kv(self) -> KVManager: ...
     @property
     def streams(self) -> StreamsManager: ...
+
+class JetStreamMessage:
+    @property
+    def subject(self) -> str: ...
+    @property
+    def reply(self) -> str | None: ...
+    @property
+    def payload(self) -> bytes: ...
+    @property
+    def headers(self) -> dict[str, Any]: ...
+    async def ack(self) -> None: ...

python/natsrpy/_natsrpy_rs/js/consumers.pyi

@@ -1,5 +1,7 @@
 from datetime import timedelta
 
+from natsrpy._natsrpy_rs.js import JetStreamMessage
+
 class DeliverPolicy:
     ALL: DeliverPolicy
     LAST: DeliverPolicy
@@ -146,4 +148,16 @@ class PushConsumerConfig:
     ) -> None: ...
 
 class PushConsumer: ...
-class PullConsumer: ...
+
+class PullConsumer:
+    async def fetch(
+        self,
+        max_messages: int | None = None,
+        group: str | None = None,
+        priority: int | None = None,
+        max_bytes: int | None = None,
+        heartbeat: timedelta | None = None,
+        expires: timedelta | None = None,
+        min_pending: int | None = None,
+        min_ack_pending: int | None = None,
+    ) -> list[JetStreamMessage]: ...

src/exceptions/rust_err.rs

@@ -64,6 +64,8 @@ pub enum NatsrpyError {
     PullMessageError(#[from] async_nats::jetstream::consumer::pull::MessagesError),
     #[error(transparent)]
     PullConsumerError(#[from] async_nats::jetstream::stream::ConsumerError),
+    #[error(transparent)]
+    PullConsumerBatchError(#[from] async_nats::jetstream::consumer::pull::BatchError),
 }
 
 impl From<NatsrpyError> for pyo3::PyErr {

src/js/consumers/pull/consumer.rs

@@ -1,7 +1,11 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
+use futures_util::StreamExt;
+use pyo3::{Bound, PyAny, Python};
 use tokio::sync::RwLock;
 
+use crate::{exceptions::rust_err::NatsrpyResult, utils::natsrpy_future};
+
 type NatsPullConsumer =
     async_nats::jetstream::consumer::Consumer<async_nats::jetstream::consumer::pull::Config>;
 
@@ -20,13 +24,68 @@ impl PullConsumer {
     }
 }
 
-#[pyo3::pyclass]
-pub struct PullMessageIterator {
-    inner: Arc<RwLock<async_nats::jetstream::consumer::pull::Batch>>,
-}
-
-#[pyo3::pymethods]
-impl PullConsumer {}
-
 #[pyo3::pymethods]
-impl PullMessageIterator {}
+impl PullConsumer {
+    #[pyo3(signature=(
+        max_messages=None,
+        group=None,
+        priority=None,
+        max_bytes=None,
+        heartbeat=None,
+        expires=None,
+        min_pending=None,
+        min_ack_pending=None,
+    ))]
+    pub fn fetch<'py>(
+        &self,
+        py: Python<'py>,
+        max_messages: Option<usize>,
+        group: Option<String>,
+        priority: Option<usize>,
+        max_bytes: Option<usize>,
+        heartbeat: Option<Duration>,
+        expires: Option<Duration>,
+        min_pending: Option<usize>,
+        min_ack_pending: Option<usize>,
+    ) -> NatsrpyResult<Bound<'py, PyAny>> {
+        let ctx = self.consumer.clone();
+        #[allow(clippy::significant_drop_tightening)]
+        natsrpy_future(py, async move {
+            // Because we borrow created value
+            // later for modifications.
+            let consumer = ctx.read().await;
+            let mut fetch_builder = consumer.fetch();
+            if let Some(max_messages) = max_messages {
+                fetch_builder = fetch_builder.max_messages(max_messages);
+            }
+            if let Some(group) = group {
+                fetch_builder = fetch_builder.group(group);
+            }
+            if let Some(priority) = priority {
+                fetch_builder = fetch_builder.priority(priority);
+            }
+            if let Some(max_bytes) = max_bytes {
+                fetch_builder = fetch_builder.max_bytes(max_bytes);
+            }
+            if let Some(heartbeat) = heartbeat {
+                fetch_builder = fetch_builder.heartbeat(heartbeat);
+            }
+            if let Some(expires) = expires {
+                fetch_builder = fetch_builder.expires(expires);
+            }
+            if let Some(min_pending) = min_pending {
+                fetch_builder = fetch_builder.min_pending(min_pending);
+            }
+            if let Some(min_ack_pending) = min_ack_pending {
+                fetch_builder = fetch_builder.min_ack_pending(min_ack_pending);
+            }
+            let mut messages = fetch_builder.messages().await?;
+            let mut ret_messages = Vec::new();
+            while let Some(msg) = messages.next().await {
+                let raw_msg = msg?;
+                ret_messages.push(crate::js::message::JetStreamMessage::from(raw_msg));
+            }
+            Ok(ret_messages)
+        })
+    }
+}

src/js/message.rs

@@ -1 +1,61 @@
+use pyo3::{Bound, Py, PyAny, Python, types::PyDict};
+use std::sync::Arc;
+use tokio::sync::RwLock;
 
+use crate::{
+    exceptions::rust_err::NatsrpyResult,
+    utils::{headers::NatsrpyHeadermapExt, natsrpy_future},
+};
+
+#[pyo3::pyclass]
+pub struct JetStreamMessage {
+    message: async_nats::Message,
+    headers: Option<Py<PyDict>>,
+    acker: Arc<RwLock<async_nats::jetstream::message::Acker>>,
+}
+
+impl From<async_nats::jetstream::Message> for JetStreamMessage {
+    fn from(value: async_nats::jetstream::Message) -> Self {
+        let (message, acker) = value.split();
+        Self {
+            message,
+            headers: None,
+            acker: Arc::new(RwLock::new(acker)),
+        }
+    }
+}
+
+#[pyo3::pymethods]
+impl JetStreamMessage {
+    #[getter]
+    pub fn subject(&self) -> &str {
+        self.message.subject.as_str()
+    }
+    #[getter]
+    pub fn reply(&self) -> Option<&str> {
+        self.message.reply.as_ref().map(async_nats::Subject::as_str)
+    }
+    #[getter]
+    pub fn payload(&self) -> &[u8] {
+        &self.message.payload
+    }
+    #[getter]
+    pub fn headers(&mut self, py: Python<'_>) -> NatsrpyResult<Py<PyDict>> {
+        if let Some(headers) = &self.headers {
+            Ok(headers.clone_ref(py))
+        } else {
+            let headermap = self.message.headers.clone().unwrap_or_default();
+            let headers = headermap.to_pydict(py)?.unbind();
+            self.headers = Some(headers.clone_ref(py));
+            Ok(headers)
+        }
+    }
+
+    pub fn ack<'py>(&self, py: Python<'py>) -> NatsrpyResult<Bound<'py, PyAny>> {
+        let acker_guard = self.acker.clone();
+        natsrpy_future(py, async move {
+            acker_guard.read().await.ack().await?;
+            Ok(())
+        })
+    }
+}

src/js/mod.rs

@@ -2,13 +2,16 @@ pub mod consumers;
 pub mod jetstream;
 pub mod kv;
 pub mod managers;
+pub mod message;
 pub mod stream;
 
 #[pyo3::pymodule(submodule, name = "js")]
 pub mod pymod {
     // Classes
     #[pymodule_export]
     pub use super::jetstream::JetStream;
+    #[pymodule_export]
+    pub use super::message::JetStreamMessage;
 
     // SubModules
     #[pymodule_export]

src/js/stream.rs

@@ -817,7 +817,7 @@ impl StreamMessage {
         Ok(Self {
             subject: msg.subject.to_string(),
             payload: PyBytes::new(py, &msg.payload).unbind(),
-            headers: msg.headers.to_pydict(py)?,
+            headers: msg.headers.to_pydict(py)?.unbind(),
             sequence: msg.sequence,
             time: time.unbind(),
         })

src/message.rs

@@ -23,7 +23,7 @@ impl TryFrom<async_nats::Message> for Message {
     fn try_from(value: async_nats::Message) -> Result<Self, Self::Error> {
         Python::attach(move |gil| {
             let headers = match value.headers {
-                Some(headermap) => headermap.to_pydict(gil)?,
+                Some(headermap) => headermap.to_pydict(gil)?.unbind(),
                 None => PyDict::new(gil).unbind(),
             };
             Ok(Self {

src/utils/headers.rs

@@ -1,13 +1,13 @@
 use pyo3::{
-    Bound, Py, Python,
+    Bound, Python,
     types::{PyAnyMethods, PyDict},
 };
 
 use crate::exceptions::rust_err::NatsrpyResult;
 
 pub trait NatsrpyHeadermapExt: Sized {
     fn from_pydict(pydict: Bound<PyDict>) -> NatsrpyResult<Self>;
-    fn to_pydict(&self, py: Python) -> NatsrpyResult<Py<PyDict>>;
+    fn to_pydict<'py>(&self, py: Python<'py>) -> NatsrpyResult<Bound<'py, PyDict>>;
 }
 
 impl NatsrpyHeadermapExt for async_nats::HeaderMap {
@@ -30,7 +30,7 @@ impl NatsrpyHeadermapExt for async_nats::HeaderMap {
         Ok(headermap)
     }
 
-    fn to_pydict(&self, py: Python) -> NatsrpyResult<Py<PyDict>> {
+    fn to_pydict<'py>(&self, py: Python<'py>) -> NatsrpyResult<Bound<'py, PyDict>> {
         let dict = PyDict::new(py);
         for (header_name, header_val) in self.iter() {
             let py_val = header_val
@@ -46,6 +46,6 @@ impl NatsrpyHeadermapExt for async_nats::HeaderMap {
             }
             dict.set_item(header_name.to_string(), py_val)?;
         }
-        Ok(dict.unbind())
+        Ok(dict)
     }
 }