natsrpy/src/js/managers/counters.rs at d4bc5a4f5f0ddb8b97b5a73374eefb472e42224d · taskiq-python/natsrpy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
use std::sync::Arc;

use crate::{
    exceptions::rust_err::NatsrpyError,
    js::counters::{Counters, CountersConfig},
};
use pyo3::{Bound, PyAny, Python};
use tokio::sync::RwLock;

use crate::{exceptions::rust_err::NatsrpyResult, utils::natsrpy_future};

#[pyo3::pyclass]
pub struct CountersManager {
    ctx: Arc<RwLock<async_nats::jetstream::Context>>,
}

impl CountersManager {
    pub const fn new(ctx: Arc<RwLock<async_nats::jetstream::Context>>) -> Self {
        Self { ctx }
    }
}

#[pyo3::pymethods]
impl CountersManager {
    pub fn create<'py>(
        &self,
        py: Python<'py>,
        config: CountersConfig,
    ) -> NatsrpyResult<Bound<'py, PyAny>> {
        let ctx = self.ctx.clone();
        natsrpy_future(py, async move {
            let js = ctx.read().await;
            Ok(Counters::new(
                js.create_stream(async_nats::jetstream::stream::Config::try_from(config)?)
                    .await?,
                ctx.clone(),
            ))
        })
    }

    pub fn create_or_update<'py>(
        &self,
        py: Python<'py>,
        config: CountersConfig,
    ) -> NatsrpyResult<Bound<'py, PyAny>> {
        let ctx = self.ctx.clone();
        natsrpy_future(py, async move {
            let info = ctx
                .read()
                .await
                .create_or_update_stream(async_nats::jetstream::stream::Config::try_from(config)?)
                .await?;
            Ok(Counters::new(
                ctx.read().await.get_stream(info.config.name).await?,
                ctx.clone(),
            ))
        })
    }

    pub fn get<'py>(&self, py: Python<'py>, name: String) -> NatsrpyResult<Bound<'py, PyAny>> {
        let ctx = self.ctx.clone();
        natsrpy_future(py, async move {
            let stream = ctx.read().await.get_stream(&name).await?;
            let config = stream.get_info().await?.config;
            if !config.allow_direct {
                return Err(NatsrpyError::SessionError(format!(
                    "Stream {name} doesn't allow direct get.",
                )));
            }
            if !config.allow_message_counter {
                return Err(NatsrpyError::SessionError(format!(
                    "Stream {name} doesn't allow message counters.",
                )));
            }
            Ok(Counters::new(stream, ctx.clone()))
        })
    }

    pub fn delete<'py>(&self, py: Python<'py>, name: String) -> NatsrpyResult<Bound<'py, PyAny>> {
        let ctx = self.ctx.clone();
        natsrpy_future(py, async move {
            let js = ctx.read().await;
            Ok(js.delete_stream(name).await?.success)
        })
    }

    pub fn update<'py>(
        &self,
        py: Python<'py>,
        config: CountersConfig,
    ) -> NatsrpyResult<Bound<'py, PyAny>> {
        let ctx = self.ctx.clone();
        natsrpy_future(py, async move {
            let info = ctx
                .read()
                .await
                .update_stream(async_nats::jetstream::stream::Config::try_from(config)?)
                .await?;
            Ok(Counters::new(
                ctx.read().await.get_stream(info.config.name).await?,
                ctx.clone(),
            ))
        })
    }
}