Some jetstream updates.

Author: s3rius

Cargo.toml

@@ -13,7 +13,7 @@ async-nats = "0.46"
 bytes = "1.11.1"
 futures-util = "0.3.32"
 log = "0.4.29"
-pyo3 = "0.28"
+pyo3 = { version = "0.28", features = ["abi3"] }
 pyo3-async-runtimes = { version = "0.28", features = ["tokio-runtime"] }
 pyo3-log = "0.13.3"
 thiserror = "2.0.18"

src/jetstream/jetstream.rs

@@ -1,9 +1,10 @@
+use std::ops::Deref;
 use std::sync::Arc;
+use std::time::Duration;
 
 use async_nats::Subject;
 use async_nats::client::traits::Publisher;
 use async_nats::connection::State;
-use pyo3::Py;
 use pyo3::types::{PyBytesMethods, PyDict};
 use pyo3::{Bound, PyAny, Python, pyclass, pymethods, types::PyBytes};
 use tokio::sync::RwLock;
@@ -13,6 +14,44 @@ use crate::jetstream::kv::KeyValue;
 use crate::utils::headers::NatsrpyHeadermapExt;
 use crate::{exceptions::rust_err::NatsrpyResult, utils::natsrpy_future};
 
+#[pyclass]
+pub enum StorageType {
+    File,
+    Memory,
+}
+
+impl From<&StorageType> for async_nats::jetstream::stream::StorageType {
+    fn from(value: &StorageType) -> Self {
+        match value {
+            StorageType::File => Self::File,
+            StorageType::Memory => Self::Memory,
+        }
+    }
+}
+
+#[pyclass]
+pub struct Republish {
+    pub source: String,
+    pub destination: String,
+    pub headers_only: bool,
+}
+impl From<&Republish> for async_nats::jetstream::stream::Republish {
+    fn from(value: &Republish) -> Self {
+        Self {
+            source: value.source.clone(),
+            destination: value.destination.clone(),
+            headers_only: value.headers_only.clone(),
+        }
+    }
+}
+
+#[pyclass]
+pub struct Source {
+    pub name: String,
+    pub filter_subject: Option<String>,
+    pub external: bool,
+}
+
 #[pyclass]
 pub struct JetStream {
     ctx: Arc<RwLock<async_nats::jetstream::Context>>,
@@ -29,11 +68,11 @@ impl JetStream {
 #[pymethods]
 impl JetStream {
     #[pyo3(signature = (
-        subject, 
-        payload, 
-        *, 
-        headers=None, 
-        reply=None, 
+        subject,
+        payload,
+        *,
+        headers=None,
+        reply=None,
         err_on_disconnect = false
     ))]
     pub fn publish<'a>(
@@ -70,22 +109,58 @@ impl JetStream {
         &self,
         py: Python<'a>,
         bucket: String,
+        description: Option<String>,
+        max_value_size: Option<i32>,
+        history: Option<i64>,
+        max_age: Option<f32>,
+        max_bytes: Option<i64>,
+        storage: Option<Bound<'a, StorageType>>,
+        num_replicas: Option<usize>,
+        republish: Option<Bound<'a, Republish>>,
+        mirror: Option<Bound<'a, Source>>,
+        // sources: Option<Vec<Source>>,
+        mirror_direct: Option<bool>,
+        compression: Option<bool>,
+        // placement: Option<Place>,
+        limit_markers: Option<f32>,
     ) -> NatsrpyResult<Bound<'a, PyAny>> {
         let ctx = self.ctx.clone();
+        let storage = storage.map(|val| val.borrow().deref().into());
+        let republish = republish.map(|val| val.borrow().deref().into());
         natsrpy_future(py, async move {
             let js = ctx.read().await;
-            Ok(KeyValue::new(js.create_key_value(async_nats::jetstream::kv::Config{
-                bucket: bucket,
-                // todo!("Add other config options")
+            let mut config = async_nats::jetstream::kv::Config {
+                bucket: bucket.clone(),
                 ..Default::default()
-            }).await?))
+            };
+            description
+                .into_iter()
+                .for_each(|descr| config.description = descr);
+            max_value_size
+                .into_iter()
+                .for_each(|val| config.max_value_size = val);
+            history.into_iter().for_each(|val| config.history = val);
+            max_age
+                .into_iter()
+                .for_each(|val| config.max_age = Duration::from_secs_f32(val));
+            max_bytes.into_iter().for_each(|val| config.max_bytes = val);
+            num_replicas
+                .into_iter()
+                .for_each(|val| config.num_replicas = val);
+            mirror_direct
+                .into_iter()
+                .for_each(|val| config.mirror_direct = val);
+            compression
+                .into_iter()
+                .for_each(|val| config.compression = val);
+            storage.into_iter().for_each(|val| config.storage = val);
+            config.republish = republish;
+            config.limit_markers = limit_markers.map(Duration::from_secs_f32);
+
+            Ok(KeyValue::new(js.create_key_value(config).await?))
         })
     }
-    pub fn get_kv<'a>(
-        &self,
-        py: Python<'a>,
-        bucket: String,
-    ) -> NatsrpyResult<Bound<'a, PyAny>> {
+    pub fn get_kv<'a>(&self, py: Python<'a>, bucket: String) -> NatsrpyResult<Bound<'a, PyAny>> {
         let ctx = self.ctx.clone();
         natsrpy_future(py, async move {
             let js = ctx.read().await;

src/jetstream/mod.rs

@@ -9,4 +9,7 @@ use pyo3::pymodule;
 pub mod pymod {
     #[pymodule_export]
     use super::jetstream::JetStream;
+
+    #[pymodule_export]
+    use super::jetstream::StorageType;
 }

src/nats_cls.rs

@@ -205,11 +205,11 @@ impl NatsCls {
     }
 
     #[pyo3(signature = (
-        *, 
-        domain=None, 
+        *,
+        domain=None,
         api_prefix=None,
-        timeout=None, 
-        ack_timeout=None, 
+        timeout=None,
+        ack_timeout=None,
         concurrency_limit = None,
         max_ack_inflight=None,
         backpressure_on_inflight=None,
@@ -247,11 +247,13 @@ impl NatsCls {
                 builder = builder.backpressure_on_inflight(backpressure_on_inflight);
             }
             if domain.is_some() && api_prefix.is_some() {
-                return Err(NatsrpyError::InvalidArgument(String::from("Either domain or api_prefix should be specified, not both.")));
+                return Err(NatsrpyError::InvalidArgument(String::from(
+                    "Either domain or api_prefix should be specified, not both.",
+                )));
             }
             let js = if let Some(api_prefix) = api_prefix {
                 builder.api_prefix(api_prefix).build(session.clone())
-            } else if let Some(domain) = domain{
+            } else if let Some(domain) = domain {
                 builder.domain(domain).build(session.clone())
             } else {
                 builder.build(session.clone())

src/subscription.rs

@@ -4,7 +4,7 @@ use std::{sync::Arc, time::Duration};
 
 use pyo3::Py;
 use pyo3::types::PyBytes;
-use pyo3::{Bound, IntoPyObjectExt, PyAny, PyRef, Python, pyclass, pymethods};
+use pyo3::{Bound, PyAny, PyRef, Python, pyclass, pymethods};
 use tokio::sync::Mutex;
 
 use crate::exceptions::rust_err::NatsrpyError;