Subscription updates.

Author: s3rius

python/natsrpy/_natsrpy_rs/__init__.pyi

@@ -1,3 +1,4 @@
+from collections.abc import Awaitable, Callable
 from datetime import timedelta
 from typing import Any
 
@@ -39,7 +40,11 @@ class Nats:
     async def request(self, subject: str, payload: bytes) -> None: ...
     async def drain(self) -> None: ...
     async def flush(self) -> None: ...
-    async def subscribe(self, subject: str) -> Subscription: ...
+    async def subscribe(
+        self,
+        subject: str,
+        callback: Callable[[Message], Awaitable[None]] | None = None,
+    ) -> Subscription: ...
     async def jetstream(self) -> JetStream: ...
 
 __all__ = ["Message", "Nats", "Subscription"]

src/message.rs

@@ -17,12 +17,12 @@ pub struct Message {
     pub length: usize,
 }
 
-impl TryFrom<async_nats::Message> for Message {
+impl TryFrom<&async_nats::Message> for Message {
     type Error = NatsrpyError;
 
-    fn try_from(value: async_nats::Message) -> Result<Self, Self::Error> {
+    fn try_from(value: &async_nats::Message) -> Result<Self, Self::Error> {
         Python::attach(move |gil| {
-            let headers = match value.headers {
+            let headers = match &value.headers {
                 Some(headermap) => headermap.to_pydict(gil)?.unbind(),
                 None => PyDict::new(gil).unbind(),
             };
@@ -32,13 +32,21 @@ impl TryFrom<async_nats::Message> for Message {
                 payload: PyBytes::new(gil, &value.payload).unbind(),
                 headers,
                 status: value.status.map(Into::<u16>::into),
-                description: value.description,
+                description: value.description.clone(),
                 length: value.length,
             })
         })
     }
 }
 
+impl TryFrom<async_nats::Message> for Message {
+    type Error = NatsrpyError;
+
+    fn try_from(value: async_nats::Message) -> Result<Self, Self::Error> {
+        Self::try_from(&value)
+    }
+}
+
 #[pyo3::pymethods]
 impl Message {
     #[must_use]

src/nats_cls.rs

@@ -1,6 +1,6 @@
 use async_nats::{Subject, client::traits::Publisher, message::OutboundMessage};
 use pyo3::{
-    Bound, PyAny, Python,
+    Bound, Py, PyAny, Python,
     types::{PyBytes, PyBytesMethods, PyDict},
 };
 use std::{sync::Arc, time::Duration};
@@ -197,16 +197,21 @@ impl NatsCls {
         })
     }
 
+    #[pyo3(signature=(subject, callback=None))]
     pub fn subscribe<'py>(
         &self,
         py: Python<'py>,
         subject: String,
+        callback: Option<Py<PyAny>>,
     ) -> NatsrpyResult<Bound<'py, PyAny>> {
         log::debug!("Subscribing to '{subject}'");
         let session = self.nats_session.clone();
         natsrpy_future(py, async move {
             if let Some(session) = session.read().await.as_ref() {
-                Ok(Subscription::new(session.subscribe(subject).await?))
+                Ok(Subscription::new(
+                    session.subscribe(subject).await?,
+                    callback,
+                )?)
             } else {
                 Err(NatsrpyError::NotInitialized)
             }

src/subscription.rs

@@ -1,7 +1,7 @@
 use futures_util::StreamExt;
 use std::sync::Arc;
 
-use pyo3::{Bound, PyAny, PyRef, Python};
+use pyo3::{Bound, Py, PyAny, PyRef, Python};
 use tokio::sync::Mutex;
 
 use crate::{
@@ -12,14 +12,56 @@ use crate::{
 #[pyo3::pyclass]
 pub struct Subscription {
     inner: Option<Arc<Mutex<async_nats::Subscriber>>>,
+    reading_task: Option<tokio::task::AbortHandle>,
+}
+
+async fn process_message(message: async_nats::message::Message, py_callback: Py<PyAny>) {
+    let task = async || -> NatsrpyResult<()> {
+        let message = crate::message::Message::try_from(&message)?;
+        let awaitable = Python::attach(|gil| -> NatsrpyResult<_> {
+            let res = py_callback.call1(gil, (message,))?;
+            let rust_task = pyo3_async_runtimes::tokio::into_future(res.into_bound(gil))?;
+            Ok(rust_task)
+        })?;
+        awaitable.await?;
+        Ok(())
+    };
+    if let Err(err) = task().await {
+        log::error!("Cannot process message {message:?}. Error: {err}");
+    }
+}
+
+async fn start_py_sub(
+    sub: Arc<Mutex<async_nats::Subscriber>>,
+    py_callback: Py<PyAny>,
+    locals: pyo3_async_runtimes::TaskLocals,
+) {
+    while let Some(message) = sub.lock().await.next().await {
+        let py_cb = Python::attach(|py| py_callback.clone_ref(py));
+        tokio::spawn(pyo3_async_runtimes::tokio::scope(
+            locals.clone(),
+            process_message(message, py_cb),
+        ));
+    }
 }
 
 impl Subscription {
-    #[must_use]
-    pub fn new(sub: async_nats::Subscriber) -> Self {
-        Self {
-            inner: Some(Arc::new(Mutex::new(sub))),
-        }
+    pub fn new(sub: async_nats::Subscriber, callback: Option<Py<PyAny>>) -> NatsrpyResult<Self> {
+        let sub = Arc::new(Mutex::new(sub));
+        let cb_sub = sub.clone();
+        let task_locals = Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?;
+        let task_handle = callback.map(move |cb| {
+            tokio::task::spawn(pyo3_async_runtimes::tokio::scope(
+                task_locals.clone(),
+                start_py_sub(cb_sub, cb, task_locals),
+            ))
+            .abort_handle()
+        });
+
+        Ok(Self {
+            inner: Some(sub),
+            reading_task: task_handle,
+        })
     }
 }
 
@@ -38,11 +80,15 @@ impl Subscription {
         let Some(inner) = self.inner.clone() else {
             unreachable!("Subscription used after del")
         };
+        if self.reading_task.is_some() {
+            log::warn!(
+                "Callback is set. Getting messages from this subscription might produce unpredictable results."
+            );
+        }
         natsrpy_future_with_timeout(py, timeout, async move {
             let Some(message) = inner.lock().await.next().await else {
                 return Err(NatsrpyError::AsyncStopIteration);
             };
-
             crate::message::Message::try_from(message)
         })
     }
@@ -96,6 +142,9 @@ impl Drop for Subscription {
     fn drop(&mut self) {
         pyo3_async_runtimes::tokio::get_runtime().block_on(async move {
             self.inner = None;
+            if let Some(reading) = self.reading_task.take() {
+                reading.abort();
+            }
         });
     }
 }